wehub-resource-sync
246 lines
9.4 KiB
Python
246 lines
9.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 pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
|
|
|
|
def check_throws(f):
|
|
try:
|
|
f()
|
|
except RuntimeError:
|
|
pass
|
|
else:
|
|
raise AssertionError("Should have raised an exception but didn't.")
|
|
|
|
|
|
def test_const_fold():
|
|
def check(f, *args):
|
|
x = f(*[tvm.tirx.const(x, "int32") for x in args])
|
|
y = f(*args)
|
|
if not isinstance(x, tvm.tirx.IntImm) or x.value != int(y):
|
|
raise ValueError(f"check error: {x} vs {y} ")
|
|
|
|
tmod = tvm.tirx.truncmod
|
|
check(lambda x, y: x + y, 3, 4)
|
|
check(lambda x, y: x * y, 3, 12)
|
|
check(lambda x, y: x * y - 10, 3, 12)
|
|
check(lambda x, y: x - tmod(y, 10), 3, 12)
|
|
check(lambda x, y: x // y + 10, 100, 12)
|
|
check(lambda x, y: x & y + 10, 112, 128)
|
|
check(lambda x, y: x > y, 112, 128)
|
|
check(lambda x, y: x < y, 112, 128)
|
|
check(lambda x, y: x <= y, 112, 128)
|
|
check(lambda x, y: x >= y, 112, 128)
|
|
check(lambda x, y: (x | y) ^ 10, 112, 128)
|
|
|
|
|
|
def test_const_fold2():
|
|
x = tvm.tirx.Var("x", "int32")
|
|
tmod = tvm.tirx.truncmod
|
|
tdiv = tvm.tirx.truncdiv
|
|
assert (x + 0).same_as(x)
|
|
assert (0 + x).same_as(x)
|
|
assert (x - 0).same_as(x)
|
|
assert tmod(x, 1).value == 0
|
|
assert (x * 1).same_as(x)
|
|
assert (1 * x).same_as(x)
|
|
assert isinstance(tdiv(1, x), tvm.tirx.Div)
|
|
|
|
|
|
def test_const_fold3():
|
|
# Test that using ints with logic operations is forbidden
|
|
x = tvm.tirx.Var("x", "int32")
|
|
for val in [0, 1]:
|
|
for func in [tvm.tirx.all, tvm.tirx.any]:
|
|
check_throws(lambda: func(tvm.tirx.const(val, "bool"), x))
|
|
check_throws(lambda: func(x, tvm.tirx.const(val, "bool")))
|
|
|
|
# Test const folding when both arguments are const
|
|
for tvm_func, py_func in [
|
|
(tvm.tirx.all, lambda a, b: a and b),
|
|
(tvm.tirx.any, lambda a, b: a or b),
|
|
]:
|
|
for v1 in [0, 1]:
|
|
for v2 in [0, 1]:
|
|
tvm.ir.assert_structural_equal(
|
|
tvm_func(tvm.tirx.const(v1, "bool"), tvm.tirx.const(v2, "bool")),
|
|
tvm.tirx.const(py_func(v1, v2), "bool"),
|
|
)
|
|
|
|
x = tvm.tirx.Var("x", "bool")
|
|
true = tvm.tirx.const(1, "bool")
|
|
false = tvm.tirx.const(0, "bool")
|
|
|
|
assert tvm.tirx.all(x, true).same_as(x)
|
|
assert tvm.tirx.all(true, x).same_as(x)
|
|
assert tvm.tirx.any(x, false).same_as(x)
|
|
assert tvm.tirx.any(false, x).same_as(x)
|
|
|
|
assert tvm.tirx.all(x, false).same_as(false)
|
|
assert tvm.tirx.all(false, x).same_as(false)
|
|
assert tvm.tirx.any(x, true).same_as(true)
|
|
assert tvm.tirx.any(true, x).same_as(true)
|
|
|
|
|
|
def test_const_fold4():
|
|
x1 = tvm.tirx.const(4, "int32")
|
|
x2 = x1 + 5
|
|
tdiv = tvm.tirx.truncdiv
|
|
assert isinstance(x2, tvm.tirx.IntImm) and x2.value == 9
|
|
x3 = tdiv(x2, 3)
|
|
assert isinstance(x3, tvm.tirx.IntImm) and x3.value == 3
|
|
x4 = x3 + 0.55
|
|
assert isinstance(x4, tvm.tirx.FloatImm) and abs(x4.value - 3.55) < 1e-6
|
|
x5 = tvm.tirx.ceil(x4)
|
|
assert isinstance(x5, tvm.tirx.FloatImm) and x5.value == 4
|
|
x6 = x5.astype("int")
|
|
assert isinstance(x6, tvm.tirx.IntImm) and x6.value == 4, f"x6={x6}"
|
|
y = (tvm.tirx.round((tvm.tirx.const(6.5, "float32") - 1) / 1.5) + 2).astype("int")
|
|
assert isinstance(y, tvm.tirx.IntImm) and y.value == 6
|
|
|
|
|
|
def test_binary_dtype_match():
|
|
def verify_general_dtype_support(f, is_conditional=False):
|
|
rules = [
|
|
[("bool", "int32"), "int32"],
|
|
[("int32", "float32"), "float32"],
|
|
[("int32", "int64"), "int64"],
|
|
[("uint32", "int8"), "uint32"],
|
|
[("uint32", "int32"), "uint32"],
|
|
]
|
|
for (lhs_dtype, rhs_dtype), out_dtype in rules:
|
|
lhs = tvm.tirx.Var("lhs", lhs_dtype)
|
|
rhs = tvm.tirx.Var("rhs", rhs_dtype)
|
|
out = f(lhs, rhs)
|
|
if not is_conditional:
|
|
assert out.ty.dtype == out_dtype
|
|
else:
|
|
assert out.ty.dtype == "bool"
|
|
if hasattr(out, "a"):
|
|
assert out.a.ty.dtype == out_dtype
|
|
assert out.b.ty.dtype == out_dtype
|
|
elif hasattr(out, "args"):
|
|
# CallOp
|
|
assert out.args[0].ty.dtype == out_dtype
|
|
assert out.args[1].ty.dtype == out_dtype
|
|
else:
|
|
raise ValueError("Unknown binary op format!")
|
|
|
|
def verify_callop_float_only(f):
|
|
for lhs_dtype in ["int32", "float32", "float64"]:
|
|
for rhs_dtype in ["int32", "float32", "float64"]:
|
|
lhs = tvm.tirx.Var("lhs", lhs_dtype)
|
|
rhs = tvm.tirx.Var("rhs", rhs_dtype)
|
|
if "float" not in lhs_dtype and "float" not in rhs_dtype:
|
|
check_throws(lambda: f(lhs, rhs))
|
|
elif "float" in lhs_dtype:
|
|
out = f(lhs, rhs)
|
|
|
|
# Upcasting for floating point types
|
|
dtypes = [lhs_dtype, rhs_dtype]
|
|
if "float64" in dtypes:
|
|
target_dtype = "float64"
|
|
elif "float32" in dtypes:
|
|
target_dtype = "float32"
|
|
else:
|
|
target_dtype = "int32"
|
|
assert out.ty.dtype == target_dtype
|
|
|
|
# Final inputs are the right type
|
|
assert out.args[0].ty.dtype == target_dtype
|
|
assert out.args[1].ty.dtype == target_dtype
|
|
else:
|
|
out = f(lhs, rhs)
|
|
assert out.ty.dtype == rhs_dtype
|
|
assert out.args[0].ty.dtype == rhs_dtype
|
|
assert out.args[1].ty.dtype == rhs_dtype
|
|
|
|
verify_general_dtype_support(lambda a, b: a + b)
|
|
verify_general_dtype_support(lambda a, b: a * b)
|
|
verify_general_dtype_support(lambda a, b: a >= b, is_conditional=True)
|
|
verify_general_dtype_support(lambda a, b: a <= b, is_conditional=True)
|
|
verify_callop_float_only(lambda a, b: tvm.tirx.power(a, b))
|
|
|
|
# verify bool & int32 constant folding
|
|
assert tvm.tirx.const(1) == tvm.tirx.const(True)
|
|
assert tvm.tirx.const(2) != tvm.tirx.const(True)
|
|
|
|
|
|
def test_if_then_else():
|
|
cases = [
|
|
[(tvm.tirx.Var("cond", "bool"), "bool", "int32"), "int32"],
|
|
[(True, "int32", "float32"), "float32"],
|
|
[(False, "int32", "int64"), "int64"],
|
|
[(tvm.tirx.Var("cond", "bool"), "uint32", "int32"), "uint32"],
|
|
[(tvm.tirx.Var("cond", "int32"), "uint32", "int32"), "uint32"],
|
|
]
|
|
for (cond, lhs_dtype, rhs_dtype), out_dtype in cases:
|
|
lhs = tvm.tirx.Var("lhs", lhs_dtype)
|
|
rhs = tvm.tirx.Var("rhs", rhs_dtype)
|
|
if cond is True or cond is False:
|
|
out = tvm.tirx.if_then_else(cond, lhs, rhs)
|
|
out2 = tvm.tirx.if_then_else(not cond, rhs, lhs)
|
|
out3 = tvm.tirx.if_then_else(not cond, lhs, rhs)
|
|
tvm.ir.assert_structural_equal(out, out2) == 1
|
|
if cond:
|
|
tvm.ir.assert_structural_equal(out, lhs.astype(out_dtype)) == 1
|
|
tvm.ir.assert_structural_equal(out3, rhs.astype(out_dtype)) == 1
|
|
else:
|
|
tvm.ir.assert_structural_equal(out, rhs.astype(out_dtype)) == 1
|
|
tvm.ir.assert_structural_equal(out3, lhs.astype(out_dtype)) == 1
|
|
elif cond.ty.dtype == "bool":
|
|
out = tvm.tirx.if_then_else(cond, lhs, rhs)
|
|
assert out.ty.dtype == out_dtype
|
|
assert out.args[1].ty.dtype == out_dtype
|
|
assert out.args[2].ty.dtype == out_dtype
|
|
elif cond.ty.dtype != "bool":
|
|
check_throws(lambda: tvm.tirx.if_then_else(cond, lhs, rhs))
|
|
else:
|
|
raise ValueError("Unknown combinations")
|
|
|
|
|
|
@pytest.mark.parametrize("num_args", list(range(2, 10)))
|
|
def test_comm_reducer(num_args):
|
|
"""Handle all arguments in tirx comm_reducer
|
|
|
|
The `tirx.comm_reducer` API has two distinct usages. It can reduce
|
|
a tensor along a specified axis, similar to numpy.max, or it can
|
|
reduce several arguments together, simililar to Python's built-in
|
|
max(). This choice is based on the type of the second argument.
|
|
|
|
If the `tirx.comm_reducer` is reducing all arguments, then all
|
|
arguments should be used. In the past, the introduction of new
|
|
arguments intended for use when reducing along a tensor axis has
|
|
failed to forward these arguments when reducing along a list of
|
|
items.
|
|
"""
|
|
assert tvm.tirx.max(*range(num_args)) == num_args - 1
|
|
|
|
|
|
def test_llvm_intrin():
|
|
with pytest.raises(ValueError, match=r"Unknown llvm intrinsic function llvm.dummy"):
|
|
a = tvm.tirx.call_llvm_intrin("int32x4", "llvm.dummy")
|
|
with pytest.raises(ValueError, match=r"Unknown llvm intrinsic function llvm.dummy"):
|
|
a = tvm.tirx.call_llvm_pure_intrin("int32x4", "llvm.dummy")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|