wehub-resource-sync
1354 lines
41 KiB
Python
1354 lines
41 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: E741, F841
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.script
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.ir.base import assert_structural_equal
|
|
from tvm.relax.transform.transform import CanonicalizeBindings
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def verify(input, expected):
|
|
tvm.ir.assert_structural_equal(CanonicalizeBindings()(input), expected)
|
|
|
|
|
|
def test_simple_assignments():
|
|
@I.ir_module
|
|
class TestChainAssignments:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
y = x
|
|
z = y
|
|
q = z
|
|
p = q
|
|
o = p
|
|
return o
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
return x
|
|
|
|
verify(TestChainAssignments, Expected)
|
|
|
|
|
|
def test_dataflow_block():
|
|
@I.ir_module
|
|
class TestDataflowAssignments:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.const(1)
|
|
z = y
|
|
o = z
|
|
p = o
|
|
m = p
|
|
n = m
|
|
R.output(n)
|
|
return n
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
n = R.const(1)
|
|
R.output(n)
|
|
return n
|
|
|
|
verify(TestDataflowAssignments, Expected)
|
|
|
|
|
|
def test_assign_to_output_in_dataflow_block():
|
|
@I.ir_module
|
|
class TestDataflowAssignments:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = x # is not a dataflow var
|
|
z = y
|
|
o = z
|
|
p = o
|
|
m = p
|
|
n = m
|
|
R.output(n)
|
|
return n
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
# we get a dataflow block where the
|
|
# only assignment is n = x, which we can eliminate,
|
|
# resulting in an empty block that is normalized away
|
|
return x
|
|
|
|
verify(TestDataflowAssignments, Expected)
|
|
|
|
|
|
def test_ops():
|
|
@I.ir_module
|
|
class TestOps:
|
|
@R.function
|
|
def main(x: R.Tensor, y: R.Tensor):
|
|
w = y
|
|
q = x
|
|
z = R.add(w, q)
|
|
return R.add(q, z)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor, y: R.Tensor):
|
|
z = R.add(y, x)
|
|
return R.add(x, z)
|
|
|
|
verify(TestOps, Expected)
|
|
|
|
|
|
@pytest.mark.xfail(reason="The lhs and rhs of an assignment should have the same type.")
|
|
def test_casting():
|
|
@I.ir_module
|
|
class TestCasting:
|
|
@R.function
|
|
def main(x: R.Tensor) -> R.Any:
|
|
y = x
|
|
# z will be treated as Any even though it's a tensor
|
|
z: R.Any = y
|
|
return z
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor) -> R.Any:
|
|
# Cannot unify because the cast indicates user intent
|
|
z: R.Any = x
|
|
return z
|
|
|
|
verify(TestCasting, Expected)
|
|
|
|
|
|
def test_match_cast():
|
|
@I.ir_module
|
|
class TestMatchCast:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
q = x
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(q, R.Tensor((m, n)))
|
|
w = z
|
|
return w
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
# can't get rid of z because its ty is different from x's
|
|
m, n = T.int64(), T.int64()
|
|
z = R.match_cast(x, R.Tensor((m, n)))
|
|
return z
|
|
|
|
verify(TestMatchCast, Expected)
|
|
|
|
|
|
def test_same_shape():
|
|
@I.ir_module
|
|
class TestSameShape:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"), "float32")):
|
|
m, n = T.int64(), T.int64()
|
|
y = x
|
|
# trivial check
|
|
z = R.match_cast(x, R.Tensor((m, n), "float32"))
|
|
w = z
|
|
q = R.add(w, y)
|
|
return R.add(q, w)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"), "float32")):
|
|
# the trivial check is canonicalized into a var binding
|
|
# and then eliminated
|
|
q = R.add(x, x)
|
|
return R.add(q, x)
|
|
|
|
verify(TestSameShape, Expected)
|
|
|
|
|
|
def test_change_shape():
|
|
@I.ir_module
|
|
class TestChangeShape:
|
|
@R.function
|
|
def main(x: R.Tensor(ndim=2)):
|
|
y = x
|
|
# The MatchCast is non-trivial, as it introduces new shape
|
|
# vars. Because the input tensor has an unknown shape
|
|
# rather than a symbolic shape, these new shape vars
|
|
# cannot be expressed in terms of previous variables.
|
|
# Therefore, the match cast must be retained.
|
|
o, p = T.int64(), T.int64()
|
|
z = R.match_cast(x, R.Tensor((o, p)))
|
|
w = z
|
|
q = R.add(w, y)
|
|
return R.add(q, w)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(ndim=2)):
|
|
o, p = T.int64(), T.int64()
|
|
z = R.match_cast(x, R.Tensor((o, p)))
|
|
# the ty field on q will need to be updated
|
|
q = R.add(z, x)
|
|
return R.add(q, z)
|
|
|
|
verify(TestChangeShape, Expected)
|
|
|
|
|
|
def test_replace_symbolic_variable_and_remove_match_cast():
|
|
@I.ir_module
|
|
class TestChangeShape:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"))):
|
|
y = x
|
|
# The MatchCast is non-trivial, as it introduces new shape
|
|
# vars. However, the new shape vars are redundant, and
|
|
# are replaced by canonicalization. After replacing the
|
|
# new shape vars, the MatchCast is trivial and may be
|
|
# removed.
|
|
o, p = T.int64(), T.int64()
|
|
z = R.match_cast(x, R.Tensor((o, p)))
|
|
w = z
|
|
q = R.add(w, y)
|
|
return R.add(q, w)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"))):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
q: R.Tensor([m, n]) = R.add(x, x)
|
|
return R.add(q, x)
|
|
|
|
verify(TestChangeShape, Expected)
|
|
|
|
|
|
def test_replace_symbolic_variable_and_remove_match_cast_of_tuple():
|
|
"""Symbolic variables may be defined in R.match_cast of tuple
|
|
|
|
This test is similar to
|
|
`test_replace_symbolic_variable_and_remove_match_cast`, except
|
|
that the MatchCast is performed on a Relax tuple.
|
|
|
|
This is a regression test. Earlier implementations only inferred
|
|
TIR variables from `R.match_cast` of tensors, shapes, and prim
|
|
values, but omitted tuples.
|
|
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tuple(R.Tensor(("m", "n")))):
|
|
y = x
|
|
o, p = T.int64(), T.int64()
|
|
z = R.match_cast(x, R.Tuple(R.Tensor((o, p))))
|
|
w = z
|
|
q = R.add(w[0], y[0])
|
|
return R.add(q, w[0])
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tuple(R.Tensor(("m", "n")))):
|
|
q = R.add(x[0], x[0])
|
|
return R.add(q, x[0])
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_unwrap_tuple():
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor, y: R.Tensor):
|
|
tuple_var = (x, y)
|
|
w = tuple_var[0]
|
|
q = tuple_var[1]
|
|
z = R.add(w, q)
|
|
return R.add(q, z)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor, y: R.Tensor):
|
|
tuple_var = (x, y)
|
|
z = R.add(x, y)
|
|
return R.add(y, z)
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_basic_folding_example():
|
|
@I.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main() -> R.Tensor((), "int32"):
|
|
with R.dataflow():
|
|
y = R.const(1)
|
|
n = y
|
|
R.output(n)
|
|
return n
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main() -> R.Tensor((), "int32"):
|
|
with R.dataflow():
|
|
n = R.const(1)
|
|
R.output(n)
|
|
return n
|
|
|
|
verify(Input, Expected)
|
|
|
|
|
|
def test_fold_match_cast():
|
|
@I.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main() -> R.Tensor((), "int32"):
|
|
with R.dataflow():
|
|
y = R.const(1)
|
|
n = R.match_cast(y, R.Tensor((), "int32"))
|
|
R.output(n)
|
|
return n
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main() -> R.Tensor((), "int32"):
|
|
with R.dataflow():
|
|
# the cast is trivial, so it is removed
|
|
n = R.const(1)
|
|
R.output(n)
|
|
return n
|
|
|
|
verify(Input, Expected)
|
|
|
|
|
|
def test_fold_variables_from_match_cast():
|
|
"""Symbolic variables in R.match_cast may be inferred
|
|
|
|
If the argument to `R.match_cast` has known shape parameters, they
|
|
may be used to infer symbolic shape parameters.
|
|
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
state: R.Tensor([16], dtype="float32"),
|
|
A: R.Tensor([16, 16], dtype="float32"),
|
|
B: R.Tensor([16, 16], dtype="float32"),
|
|
):
|
|
N1 = T.int64()
|
|
M = T.int64()
|
|
N2 = T.int64()
|
|
|
|
# The symbolic variables `N1`, `N2` and `M` are defined by
|
|
# these `R.match_cast` statements. Since the inputs have
|
|
# a known shape, the values of these symbolic variables
|
|
# may be inferred.
|
|
lhs_A = R.match_cast(A, R.Tensor([N1, M], dtype="float32"))
|
|
lhs_B = R.match_cast(B, R.Tensor([N2, M], dtype="float32"))
|
|
rhs = R.match_cast(state, R.Tensor([M], dtype="float32"))
|
|
|
|
# The symbolic shapes propagate downstream.
|
|
lhs: R.Tensor([N1 + N2, M], "float32") = R.concat((lhs_A, lhs_B), axis=0)
|
|
proj_concat: R.Tensor([N1 + N2], "float32") = R.matmul(lhs, rhs)
|
|
proj_A = R.strided_slice(
|
|
proj_concat,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(N1),),
|
|
assume_inbound=False,
|
|
)
|
|
proj_B = R.strided_slice(
|
|
proj_concat,
|
|
[R.prim_value(0)],
|
|
[R.prim_value(N1)],
|
|
[R.prim_value(N1 + N2)],
|
|
assume_inbound=False,
|
|
)
|
|
return (proj_A, proj_B)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
state: R.Tensor([16], dtype="float32"),
|
|
A: R.Tensor([16, 16], dtype="float32"),
|
|
B: R.Tensor([16, 16], dtype="float32"),
|
|
):
|
|
# The function no longer depends on symbolic variables.
|
|
# Shape inference is now propagated using the
|
|
# statically-known shapes.
|
|
|
|
lhs: R.Tensor([32, 16], dtype="float32") = R.concat((A, B), axis=0)
|
|
proj_concat: R.Tensor([32], dtype="float32") = R.matmul(lhs, state)
|
|
proj_A: R.Tensor([16], dtype="float32") = R.strided_slice(
|
|
proj_concat,
|
|
[R.prim_value(0)],
|
|
[R.prim_value(0)],
|
|
[R.prim_value(16)],
|
|
assume_inbound=False,
|
|
)
|
|
proj_B: R.Tensor([16], dtype="float32") = R.strided_slice(
|
|
proj_concat,
|
|
[R.prim_value(0)],
|
|
[R.prim_value(16)],
|
|
[R.prim_value(32)],
|
|
assume_inbound=False,
|
|
)
|
|
return (proj_A, proj_B)
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_inconsistent_match_cast_raises_error():
|
|
"""Symbolic variables from R.match_cast must be consistent
|
|
|
|
All match cast statements must provide consistent definitions for
|
|
symbolic variables. In this test, the value of `M` would be
|
|
inferred as 16 from either `state` or `A`, but would be inferred
|
|
as 32 from `B`.
|
|
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
state: R.Tensor([16], dtype="float32"),
|
|
A: R.Tensor([16, 16], dtype="float32"),
|
|
B: R.Tensor([32, 32], dtype="float32"),
|
|
):
|
|
N1 = T.int64()
|
|
M = T.int64()
|
|
N2 = T.int64()
|
|
|
|
# These R.match_cast statements define inconsistent values
|
|
# for the symbolic shape parameters.
|
|
lhs_A = R.match_cast(A, R.Tensor([N1, M], dtype="float32"))
|
|
lhs_B = R.match_cast(B, R.Tensor([N2, M], dtype="float32"))
|
|
rhs = R.match_cast(state, R.Tensor([M], dtype="float32"))
|
|
|
|
lhs: R.Tensor([N1 + N2, M], "float32") = R.concat((lhs_A, lhs_B), axis=0)
|
|
proj_concat: R.Tensor([N1 + N2], "float32") = R.matmul(lhs, rhs)
|
|
proj_A = R.strided_slice(
|
|
proj_concat,
|
|
(R.prim_value(0),),
|
|
(R.prim_value(0),),
|
|
(R.prim_value(N1),),
|
|
assume_inbound=False,
|
|
)
|
|
proj_B = R.strided_slice(
|
|
proj_concat,
|
|
[R.prim_value(0)],
|
|
[R.prim_value(N1)],
|
|
[R.prim_value(N1 + N2)],
|
|
assume_inbound=False,
|
|
)
|
|
return (proj_A, proj_B)
|
|
|
|
with pytest.raises(ValueError, match="MatchCast statements must be consistent"):
|
|
CanonicalizeBindings()(Before)
|
|
|
|
|
|
def test_match_cast_may_have_distinct_values_in_branches():
|
|
"""Conditional branches may have different values of symbolic variables
|
|
|
|
Here, the value of `N` can be inferred as 16 within the `if`
|
|
branch and as 32 within the `else` branch.
|
|
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(
|
|
state: R.Tensor(["N"], dtype="float32"),
|
|
A: R.Tensor(["M", 16], dtype="float32"),
|
|
B: R.Tensor(["M", 32], dtype="float32"),
|
|
scale: R.Prim("float32"),
|
|
):
|
|
N = T.int64()
|
|
M = T.int64()
|
|
|
|
if N == 16:
|
|
weights: R.Tensor([M, 16], "float32") = A * scale
|
|
weights: R.Tensor([M, N], "float32") = R.match_cast(
|
|
weights, R.Tensor([M, N], "float32")
|
|
)
|
|
weights: R.Tensor([M, N], "float32") = weights * scale
|
|
else:
|
|
weights: R.Tensor([M, 32], "float32") = B * scale
|
|
weights: R.Tensor([M, N], "float32") = R.match_cast(
|
|
weights, R.Tensor([M, N], "float32")
|
|
)
|
|
weights: R.Tensor([M, N], "float32") = weights * scale
|
|
|
|
weights: R.Tensor([M, N], "float32") = weights * scale
|
|
|
|
out: R.Tensor([M], "float32") = R.matmul(weights, state)
|
|
|
|
return out
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(
|
|
state: R.Tensor(["N"], dtype="float32"),
|
|
A: R.Tensor(["M", 16], dtype="float32"),
|
|
B: R.Tensor(["M", 32], dtype="float32"),
|
|
scale: R.Prim("float32"),
|
|
):
|
|
N = T.int64()
|
|
M = T.int64()
|
|
|
|
if N == 16:
|
|
# Prior to the R.match_cast, the
|
|
weights: R.Tensor([M, 16], "float32") = A * scale
|
|
# The scaled weights within the branch may perform
|
|
# shape inference knowing that N==16.
|
|
weights: R.Tensor([M, 16], "float32") = weights * scale
|
|
# The match cast on exiting the if branch restores the
|
|
weights = R.match_cast(weights, R.Tensor([M, N], "float32"))
|
|
else:
|
|
# Prior to the R.match_cast, the
|
|
weights: R.Tensor([M, 32], "float32") = B * scale
|
|
# Within the else-branch, the R.match_cast implies
|
|
# that N==32. While this conflicts with the earlier
|
|
# definition, the two occur in separate branches, so
|
|
# this is legal.
|
|
# The scaled weights within the branch may perform
|
|
# shape inference knowing that N==32.
|
|
weights: R.Tensor([M, 32], "float32") = weights * scale
|
|
weights = R.match_cast(weights, R.Tensor([M, N], "float32"))
|
|
|
|
# Outside of the conditional, we no longer have a known
|
|
# value for N, so this shape inference must be done using
|
|
# a dynamic shape for `N`.
|
|
weights: R.Tensor([M, N], "float32") = weights * scale
|
|
|
|
# After the conditional branch, we no longer have a known
|
|
# value of N, so this shape inference must use the dynamic
|
|
# shape.
|
|
out: R.Tensor([M], "float32") = R.matmul(weights, state)
|
|
|
|
return out
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_multiple_outputs():
|
|
@I.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main():
|
|
with R.dataflow():
|
|
x = R.const(1)
|
|
y = R.const(1)
|
|
z = R.const(1)
|
|
l = x
|
|
m = y
|
|
n = z
|
|
R.output(l, m, n)
|
|
return (l, m, n)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main():
|
|
with R.dataflow():
|
|
l = R.const(1)
|
|
m = R.const(1)
|
|
n = R.const(1)
|
|
R.output(l, m, n)
|
|
return (l, m, n)
|
|
|
|
verify(Input, Expected)
|
|
|
|
|
|
def test_single_output_multiple_nondataflow():
|
|
"""Non-dataflow vars being updated may also be part trivial bindings
|
|
|
|
Like `test_multiple_outputs`, but only `n` is used in the return
|
|
statement.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Input:
|
|
@R.function
|
|
def main():
|
|
with R.dataflow():
|
|
x = R.const(1)
|
|
y = R.const(1)
|
|
z = R.const(1)
|
|
l = x
|
|
m = y
|
|
n = z
|
|
R.output(l, m, n)
|
|
return n
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main():
|
|
with R.dataflow():
|
|
l = R.const(1)
|
|
m = R.const(1)
|
|
n = R.const(1)
|
|
R.output(n)
|
|
return n
|
|
|
|
verify(Input, Expected)
|
|
|
|
|
|
def test_fold_const_to_output():
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main() -> R.Tensor((), "int32"):
|
|
with R.dataflow():
|
|
n = R.const(1)
|
|
R.output(n)
|
|
return n
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main() -> R.Tensor((), "int32"):
|
|
with R.dataflow():
|
|
n = R.const(1)
|
|
R.output(n)
|
|
return R.const(1)
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_canonicalize_var_to_dataflow_var_if_legal():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
DataflowVar instances may only be used inside a DataflowBlock. If
|
|
a trivial binding `y = x` occurs, where `x` is a `DataflowVar` and
|
|
`y` is a `Var`, replacing `y` with `x` may result in usage of a
|
|
`DataflowVar` outside of a `DataflowBlock`.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
z = R.add(y, R.const(1))
|
|
R.output(y, z)
|
|
return z
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
z = R.add(y, R.const(1))
|
|
R.output(z)
|
|
return z
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_update_dataflow_computations_if_var_replacement_occurs():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
DataflowBlocks may produce additional outputs after the first
|
|
output Var, and these additional outputs may be in terms of the
|
|
first output. Computations that depend on a replaced var must be
|
|
updated to remain well-formed.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, R.const(1))
|
|
gv1 = lv1
|
|
gv2 = R.add(lv1, R.const(1))
|
|
R.output(gv1, gv2)
|
|
return (gv1, gv2)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
# lv1 has been replaced with gv1
|
|
gv1 = R.add(x, R.const(1))
|
|
# So gv1 must be used in the computation of gv2
|
|
gv2 = R.add(gv1, R.const(1))
|
|
R.output(gv1, gv2)
|
|
return (gv1, gv2)
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_update_dataflow_computations_if_var_replacement_occurs_after_usage():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
Like test_update_dataflow_computations_if_var_replacement_occurs,
|
|
but the usage of a DataflowVar occurs before the trivial binding
|
|
that causes it to be replaced.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, R.const(1))
|
|
gv2 = R.add(lv1, R.const(1))
|
|
gv1 = lv1
|
|
R.output(gv1, gv2)
|
|
return (gv1, gv2)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
# lv1 has been replaced with gv1
|
|
gv1 = R.add(x, R.const(1))
|
|
# So gv1 must be used in the computation of gv2
|
|
gv2 = R.add(gv1, R.const(1))
|
|
# Even though the trivial binding of "gv1 = lv1"
|
|
# occurred in this position.
|
|
R.output(gv1, gv2)
|
|
return (gv1, gv2)
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_replace_var_with_dataflow_if_all_usage_within_dataflow_block():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
Like `test_update_dataflow_computations_if_var_replacement_occurs`,
|
|
except that `gv1` is not part of the function's return value. When
|
|
deciding which variable to replace, the following logic is applied:
|
|
|
|
1. Normally, when encountering `x = y`, replace usage of `x` with `y`.
|
|
|
|
2. Unless the trivial binding is a `var_x = dataflow_y`, in which case
|
|
replace `dataflow_y` with `var_x` at the point of definition. This
|
|
prevents usage of `dataflow_y` from escaping the dataflow block.
|
|
|
|
3. Unless `var_x` has no usage outside the dataflow block, in which
|
|
case we replace usage of `var_x` with `dataflow_y`.
|
|
|
|
The third rule ensures that canonicalization can occur in a single
|
|
step. Otherwise, the output of this test case would contain a
|
|
non-dataflow var defined within a dataflow block, and only used within
|
|
that dataflow block. (Equivalent to the input for the test case
|
|
`test_canonicalize_var_to_dataflow_var_if_legal`.)
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
lv1 = R.add(x, R.const(1))
|
|
gv1 = lv1
|
|
gv2 = R.add(lv1, R.const(1))
|
|
R.output(gv1, gv2)
|
|
return gv2
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
gv1 = R.add(x, R.const(1))
|
|
gv2 = R.add(gv1, R.const(1))
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_var_to_dataflow_with_trivial_binding():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
Like
|
|
`test_replace_var_with_dataflow_if_all_usage_within_dataflow_block`,
|
|
except the non-DataflowVar is on the right-hand side of the trivial
|
|
binding.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
gv1 = R.add(x, R.const(1))
|
|
lv1 = gv1
|
|
gv2 = R.add(lv1, R.const(1))
|
|
R.output(gv1, gv2)
|
|
return gv2
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
gv1 = R.add(x, R.const(1))
|
|
gv2 = R.add(gv1, R.const(1))
|
|
R.output(gv2)
|
|
return gv2
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_with_updated_ty():
|
|
"""CanonicalizeBindings and Normalizer may both replace a Var
|
|
|
|
If the CanonicalizeBindings pass has no replacements to make for a
|
|
variable, it must still delegate to the ExprMutator. This is because
|
|
a variable replacement may have occurred as part of the IRNormalizer,
|
|
in order to provide better type.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function(private=True)
|
|
def main(A: R.Tensor(("n", 16), dtype="int32")) -> R.Tensor(("n", 16), dtype="int32"):
|
|
# CanonicalizeBindings recognizes this trivial binding, and
|
|
# replaces `B` with `A`.
|
|
B = A
|
|
# The value is updated from `R.add(B,B)` to `R.add(A,A)`.
|
|
# Changing the value triggers struct inference, allowing the
|
|
# shape to be updated to `[n,16]`. This requires a variable
|
|
# replacement, which is tracked by the `ExprMutator`.
|
|
C: R.Tensor(dtype="int32", ndim=2) = R.add(B, B)
|
|
# Replacement of `C` is not explicitly tracked by
|
|
# CanonicalizeBindings. However, if CanonicalizeBindings just
|
|
# returns `GetRef<Var>(var)`, `ExprMutator` cannot apply the
|
|
# replacement, and this will try to return the old
|
|
# version of `C` with `ndim=2`.
|
|
return C
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function(private=True)
|
|
def main(A: R.Tensor(("n", 16), dtype="int32")) -> R.Tensor(("n", 16), dtype="int32"):
|
|
n = T.int64()
|
|
C: R.Tensor([n, 16], "int32") = R.add(A, A)
|
|
return C
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_trivial_binding_to_dataflow_var():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
DataflowVar instances may only be used inside a DataflowBlock. If
|
|
a trivial binding `y = x` occurs, where `x` is a `DataflowVar` and
|
|
`y` is a `Var`, replacing `y` with `x` may result in usage of a
|
|
`DataflowVar` outside of a `DataflowBlock`.
|
|
|
|
If a binding exists solely to convert from DataflowVar into Var,
|
|
then canonicalization replaces the earlier DataflowVar with a Var.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
z = y
|
|
R.output(z)
|
|
return z
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
R.output(y)
|
|
return y
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_multiple_trivial_binding_to_dataflow_var():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
Like test_canonicalize_trivial_binding_to_dataflow_var, but there
|
|
exist multiple trivial bindings to the DataflowVar.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(w: R.Tensor):
|
|
with R.dataflow():
|
|
x = R.add(w, R.const(1))
|
|
y = x
|
|
z = x
|
|
R.output(y, z)
|
|
return (y, z)
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(w: R.Tensor):
|
|
with R.dataflow():
|
|
x = R.add(w, R.const(1))
|
|
R.output(x)
|
|
return (x, x)
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_trivial_var_binding_inside_dataflow_block():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock
|
|
|
|
Canonicalization handles cases where a Var could be replaced by a
|
|
DataflowVar, and where a Var is a trivial binding. If these two
|
|
cases both occur, should produce reasonable results.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
z = y
|
|
R.output(y, z)
|
|
return z
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
R.output(y)
|
|
return y
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_across_non_dataflow_tuple():
|
|
"""Canonicalize Var to DataflowVar inside DataflowBlock"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
z = (y,)
|
|
gv = R.add(z[0], R.const(1))
|
|
R.output(z, gv)
|
|
return gv
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
with R.dataflow():
|
|
y = R.add(x, R.const(1))
|
|
z = (y,)
|
|
gv = R.add(y, R.const(1))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_var_used_in_distinct_df_blocks():
|
|
"""If a var is used only in dataflow blocks,
|
|
but outside of the one where it was originally defined,
|
|
it should be exposed as an output."""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
z = R.add(x, y)
|
|
w = R.multiply(z, y)
|
|
v = R.add(w, x)
|
|
# v must remain exposed!
|
|
R.output(v)
|
|
_ = R.print(format="Hi mom!")
|
|
with R.dataflow():
|
|
a = R.multiply(v, v)
|
|
b = R.add(a, a)
|
|
c = R.subtract(b, a)
|
|
d = R.add(c, c)
|
|
R.output(d)
|
|
return d
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Before, after)
|
|
|
|
|
|
def test_inner_function():
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
|
|
@R.function(pure=False)
|
|
def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
z = R.add(x, y)
|
|
w = R.multiply(x, z)
|
|
v = R.add(y, w)
|
|
R.output(z, w, v)
|
|
_ = R.print(format="oops")
|
|
with R.dataflow():
|
|
a = R.multiply(v, v)
|
|
b = R.add(a, a)
|
|
c = R.multiply(a, b)
|
|
R.output(a, b, c)
|
|
return c
|
|
|
|
z = R.add(x, y)
|
|
w = R.multiply(z, z)
|
|
v = R.divide(w, z)
|
|
R.output(inner_func, z, v, w)
|
|
q = inner_func(w, v)
|
|
with R.dataflow():
|
|
a = R.multiply(q, q)
|
|
b = R.add(a, a)
|
|
c = R.multiply(b, a)
|
|
R.output(a, b, c)
|
|
return c
|
|
|
|
# expected: we do not need to expose all the outputs
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function(pure=False)
|
|
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
|
|
@R.function(pure=False)
|
|
def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
z = R.add(x, y)
|
|
w = R.multiply(x, z)
|
|
v = R.add(y, w)
|
|
R.output(v)
|
|
_ = R.print(format="oops")
|
|
with R.dataflow():
|
|
a = R.multiply(v, v)
|
|
b = R.add(a, a)
|
|
c = R.multiply(a, b)
|
|
R.output(c)
|
|
return c
|
|
|
|
z = R.add(x, y)
|
|
w = R.multiply(z, z)
|
|
v = R.divide(w, z)
|
|
R.output(inner_func, v, w)
|
|
q = inner_func(w, v)
|
|
with R.dataflow():
|
|
a = R.multiply(q, q)
|
|
b = R.add(a, a)
|
|
c = R.multiply(b, a)
|
|
R.output(c)
|
|
return c
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalize_inside_branches():
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
z = R.add(x, y)
|
|
R.output(z)
|
|
if R.const(True):
|
|
with R.dataflow():
|
|
w = R.add(z, z)
|
|
v = R.multiply(w, w)
|
|
# w does not need to be output
|
|
R.output(w, v)
|
|
q = v
|
|
else:
|
|
with R.dataflow():
|
|
w = R.multiply(z, z)
|
|
v = R.add(w, w)
|
|
R.output(w, v)
|
|
q = v
|
|
return q
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
|
|
with R.dataflow():
|
|
z = R.add(x, y)
|
|
R.output(z)
|
|
if R.const(True):
|
|
with R.dataflow():
|
|
w = R.add(z, z)
|
|
v = R.multiply(w, w)
|
|
R.output(v)
|
|
q = v
|
|
else:
|
|
with R.dataflow():
|
|
w = R.multiply(z, z)
|
|
v = R.add(w, w)
|
|
R.output(v)
|
|
q = v
|
|
return q
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_canonicalization_causes_ty_update():
|
|
"""Regression test for failure mode causing undefined variable
|
|
|
|
The ExprMutator is only allowed to update a variable's type
|
|
if the value bound to it has new type. When
|
|
CanonicalizeBindings replaces a trivial binding, this may provide
|
|
better type as a result. If this happens, the
|
|
|
|
In previous implementations, ExprMutator::ReEmitBinding defined a
|
|
remap for `binding->var`, even if the derived class defined a
|
|
replacement by overriding `VisitVarDef`. If the derived class
|
|
defines a new variable binding by overriding `VisitVarDef`, and
|
|
also causes a variable replacement by overriding `VisitExpr` and
|
|
returning a type with different type, then `ExprMutator`
|
|
must check for both `binding->var` *AND* `new_var`. The
|
|
former may be present in the unmodified graph, and the latter may
|
|
be produced by the derived class before delegating to the base
|
|
class.
|
|
"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def transform_params(
|
|
A: R.Tensor(("vocab_size", 4096), dtype="float16"),
|
|
B: R.Tensor((6144, 4096), dtype="float16"),
|
|
):
|
|
with R.dataflow():
|
|
# Trivial binding of `DataFlow = NonDataFlow`.
|
|
# Wherever `C` is used, Canonicalization will attempt
|
|
# to replace it with `B`.
|
|
C = B
|
|
|
|
# RHS contains `(A,C)`, which CanonicalizeBindings
|
|
# replaces with `(A,B)`. Because this changes the
|
|
# RHS, a new LHS (and new type!) will be
|
|
# generated.
|
|
D: R.Tuple(
|
|
R.Tensor(dtype="float16", ndim=2),
|
|
R.Tensor((6144, 4096), dtype="float16"),
|
|
) = (A, C)
|
|
|
|
# Trivial binding of `NonDataFlow = DataFlow`. The
|
|
# definition of `D` will be replaced with a definition
|
|
# of `E`. This definition of `E` will then be updated
|
|
# to have a known shape.
|
|
E = D
|
|
R.output(E)
|
|
|
|
# By the time `E` is encountered at a usage site, the
|
|
# `ExprMutator` must have a replacement for the old
|
|
# version of `E` with `ndim=2` to the new versions of `E`
|
|
# with `shape=[vocab_size,4096]`.
|
|
return E
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def transform_params(
|
|
A: R.Tensor(("vocab_size", 4096), dtype="float16"),
|
|
B: R.Tensor((6144, 4096), dtype="float16"),
|
|
):
|
|
vocab_size = T.int64()
|
|
with R.dataflow():
|
|
E: R.Tuple(
|
|
R.Tensor((vocab_size, 4096), dtype="float16"),
|
|
R.Tensor((6144, 4096), dtype="float16"),
|
|
) = (A, B)
|
|
|
|
R.output(E)
|
|
return E
|
|
|
|
after = relax.transform.CanonicalizeBindings()(Before)
|
|
assert_structural_equal(Expected, after)
|
|
|
|
|
|
def test_unwrap_tuple_of_constant():
|
|
@I.ir_module
|
|
class TestChainAssignments:
|
|
@R.function
|
|
def main():
|
|
tup = (R.const(0, "int64"), R.const(1, "int64"))
|
|
x = tup[0]
|
|
y = tup[1]
|
|
z = R.add(x, y)
|
|
return z
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main():
|
|
tup = (R.const(0, "int64"), R.const(1, "int64"))
|
|
x = tup[0]
|
|
y = tup[1]
|
|
z = R.add(R.const(0, "int64"), R.const(1, "int64"))
|
|
return z
|
|
|
|
verify(TestChainAssignments, Expected)
|
|
|
|
|
|
def test_trivial_binding_of_replaced_non_dataflow_var():
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(param_tuple: R.Tuple([R.Tensor])):
|
|
with R.dataflow():
|
|
A = param_tuple[0]
|
|
B = A
|
|
C = R.add(A, B)
|
|
R.output(A, B, C)
|
|
return C
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(param_tuple: R.Tuple([R.Tensor])):
|
|
with R.dataflow():
|
|
A = param_tuple[0]
|
|
C = R.add(A, A)
|
|
R.output(C)
|
|
return C
|
|
|
|
After = CanonicalizeBindings()(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
def _get_binding_names(mod):
|
|
return [binding.var.name_hint for binding in mod["main"].body.blocks[0].bindings]
|
|
|
|
expected_names = _get_binding_names(Expected)
|
|
after_names = _get_binding_names(After)
|
|
|
|
assert after_names == expected_names
|
|
|
|
|
|
def test_trace_tuple_through_round_trip():
|
|
"""Canonicalize to the orignal tuple, without unwrap/rewrap."""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(param_tuple: R.Tuple([R.Tensor, R.Tensor, R.Tensor])):
|
|
with R.dataflow():
|
|
A = param_tuple[0]
|
|
B = param_tuple[1]
|
|
C = param_tuple[2]
|
|
output = (A, B, C)
|
|
R.output(output)
|
|
return output
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(param_tuple: R.Tuple([R.Tensor, R.Tensor, R.Tensor])):
|
|
with R.dataflow():
|
|
A = param_tuple[0]
|
|
B = param_tuple[1]
|
|
C = param_tuple[2]
|
|
R.output()
|
|
|
|
return param_tuple
|
|
|
|
After = CanonicalizeBindings()(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
def test_trace_partial_tuple_through_round_trip():
|
|
"""Canonicalize to the orignal tuple, without unwrap/rewrap."""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(param_tuple: R.Tuple([R.Tensor, R.Tensor, R.Tensor])):
|
|
with R.dataflow():
|
|
A = param_tuple[0]
|
|
B = param_tuple[1]
|
|
output = (A, B)
|
|
R.output(output)
|
|
return output
|
|
|
|
Expected = Before
|
|
|
|
After = CanonicalizeBindings()(Before)
|
|
tvm.ir.assert_structural_equal(After, Expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|