wehub-resource-sync
1139 lines
46 KiB
Python
1139 lines
46 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.
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.relax.transform import ToMixedPrecision
|
|
from tvm.script.parser import ir as I
|
|
from tvm.script.parser import relax as R
|
|
from tvm.script.parser import tirx as T
|
|
|
|
|
|
def _assert_test(input, expected=None, expected2=None):
|
|
if expected:
|
|
mod = ToMixedPrecision()(input)
|
|
tvm.ir.assert_structural_equal(mod, expected)
|
|
|
|
if expected2:
|
|
mod = ToMixedPrecision(out_dtype="float16")(input)
|
|
tvm.ir.assert_structural_equal(mod, expected2)
|
|
|
|
|
|
def test_conv2d():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3), "float32")
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float32") = R.astype(lv2, dtype="float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_conv2d_relu():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3), "float32")
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
|
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(lv)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv_1: R.Tensor((2, 4, 26, 26), dtype="float16") = R.astype(lv2, dtype="float16")
|
|
lv3: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.relu(lv_1)
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv_1: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
lv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.relu(lv_1)
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float32") = R.astype(lv2, dtype="float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_unknown_dtype_is_not_rewritten():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype=None)) -> R.Tensor((2, 3), dtype=None):
|
|
with R.dataflow():
|
|
gv: R.Tensor((2, 3), dtype=None) = R.nn.relu(x)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
mod = ToMixedPrecision()(Input)
|
|
tvm.ir.assert_structural_equal(mod, Input)
|
|
|
|
|
|
def test_relu_conv2d_relu():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3), "float32")
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
x0: R.Tensor((2, 3, 28, 28), "float32") = R.nn.relu(x)
|
|
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x0, w, out_dtype="float32")
|
|
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
x0: R.Tensor((2, 3, 28, 28), dtype="float32") = R.nn.relu(x)
|
|
lv1: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x0, dtype="float16")
|
|
lv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv1,
|
|
lv,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float16") = R.astype(lv2, dtype="float16")
|
|
lv3: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.relu(gv)
|
|
gv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 4, 26, 26), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
x0: R.Tensor((2, 3, 28, 28), dtype="float32") = R.nn.relu(x)
|
|
lv1: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x0, dtype="float16")
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.conv2d(
|
|
lv1,
|
|
lv,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
lv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.relu(gv)
|
|
gv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.astype(lv2, dtype="float32")
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_conv2d_relu_conv2d():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"),
|
|
w: R.Tensor((4, 3, 3, 3), "float32"),
|
|
w2: R.Tensor((4, 4, 3, 3), "float32"),
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
|
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.relu(gv)
|
|
gv3: R.Tensor((2, 4, 24, 24), "float32") = R.nn.conv2d(gv2, w2, out_dtype="float32")
|
|
R.output(gv3)
|
|
return gv3
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"),
|
|
w: R.Tensor((4, 3, 3, 3), dtype="float32"),
|
|
w2: R.Tensor((4, 4, 3, 3), dtype="float32"),
|
|
) -> R.Tensor((2, 4, 24, 24), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((4, 4, 3, 3), dtype="float16") = R.astype(w2, dtype="float16")
|
|
lv3: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float16") = R.astype(lv3, dtype="float16")
|
|
gv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.relu(gv)
|
|
gv3: R.Tensor((2, 4, 24, 24), dtype="float32") = R.nn.conv2d(
|
|
gv2,
|
|
lv2,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
R.output(gv3)
|
|
return gv3
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"),
|
|
w: R.Tensor((4, 3, 3, 3), dtype="float32"),
|
|
w2: R.Tensor((4, 4, 3, 3), dtype="float32"),
|
|
) -> R.Tensor((2, 4, 24, 24), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((4, 4, 3, 3), dtype="float16") = R.astype(w2, dtype="float16")
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.relu(gv)
|
|
lv3: R.Tensor((2, 4, 24, 24), dtype="float16") = R.nn.conv2d(
|
|
gv2,
|
|
lv2,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv3: R.Tensor((2, 4, 24, 24), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
R.output(gv3)
|
|
return gv3
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_gemm_add_silu():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 320), "float32"),
|
|
w1: R.Tensor((320, 1280), "float32"),
|
|
w2: R.Tensor((2, 1280), "float32"),
|
|
) -> R.Tensor(None, "float32", ndim=2):
|
|
with R.dataflow():
|
|
gv0: R.Tensor((2, 1280), "float32") = R.matmul(x, w1, out_dtype="float32")
|
|
gv1: R.Tensor((2, 1280), "float32") = R.add(gv0, w2)
|
|
gv2: R.Tensor((2, 1280), "float32") = R.nn.silu(gv1)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 320), dtype="float32"),
|
|
w1: R.Tensor((320, 1280), dtype="float32"),
|
|
w2: R.Tensor((2, 1280), dtype="float32"),
|
|
) -> R.Tensor((2, 1280), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 320), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((320, 1280), dtype="float16") = R.astype(w1, dtype="float16")
|
|
lv2: R.Tensor((2, 1280), dtype="float32") = R.matmul(lv, lv1, out_dtype="float32")
|
|
gv0: R.Tensor((2, 1280), dtype="float16") = R.astype(lv2, dtype="float16")
|
|
lv3: R.Tensor((2, 1280), dtype="float32") = R.astype(gv0, dtype="float32")
|
|
gv1: R.Tensor((2, 1280), dtype="float32") = R.add(lv3, w2)
|
|
gv2: R.Tensor((2, 1280), dtype="float32") = R.nn.silu(gv1)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 320), dtype="float32"),
|
|
w1: R.Tensor((320, 1280), dtype="float32"),
|
|
w2: R.Tensor((2, 1280), dtype="float32"),
|
|
) -> R.Tensor((2, 1280), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 320), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((320, 1280), dtype="float16") = R.astype(w1, dtype="float16")
|
|
gv0: R.Tensor((2, 1280), dtype="float16") = R.matmul(lv, lv1, out_dtype="float16")
|
|
lv2: R.Tensor((2, 1280), dtype="float32") = R.astype(gv0, dtype="float32")
|
|
gv1: R.Tensor((2, 1280), dtype="float32") = R.add(lv2, w2)
|
|
gv2: R.Tensor((2, 1280), dtype="float32") = R.nn.silu(gv1)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_tuple():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"),
|
|
w: R.Tensor((4, 3, 3, 3), "float32"),
|
|
w_2: R.Tensor((4, 4, 3, 3), "float32"),
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
|
gv2: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
|
gv3 = (gv, gv2)
|
|
gv4 = (gv3, gv2)
|
|
gv5 = gv4[0]
|
|
gv6 = gv5[0]
|
|
gv7 = R.nn.conv2d(gv6, w_2, out_dtype="float32")
|
|
R.output(gv7)
|
|
return gv7
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"),
|
|
w: R.Tensor((4, 3, 3, 3), dtype="float32"),
|
|
w_2: R.Tensor((4, 4, 3, 3), dtype="float32"),
|
|
) -> R.Tensor((2, 4, 24, 24), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((4, 4, 3, 3), dtype="float16") = R.astype(w_2, dtype="float16")
|
|
lv3: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float16") = R.astype(lv3, dtype="float16")
|
|
lv4: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.astype(lv4, dtype="float16")
|
|
gv3: R.Tuple(
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
) = (gv, gv2)
|
|
gv4: R.Tuple(
|
|
R.Tuple(
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
) = (gv3, gv2)
|
|
gv5: R.Tuple(
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
) = gv4[0]
|
|
gv6: R.Tensor((2, 4, 26, 26), dtype="float16") = gv5[0]
|
|
gv7: R.Tensor((2, 4, 24, 24), dtype="float32") = R.nn.conv2d(
|
|
gv6,
|
|
lv2,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
R.output(gv7)
|
|
return gv7
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"),
|
|
w: R.Tensor((4, 3, 3, 3), dtype="float32"),
|
|
w_2: R.Tensor((4, 4, 3, 3), dtype="float32"),
|
|
) -> R.Tensor((2, 4, 24, 24), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((4, 4, 3, 3), dtype="float16") = R.astype(w_2, dtype="float16")
|
|
gv: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv2: R.Tensor((2, 4, 26, 26), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv3: R.Tuple(
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
) = (gv, gv2)
|
|
gv4: R.Tuple(
|
|
R.Tuple(
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
) = (gv3, gv2)
|
|
gv5: R.Tuple(
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
R.Tensor((2, 4, 26, 26), dtype="float16"),
|
|
) = gv4[0]
|
|
gv6: R.Tensor((2, 4, 26, 26), dtype="float16") = gv5[0]
|
|
lv3: R.Tensor((2, 4, 24, 24), dtype="float16") = R.nn.conv2d(
|
|
gv6,
|
|
lv2,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv7: R.Tensor((2, 4, 24, 24), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
R.output(gv7)
|
|
return gv7
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_concat_matmul():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
lv10: R.Tensor((2, 160), "float32"),
|
|
lv12: R.Tensor((2, 160), "float32"),
|
|
w: R.Tensor((320, 1280), "float32"),
|
|
) -> R.Tensor(None, "float32", ndim=2):
|
|
with R.dataflow():
|
|
lv13: R.Tensor((2, 320), "float32") = R.concat((lv10, lv12), axis=-1)
|
|
lv14: R.Tensor((2, 1280), "float32") = R.matmul(lv13, w, out_dtype="float32")
|
|
R.output(lv14)
|
|
return lv14
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
lv10: R.Tensor((2, 160), dtype="float32"),
|
|
lv12: R.Tensor((2, 160), dtype="float32"),
|
|
w: R.Tensor((320, 1280), dtype="float32"),
|
|
) -> R.Tensor((2, 1280), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((320, 1280), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv13: R.Tensor((2, 320), dtype="float32") = R.concat((lv10, lv12), axis=-1)
|
|
lv1: R.Tensor((2, 320), dtype="float16") = R.astype(lv13, dtype="float16")
|
|
lv14: R.Tensor((2, 1280), dtype="float32") = R.matmul(lv1, lv, out_dtype="float32")
|
|
R.output(lv14)
|
|
return lv14
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
lv10: R.Tensor((2, 160), dtype="float32"),
|
|
lv12: R.Tensor((2, 160), dtype="float32"),
|
|
w: R.Tensor((320, 1280), dtype="float32"),
|
|
) -> R.Tensor((2, 1280), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((320, 1280), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv13: R.Tensor((2, 320), dtype="float32") = R.concat((lv10, lv12), axis=-1)
|
|
lv1: R.Tensor((2, 320), dtype="float16") = R.astype(lv13, dtype="float16")
|
|
lv2: R.Tensor((2, 1280), dtype="float16") = R.matmul(lv1, lv, out_dtype="float16")
|
|
lv14: R.Tensor((2, 1280), dtype="float32") = R.astype(lv2, dtype="float32")
|
|
R.output(lv14)
|
|
return lv14
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_conv2d_softmax():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((3, 3, 3, 3), "float32")
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
gv: R.Tensor((2, 3, 28, 28), "float32") = R.nn.conv2d(x, w, padding=(1, 1))
|
|
gv1: R.Tensor((2, 3, 28, 28), "float32") = R.nn.softmax(x, axis=1)
|
|
gv2 = R.add(gv, gv1)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((3, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 3, 28, 28), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv1: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv2: R.Tensor((2, 3, 28, 28), dtype="float32") = R.nn.conv2d(
|
|
lv1,
|
|
lv,
|
|
strides=[1, 1],
|
|
padding=[1, 1, 1, 1],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
gv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(lv2, dtype="float16")
|
|
gv1: R.Tensor((2, 3, 28, 28), dtype="float32") = R.nn.softmax(x, axis=1)
|
|
lv3: R.Tensor((2, 3, 28, 28), dtype="float32") = R.astype(gv, dtype="float32")
|
|
gv2: R.Tensor((2, 3, 28, 28), dtype="float32") = R.add(lv3, gv1)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((3, 3, 3, 3), dtype="float32")
|
|
) -> R.Tensor((2, 3, 28, 28), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((3, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv1: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
gv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.nn.conv2d(
|
|
lv1,
|
|
lv,
|
|
strides=[1, 1],
|
|
padding=[1, 1, 1, 1],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
gv1: R.Tensor((2, 3, 28, 28), dtype="float32") = R.nn.softmax(x, axis=1)
|
|
lv2: R.Tensor((2, 3, 28, 28), dtype="float32") = R.astype(gv, dtype="float32")
|
|
gv2: R.Tensor((2, 3, 28, 28), dtype="float32") = R.add(lv2, gv1)
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_conv2d_bias_conv2d():
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
z: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w0: R.Tensor((512, 4, 3, 3), dtype="float16"),
|
|
w1: R.Tensor((512,), dtype="float16"),
|
|
w2: R.Tensor((4, 4, 1, 1), dtype="float16"),
|
|
w3: R.Tensor((4,), dtype="float16"),
|
|
) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((512, 4, 3, 3), dtype="float32") = R.wrap_param(w0, dtype="float32")
|
|
lv1: R.Tensor((512,), dtype="float32") = R.wrap_param(w1, dtype="float32")
|
|
lv140: R.Tensor((4, 4, 1, 1), dtype="float32") = R.wrap_param(w2, dtype="float32")
|
|
lv141: R.Tensor((4,), dtype="float32") = R.wrap_param(w3, dtype="float32")
|
|
lv142: R.Tensor((1, 4, 64, 64), dtype="float32") = R.nn.conv2d(
|
|
z,
|
|
lv140,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv143: R.Tensor((1, 4, 1, 1), dtype="float32") = R.reshape(lv141, (1, 4, 1, 1))
|
|
lv144: R.Tensor((1, 4, 64, 64), dtype="float32") = R.add(lv142, lv143)
|
|
lv145: R.Tensor((1, 512, 64, 64), dtype="float32") = R.nn.conv2d(
|
|
lv144,
|
|
lv,
|
|
strides=[1, 1],
|
|
padding=[1, 1, 1, 1],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv146: R.Tensor((1, 512, 1, 1), dtype="float32") = R.reshape(lv1, (1, 512, 1, 1))
|
|
lv147: R.Tensor((1, 512, 64, 64), dtype="float32") = R.add(lv145, lv146)
|
|
gv: R.Tensor((1, 512, 64, 64), dtype="float32") = lv147
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
z: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w0: R.Tensor((512, 4, 3, 3), dtype="float16"),
|
|
w1: R.Tensor((512,), dtype="float16"),
|
|
w2: R.Tensor((4, 4, 1, 1), dtype="float16"),
|
|
w3: R.Tensor((4,), dtype="float16"),
|
|
) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(z, dtype="float16")
|
|
lv_1: R.Tensor((512, 4, 3, 3), dtype="float16") = w0
|
|
lv1: R.Tensor((512,), dtype="float16") = w1
|
|
lv140: R.Tensor((4, 4, 1, 1), dtype="float16") = w2
|
|
lv141: R.Tensor((4,), dtype="float16") = w3
|
|
lv1_1: R.Tensor((1, 4, 64, 64), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv140,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv142: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(lv1_1, dtype="float16")
|
|
lv143: R.Tensor((1, 4, 1, 1), dtype="float16") = R.reshape(lv141, (1, 4, 1, 1))
|
|
lv144: R.Tensor((1, 4, 64, 64), dtype="float16") = R.add(lv142, lv143)
|
|
lv2: R.Tensor((1, 512, 64, 64), dtype="float32") = R.nn.conv2d(
|
|
lv144,
|
|
lv_1,
|
|
strides=[1, 1],
|
|
padding=[1, 1, 1, 1],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv145: R.Tensor((1, 512, 64, 64), dtype="float16") = R.astype(lv2, dtype="float16")
|
|
lv146: R.Tensor((1, 512, 1, 1), dtype="float16") = R.reshape(lv1, (1, 512, 1, 1))
|
|
lv147: R.Tensor((1, 512, 64, 64), dtype="float16") = R.add(lv145, lv146)
|
|
gv: R.Tensor((1, 512, 64, 64), dtype="float32") = R.astype(lv147, dtype="float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected2:
|
|
@R.function
|
|
def main(
|
|
z: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w0: R.Tensor((512, 4, 3, 3), dtype="float16"),
|
|
w1: R.Tensor((512,), dtype="float16"),
|
|
w2: R.Tensor((4, 4, 1, 1), dtype="float16"),
|
|
w3: R.Tensor((4,), dtype="float16"),
|
|
) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(z, dtype="float16")
|
|
lv_1: R.Tensor((512, 4, 3, 3), dtype="float16") = w0
|
|
lv1: R.Tensor((512,), dtype="float16") = w1
|
|
lv140: R.Tensor((4, 4, 1, 1), dtype="float16") = w2
|
|
lv141: R.Tensor((4,), dtype="float16") = w3
|
|
lv142: R.Tensor((1, 4, 64, 64), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv140,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
lv143: R.Tensor((1, 4, 1, 1), dtype="float16") = R.reshape(
|
|
lv141, R.shape([1, 4, 1, 1])
|
|
)
|
|
lv144: R.Tensor((1, 4, 64, 64), dtype="float16") = R.add(lv142, lv143)
|
|
lv145: R.Tensor((1, 512, 64, 64), dtype="float16") = R.nn.conv2d(
|
|
lv144,
|
|
lv_1,
|
|
strides=[1, 1],
|
|
padding=[1, 1, 1, 1],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float16",
|
|
)
|
|
lv146: R.Tensor((1, 512, 1, 1), dtype="float16") = R.reshape(
|
|
lv1, R.shape([1, 512, 1, 1])
|
|
)
|
|
lv147: R.Tensor((1, 512, 64, 64), dtype="float16") = R.add(lv145, lv146)
|
|
gv: R.Tensor((1, 512, 64, 64), dtype="float32") = R.astype(lv147, dtype="float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
binding = {
|
|
"w0": np.random.uniform(size=(512, 4, 3, 3)).astype("float16"),
|
|
"w1": np.random.uniform(size=(512,)).astype("float16"),
|
|
"w2": np.random.uniform(size=(4, 4, 1, 1)).astype("float16"),
|
|
"w3": np.random.uniform(size=(4,)).astype("float16"),
|
|
}
|
|
binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()}
|
|
Input = relax.transform.BindParams("main", binding)(Input)
|
|
Expected = relax.transform.BindParams("main", binding)(Expected)
|
|
Expected2 = relax.transform.BindParams("main", binding)(Expected2)
|
|
_assert_test(Input, Expected, Expected2)
|
|
|
|
|
|
def test_tuple_get():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((512, 1, 1), dtype="float32"),
|
|
) -> R.Tensor((1, 256, 64, 64), dtype="float32"):
|
|
with R.dataflow():
|
|
conv = R.nn.conv2d(
|
|
x,
|
|
w,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 1, 1],
|
|
)
|
|
bias_out = R.add(conv, bias)
|
|
split = R.split(bias_out, indices_or_sections=2, axis=1)
|
|
out = R.add(split[0], split[1])
|
|
R.output(out)
|
|
return out
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((512, 1, 1), dtype="float32"),
|
|
):
|
|
with R.dataflow():
|
|
lv = R.astype(x, dtype="float16")
|
|
lv1 = R.astype(w, dtype="float16")
|
|
conv = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 1, 1],
|
|
out_dtype="float16",
|
|
)
|
|
lv2 = R.astype(conv, dtype="float32")
|
|
bias_out = R.add(lv2, bias)
|
|
split = R.split(bias_out, indices_or_sections=2, axis=1)
|
|
lv3 = split[0]
|
|
lv4 = split[1]
|
|
out = R.add(lv3, lv4)
|
|
R.output(out)
|
|
return out
|
|
|
|
_assert_test(Module, expected2=Expected)
|
|
|
|
|
|
def test_conv2d_bias_fp32():
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((512,), dtype="float32"),
|
|
) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv142: R.Tensor((1, 512, 62, 62), dtype="float32") = R.nn.conv2d(
|
|
x,
|
|
w,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
out_dtype="float32",
|
|
)
|
|
lv143: R.Tensor((1, 512, 1, 1), dtype="float32") = R.reshape(bias, (1, 512, 1, 1))
|
|
lv144: R.Tensor((1, 512, 62, 62), dtype="float32") = R.add(lv142, lv143)
|
|
R.output(lv144)
|
|
return lv144
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((512,), dtype="float32"),
|
|
) -> R.Tensor((1, 512, 62, 62), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((512, 4, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv142: R.Tensor((1, 512, 62, 62), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
out_dtype="float16",
|
|
)
|
|
lv2: R.Tensor((512,), dtype="float16") = R.astype(bias, dtype="float16")
|
|
lv143: R.Tensor((1, 512, 1, 1), dtype="float16") = R.reshape(
|
|
lv2, R.shape([1, 512, 1, 1])
|
|
)
|
|
lv3: R.Tensor((1, 512, 62, 62), dtype="float16") = R.add(lv142, lv143)
|
|
lv144: R.Tensor((1, 512, 62, 62), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
R.output(lv144)
|
|
return lv144
|
|
|
|
@tvm.script.ir_module
|
|
class Expected_no_bias_cast:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((512,), dtype="float32"),
|
|
) -> R.Tensor((1, 512, 62, 62), dtype="float32"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((512, 4, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv142: R.Tensor((1, 512, 62, 62), dtype="float16") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
out_dtype="float16",
|
|
)
|
|
lv143: R.Tensor((1, 512, 1, 1), dtype="float32") = R.reshape(
|
|
bias, R.shape([1, 512, 1, 1])
|
|
)
|
|
lv2: R.Tensor((1, 512, 62, 62), dtype="float32") = R.astype(lv142, dtype="float32")
|
|
lv144: R.Tensor((1, 512, 62, 62), dtype="float32") = R.add(lv2, lv143)
|
|
R.output(lv144)
|
|
return lv144
|
|
|
|
binding_np = {
|
|
"w": np.random.uniform(size=(512, 4, 3, 3)).astype("float32"),
|
|
"bias": np.random.uniform(size=(512,)).astype("float32"),
|
|
}
|
|
binding = {k: tvm.runtime.tensor(v) for k, v in binding_np.items()}
|
|
|
|
Input_bound = relax.transform.BindParams("main", binding)(Input)
|
|
Expected = relax.transform.BindParams("main", binding)(Expected)
|
|
|
|
_assert_test(Input_bound, expected2=Expected)
|
|
|
|
binding_np["bias"][0] = 70000 # Out of fp16 range
|
|
binding = {k: tvm.runtime.tensor(v) for k, v in binding_np.items()}
|
|
Input_bound = relax.transform.BindParams("main", binding)(Input)
|
|
Expected_no_bias_cast = relax.transform.BindParams("main", binding)(Expected_no_bias_cast)
|
|
|
|
_assert_test(Input_bound, expected2=Expected_no_bias_cast)
|
|
|
|
|
|
def test_convert_sig():
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float32"),
|
|
bias: R.Tensor((512,), dtype="float32"),
|
|
) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
|
|
# block 0
|
|
with R.dataflow():
|
|
lv142: R.Tensor((1, 512, 62, 62), dtype="float32") = R.nn.conv2d(
|
|
x,
|
|
w,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
out_dtype="float32",
|
|
)
|
|
lv143: R.Tensor((1, 512, 1, 1), dtype="float32") = R.reshape(bias, (1, 512, 1, 1))
|
|
lv144: R.Tensor((1, 512, 62, 62), dtype="float32") = R.add(lv142, lv143)
|
|
R.output(lv144)
|
|
return lv144
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((1, 4, 64, 64), dtype="float32"),
|
|
w: R.Tensor((512, 4, 3, 3), dtype="float16"),
|
|
bias: R.Tensor((512,), dtype="float16"),
|
|
) -> R.Tensor((1, 512, 62, 62), dtype="float32"):
|
|
with R.dataflow():
|
|
lv = R.astype(x, dtype="float16")
|
|
lv142 = R.nn.conv2d(
|
|
lv, w, strides=[1, 1], padding=[0, 0, 0, 0], out_dtype="float16"
|
|
)
|
|
lv143 = R.reshape(bias, R.shape([1, 512, 1, 1]))
|
|
lv1 = R.add(lv142, lv143)
|
|
lv144 = R.astype(lv1, dtype="float32")
|
|
R.output(lv144)
|
|
return lv144
|
|
|
|
mod = ToMixedPrecision(out_dtype="float16", fp16_input_names=["w", "bias"])(Input)
|
|
tvm.ir.assert_structural_equal(mod, Expected)
|
|
|
|
|
|
def test_call_tir_with_float16_args():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(A: R.Tensor([64], "float16")):
|
|
cls = Before
|
|
with R.dataflow():
|
|
B = R.call_tir(cls.tir_identity, [A], out_ty=R.Tensor([64], "float16"))
|
|
C = R.call_tir(cls.tir_identity, [B], out_ty=R.Tensor([64], "float16"))
|
|
R.output(C)
|
|
return C
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def tir_identity(
|
|
Input: T.Buffer(64, "float16"),
|
|
Output: T.Buffer(64, "float16"),
|
|
):
|
|
for i in range(64):
|
|
with T.sblock("copy"):
|
|
vi = T.axis.remap("S", [i])
|
|
Output[vi] = Input[vi]
|
|
|
|
Expected = Before
|
|
|
|
After = ToMixedPrecision()(Before)
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
def test_dynamic_strided_slice():
|
|
@I.ir_module(s_tir=True)
|
|
class Input:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), "float32"),
|
|
w: R.Tensor((4, 3, 3, 3), "float32"),
|
|
begin: R.Tensor((4,), "int64"),
|
|
end: R.Tensor((4,), "int64"),
|
|
strides: R.Tensor((4,), "int64"),
|
|
) -> R.Tensor(None, "float32", ndim=4):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w, out_dtype="float32")
|
|
gv = R.dynamic_strided_slice(lv, begin, end, strides)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((2, 3, 28, 28), dtype="float32"),
|
|
w: R.Tensor((4, 3, 3, 3), dtype="float32"),
|
|
begin: R.Tensor((4,), dtype="int64"),
|
|
end: R.Tensor((4,), dtype="int64"),
|
|
strides: R.Tensor((4,), dtype="int64"),
|
|
) -> R.Tensor(None, dtype="float32", ndim=4):
|
|
with R.dataflow():
|
|
lv: R.Tensor((2, 3, 28, 28), dtype="float16") = R.astype(x, dtype="float16")
|
|
lv1: R.Tensor((4, 3, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
|
|
lv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.nn.conv2d(
|
|
lv,
|
|
lv1,
|
|
strides=[1, 1],
|
|
padding=[0, 0, 0, 0],
|
|
dilation=[1, 1],
|
|
groups=1,
|
|
data_layout="NCHW",
|
|
kernel_layout="OIHW",
|
|
out_layout="NCHW",
|
|
out_dtype="float32",
|
|
)
|
|
lv3: R.Tensor((2, 4, 26, 26), dtype="float16") = R.astype(lv2, dtype="float16")
|
|
lv4: R.Tensor((2, 4, 26, 26), dtype="float32") = R.astype(lv3, dtype="float32")
|
|
gv: R.Tensor(None, dtype="float32", ndim=4) = R.dynamic_strided_slice(
|
|
lv4, begin, end, strides
|
|
)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
_assert_test(Input, Expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|