chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,186 @@
|
||||
# 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
|
||||
import random
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm import arith, ir, testing, tirx
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def test_solution_consistency():
|
||||
seed = random.randrange(sys.maxsize)
|
||||
print(
|
||||
"\nThis test is intentionally non-deterministic, "
|
||||
f"if it fails please report it in GitHub issue together with this seed {seed}\n"
|
||||
)
|
||||
random.seed(seed)
|
||||
|
||||
def _check(num_vars, num_formulas, coef=(-5, 5), bounds=(-20, 20)):
|
||||
variables = [tvm.tirx.Var("x" + str(i), "int32") for i in range(num_vars)]
|
||||
|
||||
relations = []
|
||||
for i in range(num_formulas):
|
||||
s1 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
|
||||
s1 += random.randint(coef[0], coef[1])
|
||||
s2 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
|
||||
s2 += random.randint(coef[0], coef[1])
|
||||
if random.random() < 0.7:
|
||||
op = tvm.tirx.EQ
|
||||
else:
|
||||
# we also make sure it can correctly handle inequalities
|
||||
op = random.choice([tvm.tirx.LE, tvm.tirx.LT, tvm.tirx.GE, tvm.tirx.GT])
|
||||
relations.append(op(s1, s2))
|
||||
|
||||
vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in variables}
|
||||
solution = arith.solve_linear_equations(relations, variables, vranges)
|
||||
|
||||
testing.check_int_constraints_trans_consistency(solution)
|
||||
|
||||
# leaving some variables as parameters should also be ok
|
||||
for k in [1, 2]:
|
||||
if len(variables) > k:
|
||||
solution = arith.solve_linear_equations(relations, variables[:-k], vranges)
|
||||
param_ranges = {v: vranges[v] for v in variables[-k:]}
|
||||
testing.check_int_constraints_trans_consistency(solution, param_ranges)
|
||||
|
||||
for i in range(2):
|
||||
_check(num_vars=1, num_formulas=1)
|
||||
for i in range(2):
|
||||
_check(num_vars=1, num_formulas=2)
|
||||
|
||||
for i in range(2):
|
||||
_check(num_vars=2, num_formulas=1)
|
||||
for i in range(2):
|
||||
_check(num_vars=2, num_formulas=2)
|
||||
for i in range(2):
|
||||
_check(num_vars=2, num_formulas=3)
|
||||
|
||||
for i in range(3):
|
||||
_check(num_vars=3, num_formulas=3, coef=(-2, 2))
|
||||
for i in range(3):
|
||||
_check(num_vars=3, num_formulas=4, coef=(-2, 2))
|
||||
|
||||
for i in range(3):
|
||||
_check(num_vars=4, num_formulas=3, coef=(-1, 1))
|
||||
|
||||
for i in range(3):
|
||||
_check(num_vars=10, num_formulas=2, coef=(-1, 1), bounds=(0, 4))
|
||||
for i in range(3):
|
||||
_check(num_vars=10, num_formulas=3, coef=(0, 1), bounds=(0, 4))
|
||||
|
||||
|
||||
def test_empty_var_to_solve():
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
equations = [
|
||||
tvm.tirx.EQ(x + y, 20),
|
||||
tvm.tirx.EQ(x - y, 10),
|
||||
]
|
||||
solution = arith.solve_linear_equations(equations)
|
||||
assert len(solution.src_to_dst) == 0
|
||||
assert len(solution.dst_to_src) == 0
|
||||
assert len(solution.src.variables) == 0
|
||||
assert len(solution.src.ranges) == 0
|
||||
assert tvm_ffi.structural_equal(solution.src.relations, equations)
|
||||
assert tvm_ffi.structural_equal(solution.src, solution.dst)
|
||||
|
||||
|
||||
def test_unique_solution():
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
|
||||
solution = arith.solve_linear_equations(
|
||||
[
|
||||
tvm.tirx.EQ(x + y, 20),
|
||||
tvm.tirx.EQ(x - y, 10),
|
||||
],
|
||||
[x, y],
|
||||
)
|
||||
assert list(solution.dst.variables) == []
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[x], T.int32(15))
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[y], T.int32(5))
|
||||
|
||||
|
||||
def test_low_rank():
|
||||
x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32")
|
||||
ranges = {}
|
||||
|
||||
solution = arith.solve_linear_equations(
|
||||
[
|
||||
tvm.tirx.EQ(x + y + z, 15),
|
||||
tvm.tirx.EQ(x + y, 10),
|
||||
],
|
||||
[x, y, z],
|
||||
ranges,
|
||||
)
|
||||
[n0] = solution.dst.variables
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0 + 10)
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[z], T.int32(5))
|
||||
|
||||
|
||||
def test_infer_range():
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
ranges = {
|
||||
x: tvm.ir.Range.from_min_extent(-5, 10),
|
||||
y: tvm.ir.Range.from_min_extent(0, 10),
|
||||
}
|
||||
|
||||
solution = arith.solve_linear_equations(
|
||||
[
|
||||
tvm.tirx.EQ(x + y, 0),
|
||||
],
|
||||
[x, y],
|
||||
ranges,
|
||||
)
|
||||
[n0] = solution.dst.variables
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0)
|
||||
assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
|
||||
# inferred from y's range
|
||||
assert tvm_ffi.structural_equal(solution.dst.ranges[n0].min, T.int32(-9))
|
||||
assert tvm_ffi.structural_equal(solution.dst.ranges[n0].extent, T.int32(10))
|
||||
# additional inequality is added into the system for x
|
||||
[ineq] = solution.dst.relations
|
||||
assert isinstance(ineq, tvm.tirx.LE)
|
||||
assert tvm_ffi.structural_equal(ineq.a, T.int32(-5))
|
||||
assert tvm_ffi.structural_equal(ineq.b, n0)
|
||||
|
||||
|
||||
def test_ill_formed():
|
||||
x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
|
||||
|
||||
solution = arith.solve_linear_equations(
|
||||
[
|
||||
tvm.tirx.EQ(x + y, 0),
|
||||
tvm.tirx.EQ(x - y, 0),
|
||||
tvm.tirx.EQ(x, 5),
|
||||
],
|
||||
[x, y],
|
||||
{},
|
||||
)
|
||||
assert list(solution.dst.variables) == []
|
||||
[rel] = solution.dst.relations
|
||||
ir.assert_structural_equal(rel, tirx.const(False))
|
||||
assert len(solution.src_to_dst) == 0
|
||||
assert len(solution.dst_to_src) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user