wehub-resource-sync
598 lines
21 KiB
Python
598 lines
21 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
|
|
import math
|
|
import random
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import tirx
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype, literals",
|
|
[
|
|
["int8", [-128, 0, 127]],
|
|
["uint8", [0, 255]],
|
|
["int32", [-2147483648, 2147483647]],
|
|
["uint32", [0, 4294967295]],
|
|
["int64", [-9223372036854775808, 9223372036854775807]],
|
|
["uint64", [0, 9223372036854775807]],
|
|
],
|
|
)
|
|
def test_tir_make_intimm(dtype, literals):
|
|
for l in literals:
|
|
imm = tirx.const(l, dtype)
|
|
assert imm.value == l, imm
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype, literals",
|
|
[
|
|
["int8", [-129, 128]],
|
|
["uint8", [-1, 256]],
|
|
["int32", [-2147483650, 2147483648]],
|
|
["uint32", [-1, 4294967296]],
|
|
["uint64", [-1, 18446744073709551616]],
|
|
],
|
|
)
|
|
def test_tir_invalid_intimm(dtype, literals):
|
|
for l in literals:
|
|
# Out-of-range positive literals raise a builtin ValueError from
|
|
# the IntImm range check; negative-into-unsigned raises an
|
|
# InternalError ("cannot make uint from negative value") which is a
|
|
# RuntimeError subclass. Accept either.
|
|
with pytest.raises((RuntimeError, ValueError)):
|
|
tirx.const(l, dtype)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype, literals",
|
|
[
|
|
[
|
|
"uint64",
|
|
{
|
|
9223372036854775807: 9223372036854775807,
|
|
18446744073709551615: 18446744073709551615,
|
|
},
|
|
],
|
|
],
|
|
)
|
|
def test_tir_large_py_int_literals(dtype, literals):
|
|
"""
|
|
For large uint value, use LargeUIntImm intrin,
|
|
"""
|
|
for l in literals:
|
|
x = tirx.const(l, dtype)
|
|
if isinstance(x, tirx.IntImm | tirx.FloatImm):
|
|
assert x.value == literals[l]
|
|
else:
|
|
# LargeUIntImm(low32, hi32)
|
|
assert (int(x.args[1]) << 32) + int(x.args[0]) == literals[l]
|
|
|
|
|
|
def test_tir_intimm_overflow():
|
|
assert int(tirx.const(255, "uint8") + tirx.const(1, "uint8")) == 0
|
|
assert int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) == -(2**31)
|
|
assert int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) == 0
|
|
assert int(tirx.const(2**63 - 1, "int64") + tirx.const(1, "int64")) == -(2**63)
|
|
assert int(tirx.const(2**32, "uint64") * tirx.const(2**32, "uint64")) == 0
|
|
# customized int types
|
|
assert int(tirx.const(7, "int4") + tirx.const(1, "int4")) == -8
|
|
assert int(tirx.const(2**39 - 1, "int40") + tirx.const(1, "int40")) == -(2**39)
|
|
|
|
|
|
def compare_float_value(value, expect, msg):
|
|
if math.isfinite(value):
|
|
assert np.abs(value - expect) < 1e-5, f"{value} vs {expect}, {msg}"
|
|
elif math.isnan(value):
|
|
assert math.isnan(expect), f"{value} vs {expect}, {msg}"
|
|
elif math.isinf(value):
|
|
assert math.isinf(expect), f"{value} vs {expect}, {msg}"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype, literals",
|
|
[
|
|
["float16", [-65504.0, 3.14, 65504.0, np.inf, np.nan]],
|
|
["bfloat16", [-3.38953139e38, 3.38953139e38, 3.14]],
|
|
["float32", [np.finfo("float32").min, 3.14, np.finfo("float32").max, np.inf, np.nan]],
|
|
["float64", [np.finfo("float64").min, 3.14, np.finfo("float64").max, np.inf, np.nan]],
|
|
],
|
|
)
|
|
def test_tir_make_floatimm(dtype, literals):
|
|
for l in literals:
|
|
imm = tirx.const(l, dtype)
|
|
compare_float_value(imm.value, l, "imm value should match feed value")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype, literals",
|
|
[
|
|
["float16", [-65505.0, 65505.0]],
|
|
["float32", [-3.402e39, 3.402e39]],
|
|
],
|
|
)
|
|
def test_tir_invalid_floatimm(dtype, literals):
|
|
"""Currently only fp16 and fp32 have range check."""
|
|
for l in literals:
|
|
# FloatImm out-of-range raises a builtin ValueError.
|
|
with pytest.raises(ValueError):
|
|
tirx.const(l, dtype)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
|
|
@pytest.mark.parametrize("literal", [3.14, np.nan, np.inf])
|
|
def test_tir_special_floatimms(dtype, literal):
|
|
x = tirx.const(literal, dtype)
|
|
compare_float_value(x.value, literal, "imm value should match feed value")
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_tir_too_large_literal_f64():
|
|
# Behavior check: if literal f64 value is out of dtype range, the
|
|
# object is still constructed, and eval to infinity.
|
|
@T.prim_func(s_tir=True)
|
|
def imm_overflow_fp64() -> T.float64:
|
|
T.evaluate(T.ret(T.float64(1.7976e309), dtype="float64"))
|
|
|
|
f = tvm.compile(imm_overflow_fp64, target="llvm")
|
|
assert math.isinf(f())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"literal, expect_dtype",
|
|
[
|
|
(256, "int32"),
|
|
(2147483647, "int32"),
|
|
(-2147483648, "int32"),
|
|
(2147483648, "int64"),
|
|
(-2147483649, "int64"),
|
|
(3.14159, "float32"),
|
|
(np.finfo("float32").min, "float32"),
|
|
(np.finfo("float32").max, "float32"),
|
|
(-3.402e39, "float64"),
|
|
(3.402e39, "float64"),
|
|
],
|
|
)
|
|
def test_tir_const_auto_dtype(literal, expect_dtype):
|
|
x = tirx.const(literal, dtype=None)
|
|
assert x.ty.dtype == expect_dtype
|
|
assert x.value == literal
|
|
|
|
|
|
def check_tir_const_fold(
|
|
dtype, foldf, calcf, x_range=None, y_range=None, expect=None, skip_overflow=False
|
|
):
|
|
"""Helper to check constant folding behavior
|
|
|
|
Parameters
|
|
----------
|
|
dtype: str
|
|
Datatype of constants
|
|
|
|
foldf: (x, y) -> z
|
|
Folding function to call
|
|
|
|
calcf: (x, y) -> z
|
|
Compiled calculation function to call
|
|
|
|
x_range: Union[int, float, tuple]
|
|
Single value or value range [min, max]
|
|
|
|
y_range: Union[int, float, tuple]
|
|
Single value or value range [min, max]
|
|
|
|
expect: Union[int, float]
|
|
Expected calculation result
|
|
|
|
skip_overflow: bool
|
|
Skip assertion if the overflow happens
|
|
"""
|
|
seed = random.randint(0, 2147483648)
|
|
np.random.seed(seed)
|
|
ninfo = np.finfo(dtype) if dtype.startswith("float") else np.iinfo(dtype)
|
|
|
|
if x_range is None:
|
|
x_range = (ninfo.min, ninfo.max)
|
|
if isinstance(x_range, int | float):
|
|
x = x_range
|
|
elif dtype.startswith("int") or dtype.startswith("uint"):
|
|
x = np.random.randint(x_range[0], x_range[1] + 1, dtype=dtype)
|
|
else:
|
|
x = np.random.uniform(x_range[0], x_range[1])
|
|
|
|
if y_range is None:
|
|
y_range = (ninfo.min, ninfo.max)
|
|
if isinstance(y_range, int | float):
|
|
y = y_range
|
|
elif dtype.startswith("int") or dtype.startswith("uint"):
|
|
y = np.random.randint(y_range[0], y_range[1] + 1, dtype=dtype)
|
|
else:
|
|
y = np.random.uniform(y_range[0], y_range[1])
|
|
|
|
if skip_overflow:
|
|
py_res = foldf(x, y)
|
|
if isinstance(py_res, tirx.IntImm | tirx.FloatImm):
|
|
py_res = py_res.value
|
|
if not (ninfo.min <= py_res <= ninfo.max):
|
|
# If the result overflow, certain arithmetics is non-defined
|
|
# thus we intentionally do not make the test failed.
|
|
return
|
|
|
|
fold_res = foldf(tirx.const(x, dtype), tirx.const(y, dtype))
|
|
calc_res = calcf(x, y)
|
|
|
|
flaky_msg = (
|
|
f"{dtype} ({x}, {y}, {expect}) const folding check failed.\n"
|
|
+ "This test is intentionally non-deterministic, "
|
|
+ f"if it fails please report it in GitHub issue together with this seed {seed}\n"
|
|
)
|
|
if dtype.startswith("float"):
|
|
compare_float_value(calc_res, fold_res.value, flaky_msg)
|
|
if expect:
|
|
compare_float_value(expect, calc_res, flaky_msg)
|
|
else:
|
|
assert calc_res == fold_res.value, flaky_msg
|
|
if expect:
|
|
assert expect == calc_res, flaky_msg
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_tir_floatimm_const_fold():
|
|
"""Behavior check: folding fp32 match platform f32 arithmetic"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def float_imm_multiply(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
|
|
z[()] = x * y
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def float_imm_add(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
|
|
z[()] = x + y
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def float_imm_sub(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
|
|
z[()] = x - y
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def float_imm_div(x: T.float32, y: T.float32, z: T.Buffer((), "float32")):
|
|
z[()] = x / y
|
|
|
|
def __wrap_build(f):
|
|
lib = tvm.compile(f, target="llvm")
|
|
z = tvm.runtime.tensor(np.zeros([]).astype("float32"))
|
|
|
|
def _func(x, y):
|
|
lib(x, y, z)
|
|
return z.numpy()
|
|
|
|
return _func
|
|
|
|
fmul = __wrap_build(float_imm_multiply)
|
|
fadd = __wrap_build(float_imm_add)
|
|
fsub = __wrap_build(float_imm_sub)
|
|
fdiv = __wrap_build(float_imm_div)
|
|
|
|
# overflow
|
|
check_tir_const_fold("float32", lambda x, y: x * y, fmul, 3.0e30, 3.0e30, np.inf)
|
|
check_tir_const_fold("float32", lambda x, y: x * y, fmul, 3.0e30, -3.0e30, -np.inf)
|
|
check_tir_const_fold("float32", lambda x, y: x / y, fdiv, 3.0e30, 3.0e-30, np.inf)
|
|
|
|
# divide by zero
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("float32", lambda x, y: x / y, fdiv, 1.0, 0.0)
|
|
|
|
# nan and inf
|
|
check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, np.nan, np.nan)
|
|
check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, np.inf, np.inf)
|
|
check_tir_const_fold("float32", lambda x, y: x + y, fadd, 1.0, -np.inf, -np.inf)
|
|
|
|
# randomized check
|
|
check_tir_const_fold("float32", lambda x, y: x * y, fmul)
|
|
check_tir_const_fold("float32", lambda x, y: x + y, fadd)
|
|
check_tir_const_fold("float32", lambda x, y: x - y, fsub)
|
|
check_tir_const_fold(
|
|
"float32", lambda x, y: x / y, fdiv, y_range=(0.01, np.finfo("float32").max)
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_tir_int8_const_fold():
|
|
"""Behavior check: folding i8 operation match platform i8 arithmetic"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_multiply(x: T.int8, y: T.int8) -> T.int8:
|
|
T.evaluate(T.ret(x * y, dtype="int8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_add(x: T.int8, y: T.int8) -> T.int8:
|
|
T.evaluate(T.ret(x + y, dtype="int8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_sub(x: T.int8, y: T.int8) -> T.int8:
|
|
T.evaluate(T.ret(x - y, dtype="int8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_truncdiv(x: T.int8, y: T.int8) -> T.int8:
|
|
T.evaluate(T.ret(T.truncdiv(x, y), dtype="int8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_floordiv(x: T.int8, y: T.int8) -> T.int8:
|
|
T.evaluate(T.ret(T.floordiv(x, y), dtype="int8"))
|
|
|
|
fmul = tvm.compile(imm_multiply, target="llvm")
|
|
fadd = tvm.compile(imm_add, target="llvm")
|
|
fsub = tvm.compile(imm_sub, target="llvm")
|
|
ffloordiv = tvm.compile(imm_floordiv, target="llvm")
|
|
ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
|
|
|
|
# overflow
|
|
check_tir_const_fold("int8", lambda x, y: x + y, fadd, 127, 1, -128)
|
|
check_tir_const_fold("int8", lambda x, y: x * y, fmul, 127, 127, 1)
|
|
|
|
# divide by zero
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("int8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
|
|
|
|
# i8 mod folding is not implemented
|
|
assert not isinstance(tirx.floormod(tirx.const(7, "int8"), tirx.const(3, "int8")), tirx.IntImm)
|
|
assert not isinstance(tirx.truncmod(tirx.const(7, "int8"), tirx.const(3, "int8")), tirx.IntImm)
|
|
|
|
# randomized check
|
|
check_tir_const_fold("int8", lambda x, y: x * y, fmul)
|
|
check_tir_const_fold("int8", lambda x, y: x + y, fadd)
|
|
check_tir_const_fold("int8", lambda x, y: x - y, fsub)
|
|
check_tir_const_fold(
|
|
"int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("int8").max)
|
|
)
|
|
check_tir_const_fold(
|
|
"int8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, y_range=(1, np.iinfo("int8").max)
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_tir_uint8_const_fold():
|
|
"""Behavior check: folding u8 operation match platform u8 arithmetic"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_multiply(x: T.uint8, y: T.uint8) -> T.uint8:
|
|
T.evaluate(T.ret(x * y, dtype="uint8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_add(x: T.uint8, y: T.uint8) -> T.uint8:
|
|
T.evaluate(T.ret(x + y, dtype="uint8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_sub(x: T.uint8, y: T.uint8) -> T.uint8:
|
|
T.evaluate(T.ret(x - y, dtype="uint8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_truncdiv(x: T.uint8, y: T.uint8) -> T.uint8:
|
|
T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint8"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_floordiv(x: T.uint8, y: T.uint8) -> T.uint8:
|
|
T.evaluate(T.ret(T.floordiv(x, y), dtype="uint8"))
|
|
|
|
fmul = tvm.compile(imm_multiply, target="llvm")
|
|
fadd = tvm.compile(imm_add, target="llvm")
|
|
fsub = tvm.compile(imm_sub, target="llvm")
|
|
ffloordiv = tvm.compile(imm_floordiv, target="llvm")
|
|
ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
|
|
|
|
# overflow
|
|
check_tir_const_fold("uint8", lambda x, y: x + y, fadd, 255, 1, 0)
|
|
|
|
# zero sub
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("uint8", lambda x, y: x - y, fsub, 0, 10)
|
|
|
|
# divide by zero
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("uint8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
|
|
|
|
# u8 floormod folding is overflow-free and implemented
|
|
folded_floormod = tirx.floormod(tirx.const(7, "uint8"), tirx.const(3, "uint8"))
|
|
assert isinstance(folded_floormod, tirx.IntImm)
|
|
assert int(folded_floormod) == 1
|
|
|
|
# u8 truncmod folding is not implemented
|
|
assert not isinstance(
|
|
tirx.truncmod(tirx.const(7, "uint8"), tirx.const(3, "uint8")), tirx.IntImm
|
|
)
|
|
|
|
# randomized check
|
|
check_tir_const_fold("uint8", lambda x, y: x * y, fmul)
|
|
check_tir_const_fold("uint8", lambda x, y: x + y, fadd)
|
|
check_tir_const_fold("uint8", lambda x, y: x - y, fsub)
|
|
check_tir_const_fold(
|
|
"uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("uint8").max)
|
|
)
|
|
check_tir_const_fold(
|
|
"uint8", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, y_range=(1, np.iinfo("uint8").max)
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_tir_int32_const_fold():
|
|
"""Behavior check: folding i32 operation match platform i32 arithmetic"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_multiply(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(x * y, dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_add(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(x + y, dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_sub(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(x - y, dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_truncdiv(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(T.truncdiv(x, y), dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_truncmod(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(T.truncmod(x, y), dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_floordiv(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(T.floordiv(x, y), dtype="int32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_floormod(x: T.int32, y: T.int32) -> T.int32:
|
|
T.evaluate(T.ret(T.floormod(x, y), dtype="int32"))
|
|
|
|
fmul = tvm.compile(imm_multiply, target="llvm")
|
|
fadd = tvm.compile(imm_add, target="llvm")
|
|
fsub = tvm.compile(imm_sub, target="llvm")
|
|
ffloordiv = tvm.compile(imm_floordiv, target="llvm")
|
|
ffloormod = tvm.compile(imm_floormod, target="llvm")
|
|
ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
|
|
ftruncmod = tvm.compile(imm_truncmod, target="llvm")
|
|
|
|
# i32 overflow is not specified, only check for range
|
|
assert -(2**31) <= int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) < 2**31
|
|
assert -(2**31) <= int(tirx.const(-(2**31), "int32") - tirx.const(1, "int32")) < 2**31
|
|
|
|
# divide by zero
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("int32", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("int32", lambda x, y: tirx.floormod(x, y), ffloormod, 1, 0)
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("int32", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("int32", lambda x, y: tirx.truncmod(x, y), ftruncmod, 1, 0)
|
|
|
|
# randomized check
|
|
check_tir_const_fold("int32", lambda x, y: x * y, fmul, skip_overflow=True)
|
|
check_tir_const_fold("int32", lambda x, y: x + y, fadd, skip_overflow=True)
|
|
check_tir_const_fold("int32", lambda x, y: x - y, fsub, skip_overflow=True)
|
|
check_tir_const_fold(
|
|
"int32",
|
|
lambda x, y: tirx.floordiv(x, y),
|
|
ffloordiv,
|
|
y_range=(1, np.iinfo("int32").max),
|
|
skip_overflow=True,
|
|
)
|
|
check_tir_const_fold(
|
|
"int32",
|
|
lambda x, y: tirx.truncdiv(x, y),
|
|
ftruncdiv,
|
|
y_range=(1, np.iinfo("int32").max),
|
|
skip_overflow=True,
|
|
)
|
|
check_tir_const_fold(
|
|
"int32",
|
|
lambda x, y: tirx.floormod(x, y),
|
|
ffloormod,
|
|
y_range=(1, np.iinfo("int32").max),
|
|
skip_overflow=False,
|
|
)
|
|
check_tir_const_fold(
|
|
"int32",
|
|
lambda x, y: tirx.truncmod(x, y),
|
|
ftruncmod,
|
|
y_range=(1, np.iinfo("int32").max),
|
|
skip_overflow=False,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
|
def test_tir_uint32_const_fold():
|
|
"""Behavior check: folding u32 operation match platform u32 arithmetic"""
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_multiply(x: T.uint32, y: T.uint32) -> T.uint32:
|
|
T.evaluate(T.ret(x * y, dtype="uint32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_add(x: T.uint32, y: T.uint32) -> T.uint32:
|
|
T.evaluate(T.ret(x + y, dtype="uint32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_sub(x: T.uint32, y: T.uint32) -> T.uint32:
|
|
T.evaluate(T.ret(x - y, dtype="uint32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_truncdiv(x: T.uint32, y: T.uint32) -> T.uint32:
|
|
T.evaluate(T.ret(T.truncdiv(x, y), dtype="uint32"))
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def imm_floordiv(x: T.uint32, y: T.uint32) -> T.uint32:
|
|
T.evaluate(T.ret(T.floordiv(x, y), dtype="uint32"))
|
|
|
|
fmul = tvm.compile(imm_multiply, target="llvm")
|
|
fadd = tvm.compile(imm_add, target="llvm")
|
|
fsub = tvm.compile(imm_sub, target="llvm")
|
|
ffloordiv = tvm.compile(imm_floordiv, target="llvm")
|
|
ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
|
|
|
|
# u32 overflow is not specified, only check for range
|
|
assert 0 <= int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) < 2**32
|
|
|
|
# divide by zero
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("uint32", lambda x, y: tirx.floordiv(x, y), ffloordiv, 1, 0)
|
|
with pytest.raises(RuntimeError):
|
|
check_tir_const_fold("uint32", lambda x, y: tirx.truncdiv(x, y), ftruncdiv, 1, 0)
|
|
|
|
# u32 floormod folding is overflow-free and implemented
|
|
folded_floormod = tirx.floormod(tirx.const(7, "uint32"), tirx.const(3, "uint32"))
|
|
assert isinstance(folded_floormod, tirx.IntImm)
|
|
assert int(folded_floormod) == 1
|
|
|
|
# u32 truncmod folding is not implemented
|
|
assert not isinstance(
|
|
tirx.truncmod(tirx.const(7, "uint32"), tirx.const(3, "uint32")), tirx.IntImm
|
|
)
|
|
|
|
# randomized check
|
|
check_tir_const_fold("uint32", lambda x, y: x * y, fmul, skip_overflow=True)
|
|
check_tir_const_fold("uint32", lambda x, y: x + y, fadd, skip_overflow=True)
|
|
check_tir_const_fold("uint32", lambda x, y: x - y, fsub, skip_overflow=True)
|
|
check_tir_const_fold(
|
|
"uint32",
|
|
lambda x, y: tirx.floordiv(x, y),
|
|
ffloordiv,
|
|
y_range=(1, np.iinfo("uint32").max),
|
|
skip_overflow=False,
|
|
)
|
|
check_tir_const_fold(
|
|
"uint32",
|
|
lambda x, y: tirx.truncdiv(x, y),
|
|
ftruncdiv,
|
|
y_range=(1, np.iinfo("uint32").max),
|
|
skip_overflow=False,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|