Use FlattenMapper before codegen

- Smaller expressions to crunch through
- Throws off flop counts

Author: inducer

sumpy/codegen.py

@@ -36,6 +36,7 @@
 import pymbolic.primitives as prim
 from loopy.kernel.instruction import Assignment, CallInstruction, make_assignment
 from pymbolic.mapper import CSECachingMapperMixin, IdentityMapper, P
+from pymbolic.mapper.flattener import FlattenMapper
 from pymbolic.typing import ArithmeticExpression, Expression
 from pytools import memoize_method
 
@@ -753,6 +754,7 @@ def to_loopy_insns(
     sympy_conv = SympyToPymbolicMapper()
     pymbolic_assignments = [(name, sympy_conv(expr)) for name, expr in assignments]
 
+    flat = FlattenMapper()
     bdr = BesselDerivativeReplacer()
     btog = BesselTopOrderGatherer()
     vcr = VectorComponentRewriter(vector_names)
@@ -762,6 +764,7 @@ def to_loopy_insns(
     cmr = ComplexRewriter(complex_dtype)
 
     def cmb_mapper(expr: Expression, /) -> Expression:
+        expr = flat(expr)
         expr = bdr(expr)
         expr = vcr(expr)
         expr = pwr(expr)

sumpy/version.py

@@ -41,4 +41,4 @@ def _parse_version(version: str) -> tuple[tuple[int, ...], str]:
 VERSION, VERSION_STATUS = _parse_version(VERSION_TEXT)
 
 _GIT_REVISION = find_module_git_revision(__file__, n_levels_up=1)
-KERNEL_VERSION = (*VERSION, _GIT_REVISION, 1)
+KERNEL_VERSION = (*VERSION, _GIT_REVISION, 2)