diff --git a/loopy/codegen/loop.py b/loopy/codegen/loop.py index 81686878d55ddebc49f2122d473eb8fcf8d2ec32..ad80475c1d27f67b3df8a885f60dd96ff28efe6a 100644 --- a/loopy/codegen/loop.py +++ b/loopy/codegen/loop.py @@ -410,12 +410,17 @@ def generate_sequential_loop_dim_code(codegen_state, sched_index): # {{{ find implemented loop, build inner code - from loopy.isl_helpers import make_loop_bounds_from_pwaffs + from loopy.symbolic import pw_aff_to_pw_aff_implemented_by_expr + impl_lbound = pw_aff_to_pw_aff_implemented_by_expr(lbound) + impl_ubound = pw_aff_to_pw_aff_implemented_by_expr(ubound) # impl_loop may be overapproximated + from loopy.isl_helpers import make_loop_bounds_from_pwaffs impl_loop = make_loop_bounds_from_pwaffs( dom_and_slab.space, - loop_iname, lbound, ubound) + loop_iname, + impl_lbound, + impl_ubound) for iname in moved_inames: dt, idx = impl_loop.get_var_dict()[iname] diff --git a/loopy/symbolic.py b/loopy/symbolic.py index 430c651589939a1001432bd8db413cb5902b14a6..52fd6e57f92e7f9599a3a0fb4256f97347708303 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -1150,6 +1150,23 @@ def pw_aff_to_expr(pw_aff, int_ok=False): return expr + +def pw_aff_to_pw_aff_implemented_by_expr(pw_aff): + pieces = pw_aff.get_pieces() + + rest = isl.Set.universe(pw_aff.space.params()) + aff_set, aff = pieces[0] + impl_pw_aff = isl.PwAff.alloc(aff_set, aff) + rest = rest.intersect_params(aff_set.complement()) + + for aff_set, aff in pieces[1:-1]: + impl_pw_aff = impl_pw_aff.union_max( + isl.PwAff.alloc(aff_set, aff)) + rest = rest.intersect_params(aff_set.complement()) + + _, aff = pieces[-1] + return impl_pw_aff.union_max(isl.PwAff.alloc(rest, aff)).coalesce() + # }}}