chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,456 @@
|
||||
# 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 tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.arith.analyzer import Analyzer
|
||||
|
||||
|
||||
class IntSetChecker:
|
||||
def __init__(self):
|
||||
self.analyzer = tvm.arith.Analyzer()
|
||||
|
||||
def verify(self, data, dmap, expected):
|
||||
res = self.analyzer.int_set(data, dmap)
|
||||
|
||||
def err_msg():
|
||||
return f"\ndata={data}\ndmap={dmap}\nres={res}\nexpected={expected}"
|
||||
|
||||
assert self.analyzer.can_prove_equal(res.min_value, expected[0]), err_msg()
|
||||
assert self.analyzer.can_prove_equal(res.max_value, expected[1]), err_msg()
|
||||
|
||||
|
||||
def test_basic():
|
||||
s = tvm.arith.IntervalSet(2, 3)
|
||||
assert s.min_value.value == 2
|
||||
assert s.max_value.value == 3
|
||||
|
||||
s = tvm.arith.IntSet.single_point(2)
|
||||
assert s.min_value.value == 2
|
||||
assert s.max_value.value == 2
|
||||
|
||||
|
||||
def test_vector():
|
||||
base = 10
|
||||
stride = 3
|
||||
lanes = 2
|
||||
s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, stride, lanes))
|
||||
assert s.min_value.value == base
|
||||
assert s.max_value.value == base + stride * (lanes - 1)
|
||||
|
||||
|
||||
def test_scalable_vector():
|
||||
base = 5
|
||||
s = tvm.arith.IntSet.vector(tvm.tirx.Ramp(base, 2, tvm.tirx.vscale() * 4))
|
||||
|
||||
assert s.min_value.value == base
|
||||
assert s.max_value.same_as(tvm.arith.int_set.pos_inf())
|
||||
|
||||
|
||||
def test_add_sub():
|
||||
ck = IntSetChecker()
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10)}, (y, 10 + y))
|
||||
ck.verify(x + y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (1, 21))
|
||||
ck.verify(x - y, {x: tvm.arith.IntervalSet(0, 10), y: tvm.arith.IntervalSet(1, 11)}, (-11, 9))
|
||||
|
||||
|
||||
def test_mul_div():
|
||||
ck = IntSetChecker()
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
|
||||
tdiv = tvm.tirx.truncdiv
|
||||
ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True)
|
||||
ck.verify(x * y, {x: tvm.arith.IntervalSet(0, 10)}, (0, 10 * y))
|
||||
ck.verify(x * 2, {x: tvm.arith.IntervalSet(1, 10)}, (2, 20))
|
||||
ck.verify(x * -2, {x: tvm.arith.IntervalSet(1, 10)}, (-20, -2))
|
||||
|
||||
ck.verify(tdiv(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, tdiv(10, y)))
|
||||
ck.verify(tdiv(x, 2), {x: tvm.arith.IntervalSet(1, 10)}, (0, 5))
|
||||
|
||||
fld = tvm.tirx.floordiv
|
||||
ck.verify(fld(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, fld(10, y)))
|
||||
ck.verify(fld(x, 2), {x: tvm.arith.IntervalSet(-1, 10)}, (-1, 5))
|
||||
|
||||
|
||||
def test_mod():
|
||||
ck = IntSetChecker()
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
tmod = tvm.tirx.truncmod
|
||||
ck.analyzer.update(y, tvm.arith.ConstIntBound(1, 100), override=True)
|
||||
ck.verify(tmod(x, y), {x: tvm.arith.IntervalSet(0, 10)}, (0, y - 1))
|
||||
ck.verify(tmod(x, 10), {x: tvm.arith.IntervalSet(1, 10)}, (0, 9))
|
||||
|
||||
flm = tvm.tirx.floormod
|
||||
ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(-10, 10)}, (0, 9))
|
||||
ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 5)}, (3, 5))
|
||||
ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(13, 15)}, (3, 5))
|
||||
ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 15)}, (0, 9))
|
||||
ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 11)}, (0, 9))
|
||||
ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(1, 21)}, (0, 9))
|
||||
|
||||
fld = tvm.tirx.floordiv
|
||||
z = tvm.tirx.Var("z", "int32")
|
||||
ck.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 3))
|
||||
ck.verify(
|
||||
flm(y, 8),
|
||||
{y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)},
|
||||
(
|
||||
z * 8 + x * 4 - 8 * fld(z * 8 + x * 4, 8),
|
||||
z * 8 + x * 4 + 3 - 8 * fld(z * 8 + x * 4, 8),
|
||||
),
|
||||
)
|
||||
ck1 = IntSetChecker()
|
||||
ck1.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 2))
|
||||
ck1.verify(
|
||||
flm(y, 8), {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, (x * 4, x * 4 + 3)
|
||||
)
|
||||
|
||||
|
||||
def test_max_min():
|
||||
ck = IntSetChecker()
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
ck.verify(tvm.tirx.max(x, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (1, 11))
|
||||
ck.verify(tvm.tirx.min(x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 9))
|
||||
ck.verify(tvm.tirx.min(x, y), {}, (tvm.tirx.min(x, y), tvm.tirx.min(x, y)))
|
||||
ck.verify(tvm.tirx.max(x, y), {}, (tvm.tirx.max(x, y), tvm.tirx.max(x, y)))
|
||||
|
||||
|
||||
def test_select():
|
||||
ck = IntSetChecker()
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
ck.verify(tvm.tirx.Select(x > 0, x - 1, x + 1), {x: tvm.arith.IntervalSet(0, 10)}, (-1, 11))
|
||||
|
||||
|
||||
def check_region_bound(expect_region, var_dom, mode, predicate=None):
|
||||
"""Helper to check region bound estimation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
expect_region: dict
|
||||
The keys are of form (begin, end) or Expr as a single point. The values are
|
||||
expected estimated region or region dict on different bindings.
|
||||
|
||||
var_dom: dict
|
||||
Map var to iteration domain range.
|
||||
|
||||
mode: str
|
||||
Specify "lowerbound", "upperbound" or else use strict bound estimation.
|
||||
|
||||
predicate: Expr
|
||||
Extra predicate, defaults to True.
|
||||
"""
|
||||
if predicate is None:
|
||||
predicate = tvm.tirx.IntImm("bool", 1)
|
||||
region = []
|
||||
expect = []
|
||||
for k, v in expect_region.items():
|
||||
if not isinstance(k, tuple | list):
|
||||
k = (k, k + 1)
|
||||
region.append(tvm.ir.Range.from_min_extent(k[0], Analyzer().simplify(k[1] - k[0])))
|
||||
expect.append(v)
|
||||
if mode == "lowerbound":
|
||||
result = tvm.arith.estimate_region_lower_bound(
|
||||
region=region, var_dom=var_dom, predicate=predicate
|
||||
)
|
||||
elif mode == "upperbound":
|
||||
result = tvm.arith.estimate_region_upper_bound(
|
||||
region=region, var_dom=var_dom, predicate=predicate
|
||||
)
|
||||
else:
|
||||
result = tvm.arith.estimate_region_strict_bound(
|
||||
region=region, var_dom=var_dom, predicate=predicate
|
||||
)
|
||||
if result is None:
|
||||
assert all([_ is None for _ in expect])
|
||||
return
|
||||
assert len(result) == len(expect)
|
||||
for intset, expect_desc in zip(result, expect):
|
||||
if isinstance(expect_desc, dict):
|
||||
# check range on different free var bindings
|
||||
for binding in expect_desc:
|
||||
analyzer = Analyzer()
|
||||
for k, v in binding:
|
||||
analyzer.bind(k, v)
|
||||
expect_begin, expect_end = expect_desc[binding]
|
||||
result_begin = analyzer.simplify(intset.min_value, 3)
|
||||
result_end = analyzer.simplify(intset.max_value + 1, 3)
|
||||
assert analyzer.can_prove_equal(result_begin - expect_begin, 0), (
|
||||
f"{result_begin} vs {expect_begin}"
|
||||
)
|
||||
assert analyzer.can_prove_equal(result_end - expect_end, 0), (
|
||||
f"{result_end} vs {expect_end}"
|
||||
)
|
||||
else:
|
||||
# check range
|
||||
expect_begin, expect_end = expect_desc
|
||||
analyzer = Analyzer()
|
||||
assert analyzer.can_prove_equal(intset.min_value - expect_begin, 0), (
|
||||
f"{intset.min_value} vs {expect_begin}"
|
||||
)
|
||||
assert analyzer.can_prove_equal(intset.max_value - expect_end + 1, 0), (
|
||||
f"{intset.max_value} vs {expect_end - 1}"
|
||||
)
|
||||
|
||||
|
||||
def test_region_bound_not_independent():
|
||||
# (i, i+2) and (i+2, i+4) are dependent, this the lowerbound is not available
|
||||
i = tvm.tirx.Var("i", "int32")
|
||||
var_dom = {
|
||||
i: tvm.ir.Range(begin=0, end=64),
|
||||
}
|
||||
check_region_bound({(i, i + 2): None, (i + 2, i + 4): None}, var_dom, mode="lowerbound")
|
||||
check_region_bound({(i, i + 2): (0, 65), (i + 2, i + 4): (2, 67)}, var_dom, mode="upperbound")
|
||||
|
||||
# when only a subset of access indices are affine
|
||||
i, j, k = tvm.tirx.Var("i", "int32"), tvm.tirx.Var("j", "int32"), tvm.tirx.Var("k", "int32")
|
||||
var_dom = {
|
||||
i: tvm.ir.Range(begin=0, end=16),
|
||||
j: tvm.ir.Range(begin=0, end=16),
|
||||
k: tvm.ir.Range(begin=0, end=16),
|
||||
}
|
||||
check_region_bound(
|
||||
{i // 4: None, j * 4 + i % 4: None, tirx.truncdiv(k, 2): None},
|
||||
var_dom,
|
||||
predicate=j * 4 + i % 4 > 3,
|
||||
mode="lowerbound",
|
||||
)
|
||||
check_region_bound(
|
||||
{i // 4: (0, 4), j * 4 + i % 4: (4, 64), tirx.truncdiv(k, 2): (0, 8)},
|
||||
var_dom,
|
||||
predicate=j * 4 + i % 4 > 3,
|
||||
mode="upperbound",
|
||||
)
|
||||
|
||||
|
||||
def test_region_bound_stride_too_wide():
|
||||
i = tvm.tirx.Var("i", "int32")
|
||||
var_dom = {i: tvm.ir.Range(begin=0, end=64)}
|
||||
check_region_bound({(i * 4, i * 4 + 2): None}, var_dom, mode="lowerbound")
|
||||
check_region_bound({(i * 4, i * 4 + 2): (0, 254)}, var_dom, mode="upperbound")
|
||||
|
||||
|
||||
def test_region_bound_small_stride():
|
||||
i = tvm.tirx.Var("i", "int32")
|
||||
var_dom = {
|
||||
i: tvm.ir.Range(begin=0, end=64),
|
||||
}
|
||||
check_region_bound({(i * 4, i * 4 + 8): (0, 260)}, var_dom, mode="lowerbound")
|
||||
|
||||
|
||||
def test_region_lower_bound_split_predicate():
|
||||
x_o = tvm.tirx.Var("xo", "int32")
|
||||
x_i = tvm.tirx.Var("xi", "int32")
|
||||
x = x_o * 4 + x_i
|
||||
var_dom = {
|
||||
x_o: tvm.ir.Range(begin=0, end=16),
|
||||
x_i: tvm.ir.Range(begin=0, end=4),
|
||||
}
|
||||
check_region_bound({(x * 4, x * 4 + 8): (0, 256)}, var_dom, predicate=x < 63, mode="lowerbound")
|
||||
|
||||
check_region_bound(
|
||||
{(x * 4, x * 4 + 8): (0, 256), (x * 3, x * 3 + 5): (0, 191)},
|
||||
var_dom,
|
||||
predicate=x < 63,
|
||||
mode="upperbound",
|
||||
)
|
||||
|
||||
|
||||
def test_region_lower_bound_multiple_variables():
|
||||
div = tvm.tirx.floordiv
|
||||
mod = tvm.tirx.floormod
|
||||
x = tvm.tirx.Var("x", "int32")
|
||||
wid = tvm.tirx.Var("wid", "int32")
|
||||
i = div(x, 16)
|
||||
j = div(mod(x, 16), 4) * 8 + mod(x, 4) + div(wid, 32) * 4
|
||||
k = wid % 32
|
||||
var_dom = {
|
||||
x: tvm.ir.Range(begin=0, end=32),
|
||||
wid: tvm.ir.Range(begin=0, end=64),
|
||||
}
|
||||
check_region_bound({i: (0, 2), j: (0, 32), k: (0, 32)}, var_dom, mode="lowerbound")
|
||||
|
||||
|
||||
def test_region_lower_bound_negative_scale():
|
||||
i = tvm.tirx.Var("i", "int32")
|
||||
j = tvm.tirx.Var("j", "int32")
|
||||
var_dom = {
|
||||
i: tvm.ir.Range(begin=0, end=4),
|
||||
j: tvm.ir.Range(begin=0, end=4),
|
||||
}
|
||||
check_region_bound(
|
||||
{(1 - i, 5 - i): (-2, 5), (20 - j * 4, 36 - j * 4): (8, 36)}, var_dom, mode="lowerbound"
|
||||
)
|
||||
|
||||
|
||||
def test_region_lower_bound_for_non_perfect_tile():
|
||||
h1 = tvm.tirx.Var("h1", "int32")
|
||||
h2 = tvm.tirx.Var("h2", "int32")
|
||||
h3 = tvm.tirx.Var("h3", "int32")
|
||||
|
||||
# non-uniform tiling, single inner variable
|
||||
var_dom = {
|
||||
h2: tvm.ir.Range(begin=0, end=10),
|
||||
}
|
||||
check_region_bound(
|
||||
{
|
||||
h3 * 8 + h2: {
|
||||
(): (
|
||||
tvm.tirx.max(h3 * 8, 1),
|
||||
tvm.tirx.min(0, h3 * 8 - 214) + 224,
|
||||
),
|
||||
((h3, 0),): (1, 10), # h3 == 0: region is [1, 10)
|
||||
((h3, 10),): (h3 * 8, h3 * 8 + 10), # 0 < h3 <= 26: region is [h3 * 8, h3 * 8 + 10)
|
||||
((h3, 27),): (h3 * 8, 224), # h3 > 26: region is [h3 * 8, 224)
|
||||
}
|
||||
},
|
||||
var_dom,
|
||||
predicate=tvm.tirx.all(1 <= h3 * 8 + h2, h3 * 8 + h2 < 224),
|
||||
mode="lowerbound",
|
||||
)
|
||||
|
||||
# non-uniform tiling, two inner variables
|
||||
var_dom = {
|
||||
h1: tvm.ir.Range(begin=0, end=5),
|
||||
h2: tvm.ir.Range(begin=0, end=2),
|
||||
}
|
||||
check_region_bound(
|
||||
{
|
||||
h3 * 8 + h2 * 5 + h1: {
|
||||
(): (
|
||||
tvm.tirx.max(h3 * 8, 1),
|
||||
tvm.tirx.min(0, h3 * 8 - 214) + 224,
|
||||
),
|
||||
((h3, 0),): (1, 10),
|
||||
((h3, 10),): (h3 * 8, h3 * 8 + 10),
|
||||
((h3, 27),): (h3 * 8, 224),
|
||||
}
|
||||
},
|
||||
var_dom,
|
||||
predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h2 * 5 + h1 < 224),
|
||||
mode="lowerbound",
|
||||
)
|
||||
|
||||
# lowerbound should fail on incompatible predicates
|
||||
check_region_bound(
|
||||
{h3 * 8 + h2 * 5 + h1: None},
|
||||
var_dom,
|
||||
predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h1 * 2 + h2 < 224),
|
||||
mode="lowerbound",
|
||||
)
|
||||
check_region_bound(
|
||||
{h3 * 8 + h2 * 5 + h1: (h3 * 8, h3 * 8 + 10)},
|
||||
var_dom,
|
||||
predicate=tvm.tirx.all(1 <= h3 * 8 + h2 * 5 + h1, h3 * 8 + h1 * 2 + h2 < 224),
|
||||
mode="upperbound",
|
||||
)
|
||||
|
||||
|
||||
def test_region_lower_bound_unfusable():
|
||||
var_dom = {
|
||||
tvm.tirx.Var("i", "int32"): tvm.ir.Range(8),
|
||||
tvm.tirx.Var("j", "int32"): tvm.ir.Range(4),
|
||||
}
|
||||
i, j = var_dom
|
||||
check_region_bound({(i + j) // 2: (0, 6)}, var_dom, mode="lowerbound")
|
||||
|
||||
|
||||
def test_union_lower_bound():
|
||||
neg_inf = tvm.arith.int_set.neg_inf()
|
||||
pos_inf = tvm.arith.int_set.pos_inf()
|
||||
set_0 = tvm.arith.IntervalSet(min_value=neg_inf, max_value=0)
|
||||
set_1 = tvm.arith.IntervalSet(min_value=1, max_value=pos_inf)
|
||||
result = tvm.arith.int_set.union_lower_bound([set_0, set_1])
|
||||
assert result.min_value.same_as(neg_inf)
|
||||
assert result.max_value.same_as(pos_inf)
|
||||
set_2 = tvm.arith.IntervalSet(min_value=pos_inf, max_value=neg_inf)
|
||||
result = tvm.arith.int_set.union_lower_bound([set_0, set_1, set_2])
|
||||
assert result.min_value.same_as(neg_inf)
|
||||
assert result.max_value.same_as(pos_inf)
|
||||
|
||||
|
||||
def test_modular_set():
|
||||
ck = IntSetChecker()
|
||||
x = tvm.tirx.Var("x", "int32")
|
||||
y = tvm.tirx.Var("y", "int32")
|
||||
expr = (x * 2048 + y * 16) % 7168
|
||||
ck.verify(
|
||||
expr, {x: tvm.arith.IntervalSet(0, 128), y: tvm.arith.IntervalSet(0, 3584)}, (0, 7152)
|
||||
)
|
||||
|
||||
|
||||
def test_relax_deep_variable_dependency_chain():
|
||||
"""Regression test for exponential variable-relaxation blowup.
|
||||
|
||||
When a variable's interval bound references another variable that is also in
|
||||
the domain map, the evaluator relaxes it transitively. A diamond-shaped
|
||||
chain -- where each variable's bound references the next one in *both* its
|
||||
min and its max -- used to be re-expanded along every path, costing
|
||||
O(2^depth) and hanging indefinitely. The relaxation is now memoized per
|
||||
variable, so this completes in linear time.
|
||||
"""
|
||||
ck = IntSetChecker()
|
||||
n = 64 # 2^64 expansions without memoization; trivially fast with it.
|
||||
xs = [tvm.tirx.Var(f"x{i}", "int32") for i in range(n + 1)]
|
||||
dmap = {xs[i]: tvm.arith.IntervalSet(xs[i + 1] - 1, xs[i + 1] + 1) for i in range(n)}
|
||||
dmap[xs[n]] = tvm.arith.IntervalSet(0, 100)
|
||||
# x0 relaxes through the whole chain: [0 - n, 100 + n].
|
||||
ck.verify(xs[0], dmap, (-n, 100 + n))
|
||||
|
||||
|
||||
def test_relax_cyclic_variable_dependency():
|
||||
"""A cyclic variable dependency must terminate (and stay symbolic)."""
|
||||
ana = tvm.arith.Analyzer()
|
||||
x = tvm.tirx.Var("x", "int32")
|
||||
y = tvm.tirx.Var("y", "int32")
|
||||
# x depends on y and y depends on x: relaxation must not loop forever.
|
||||
dmap = {x: tvm.arith.IntervalSet(y, y), y: tvm.arith.IntervalSet(x, x)}
|
||||
res = ana.int_set(x, dmap)
|
||||
assert res is not None
|
||||
|
||||
|
||||
def test_estimate_region_accepts_external_analyzer():
|
||||
i = tvm.tirx.Var("i", "int32")
|
||||
tile = tvm.tirx.Var("tile", "int32")
|
||||
region = [tvm.ir.Range.from_min_extent(i % tile, 1)]
|
||||
dom = {i: tvm.ir.Range(0, 16)}
|
||||
|
||||
# Without knowing `tile`, the affine detection fails for exact bounds.
|
||||
assert tvm.arith.estimate_region_lower_bound(region, dom, True) is None
|
||||
assert tvm.arith.estimate_region_strict_bound(region, dom, True) is None
|
||||
upper_without_analyzer = tvm.arith.estimate_region_upper_bound(region, dom, True)
|
||||
|
||||
analyzer = tvm.arith.Analyzer()
|
||||
analyzer.bind(tile, tvm.tirx.const(4, "int32"))
|
||||
# The external binding lets the affine detection succeed.
|
||||
for estimate_region in [
|
||||
tvm.arith.estimate_region_lower_bound,
|
||||
tvm.arith.estimate_region_strict_bound,
|
||||
tvm.arith.estimate_region_upper_bound,
|
||||
]:
|
||||
result = estimate_region(region, dom, True, analyzer=analyzer)
|
||||
assert result is not None
|
||||
assert analyzer.can_prove_equal(result[0].min_value, 0)
|
||||
assert analyzer.can_prove_equal(result[0].max_value, 3)
|
||||
|
||||
# The upper-bound fallback without analyzer is safe but much wider.
|
||||
assert not analyzer.can_prove_equal(upper_without_analyzer[0].min_value, 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user