diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2c6c77e77..234988280 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,6 +7,10 @@ on: schedule: - cron: '17 3 * * 0' +concurrency: + group: ${{ github.head_ref || github.ref_name }} + cancel-in-progress: true + jobs: ruff: name: Ruff diff --git a/loopy/codegen/loop.py b/loopy/codegen/loop.py index 645a57e31..c64c2ea67 100644 --- a/loopy/codegen/loop.py +++ b/loopy/codegen/loop.py @@ -28,6 +28,7 @@ from loopy.codegen.control import build_loop_nest from loopy.codegen.result import merge_codegen_results from loopy.diagnostic import LoopyError, warn +from loopy.symbolic import flatten # {{{ conditional-reducing slab decomposition @@ -309,7 +310,7 @@ def set_up_hw_parallel_loops(codegen_state, schedule_index, next_func, codegen_state = codegen_state.intersect(slab) from loopy.symbolic import pw_aff_to_expr - hw_axis_expr = hw_axis_expr + pw_aff_to_expr(lower_bound) + hw_axis_expr = flatten(hw_axis_expr + pw_aff_to_expr(lower_bound)) # }}} diff --git a/loopy/expression.py b/loopy/expression.py index 224521b03..2581ec022 100644 --- a/loopy/expression.py +++ b/loopy/expression.py @@ -23,7 +23,7 @@ import numpy as np -from pymbolic.mapper import RecursiveMapper +from pymbolic.mapper import Mapper from loopy.codegen import UnvectorizableError from loopy.diagnostic import LoopyError @@ -55,7 +55,7 @@ def dtype_to_type_context(target, dtype): # {{{ vectorizability checker -class VectorizabilityChecker(RecursiveMapper): +class VectorizabilityChecker(Mapper): """The return value from this mapper is a :class:`bool` indicating whether the result of the expression is vectorized along :attr:`vec_iname`. If the expression is not vectorizable, the mapper raises diff --git a/loopy/symbolic.py b/loopy/symbolic.py index 22dbd3bf5..6727423a8 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -54,6 +54,7 @@ ) from pymbolic.mapper.dependency import CachedDependencyMapper as DependencyMapperBase from pymbolic.mapper.evaluator import CachedEvaluationMapper as EvaluationMapperBase +from pymbolic.mapper.flattener import FlattenMapper as FlattenMapperBase from pymbolic.mapper.stringifier import StringifyMapper as StringifyMapperBase from pymbolic.mapper.substitutor import ( CachedSubstitutionMapper as SubstitutionMapperBase, @@ -195,6 +196,14 @@ def map_resolved_function(self, expr, *args, **kwargs): map_fortran_division = IdentityMapperBase.map_quotient +class FlattenMapper(FlattenMapperBase, IdentityMapperMixin): + pass + + +def flatten(expr: ExpressionT) -> ExpressionT: + return FlattenMapper()(expr) + + class IdentityMapper(IdentityMapperBase, IdentityMapperMixin): pass @@ -1833,7 +1842,7 @@ def aff_to_expr(aff: isl.Aff) -> ExpressionT: if coeff: result += coeff*aff_to_expr(aff.get_div(i)) - return result // denom + return flatten(result // denom) def pw_aff_to_expr(pw_aff: isl.PwAff, int_ok: bool = False) -> ExpressionT: @@ -2178,14 +2187,17 @@ def qpolynomial_to_expr(qpoly): assert all(isinstance(num, int) for num in numerators) assert isinstance(common_denominator, int) + # FIXME: Delete if in favor of the general case once we depend on pymbolic 2024.1. if common_denominator == 1: - return sum(num * monomial + res = sum(num * monomial for num, monomial in zip(numerators, monomials)) else: - return FloorDiv(sum(num * monomial + res = FloorDiv(sum(num * monomial for num, monomial in zip(numerators, monomials)), common_denominator) + return flatten(res) + # }}} diff --git a/loopy/target/c/codegen/expression.py b/loopy/target/c/codegen/expression.py index 276407fc1..810ac4379 100644 --- a/loopy/target/c/codegen/expression.py +++ b/loopy/target/c/codegen/expression.py @@ -28,7 +28,7 @@ import islpy as isl import pymbolic.primitives as p from pymbolic import var -from pymbolic.mapper import IdentityMapper, RecursiveMapper +from pymbolic.mapper import IdentityMapper, Mapper from pymbolic.mapper.stringifier import ( PREC_BITWISE_AND, PREC_BITWISE_OR, @@ -124,9 +124,8 @@ def wrap_in_typecast(self, actual_type: LoopyType, needed_type: LoopyType, s): return s - def rec(self, expr, type_context=None, needed_type: Optional[LoopyType] = None): - result = RecursiveMapper.rec(self, expr, type_context) - + def rec(self, expr, type_context=None, needed_type: Optional[LoopyType] = None): # type: ignore[override] + result = Mapper.rec(self, expr, type_context) if needed_type is None: return result else: @@ -604,7 +603,7 @@ def map_nan(self, expr, type_context): # {{{ C expression to code mapper -class CExpressionToCodeMapper(RecursiveMapper): +class CExpressionToCodeMapper(Mapper): # {{{ helpers