wehub-resource-sync
324 lines
11 KiB
Python
324 lines
11 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: RUF012
|
|
|
|
import contextlib
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.arith import ConstIntBound
|
|
|
|
NEG_INF = ConstIntBound.NEG_INF
|
|
POS_INF = ConstIntBound.POS_INF
|
|
|
|
|
|
class TestCase:
|
|
def __init__(self, expr, expected_bounds, known_bounds=None, constraint=None):
|
|
self.expr = expr
|
|
self.expected_bounds = expected_bounds
|
|
if known_bounds is None:
|
|
self.known_bounds = {}
|
|
else:
|
|
self.known_bounds = known_bounds
|
|
|
|
self.constraint = constraint
|
|
|
|
@property
|
|
def __name__(self):
|
|
return str(self.expr)
|
|
|
|
|
|
class BaseCompare:
|
|
def test_const_bounds(self, test_case):
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
for var, bounds in test_case.known_bounds.items():
|
|
analyzer.update(var, ConstIntBound(*bounds))
|
|
assert analyzer.const_int_bound_is_bound(var)
|
|
|
|
with contextlib.ExitStack() as stack:
|
|
if test_case.constraint is not None:
|
|
stack.enter_context(analyzer.constraint_scope(test_case.constraint))
|
|
|
|
bounds = analyzer.const_int_bound(test_case.expr)
|
|
|
|
if test_case.expected_bounds[0] is None:
|
|
assert bounds.max_value == test_case.expected_bounds[1]
|
|
elif test_case.expected_bounds[1] is None:
|
|
assert bounds.min_value == test_case.expected_bounds[0]
|
|
else:
|
|
assert (bounds.min_value, bounds.max_value) == test_case.expected_bounds
|
|
|
|
|
|
class TestDataType(BaseCompare):
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tvm.tirx.Var("x", "int64"), (NEG_INF, POS_INF)),
|
|
TestCase(tvm.tirx.Var("x", "int8"), (-128, 127)),
|
|
TestCase(tvm.tirx.Var("x", "uint8"), (0, 255)),
|
|
TestCase(tvm.tirx.Var("x", "int32"), (-(2**31), 2**31 - 1)),
|
|
)
|
|
|
|
|
|
def test_plain_var_non_negative_bound_requires_context():
|
|
var = tvm.tirx.Var("x", "int64")
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
assert analyzer.const_int_bound(var).min_value == NEG_INF
|
|
with analyzer.constraint_scope(var >= 0):
|
|
assert analyzer.const_int_bound(var).min_value == 0
|
|
assert analyzer.const_int_bound(var).min_value == NEG_INF
|
|
|
|
|
|
class TestCastBound(BaseCompare):
|
|
x = tvm.tirx.Var("x", "int8")
|
|
tmod = tvm.tirx.truncmod
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tmod(x, 3).astype("uint32"), (0, 2)),
|
|
TestCase(tmod(x, 3).astype("float32").astype("int32"), (-2, 2)),
|
|
)
|
|
|
|
|
|
class TestAddSubBound(BaseCompare):
|
|
x = tvm.tirx.Var("x", "int64")
|
|
y = tvm.tirx.Var("y", "int64")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(x + y, (NEG_INF, POS_INF)),
|
|
TestCase(x + y, (1, 14), known_bounds={x: (0, 4), y: (1, 10)}),
|
|
TestCase(x - y, (-10, 3), known_bounds={x: (0, 4), y: (1, 10)}),
|
|
TestCase(x - y, (-10, POS_INF), known_bounds={x: (0, POS_INF), y: (1, 10)}),
|
|
TestCase(1 - x, (NEG_INF, 1), known_bounds={x: (0, POS_INF), y: (1, 10)}),
|
|
)
|
|
|
|
|
|
@pytest.mark.xfail(reason="Not currently supported")
|
|
class TestBoundsUsingReciprocals(BaseCompare):
|
|
"""Special handling for differences of reciprocals
|
|
|
|
These terms can appear when comparing the number of operations for
|
|
different orderings of matrix multiplications, with A, B, and C
|
|
known to be positive values.
|
|
|
|
In these cases, comparing `(A+B)*C < A*B` is equivalent to
|
|
`1/A + 1/B < 1/C`. Working in terms of the reciprocals
|
|
allows the ConstIntBound analyzer to provide a tighter
|
|
bound for these differences than would otherwise be
|
|
available.
|
|
|
|
For `(A+B)*C - A*B`, the normal bottom-up integer bounds are unable to
|
|
provide the bounds required to provide these inequalities, because they
|
|
treat the terms as uncorrelated. That is, they assume that `(A+B)*C` may
|
|
achieve its minimum while `A*B` simultaneously achieves its maximum.
|
|
"""
|
|
|
|
A, B, C = [tvm.tirx.Var(letter, "int64") for letter in "ABC"]
|
|
|
|
symmetric_bounds = {A: (1, 4095), B: (1, 4095), C: (2048, 2048)}
|
|
asymmetric_bounds = {A: (1, 1024), B: (1, POS_INF), C: (2048, 2048)}
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase((A + B) * C - A * B, (2048, None), known_bounds=symmetric_bounds),
|
|
TestCase((A + B) * C - B * A, (2048, None), known_bounds=symmetric_bounds),
|
|
TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),
|
|
TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),
|
|
TestCase((A + B) * C - A * B, (2048, None), known_bounds=asymmetric_bounds),
|
|
TestCase((A + B) * C - B * A, (2048, None), known_bounds=asymmetric_bounds),
|
|
TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),
|
|
TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),
|
|
)
|
|
|
|
|
|
class TestMulBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(x * y + 20, (0, 60), {x: (-2, 4), y: (4, 10)}),
|
|
TestCase(x * y, (-32, 24), {x: (-3, 4), y: (-8, 2)}),
|
|
TestCase(x * y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-8, 2)}),
|
|
)
|
|
|
|
|
|
class TestTruncDivBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
expr = tvm.tirx.truncdiv(x, y)
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(expr, (-2, None), {x: (-9, 4), y: (4, 10)}),
|
|
TestCase(expr, (-4, 9), {x: (-9, 4), y: (-2, 0)}),
|
|
TestCase(expr, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),
|
|
TestCase(expr, (-9, 9), {x: (-9, 4), y: (-4, 12)}),
|
|
)
|
|
|
|
|
|
class TestTruncModBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
expr = tvm.tirx.truncmod(x, y)
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(expr, (-9, 4), {x: (-9, 4), y: (4, 10)}),
|
|
TestCase(expr, (-9, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),
|
|
TestCase(expr, (0, 9), {x: (1, POS_INF), y: (4, 10)}),
|
|
)
|
|
|
|
|
|
class TestFloorDivBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
ux = tvm.tirx.Var("x", "uint32")
|
|
uy = tvm.tirx.Var("y", "uint32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(x // y, (-9 // 4, None), {x: (-9, 4), y: (4, 10)}),
|
|
TestCase(x // y, (-4, 9), {x: (-9, 4), y: (-2, 0)}),
|
|
TestCase(x // y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),
|
|
TestCase(x // y, (-9, 9), {x: (-9, 4), y: (-4, 12)}),
|
|
TestCase(ux // uy, (0, 4), {ux: (1, 4), uy: (0, 12)}),
|
|
)
|
|
|
|
|
|
class TestFloorModBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(x % y, (0, 9), {x: (-9, 4), y: (4, 10)}),
|
|
TestCase(x % y, (0, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),
|
|
TestCase(x % y, (0, 9), {x: (1, POS_INF), y: (4, 10)}),
|
|
)
|
|
|
|
|
|
class TestMinMaxBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tvm.tirx.min(x, y), (-9, 10), {x: (-9, 11), y: (4, 10)}),
|
|
TestCase(tvm.tirx.min(x, y), (NEG_INF, 10), {x: (NEG_INF, POS_INF), y: (4, 10)}),
|
|
TestCase(tvm.tirx.max(x, y), (4, POS_INF), {x: (NEG_INF, POS_INF), y: (4, 10)}),
|
|
TestCase(tvm.tirx.max(x, y), (4, POS_INF), {x: (1, POS_INF), y: (4, 10)}),
|
|
)
|
|
|
|
|
|
class TestSelectBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(
|
|
tvm.tirx.Select(x > 1, (y < 0).astype("int32"), y + 1),
|
|
(0, 11),
|
|
{x: (-9, 11), y: (4, 10)},
|
|
),
|
|
)
|
|
|
|
|
|
class TestShiftAndBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(x >> y, (-3, 2), {x: (-9, 11), y: (2, 10)}),
|
|
TestCase(x & y, (0, 10), {x: (-9, 11), y: (2, 10)}),
|
|
TestCase(x & y, (0, 10), {x: (10, 11), y: (2, 10)}),
|
|
)
|
|
|
|
|
|
class TestMixIndexBound(BaseCompare):
|
|
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
|
tdiv = tvm.tirx.truncdiv
|
|
tmod = tvm.tirx.truncmod
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tmod(x, 8) + tdiv(x, 8) * 8, (0, 24 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),
|
|
TestCase(y + x * 3, (0, 24 * 3 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),
|
|
TestCase(
|
|
tmod(x, 7) + tdiv(x, 7) * 7, (0, (23 // 7) * 7 + 6), {x: (0, 24 - 1), y: (0, 3 - 1)}
|
|
),
|
|
)
|
|
|
|
|
|
class TestLetBound(BaseCompare):
|
|
x = tvm.tirx.Var("x", "int32")
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tvm.tirx.Let(x, 1, x + 1), (2, 2)),
|
|
)
|
|
|
|
|
|
class TestFloorModNegativeDivisor(BaseCompare):
|
|
flm, fld = tvm.tirx.floormod, tvm.tirx.floordiv
|
|
a, b = tvm.tirx.Var("a", "int32"), tvm.tirx.Var("b", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(a % b, (-4, 6), {a: (0, 6), b: (-5, 7)}),
|
|
)
|
|
|
|
|
|
class TestDivModAssumeNoZeroDivisor(BaseCompare):
|
|
"""Divmod non negative expression makes assumption that divide by
|
|
zero won't occur this assumption is important to get best result
|
|
from symbolic shape programs
|
|
"""
|
|
|
|
a, b = tvm.tirx.Var("a", "int32"), tvm.tirx.Var("b", "int32")
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(a // b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),
|
|
TestCase(a % b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),
|
|
)
|
|
|
|
|
|
class TestMultipleCondition(BaseCompare):
|
|
a = tvm.tirx.Var("a", "int32")
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(
|
|
a % 58 - 1,
|
|
(0, None),
|
|
known_bounds={a: (0, 128)},
|
|
constraint=tvm.tirx.all(1 <= a % 58, a % 58 < 57),
|
|
),
|
|
)
|
|
|
|
|
|
class TestBroadcastBound(BaseCompare):
|
|
a = tvm.tirx.Var("a", "int32")
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tvm.tirx.Broadcast(a, 4), (0, 128), {a: (0, 128)}),
|
|
)
|
|
|
|
|
|
class TestRampBound(BaseCompare):
|
|
a = tvm.tirx.Var("a", "int32")
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(tvm.tirx.Ramp(a, 2, 4) + 2, (2, 128 + 2 * 3 + 2), {a: (0, 128)}),
|
|
)
|
|
|
|
|
|
class TestModularSetBound(BaseCompare):
|
|
analyzer = tvm.arith.Analyzer()
|
|
tx = tvm.tirx.Var("tx", "int32")
|
|
bx = tvm.tirx.Var("bx", "int32")
|
|
|
|
expr = (bx * 2048 + tx * 16) % 7168
|
|
|
|
test_case = tvm.testing.parameter(
|
|
TestCase(expr, (0, 7152), {bx: (0, 3584), tx: (0, 128)}),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|