wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/python/relax/test_backend_transform_shape_lower.py
+819
View File
@@ -0,0 +1,819 @@
# 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()