wehub-resource-sync
273 lines
11 KiB
Python
273 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: F401, F841
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.tirx.buffer import decl_buffer
|
|
|
|
|
|
def test_deduce():
|
|
a = tvm.tirx.Var("a", "int32")
|
|
b = tvm.tirx.Var("b", "int32")
|
|
c = tvm.tirx.Var("c", "int32")
|
|
d = tvm.tirx.Var("d", "int32")
|
|
|
|
b_s = tvm.arith.IntervalSet(2, 3)
|
|
c_s = tvm.arith.IntervalSet(10, 15)
|
|
d_s = tvm.arith.IntervalSet(-3, -1)
|
|
zero = tvm.tirx.const(0, "int32")
|
|
|
|
fdiv = tvm.tirx.floordiv
|
|
|
|
e0 = (-b) * a + c - d
|
|
res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
|
|
ans0 = fdiv(d - c, b * -1)
|
|
tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
|
|
|
|
# expression containing variable a is on rhs
|
|
res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
|
|
tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
|
|
|
|
e0 = d * a + c - d
|
|
res0 = tvm.arith.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
|
|
ans0 = fdiv(d - c, d)
|
|
tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
|
|
|
|
# expression containing variable a is on rhs
|
|
res0 = tvm.arith.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
|
|
tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
|
|
|
|
e1 = a * 4 + b < c
|
|
res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
|
|
ans1 = fdiv(c - 1 - b, 4)
|
|
tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
|
|
|
|
# expression containing variable a is on rhs
|
|
e1 = c > a * 4 + b
|
|
res1 = tvm.arith.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
|
|
tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
|
|
|
|
e2 = tvm.tirx.max(5, a * 4) < 0
|
|
res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
|
|
assert str(res2.max_value) == "neg_inf"
|
|
assert str(res2.min_value) == "pos_inf"
|
|
|
|
# expression containing variable a is on rhs
|
|
e2 = zero < tvm.tirx.max(5, a * 4)
|
|
res2 = tvm.arith.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
|
|
assert str(res2.max_value) == "neg_inf"
|
|
assert str(res2.min_value) == "pos_inf"
|
|
|
|
e3 = (-b) + a * c - d
|
|
res3 = tvm.arith.deduce_bound(a, e3 >= 0, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
|
|
ans3 = fdiv(2, c) + 1
|
|
tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
|
|
|
|
res3 = tvm.arith.deduce_bound(a, zero <= e3, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
|
|
tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
|
|
|
|
# tests for `EQ` op
|
|
res4 = tvm.arith.deduce_bound(a, a == b, {}, {})
|
|
tvm.testing.assert_prim_expr_equal(res4.max_value, b)
|
|
tvm.testing.assert_prim_expr_equal(res4.min_value, b)
|
|
|
|
# Unsatisfiable `EQ`, variable as one of the Operand
|
|
res5 = tvm.arith.deduce_bound(a, (a == b), {b: b_s}, {b: b_s})
|
|
assert str(res5.max_value) == "neg_inf"
|
|
assert str(res5.min_value) == "pos_inf"
|
|
|
|
# variable `a` on the RHS side
|
|
res6 = tvm.arith.deduce_bound(a, 10 == a, {}, {})
|
|
tvm.testing.assert_prim_expr_equal(res6.max_value, 10)
|
|
tvm.testing.assert_prim_expr_equal(res6.min_value, 10)
|
|
|
|
# Add, Sub in `EQ`
|
|
e4 = (a - c) == (b + d)
|
|
ans4 = b + d + c
|
|
res7 = tvm.arith.deduce_bound(a, e4, {b: b_s, c: c_s, d: d_s}, {})
|
|
tvm.testing.assert_prim_expr_equal(res7.max_value, ans4)
|
|
tvm.testing.assert_prim_expr_equal(res7.min_value, ans4)
|
|
|
|
# Satisfiable Mul in `EQ` with negative sign
|
|
res8 = tvm.arith.deduce_bound(a, (5 * a == -10), {}, {})
|
|
tvm.testing.assert_prim_expr_equal(res8.max_value, -2)
|
|
tvm.testing.assert_prim_expr_equal(res8.min_value, -2)
|
|
|
|
# Unsatisfiable Mul in `EQ`
|
|
e5 = 4 * a == b
|
|
res9 = tvm.arith.deduce_bound(a, e5, {b: b_s}, {})
|
|
assert str(res9.max_value) == "neg_inf"
|
|
assert str(res9.min_value) == "pos_inf"
|
|
|
|
res10 = tvm.arith.deduce_bound(a, (b * a == b), {b: b_s}, {})
|
|
# simplifier is now able to prove symbolic relation (b * a % b == 0)
|
|
tvm.testing.assert_prim_expr_equal(res10.max_value, 1)
|
|
tvm.testing.assert_prim_expr_equal(res10.min_value, 1)
|
|
|
|
|
|
def test_check():
|
|
a = tvm.tirx.Var("a", "int32")
|
|
b = tvm.tirx.Var("b", "int32")
|
|
c = tvm.tirx.Var("c", "int32")
|
|
d = tvm.tirx.Var("d", "int32")
|
|
|
|
b_s = tvm.arith.IntervalSet(2, 3)
|
|
c_s = tvm.arith.IntervalSet(5, 7)
|
|
d_s = tvm.arith.IntervalSet(-3, -1)
|
|
|
|
# no compare operator
|
|
res1 = tvm.arith.deduce_bound(a, a + b, {b: b_s}, {})
|
|
assert res1.is_nothing()
|
|
|
|
# multiple compare operators
|
|
res2 = tvm.arith.deduce_bound(a, (a + b > 3).astype(c.ty) > c, {b: b_s, c: c_s}, {})
|
|
assert res2.is_nothing()
|
|
|
|
# multiple target variable
|
|
res2 = tvm.arith.deduce_bound(a, a * 2 - a > b, {b: b_s}, {})
|
|
assert res2.is_nothing()
|
|
|
|
|
|
def test_deduce_basic():
|
|
def test_basic(a1, a2, coff):
|
|
a = tvm.tirx.Var("a", "int32")
|
|
b = tvm.tirx.Var("b", "int32")
|
|
b_s = tvm.arith.IntervalSet(a1, a2)
|
|
e0 = b + a * coff + 3
|
|
|
|
res1 = tvm.arith.deduce_bound(a, e0 < 17, {b: b_s}, {b: b_s})
|
|
[x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) < 17, True)
|
|
|
|
# expression containing variable a is on rhs
|
|
res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") < e0, {b: b_s}, {b: b_s})
|
|
[x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) > 17, True)
|
|
|
|
# expression containing variable a is on rhs
|
|
res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(17, "int32") >= e0, {b: b_s}, {b: b_s})
|
|
[x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
|
|
|
|
tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) <= 17, True)
|
|
|
|
res1 = tvm.arith.deduce_bound(a, e0 >= 17, {b: b_s}, {b: b_s})
|
|
[x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) >= 17, True)
|
|
|
|
test_basic(0, 4, 4)
|
|
test_basic(1, 5, 4)
|
|
test_basic(2, 6, 4)
|
|
test_basic(0, 4, -4)
|
|
test_basic(1, 5, -4)
|
|
test_basic(2, 6, -4)
|
|
|
|
|
|
def test_deduce_complex():
|
|
def test_complex(a1, a2, coff):
|
|
a = tvm.tirx.Var("a", "int32")
|
|
b = tvm.tirx.Var("b", "int32")
|
|
b_s = tvm.arith.IntervalSet(a1, a2)
|
|
e0 = (b * 3 + a * coff) * 4
|
|
|
|
res1 = tvm.arith.deduce_bound(a, e0 < 63, {b: b_s}, {b: b_s})
|
|
[t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) < 63, True)
|
|
|
|
# expression containing variable a is on rhs
|
|
res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") >= e0, {b: b_s}, {b: b_s})
|
|
[t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) <= 63, True)
|
|
|
|
res1 = tvm.arith.deduce_bound(a, e0 > 63, {b: b_s}, {b: b_s})
|
|
[t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) > 63, True)
|
|
|
|
# expression containing variable a is on rhs
|
|
res1 = tvm.arith.deduce_bound(a, tvm.tirx.const(63, "int32") <= e0, {b: b_s}, {b: b_s})
|
|
[t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
|
|
tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) >= 63, True)
|
|
|
|
test_complex(0, 4, 4)
|
|
test_complex(0, 4, -4)
|
|
test_complex(2, 6, 4)
|
|
test_complex(0, 4, -4)
|
|
test_complex(1, 5, -4)
|
|
test_complex(2, 6, -4)
|
|
|
|
|
|
def test_deduce_non_support():
|
|
a = tvm.tirx.Var("a", "int32")
|
|
|
|
def test_non_support(lhs):
|
|
res = tvm.arith.deduce_bound(a, lhs < 10, {}, {})
|
|
assert res.is_nothing()
|
|
|
|
test_non_support(tvm.tirx.floormod(a, 16))
|
|
test_non_support(tvm.tirx.Min(a, 16))
|
|
test_non_support(tvm.tirx.Max(a, 16))
|
|
test_non_support(tvm.tirx.LE(a, 16))
|
|
test_non_support(tvm.tirx.LT(a, 16))
|
|
test_non_support(tvm.tirx.GE(a, 16))
|
|
test_non_support(tvm.tirx.GT(a, 16))
|
|
test_non_support(tvm.tirx.EQ(a, 16))
|
|
test_non_support(tvm.tirx.NE(a, 16))
|
|
test_non_support(tvm.tirx.log(a))
|
|
test_non_support(tvm.tirx.BufferLoad(decl_buffer([16], "int32"), [a]))
|
|
|
|
|
|
def test_deduce_floordiv():
|
|
def do_test(gen_expr, dom_map, expect_min, expect_max):
|
|
a = tvm.tirx.Var("a", "int32")
|
|
expr = gen_expr(a)
|
|
res = tvm.arith.deduce_bound(a, expr, dom_map, dom_map)
|
|
if isinstance(expect_min, str):
|
|
assert str(res.min_value) == expect_min
|
|
else:
|
|
tvm.testing.assert_prim_expr_equal(res.min_value, expect_min)
|
|
if isinstance(expect_max, str):
|
|
assert str(res.max_value) == expect_max
|
|
else:
|
|
tvm.testing.assert_prim_expr_equal(res.max_value, expect_max)
|
|
|
|
# test basic cases
|
|
do_test(lambda a: a // 8 > 3, {}, 32, "pos_inf")
|
|
do_test(lambda a: a // 8 >= 3, {}, 24, "pos_inf")
|
|
do_test(lambda a: a // 8 < 3, {}, "neg_inf", 23)
|
|
do_test(lambda a: a // 8 <= 3, {}, "neg_inf", 31)
|
|
do_test(lambda a: a // 8 == 3, {}, "pos_inf", "neg_inf")
|
|
do_test(lambda a: a // 8 > -3, {}, -16, "pos_inf")
|
|
do_test(lambda a: a // 8 >= -3, {}, -24, "pos_inf")
|
|
do_test(lambda a: a // -8 > 3, {}, "neg_inf", -32)
|
|
do_test(lambda a: a // -8 >= 3, {}, "neg_inf", -24)
|
|
do_test(lambda a: a // -8 < 3, {}, -23, "pos_inf")
|
|
do_test(lambda a: a // -8 <= 3, {}, -31, "pos_inf")
|
|
do_test(lambda a: 8 // a >= 2, {}, "pos_inf", "neg_inf")
|
|
|
|
# test nested cases
|
|
b = tvm.tirx.Var("b", "int32")
|
|
bs = {b: tvm.arith.IntervalSet(2, 6)}
|
|
do_test(lambda a: b * 3 + a // 8 < 63, bs, "neg_inf", 359)
|
|
do_test(lambda a: b * 3 + a // 8 <= 63, bs, "neg_inf", 367)
|
|
do_test(lambda a: b * 3 + a // 8 > 63, bs, 464, "pos_inf")
|
|
do_test(lambda a: b * 3 + a // 8 >= 63, bs, 456, "pos_inf")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|