wehub-resource-sync
2397 lines
75 KiB
Python
2397 lines
75 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: F401, F821, F841
|
|
import sys
|
|
from typing import Optional, Union
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.script
|
|
import tvm.testing
|
|
from tvm import IRModule, relax, tirx, topi
|
|
from tvm.ir import DummyGlobalInfo, VDevice
|
|
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 _check(
|
|
parsed: relax.Function | IRModule,
|
|
expect: relax.Function | IRModule | None = None,
|
|
):
|
|
test = parsed.script(show_meta=True)
|
|
roundtrip_mod = tvm.script.from_source(test)
|
|
tvm.ir.assert_structural_equal(parsed, roundtrip_mod)
|
|
if isinstance(parsed, IRModule) and isinstance(roundtrip_mod, IRModule):
|
|
relax.analysis.well_formed(parsed)
|
|
relax.analysis.well_formed(roundtrip_mod)
|
|
if expect:
|
|
tvm.ir.assert_structural_equal(parsed, expect)
|
|
|
|
|
|
def test_simple_func():
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
|
|
R.func_attr({"Primitive": True})
|
|
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
gv1 = R.call_dps_packed("extern_dps_func", gv0, R.Tensor((128, 128), dtype="float32"))
|
|
return gv1
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,), attrs={"Primitive": True}):
|
|
y = bb.emit(relax.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32")))
|
|
out = bb.emit(
|
|
relax.call_dps_packed("extern_dps_func", y, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
bb.emit_func_output(out)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_error_report():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
# error: a = b = c is not allowed.
|
|
gv0 = gv1 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
return gv0
|
|
|
|
|
|
def test_mismatch_cast_dims_and_ndim():
|
|
with pytest.raises(Exception):
|
|
|
|
@R.function
|
|
def f(
|
|
x: R.Tensor((2, 3), "float32", ndim=3),
|
|
): # error: ndim and the shape dims are mismatch
|
|
return x
|
|
|
|
|
|
def test_unexpected_num_kw_args():
|
|
with pytest.raises(Exception):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor(dtype="float32", ndim=1, foo=2)): # error: unexpected kw args foo
|
|
return x
|
|
|
|
|
|
def test_unexpected_ndim():
|
|
with pytest.raises(Exception):
|
|
|
|
@R.function
|
|
# error: dim is expected to be non-negative int or -1 for unknown
|
|
def f(x: R.Tensor(dtype="float32", ndim=-2)):
|
|
return x
|
|
|
|
|
|
def test_unexpected_ndim_type():
|
|
with pytest.raises(Exception):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor(dtype="float32", ndim="1")): # error: dim is expected to be int
|
|
return x
|
|
|
|
|
|
def test_unexpected_tir_cast_args():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor(("m",), "float32")):
|
|
m = T.int64()
|
|
# tirx.cast expects 2 arguments, but got 3
|
|
return R.call_tir("foo", (x,), R.Tensor((T.cast("int32", m, 1),), dtype="float32"))
|
|
|
|
|
|
def test_unexpected_tir_args():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@tvm.script.ir_module
|
|
class TestWellCallTIR:
|
|
@T.prim_func(s_tir=True)
|
|
def tir_addone(A: T.Buffer((16, 16), "int32"), B: T.Buffer((16, 16), "int32")) -> None:
|
|
T.func_attr({"global_symbol": "tir_addone"})
|
|
for i, j in T.grid(16, 16):
|
|
with T.sblock("tir_addone"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] + T.int32(1)
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "m"), "float32")):
|
|
m = T.int64()
|
|
# tirx.max expects 2 arguments, but got 1
|
|
gv = R.call_tir(tir_addone, (x,), R.Tensor((T.max(16),), dtype="float32"))
|
|
return gv
|
|
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor(("m", "n"), "float32")):
|
|
m = T.int64()
|
|
# call_tir expected a tirx prim_func
|
|
return relax.call_tir("extern_func", (x,), R.Tensor((T.max(m),), dtype="float32"))
|
|
|
|
|
|
def test_func_type_annotation_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(x, y): # error: the parameter type annotation is missing
|
|
z = R.add(x, y)
|
|
y = z
|
|
return y
|
|
|
|
|
|
def test_if_mismatch_var_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
|
|
if cond:
|
|
w = R.add(x, x)
|
|
y = R.multiply(w, w)
|
|
else:
|
|
w = R.multiply(x, x)
|
|
z = R.add(w, w) # error: The binding var is expected to `y`
|
|
return z
|
|
|
|
|
|
def test_unassigned_call_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor):
|
|
R.add(x, x)
|
|
return x
|
|
|
|
|
|
def test_incorrect_tensor_shape():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor([16])):
|
|
y: R.Tensor(16) = R.add(x, x)
|
|
return y
|
|
|
|
|
|
def test_simple_module():
|
|
@I.ir_module(s_tir=True)
|
|
class TestModule:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def tir_func(
|
|
x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i, j in T.grid(T.int64(128), T.int64(128)):
|
|
with T.sblock():
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
y[vi, vj] = x[vi, vj] + 1.0
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
|
|
cls = TestModule
|
|
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
|
|
return gv0
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,), {"global_symbol": "foo"}):
|
|
out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
|
|
bb.emit_func_output(out)
|
|
|
|
_check(TestModule, bb.get())
|
|
|
|
|
|
def test_emit_te_primfunc_attrs():
|
|
@I.ir_module(s_tir=True)
|
|
class TestModule:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def plus_one(
|
|
x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
):
|
|
T.func_attr({"some_attr": "foo", "another_attr": True, "tirx.noalias": True})
|
|
for i, j in T.grid(T.int64(128), T.int64(128)):
|
|
with T.sblock():
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
y[vi, vj] = x[vi, vj] + 1.0
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
|
|
cls = TestModule
|
|
gv0 = R.call_tir(cls.plus_one, x, R.Tensor((128, 128), dtype="float32"))
|
|
return gv0
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,), {"global_symbol": "foo"}):
|
|
out = bb.emit_te(
|
|
lambda x: x + 1,
|
|
x,
|
|
primfunc_name_hint="plus_one",
|
|
primfunc_attrs={"some_attr": "foo", "another_attr": True},
|
|
)
|
|
bb.emit_func_output(out)
|
|
_check(TestModule, bb.get())
|
|
|
|
|
|
def test_emit_te():
|
|
@I.ir_module(s_tir=True)
|
|
class EmitTE:
|
|
@R.function
|
|
def main(x: R.Tensor((10, 20), "float32")) -> R.Tensor((10, 20), dtype="float32"):
|
|
lv1 = R.emit_te(topi.add, x, x)
|
|
out = R.emit_te(topi.multiply, lv1, lv1)
|
|
return out
|
|
|
|
bb = relax.BlockBuilder()
|
|
x = relax.Var("x", relax.TensorType([10, 20], "float32"))
|
|
with bb.function("main", [x], {"global_symbol": "main"}):
|
|
lv1 = bb.emit_te(topi.add, x, x)
|
|
out = bb.emit_te(topi.multiply, lv1, lv1)
|
|
bb.emit_func_output(out)
|
|
|
|
_check(EmitTE, bb.get())
|
|
|
|
|
|
def test_module_with_attr_and_global_info():
|
|
@I.ir_module(s_tir=True)
|
|
class TestModule:
|
|
I.module_attrs({"attr": 10})
|
|
I.module_global_infos(
|
|
{
|
|
"dummy": [
|
|
I.dummy_global_info(), # dummy[0]
|
|
I.dummy_global_info(), # dummy[1]
|
|
]
|
|
}
|
|
)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def tir_func(
|
|
x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i, j in T.grid(T.int64(128), T.int64(128)):
|
|
with T.sblock():
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
y[vi, vj] = x[vi, vj] + 1.0
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
|
|
cls = TestModule
|
|
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
|
|
return gv0
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,), {"global_symbol": "foo"}):
|
|
out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
|
|
bb.emit_func_output(out)
|
|
mod = bb.get()
|
|
mod.update_global_info("dummy", [DummyGlobalInfo(), DummyGlobalInfo()])
|
|
mod = mod.with_attr("attr", 10)
|
|
_check(TestModule, mod)
|
|
|
|
|
|
def test_global_info_vdevice():
|
|
vdevices = [
|
|
VDevice("llvm"),
|
|
VDevice("cuda", 0),
|
|
VDevice({"kind": "cuda", "arch": "sm_80"}, 0),
|
|
VDevice("metal", 0, "global"),
|
|
]
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class TestModule:
|
|
I.module_attrs({"attr": 10})
|
|
I.module_global_infos(
|
|
{
|
|
"vdevice": [
|
|
I.vdevice("llvm"),
|
|
I.vdevice("cuda", 0),
|
|
I.vdevice({"kind": "cuda", "arch": "sm_80"}, 0),
|
|
I.vdevice("metal", 0, "global"),
|
|
]
|
|
}
|
|
)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def tir_func(
|
|
x: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
y: T.Buffer((T.int64(128), T.int64(128)), "float32"),
|
|
):
|
|
T.func_attr({"tirx.noalias": True})
|
|
for i, j in T.grid(T.int64(128), T.int64(128)):
|
|
with T.sblock():
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
y[vi, vj] = x[vi, vj] + 1.0
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
|
|
cls = TestModule
|
|
gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32"))
|
|
return gv0
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func")
|
|
bb.emit_func_output(out)
|
|
mod = bb.get()
|
|
mod.update_global_info("vdevice", vdevices)
|
|
mod = mod.with_attr("attr", 10)
|
|
_check(TestModule, mod)
|
|
|
|
|
|
def test_relax_tensor_op():
|
|
@R.function
|
|
def foo(x: R.Tensor((4, 4), "float32")) -> R.Tensor((4, 4), "float32"):
|
|
y = R.add(x, x)
|
|
z = R.multiply(x, y)
|
|
return z
|
|
|
|
x = relax.Var("x", R.Tensor((4, 4), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
y = bb.emit(relax.op.add(x, x))
|
|
z = bb.emit(relax.op.multiply(x, y))
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_relax_base_op():
|
|
@R.function
|
|
def foo(x: R.Tensor((4, 4), "float32")):
|
|
alloc = R.builtin.alloc_tensor(R.shape([4, 4]), runtime_device_index=0, dtype="float32")
|
|
shape = R.shape_of(alloc)
|
|
return shape
|
|
|
|
x = relax.Var("x", R.Tensor((4, 4), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
alloc = bb.emit(relax.op.builtin.alloc_tensor(relax.ShapeExpr((4, 4)), "float32", 0))
|
|
shape = bb.emit(relax.op.shape_of(alloc))
|
|
bb.emit_func_output(shape)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_relax_shape_to_tensor():
|
|
@R.function
|
|
def foo(x: R.Shape((4, 4))):
|
|
tensor = R.shape_to_tensor(x)
|
|
return tensor
|
|
|
|
x = relax.Var("x", R.Shape((4, 4)))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
tensor = bb.emit(relax.op.shape_to_tensor(x))
|
|
bb.emit_func_output(tensor)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_symbolic_shape():
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
|
|
return gv0
|
|
|
|
@R.function
|
|
def bar(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
|
|
return gv0
|
|
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def mismatch_dtype(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
m = T.int64()
|
|
n = T.int32() # The shape dtype should be int64
|
|
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
|
|
return gv0
|
|
|
|
def _expected(name: str):
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = relax.Var("x", R.Tensor([m, n], "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function(name, (x,)):
|
|
out = bb.emit(
|
|
relax.call_dps_packed("extern_func", x, R.Tensor((m, n), dtype="float32"))
|
|
)
|
|
bb.emit_func_output(out)
|
|
return bb.get()[name]
|
|
|
|
_check(foo, _expected("foo"))
|
|
_check(bar, _expected("bar"))
|
|
|
|
|
|
def test_shadowing():
|
|
@R.function
|
|
def foo(x: R.Tensor((4, 4), "float32")):
|
|
y = R.add(x, x)
|
|
z = R.multiply(x, y)
|
|
y = R.add(x, y)
|
|
y = z
|
|
y = R.multiply(y, x)
|
|
z = y
|
|
return z
|
|
|
|
x = relax.Var("x", R.Tensor((4, 4), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
y = bb.emit(relax.op.add(x, x))
|
|
z = bb.emit(relax.op.multiply(x, y))
|
|
y = bb.emit(relax.op.add(x, y))
|
|
y = bb.emit(z)
|
|
y = bb.emit(relax.op.multiply(y, x))
|
|
z = bb.emit(y)
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_match_cast():
|
|
@R.function
|
|
def foo(x: R.Tensor("float32"), y: R.Tensor("float32")):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
x0 = R.match_cast(x, R.Tensor([m], "float32"))
|
|
with R.dataflow():
|
|
y0 = R.match_cast(y, R.Tensor([n], "float32"))
|
|
gv = y0
|
|
R.output(gv)
|
|
return (x0, R.shape([m, n * 2]))
|
|
|
|
x = relax.Var("x", R.Tensor("float32"))
|
|
y = relax.Var("y", R.Tensor("float32"))
|
|
m = tirx.Var("m", dtype="int64")
|
|
n = tirx.Var("n", dtype="int64")
|
|
y2 = relax.Var("y", R.Tensor([n], "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x, y)):
|
|
x0 = bb.match_cast(x, R.Tensor([m], "float32"))
|
|
with bb.dataflow():
|
|
y0 = bb.match_cast(y, R.Tensor([n], "float32"))
|
|
bb.emit_output(y0)
|
|
bb.emit_func_output(relax.Tuple([x0, relax.ShapeExpr([m, n * 2])]))
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_tuple_return():
|
|
@R.function
|
|
def foo(x: R.Tensor((4, 4), "float32")):
|
|
gv0 = R.call_dps_packed("extern_func_0", x, R.Tensor((4, 4), dtype="float32"))
|
|
gv1 = R.call_dps_packed("extern_func_1", x, R.Tensor((4, 4), dtype="float32"))
|
|
return (gv0, gv1)
|
|
|
|
x = relax.Var("x", R.Tensor((4, 4), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
gv0 = bb.emit(relax.call_dps_packed("extern_func_0", x, R.Tensor((4, 4), dtype="float32")))
|
|
gv1 = bb.emit(relax.call_dps_packed("extern_func_1", x, R.Tensor((4, 4), dtype="float32")))
|
|
bb.emit_func_output(relax.Tuple((gv0, gv1)))
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_tuple_return_2():
|
|
@R.function
|
|
def foo(x: R.Tensor("float32", ndim=2)):
|
|
n, m = T.int64(), T.int64()
|
|
x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
|
|
return (x0, R.shape([n + 1, m, 1]))
|
|
|
|
x = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
x0 = bb.match_cast(x, R.Tensor((n, m), "float32"))
|
|
bb.emit_func_output(relax.Tuple([x0, relax.ShapeExpr([n + 1, m, 1])]))
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_tuple_binding():
|
|
@R.function
|
|
def foo(x: R.Tensor("float32", ndim=2)):
|
|
n, m = T.int64(), T.int64()
|
|
x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
|
|
t0 = (x, x0)
|
|
t1 = (x, R.shape([n, m]), t0)
|
|
return t1
|
|
|
|
x = relax.Var("x", R.Tensor("float32", ndim=2))
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
x0 = bb.match_cast(x, R.Tensor((n, m), "float32"))
|
|
t0 = bb.emit(relax.Tuple([x, x0]))
|
|
t1 = bb.emit(relax.Tuple([x, relax.ShapeExpr([n, m]), t0]))
|
|
bb.emit_func_output(t1)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_tuple_get_item():
|
|
@R.function
|
|
def foo(x: R.Tensor, y: R.Tensor):
|
|
t1 = R.tuple(x, y)
|
|
t2 = (x, y)
|
|
a = t1[0]
|
|
b = R.TupleGetItem(t2, 1)
|
|
c = R.add(a, b)
|
|
return c
|
|
|
|
x = relax.Var("x", R.Tensor())
|
|
y = relax.Var("y", R.Tensor())
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x, y)):
|
|
t1 = bb.emit(relax.Tuple([x, y]))
|
|
t2 = bb.emit(relax.Tuple([x, y]))
|
|
a = bb.emit(relax.TupleGetItem(t1, 0))
|
|
b = bb.emit(relax.TupleGetItem(t2, 1))
|
|
c = bb.emit(relax.op.add(a, b))
|
|
bb.emit_func_output(c)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_dataflow_block():
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
with R.dataflow():
|
|
lv0 = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
lv1 = R.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
|
|
gv = lv1
|
|
R.output(gv)
|
|
return gv
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(
|
|
relax.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
lv1 = bb.emit(
|
|
relax.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
gv = bb.emit_output(lv1)
|
|
bb.emit_func_output(gv)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_dataflow_block_advanced():
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
gv0 = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
gv1 = R.call_dps_packed("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
|
|
with R.dataflow():
|
|
m = T.int64()
|
|
n = T.int64()
|
|
lv0 = R.call_dps_packed("extern_func", gv1, R.Tensor((128, 128), dtype="float32"))
|
|
lv1 = R.match_cast(lv0, R.Tensor((m, n), "float32"))
|
|
gv2 = R.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
|
|
gv2 = R.call_dps_packed("extern_func", gv2, R.Tensor((128, 128), dtype="float32"))
|
|
gv3 = R.match_cast(gv2, R.Tensor((m, n), "float32"))
|
|
gv3 = R.match_cast(lv0, R.Tensor((m, n), "float32"))
|
|
gv4 = gv3
|
|
gv5 = gv2
|
|
R.output(gv5, gv4)
|
|
gv6 = R.call_dps_packed("extern_func", gv5, R.Tensor((128, 128), dtype="float32"))
|
|
gv7 = R.call_dps_packed("extern_func", gv6, R.Tensor((128, 128), dtype="float32"))
|
|
return gv7
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
m = tirx.Var("m", dtype="int64")
|
|
n = tirx.Var("n", dtype="int64")
|
|
with bb.function("foo", (x,)):
|
|
gv0 = bb.emit(
|
|
relax.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
gv1 = bb.emit(
|
|
relax.call_dps_packed("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(
|
|
relax.call_dps_packed("extern_func", gv1, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
lv1 = bb.match_cast(lv0, R.Tensor((m, n), "float32"))
|
|
gv2 = bb.emit(
|
|
relax.call_dps_packed("extern_func", lv0, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
gv21 = bb.emit(
|
|
relax.call_dps_packed("extern_func", gv2, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
gv3 = bb.match_cast(gv21, R.Tensor((m, n), "float32"))
|
|
gv31 = bb.match_cast(lv0, R.Tensor((m, n), "float32"))
|
|
gv32 = bb.emit_output(gv31)
|
|
gv22 = bb.emit_output(gv21)
|
|
gv4 = bb.emit(
|
|
relax.call_dps_packed("extern_func", gv22, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
gv5 = bb.emit(
|
|
relax.call_dps_packed("extern_func", gv4, R.Tensor((128, 128), dtype="float32"))
|
|
)
|
|
bb.emit_func_output(gv5)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_dataflow_binding_after_output():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
with R.dataflow():
|
|
gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
R.output(gv)
|
|
lv = R.call_tir("extern_func", gv, R.Tensor((128, 128), dtype="float32"))
|
|
return gv
|
|
|
|
|
|
def test_dataflow_output_global_var():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
with R.dataflow():
|
|
gv1 = R.call_tir("extern_func", gv0, R.Tensor((128, 128), dtype="float32"))
|
|
R.output(gv0, gv1)
|
|
return gv1
|
|
|
|
|
|
def test_dataflow_multiple_output():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
with R.dataflow():
|
|
gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
R.output(gv)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
|
|
def test_dataflow_output_outside_dataflow_block():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
|
|
gv = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
|
|
def test_dataflow_scope_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(x: R.Tensor(ndim=2)):
|
|
with R.dataflow():
|
|
y = R.add(x, x)
|
|
z = R.multiply(y, x)
|
|
w = R.add(z, x)
|
|
R.output(y, w)
|
|
t = R.multiply(y, z) # z is not in the outer scope
|
|
return t
|
|
|
|
|
|
def test_return_without_binding():
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")):
|
|
return x
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
bb.emit_func_output(x)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_multiple_return():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")):
|
|
return x
|
|
return x
|
|
|
|
|
|
def test_function_without_return():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")):
|
|
gv0 = R.call_tir("extern_func", x, R.Tensor((128, 128), dtype="float32"))
|
|
|
|
|
|
def test_tensor_type_without_args():
|
|
@R.function
|
|
def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
|
v = R.call_dps_packed("extern_relu", x, R.Tensor((32, 32), dtype="float32"))
|
|
return v
|
|
|
|
x = relax.Var("x", R.Tensor((32, 32), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x)):
|
|
v = bb.emit(relax.call_dps_packed("extern_relu", x, R.Tensor((32, 32), dtype="float32")))
|
|
bb.emit_func_output(v)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_tensor_with_vdevice():
|
|
vdevices = [
|
|
VDevice("llvm"),
|
|
VDevice("cuda", 0),
|
|
VDevice("metal", 0, "global"),
|
|
VDevice({"kind": "cuda", "arch": "sm_80"}, 0),
|
|
]
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class TestModule:
|
|
I.module_attrs({"attr": 10})
|
|
I.module_global_infos(
|
|
{
|
|
"vdevice": [
|
|
I.vdevice("llvm"),
|
|
I.vdevice("cuda", 0),
|
|
I.vdevice("metal", 0, "global"),
|
|
I.vdevice({"kind": "cuda", "arch": "sm_80"}, 0),
|
|
]
|
|
}
|
|
)
|
|
|
|
@R.function
|
|
def foo(
|
|
a: R.Tensor((128, 128), "float32", "cuda:1"),
|
|
b: R.Tensor((128, 128), "float32", "llvm"),
|
|
c: R.Tensor((128, 128), "float32", "vdevice:3"),
|
|
) -> R.Tensor((128, 128), "float32", "cuda:1"):
|
|
s = R.add(a, c)
|
|
return s
|
|
|
|
a = relax.Var("a", R.Tensor((128, 128), "float32", vdevices[3]))
|
|
b = relax.Var("b", R.Tensor((128, 128), "float32", vdevices[0]))
|
|
c = relax.Var("c", R.Tensor((128, 128), "float32", vdevices[3]))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (a, b, c)):
|
|
out = bb.emit(relax.op.add(a, c))
|
|
bb.emit_func_output(out)
|
|
mod = bb.get()
|
|
mod = mod.with_attr("attr", 10)
|
|
mod.update_global_info("vdevice", vdevices)
|
|
|
|
_check(TestModule, mod)
|
|
|
|
|
|
def test_direct_return():
|
|
@R.function
|
|
def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
|
|
return x
|
|
|
|
x = relax.Var("x", R.Tensor((32, 32), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x)):
|
|
bb.emit_func_output(x)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_call_packed():
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
|
z = R.call_packed("vm.builtin.copy", x, ty_args=R.Tensor((32, 32), "float32"))
|
|
return z
|
|
|
|
x = relax.Var("x", R.Tensor((32, 32), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x), pure=False):
|
|
z = bb.emit(
|
|
relax.Call(
|
|
relax.ExternFunc("vm.builtin.copy"),
|
|
(x,),
|
|
None,
|
|
ty_args=[R.Tensor((32, 32), "float32")],
|
|
)
|
|
)
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_call_packed_without_ty_args():
|
|
@R.function(pure=False)
|
|
def foo(x: R.Any) -> R.Any:
|
|
z = R.call_packed("test", x)
|
|
return z
|
|
|
|
x = relax.Var("x", R.Any())
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x), pure=False):
|
|
z = bb.emit(
|
|
relax.Call(
|
|
relax.ExternFunc("test"),
|
|
(x,),
|
|
None,
|
|
ty_args=[],
|
|
)
|
|
)
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_object_proxy_compat_alias():
|
|
@R.function
|
|
def foo(x: R.Object) -> R.Object:
|
|
return x
|
|
|
|
assert isinstance(foo.ret_ty, relax.AnyType)
|
|
|
|
|
|
def test_annotation():
|
|
@R.function(pure=False)
|
|
def foo(
|
|
x: R.Tensor((32, "m"), "float32"),
|
|
y: R.Tensor(("m",), "float32"),
|
|
r: R.Tensor(dtype="int64"),
|
|
) -> R.Any:
|
|
m = T.int64()
|
|
z: R.Tensor((32, m), "float32") = R.multiply(x, y)
|
|
w: R.Tensor(ndim=2) = R.multiply(z, z)
|
|
q: R.Tensor = R.add(w, w)
|
|
t = R.add(w, z)
|
|
sh: R.Shape = R.call_packed("shape_of", x, ty_args=R.Shape)
|
|
lv: R.Tensor(sh, dtype="float32") = R.reshape(x, sh)
|
|
o: R.Any = R.call_packed("contrib.tensor_array_stack", x, y, ty_args=R.Any)
|
|
return o
|
|
|
|
def _check_ty(binding, expected_ty):
|
|
tvm.ir.assert_structural_equal(binding.var.ty, expected_ty)
|
|
tvm.ir.assert_structural_equal(binding.value.ty, expected_ty)
|
|
|
|
# Cannot use block builder here because we need to check the annotated type,
|
|
# which may be inconsistent with deduced type.
|
|
assert isinstance(foo.ret_ty, relax.AnyType)
|
|
m = relax.get_shape_of(foo.params[0])[1]
|
|
bindings = foo.body.blocks[0].bindings
|
|
sh = bindings[4].var
|
|
|
|
_check_ty(bindings[0], relax.TensorType([32, m], "float32"))
|
|
_check_ty(bindings[1], relax.TensorType(dtype=None, ndim=2))
|
|
_check_ty(bindings[2], relax.TensorType(dtype=None, ndim=-1))
|
|
_check_ty(bindings[3], relax.TensorType(dtype=None, ndim=2))
|
|
_check_ty(bindings[4], relax.ShapeType(ndim=-1))
|
|
_check_ty(bindings[5], relax.TensorType(sh))
|
|
_check_ty(bindings[6], relax.AnyType())
|
|
|
|
|
|
def test_annotate_override():
|
|
@R.function
|
|
def foo(x: R.Tensor):
|
|
y = x
|
|
# z will be treated as Any even though it's a tensor
|
|
z: R.Any = R.add(x, y)
|
|
return z
|
|
|
|
assert isinstance(foo.ret_ty, relax.AnyType)
|
|
y_bind, z_bind = foo.body.blocks[0].bindings
|
|
assert isinstance(y_bind.var.ty, relax.TensorType)
|
|
assert isinstance(z_bind.var.ty, relax.AnyType)
|
|
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def test(x: R.Tensor):
|
|
# Error: x is of Tensor Type, which can not annotate to R.Shape.
|
|
z: R.Shape = x
|
|
return z
|
|
|
|
@R.function
|
|
def bar(x: R.Tensor):
|
|
# x is of Tensor Type, the annotation of `z` is ignored.
|
|
z: R.Any = x
|
|
return z
|
|
|
|
assert isinstance(bar.ret_ty, relax.TensorType)
|
|
(z_bind,) = bar.body.blocks[0].bindings
|
|
assert isinstance(z_bind.var.ty, relax.TensorType)
|
|
|
|
|
|
def test_call_dps_packed_empty_shape():
|
|
@R.function
|
|
def foo(x: R.Tensor((), "float32")):
|
|
z = R.call_dps_packed("scalar_add", x, R.Tensor((), dtype="float32"))
|
|
return z
|
|
|
|
(z_bind,) = foo.body.blocks[0].bindings
|
|
shape_expr = z_bind.value.ty_args[0].shape
|
|
|
|
assert isinstance(shape_expr, relax.ShapeExpr)
|
|
assert len(shape_expr.values) == 0
|
|
|
|
|
|
def test_call_tir_empty_tuple_arg():
|
|
bb = relax.BlockBuilder()
|
|
dummy_param = relax.Var("dummy_param", R.Tensor(()))
|
|
with bb.function("foo", [dummy_param], {"global_symbol": "foo"}):
|
|
output = bb.emit_te(topi.full, shape=(16, 32), dtype="float32", fill_value=1.0)
|
|
bb.emit_func_output(output)
|
|
|
|
_check(bb.get())
|
|
|
|
|
|
def test_call_tir_with_tir_var():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(
|
|
dumb_param: R.Tensor(("n",), "float32"), x: R.Tensor(("n * 2",), "float32")
|
|
) -> R.Tensor(("n * 2",), "float32"):
|
|
n = T.int64()
|
|
cls = Module
|
|
y = R.call_tir(cls.copy, x, R.Tensor((n * 2,), dtype="float32"), tir_vars=(n,))
|
|
return y
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def copy(var_x: T.handle, var_y: T.handle, n: T.int64):
|
|
X = T.match_buffer(var_x, (n * 2,), dtype="float32")
|
|
Y = T.match_buffer(var_y, (n * 2,), dtype="float32")
|
|
for i in T.grid(n * 2):
|
|
with T.sblock("block"):
|
|
vi = T.axis.remap("S", [i])
|
|
Y[vi] = X[vi]
|
|
|
|
_check(Module)
|
|
|
|
|
|
def test_call_tir_with_grad():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def identity_tir(a: T.handle, b: T.handle) -> None:
|
|
A = T.match_buffer(a, [54, 96])
|
|
B = T.match_buffer(b, [54, 96])
|
|
|
|
for i, j in T.grid(54, 96):
|
|
with T.sblock("compute"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj]
|
|
|
|
@R.function
|
|
def main(v0: R.Tensor([54, 96], "float32")):
|
|
cls = Module
|
|
out = R.call_tir_with_grad(
|
|
cls.identity_tir,
|
|
(v0,),
|
|
R.Tensor((54, 96), "float32"),
|
|
te_grad_name="identity_k_grad",
|
|
te_grad_kwargs={"k": 1.0},
|
|
)
|
|
return out
|
|
|
|
_check(Module)
|
|
|
|
|
|
def test_call_tir_inplace():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def copy(
|
|
A: T.Buffer((2, 3), "int32"),
|
|
B: T.Buffer((2, 3), "int32"),
|
|
out1: T.Buffer((2, 3), "int32"),
|
|
):
|
|
# copies the contents of B into A and out1
|
|
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(B[ax0, ax1])
|
|
T.writes(A[ax0, ax1], out1[ax0, ax1])
|
|
A[ax0, ax1] = B[ax0, ax1]
|
|
out1[ax0, ax1] = B[ax0, ax1]
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")) -> R.Tuple(
|
|
R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")
|
|
):
|
|
res = R.call_tir_inplace(
|
|
Module.copy,
|
|
(x, y),
|
|
[0, -1],
|
|
[R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")],
|
|
)
|
|
return res
|
|
|
|
_check(Module)
|
|
|
|
|
|
def test_call_tir_inplace_with_tuple_var_raises_error():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")):
|
|
cls = Module
|
|
args = (x, y)
|
|
res = R.call_tir_inplace(
|
|
cls.copy,
|
|
# The `args` tuple must be an in-line tuple, not a
|
|
# reference to a tuple. This error should be
|
|
# caught and raised during parsing.
|
|
args,
|
|
inplace_indices=[0, -1],
|
|
out_ty=[R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")],
|
|
)
|
|
return res
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def copy(
|
|
A: T.Buffer((2, 3), "int32"),
|
|
B: T.Buffer((2, 3), "int32"),
|
|
out1: T.Buffer((2, 3), "int32"),
|
|
):
|
|
# copies the contents of B into A and out1
|
|
T.func_attr({"tirx.noalias": True})
|
|
for iters in T.grid(T.int64(2), T.int64(3)):
|
|
with T.sblock("T_zeros"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
A[i, j] = B[i, j]
|
|
out1[i, j] = B[i, j]
|
|
|
|
|
|
def test_local_function():
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")) -> R.Tensor(
|
|
(2, 3), "float32"
|
|
):
|
|
@R.function
|
|
def outer_func(
|
|
c1: R.Tensor((2, 3), "float32"),
|
|
) -> R.Callable((R.Tensor(None, "float32", ndim=2),), R.Tensor(None, "float32", ndim=2)):
|
|
@R.function
|
|
def inner_func(x1: R.Tensor((2, 3), "float32")):
|
|
s: R.Tensor((2, 3), "float32") = R.add(x1, c1)
|
|
return s
|
|
|
|
return inner_func
|
|
|
|
in_call = outer_func(x)
|
|
res = in_call(y)
|
|
return res
|
|
|
|
main_bindings = main.body.blocks[0].bindings
|
|
assert len(main_bindings) == 3
|
|
outer_func = main_bindings[0].value
|
|
assert isinstance(outer_func, relax.Function)
|
|
|
|
outer_func_bindings = outer_func.body.blocks[0].bindings
|
|
assert len(outer_func_bindings) == 1
|
|
inner_func = outer_func_bindings[0].value
|
|
assert isinstance(inner_func, relax.Function)
|
|
|
|
|
|
def test_inline_prim_func():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class TestModule:
|
|
@R.function
|
|
def f(x: R.Tensor((128, 128), "float32"), y: R.Tensor((128, 128), "float32")):
|
|
@T.prim_func(s_tir=True)
|
|
def my_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128))
|
|
B = T.match_buffer(b, (128, 128))
|
|
C = T.match_buffer(c, (128, 128))
|
|
|
|
for i, j, k in T.grid(128, 128, 128):
|
|
with T.sblock():
|
|
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
|
with T.init():
|
|
C[vi, vj] = 0.0
|
|
C[vi, vj] += A[vi, vk] * B[vj, vk]
|
|
|
|
z = relax.call_tir(my_matmul, (x, y), R.Tensor((128, 128), dtype="float32"))
|
|
return z
|
|
|
|
|
|
def test_cross_function_call():
|
|
@I.ir_module(s_tir=True)
|
|
class Mod0:
|
|
@R.function
|
|
def foo(x: R.Tensor((10, 5), "float32")):
|
|
s = R.add(x, x)
|
|
return s
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((10, 5), "float32")):
|
|
cls = Mod0
|
|
inner = cls.foo
|
|
gv1 = inner(x)
|
|
gv2 = Mod0.foo(x)
|
|
return (inner, gv1, gv2)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Mod1:
|
|
@R.function
|
|
def main(x: R.Tensor((10, 5), "float32")):
|
|
cls = Mod1
|
|
inner = cls.foo
|
|
gv1 = inner(x)
|
|
gv2 = Mod1.foo(x)
|
|
return (inner, gv1, gv2)
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor((10, 5), "float32")) -> R.Tensor((10, 5), "float32"):
|
|
s = R.add(x, x)
|
|
return s
|
|
|
|
|
|
def test_if_branch():
|
|
@R.function
|
|
def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")) -> R.Tensor((1,), "float32"):
|
|
if cond:
|
|
w = R.add(x, x)
|
|
y = R.multiply(w, w)
|
|
else:
|
|
w = R.multiply(x, x)
|
|
y = R.add(w, w)
|
|
return y
|
|
|
|
cond, x = foo.params
|
|
y_bind = foo.body.blocks[0].bindings[0]
|
|
y, ite = y_bind.var, y_bind.value
|
|
|
|
assert isinstance(y, relax.Var)
|
|
assert y.name_hint == "y"
|
|
|
|
assert isinstance(ite, relax.If)
|
|
assert isinstance(ite.true_branch, relax.SeqExpr)
|
|
assert isinstance(ite.false_branch, relax.SeqExpr)
|
|
|
|
def check_call(call, op, args):
|
|
assert isinstance(call, relax.Call)
|
|
if isinstance(op, str):
|
|
assert call.op.name == op
|
|
else:
|
|
assert call.op == op
|
|
tvm.ir.assert_structural_equal(call.args, args)
|
|
|
|
w_bind = ite.true_branch.blocks[0].bindings[0]
|
|
# the seq exprts in the branches are normalized to bind any call
|
|
# in the seq expr "body" to a var
|
|
y_bind = ite.true_branch.blocks[-1].bindings[-1]
|
|
assert w_bind.var.name_hint == "w"
|
|
check_call(w_bind.value, "relax.add", [x, x])
|
|
check_call(y_bind.value, "relax.multiply", [w_bind.var, w_bind.var])
|
|
|
|
w_bind = ite.false_branch.blocks[0].bindings[0]
|
|
y_bind = ite.false_branch.blocks[-1].bindings[-1]
|
|
assert w_bind.var.name_hint == "w"
|
|
check_call(w_bind.value, "relax.multiply", [x, x])
|
|
check_call(y_bind.value, "relax.add", [w_bind.var, w_bind.var])
|
|
|
|
|
|
def test_if_branch_with_match_cast():
|
|
"""The last branch of a relax::If node may be a MatchCast
|
|
|
|
This is a regression test. In previous implementations, using
|
|
R.match_cast as the last binding would cause a segfault while
|
|
parsing.
|
|
"""
|
|
|
|
@R.function
|
|
def func(A: R.Tensor([16, 16]), is_bfloat16: R.Prim("bool")):
|
|
if is_bfloat16:
|
|
A = R.match_cast(A, R.Tensor([16, 16], "bfloat16"))
|
|
B = A.astype("float16")
|
|
else:
|
|
B = R.match_cast(A, R.Tensor([16, 16], "float16"))
|
|
return B
|
|
|
|
A, is_bfloat16 = func.params
|
|
(block,) = func.body.blocks
|
|
(B_binding,) = block.bindings
|
|
|
|
B_var = B_binding.var
|
|
assert isinstance(B_var, relax.Var)
|
|
assert B_var.name_hint == "B"
|
|
|
|
if_then_else = B_binding.value
|
|
assert isinstance(if_then_else, relax.If)
|
|
assert isinstance(if_then_else.true_branch, relax.SeqExpr)
|
|
assert isinstance(if_then_else.false_branch, relax.SeqExpr)
|
|
|
|
else_branch = if_then_else.false_branch
|
|
(else_block,) = else_branch.blocks
|
|
|
|
assert isinstance(else_block.bindings[-1], relax.MatchCast)
|
|
|
|
# If the `R.match_cast` were removed, the function would infer the
|
|
# return value as `R.Tensor([16,16])`, with an unknown dtype.
|
|
# With the `R.match_cast` retained, the output dtype is known.
|
|
tvm.ir.assert_structural_equal(func.ret_ty, R.Tensor([16, 16], "float16"))
|
|
|
|
|
|
def test_if_inside_dataflow():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
|
|
with R.dataflow():
|
|
if cond:
|
|
w = R.add(x, x)
|
|
y = R.multiply(w, w)
|
|
else:
|
|
w = R.multiply(x, x)
|
|
y = R.add(w, w)
|
|
R.output(y)
|
|
return y
|
|
|
|
|
|
def test_var_if_scoping_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def f(cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32")):
|
|
if cond:
|
|
w = R.add(x, x)
|
|
y = R.multiply(w, w)
|
|
else:
|
|
w = R.multiply(x, x)
|
|
y = R.add(w, w)
|
|
return w # error: The w is not defined in the outer scope
|
|
|
|
|
|
def test_scalar_tensor_as_branch_condition():
|
|
"""Branch condition can be 0-d tensor"""
|
|
|
|
@R.function
|
|
def func(cond: R.Tensor([], "bool"), x: R.Tensor((1,), "float32")):
|
|
if cond:
|
|
out = R.add(x, x)
|
|
else:
|
|
out = R.multiply(x, x)
|
|
return out
|
|
|
|
if_else = func.body.blocks[0].bindings[0].value
|
|
assert isinstance(if_else.cond, relax.Var)
|
|
tvm.ir.assert_structural_equal(if_else.cond.ty, R.Tensor([], "bool"))
|
|
|
|
|
|
def test_prim_annotation_requires_dtype():
|
|
with pytest.raises(TypeError, match="missing 1 required positional argument: 'dtype'"):
|
|
R.Prim()
|
|
|
|
with pytest.raises(TypeError, match="unexpected keyword argument 'value'"):
|
|
R.Prim(value="n")
|
|
|
|
|
|
def test_prim_value_as_branch_condition():
|
|
"""In addition to scalar tensor, can use R.Prim condition"""
|
|
|
|
@R.function
|
|
def func(cond: R.Prim("bool"), x: R.Tensor((1,), "float32")):
|
|
if cond:
|
|
out = R.add(x, x)
|
|
else:
|
|
out = R.multiply(x, x)
|
|
return out
|
|
|
|
if_else = func.body.blocks[0].bindings[0].value
|
|
assert isinstance(if_else.cond, relax.Var)
|
|
tvm.ir.assert_structural_equal(if_else.cond.ty, R.Prim("bool"))
|
|
|
|
|
|
def test_computed_prim_value_as_branch_condition():
|
|
"""The R.Prim condition may be computed within the function"""
|
|
|
|
@R.function
|
|
def func(x: R.Tensor(["N"], "float32")):
|
|
N = T.int64()
|
|
if R.prim_value(N % 16 == 0):
|
|
out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
|
|
else:
|
|
out = R.call_pure_packed("slow_non_vectorized_impl", x, ty_args=[x.ty])
|
|
return out
|
|
|
|
N = func.params[0].ty.shape[0]
|
|
if_else = func.body.blocks[0].bindings[0].value
|
|
assert tvm.ir.is_prim_expr(if_else.cond)
|
|
tvm.ir.assert_structural_equal(N % 16 == 0, if_else.cond)
|
|
tvm.ir.assert_structural_equal(if_else.cond.ty, R.Prim("bool"))
|
|
|
|
|
|
def test_tir_expr_as_branch_condition():
|
|
"""Syntactic sugar, use Expr directly"""
|
|
|
|
@R.function(private=True)
|
|
def sugared(x: R.Tensor(["N"], "float32")):
|
|
N = T.int64()
|
|
if N % 16 == 0:
|
|
out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
|
|
else:
|
|
out = R.call_pure_packed("slow_non_vectorized_impl", x, ty_args=[x.ty])
|
|
return out
|
|
|
|
@R.function(private=True)
|
|
def unsugared(x: R.Tensor(["N"], "float32")):
|
|
N = T.int64()
|
|
if R.prim_value(N % 16 == 0):
|
|
out = R.call_pure_packed("fast_vectorized_impl", x, ty_args=[x.ty])
|
|
else:
|
|
out = R.call_pure_packed("slow_non_vectorized_impl", x, ty_args=[x.ty])
|
|
return out
|
|
|
|
tvm.ir.assert_structural_equal(unsugared, sugared)
|
|
|
|
|
|
def test_scalar_tensor_as_assert_condition():
|
|
"""Branch condition can be 0-d tensor"""
|
|
|
|
@R.function(pure=False)
|
|
def func(cond: R.Tensor([], "bool"), x: R.Tensor((1,), "float32")):
|
|
_ = R.assert_op(cond)
|
|
out = R.add(x, x)
|
|
return out
|
|
|
|
assert_op = func.body.blocks[0].bindings[0].value
|
|
condition = assert_op.args[0]
|
|
assert isinstance(condition, relax.Var)
|
|
tvm.ir.assert_structural_equal(condition.ty, R.Tensor([], "bool"))
|
|
|
|
|
|
def test_prim_value_as_assert_condition():
|
|
"""In addition to scalar tensor, can use R.Prim condition"""
|
|
|
|
@R.function(pure=False)
|
|
def func(cond: R.Prim("bool"), x: R.Tensor((1,), "float32")):
|
|
_ = R.assert_op(cond)
|
|
out = R.add(x, x)
|
|
return out
|
|
|
|
assert_op = func.body.blocks[0].bindings[0].value
|
|
condition = assert_op.args[0]
|
|
assert isinstance(condition, relax.Var)
|
|
tvm.ir.assert_structural_equal(condition.ty, R.Prim("bool"))
|
|
|
|
|
|
def test_computed_prim_value_as_assert_condition():
|
|
"""The R.Prim condition may be computed within the function"""
|
|
|
|
@R.function(pure=False)
|
|
def func(x: R.Tensor(["N"], "float32")):
|
|
N = T.int64()
|
|
_ = R.assert_op(R.prim_value(N % 16 == 0))
|
|
out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
|
|
return out
|
|
|
|
N = func.params[0].ty.shape[0]
|
|
assert_op = func.body.blocks[0].bindings[0].value
|
|
condition = assert_op.args[0]
|
|
assert tvm.ir.is_prim_expr(condition)
|
|
tvm.ir.assert_structural_equal(N % 16 == 0, condition)
|
|
tvm.ir.assert_structural_equal(condition.ty, R.Prim("bool"))
|
|
|
|
|
|
def test_tir_expr_as_assert_condition():
|
|
"""Syntactic sugar, use Expr directly"""
|
|
|
|
@R.function(pure=False, private=True)
|
|
def sugared(x: R.Tensor(["N"], "float32")):
|
|
N = T.int64()
|
|
_ = R.assert_op(N % 16 == 0)
|
|
out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
|
|
return out
|
|
|
|
@R.function(pure=False, private=True)
|
|
def unsugared(x: R.Tensor(["N"], "float32")):
|
|
N = T.int64()
|
|
_ = R.assert_op(R.prim_value(N % 16 == 0))
|
|
out = R.call_packed("fast_vectorized_impl", x, ty_args=[x.ty])
|
|
return out
|
|
|
|
tvm.ir.assert_structural_equal(unsugared, sugared)
|
|
|
|
|
|
def test_erase_to_well_defined_removes_internal_vars():
|
|
@R.function
|
|
def foo(x: R.Tensor):
|
|
q = x
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(q, R.Tensor((m, n)))
|
|
w = z
|
|
return w
|
|
|
|
tvm.ir.assert_structural_equal(foo.ret_ty, R.Tensor(ndim=2))
|
|
assert foo.ret_ty.shape is None
|
|
_check(foo)
|
|
|
|
|
|
def test_erase_to_well_defined_keeps_variables_exposed_by_tensor_shape():
|
|
@R.function
|
|
def foo(x: R.Tensor(["m", "n"])):
|
|
q = x
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(q, R.Tensor((m, n)))
|
|
w = z
|
|
return w
|
|
|
|
assert foo.ret_ty.shape is not None
|
|
_check(foo)
|
|
|
|
|
|
def test_erase_to_well_defined_keeps_variants_exposed_by_shape_expr():
|
|
@R.function
|
|
def foo(x: R.Tensor, _: R.Shape(["m", "n"])):
|
|
q = x
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(q, R.Tensor((m, n)))
|
|
w = z
|
|
return w
|
|
|
|
assert foo.ret_ty.shape is not None
|
|
_check(foo)
|
|
|
|
|
|
def test_erase_to_well_defined_infers_from_shape_expr():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
# The subroutine's symbolic variables are only in-scope for the subroutine.
|
|
@R.function
|
|
def subroutine(x: R.Tensor, _: R.Shape(["m", "n"])) -> R.Tensor(["m", "n"]):
|
|
q = x
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(q, R.Tensor((m, n)))
|
|
w = z
|
|
return w
|
|
|
|
# However, struct inference can make the symbolic variables in
|
|
# the main function to the symbolic variables in the
|
|
# subroutine. Therefore, the shape of the tensor returned
|
|
# from main can have a well-defined shape.
|
|
@R.function
|
|
def main(x: R.Tensor, shape: R.Shape(["m", "n"])):
|
|
output = Module.subroutine(x, shape)
|
|
return output
|
|
|
|
assert Module["main"].ret_ty.shape is not None
|
|
_check(Module)
|
|
|
|
|
|
def test_empty_tuple():
|
|
@R.function
|
|
def foo(x: R.Tuple()):
|
|
y: R.Tuple() = R.tuple()
|
|
return y
|
|
|
|
x = relax.Var("x", relax.TupleType([]))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x,)):
|
|
y = bb.emit(relax.Tuple([]))
|
|
bb.emit_func_output(y)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_symbolic_vars_in_tensor_shape_with_usage_first():
|
|
"""First param may use symbolic variable defined in second param"""
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor(("m + 1",), "float32"), y: R.Tensor(("m", 1), "float32")):
|
|
z = R.add(x, y)
|
|
return z
|
|
|
|
m = tirx.Var("m", "int64")
|
|
x = relax.Var("x", relax.TensorType([m + 1], "float32"))
|
|
y = relax.Var("y", relax.TensorType([m, 1], "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x, y)):
|
|
z = bb.emit(relax.op.add(x, y))
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_symbolic_vars_in_tensor_shape_with_definition_first():
|
|
"""Second param may use symbolic variable defined in first param"""
|
|
|
|
@R.function
|
|
def bar(x: R.Tensor(("m",), "float32"), y: R.Tensor(("T.max(m, 20)",), "float32")) -> R.Tensor(
|
|
("T.max(m, 20) + 1",), "float32"
|
|
):
|
|
m = T.int64()
|
|
z = R.call_dps_packed("test_intrin", (x, y), R.Tensor((T.max(m, 20) + 1,), dtype="float32"))
|
|
return z
|
|
|
|
m = tirx.Var("m", "int64")
|
|
x = relax.Var("x", relax.TensorType([m], "float32"))
|
|
y = relax.Var("y", relax.TensorType([tirx.max(m, 20)], "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("bar", (x, y)):
|
|
z = bb.emit(
|
|
relax.call_dps_packed(
|
|
"test_intrin", (x, y), R.Tensor((tirx.max(m, 20) + 1,), dtype="float32")
|
|
)
|
|
)
|
|
bb.emit_func_output(z)
|
|
|
|
_check(bar, bb.get()["bar"])
|
|
|
|
|
|
def test_symbolic_vars_in_shape():
|
|
"""Symbolic variable may be defined in R.Shape"""
|
|
|
|
@R.function
|
|
def baz(x: R.Shape(("m",)), y: R.Tensor(("m * 2",), "float32")):
|
|
m = T.int64()
|
|
z = R.call_dps_packed("test_intrin", y, R.Tensor((m * 2,), dtype="float32"))
|
|
return z
|
|
|
|
m = tirx.Var("m", "int64")
|
|
x = relax.Var("x", relax.ShapeType([m]))
|
|
y = relax.Var("y", relax.TensorType([m * 2], "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("baz", (x, y)):
|
|
z = bb.emit(relax.call_dps_packed("test_intrin", (y), R.Tensor((m * 2,), dtype="float32")))
|
|
bb.emit_func_output(z)
|
|
|
|
_check(baz, bb.get()["baz"])
|
|
|
|
|
|
def test_undefined_symbolic_var_raises_error():
|
|
"""An undefined symbolic variable in an error
|
|
|
|
A symbolic variables is defined at the first site where it appears
|
|
as a shape parameter without any modification. TVMScript does not
|
|
support solving for a symbolic variable in terms of the argument
|
|
shape. That is, this test case raises an error, and will not
|
|
attempt to define `m` as either `x.shape[0]-1` or `x.shape[1]//2`.
|
|
"""
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function
|
|
def foo(x: R.Tensor(("m + 1", "m * 2"), "float32")): # name 'm' is not defined
|
|
z = R.add(x, x)
|
|
return z
|
|
|
|
|
|
def test_arith_operators():
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), "float32"), y: R.Tensor(("m", "n"), "float32")):
|
|
a0 = -x
|
|
a1 = x + y
|
|
a2 = x - y
|
|
a3 = x * y
|
|
a4 = x / y
|
|
a5 = x // y
|
|
a6 = x**y
|
|
|
|
c0 = x > y
|
|
c1 = x < y
|
|
c2 = x >= y
|
|
c3 = x <= y
|
|
|
|
tuple_expr = ((x, x), y)
|
|
t0 = tuple_expr[0]
|
|
t1 = tuple_expr[1]
|
|
t2 = tuple_expr[0][0] # <= Will normalize to two bindings
|
|
return (a0, a1, a2, a3, a4, a5, a6, c0, c1, c2, c3, t0, t1, t2)
|
|
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = relax.Var("x", relax.TensorType([m, n], "float32"))
|
|
y = relax.Var("y", relax.TensorType([m, n], "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x, y)):
|
|
a0 = bb.emit(relax.op.negative(x))
|
|
a1 = bb.emit(relax.op.add(x, y))
|
|
a2 = bb.emit(relax.op.subtract(x, y))
|
|
a3 = bb.emit(relax.op.multiply(x, y))
|
|
a4 = bb.emit(relax.op.divide(x, y))
|
|
a5 = bb.emit(relax.op.floor_divide(x, y))
|
|
a6 = bb.emit(relax.op.power(x, y))
|
|
|
|
c0 = bb.emit(relax.op.greater(x, y))
|
|
c1 = bb.emit(relax.op.less(x, y))
|
|
c2 = bb.emit(relax.op.greater_equal(x, y))
|
|
c3 = bb.emit(relax.op.less_equal(x, y))
|
|
|
|
tuple_expr = bb.emit(relax.Tuple((relax.Tuple((x, x)), y)))
|
|
t0 = bb.emit(relax.TupleGetItem(tuple_expr, 0))
|
|
t1 = bb.emit(relax.TupleGetItem(tuple_expr, 1))
|
|
tmp = bb.emit(relax.TupleGetItem(tuple_expr, 0))
|
|
t2 = bb.emit(relax.TupleGetItem(tmp, 0))
|
|
bb.emit_func_output(relax.Tuple((a0, a1, a2, a3, a4, a5, a6, c0, c1, c2, c3, t0, t1, t2)))
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_memory_ops():
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), dtype="float32")):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
storage = R.memory.alloc_storage(
|
|
R.shape([4 * m * n]), virtual_device_index=0, storage_scope="global", dtype="float32"
|
|
)
|
|
alloc = R.memory.alloc_tensor(storage, offset=0, shape=R.shape([m, n]), dtype="float32")
|
|
tensor = R.builtin.alloc_tensor(R.shape([m, n]), dtype="float32", runtime_device_index=0)
|
|
gv = tensor
|
|
return alloc, gv
|
|
|
|
_check(foo)
|
|
|
|
|
|
def test_vm_ops():
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor(("m", "n"), dtype="float32")):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
storage = R.vm.alloc_storage(R.shape([4 * m * n]), runtime_device_index=0, dtype="uint8")
|
|
alloc = R.vm.alloc_tensor(storage, offset=0, shape=R.shape([m, n]), dtype="float32")
|
|
tensor = R.builtin.alloc_tensor(R.shape([m, n]), dtype="float32", runtime_device_index=0)
|
|
tir_dym = R.vm.call_tir_dyn("te_func", (x, tensor, R.ShapeExpr((m, n))))
|
|
return alloc, tir_dym
|
|
|
|
_check(foo)
|
|
|
|
|
|
def test_builtin_ops():
|
|
@R.function
|
|
def foo(x: R.Tensor(("m", "n"), dtype="float32")):
|
|
tensor = R.builtin.stop_lift_params(x)
|
|
gv = tensor
|
|
return gv
|
|
|
|
_check(foo)
|
|
|
|
|
|
def test_prim_value():
|
|
@R.function(pure=False)
|
|
def foo():
|
|
gv = R.call_packed("test", 1, ty_args=R.Tensor((32, 32), "float32"))
|
|
return gv
|
|
|
|
_check(foo)
|
|
|
|
|
|
def test_string_imm():
|
|
@R.function(pure=False)
|
|
def foo():
|
|
gv = R.call_packed("test", "hello", ty_args=R.Tensor((32, 32), "float32"))
|
|
return gv
|
|
|
|
_check(foo)
|
|
|
|
|
|
def test_datatype_imm():
|
|
@R.function(pure=False)
|
|
def foo():
|
|
gv = R.call_packed("test", R.dtype("float32"), ty_args=R.Tensor((32, 32), "float32"))
|
|
return gv
|
|
|
|
_check(foo)
|
|
|
|
|
|
def test_function_void_return_type():
|
|
@tvm.script.ir_module
|
|
class Foo:
|
|
@R.function
|
|
def main(x: R.Tensor((3, 3), dtype="float32")):
|
|
res = Foo.mul(x)
|
|
return res
|
|
|
|
@R.function
|
|
def mul(x: R.Tensor((3, 3), dtype="float32")):
|
|
res = R.multiply(x, x)
|
|
return res
|
|
|
|
_check(Foo)
|
|
# Since the return type of function `mul` is not annotated,
|
|
# the function `main` regards it as a generic return type.
|
|
assert isinstance(Foo["main"].ret_ty, relax.AnyType)
|
|
assert isinstance(Foo["mul"].ret_ty, relax.TensorType)
|
|
|
|
@tvm.script.ir_module
|
|
class Bar:
|
|
@R.function
|
|
def main(x1: R.Tensor((3, 3), dtype="float32")):
|
|
res1 = Bar.mul(x1)
|
|
return res1
|
|
|
|
@R.function
|
|
def mul(x: R.Tensor((3, 3), dtype="float32")) -> None:
|
|
res = R.multiply(x, x)
|
|
return res
|
|
|
|
# Since the return type of function `mul` is not annotated,
|
|
# the function `main` regards it as a generic return type.
|
|
_check(Bar)
|
|
tvm.ir.assert_structural_equal(Bar["main"].ret_ty, relax.TupleType([]))
|
|
tvm.ir.assert_structural_equal(Bar["mul"].ret_ty, relax.TupleType([]))
|
|
|
|
|
|
def test_class_normalize():
|
|
@tvm.script.ir_module
|
|
class InputModule:
|
|
@R.function
|
|
def mul_add(x: R.Tensor) -> R.Tensor:
|
|
return R.multiply(R.add(x, x), R.add(x, x))
|
|
|
|
# The parser automatically normalizes the input AST to the following ANF form
|
|
@tvm.script.ir_module
|
|
class OutputModule:
|
|
@R.function
|
|
def mul_add(x: R.Tensor) -> R.Tensor:
|
|
gv = R.add(x, x)
|
|
gv1 = R.add(x, x)
|
|
return R.multiply(gv, gv1)
|
|
|
|
_check(InputModule, OutputModule)
|
|
|
|
|
|
def test_context_aware_parsing(monkeypatch):
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def add(
|
|
X: T.Buffer([T.int64(2), T.int64(4)], "float32"),
|
|
Y: T.Buffer((), "float32"),
|
|
Z: T.Buffer([T.int64(2), T.int64(4)], "float32"),
|
|
):
|
|
T.evaluate(0)
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((2, 4), dtype="float32")) -> R.Tensor((10,), dtype="float32"):
|
|
R.func_attr({"relax.force_pure": True})
|
|
cls = Module
|
|
alloc = R.builtin.alloc_tensor(R.shape([2, 4]), dtype="float32", runtime_device_index=0)
|
|
_: R.Tuple() = cls.add(x, R.const(1, "float32"), alloc)
|
|
return alloc
|
|
|
|
_check(Module)
|
|
|
|
# Break the env settings, but context-aware parsing can still handle it
|
|
def _break_env(self, *args):
|
|
raise RuntimeError("Fail to pass context-aware parsing")
|
|
|
|
monkeypatch.setattr(tvm.ir.GlobalVar, "__call__", _break_env)
|
|
|
|
_check(Module)
|
|
|
|
|
|
def test_unit_tuple_on_rhs_of_assign():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(input: R.Tensor((5, 5))) -> R.Tuple(R.Tensor((5, 5))):
|
|
gv = (input,)
|
|
return gv
|
|
|
|
_check(Module)
|
|
|
|
|
|
def test_empty_tuple_on_rhs_of_assign():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(input: R.Tensor((5, 5))) -> R.Tuple():
|
|
gv = ()
|
|
return gv
|
|
|
|
_check(Module)
|
|
|
|
|
|
def test_global_var_ty():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def foo(x: R.Tensor((128, 128), "float32")):
|
|
gv0 = R.emit_te(topi.add, x, x)
|
|
return gv0
|
|
|
|
target_ty = R.Callable(
|
|
(R.Tensor((128, 128), dtype="float32"),), R.Tensor((128, 128), dtype="float32")
|
|
)
|
|
gv = Module.get_global_var("foo")
|
|
tvm.ir.assert_structural_equal(gv.ty, target_ty)
|
|
tvm.ir.assert_structural_equal(Module["foo"].ty, target_ty)
|
|
_check(Module)
|
|
|
|
|
|
def test_assert_op():
|
|
@I.ir_module(s_tir=True)
|
|
class AssertOp:
|
|
@R.function(pure=False)
|
|
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
|
|
return x
|
|
|
|
_check(AssertOp)
|
|
|
|
|
|
def test_assert_outside_of_class():
|
|
@R.function(pure=False)
|
|
def func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
|
|
return x
|
|
|
|
# this just makes sure that the machinery regarding the pure attribute parses
|
|
# in the case where the function is outside of a class too
|
|
_check(func)
|
|
|
|
|
|
def test_impure_inner_function():
|
|
@R.function
|
|
def f(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
# we will not actually call it
|
|
@R.function(pure=False)
|
|
def g(y: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
z = R.assert_op(R.const(False, dtype="bool"), y, format="y: {}")
|
|
return y
|
|
|
|
return x
|
|
|
|
assert f.is_pure
|
|
# definition of g
|
|
assert not f.body.blocks[0].bindings[0].value.is_pure
|
|
|
|
# make sure we are not incorrectly passing state for inner functions
|
|
_check(f)
|
|
|
|
|
|
def test_impure_inner_function_in_class():
|
|
@I.ir_module(s_tir=True)
|
|
class ImpureInner:
|
|
@R.function
|
|
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
# we will not actually call it
|
|
@R.function(pure=False)
|
|
def g(y: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
z = R.assert_op(R.const(False, dtype="bool"), y, format="y: {}")
|
|
return y
|
|
|
|
return x
|
|
|
|
assert ImpureInner["main"].is_pure
|
|
# definition of g
|
|
assert not ImpureInner["main"].body.blocks[0].bindings[0].value.is_pure
|
|
|
|
# make sure we are not incorrectly passing state for inner functions
|
|
_check(ImpureInner)
|
|
|
|
|
|
def test_print():
|
|
@I.ir_module(s_tir=True)
|
|
class Print:
|
|
@R.function(pure=False)
|
|
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.print(x, format="x: {}")
|
|
return x
|
|
|
|
_check(Print)
|
|
|
|
|
|
def test_parse_multiple_pure_and_impure_funcs():
|
|
@I.ir_module(s_tir=True)
|
|
class Mixture:
|
|
@R.function(pure=False)
|
|
def print(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.print(x, format="x: {}")
|
|
return x
|
|
|
|
@R.function(pure=False)
|
|
def assert_func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
|
|
return x
|
|
|
|
@R.function
|
|
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
return x
|
|
|
|
assert not Mixture["print"].is_pure
|
|
assert not Mixture["assert_func"].is_pure
|
|
assert Mixture["main"].is_pure
|
|
_check(Mixture)
|
|
|
|
|
|
def test_function_with_void_return_type_may_be_used_as_statements():
|
|
"""Void return of calls do not need to be assigned"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Unsugared:
|
|
@R.function(pure=False)
|
|
def print(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.print(x, format="x: {}")
|
|
return x
|
|
|
|
@R.function(pure=False)
|
|
def assert_func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
|
|
return x
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Sugared:
|
|
@R.function(pure=False)
|
|
def print(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.print(x, format="x: {}")
|
|
return x
|
|
|
|
@R.function(pure=False)
|
|
def assert_func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.assert_op(R.const(False, dtype="bool"), x, format="x: {}")
|
|
return x
|
|
|
|
tvm.ir.assert_structural_equal(Unsugared, Sugared)
|
|
|
|
|
|
def test_function_with_non_void_return_type_must_be_assigned():
|
|
"""Non-void results must be assigned to a variable"""
|
|
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function(pure=False)
|
|
def func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.add(x, x)
|
|
return x
|
|
|
|
|
|
def test_function_with_void_return_type_in_if_else():
|
|
"""Last statement in if/else may be a void return"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Unsugared:
|
|
@R.function(pure=False)
|
|
def conditional(x: R.Tensor((), "int32"), condition: R.Tensor((), "bool")) -> R.Tensor(
|
|
(), "int32"
|
|
):
|
|
if condition:
|
|
y = R.print(x, format="True condition: {}")
|
|
else:
|
|
y = R.print(x, format="False condition: {}")
|
|
return x
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Sugared:
|
|
@R.function(pure=False)
|
|
def conditional(x: R.Tensor((), "int32"), condition: R.Tensor((), "bool")) -> R.Tensor(
|
|
(), "int32"
|
|
):
|
|
if condition:
|
|
R.print(x, format="True condition: {}")
|
|
else:
|
|
R.print(x, format="False condition: {}")
|
|
return x
|
|
|
|
_check(Sugared, Unsugared)
|
|
|
|
|
|
def test_call_pure_packed():
|
|
@R.function
|
|
def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
|
z = R.call_pure_packed("vm.builtin.copy", x, ty_args=R.Tensor((32, 32), "float32"))
|
|
return z
|
|
|
|
x = relax.Var("x", R.Tensor((32, 32), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", (x)):
|
|
z = bb.emit(
|
|
R.call_pure_packed("vm.builtin.copy", x, ty_args=[R.Tensor((32, 32), "float32")])
|
|
)
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_call_pure_packed_returning_object():
|
|
@R.function
|
|
def foo() -> R.Any:
|
|
z = R.call_pure_packed("dummy_func", ty_args=R.Any)
|
|
return z
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("foo", params=[]):
|
|
z = bb.emit(R.call_pure_packed("dummy_func", ty_args=[relax.AnyType()]))
|
|
bb.emit_func_output(z)
|
|
|
|
_check(foo, bb.get()["foo"])
|
|
|
|
|
|
def test_private_function():
|
|
@I.ir_module(s_tir=True)
|
|
class Addition:
|
|
@R.function(private=True)
|
|
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
y = R.add(x, x)
|
|
return y
|
|
|
|
x = relax.Var("x", R.Tensor((), "int32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("main", (x), private=True):
|
|
y = bb.emit(R.add(x, x))
|
|
bb.emit_func_output(y)
|
|
|
|
_check(Addition, bb.get())
|
|
|
|
|
|
def test_private_function_with_global_symbol_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Addition:
|
|
@R.function(private=True)
|
|
def main(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
# it is an error to simultaneously mark a function private
|
|
# and give it a global symbol manually
|
|
R.func_attr({"global_symbol": "main"})
|
|
y = R.add(x, x)
|
|
return y
|
|
|
|
# should not execute
|
|
_check(Addition)
|
|
|
|
|
|
def test_private_function_with_global_symbol_no_module_fail():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.function(private=True)
|
|
def func(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"):
|
|
R.func_attr({"global_symbol": "main"})
|
|
y = R.add(x, x)
|
|
return y
|
|
|
|
# should not execute
|
|
_check(func)
|
|
|
|
|
|
def test_macro_hygienic():
|
|
x = R.prim_value(2)
|
|
|
|
@R.macro(hygienic=True)
|
|
def alloc_and_shape(dtype: str):
|
|
alloc = R.builtin.alloc_tensor(R.shape([4, 4]), runtime_device_index=x, dtype=dtype)
|
|
shape = R.shape_of(alloc)
|
|
return shape
|
|
|
|
x = R.prim_value(1)
|
|
|
|
@R.function(private=True)
|
|
def func(z: R.Tensor((4, 4), "float32")):
|
|
shape = alloc_and_shape(dtype="float32")
|
|
return shape
|
|
|
|
@R.function(private=True)
|
|
def expect(z: R.Tensor((4, 4), dtype="float32")) -> R.Shape([4, 4]):
|
|
alloc: R.Tensor((4, 4), dtype="float32") = R.builtin.alloc_tensor(
|
|
R.shape([4, 4]),
|
|
R.dtype("float32"),
|
|
R.prim_value(2), # Make sure prim_value is 2
|
|
)
|
|
shape: R.Shape([4, 4]) = R.shape_of(alloc)
|
|
shape_1: R.Shape([4, 4]) = shape
|
|
return shape_1
|
|
|
|
_check(func, expect)
|
|
|
|
|
|
def test_macro_non_hygienic():
|
|
global global_x_var # Lookup doesn't find this variable if it's not global
|
|
|
|
global_x_var = R.prim_value(2)
|
|
|
|
@R.macro(hygienic=False)
|
|
def alloc_and_shape(dtype: str):
|
|
alloc = R.builtin.alloc_tensor(
|
|
R.shape([4, 4]), runtime_device_index=global_x_var, dtype=dtype
|
|
)
|
|
shape = R.shape_of(alloc)
|
|
return shape
|
|
|
|
global_x_var = R.prim_value(1)
|
|
|
|
@R.function(private=True)
|
|
def func(z: R.Tensor((4, 4), "float32")):
|
|
shape = alloc_and_shape(dtype="float32")
|
|
return shape
|
|
|
|
@R.function(private=True)
|
|
def expect(z: R.Tensor((4, 4), dtype="float32")) -> R.Shape([4, 4]):
|
|
alloc: R.Tensor((4, 4), dtype="float32") = R.builtin.alloc_tensor(
|
|
R.shape([4, 4]),
|
|
R.dtype("float32"),
|
|
R.prim_value(1), # Make sure prim_value is 1
|
|
)
|
|
shape: R.Shape([4, 4]) = R.shape_of(alloc)
|
|
shape_1: R.Shape([4, 4]) = shape
|
|
return shape_1
|
|
|
|
_check(func, expect)
|
|
|
|
|
|
def test_macro_no_variable_leak():
|
|
with pytest.raises(tvm.error.DiagnosticError):
|
|
|
|
@R.macro(hygienic=True)
|
|
def add_two(value):
|
|
x = value + R.const(1) # `x` defined in macro
|
|
y = x + R.const(1)
|
|
return y
|
|
|
|
@R.function(private=True)
|
|
def func(t: R.Tensor((), "int32")):
|
|
u = add_two(t)
|
|
return x # Should be undefined here
|
|
|
|
|
|
def test_reused_extern_func():
|
|
"""ExternFunc lookups can become bindings in EliminateCommonSubexpr"""
|
|
|
|
@R.function(private=True)
|
|
def parsed(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"):
|
|
func = R.ExternFunc("extern_func")
|
|
gv0 = R.call_dps_packed(func, x, R.Tensor((128, 128), dtype="float32"))
|
|
gv1 = R.call_dps_packed(func, gv0, R.Tensor((128, 128), dtype="float32"))
|
|
return gv1
|
|
|
|
x = relax.Var("x", R.Tensor((128, 128), "float32"))
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("main", [x], private=True):
|
|
func = bb.emit(relax.ExternFunc("extern_func"))
|
|
y = bb.emit(relax.call_dps_packed(func, x, out_ty=R.Tensor((128, 128), "float32")))
|
|
z = bb.emit(relax.call_dps_packed(func, y, out_ty=R.Tensor((128, 128), "float32")))
|
|
bb.emit_func_output(z)
|
|
|
|
expected = bb.get()["main"]
|
|
|
|
_check(parsed, expected)
|
|
|
|
|
|
def test_extern_func_in_module():
|
|
"""Module-level parsing may produce function bindings"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class parsed_module:
|
|
my_ext = R.ExternFunc("my_ext")
|
|
|
|
@R.function
|
|
def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
|
|
return a
|
|
|
|
@R.function
|
|
def func(a: R.Tensor((10, 10))) -> R.Tensor((10, 10)):
|
|
return a
|
|
|
|
expected = tvm.IRModule({"my_ext": relax.ExternFunc("my_ext"), "func": func})
|
|
|
|
_check(parsed_module, expected)
|
|
|
|
|
|
def test_define_relax_function_using_global_var():
|
|
"""A @R.function may call a GlobalVar
|
|
|
|
When parsing a @R.function, the function's body may reference
|
|
GlobalVar instances available in the calling python scope. The
|
|
resulting function should pass TVMScript's well-formed check, as
|
|
the GlobalVar may be available in the IRModule for which the
|
|
function is being defined.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class DefinedAllAtOnce:
|
|
@R.function
|
|
def main(A: R.Tensor, B: R.Tensor):
|
|
return DefinedAllAtOnce.subroutine(A, B)
|
|
|
|
@R.function(private=True)
|
|
def subroutine(A: R.Tensor, B: R.Tensor) -> R.Tensor:
|
|
return R.matmul(A, B)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class MainDefinedLater:
|
|
@R.function(private=True)
|
|
def subroutine(A: R.Tensor, B: R.Tensor) -> R.Tensor:
|
|
return R.matmul(A, B)
|
|
|
|
subroutine_gvar = MainDefinedLater.get_global_var("subroutine")
|
|
|
|
@R.function
|
|
def main(A: R.Tensor, B: R.Tensor):
|
|
return subroutine_gvar(A, B)
|
|
|
|
MainDefinedLater["main"] = main
|
|
|
|
tvm.ir.assert_structural_equal(DefinedAllAtOnce, MainDefinedLater)
|
|
|
|
|
|
def test_function_attributes_are_defined():
|
|
"""func.attrs defaults to an empty DictAttrs"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def main(x: R.Tensor, shape: R.Shape(["m", "n"])):
|
|
output = Module.subroutine(x, shape)
|
|
return output
|
|
|
|
@R.function
|
|
def subroutine(x: R.Tensor, _: R.Shape(["m", "n"])) -> R.Tensor(["m", "n"]):
|
|
q = x
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(q, R.Tensor((m, n)))
|
|
w = z
|
|
return w
|
|
|
|
for gvar, func in Module.functions.items():
|
|
assert func.attrs is not None
|
|
|
|
|
|
@pytest.mark.xfail(reason="Bug: Implicit bounds not provided when parsing")
|
|
def test_function_symbolic_variables_are_annotated():
|
|
"""Symbolic variables must be exposed for struct inference
|
|
|
|
Because Relax struct inference is performed while the function is
|
|
being built, all constraints on symbolic variables that are used
|
|
for simplifications must be provided to the analyzer.
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def inferred_ty(A: R.Tensor(["extent"])):
|
|
extent = T.int64()
|
|
output = R.strided_slice(A, [0], [0], [extent - 1])
|
|
return output
|
|
|
|
@R.function(private=True)
|
|
def expected(A: R.Tensor(["extent"])) -> R.Tensor(["extent-1"]):
|
|
extent = T.int64()
|
|
output: R.Tensor([extent - 1]) = R.strided_slice(A, [0], [0], [extent - 1])
|
|
return output
|
|
|
|
tvm.ir.assert_structural_equal(inferred_ty, expected)
|
|
|
|
|
|
def test_constant_prim_expr_alias_is_not_symbolic_declaration():
|
|
"""Constant Expr locals are constants, not declarations."""
|
|
|
|
@R.function(private=True)
|
|
def func(A: R.Tensor([4], "float32")):
|
|
extent = T.int64(4)
|
|
output: R.Tensor([extent], "float32") = A
|
|
return output
|
|
|
|
tvm.ir.assert_structural_equal(func.ret_ty.shape.values[0], T.int64(4))
|
|
|
|
|
|
def test_conditional_may_use_symbolic_variables_from_function_scope():
|
|
"""Symbolic variables from function scope may be used in branch
|
|
|
|
This is a regression test. In earlier implementations, the
|
|
branches of `relax::If` were normalized with
|
|
`EraseToWellDefinedInScope`, using a fresh variable scope. While
|
|
this had the intended behavior of preventing variables defined in
|
|
a single branch from being usable outside of the conditional, it
|
|
also caused the conditional's branches to treat function-scope
|
|
symbolic variables as if they were undefined.
|
|
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def explicit_ty(
|
|
A: R.Tensor(["N"], "float32"),
|
|
B: R.Tensor(["N"], "float32"),
|
|
cond: R.Prim("bool"),
|
|
) -> R.Tensor(["N"], "float32"):
|
|
N = T.int64()
|
|
|
|
if cond:
|
|
out: R.Tensor([N], "float32") = A + B
|
|
else:
|
|
out: R.Tensor([N], "float32") = A * B
|
|
|
|
return out
|
|
|
|
@R.function(private=True)
|
|
def inferred_ty(
|
|
A: R.Tensor(["N"], "float32"),
|
|
B: R.Tensor(["N"], "float32"),
|
|
cond: R.Prim("bool"),
|
|
):
|
|
N = T.int64()
|
|
if cond:
|
|
out = A + B
|
|
else:
|
|
out = A * B
|
|
|
|
return out
|
|
|
|
tvm.ir.assert_structural_equal(explicit_ty, inferred_ty)
|
|
|
|
|
|
def test_return_from_dataflow_block():
|
|
"""Return statements imply
|
|
|
|
The `R.output` statement in a `R.dataflow()` block marks a
|
|
variable that should be a `relax.Var` instead of a
|
|
`relax.DataflowVar`, allowing it to be used outside of the
|
|
`DataflowBlock` that defined it. A relax function's output is not
|
|
part of any binding, and must not contain any `DataflowVar`, so
|
|
these are exposed implicitly.
|
|
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def output_then_return(A: R.Tensor([16], "float16")):
|
|
with R.dataflow():
|
|
B = R.add(A, A)
|
|
C = R.multiply(B, B)
|
|
R.output(C)
|
|
|
|
return C
|
|
|
|
@R.function(private=True)
|
|
def return_inside_dataflow(A: R.Tensor([16], "float16")):
|
|
with R.dataflow():
|
|
B = R.add(A, A)
|
|
C = R.multiply(B, B)
|
|
return C
|
|
|
|
tvm.ir.assert_structural_equal(output_then_return, return_inside_dataflow)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|