chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,323 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user