wehub-resource-sync
749 lines
28 KiB
Python
749 lines
28 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, F401, F841
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.ir import assert_structural_equal
|
|
from tvm.relax.frontend import nn
|
|
from tvm.relax.frontend.nn import core, modules, spec
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_relu():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((3, 3), dtype="float32"),
|
|
_io: R.Any,
|
|
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
relu: R.Tensor((3, 3), dtype="float32") = R.nn.relu(x)
|
|
gv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)) = relu, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.ReLU()
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((3, 3), "float32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_silu():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((3, 3), dtype="float32"),
|
|
_io: R.Any,
|
|
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
silu: R.Tensor((3, 3), dtype="float32") = R.nn.silu(x)
|
|
gv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)) = silu, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.SiLU()
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((3, 3), "float32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_gelu():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((3, 3), dtype="float32"),
|
|
_io: R.Any,
|
|
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
gelu: R.Tensor((3, 3), dtype="float32") = R.nn.gelu(x)
|
|
gv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)) = gelu, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.GELU()
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((3, 3), "float32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_identity():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((3, 3), dtype="float32"),
|
|
_io: R.Any,
|
|
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
gv1: R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tuple(R.Any)) = x, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Identity()
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((3, 3), "float32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_linear():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((1, 4), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((8, 4), dtype="float32"),
|
|
bias: R.Tensor((8,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 8), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
permute_dims: R.Tensor((4, 8), dtype="float32") = R.permute_dims(weight, axes=None)
|
|
matmul: R.Tensor((1, 8), dtype="float32") = R.matmul(x, permute_dims)
|
|
add: R.Tensor((1, 8), dtype="float32") = R.add(matmul, bias)
|
|
gv1: R.Tuple(R.Tensor((1, 8), dtype="float32"), R.Tuple(R.Any)) = add, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Linear(4, 8)
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((1, 4), "float32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_conv1d():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((1, 3, 32), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((32, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((32,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 32, 30), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 32, 30), dtype="float32") = R.nn.conv1d(
|
|
x,
|
|
weight,
|
|
strides=[1],
|
|
padding=[0, 0],
|
|
dilation=[1],
|
|
groups=1,
|
|
data_layout="NCW",
|
|
kernel_layout="OIW",
|
|
out_layout="NCW",
|
|
)
|
|
lv2: R.Tensor((1, 32, 1), dtype="float32") = R.reshape(bias, R.shape([1, 32, 1]))
|
|
conv1d: R.Tensor((1, 32, 30), dtype="float32") = R.add(lv1, lv2)
|
|
gv1: R.Tuple(R.Tensor((1, 32, 30), dtype="float32"), R.Tuple(R.Any)) = conv1d, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Conv1D(3, 32, 3, bias=True)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": spec.Tensor([1, 3, 32], "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_conv1d_transpose():
|
|
# fmt: off
|
|
@R.function
|
|
def forward(x: R.Tensor((1, 3, 30), dtype="float32"), _io: R.Any, weight: R.Tensor((3, 32, 3), dtype="float32"), bias: R.Tensor((32,), dtype="float32")) -> R.Tuple(R.Tensor((1, 32, 32), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 32, 32), dtype="float32") = R.nn.conv1d_transpose(x, weight, strides=[1], padding=[0, 0], output_padding=[0], dilation=[1], groups=1, data_layout="NCW", kernel_layout="IOW", out_layout="NCW")
|
|
lv2: R.Tensor((1, 32, 1), dtype="float32") = R.reshape(bias, R.shape([1, 32, 1]))
|
|
conv1d_transpose: R.Tensor((1, 32, 32), dtype="float32") = R.add(lv1, lv2)
|
|
gv1: R.Tuple(R.Tensor((1, 32, 32), dtype="float32"), R.Tuple(R.Any)) = conv1d_transpose, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
# fmt: on
|
|
|
|
mod = modules.ConvTranspose1D(3, 32, 3, bias=True)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": spec.Tensor([1, 3, 30], "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_layer_norm():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((2, 4, 8), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((8,), dtype="float32"),
|
|
bias: R.Tensor((8,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
layer_norm: R.Tensor((2, 4, 8), dtype="float32") = R.nn.layer_norm(
|
|
x, weight, bias, axes=[-1], epsilon=1.0000000000000001e-05, center=True, scale=True
|
|
)
|
|
gv1: R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Any)) = (
|
|
layer_norm,
|
|
(_io,),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.LayerNorm(8)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={"forward": {"x": spec.Tensor((2, 4, 8), "float32")}}, debug=True
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_conv2d():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((1, 3, 32, 32), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((32, 3, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((32,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 32, 30, 30), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 32, 30, 30), dtype="float32") = R.nn.conv2d(x, weight)
|
|
lv2: R.Tensor((1, 32, 1, 1), dtype="float32") = R.reshape(bias, R.shape([1, 32, 1, 1]))
|
|
conv2d: R.Tensor((1, 32, 30, 30), dtype="float32") = R.add(lv1, lv2)
|
|
gv1: R.Tuple(R.Tensor((1, 32, 30, 30), dtype="float32"), R.Tuple(R.Any)) = (
|
|
conv2d,
|
|
(_io,),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Conv2D(3, 32, 3, bias=True)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": spec.Tensor([1, 3, 32, 32], "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_conv3d():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((1, 3, 32, 32, 32), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((32, 3, 3, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((32,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 32, 30, 30, 30), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((1, 32, 30, 30, 30), dtype="float32") = R.nn.conv3d(x, weight)
|
|
lv2: R.Tensor((1, 32, 1, 1, 1), dtype="float32") = R.reshape(
|
|
bias, R.shape([1, 32, 1, 1, 1])
|
|
)
|
|
conv3d: R.Tensor((1, 32, 30, 30, 30), dtype="float32") = R.add(lv1, lv2)
|
|
gv1: R.Tuple(R.Tensor((1, 32, 30, 30, 30), dtype="float32"), R.Tuple(R.Any)) = (
|
|
conv3d,
|
|
(_io,),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Conv3D(3, 32, 3, bias=True)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": spec.Tensor([1, 3, 32, 32, 32], "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_conv2d_dynamic():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor(("n", "c", "h", "w"), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((32, "in_channels", 3, 3), dtype="float32"),
|
|
bias: R.Tensor((32,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor(("n", 32, "h - 2", "w - 2"), dtype="float32"), R.Tuple(R.Any)):
|
|
n = T.int64()
|
|
h = T.int64()
|
|
w = T.int64()
|
|
c = T.int64()
|
|
in_channels = T.int64()
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((n, 32, h - 2, w - 2), dtype="float32") = R.nn.conv2d(x, weight)
|
|
lv2: R.Tensor((1, 32, 1, 1), dtype="float32") = R.reshape(bias, R.shape([1, 32, 1, 1]))
|
|
conv2d: R.Tensor((n, 32, h - 2, w - 2), dtype="float32") = R.add(lv1, lv2)
|
|
gv1: R.Tuple(R.Tensor((n, 32, h - 2, w - 2), dtype="float32"), R.Tuple(R.Any)) = (
|
|
conv2d,
|
|
(_io,),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Conv2D(tvm.tirx.Var("in_channels", "int64"), 32, 3, bias=True)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": spec.Tensor(["n", "c", "h", "w"], "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_rms_norm():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((2, 4, 8), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((8,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
rms_norm: R.Tensor((2, 4, 8), dtype="float32") = R.nn.rms_norm(
|
|
x, weight, axes=[2], epsilon=1.0000000000000001e-05
|
|
)
|
|
gv1: R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Any)) = rms_norm, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.RMSNorm(8, [2], bias=False)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={"forward": {"x": spec.Tensor((2, 4, 8), "float32")}}, debug=True
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_group_norm():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((2, 4, 8), dtype="float32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((4,), dtype="float32"),
|
|
bias: R.Tensor((4,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
group_norm: R.Tensor((2, 4, 8), dtype="float32") = R.nn.group_norm(
|
|
x, weight, bias, num_groups=2, channel_axis=1, axes=[2]
|
|
)
|
|
gv1: R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Any)) = (
|
|
group_norm,
|
|
(_io,),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.GroupNorm(num_groups=2, num_channels=4)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={"forward": {"x": spec.Tensor((2, 4, 8), "float32")}}, debug=True
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_embedding_1d():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((4,), dtype="int32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((8, 16), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((4, 16), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
take: R.Tensor((4, 16), dtype="float32") = R.take(weight, x, axis=0)
|
|
gv1: R.Tuple(R.Tensor((4, 16), dtype="float32"), R.Tuple(R.Any)) = take, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Embedding(8, 16, "float32")
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((4,), "int32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_embedding_2d():
|
|
@R.function
|
|
def forward(
|
|
x: R.Tensor((1, 4), dtype="int32"),
|
|
_io: R.Any,
|
|
weight: R.Tensor((4, 8), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((1, 4, 8), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
reshape: R.Tensor((4,), dtype="int32") = R.reshape(x, R.shape([4]))
|
|
take: R.Tensor((4, 8), dtype="float32") = R.take(weight, reshape, axis=0)
|
|
reshape1: R.Tensor((1, 4, 8), dtype="float32") = R.reshape(take, R.shape([1, 4, 8]))
|
|
gv1: R.Tuple(R.Tensor((1, 4, 8), dtype="float32"), R.Tuple(R.Any)) = reshape1, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Embedding(4, 8, "float32")
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((1, 4), "int32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_timestep_embedding():
|
|
@R.function
|
|
def forward(
|
|
sample: R.Tensor((32, 32), dtype="float32"),
|
|
condition: R.Tensor((32, 16), dtype="float32"),
|
|
_io: R.Any,
|
|
linear_1_weight: R.Tensor((32, 32), dtype="float32"),
|
|
linear_1_bias: R.Tensor((32,), dtype="float32"),
|
|
cond_proj_weight: R.Tensor((32, 16), dtype="float32"),
|
|
linear_2_weight: R.Tensor((32, 32), dtype="float32"),
|
|
linear_2_bias: R.Tensor((32,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((32, 32), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 3})
|
|
with R.dataflow():
|
|
permute_dims: R.Tensor((16, 32), dtype="float32") = R.permute_dims(
|
|
cond_proj_weight, axes=None
|
|
)
|
|
matmul: R.Tensor((32, 32), dtype="float32") = R.matmul(condition, permute_dims)
|
|
add: R.Tensor((32, 32), dtype="float32") = R.add(sample, matmul)
|
|
permute_dims1: R.Tensor((32, 32), dtype="float32") = R.permute_dims(
|
|
linear_1_weight, axes=None
|
|
)
|
|
matmul1: R.Tensor((32, 32), dtype="float32") = R.matmul(add, permute_dims1)
|
|
add1: R.Tensor((32, 32), dtype="float32") = R.add(matmul1, linear_1_bias)
|
|
silu: R.Tensor((32, 32), dtype="float32") = R.nn.silu(add1)
|
|
permute_dims2: R.Tensor((32, 32), dtype="float32") = R.permute_dims(
|
|
linear_2_weight, axes=None
|
|
)
|
|
matmul2: R.Tensor((32, 32), dtype="float32") = R.matmul(silu, permute_dims2)
|
|
add2: R.Tensor((32, 32), dtype="float32") = R.add(matmul2, linear_2_bias)
|
|
gv1: R.Tuple(R.Tensor((32, 32), dtype="float32"), R.Tuple(R.Any)) = add2, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.TimestepEmbedding(32, 32, cond_proj_dim=16)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"sample": spec.Tensor((32, 32), "float32"),
|
|
"condition": spec.Tensor((32, 16), "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_timesteps():
|
|
@R.function
|
|
def forward(x: R.Tensor((3,), dtype="float32"), _io: R.Any) -> R.Tuple(
|
|
R.Tensor((3, 10), dtype="float32"), R.Tuple(R.Any)
|
|
):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((3,), dtype="float32") = R.astype(x, dtype="float32")
|
|
lv2: R.Tensor((3, 1), dtype="float32") = R.expand_dims(lv1, axis=[1])
|
|
lv3: R.Tensor((5,), dtype="float32") = R.arange(
|
|
R.prim_value(0), R.prim_value(5), R.prim_value(1), dtype="float32"
|
|
)
|
|
lv4: R.Tensor((5,), dtype="float32") = R.multiply(
|
|
R.const(-9.2103404998779297, "float32"), lv3
|
|
)
|
|
lv5: R.Tensor((5,), dtype="float32") = R.divide(lv4, R.const(4, "float32"))
|
|
lv6: R.Tensor((5,), dtype="float32") = R.exp(lv5)
|
|
lv7: R.Tensor((1, 5), dtype="float32") = R.expand_dims(lv6, axis=[0])
|
|
lv8: R.Tensor((3, 5), dtype="float32") = R.multiply(lv2, lv7)
|
|
lv9: R.Tensor((3, 5), dtype="float32") = R.sin(lv8)
|
|
lv10: R.Tensor((3, 5), dtype="float32") = R.cos(lv8)
|
|
lv11: R.Tensor((3, 10), dtype="float32") = R.concat((lv9, lv10), axis=-1)
|
|
get_timestep_embedding: R.Tensor((3, 10), dtype="float32") = R.astype(
|
|
lv11, dtype="float32"
|
|
)
|
|
gv1: R.Tuple(R.Tensor((3, 10), dtype="float32"), R.Tuple(R.Any)) = (
|
|
get_timestep_embedding,
|
|
(_io,),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Timesteps(10)
|
|
tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((3,), "float32")}}, debug=True)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_kv_cache():
|
|
@I.ir_module
|
|
class Module:
|
|
@R.function
|
|
def _initialize_effect() -> R.Tuple(R.Any, R.Any):
|
|
with R.dataflow():
|
|
_io: R.Any = R.null_value()
|
|
lv: R.Tensor((8, 2, 4), dtype="float32") = R.zeros(
|
|
R.shape([8, 2, 4]), dtype="float32"
|
|
)
|
|
cache: R.Any = R.call_pure_packed(
|
|
"vm.builtin.attention_kv_cache_create",
|
|
lv,
|
|
R.shape([8, 2, 4]),
|
|
R.prim_value(0),
|
|
ty_args=[R.Any()],
|
|
)
|
|
lv1 = _io, cache
|
|
gv = lv1
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def forward(x: R.Tensor((2, 4), dtype="float32"), _io: R.Any, cache: R.Any) -> R.Tuple(
|
|
R.Tensor((4, 2, 4), dtype="float32"), R.Tuple(R.Any, R.Any)
|
|
):
|
|
R.func_attr({"num_input": 3})
|
|
with R.dataflow():
|
|
lv2: R.Any = R.call_inplace_packed(
|
|
"vm.builtin.attention_kv_cache_append",
|
|
cache,
|
|
x,
|
|
inplace_indices=[0],
|
|
ty_args=[R.Any()],
|
|
)
|
|
lv3: R.Tensor((4, 2, 4), dtype="float32") = R.call_pure_packed(
|
|
"vm.builtin.attention_kv_cache_view",
|
|
lv2,
|
|
R.shape([4, 2, 4]),
|
|
ty_args=(R.Tensor((4, 2, 4), dtype="float32"),),
|
|
)
|
|
gv1: R.Tuple(R.Tensor((4, 2, 4), dtype="float32"), R.Tuple(R.Any, R.Any)) = (
|
|
lv3,
|
|
(_io, lv2),
|
|
)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
class KVCacheTest(modules.Module):
|
|
def __init__(self) -> None:
|
|
self.cache = modules.KVCache(8, [2, 4])
|
|
|
|
def forward(self, x: core.Tensor) -> core.Tensor:
|
|
self.cache.append(x)
|
|
return self.cache.view(4)
|
|
|
|
tvm_mod, _ = KVCacheTest().export_tvm(
|
|
spec={"forward": {"x": spec.Tensor((2, 4), "float32")}}, debug=True
|
|
)
|
|
assert_structural_equal(tvm_mod, Module, True)
|
|
|
|
|
|
def test_attention():
|
|
@R.function
|
|
def forward(
|
|
hidden_states: R.Tensor((2, 4096, 640), dtype="float32"),
|
|
encoder_hidden_states: R.Tensor((2, 77, 2048), dtype="float32"),
|
|
_io: R.Any,
|
|
to_q_weight: R.Tensor((640, 640), dtype="float32"),
|
|
to_k_weight: R.Tensor((640, 2048), dtype="float32"),
|
|
to_v_weight: R.Tensor((640, 2048), dtype="float32"),
|
|
group_norm_weight: R.Tensor((640,), dtype="float32"),
|
|
group_norm_bias: R.Tensor((640,), dtype="float32"),
|
|
to_out_0_weight: R.Tensor((640, 640), dtype="float32"),
|
|
to_out_0_bias: R.Tensor((640,), dtype="float32"),
|
|
) -> R.Tuple(R.Tensor((2, 4096, 640), dtype="float32"), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 3})
|
|
with R.dataflow():
|
|
group_norm: R.Tensor((2, 4096, 640), dtype="float32") = R.nn.group_norm(
|
|
hidden_states,
|
|
group_norm_weight,
|
|
group_norm_bias,
|
|
num_groups=8,
|
|
channel_axis=2,
|
|
axes=[1],
|
|
epsilon=1.0000000000000001e-05,
|
|
center=True,
|
|
scale=True,
|
|
)
|
|
permute_dims: R.Tensor((640, 640), dtype="float32") = R.permute_dims(
|
|
to_q_weight, axes=None
|
|
)
|
|
matmul: R.Tensor((2, 4096, 640), dtype="float32") = R.matmul(group_norm, permute_dims)
|
|
permute_dims1: R.Tensor((2048, 640), dtype="float32") = R.permute_dims(
|
|
to_k_weight, axes=None
|
|
)
|
|
matmul1: R.Tensor((2, 77, 640), dtype="float32") = R.matmul(
|
|
encoder_hidden_states, permute_dims1
|
|
)
|
|
permute_dims2: R.Tensor((2048, 640), dtype="float32") = R.permute_dims(
|
|
to_v_weight, axes=None
|
|
)
|
|
matmul2: R.Tensor((2, 77, 640), dtype="float32") = R.matmul(
|
|
encoder_hidden_states, permute_dims2
|
|
)
|
|
reshape: R.Tensor((2, 4096, 10, 64), dtype="float32") = R.reshape(
|
|
matmul, R.shape([2, 4096, 10, 64])
|
|
)
|
|
reshape1: R.Tensor((2, 77, 10, 64), dtype="float32") = R.reshape(
|
|
matmul1, R.shape([2, 77, 10, 64])
|
|
)
|
|
reshape2: R.Tensor((2, 77, 10, 64), dtype="float32") = R.reshape(
|
|
matmul2, R.shape([2, 77, 10, 64])
|
|
)
|
|
scaled_dot_product_attention: R.Tensor((2, 4096, 10, 64), dtype="float32") = (
|
|
R.nn.attention(reshape, reshape1, reshape2, scale=None, causal_mask=None)
|
|
)
|
|
reshape3: R.Tensor((2, 4096, 640), dtype="float32") = R.reshape(
|
|
scaled_dot_product_attention, R.shape([2, 4096, 640])
|
|
)
|
|
permute_dims3: R.Tensor((640, 640), dtype="float32") = R.permute_dims(
|
|
to_out_0_weight, axes=None
|
|
)
|
|
matmul3: R.Tensor((2, 4096, 640), dtype="float32") = R.matmul(reshape3, permute_dims3)
|
|
add: R.Tensor((2, 4096, 640), dtype="float32") = R.add(matmul3, to_out_0_bias)
|
|
gv1: R.Tuple(R.Tensor((2, 4096, 640), dtype="float32"), R.Tuple(R.Any)) = add, (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
|
|
mod = modules.Attention(query_dim=640, cross_attention_dim=2048, heads=10, norm_num_groups=8)
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"hidden_states": spec.Tensor((2, 4096, 640), "float32"),
|
|
"encoder_hidden_states": spec.Tensor((2, 77, 2048), "float32"),
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
assert_structural_equal(tvm_mod["forward"], forward, True)
|
|
|
|
|
|
def test_nn_module_tuple_input():
|
|
class Layer(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def forward(self, x: tuple[nn.Tensor, nn.Tensor]):
|
|
x0 = x[0]
|
|
x1 = x[1]
|
|
y0 = nn.add(x0, x1)
|
|
y1 = nn.subtract(x0, x1)
|
|
return (y0, y1)
|
|
|
|
# fmt: off
|
|
@R.function
|
|
def forward(x: R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((10, 5), dtype="float32")), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((10, 5), dtype="float32")), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((10, 5), dtype="float32") = x[0]
|
|
lv2: R.Tensor((10, 5), dtype="float32") = x[1]
|
|
add: R.Tensor((10, 5), dtype="float32") = R.add(lv1, lv2)
|
|
subtract: R.Tensor((10, 5), dtype="float32") = R.subtract(lv1, lv2)
|
|
gv1: R.Tuple(R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((10, 5), dtype="float32")), R.Tuple(R.Any)) = (add, subtract), (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
# fmt: on
|
|
|
|
mod = Layer()
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": (spec.Tensor([10, 5], dtype="float32"), spec.Tensor([10, 5], dtype="float32"))
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
|
|
assert_structural_equal(tvm_mod["forward"], forward)
|
|
|
|
|
|
def test_nn_module_list_input():
|
|
class Layer(nn.Module):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def forward(self, x: list[nn.Tensor]):
|
|
x0 = x[0]
|
|
x1 = x[1]
|
|
y0 = nn.add(x0, x1)
|
|
y1 = nn.subtract(x0, x1)
|
|
return [y0, y1]
|
|
|
|
# fmt: off
|
|
@R.function
|
|
def forward(x: R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((10, 5), dtype="float32")), _io: R.Any) -> R.Tuple(R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((10, 5), dtype="float32")), R.Tuple(R.Any)):
|
|
R.func_attr({"num_input": 2})
|
|
with R.dataflow():
|
|
lv1: R.Tensor((10, 5), dtype="float32") = x[0]
|
|
lv2: R.Tensor((10, 5), dtype="float32") = x[1]
|
|
add: R.Tensor((10, 5), dtype="float32") = R.add(lv1, lv2)
|
|
subtract: R.Tensor((10, 5), dtype="float32") = R.subtract(lv1, lv2)
|
|
gv1: R.Tuple(R.Tuple(R.Tensor((10, 5), dtype="float32"), R.Tensor((10, 5), dtype="float32")), R.Tuple(R.Any)) = (add, subtract), (_io,)
|
|
R.output(gv1)
|
|
return gv1
|
|
# fmt: on
|
|
|
|
mod = Layer()
|
|
tvm_mod, _ = mod.export_tvm(
|
|
spec={
|
|
"forward": {
|
|
"x": [spec.Tensor([10, 5], dtype="float32"), spec.Tensor([10, 5], dtype="float32")]
|
|
}
|
|
},
|
|
debug=True,
|
|
)
|
|
|
|
assert_structural_equal(tvm_mod["forward"], forward)
|
|
|
|
|
|
def test_module_list():
|
|
class Module(nn.Module):
|
|
def __init__(self):
|
|
self.layers = nn.ModuleList(
|
|
[nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(2)]) for _ in range(1)]
|
|
)
|
|
|
|
def forward(self, x: nn.Tensor):
|
|
return self.layers(x)
|
|
|
|
mod = Module()
|
|
named_params = dict(mod.named_parameters())
|
|
assert ["layers.0.0.weight", "layers.0.1.weight"] == sorted(list(named_params.keys()))
|
|
|
|
|
|
def test_module_dict():
|
|
class Module(nn.Module):
|
|
def __init__(self):
|
|
self.layers = nn.ModuleDict(
|
|
{"linear0": nn.Linear(4, 4, bias=False), "linear1": nn.Linear(4, 4, bias=False)}
|
|
)
|
|
|
|
def forward(self, x: nn.Tensor):
|
|
x = self.layers["linear0"](x)
|
|
x = self.layers["linear1"](x)
|
|
return x
|
|
|
|
mod = Module()
|
|
named_params = dict(mod.named_parameters())
|
|
assert ["layers.linear0.weight", "layers.linear1.weight"] == sorted(list(named_params.keys()))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|