Shushi Hong
eafcba1c44
Fixes #19887. This PR fixes several Relax TensorRT BYOC issues exposed by YOLO-style models: - adds TensorRT support for SiLU and resize2d - preserves operand and TupleGetItem ordering during codegen - fixes cyclic and unsafe Tuple/TGI region merging - handles static Shape bindings and nested packed-function outputs - normalizes PrimType dtype arguments passed to relax.arange With these changes, yolo11n-seg can be merged into a single TensorRT region, while yolo11n can be imported and partitioned successfully.
763 lines
27 KiB
Python
763 lines
27 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 pytest
|
|
|
|
import tvm
|
|
import tvm.script
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def _make_dataflow_block_with_dual_role_var():
|
|
prim_ty = tvm.ir.PrimType("int64")
|
|
source = tvm.ir.Var("source", prim_ty)
|
|
data = tvm.ir.Var("data", relax.TensorType(dtype="float32", ndim=1))
|
|
extent = tvm.ir.Var("extent", prim_ty)
|
|
matched_ty = relax.TensorType([extent], "float32")
|
|
matched = relax.DataflowVar("matched", matched_ty)
|
|
|
|
block = relax.DataflowBlock(
|
|
[
|
|
relax.VarBinding(extent, source),
|
|
relax.MatchCast(matched, data, matched_ty),
|
|
]
|
|
)
|
|
func = relax.Function([source, data], relax.SeqExpr([block], extent), prim_ty)
|
|
return relax.transform.Normalize()(tvm.IRModule.from_expr(func))
|
|
|
|
|
|
def test_dataflowblock_pass_rejects_deleting_binding_role_of_symbolic_var():
|
|
mod = _make_dataflow_block_with_dual_role_var()
|
|
|
|
@relax.transform.dataflowblock_pass(opt_level=0)
|
|
def delete_ordinary_binding(block, _mod, _ctx):
|
|
return relax.DataflowBlock([block.bindings[1]])
|
|
|
|
with pytest.raises(tvm.error.InternalError, match="global-scope Var"):
|
|
delete_ordinary_binding(mod)
|
|
|
|
|
|
def test_dataflowblock_pass_rejects_rewriting_match_cast_role_of_binding_var():
|
|
mod = _make_dataflow_block_with_dual_role_var()
|
|
|
|
@relax.transform.dataflowblock_pass(opt_level=0)
|
|
def rewrite_match_cast_symbol(block, _mod, _ctx):
|
|
old_match_cast = block.bindings[1]
|
|
replacement = tvm.ir.Var("replacement", "int64")
|
|
replacement_ty = relax.TensorType([replacement], "float32")
|
|
new_match_cast = relax.MatchCast(
|
|
relax.DataflowVar(old_match_cast.var.name, replacement_ty),
|
|
old_match_cast.value,
|
|
replacement_ty,
|
|
)
|
|
return relax.DataflowBlock([block.bindings[0], new_match_cast])
|
|
|
|
with pytest.raises(tvm.error.InternalError, match="symbolic Var declared by a MatchCast"):
|
|
rewrite_match_cast_symbol(mod)
|
|
|
|
|
|
def test_to_non_dataflow():
|
|
@tvm.script.ir_module
|
|
class TestToNonDataflow:
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), "float32")):
|
|
m, n = T.int64(), T.int64()
|
|
with R.dataflow():
|
|
lv0 = R.call_dps_packed(
|
|
"test.op.identity",
|
|
(x,),
|
|
R.Tensor(
|
|
(m, n),
|
|
dtype="float32",
|
|
),
|
|
)
|
|
gv0 = R.call_dps_packed(
|
|
"test.op.identity",
|
|
(lv0,),
|
|
R.Tensor(
|
|
(m, n),
|
|
dtype="float32",
|
|
),
|
|
)
|
|
R.output(gv0)
|
|
return gv0
|
|
|
|
mod = TestToNonDataflow
|
|
|
|
old_vars = []
|
|
|
|
def fvisit(e):
|
|
if isinstance(e, relax.Var) and not tvm.ir.is_prim_expr(e):
|
|
nonlocal old_vars
|
|
old_vars.append(e)
|
|
|
|
relax.analysis.post_order_visit(mod["foo"], fvisit)
|
|
x, lv0, gv0 = old_vars
|
|
|
|
new_mod = relax.transform.ToNonDataflow()(mod)
|
|
|
|
new_vars = []
|
|
|
|
def fvisit(e):
|
|
if isinstance(e, relax.Var) and not tvm.ir.is_prim_expr(e):
|
|
nonlocal new_vars
|
|
new_vars.append(e)
|
|
|
|
relax.analysis.post_order_visit(new_mod["foo"], fvisit)
|
|
|
|
assert x == new_vars[0]
|
|
assert lv0 != new_vars[1]
|
|
assert isinstance(lv0, relax.DataflowVar)
|
|
assert not isinstance(new_vars[1], relax.DataflowVar)
|
|
|
|
assert isinstance(gv0, relax.Var)
|
|
assert isinstance(new_vars[2], relax.Var)
|
|
assert gv0 == new_vars[2]
|
|
|
|
|
|
def test_call_tir_rewrite():
|
|
@tvm.script.ir_module
|
|
class TestCallTIRRewrite:
|
|
@T.prim_func(s_tir=True)
|
|
def exp(A_handle: T.handle, B_handle: T.handle):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
A = T.match_buffer(A_handle, (m, n), "float32")
|
|
B = T.match_buffer(B_handle, (m, n), "float32")
|
|
T.evaluate(0)
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), "float32")):
|
|
# we expect RemovePurityChecking to have been used before this point
|
|
R.func_attr({"relax.force_pure": True})
|
|
m, n = T.int64(), T.int64()
|
|
gv0 = R.call_tir(TestCallTIRRewrite.exp, (x,), R.Tensor((m, n), dtype="float32"))
|
|
return gv0
|
|
|
|
mod = TestCallTIRRewrite
|
|
|
|
# before rewrite
|
|
v0 = mod["foo"].body.blocks[0].bindings[0].var
|
|
s0 = mod["foo"].body.blocks[0].bindings[0].value
|
|
assert isinstance(s0, relax.Call)
|
|
assert s0.op.name == "relax.call_tir"
|
|
|
|
# after rewrite
|
|
new_mod = relax.transform.CallTIRRewrite()(mod)
|
|
func = new_mod["foo"]
|
|
|
|
block = func.body.blocks[0]
|
|
assert not isinstance(block, relax.DataflowBlock)
|
|
|
|
s1 = block.bindings[0].value
|
|
assert isinstance(s1, relax.Call)
|
|
assert s1.op.name == "relax.builtin.alloc_tensor"
|
|
assert isinstance(s1.args[0], relax.ShapeExpr)
|
|
tvm.ir.assert_structural_equal(s1.args[0], s0.ty_args[0].shape)
|
|
s2 = block.bindings[1].value
|
|
tvm.ir.expr.GlobalVar
|
|
assert s2.op.name_hint == "exp"
|
|
|
|
|
|
def test_call_tir_rewrite_with_interspersed_primitive_argument():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def scale_add(
|
|
A: T.Buffer((16,), "float32"),
|
|
scale: T.float32,
|
|
C: T.Buffer((16,), "float32"),
|
|
B: T.Buffer((16,), "float32"),
|
|
):
|
|
for i in range(16):
|
|
B[i] = A[i] + scale * C[i]
|
|
|
|
@R.function
|
|
def main(
|
|
A: R.Tensor((16,), "float32"),
|
|
scale: R.Prim("float32"),
|
|
C: R.Tensor((16,), "float32"),
|
|
) -> R.Tensor((16,), "float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
B = R.call_tir(Module.scale_add, (A, scale, C), R.Tensor((16,), "float32"))
|
|
return B
|
|
|
|
after = relax.transform.CallTIRRewrite()(Module)
|
|
func = after["main"]
|
|
bindings = func.body.blocks[0].bindings
|
|
output_buffer = bindings[0].var
|
|
call = bindings[1].value
|
|
|
|
assert call.op.name_hint == "scale_add"
|
|
tvm.ir.assert_structural_equal(
|
|
call.args,
|
|
[func.params[0], func.params[1], func.params[2], output_buffer],
|
|
)
|
|
|
|
|
|
def test_transform_remove_purity_checking():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def base(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.add(x, x)
|
|
z = R.add(x, y)
|
|
return z
|
|
|
|
@R.function
|
|
def use_call_pure_packed(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.add(x, x)
|
|
z = R.call_pure_packed("vm.builtin.copy", y, ty_args=(R.Tensor((), dtype="int32")))
|
|
return z
|
|
|
|
@R.function
|
|
def use_invoke_pure_closure(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
closure = R.make_closure(Before.base, ())
|
|
res = R.invoke_pure_closure(closure, (x,), ty_args=R.Tensor((), "int32"))
|
|
return res
|
|
|
|
@R.function(pure=False)
|
|
def impure_func() -> R.Any:
|
|
y = R.print(format="I am impure!")
|
|
return y
|
|
|
|
@R.function
|
|
def nested_pure_func() -> R.Tensor((), "int32"):
|
|
@R.function
|
|
def nested(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.add(x, x)
|
|
q = R.call_pure_packed("vm.builtin.copy", y, ty_args=(R.Tensor((), dtype="int32")))
|
|
return q
|
|
|
|
z = R.const(1, dtype="int32")
|
|
w = nested(z)
|
|
return w
|
|
|
|
@R.function(pure=False)
|
|
def nested_impure_func() -> R.Tensor((), "int32"):
|
|
@R.function(pure=False)
|
|
def nested() -> R.Any:
|
|
x = R.print(format="Oops!")
|
|
return x
|
|
|
|
y = R.const(1, dtype="int32")
|
|
z = nested()
|
|
return y
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def base(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
y = R.add(x, x)
|
|
z = R.add(x, y)
|
|
return z
|
|
|
|
@R.function
|
|
def use_call_pure_packed(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
y = R.add(x, x)
|
|
z = R.call_packed("vm.builtin.copy", y, ty_args=(R.Tensor((), dtype="int32")))
|
|
return z
|
|
|
|
@R.function
|
|
def use_invoke_pure_closure(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
closure = R.make_closure(Expected.base, ())
|
|
res = R.invoke_closure(closure, (x,), ty_args=R.Tensor((), "int32"))
|
|
return res
|
|
|
|
@R.function(pure=False)
|
|
def impure_func() -> R.Any:
|
|
y = R.print(format="I am impure!")
|
|
return y
|
|
|
|
@R.function
|
|
def nested_pure_func() -> R.Tensor((), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
|
|
@R.function
|
|
def nested(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
y = R.add(x, x)
|
|
q = R.call_packed("vm.builtin.copy", y, ty_args=(R.Tensor((), dtype="int32")))
|
|
return q
|
|
|
|
z = R.const(1, dtype="int32")
|
|
w = nested(z)
|
|
return w
|
|
|
|
@R.function(pure=False)
|
|
def nested_impure_func() -> R.Tensor((), "int32"):
|
|
@R.function(pure=False)
|
|
def nested() -> R.Any:
|
|
x = R.print(format="Oops!")
|
|
return x
|
|
|
|
y = R.const(1, dtype="int32")
|
|
z = nested()
|
|
return y
|
|
|
|
new_mod = relax.transform.RemovePurityChecking()(Before)
|
|
tvm.ir.assert_structural_equal(new_mod, Expected)
|
|
|
|
|
|
def test_call_dps_packed_rewrite():
|
|
@tvm.script.ir_module
|
|
class TestCallDPSPackedRewrite:
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), "float32")):
|
|
# we expect RemovePurityChecking to have been used before this point
|
|
R.func_attr({"relax.force_pure": True})
|
|
m, n = T.int64(), T.int64()
|
|
gv0 = R.call_dps_packed("test.op.identity", (x,), R.Tensor((m, n), dtype="float32"))
|
|
return gv0
|
|
|
|
mod = TestCallDPSPackedRewrite
|
|
|
|
# before rewrite
|
|
v0 = mod["foo"].body.blocks[0].bindings[0].var
|
|
s0 = mod["foo"].body.blocks[0].bindings[0].value
|
|
assert isinstance(s0, relax.Call)
|
|
assert s0.op.name == "relax.call_dps_packed"
|
|
|
|
# CallTIRRewrite also works for call_dps_packed
|
|
new_mod = relax.transform.CallTIRRewrite()(mod)
|
|
func = new_mod["foo"]
|
|
|
|
block = func.body.blocks[0]
|
|
assert not isinstance(block, relax.DataflowBlock)
|
|
|
|
s1 = block.bindings[0].value
|
|
assert isinstance(s1, relax.Call)
|
|
assert s1.op.name == "relax.builtin.alloc_tensor"
|
|
assert isinstance(s1.args[0], relax.ShapeExpr)
|
|
tvm.ir.assert_structural_equal(s1.args[0], s0.ty_args[0].shape)
|
|
s2 = block.bindings[1].value
|
|
assert s2.op.global_symbol == "test.op.identity"
|
|
|
|
|
|
def test_call_dps_packed_rewrite_nested_tuple_output():
|
|
"""Flatten nested outputs for the packed ABI, then rebuild their Relax structure."""
|
|
input_ty = relax.TensorType((2, 3), "float32")
|
|
flat_output_types = [
|
|
relax.TensorType((2, 3), "float32"),
|
|
relax.TensorType((4,), "int32"),
|
|
relax.TensorType((5, 6), "float16"),
|
|
]
|
|
output_ty = tvm.ir.TupleType([flat_output_types[0], tvm.ir.TupleType(flat_output_types[1:])])
|
|
|
|
x = relax.Var("x", input_ty)
|
|
call = relax.Call(
|
|
tvm.ir.Op.get("relax.call_dps_packed"),
|
|
[relax.ExternFunc("test.op.nested_outputs"), relax.Tuple([x])],
|
|
ty_args=[output_ty],
|
|
)
|
|
builder = relax.BlockBuilder()
|
|
with builder.function("main", [x], attrs={"relax.force_pure": True}):
|
|
out = builder.emit(call)
|
|
builder.emit_func_output(out)
|
|
|
|
after = relax.transform.CallTIRRewrite()(builder.get())
|
|
relax.analysis.well_formed(after)
|
|
func = after["main"]
|
|
block = func.body.blocks[0]
|
|
|
|
alloc_bindings = block.bindings[:3]
|
|
for binding, expected_ty in zip(alloc_bindings, flat_output_types):
|
|
assert binding.value.op.name == "relax.builtin.alloc_tensor"
|
|
tvm.ir.assert_structural_equal(binding.var.ty, expected_ty)
|
|
|
|
packed_call = block.bindings[3].value
|
|
assert packed_call.op.global_symbol == "test.op.nested_outputs"
|
|
assert packed_call.args[0].same_as(func.params[0])
|
|
assert all(
|
|
arg.same_as(binding.var) for arg, binding in zip(packed_call.args[1:], alloc_bindings)
|
|
)
|
|
|
|
rebuilt = block.bindings[4].value
|
|
assert rebuilt.fields[0].same_as(alloc_bindings[0].var)
|
|
assert rebuilt.fields[1].fields[0].same_as(alloc_bindings[1].var)
|
|
assert rebuilt.fields[1].fields[1].same_as(alloc_bindings[2].var)
|
|
|
|
|
|
def test_call_tir_inplace_simple():
|
|
# simple case: one inplace argument
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@T.prim_func(s_tir=True)
|
|
def zeros(A: T.Buffer((2, 3), "int32")):
|
|
# just overwrites A with 0s
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.writes(A[ax0, ax1])
|
|
A[ax0, ax1] = T.int32(0)
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
|
|
# we expect RemovePurityChecking to have been used before this point
|
|
R.func_attr({"relax.force_pure": True})
|
|
gv0 = R.call_tir_inplace(Input.zeros, x, 0, R.Tensor((2, 3), dtype="int32"))
|
|
return gv0
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def zeros(A: T.Buffer((2, 3), "int32")):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.writes(A[ax0, ax1])
|
|
A[ax0, ax1] = T.int32(0)
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
_ = Expected.zeros(x)
|
|
gv0 = x
|
|
return gv0
|
|
|
|
new_mod = relax.transform.CallTIRRewrite()(Input)
|
|
tvm.ir.assert_structural_equal(Expected["foo"], new_mod["foo"], map_free_vars=True)
|
|
|
|
|
|
def test_call_tir_inplace_multiple_args():
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@T.prim_func(s_tir=True)
|
|
def copy(
|
|
A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32"), C: T.Buffer((2, 3), "int32")
|
|
):
|
|
# copies the contents of C into A and B
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(C[ax0, ax1])
|
|
T.writes(A[ax0, ax1], B[ax0, ax1])
|
|
A[ax0, ax1] = C[ax0, ax1]
|
|
B[ax0, ax1] = C[ax0, ax1]
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
|
|
) -> R.Tuple(R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
gv0 = R.call_tir_inplace(
|
|
Input.copy,
|
|
(x, y, z),
|
|
[0, 1],
|
|
[R.Tensor((2, 3), dtype="int32"), R.Tensor((2, 3), dtype="int32")],
|
|
)
|
|
return gv0
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def copy(
|
|
A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32"), C: T.Buffer((2, 3), "int32")
|
|
):
|
|
# copies the contents of C into A and B
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(C[ax0, ax1])
|
|
T.writes(A[ax0, ax1], B[ax0, ax1])
|
|
A[ax0, ax1] = C[ax0, ax1]
|
|
B[ax0, ax1] = C[ax0, ax1]
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
|
|
) -> R.Tuple(R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
_ = Expected.copy(x, y, z)
|
|
gv0 = (x, y)
|
|
return gv0
|
|
|
|
new_mod = relax.transform.CallTIRRewrite()(Input)
|
|
tvm.ir.assert_structural_equal(Expected["foo"], new_mod["foo"], map_free_vars=True)
|
|
|
|
|
|
def test_call_tir_inplace_some_new():
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@T.prim_func(s_tir=True)
|
|
def copy(
|
|
A: T.Buffer((2, 3), "int32"),
|
|
B: T.Buffer((2, 3), "int32"),
|
|
C: T.Buffer((2, 3), "int32"),
|
|
out1: T.Buffer((2, 3), "int32"),
|
|
out2: T.Buffer((2, 3), "int32"),
|
|
):
|
|
# copies the contents of C into A, out1, and out2
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(C[ax0, ax1])
|
|
T.writes(A[ax0, ax1], out1[ax0, ax1], out2[ax0, ax1])
|
|
A[ax0, ax1] = C[ax0, ax1]
|
|
out1[ax0, ax1] = C[ax0, ax1]
|
|
out2[ax0, ax1] = C[ax0, ax1]
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
|
|
) -> R.Tuple(
|
|
R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), dtype="int32")
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
gv0 = R.call_tir_inplace(
|
|
Input.copy,
|
|
(x, y, z),
|
|
[0, -1, -1],
|
|
[
|
|
R.Tensor((2, 3), dtype="int32"),
|
|
R.Tensor((2, 3), dtype="int32"),
|
|
R.Tensor((2, 3), dtype="int32"),
|
|
],
|
|
)
|
|
return gv0
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def copy(
|
|
A: T.Buffer((2, 3), "int32"),
|
|
B: T.Buffer((2, 3), "int32"),
|
|
C: T.Buffer((2, 3), "int32"),
|
|
out1: T.Buffer((2, 3), "int32"),
|
|
out2: T.Buffer((2, 3), "int32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i0, i1 in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(C[ax0, ax1])
|
|
T.writes(A[ax0, ax1], out1[ax0, ax1], out2[ax0, ax1])
|
|
A[ax0, ax1] = C[ax0, ax1]
|
|
out1[ax0, ax1] = C[ax0, ax1]
|
|
out2[ax0, ax1] = C[ax0, ax1]
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
|
|
) -> R.Tuple(
|
|
R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), dtype="int32")
|
|
):
|
|
R.func_attr({"relax.force_pure": True})
|
|
gv0: R.Tensor((2, 3), dtype="int32") = R.emit_with_ty(
|
|
"relax.builtin.alloc_tensor",
|
|
(R.shape([2, 3]), R.dtype("int32"), R.prim_value(0), R.str("global")),
|
|
(R.Tensor((2, 3), dtype="int32"),),
|
|
)
|
|
gv1: R.Tensor((2, 3), dtype="int32") = R.emit_with_ty(
|
|
"relax.builtin.alloc_tensor",
|
|
(R.shape([2, 3]), R.dtype("int32"), R.prim_value(0), R.str("global")),
|
|
(R.Tensor((2, 3), dtype="int32"),),
|
|
)
|
|
_ = Expected.copy(x, y, z, gv0, gv1)
|
|
gv2 = (x, gv0, gv1)
|
|
return gv2
|
|
|
|
new_mod = relax.transform.CallTIRRewrite()(Input)
|
|
tvm.ir.assert_structural_equal(Expected["foo"], new_mod["foo"], map_free_vars=True)
|
|
|
|
|
|
def test_call_tir_inplace_repeated_input():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@T.prim_func(s_tir=True)
|
|
def func(
|
|
A: T.Buffer((2, 3), "int32"),
|
|
B: T.Buffer((2, 3), "int32"),
|
|
C: T.Buffer((2, 3), "int32"),
|
|
):
|
|
T.evaluate(0)
|
|
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), "int32"),
|
|
y: R.Tensor((2, 3), "int32"),
|
|
z: R.Tensor((2, 3), "int32"),
|
|
) -> R.Tuple(R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
gv0 = R.call_tir_inplace(
|
|
Input.func,
|
|
(x, y, z),
|
|
# repeated 0 -> that's an error
|
|
[0, 0],
|
|
[R.Tensor((2, 3), dtype="int32"), R.Tensor((2, 3), dtype="int32")],
|
|
)
|
|
return gv0
|
|
|
|
|
|
def test_call_tir_inplace_all_new():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@tvm.script.ir_module
|
|
class Input:
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((2, 3), "int32")):
|
|
T.evaluate(0)
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
# cannot make the only output a fresh one
|
|
gv0 = R.call_tir_inplace(Input.func, x, -1, R.Tensor((2, 3), dtype="int32"))
|
|
return gv0
|
|
|
|
|
|
def test_inplace_mutation_with_tuple_argument_raises_error():
|
|
"""TIR PrimFuncs do not support Tuple arguments
|
|
|
|
The `R.call_tir_inplace` operator must receive an in-line tuple of
|
|
arguments, where each argument in the tuple may be expressed in
|
|
TIR. Here, `[[A]]` specifies a tuple of arguments, where the
|
|
first argument is itself a tuple. Since PrimFuncs do not support
|
|
Tuple arguments, this is invalid.
|
|
|
|
This is a regression test. In previous implementations, this
|
|
triggered a segfault rather than raising an exception.
|
|
|
|
"""
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(A: R.Tensor((16,), dtype="float32")) -> R.Tensor((16,), dtype="float32"):
|
|
cls = Module
|
|
gv1 = R.call_tir_inplace(
|
|
cls.multiply_by_two,
|
|
[[A]],
|
|
out_ty=R.Tensor((16,), dtype="float32"),
|
|
inplace_indices=[0],
|
|
)
|
|
return gv1
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def multiply_by_two(A: T.Buffer((16,), "float32")):
|
|
for i in range(16):
|
|
A[i] = A[i] * T.float32(2)
|
|
|
|
|
|
def test_inplace_mutation_with_non_tensor_argument_raises_error():
|
|
"""In-place argument must be a tensor
|
|
|
|
The `R.call_tir_inplace` operator must receive an in-line tuple of
|
|
arguments, where each argument in the tuple may be expressed in
|
|
TIR. Here, the argument `A` is not a tensor.
|
|
|
|
This is a regression test. In previous implementations, this
|
|
triggered a segfault rather than raising an exception.
|
|
|
|
"""
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(A: R.Any):
|
|
gv1 = R.call_tir_inplace(
|
|
Module.multiply_by_two,
|
|
[A],
|
|
out_ty=R.Tensor((16,), dtype="float32"),
|
|
inplace_indices=[0],
|
|
)
|
|
return gv1
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def multiply_by_two(A: T.Buffer((16,), "float32")):
|
|
for i in range(16):
|
|
A[i] = A[i] * T.float32(2)
|
|
|
|
|
|
def test_inplace_mutation_with_incompatible_tensor_shape_raises_error():
|
|
"""In-place argument must have compatible shape
|
|
|
|
The `R.call_tir_inplace` operator must receive an in-line tuple of
|
|
arguments, where the shape of each in-place argument is compatible
|
|
with the corresponding output. Here, the shape of argument `A` is
|
|
different than the output's shape (`[32]` as opposed to `[16]`).
|
|
|
|
"""
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(A: R.Tensor([32], dtype="float32")):
|
|
gv1 = R.call_tir_inplace(
|
|
Module.multiply_by_two,
|
|
[A],
|
|
out_ty=R.Tensor((16,), dtype="float32"),
|
|
inplace_indices=[0],
|
|
)
|
|
return gv1
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def multiply_by_two(A: T.Buffer((16,), "float32")):
|
|
for i in range(16):
|
|
A[i] = A[i] * T.float32(2)
|
|
|
|
|
|
def test_inplace_mutation_with_incompatible_tensor_dtype_raises_error():
|
|
"""In-place argument must have compatible dtype
|
|
|
|
The `R.call_tir_inplace` operator must receive an in-line tuple of
|
|
arguments, where the shape of each in-place argument is compatible
|
|
with the corresponding output. Here, the dtype of argument `A` is
|
|
different than the output's dtype (`int32` as opposed to `float32`).
|
|
|
|
"""
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(A: R.Tensor([16], dtype="int32")):
|
|
gv1 = R.call_tir_inplace(
|
|
Module.multiply_by_two,
|
|
[A],
|
|
out_ty=R.Tensor((16,), dtype="float32"),
|
|
inplace_indices=[0],
|
|
)
|
|
return gv1
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def multiply_by_two(A: T.Buffer((16,), "float32")):
|
|
for i in range(16):
|
|
A[i] = A[i] * T.float32(2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|