wehub-resource-sync
273 lines
8.4 KiB
Python
273 lines
8.4 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: F841
|
|
|
|
import re
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.ir.base import assert_structural_equal
|
|
from tvm.script.parser import relax as R
|
|
|
|
|
|
def test_copy_with_new_vars():
|
|
@R.function
|
|
def before(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
after = relax.utils.copy_with_new_vars(before)
|
|
assert_structural_equal(after, before)
|
|
|
|
assert len(after.params) == len(before.params)
|
|
for before_var, after_var in zip(before.params, after.params):
|
|
assert before_var != after_var
|
|
|
|
|
|
def test_copy_with_new_vars_copied_symbolic_vars():
|
|
@R.function
|
|
def before(x: R.Tensor(("m",), "float32"), y: R.Tensor(("m",), "float32")):
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
after = relax.utils.copy_with_new_vars(before)
|
|
assert_structural_equal(after, before)
|
|
|
|
assert len(after.params) == len(before.params)
|
|
for before_var, after_var in zip(before.params, after.params):
|
|
assert before_var != after_var
|
|
assert before_var.ty.shape[0] != after_var.ty.shape[0]
|
|
|
|
|
|
def test_copy_with_new_vars_on_ir_module():
|
|
@tvm.script.ir_module
|
|
class Actual:
|
|
@R.function
|
|
def func(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def func(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
@R.function
|
|
def func_copied(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
Actual["func_copied"] = relax.utils.copy_with_new_vars(Actual["func"]).with_attr(
|
|
"global_symbol", "func_copied"
|
|
)
|
|
|
|
# Assertion will fail if the f_copied contains the same VarNode that's used in
|
|
# the original function, due to var mapping during structural equal.
|
|
assert_structural_equal(Actual, Expected)
|
|
|
|
|
|
def test_copy_with_new_vars_on_ir_module_nested_function():
|
|
@tvm.script.ir_module
|
|
class Actual:
|
|
@R.function
|
|
def func(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
@R.function
|
|
def inner(x: R.Tensor((3,), "float32")) -> R.Tensor((3,), dtype="float32"):
|
|
gv = R.add(x, x)
|
|
return gv
|
|
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def func(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
@R.function
|
|
def inner(x: R.Tensor((3,), "float32")) -> R.Tensor((3,), dtype="float32"):
|
|
gv = R.add(x, x)
|
|
return gv
|
|
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
@R.function
|
|
def func_copied(x: R.Tensor((3,), "float32"), y: R.Tensor((3,), "float32")):
|
|
@R.function
|
|
def inner(x: R.Tensor((3,), "float32")) -> R.Tensor((3,), dtype="float32"):
|
|
gv = R.add(x, x)
|
|
return gv
|
|
|
|
gv = R.add(x, y)
|
|
return gv
|
|
|
|
Actual["func_copied"] = relax.utils.copy_with_new_vars(Actual["func"]).with_attr(
|
|
"global_symbol", "func_copied"
|
|
)
|
|
|
|
assert_structural_equal(Actual, Expected)
|
|
|
|
|
|
def test_assert_structural_equal_in_seqexpr():
|
|
"""The first mismatch is correctly identified."""
|
|
|
|
@R.function(private=True)
|
|
def func_1(A: R.Tensor([16, 16], "float32")):
|
|
B = R.concat([A, A])
|
|
return B
|
|
|
|
@R.function(private=True)
|
|
def func_2(A: R.Tensor([16, 16], "float32")):
|
|
B = R.add(A, A)
|
|
C = R.add(B, B)
|
|
return B
|
|
|
|
with pytest.raises(
|
|
ValueError,
|
|
match=re.escape("<root>.ty.ret.shape.ty.values[0].value"),
|
|
):
|
|
assert_structural_equal(func_1, func_2)
|
|
|
|
|
|
def test_structural_equal_of_call_nodes():
|
|
"""relax.Call must be compared by structural equality, not reference"""
|
|
|
|
# Three identical calls to relax.op.zeros
|
|
calls_to_op_zero = [relax.op.zeros([16], "int32") for _ in range(3)]
|
|
|
|
@R.function(private=True)
|
|
def uses_same_object_twice():
|
|
A = calls_to_op_zero[0]
|
|
B = calls_to_op_zero[0]
|
|
C = R.add(A, B)
|
|
return C
|
|
|
|
@R.function(private=True)
|
|
def uses_two_different_objects():
|
|
A = calls_to_op_zero[1]
|
|
B = calls_to_op_zero[2]
|
|
C = R.add(A, B)
|
|
return C
|
|
|
|
tvm.ir.assert_structural_equal(uses_same_object_twice, uses_two_different_objects)
|
|
|
|
|
|
def test_structural_equal_with_recursive_lambda_function():
|
|
"""A recursive lambda function may be checked for structural equality
|
|
|
|
Recursive function definitions may reference the bound variable
|
|
within the value being bound. In these cases, the `DefEqual(var,
|
|
other->var)` must occur first, to ensure it is defined at point of
|
|
use.
|
|
|
|
In all other cases, checking for structural equality of the bound
|
|
value prior to the variable provides a better error message.
|
|
"""
|
|
|
|
def define_function():
|
|
@R.function
|
|
def func(n: R.Prim("int64")):
|
|
@R.function
|
|
def recursive_lambda(i_arg: R.Prim("int64")) -> R.Prim("int64"):
|
|
condition = R.equal(i_arg, R.prim_value(0))
|
|
if condition:
|
|
output = R.prim_value(0)
|
|
else:
|
|
next_i = R.subtract(i_arg, R.prim_value(1))
|
|
remainder = recursive_lambda(next_i)
|
|
output = R.add(i_arg, remainder)
|
|
return output
|
|
|
|
return recursive_lambda(n)
|
|
|
|
return func
|
|
|
|
func_1 = define_function()
|
|
func_2 = define_function()
|
|
|
|
tvm.ir.assert_structural_equal(func_1, func_2)
|
|
|
|
|
|
def test_structural_equal_with_distinct_recursive_lambda_function():
|
|
"""A recursive lambda function may be checked for structural equality
|
|
|
|
Like `test_structural_equal_with_recursive_lambda_function`, but
|
|
comparing between two distinct functions.
|
|
"""
|
|
|
|
@R.function(private=True)
|
|
def func_a(n: R.Prim("int64")):
|
|
@R.function
|
|
def recursive_lambda(i_arg: R.Prim("int64")) -> R.Prim("int64"):
|
|
condition = R.equal(i_arg, R.prim_value(0))
|
|
if condition:
|
|
output = R.prim_value(0)
|
|
# ^
|
|
# The first mismatch is here
|
|
else:
|
|
next_i = R.subtract(i_arg, R.prim_value(1))
|
|
remainder = recursive_lambda(next_i)
|
|
output = R.add(i_arg, remainder)
|
|
return output
|
|
|
|
return recursive_lambda(n)
|
|
|
|
@R.function(private=True)
|
|
def func_b(n: R.Prim("int64")):
|
|
@R.function
|
|
def recursive_lambda(i_arg: R.Prim("int64")) -> R.Prim("int64"):
|
|
condition = R.equal(i_arg, R.prim_value(0))
|
|
if condition:
|
|
output = R.prim_value(1)
|
|
# ^
|
|
# The first mismatch is here
|
|
else:
|
|
next_i = R.subtract(i_arg, R.prim_value(1))
|
|
remainder = recursive_lambda(next_i)
|
|
output = R.multiply(i_arg, remainder)
|
|
return output
|
|
|
|
return recursive_lambda(n)
|
|
|
|
# The path to the first mismatch, which should appear within the
|
|
# error message.
|
|
mismatch_path = [
|
|
"<root>",
|
|
"body",
|
|
"blocks[0]",
|
|
"bindings[0]",
|
|
"value",
|
|
"body",
|
|
"blocks[0]",
|
|
"bindings[1]",
|
|
"value",
|
|
"true_branch",
|
|
"body",
|
|
"value",
|
|
]
|
|
|
|
with pytest.raises(ValueError, match=re.escape(".".join(mismatch_path))):
|
|
tvm.ir.assert_structural_equal(func_a, func_b)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|