From 764ce21cd94c56aa6c7e193e0d7ed9d1d3926443 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 16 Feb 2026 09:32:20 -0600 Subject: [PATCH] Use FlattenMapper before codegen - Smaller expressions to crunch through - Throws off flop counts --- sumpy/codegen.py | 3 +++ sumpy/version.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/sumpy/codegen.py b/sumpy/codegen.py index f2c7b70f4..25e829f25 100644 --- a/sumpy/codegen.py +++ b/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) diff --git a/sumpy/version.py b/sumpy/version.py index 8c4f151f9..34f17e716 100644 --- a/sumpy/version.py +++ b/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)