def _check(variables, formulas, coef=(-5, 5), bounds=(-20, 20)):
vs = [te.var("x" + str(i)) for i in range(variables)]
fs = []
for i in range(formulas):
s1 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
s1 += random.randint(coef[0], coef[1])
s2 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
s2 += random.randint(coef[0], coef[1])
op = random.choice([
tir.expr.EQ, tir.expr.LE, tir.expr.LT, tir.expr.GE, tir.expr.GT
])
fs.append(op(s1, s2))
vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in vs}
before = te.all(tir.const(1, 'bool'), *fs)
after = arith._ffi_api.SolveInequalitiesAsCondition(vs, vranges, fs)
after = te.all(tir.const(1, 'bool'), *after)
testing.check_bool_expr_is_true(before == after, vranges)
solution = arith.solve_linear_inequalities(fs,
vs,
vranges,
deskew_range=True)
testing.check_int_constraints_trans_consistency(solution)
def _get_pixel_value(n, c, h, w):
if padding_mode == "zeros":
return te.if_then_else(
te.all(h >= 0, w >= 0, h < in_height, w < in_width),
data[n, c, h, w],
tir.const(0.0, dtype=data.dtype),
)
if padding_mode == "border":
h_b = te.max(te.min(h, in_height - 1), 0)
w_b = te.max(te.min(w, in_width - 1), 0)
return data[n, c, h_b, w_b]
raise AssertionError("unsupported padding_mode")
def _dilate(*indices):
not_zero = []
index_tuple = []
for i in range(n):
if not strides[i] == 1:
index_tuple.append(idx_div(indices[i], strides[i]))
not_zero.append(idx_mod(indices[i], strides[i]).equal(0))
else:
index_tuple.append(indices[i])
if not_zero:
not_zero = te.all(*not_zero)
return te.if_then_else(not_zero, padded(*index_tuple),
tir.const(0.0, padded.dtype))
return padded(*index_tuple)
def _check_forward(constraints1, constraints2, varmap, backvarmap):
ana = tvm.arith.Analyzer()
all_vranges = vranges.copy()
all_vranges.update({v: r for v, r in constraints1.ranges.items()})
# Check that the transformation is injective
cond_on_vars = tir.const(1, 'bool')
for v in constraints1.variables:
# variable mapping is consistent
v_back = ana.simplify(tir.stmt_functor.substitute(varmap[v], backvarmap))
cond_on_vars = te.all(cond_on_vars, v == v_back)
# Also we have to check that the new relations are true when old relations are true
cond_subst = tir.stmt_functor.substitute(
te.all(tir.const(1, 'bool'), *constraints2.relations), backvarmap)
# We have to include relations from vranges too
for v in constraints2.variables:
if v in constraints2.ranges:
r = constraints2.ranges[v]
range_cond = te.all(v >= r.min, v < r.min + r.extent)
range_cond = tir.stmt_functor.substitute(range_cond, backvarmap)
cond_subst = te.all(cond_subst, range_cond)
cond_subst = ana.simplify(cond_subst)
check_bruteforce(te.all(cond_subst, cond_on_vars), all_vranges,
cond=te.all(tir.const(1, 'bool'), *constraints1.relations))
def _get_pixel_value(n, c, h, w):
return te.if_then_else(
te.all(h >= 0, w >= 0, h < in_height, w < in_width),
data[n, c, h, w],
tir.const(0.0, dtype=data.dtype),
)
import tvm
from tvm import te
n = 1024
k = 3
pad = 2
A = te.placeholder((n, n), name='A')
W = te.placeholder((k, k), name='W')
m = (n - k + 2 * pad) + 1
Apad = te.compute((n + 2 * pad, n + 2 * pad),
lambda yy, xx: te.if_then_else(
te.all(yy >= pad, yy < pad + n, xx >= pad, xx < pad + n),
A[yy - pad, xx - pad], tvm.tir.const(0., "float32")),
name='Apad')
ry = te.reduce_axis((0, k), name='ry')
rx = te.reduce_axis((0, k), name='rx')
B = te.compute(
(m, m),
lambda yy, xx: te.sum(Apad[yy + ry, xx + rx] * W[ry, rx], axis=[ry, rx]),
name='B')
s = te.create_schedule(B.op)
print(tvm.lower(s, [A, W, B], simple_mode=True))
print("---------cutting line---------")
s[Apad].compute_inline()
print(tvm.lower(s, [A, W, B], simple_mode=True))
exit(0)