wehub-resource-sync
820 lines
25 KiB
Python
820 lines
25 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: F841
|
|
|
|
import tvm.script
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.ir import assert_structural_equal
|
|
from tvm.relax.testing.runtime_builtin import MakeShapeCode, MatchShapeCode
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
# note: we expected RemovePurityChecking to be run first, so we force purity in most test cases
|
|
|
|
|
|
def test_const_shape_arg():
|
|
MS = MatchShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Shape([1, 2]), y: R.Shape):
|
|
R.func_attr({"relax.force_pure": True})
|
|
return x
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def extra_func(H: T.Buffer(T.int64(4), "int64")):
|
|
"""Extra function, checks if the pass preserves it."""
|
|
H[T.int64(1)] = H[T.int64(0)] + T.int64(1)
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Shape([1, 2]), y: R.Shape):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap = R.null_value()
|
|
_ = R.call_packed("vm.builtin.check_shape_info", x, 2, "", ty_args=[R.Tuple()])
|
|
_ = R.call_packed("vm.builtin.check_shape_info", y, -1, "", ty_args=[R.Tuple()])
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
2,
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
1,
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
2,
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
return x
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def extra_func(H: T.Buffer(T.int64(4), "int64")):
|
|
H[T.int64(1)] = H[T.int64(0)] + T.int64(1)
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_static_fn_check():
|
|
"""Check static shape and function."""
|
|
MS = MatchShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(f: R.Callable([R.Any], R.Any), y: R.Shape([1, 2])):
|
|
R.func_attr({"relax.force_pure": True})
|
|
return y
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(f: R.Callable([R.Any], R.Any), y: R.Shape([1, 2])):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap = R.null_value()
|
|
_ = R.call_packed("vm.builtin.check_func_info", f, "", ty_args=[R.Tuple()])
|
|
_ = R.call_packed("vm.builtin.check_shape_info", y, 2, "", ty_args=[R.Tuple()])
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
y,
|
|
shape_heap,
|
|
2,
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
1,
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
2,
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
return y
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_simple_symbolic_shape():
|
|
MS = MatchShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", 2, "m"], "float32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
return x
|
|
|
|
sindex = {
|
|
"n": 0,
|
|
"m": 1,
|
|
}
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", 2, "m"], "float32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
[R.prim_value(2)],
|
|
ty_args=[R.Tensor(ndim=1, dtype="int64")],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
x,
|
|
3,
|
|
R.dtype("float32"),
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
3,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["n"],
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
2,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["m"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
return x
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_symbolic_compute():
|
|
MS = MatchShapeCode
|
|
MK = MakeShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "m"], "float32"), y: R.Tensor(ndim=3, dtype=None)) -> R.Shape(
|
|
ndim=3
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
m = T.int64()
|
|
k = T.int64()
|
|
z = R.match_cast(y, R.Tensor([k, m, k + 1], dtype=None))
|
|
return R.shape([k + 1, m, 2])
|
|
|
|
# slot assignment:
|
|
# 0: n, 1: m, 2:k, 3: k+1
|
|
sindex = {"n": 0, "m": 1, "k": 2, "k+1": 3}
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def shape_func(H: T.Buffer(T.int64(4), "int64")):
|
|
# generated compute function
|
|
T.func_attr({"tirx.is_host_func": True})
|
|
H[T.int64(sindex["k+1"])] = H[T.int64(sindex["k"])] + T.int64(1)
|
|
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "m"], "float32"), y: R.Tensor(ndim=3, dtype=None)) -> R.Shape(
|
|
ndim=3
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
m = T.int64()
|
|
k = T.int64()
|
|
cls = Expected
|
|
shape_heap = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
[R.prim_value(4)],
|
|
ty_args=[R.Tensor(ndim=1, dtype="int64")],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
x,
|
|
2,
|
|
R.dtype("float32"),
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
gv = R.null_value()
|
|
_ = R.call_packed("vm.builtin.check_tensor_info", y, 3, gv, "", ty_args=[R.Tuple()])
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
2,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["n"],
|
|
MS.STORE_TO_HEAP,
|
|
sindex["m"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
y,
|
|
shape_heap,
|
|
3,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["k"],
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["m"],
|
|
MS.NO_OP,
|
|
0,
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
_ = cls.shape_func(shape_heap)
|
|
# extra assertion on y's shape after shape computation
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
y,
|
|
shape_heap,
|
|
3,
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["k"],
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["m"],
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["k+1"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
z = R.match_cast(y, R.Tensor([k, m, k + 1], dtype=None))
|
|
# construct shape value for return
|
|
s = R.call_packed(
|
|
"vm.builtin.make_shape",
|
|
shape_heap,
|
|
3,
|
|
MK.LOAD_SHAPE,
|
|
sindex["k+1"],
|
|
MK.LOAD_SHAPE,
|
|
sindex["m"],
|
|
MK.USE_IMM,
|
|
2,
|
|
ty_args=[R.Shape(ndim=3)],
|
|
)
|
|
return s
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_tuple_handling():
|
|
MS = MatchShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tuple(
|
|
R.Tensor(["n", "m"], "float32"), R.Tuple(R.Shape, R.Tensor(["n", "k"], "int32"))
|
|
),
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
return x
|
|
|
|
# slot assignment:
|
|
sindex = {"n": 0, "m": 1, "k": 2}
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tuple(
|
|
R.Tensor(["n", "m"], "float32"), R.Tuple(R.Shape, R.Tensor(["n", "k"], "int32"))
|
|
),
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
[R.prim_value(3)],
|
|
ty_args=[R.Tensor(ndim=1, dtype="int64")],
|
|
)
|
|
# recursively unpack tuple for static info check
|
|
_ = R.call_packed("vm.builtin.check_tuple_info", x, 2, "", ty_args=[R.Tuple()])
|
|
t0 = x[0]
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
t0,
|
|
2,
|
|
R.dtype("float32"),
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
t1 = x[1]
|
|
_ = R.call_packed("vm.builtin.check_tuple_info", t1, 2, "", ty_args=[R.Tuple()])
|
|
t1x0 = t1[0]
|
|
_ = R.call_packed("vm.builtin.check_shape_info", t1x0, -1, "", ty_args=[R.Tuple()])
|
|
t1x1 = t1[1]
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
t1x1,
|
|
2,
|
|
R.dtype("int32"),
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
# match shape checks.
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
t0,
|
|
shape_heap,
|
|
2,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["n"],
|
|
MS.STORE_TO_HEAP,
|
|
sindex["m"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
t1x1,
|
|
shape_heap,
|
|
2,
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["n"],
|
|
MS.STORE_TO_HEAP,
|
|
sindex["k"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
return x
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_return_match_check():
|
|
"""Test when return body is not same as ret_ty, runtime match check needed."""
|
|
MS = MatchShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "m"], "float32"), y: R.Any) -> R.Tuple(
|
|
R.Tensor(["n", "m"], "float32")
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
return y
|
|
|
|
# slot assignment:
|
|
sindex = {
|
|
"n": 0,
|
|
"m": 1,
|
|
}
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "m"], "float32"), y: R.Any) -> R.Tuple(
|
|
R.Tensor(["n", "m"], "float32")
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
[R.prim_value(2)],
|
|
ty_args=[R.Tensor(ndim=1, dtype="int64")],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info", x, 2, R.dtype("float32"), "", ty_args=[R.Tuple()]
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
2,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["n"],
|
|
MS.STORE_TO_HEAP,
|
|
sindex["m"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
_ = R.call_packed("vm.builtin.check_tuple_info", y, 1, "", ty_args=[R.Tuple()])
|
|
# emit runtime function call since y do not have the right type.
|
|
y1 = R.call_packed("vm.builtin.tuple_getitem", y, 0, ty_args=[R.Any])
|
|
# run check
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
y1,
|
|
2,
|
|
R.dtype("float32"),
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
# shape check
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
y1,
|
|
shape_heap,
|
|
2,
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["n"],
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["m"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
|
|
return y
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_return_match_check_with_new_expr():
|
|
"""Like test_return_match_check, but requires a computation
|
|
|
|
When return body is not same as ret_ty, a runtime match
|
|
check is required. This match check may require a symbolic
|
|
expression to be computed.
|
|
"""
|
|
MS = MatchShapeCode
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "n"], "float32")) -> R.Tensor(["n * n"], "float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
out = R.call_packed("flatten_matrix", x, ty_args=R.Any)
|
|
return out
|
|
|
|
# slot assignment:
|
|
sindex = {
|
|
"n": 0,
|
|
"n * n": 1,
|
|
}
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(["n", "n"], "float32")) -> R.Tensor(["n * n"], "float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
[R.prim_value(2)],
|
|
ty_args=[R.Tensor(ndim=1, dtype="int64")],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info", x, 2, R.dtype("float32"), "", ty_args=[R.Tuple()]
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
2,
|
|
MS.STORE_TO_HEAP,
|
|
sindex["n"],
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["n"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
|
|
_ = Expected.shape_func(shape_heap)
|
|
|
|
out = R.call_packed("flatten_matrix", x, ty_args=R.Any)
|
|
_ = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
out,
|
|
1,
|
|
R.dtype("float32"),
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
_ = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
out,
|
|
shape_heap,
|
|
1,
|
|
MS.ASSERT_EQUAL_TO_LOAD,
|
|
sindex["n * n"],
|
|
"",
|
|
ty_args=[R.Tuple()],
|
|
)
|
|
return out
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def shape_func(H: T.Buffer(T.int64(2), "int64")):
|
|
# generated compute function
|
|
T.func_attr({"tirx.is_host_func": True})
|
|
H[T.int64(sindex["n * n"])] = H[T.int64(sindex["n"])] * H[T.int64(sindex["n"])]
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_symbolic_shape_multiple_function():
|
|
MS = MatchShapeCode
|
|
MK = MakeShapeCode
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def fn1(A: R.Tensor(("m", "n"), dtype="float32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
m = T.int64()
|
|
n = T.int64()
|
|
return A
|
|
|
|
@R.function
|
|
def fn2(A: R.Tensor(("n", "m"), dtype="float32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
n = T.int64()
|
|
m = T.int64()
|
|
return A
|
|
|
|
# slot assignment:
|
|
sindex_fn1 = {
|
|
"m": 0,
|
|
"n": 1,
|
|
}
|
|
sindex_fn2 = {
|
|
"n": 0,
|
|
"m": 1,
|
|
}
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def fn1(A: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m", "n"), dtype="float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
m = T.int64()
|
|
n = T.int64()
|
|
shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
(R.prim_value(2),),
|
|
ty_args=(R.Tensor(dtype="int64", ndim=1),),
|
|
)
|
|
_: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
A,
|
|
R.prim_value(2),
|
|
R.dtype("float32"),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_1: R.Tuple = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
A,
|
|
shape_heap,
|
|
R.prim_value(2),
|
|
MS.STORE_TO_HEAP,
|
|
sindex_fn1["m"],
|
|
MS.STORE_TO_HEAP,
|
|
sindex_fn1["n"],
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
return A
|
|
|
|
@R.function
|
|
def fn2(A: R.Tensor(("n", "m"), dtype="float32")) -> R.Tensor(("n", "m"), dtype="float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
n = T.int64()
|
|
m = T.int64()
|
|
shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
(R.prim_value(2),),
|
|
ty_args=(R.Tensor(dtype="int64", ndim=1),),
|
|
)
|
|
_2: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
A,
|
|
R.prim_value(2),
|
|
R.dtype("float32"),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_3: R.Tuple = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
A,
|
|
shape_heap,
|
|
R.prim_value(2),
|
|
MS.STORE_TO_HEAP,
|
|
sindex_fn2["n"],
|
|
MS.STORE_TO_HEAP,
|
|
sindex_fn2["m"],
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
return A
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_check_lifted_weights():
|
|
MS = MatchShapeCode
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main_transform_params(params: R.Tuple(R.Tensor((16, 16), dtype="float32"))) -> R.Tuple(
|
|
R.Tensor((16, 16), dtype="float32")
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
return params
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((16, 16), "float32"), param_0: R.Tensor((16, 16), dtype="float32")):
|
|
R.func_attr({"relax.force_pure": True, "num_input": 1})
|
|
return (x, param_0)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main_transform_params(params: R.Tuple(R.Tensor((16, 16), dtype="float32"))) -> R.Tuple(
|
|
R.Tensor((16, 16), dtype="float32")
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape_heap: R.Any = R.null_value()
|
|
_: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tuple_info",
|
|
params,
|
|
R.prim_value(1),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
gv: R.Tensor((16, 16), dtype="float32") = params[0]
|
|
_1: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
gv,
|
|
R.prim_value(2),
|
|
R.dtype("float32"),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_2: R.Tuple = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
gv,
|
|
shape_heap,
|
|
R.prim_value(2),
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
R.prim_value(16),
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
R.prim_value(16),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
return params
|
|
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((16, 16), dtype="float32"), param_0: R.Tensor((16, 16), dtype="float32")
|
|
) -> R.Tuple(R.Tensor((16, 16), dtype="float32"), R.Tensor((16, 16), dtype="float32")):
|
|
R.func_attr({"num_input": 1, "relax.force_pure": True})
|
|
shape_heap: R.Any = R.null_value()
|
|
_: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
x,
|
|
R.prim_value(2),
|
|
R.dtype("float32"),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_1: R.Tuple = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
R.prim_value(2),
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
R.prim_value(16),
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
R.prim_value(16),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
return (x, param_0)
|
|
|
|
before = Before
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
expected = Expected
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_check_weights_with_dynamic_shape():
|
|
MS = MatchShapeCode
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((16, 16), "float32"),
|
|
params: R.Tuple(R.Tensor((16, 16), dtype="float32"), R.Tensor(("n",), "float32")),
|
|
):
|
|
R.func_attr({"relax.force_pure": True, "num_input": 1})
|
|
n = T.int64()
|
|
param_0 = params[0]
|
|
param_1 = params[1]
|
|
return (x, param_0, param_1)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor((16, 16), "float32"),
|
|
params: R.Tuple(R.Tensor((16, 16), dtype="float32"), R.Tensor(("n",), "float32")),
|
|
):
|
|
n = T.int64()
|
|
R.func_attr({"num_input": 1, "relax.force_pure": True})
|
|
shape_heap: R.Tensor(dtype="int64", ndim=1) = R.call_builtin_with_ctx(
|
|
"vm.builtin.alloc_shape_heap",
|
|
(R.prim_value(1),),
|
|
ty_args=(R.Tensor(dtype="int64", ndim=1),),
|
|
)
|
|
_: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
x,
|
|
R.prim_value(2),
|
|
R.dtype("float32"),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_1: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tuple_info",
|
|
params,
|
|
R.prim_value(2),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_param_1: R.Tensor((n,), dtype="float32") = params[1]
|
|
_2: R.Tuple = R.call_packed(
|
|
"vm.builtin.check_tensor_info",
|
|
_param_1,
|
|
R.prim_value(1),
|
|
R.dtype("float32"),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_3: R.Tuple = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
x,
|
|
shape_heap,
|
|
R.prim_value(2),
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
R.prim_value(16),
|
|
MS.ASSERT_EQUAL_TO_IMM,
|
|
R.prim_value(16),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
_4: R.Tuple = R.call_packed(
|
|
"vm.builtin.match_shape",
|
|
_param_1,
|
|
shape_heap,
|
|
MS.STORE_TO_HEAP,
|
|
R.prim_value(1),
|
|
R.prim_value(0),
|
|
R.str(""),
|
|
ty_args=(R.Tuple,),
|
|
)
|
|
|
|
param_0 = params[0]
|
|
param_1 = params[1]
|
|
return (x, param_0, param_1)
|
|
|
|
before = Before
|
|
after = relax.transform.VMShapeLower(emit_err_ctx=False)(before)
|
|
print(after)
|
|
expected = Expected
|
|
assert_structural_equal(after, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|