check for non-affineness instead of subscripts

Author: majosm

pytato/target/loopy/codegen.py

@@ -37,6 +37,7 @@
 from typing_extensions import Never, override
 
 import loopy as lp
+import loopy.diagnostic as lp_diagnostic
 import loopy.symbolic as lp_symbolic
 import pymbolic.primitives as prim
 from pymbolic import ArithmeticExpression, var
@@ -505,28 +506,64 @@ def map_index_lambda(self, expr: IndexLambda,
 
         loopy_shape = shape_to_scalar_expression(expr.shape, self, state)
 
-        # If the scalar expression contains any reductions with bounds expressions
-        # that index into a binding, need to store the results of those expressions
-        # as scalar temporaries
-        subscript_detector = SubscriptDetector()
-
         redn_bounds = {
             var_name: redn.bounds[var_name]
             for var_name, redn in var_to_reduction.items()}
 
-        # FIXME: Forcing storage of expressions containing processed reductions for
-        # now; attempting to generalize to unmaterialized expressions would require
-        # handling at least two complications:
-        #   1) final inames aren't assigned until the expression is stored, so any
-        #      temporary variables defined below would need to be finalized at the
-        #      point of storage, not here
+        loopy_redn_bounds: Mapping[str, tuple[Expression, Expression]] = {}
+        for var_name, bounds in redn_bounds.items():
+            loopy_bound_list: list[Expression] = []
+            for bound in bounds:
+                loopy_bound_list.append(
+                    self.exprgen_mapper(bound, prstnt_ctx, local_ctx))
+            loopy_redn_bounds[var_name] = cast(
+                "tuple[Expression, Expression]", tuple(loopy_bound_list))
+
+        store_result = expr.tags_of_type(ImplStored)
+
+        # FIXME: Did I do this right (the space creation below)? I end up with
+        # a space that looks like Space("[_pt_data_1] -> { [_0, _r0] }"), where
+        # _pt_data_1 is the array being indexed into in the reduction bound
+
+        set_names = set(
+            tuple(f"_{idim}" for idim in range(expr.ndim))
+            + tuple(loopy_redn_bounds.keys()))
+
+        param_names: set[str] = set()
+        for sdep in map(scalar_expr.get_dependencies, loopy_shape):
+            param_names |= sdep
+
+        for bounds in loopy_redn_bounds.values():
+            for sdep in map(scalar_expr.get_dependencies, bounds):
+                param_names |= sdep
+
+        param_names -= set_names
+
+        space = isl.Space.create_from_names(
+            isl.DEFAULT_CONTEXT,
+            set=sorted(set_names),
+            params=sorted(param_names))
+
+        # If the scalar expression contains any reductions with bounds
+        # expressions that are non-affine, we need to store the results of those
+        # expressions as scalar temporaries.
+        # FIXME: For now, forcing storage of expressions containing such
+        # processed reductions; attempting to generalize to unmaterialized
+        # expressions would require handling at least two complications:
+        #   1) final inames aren't assigned until the expression is stored, so
+        #      any temporary variables defined below would need to be finalized
+        #      at the point of storage, not here
         #   2) lp.make_reduction_inames_unique does not rename the temporaries
-        #      created below, so something would need to be done to make them unique
-        #      across all index lambda evaluations.
-        store_result = expr.tags_of_type(ImplStored) or any(
-            subscript_detector(bound)
-            for bounds in redn_bounds.values()
-            for bound in bounds)
+        #      created below, so something would need to be done to make them
+        #      unique across all index lambda evaluations.
+
+        for bounds in loopy_redn_bounds.values():
+            for bound in bounds:
+                try:
+                    lp_symbolic.guarded_pwaff_from_expr(space, bound)
+                except lp_diagnostic.ExpressionToAffineConversionError:
+                    store_result = True
+                    break
 
         name: str | None = None
         inames: tuple[str, ...] | None = None
@@ -550,13 +587,15 @@ def map_index_lambda(self, expr: IndexLambda,
                 str, tuple[ArithmeticExpression, ArithmeticExpression]] = {}
             bound_prefixes = ("l", "u")
             for var_name, bounds in redn_bounds.items():
+                loopy_bounds = loopy_redn_bounds[var_name]
                 new_bounds_list: list[ArithmeticExpression] = []
-                for bound_prefix, bound in zip(bound_prefixes, bounds, strict=True):
-                    if subscript_detector(bound):
+                for bound_prefix, bound, loopy_bound in zip(
+                        bound_prefixes, bounds, loopy_bounds, strict=True):
+                    try:
+                        lp_symbolic.guarded_pwaff_from_expr(space, loopy_bound)
+                    except lp_diagnostic.ExpressionToAffineConversionError:
                         unique_name = var_to_reduction_unique_name[var_name]
                         bound_name = f"{unique_name}_{bound_prefix}bound"
-                        loopy_bound = self.exprgen_mapper(
-                            bound, prstnt_ctx, local_ctx)
                         bound_result: ImplementedResult = InlinedResult(
                             loopy_bound, expr.ndim, prstnt_ctx.depends_on)
                         bound_result = StoredResult(
@@ -793,25 +832,6 @@ def map_call(self, expr: Call, state: CodeGenState) -> None:
 }
 
 
-class SubscriptDetector(scalar_expr.CombineMapper[bool, []]):
-    """Returns *True* if a scalar expression contains any subscripts."""
-    @override
-    def combine(self, values: Iterable[bool]) -> bool:
-        return any(values)
-
-    @override
-    def map_algebraic_leaf(self, expr: prim.AlgebraicLeaf) -> bool:
-        return False
-
-    @override
-    def map_subscript(self, expr: prim.Subscript) -> bool:
-        return True
-
-    @override
-    def map_constant(self, expr: object) -> bool:
-        return False
-
-
 class ReductionCollector(scalar_expr.CombineMapper[frozenset[scalar_expr.Reduce], []]):
     """
     Constructs a :class:`frozenset` containing all instances of