diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 1718e8b8f61b4bd3808ca27fae09863c6dbbea17..eb0de5869c3dc80efd69ca1e4674b821eae7ecbe 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -113,14 +113,9 @@ class CodeGenerationState(object): self.c_code_mapper) def fix(self, iname, aff, space): - dt, pos = space.get_var_dict()[iname] - assert dt == isl.dim_type.set - - zero = isl.Aff.zero_on_domain(space) - - from islpy import align_spaces - iname_plus_lb_aff = align_spaces(aff, zero).add_coefficient( - isl.dim_type.in_, pos, -1) + from loopy.isl_helpers import iname_rel_aff + iname_plus_lb_aff = iname_rel_aff( + space, iname, "==", aff) from loopy.symbolic import pw_aff_to_expr cns = isl.Constraint.equality_from_aff(iname_plus_lb_aff) diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py index 055b65655a1feeabecd4e7001e94b4b8ec98f124..d1d42ffcfcdbd03e86389c39c2313c8684ff61d8 100644 --- a/loopy/isl_helpers.py +++ b/loopy/isl_helpers.py @@ -127,6 +127,29 @@ def make_slab(space, iname, start, stop): +def iname_rel_aff(space, iname, rel, aff): + """*aff*'s domain space is allowed to not match *space*.""" + + dt, pos = space.get_var_dict()[iname] + assert dt == isl.dim_type.set + + from islpy import align_spaces + aff = align_spaces(aff, isl.Aff.zero_on_domain(space)) + + if rel in ["==", "<="]: + return aff.add_coefficient(isl.dim_type.in_, pos, -1) + elif rel == ">=": + return aff.neg().add_coefficient(isl.dim_type.in_, pos, 1) + elif rel == "<": + return (aff-1).add_coefficient(isl.dim_type.in_, pos, -1) + elif rel == ">": + return (aff+1).neg().add_coefficient(isl.dim_type.in_, pos, 1) + else: + raise ValueError("unknown value of 'rel': %s" % rel) + + + + def static_extremum_of_pw_aff(pw_aff, constants_only, set_method, what): pieces = pw_aff.get_pieces() if len(pieces) == 1: