wehub-resource-sync
992 lines
32 KiB
Python
992 lines
32 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, F841
|
|
"""Block builder unit test"""
|
|
|
|
# The test here do not depend on tvmscript to cover most basic features
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.contrib.cblas
|
|
import tvm.testing
|
|
from tvm import relax as rx
|
|
from tvm import te, tirx, topi
|
|
from tvm.ir.base import assert_structural_equal
|
|
from tvm.relax import ExternFunc
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
from tvm.tirx.function import PrimFunc
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def register_nop():
|
|
@tvm.register_global_func("test.blockbuilder.nop")
|
|
def nop():
|
|
pass
|
|
|
|
|
|
def test_block_builder():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
bb._begin_binding_block()
|
|
gv0 = bb.emit(rx.op.add(x, y))
|
|
bb._begin_dataflow_block()
|
|
lv0 = bb.emit(rx.op.multiply(gv0, y))
|
|
gv1 = bb.emit_output(rx.op.multiply(lv0, lv0))
|
|
b0 = bb._end_block()
|
|
bb._begin_dataflow_block()
|
|
lv1 = bb.emit(rx.op.multiply(gv0, y))
|
|
gv2 = bb.emit_output(rx.op.multiply(lv1, lv1))
|
|
b1 = bb._end_block()
|
|
gv3 = bb.emit(rx.op.add(x, y))
|
|
b2 = bb._end_block()
|
|
|
|
assert isinstance(b0, rx.DataflowBlock)
|
|
assert isinstance(b1, rx.DataflowBlock)
|
|
assert not isinstance(b2, rx.DataflowBlock)
|
|
|
|
|
|
def test_emit_with_name():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
bb._begin_dataflow_block()
|
|
lv0 = bb.emit(rx.op.add(x, y), "add")
|
|
gv0 = bb.emit_output(rx.op.multiply(lv0, y), "multi")
|
|
b0 = bb._end_block()
|
|
|
|
assert b0.bindings[0].var.name_hint == "add"
|
|
assert b0.bindings[1].var.name_hint == "multi"
|
|
|
|
|
|
def test_function_single_block():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", [x, y]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(rx.op.add(x, y))
|
|
assert lv0.name_hint == "lv"
|
|
lv1 = bb.emit(rx.op.multiply(lv0, y))
|
|
assert lv1.name_hint == "lv1"
|
|
gv0 = bb.emit_output(lv1)
|
|
assert gv0.name_hint == "gv"
|
|
bb.emit_func_output(gv0)
|
|
|
|
func = bb.finalize()["func"]
|
|
assert func.params[0] == x
|
|
assert func.params[1] == y
|
|
assert func.body.body == gv0
|
|
assert_structural_equal(gv0.ty, rx.TensorType([m, n], "float16"))
|
|
assert len(func.body.blocks) == 1
|
|
assert len(func.body.blocks[0].bindings) == 3
|
|
|
|
|
|
def test_function_multi_blocks():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", [x, y]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(rx.op.add(x, y))
|
|
assert lv0.name_hint == "lv"
|
|
gv0 = bb.emit_output(lv0)
|
|
assert gv0.name_hint == "gv"
|
|
gv1 = bb.emit(rx.op.add(gv0, gv0))
|
|
assert gv1.name_hint == "gv1"
|
|
with bb.dataflow():
|
|
lv1 = bb.emit(rx.op.add(gv1, gv1))
|
|
assert lv1.name_hint == "lv1"
|
|
gv2 = bb.emit_output(gv1)
|
|
bb.emit_func_output(gv2)
|
|
|
|
func = bb.finalize()["func"]
|
|
|
|
assert_structural_equal(gv2.ty, rx.TensorType([m, n], "float16"))
|
|
assert func.params[0] == x
|
|
assert func.params[1] == y
|
|
assert func.body.body == gv2
|
|
assert len(func.body.blocks) == 3
|
|
assert len(func.body.blocks[0].bindings) == 2
|
|
assert len(func.body.blocks[1].bindings) == 1
|
|
assert len(func.body.blocks[2].bindings) == 2
|
|
|
|
|
|
def test_multi_functions():
|
|
bb = rx.BlockBuilder()
|
|
|
|
m_1 = tirx.Var("m", "int64")
|
|
n_1 = tirx.Var("n", "int64")
|
|
x_1 = rx.Var("x", rx.TensorType([m_1, n_1], "float16"))
|
|
y_1 = rx.Var("y", rx.TensorType([n_1], "float16"))
|
|
|
|
with bb.function("func1", [x_1, y_1]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(rx.op.add(x_1, y_1))
|
|
assert lv0.name_hint == "lv"
|
|
gv0 = bb.emit_output(lv0)
|
|
bb.emit_func_output(gv0)
|
|
|
|
m_2 = tirx.Var("m", "int64")
|
|
n_2 = tirx.Var("n", "int64")
|
|
x_2 = rx.Var("x", rx.TensorType([m_2, n_2], "float16"))
|
|
y_2 = rx.Var("y", rx.TensorType([n_2], "float16"))
|
|
|
|
with bb.function("func2", [x_2, y_2]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(rx.op.add(y_2, x_2))
|
|
# TODO(@yuchen): enable block builder to reset local var unique name map
|
|
assert lv0.name_hint == "lv1"
|
|
gv0 = bb.emit_output(lv0)
|
|
bb.emit_func_output(gv0)
|
|
|
|
mod = bb.finalize()
|
|
func1 = mod["func1"]
|
|
assert func1.params[0] == x_1
|
|
assert func1.params[1] == y_1
|
|
assert len(func1.body.blocks) == 1
|
|
func2 = mod["func2"]
|
|
assert func2.params[0] == x_2
|
|
assert func2.params[1] == y_2
|
|
assert len(func2.body.blocks) == 1
|
|
|
|
|
|
def test_binary_shape_type_deduction():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
k = tirx.Var("k", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, 1], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
z = rx.Var("z", rx.TensorType([5], "float16"))
|
|
w = rx.Var("w", rx.TensorType([k], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", [x, y, z, w]):
|
|
with bb.dataflow():
|
|
lv0 = bb.emit(rx.op.add(x, y))
|
|
assert_structural_equal(lv0.ty, rx.TensorType([m, n], "float16"))
|
|
|
|
lv1 = bb.emit(rx.op.multiply(x, z))
|
|
assert_structural_equal(lv1.ty, rx.TensorType([m, 5], "float16"))
|
|
|
|
lv2 = bb.emit(rx.op.multiply(z, w))
|
|
assert isinstance(lv2.ty, rx.TensorType)
|
|
assert lv2.ty.ndim == 1
|
|
assert lv2.ty.dtype == "float16"
|
|
|
|
lv3 = bb.emit(rx.op.multiply(y, w))
|
|
assert isinstance(lv3.ty, rx.TensorType)
|
|
assert lv3.ty.ndim == 1
|
|
assert lv3.ty.dtype == "float16"
|
|
|
|
gv0 = bb.emit_output(lv3)
|
|
bb.emit_func_output(gv0)
|
|
|
|
assert isinstance(gv0.ty, rx.TensorType)
|
|
assert gv0.ty.ndim == 1
|
|
assert gv0.ty.dtype == "float16"
|
|
|
|
|
|
def test_emit_match_cast():
|
|
m = tirx.Var("m", dtype="int64")
|
|
n = tirx.Var("n", dtype="int64")
|
|
x = rx.Var("tensor_value", rx.TensorType(dtype="float32", ndim=-1))
|
|
y = rx.Var("shape_value", rx.ShapeType([16, 8]))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", [x, y]):
|
|
with bb.dataflow():
|
|
# lv0: Tensor((m, n), "float32") =
|
|
# match_cast(x: Tensor(_, "float32"], [m, n))
|
|
lv0 = bb.match_cast(x, rx.TensorType([m, n], "float32"))
|
|
assert isinstance(lv0, rx.DataflowVar)
|
|
assert_structural_equal(lv0.ty, rx.TensorType([m, n], "float32"))
|
|
|
|
# lv1: Shape = match_cast(shape, rx.ShapeType([m, n]))
|
|
lv1 = bb.match_cast(y, rx.ShapeType([m, n]), "var_name")
|
|
assert lv1.ty == rx.ShapeType([m, n])
|
|
gv0 = bb.emit_output(lv1)
|
|
|
|
bb.emit_func_output(gv0)
|
|
func = bb.finalize()["func"]
|
|
block = func.body.blocks[0]
|
|
b0, b1 = block.bindings[:2]
|
|
assert isinstance(b0, rx.MatchCast)
|
|
assert isinstance(b1, rx.MatchCast)
|
|
|
|
assert b0.value == x
|
|
assert b0.ty == rx.TensorType([m, n], "float32")
|
|
assert b0.var == lv0
|
|
|
|
assert b1.value == y
|
|
assert b1.ty == rx.ShapeType([m, n])
|
|
assert b1.var == lv1
|
|
assert b1.var.name_hint == "var_name"
|
|
|
|
|
|
def test_emit_match_cast_binding_in_dataflow_block():
|
|
bb = rx.BlockBuilder()
|
|
|
|
x = rx.Var("x", rx.TensorType(dtype="float32", ndim=-1))
|
|
m = tirx.Var("m", dtype="int64")
|
|
gv = rx.Var("gv", rx.TensorType(dtype="float32", ndim=-1))
|
|
match_cast = rx.MatchCast(gv, x, rx.TensorType((m,), "float32"))
|
|
|
|
with bb.function("main", [x]):
|
|
with bb.dataflow():
|
|
bb.emit_normalized(match_cast)
|
|
bb.emit_output(gv)
|
|
bb.emit_func_output(x)
|
|
|
|
func = bb.finalize()["main"]
|
|
block = func.body.blocks[0]
|
|
b0 = block.bindings[0]
|
|
assert isinstance(b0, rx.MatchCast)
|
|
|
|
assert b0.value == x
|
|
assert isinstance(b0.ty, rx.TensorType)
|
|
assert b0.ty.shape[0] == m
|
|
assert b0.var == gv
|
|
|
|
|
|
def test_normalize():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
# Call node
|
|
add_call = rx.op.multiply(x, y)
|
|
|
|
bb.normalize(add_call)
|
|
shape = rx.get_shape_of(add_call)
|
|
|
|
assert isinstance(shape, rx.ShapeExpr)
|
|
assert shape[0] == m
|
|
assert shape[1] == n
|
|
|
|
# Tuple node
|
|
tuple_1 = rx.Tuple([x, y])
|
|
bb.normalize(tuple_1)
|
|
assert isinstance(tuple_1.ty, rx.TupleType)
|
|
assert isinstance(tuple_1.ty.fields[0], rx.TensorType)
|
|
assert isinstance(tuple_1.ty.fields[1], rx.TensorType)
|
|
|
|
# Nested Tuple
|
|
tuple_2 = rx.Tuple([x, rx.Tuple([x, y])])
|
|
bb.normalize(tuple_2)
|
|
|
|
assert isinstance(tuple_2.ty, rx.TupleType)
|
|
assert isinstance(tuple_2.ty.fields[0], rx.TensorType)
|
|
assert isinstance(tuple_2.ty.fields[1], rx.TupleType)
|
|
assert isinstance(tuple_2.ty.fields[1].fields[0], rx.TensorType)
|
|
assert isinstance(tuple_2.ty.fields[1].fields[1], rx.TensorType)
|
|
|
|
|
|
def test_tuple_indexing():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
|
|
shape_x = rx.TensorType([m, n], "float16")
|
|
shape_y = rx.TensorType([n], "float16")
|
|
relax_tuple = rx.Var("relax_tuple", rx.TupleType([shape_x, shape_y]))
|
|
|
|
assert isinstance(relax_tuple.ty, rx.TupleType)
|
|
assert isinstance(relax_tuple.ty.fields[0], rx.TensorType)
|
|
assert isinstance(relax_tuple.ty.fields[1], rx.TensorType)
|
|
|
|
# TupleGetItem will initialize type from the
|
|
# TupleType, if present.
|
|
x = relax_tuple[0]
|
|
tvm.ir.assert_structural_equal(x.ty, shape_x)
|
|
|
|
y = relax_tuple[1]
|
|
tvm.ir.assert_structural_equal(y.ty, shape_y)
|
|
|
|
# Tuple unpacking produces TupleGetItem structs
|
|
x_unpack, y_unpack = relax_tuple
|
|
tvm.ir.assert_structural_equal(x, x_unpack)
|
|
tvm.ir.assert_structural_equal(y, y_unpack)
|
|
|
|
# When TupleType is available, tuple unpacking fails immediately
|
|
# for incorrect number of arguments.
|
|
with pytest.raises(ValueError):
|
|
x_unpack, y_unpack, z_unpack = relax_tuple
|
|
|
|
|
|
def test_call_te():
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
y = rx.Var("y", rx.TensorType([n, m], "float32"))
|
|
z = rx.Var("z", rx.TensorType([n, m], "float32"))
|
|
|
|
def te_func(args, args_dict, msg):
|
|
A, B = args
|
|
C = args_dict["C"]
|
|
D = te.compute((128, 128), lambda i, j: A[i, j] + B[i, j])
|
|
E = te.compute((128, 128), lambda i, j: D[i, j] - C[i, j])
|
|
return E
|
|
|
|
with bb.function("rx_func", [x, y, z]):
|
|
with bb.dataflow():
|
|
out = bb.emit_output(bb.call_te(te_func, [x, y], {"C": z}, msg="hello"))
|
|
bb.emit_func_output(out)
|
|
|
|
mod = bb.finalize()
|
|
rx_func = mod["rx_func"]
|
|
|
|
assert rx_func.params[0] == x
|
|
assert rx_func.params[1] == y
|
|
assert rx_func.params[2] == z
|
|
assert rx_func.body.body == out
|
|
assert len(rx_func.body.blocks) == 1
|
|
assert len(rx_func.body.blocks[0].bindings) == 1
|
|
|
|
|
|
def test_call_te_unique_tensor_name():
|
|
bb = rx.BlockBuilder()
|
|
x = rx.Var("x", R.Tensor((2, 3), "float32"))
|
|
y = rx.Var("y", R.Tensor((3, 4), "float32"))
|
|
with bb.function("main", [x, y]):
|
|
gv = bb.emit_te(topi.nn.matmul, x, y)
|
|
bb.emit_func_output(gv)
|
|
|
|
f_matmul = bb.finalize()["matmul"]
|
|
param_A = f_matmul.params[0]
|
|
param_B = f_matmul.params[1]
|
|
buffer_A = f_matmul.buffer_map[param_A]
|
|
buffer_B = f_matmul.buffer_map[param_B]
|
|
assert param_A.name != param_B.name
|
|
assert buffer_A.name != buffer_B.name
|
|
assert buffer_A.data.name != buffer_B.data.name
|
|
|
|
|
|
def test_call_te_with_unsupported_shape_arg():
|
|
bb = rx.BlockBuilder()
|
|
x = rx.Var("x", rx.TensorType((200,), "float32"))
|
|
s = rx.Var("s", rx.ShapeType((200,)))
|
|
|
|
with pytest.raises(AssertionError):
|
|
with bb.function("rx_func", [x]):
|
|
out = bb.emit(bb.call_te(topi.reshape, x, s))
|
|
bb.emit_func_output(out)
|
|
|
|
|
|
def test_emit_te():
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
y = rx.Var("y", rx.TensorType([n, m], "float32"))
|
|
z = rx.Var("z", rx.TensorType([n, m], "float32"))
|
|
|
|
def te_func(args, args_dict, msg):
|
|
A, B = args
|
|
C = args_dict["C"]
|
|
D = te.compute((128, 128), lambda i, j: A[i, j] + B[i, j])
|
|
E = te.compute((128, 128), lambda i, j: D[i, j] - C[i, j])
|
|
return E
|
|
|
|
with bb.function("rx_func", [x, y, z]):
|
|
out = bb.emit_te(te_func, [x, y], {"C": z}, msg="hello")
|
|
bb.emit_func_output(out)
|
|
|
|
mod = bb.finalize()
|
|
rx_func = mod["rx_func"]
|
|
|
|
def get_tir_func():
|
|
A = te.placeholder((n, m), dtype="float32", name="A")
|
|
B = te.placeholder((n, m), dtype="float32", name="B")
|
|
C = te.placeholder((n, m), dtype="float32", name="C")
|
|
out = te_func((A, B), {"C": C}, "")
|
|
return tvm.te.create_prim_func([A, B, C, out], index_dtype_override="int64")
|
|
|
|
# check TIR structure matches expected
|
|
assert_structural_equal(mod["te_func"].body, get_tir_func().body)
|
|
|
|
# check Relax function calls TIR function with call_tir call
|
|
assert rx_func.params[0] == x
|
|
assert rx_func.params[1] == y
|
|
assert rx_func.params[2] == z
|
|
assert rx_func.body.body == out
|
|
assert len(rx_func.body.blocks) == 1
|
|
assert len(rx_func.body.blocks[0].bindings) == 1
|
|
|
|
call_node = rx_func.body.blocks[0].bindings[0].value
|
|
assert isinstance(call_node, rx.Call)
|
|
assert len(call_node.args) == 2
|
|
assert call_node.args[0].name_hint == "te_func"
|
|
assert call_node.args[1][0] == x
|
|
assert call_node.args[1][1] == y
|
|
assert call_node.args[1][2] == z
|
|
|
|
|
|
def test_emit_te_multiple():
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
y = rx.Var("y", rx.TensorType([n, m], "float32"))
|
|
z = rx.Var("z", rx.TensorType([128, m], "float32"))
|
|
|
|
def te_func(A):
|
|
B = te.compute((128, 128), lambda i, j: A[i, j] + 1)
|
|
return B
|
|
|
|
with bb.function("rx_func", [x, y, z]):
|
|
x1 = bb.emit_te(te_func, x)
|
|
y1 = bb.emit_te(te_func, y)
|
|
z1 = bb.emit_te(te_func, z)
|
|
bb.emit_func_output(z1)
|
|
|
|
mod = bb.finalize()
|
|
rx_func = mod["rx_func"]
|
|
|
|
prim_func = []
|
|
for gv in mod.get_global_vars():
|
|
if isinstance(mod[gv], PrimFunc):
|
|
prim_func.append(mod[gv])
|
|
|
|
# only two PrimFuncs were generated since two of them are equal so got deduped
|
|
assert len(prim_func) == 2
|
|
assert rx_func.body.blocks[0].bindings[0].value.args[0].name_hint == "te_func"
|
|
assert rx_func.body.blocks[0].bindings[1].value.args[0].name_hint == "te_func"
|
|
assert rx_func.body.blocks[0].bindings[2].value.args[0].name_hint == "te_func1"
|
|
|
|
|
|
def test_emit_te_multiple_output():
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
|
|
def te_func(A):
|
|
B0, B1 = te.compute((n, m), lambda i, j: (A[i, j] + 1, A[i, j] * 2), name="B")
|
|
return (B0, B1)
|
|
|
|
with bb.function("rx_func", [x]):
|
|
y = bb.emit_te(te_func, x)
|
|
z = rx.TupleGetItem(y, 0)
|
|
bb.emit_func_output([y, z])
|
|
|
|
rx_func = bb.finalize()["rx_func"]
|
|
|
|
# check call tirx output shape is a Tuple of ShapeExpr
|
|
assert rx_func.params[0] == x
|
|
call_node = rx_func.body.blocks[0].bindings[0].value
|
|
assert call_node.args[0].name_hint == "te_func"
|
|
assert isinstance(call_node.ty_args[0], rx.TupleType)
|
|
assert len(call_node.ty_args[0].fields) == 2
|
|
assert isinstance(call_node.ty_args[0].fields[0].shape, rx.ShapeExpr)
|
|
assert isinstance(call_node.ty_args[0].fields[1].shape, rx.ShapeExpr)
|
|
|
|
|
|
def test_emit_te_extern():
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
y = rx.Var("y", rx.TensorType([m, n], "float32"))
|
|
|
|
with bb.function("rx_cblas_matmul", [x, y]):
|
|
out = bb.emit_te(tvm.contrib.cblas.matmul, x, y, transa=False, transb=False)
|
|
bb.emit_func_output(out)
|
|
|
|
mod = bb.finalize()
|
|
rx_func = mod["rx_cblas_matmul"]
|
|
|
|
# check Relax function calls TIR function with call_tir call
|
|
assert rx_func.params[0] == x
|
|
assert rx_func.params[1] == y
|
|
assert len(rx_func.body.blocks) == 1
|
|
call_node = rx_func.body.blocks[0].bindings[0].value
|
|
assert isinstance(call_node, rx.Call)
|
|
assert len(call_node.args) == 2
|
|
assert call_node.args[0].name_hint == "matmul"
|
|
assert call_node.args[1][0] == x
|
|
assert call_node.args[1][1] == y
|
|
assert call_node.ty_args[0].shape[0] == n
|
|
assert call_node.ty_args[0].shape[1] == n
|
|
|
|
|
|
def test_emit_te_prim_value():
|
|
bb = rx.BlockBuilder()
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", R.Tensor([n, m], "float32"))
|
|
a_min = tirx.IntImm("int64", 0)
|
|
a_max = tirx.IntImm("int64", 6)
|
|
|
|
with bb.function("rx_clip", [x]):
|
|
out = bb.emit_te(topi.clip, x, a_min, a_max)
|
|
bb.emit_func_output(out)
|
|
|
|
rx_func = bb.finalize()["rx_clip"]
|
|
|
|
# check Relax function calls TIR function with call_tir call
|
|
assert rx_func.params[0] == x
|
|
assert len(rx_func.body.blocks) == 1
|
|
call_node = rx_func.body.blocks[0].bindings[0].value
|
|
assert isinstance(call_node, rx.Call)
|
|
assert len(call_node.args) == 2
|
|
assert call_node.args[1][0] == x
|
|
|
|
|
|
def test_nested_function_fail():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
with bb.function("func", [x, y]):
|
|
gv0 = bb.emit(rx.op.add(x, x))
|
|
with bb.function("func1", [x, y]):
|
|
gv1 = bb.emit(rx.op.add(x, x))
|
|
bb.emit_func_output(gv0)
|
|
|
|
|
|
def test_emit_func_output_twice_fail():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
with bb.function("func", [x, y]):
|
|
gv0 = bb.emit(rx.op.add(x, y))
|
|
bb.emit_func_output(gv0)
|
|
bb.emit_func_output(gv0)
|
|
|
|
|
|
def test_func_params_twice_fail():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
with bb.function("func", [x, y]):
|
|
gv0 = bb.emit(rx.op.add(x, y))
|
|
bb.emit_func_output(gv0, [x])
|
|
|
|
|
|
def test_no_func_params_fail():
|
|
m = tirx.Var("m", "int64")
|
|
n = tirx.Var("n", "int64")
|
|
x = rx.Var("x", rx.TensorType([m, n], "float16"))
|
|
y = rx.Var("y", rx.TensorType([n], "float16"))
|
|
bb = rx.BlockBuilder()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
with bb.function("func"):
|
|
gv0 = bb.emit(rx.Call(ExternFunc("test.blockbuilder.nop"), []))
|
|
bb.emit_func_output(gv0)
|
|
|
|
|
|
def test_block_builder_scope_recovery():
|
|
bb = rx.BlockBuilder()
|
|
|
|
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
|
|
x = rx.Var("x", rx.TensorType([n, m], "float32"))
|
|
y = rx.Var("y", rx.TensorType([m, n], "float32"))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
# this line fails
|
|
with bb.function("func", [x, y]):
|
|
gv0 = bb.emit(rx.op.add(x, y))
|
|
|
|
# current should be recovered
|
|
assert rx.BlockBuilder.current() is None
|
|
|
|
# second attempt to do it correctly.
|
|
with bb.function("func", [x, y]):
|
|
gv0 = bb.emit(rx.op.add(x, y))
|
|
bb.emit_func_output(gv0)
|
|
|
|
|
|
@pytest.mark.parametrize("emit_nested_tuple", [True, False])
|
|
def test_emit_nested_tuple(emit_nested_tuple):
|
|
"""Convert nested tuples when emitting relax"""
|
|
|
|
def make_function(emit_nested_tuple: bool):
|
|
bb = rx.BlockBuilder()
|
|
|
|
n_sym = tirx.Var("n", "int64")
|
|
m_sym = tirx.Var("m", "int64")
|
|
n = rx.Var("n", tvm.ir.PrimType("int64"))
|
|
m = rx.Var("m", tvm.ir.PrimType("int64"))
|
|
x = rx.Var("x", rx.TensorType([n_sym, m_sym], "float32"))
|
|
y = rx.Var("y", rx.TensorType([m_sym, n_sym], "float32"))
|
|
|
|
with bb.function("func", [n, m, x, y]):
|
|
scalars = (n, m)
|
|
if not emit_nested_tuple:
|
|
scalars = bb.emit(scalars)
|
|
output = (scalars, x, y)
|
|
bb.emit_func_output(output)
|
|
|
|
return bb.finalize()["func"]
|
|
|
|
def make_expected(emit_nested_tuple: bool):
|
|
if emit_nested_tuple:
|
|
|
|
@R.function
|
|
def func(
|
|
n_1: R.Prim("int64"),
|
|
m_1: R.Prim("int64"),
|
|
x: R.Tensor(("n", "m"), dtype="float32"),
|
|
y: R.Tensor(("m", "n"), dtype="float32"),
|
|
):
|
|
return ((n_1, m_1), x, y)
|
|
|
|
else:
|
|
|
|
@R.function
|
|
def func(
|
|
n_1: R.Prim("int64"),
|
|
m_1: R.Prim("int64"),
|
|
x: R.Tensor(("n", "m"), dtype="float32"),
|
|
y: R.Tensor(("m", "n"), dtype="float32"),
|
|
):
|
|
gv = n_1, m_1
|
|
return (gv, x, y)
|
|
|
|
return func
|
|
|
|
expected = make_expected(emit_nested_tuple)
|
|
actual = make_function(emit_nested_tuple)
|
|
|
|
tvm.ir.assert_structural_equal(expected, actual)
|
|
|
|
|
|
@pytest.mark.skip_well_formed_check_before_transform
|
|
def test_finalize_public_private_name_conflict():
|
|
# tirx call
|
|
bb = rx.BlockBuilder()
|
|
|
|
def te_zero():
|
|
return topi.full((), "int64", tirx.IntImm("int64", 0))
|
|
|
|
def te_one():
|
|
return topi.full((), "int64", tirx.IntImm("int64", 1))
|
|
|
|
with bb.function("func", []):
|
|
gv0 = bb.emit_te(te_zero, primfunc_name_hint="func")
|
|
gv1 = bb.emit_te(te_one, primfunc_name_hint="func")
|
|
bb.emit_func_output((gv0, gv1))
|
|
|
|
mod = bb.get()
|
|
assert not rx.analysis.check_well_formed(mod)
|
|
mod_final = bb.finalize()
|
|
rx.analysis.well_formed(mod_final)
|
|
|
|
# relax function call
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", [], private=True):
|
|
gvar = bb.emit_func_output(rx.const(0, "int64"))
|
|
|
|
with bb.function("func", [], private=True):
|
|
gv0 = bb.emit(rx.Call(gvar, []))
|
|
gvar1 = bb.emit_func_output(gv0)
|
|
|
|
with bb.function("func", []):
|
|
gv0 = bb.emit(rx.Call(gvar1, []))
|
|
bb.emit_func_output(gv0)
|
|
|
|
mod = bb.get()
|
|
assert not rx.analysis.check_well_formed(mod)
|
|
mod_final = bb.finalize()
|
|
rx.analysis.well_formed(mod_final)
|
|
|
|
|
|
def test_emit_nested_seqexpr_in_binding_block():
|
|
"""May emit a SeqExpr inside a BindingBlock"""
|
|
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", []):
|
|
lhs = bb.emit(rx.const(1, "int64"), "a")
|
|
rhs = bb.emit(rx.const(2, "int64"), "b")
|
|
out = bb.emit(rx.op.add(lhs, rhs), "c")
|
|
bb.emit_func_output(out)
|
|
|
|
seq_expr = bb.finalize()["func"].body
|
|
|
|
bb = rx.BlockBuilder()
|
|
with bb.function("func", [], private=True):
|
|
lhs = bb.emit(rx.const(3, "int64"), "d")
|
|
rhs = bb.emit(seq_expr, "e")
|
|
out = bb.emit(rx.op.add(lhs, rhs), "f")
|
|
bb.emit_func_output(out)
|
|
|
|
output = bb.finalize()["func"]
|
|
|
|
@R.function(private=True)
|
|
def expected():
|
|
d = R.const(3, "int64")
|
|
a = R.const(1, "int64")
|
|
b = R.const(2, "int64")
|
|
c = R.add(a, b)
|
|
e = c
|
|
f = R.add(d, e)
|
|
return f
|
|
|
|
tvm.ir.assert_structural_equal(expected, output)
|
|
|
|
|
|
def test_emit_nested_dataflow_seqexpr_in_dataflow_block():
|
|
"""May emit a SeqExpr with dataflow inside a DataflowBlock"""
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", []):
|
|
with bb.dataflow():
|
|
lhs = bb.emit(rx.const(1, "int64"), "a")
|
|
rhs = bb.emit(rx.const(2, "int64"), "b")
|
|
out = bb.emit_output(rx.op.add(lhs, rhs), "c")
|
|
bb.emit_func_output(out)
|
|
|
|
seq_expr = bb.finalize()["func"].body
|
|
|
|
bb = rx.BlockBuilder()
|
|
with bb.function("func", [], private=True):
|
|
with bb.dataflow():
|
|
lhs = bb.emit(rx.const(3, "int64"), "d")
|
|
rhs = bb.emit(seq_expr, "e")
|
|
out = bb.emit_output(rx.op.add(lhs, rhs), "f")
|
|
bb.emit_func_output(out)
|
|
|
|
output = bb.finalize()["func"]
|
|
|
|
@R.function(private=True)
|
|
def expected():
|
|
with R.dataflow():
|
|
d = R.const(3, "int64")
|
|
a = R.const(1, "int64")
|
|
b = R.const(2, "int64")
|
|
c = R.add(a, b)
|
|
e = c
|
|
f = R.add(d, e)
|
|
R.output(c, f)
|
|
return f
|
|
|
|
tvm.ir.assert_structural_equal(expected, output)
|
|
|
|
|
|
def test_emit_ill_formed_nested_seqexpr_in_dataflow_block():
|
|
"""May emit a SeqExpr inside a DataflowBlock
|
|
|
|
This produces ill-formed code, but cannot be caught at the
|
|
normalizer. See also
|
|
test_emit_well_formed_nested_seqexpr_in_dataflow_block.
|
|
|
|
"""
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", []):
|
|
lhs = bb.emit(rx.const(1, "int64"), "a")
|
|
rhs = bb.emit(rx.const(2, "int64"), "b")
|
|
out = bb.emit(rx.op.add(lhs, rhs), "c")
|
|
bb.emit_func_output(out)
|
|
|
|
seq_expr = bb.finalize()["func"].body
|
|
|
|
bb = rx.BlockBuilder()
|
|
with bb.function("func", [], private=True):
|
|
with bb.dataflow():
|
|
lhs = bb.emit(rx.const(3, "int64"), "d")
|
|
# This would be ill-formed, as it requires breaking up the
|
|
# DataflowBlock with a BindingBlock.
|
|
rhs = bb.emit(seq_expr, "e")
|
|
|
|
# We cannot throw an error at that point, because it is
|
|
# only the later usage of "d" that results in use of a
|
|
# DataflowVar outside of its home DataflowBlock.
|
|
out = bb.emit_output(rx.op.add(lhs, rhs), "f")
|
|
bb.emit_func_output(out)
|
|
|
|
output = bb.finalize()["func"]
|
|
|
|
assert not rx.analysis.check_well_formed(tvm.ir.IRModule.from_expr(output))
|
|
|
|
|
|
def test_emit_well_formed_nested_seqexpr_in_dataflow_block():
|
|
"""May emit a SeqExpr inside a DataflowBlock
|
|
|
|
This produces well-formed code, and should not have any output
|
|
produced by the normalizer. See also
|
|
test_emit_ill_formed_nested_seqexpr_in_dataflow_block.
|
|
"""
|
|
bb = rx.BlockBuilder()
|
|
|
|
with bb.function("func", []):
|
|
lhs = bb.emit(rx.const(1, "int64"), "a")
|
|
rhs = bb.emit(rx.const(2, "int64"), "b")
|
|
out = bb.emit(rx.op.add(lhs, rhs), "c")
|
|
bb.emit_func_output(out)
|
|
|
|
seq_expr = bb.finalize()["func"].body
|
|
|
|
bb = rx.BlockBuilder()
|
|
with bb.function("func", [], private=True):
|
|
with bb.dataflow():
|
|
lhs = bb.emit(rx.const(3, "int64"), "d")
|
|
# This similarly breaks up the DataflowBlock, with
|
|
# identical steps as the previous test up until this
|
|
# point.
|
|
rhs = bb.emit(seq_expr, "e")
|
|
|
|
# But the "d" variable isn't used, and so there aren't any
|
|
# usages of DataflowVar outside of their home
|
|
# DataflowBlock.
|
|
out = bb.emit_output(rhs, "f")
|
|
bb.emit_func_output(out)
|
|
|
|
output = bb.finalize()["func"]
|
|
|
|
rx.analysis.well_formed(tvm.ir.IRModule.from_expr(output))
|
|
|
|
@R.function(private=True)
|
|
def expected() -> R.Tensor((), dtype="int64"):
|
|
with R.dataflow():
|
|
d = R.const(3, "int64")
|
|
R.output()
|
|
a = R.const(1, "int64")
|
|
b = R.const(2, "int64")
|
|
c = R.add(a, b)
|
|
with R.dataflow():
|
|
e = c
|
|
f = e
|
|
R.output(f)
|
|
return f
|
|
|
|
tvm.ir.assert_structural_equal(expected, output)
|
|
|
|
|
|
def test_error_when_unwrapping_dataflowvar():
|
|
"""Checks for ill-formed use of DataflowVar at normalization
|
|
|
|
We can check for some illegal unwrapping of SeqExpr, though. If
|
|
the inlined non-dataflow SeqExpr uses a DataflowVar, that should
|
|
trigger an error when the SeqExpr is being unwrapped.
|
|
"""
|
|
bb = rx.BlockBuilder()
|
|
|
|
lhs = rx.Var("a", rx.TensorType(shape=[], dtype="int64"))
|
|
|
|
with bb.function("func", [lhs]):
|
|
rhs = rx.const(2, "int64")
|
|
out = bb.emit(rx.op.add(lhs, rhs))
|
|
bb.emit_func_output(out)
|
|
|
|
func = bb.finalize()["func"]
|
|
|
|
bb = rx.BlockBuilder()
|
|
with bb.function("func", [], private=True):
|
|
with bb.dataflow():
|
|
local_lhs = bb.emit(rx.const(3, "int64"), "local_a")
|
|
rhs = bb.emit(func.bind_params({lhs: local_lhs}).body, "f")
|
|
out = bb.emit_output(rhs, "f")
|
|
|
|
with pytest.raises(RuntimeError, match="Malformed AST"):
|
|
bb.emit_func_output(out)
|
|
|
|
|
|
def test_deduplication_when_input_contains_duplicates():
|
|
"""De-duplication of IRModules
|
|
|
|
A well-formed IRModule may contain duplicate function definitions.
|
|
This is rare, as most functions can be disambiguated by the the
|
|
function attribute `tvm::attr::kGlobalSymbol`. However, private
|
|
functions do not have this attribute, and a well-formed IRModule
|
|
may contain multiple copies of the same function.
|
|
|
|
This is a regression test. Previous implementation de-duplicated
|
|
using a `Dict[Function, GlobalVar]`, which has the failure mode
|
|
shown below. This was resolved by de-duplicating using a
|
|
`Dict[Function, Set[GlobalVar]]` instead.
|
|
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Module:
|
|
@R.function
|
|
def main(A: R.Tensor):
|
|
B = Module.subroutine_a(A)
|
|
C = Module.subroutine_b(B)
|
|
return C
|
|
|
|
@R.function(private=True)
|
|
def subroutine_a(arg: R.Tensor) -> R.Tensor:
|
|
return R.add(arg, arg)
|
|
|
|
@R.function(private=True)
|
|
def subroutine_b(arg: R.Tensor) -> R.Tensor:
|
|
return R.add(arg, arg)
|
|
|
|
@R.function(private=True)
|
|
def subroutine_c(arg: R.Tensor) -> R.Tensor:
|
|
return R.multiply(arg, arg)
|
|
|
|
# This test case is only valid when the two subroutines are
|
|
# structurally equal, and therefore allowed to be de-duplicated by
|
|
# the BlockBuilder.
|
|
tvm.ir.assert_structural_equal(Module["subroutine_a"], Module["subroutine_b"])
|
|
|
|
gvar_a = Module.get_global_var("subroutine_a")
|
|
gvar_b = Module.get_global_var("subroutine_b")
|
|
subroutine_c = Module["subroutine_c"]
|
|
|
|
bb = rx.BlockBuilder(Module)
|
|
|
|
# Add a function to the module. What we add doesn't matter, as
|
|
# this is only to initialize the de-duplication map.
|
|
bb.add_func(subroutine_c, "_unused")
|
|
# The deduplication table now maps `subroutine_ab` to either
|
|
# `gvar_a` or `gvar_b`.
|
|
|
|
# Update gvar_a.
|
|
bb.update_func(gvar_a, subroutine_c)
|
|
# The deduplication map no longer has an entry for
|
|
# `subroutine_ab`.
|
|
|
|
# Update gvar_b. The deduplication map is present (because we
|
|
# called `add_func`), but doesn't contain an entry for
|
|
# `subroutine_ab` (because it was just removed). This throws an
|
|
# error.
|
|
bb.update_func(gvar_b, subroutine_c)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|