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)