Tianqi Chen
382 lines
12 KiB
Python
382 lines
12 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
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
import tvm.topi.testing
|
|
from tvm import relax
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_basic():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
b: R.Tensor([16], "float32"),
|
|
c: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
expr = a
|
|
expr = R.add(expr, b)
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
params: R.Tuple(R.Tensor([16], "float32"), R.Tensor([16], "float32")),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
expr = a
|
|
b = params[0]
|
|
expr = R.add(expr, b)
|
|
c = params[1]
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
mod = Before
|
|
after = relax.transform.BundleModelParams()(mod)
|
|
tvm.ir.assert_structural_equal(after, Expected)
|
|
|
|
|
|
def test_no_model_params():
|
|
"""If all parameters are inputs, model params should be an empty tuple
|
|
|
|
This ensures that a caller does not need to check whether the
|
|
model has compile-time inputs, and can instead provide the output
|
|
of a lifted parameter transformation in all cases, even if that
|
|
transformation returns an empty tuple.
|
|
"""
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
b: R.Tensor([16], "float32"),
|
|
c: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 3})
|
|
expr = a
|
|
expr = R.add(expr, b)
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
b: R.Tensor([16], "float32"),
|
|
c: R.Tensor([16], "float32"),
|
|
params: R.Tuple(),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 3})
|
|
expr = a
|
|
expr = R.add(expr, b)
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
mod = Before
|
|
after = relax.transform.BundleModelParams()(mod)
|
|
tvm.ir.assert_structural_equal(after, Expected)
|
|
|
|
|
|
def test_dataflow():
|
|
"""Parameters can be substituted into a dataflow block"""
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
b: R.Tensor([16], "float32"),
|
|
c: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
with R.dataflow():
|
|
expr = a
|
|
expr = R.add(expr, b)
|
|
expr = R.add(expr, c)
|
|
R.output(expr)
|
|
return expr
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
params: R.Tuple(R.Tensor([16], "float32"), R.Tensor([16], "float32")),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
with R.dataflow():
|
|
expr = a
|
|
b = params[0]
|
|
expr = R.add(expr, b)
|
|
c = params[1]
|
|
expr = R.add(expr, c)
|
|
R.output(expr)
|
|
return expr
|
|
|
|
mod = Before
|
|
after = relax.transform.BundleModelParams()(mod)
|
|
tvm.ir.assert_structural_equal(after, Expected)
|
|
|
|
|
|
def test_variable_names():
|
|
"""Parameters retain their names within the updated function
|
|
|
|
For readability, the parameter names should be used to generate
|
|
the new variable names.
|
|
|
|
Like `test_basic`, but explicitly checks the names of bound
|
|
variables.
|
|
"""
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
b: R.Tensor([16], "float32"),
|
|
c: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
expr = a
|
|
expr = R.add(expr, b)
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
params: R.Tuple(R.Tensor([16], "float32"), R.Tensor([16], "float32")),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
expr = a
|
|
b = params[0]
|
|
expr = R.add(expr, b)
|
|
c = params[1]
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
mod = Before
|
|
after = relax.transform.BundleModelParams()(mod)
|
|
tvm.ir.assert_structural_equal(after, Expected)
|
|
|
|
for binding, expected_binding in zip(
|
|
after["main"].body.blocks[0].bindings,
|
|
Expected["main"].body.blocks[0].bindings,
|
|
):
|
|
assert binding.var.name == expected_binding.var.name
|
|
|
|
|
|
def test_bundled_param_name():
|
|
"""The tuple parameter can have an explicit name"""
|
|
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
b: R.Tensor([16], "float32"),
|
|
c: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
expr = a
|
|
expr = R.add(expr, b)
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
a: R.Tensor([16], "float32"),
|
|
custom_tuple_name: R.Tuple(R.Tensor([16], "float32"), R.Tensor([16], "float32")),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
expr = a
|
|
b = custom_tuple_name[0]
|
|
expr = R.add(expr, b)
|
|
c = custom_tuple_name[1]
|
|
expr = R.add(expr, c)
|
|
return expr
|
|
|
|
mod = Before
|
|
after = relax.transform.BundleModelParams("custom_tuple_name")(mod)
|
|
tvm.ir.assert_structural_equal(after, Expected)
|
|
|
|
for param, expected_param in zip(after["main"].params, Expected["main"].params):
|
|
assert param.name == expected_param.name
|
|
|
|
|
|
def test_nested_function_preserves_outer_rewrite_context():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor([16], "float32"),
|
|
weight: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
|
|
@R.function
|
|
def inner(
|
|
inner_x: R.Tensor([16], "float32"),
|
|
inner_weight: R.Tensor([16], "float32"),
|
|
) -> R.Tensor([16], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
return R.add(inner_x, inner_weight)
|
|
|
|
return R.add(x, weight)
|
|
|
|
after = relax.transform.BundleModelParams()(Before)
|
|
relax.analysis.well_formed(after)
|
|
|
|
func = after["main"]
|
|
assert len(func.params) == 2
|
|
assert isinstance(func.params[1].ty, relax.TupleType)
|
|
|
|
|
|
def test_bundled_tensor_preserves_retained_shape_dependency():
|
|
x = relax.Var("x", relax.TensorType(dtype="float32", ndim=1))
|
|
shape = relax.Var("shape", relax.ShapeType(ndim=1))
|
|
weight = relax.Var("weight", relax.TensorType(shape, "float32"))
|
|
|
|
bb = relax.BlockBuilder()
|
|
with bb.function("main", [x, shape, weight]):
|
|
with bb.dataflow():
|
|
out = bb.emit(relax.op.add(x, weight))
|
|
output = bb.emit_output(out)
|
|
bb.emit_func_output(output)
|
|
before = bb.get()
|
|
global_var = before.get_global_var("main")
|
|
before.update_func(global_var, before[global_var].with_attr("num_input", 2))
|
|
relax.analysis.well_formed(before)
|
|
|
|
after = relax.transform.BundleModelParams()(before)
|
|
relax.analysis.well_formed(after)
|
|
|
|
func = after[global_var]
|
|
bundled_ty = func.params[2].ty
|
|
assert isinstance(bundled_ty, relax.TupleType)
|
|
assert isinstance(bundled_ty.fields[0], relax.TensorType)
|
|
assert bundled_ty.fields[0].shape.same_as(func.params[1])
|
|
|
|
|
|
def test_primitive_model_param_is_materialized_outside_if_branches():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
cond: R.Tensor((), "bool"),
|
|
extent: R.Prim("int64"),
|
|
weight: R.Tensor(["extent"], "float32"),
|
|
) -> R.Tensor(["extent"], "float32"):
|
|
R.func_attr({"num_input": 1})
|
|
if cond:
|
|
out = R.add(weight, weight)
|
|
else:
|
|
out = R.multiply(weight, weight)
|
|
return out
|
|
|
|
after = relax.transform.BundleModelParams()(Before)
|
|
relax.analysis.well_formed(after)
|
|
|
|
func = after["main"]
|
|
extent_binding = func.body.blocks[0].bindings[0]
|
|
assert isinstance(extent_binding, relax.VarBinding)
|
|
assert tvm.ir.is_prim_var(extent_binding.var)
|
|
if_expr = next(
|
|
binding.value
|
|
for block in func.body.blocks
|
|
for binding in block.bindings
|
|
if isinstance(binding, relax.VarBinding) and isinstance(binding.value, relax.If)
|
|
)
|
|
for branch in [if_expr.true_branch, if_expr.false_branch]:
|
|
match_cast = next(
|
|
binding
|
|
for block in branch.blocks
|
|
for binding in block.bindings
|
|
if isinstance(binding, relax.MatchCast)
|
|
)
|
|
assert match_cast.ty.shape[0].same_as(extent_binding.var)
|
|
|
|
|
|
def test_primitive_model_param_remains_linked_to_dependent_tensor():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
x: R.Tensor(dtype="float32", ndim=1),
|
|
extent: R.Prim("int64"),
|
|
weight: R.Tensor(["extent"], "float32"),
|
|
):
|
|
R.func_attr({"num_input": 1})
|
|
out = R.add(x, weight)
|
|
return out
|
|
|
|
old_extent = Before["main"].params[1]
|
|
after = relax.transform.BundleModelParams()(Before)
|
|
relax.analysis.well_formed(after)
|
|
|
|
func = after["main"]
|
|
bundled_ty = func.params[1].ty
|
|
assert isinstance(bundled_ty, relax.TupleType)
|
|
assert isinstance(bundled_ty.fields[0], tvm.ir.PrimType)
|
|
assert isinstance(bundled_ty.fields[1], relax.TensorType)
|
|
assert bundled_ty.fields[1].shape is None
|
|
assert bundled_ty.fields[1].ndim == 1
|
|
|
|
bindings = list(func.body.blocks[0].bindings)
|
|
extent_index = next(
|
|
i
|
|
for i, binding in enumerate(bindings)
|
|
if isinstance(binding, relax.VarBinding) and tvm.ir.is_prim_var(binding.var)
|
|
)
|
|
match_cast_index = next(
|
|
i for i, binding in enumerate(bindings) if isinstance(binding, relax.MatchCast)
|
|
)
|
|
extent = bindings[extent_index].var
|
|
weight = bindings[match_cast_index]
|
|
assert extent_index < match_cast_index
|
|
assert weight.ty.shape[0].same_as(extent)
|
|
|
|
signature_types = [param.ty for param in func.params] + [func.ret_ty]
|
|
assert all(
|
|
all(not var.same_as(old_extent) for var in relax.analysis.tir_vars_in_type(ty))
|
|
for ty in signature_types
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|