def test_subst_into_pwaff():
from pymbolic.primitives import Variable
arg_dict = {"m": 3 * Variable("nx"), "n": 2 * Variable("ny") + 4}
space = isl.Set("[nx, ny, nz] -> { []: }").params().space
poly = isl.PwAff("[m, n] -> { [3 * m + 2 * n] : "
"m > 0 and n > 0; [7* m + 4*n] : m > 0 and n <= 0 }")
from loopy.isl_helpers import subst_into_pwaff
result = subst_into_pwaff(space, poly, arg_dict)
expected = isl.PwAff(
"[nx, ny, nz] -> { [(9nx + 4ny+8)] : nx > 0 and ny > -2;"
" [(21nx + 8ny+16)] : nx > 0 and ny <= -2 }")
assert result == expected
def test_eval_pw_qpolynomial():
pwaff = isl.PwAff("[n] -> { [(0)] : n <= 4 and n >= 1; "
"[(-1 + n - floor((3n)/4))] : n >= 5 }")
pwq = isl.PwQPolynomial.from_pw_aff(pwaff)
pwq.eval_with_dict(dict(n=10))
def test_pickling():
instances = [
isl.Aff("[n] -> { [(-1 - floor((-n)/4))] }"),
isl.PwAff("[n] -> { [(0)] : n <= 4 and n >= 1; "
"[(-1 + n - floor((3n)/4))] : n >= 5 }"),
isl.BasicSet("[n] -> {[i,j,k]: i<=j + k and (exists m: m=j+k) "
"and n mod 5 = 17}"),
isl.Set("[n] -> {[i,j,k]: (i<=j + k and (exists m: m=j+k)) or (k=j)}")
]
from pickle import dumps, loads
for inst in instances:
inst2 = loads(dumps(inst))
assert inst.space == inst2.space
assert inst.is_equal(inst2)
def test_pw_aff_to_conditional_expr():
from loopy.symbolic import pw_aff_to_expr
cond = isl.PwAff("[i] -> { [(0)] : i = 0; [(-1 + i)] : i > 0 }")
expr = pw_aff_to_expr(cond)
assert str(expr) == "0 if i == 0 else -1 + i"
def test_upcast():
a = isl.PwAff("[n] -> { [(-1 - floor((-n)/4))] }")
b = isl.Aff("[n] -> { [(-1 - floor((-n)/4))] }")
assert b.is_equal(a)
assert a.is_equal(b)
