Lower spatial constants for Loopy

sumpy/codegen.py

@@ -35,6 +35,7 @@
 import loopy as lp
 import pymbolic.primitives as prim
 from loopy.kernel.instruction import Assignment, CallInstruction, make_assignment
+from loopy.symbolic import IdentityMapper as LoopyIdentityMapper
 from pymbolic.mapper import CSECachingMapperMixin, IdentityMapper, P
 from pymbolic.mapper.flattener import FlattenMapper
 from pymbolic.typing import ArithmeticExpression, Expression
@@ -713,6 +714,13 @@ def map_common_subexpression_uncached(
 
 # {{{ to-loopy conversion
 
+class SpatialConstantRewriter(LoopyIdentityMapper[[]]):
+    """Convert Sumpy's spatial constants to variables at the Loopy boundary."""
+
+    def map_spatial_constant(self, expr: sym.SpatialConstant, /) -> prim.Variable:
+        return prim.Variable(expr.name)
+
+
 def to_loopy_insns(
         assignments: Iterable[tuple[str, sym.Basic]],
         vector_names: AbstractSet[str] | None = None,
@@ -785,6 +793,12 @@ def convert_expr(name: str, expr: Expression) -> Expression:
 
         result.extend(bessel_sub.assignments)
 
+        # Keep spatial constants through symbolic/code transformations, then
+        # lower them in every outgoing instruction, including Bessel helpers.
+        spatial_rewriter = SpatialConstantRewriter()
+        result = [insn.with_transformed_expressions(spatial_rewriter)
+                  for insn in result]
+
     logger.info("loopy instruction generation: done")
     return result
 

sumpy/e2p.py

@@ -107,12 +107,15 @@ def get_loopy_args(self):
         return gather_loopy_arguments((self.expansion, *tuple(self.kernels)))
 
     def get_kernel_scaling_assignment(self):
+        from sumpy.codegen import SpatialConstantRewriter
         from sumpy.symbolic import to_pymbolic
+
+        spatial_rewriter = SpatialConstantRewriter()
         return [lp.Assignment(id="kernel_scaling",
                     assignee="kernel_scaling",
-                    expression=to_pymbolic(
+                    expression=spatial_rewriter(to_pymbolic(
                         self.expansion.kernel.get_global_scaling_const(),
-                    ),
+                    )),
                     temp_var_type=lp.Optional(None),
                     )]
 # }}}

sumpy/test/test_codegen.py

@@ -53,6 +53,27 @@
 # }}}
 
 
+# {{{ test_spatial_constants_at_loopy_boundary
+
+def test_to_loopy_insns_rewrites_spatial_constants():
+    import pymbolic.primitives as prim
+    from pymbolic.mapper.dependency import DependencyMapper
+
+    import sumpy.symbolic as sym
+    from sumpy.codegen import to_loopy_insns
+
+    k = sym.SpatialConstant("k").as_sympy()
+    instructions = to_loopy_insns([("result", sym.BesselJ(2, k, 0))])
+    dependencies = set()
+    mapper = DependencyMapper(include_calls=False, include_subscripts=False)
+    for instruction in instructions:
+        dependencies.update(mapper(instruction.expression))
+    assert prim.Variable("k") in dependencies
+    assert all(type(dep) is prim.Variable for dep in dependencies)
+
+# }}}
+
+
 # {{{ test_line_taylor_coeff_growth
 
 def test_line_taylor_coeff_growth():

sumpy/tools.py

@@ -368,12 +368,16 @@ def default_name(self) -> str:
 
     def get_kernel_scaling_assignments(self) -> Sequence[lp.Assignment]:
         import loopy as lp
+
+        from sumpy.codegen import SpatialConstantRewriter
+
+        spatial_rewriter = SpatialConstantRewriter()
         return [
                 lp.Assignment(id=f"knl_{i}_scaling",
                     assignee=f"knl_{i}_scaling",
-                    expression=sym.to_pymbolic(
+                    expression=spatial_rewriter(sym.to_pymbolic(
                         kernel.get_global_scaling_const()
-                    ),
+                    )),
                     temp_var_type=lp.Optional(dtype),
                     tags=frozenset([ScalingAssignmentTag()]))
                 for i, (kernel, dtype) in enumerate(