wehub-resource-sync
635 lines
22 KiB
Python
635 lines
22 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 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):
|
|
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):
|
|
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_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_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()
|