wehub-resource-sync
187 lines
6.3 KiB
Python
187 lines
6.3 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
|
|
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()
|