diff --git a/src/pystencils/backend/kernelcreation/freeze.py b/src/pystencils/backend/kernelcreation/freeze.py index 63e9ea5b15f4fd15f87ff254862d1e6c17d1b626..5cdb3864c1f1d09e65686dd36da3623d837e71d5 100644 --- a/src/pystencils/backend/kernelcreation/freeze.py +++ b/src/pystencils/backend/kernelcreation/freeze.py @@ -13,7 +13,7 @@ from ...sympyextensions import ( integer_functions, ConditionalFieldAccess, ) -from ...compound_op_mapping import compound_op_to_expr +from ...reduction_op_mapping import reduction_op_to_expr from ...sympyextensions.typed_sympy import TypedSymbol, TypeCast, DynamicType from ...sympyextensions.pointers import AddressOf, mem_acc from ...sympyextensions.reduction import ReductionAssignment, ReductionOp @@ -174,15 +174,17 @@ class FreezeExpressions: assert isinstance(lhs, PsExpression) assert isinstance(rhs, PsExpression) - _str_to_compound_op: dict[str, ReductionOp] = { + # transform augmented assignment to reduction op + str_to_reduction_op: dict[str, ReductionOp] = { "+=": ReductionOp.Add, "-=": ReductionOp.Sub, "*=": ReductionOp.Mul, "/=": ReductionOp.Div, } + # reuse existing handling for transforming reduction ops to expressions return PsAssignment( - lhs, compound_op_to_expr(_str_to_compound_op[expr.op], lhs.clone(), rhs) + lhs, reduction_op_to_expr(str_to_reduction_op[expr.op], lhs.clone(), rhs) ) def map_ReductionAssignment(self, expr: ReductionAssignment): @@ -208,7 +210,7 @@ class FreezeExpressions: new_lhs = PsSymbolExpr(new_lhs_symb) # get new rhs from augmented assignment - new_rhs: PsExpression = compound_op_to_expr(op, new_lhs.clone(), rhs) + new_rhs: PsExpression = reduction_op_to_expr(op, new_lhs.clone(), rhs) # match for reduction operation and set neutral init_val init_val: PsExpression diff --git a/src/pystencils/backend/platforms/generic_cpu.py b/src/pystencils/backend/platforms/generic_cpu.py index ccef6181793af2c8fa3271edcf5f6b62716863a1..3de7cf696ff47d65950e376e4aa38c63e1528dcf 100644 --- a/src/pystencils/backend/platforms/generic_cpu.py +++ b/src/pystencils/backend/platforms/generic_cpu.py @@ -13,7 +13,7 @@ from ..functions import ( PsReductionFunction, ) from ..literals import PsLiteral -from ...compound_op_mapping import compound_op_to_expr +from ...reduction_op_mapping import reduction_op_to_expr from ...sympyextensions import ReductionOp from ...types import PsIntegerType, PsIeeeFloatType, PsScalarType, PsPointerType @@ -97,7 +97,7 @@ class GenericCpu(Platform): actual_op = ReductionOp.Add if op is ReductionOp.Sub else op # create binop and potentially select corresponding function for e.g. min or max - potential_call = compound_op_to_expr(actual_op, ptr_access, symbol_expr) + potential_call = reduction_op_to_expr(actual_op, ptr_access, symbol_expr) if isinstance(potential_call, PsCall): potential_call.dtype = symbol_expr.dtype return self.select_function(potential_call) diff --git a/src/pystencils/backend/platforms/generic_gpu.py b/src/pystencils/backend/platforms/generic_gpu.py index 8b7eead8dc276e8935f9a07c7b80b008a8ae9ca6..349e79d4b6a727e84d9f2ec6e941c231ad3fcb81 100644 --- a/src/pystencils/backend/platforms/generic_gpu.py +++ b/src/pystencils/backend/platforms/generic_gpu.py @@ -7,7 +7,7 @@ from functools import reduce from ..ast import PsAstNode from ..constants import PsConstant -from ...compound_op_mapping import compound_op_to_expr +from ...reduction_op_mapping import reduction_op_to_expr from ...sympyextensions.reduction import ReductionOp from ...types import constify, deconstify, PsPointerType, PsScalarType, PsCustomType, PsType from ...types.quick import UInt, SInt @@ -325,7 +325,7 @@ class GenericGpu(Platform): shuffles = tuple( PsAssignment( symbol_expr, - compound_op_to_expr( + reduction_op_to_expr( actual_op, symbol_expr, gen_shuffle_instr(pow(2, i - 1)) ), ) diff --git a/src/pystencils/compound_op_mapping.py b/src/pystencils/reduction_op_mapping.py similarity index 82% rename from src/pystencils/compound_op_mapping.py rename to src/pystencils/reduction_op_mapping.py index 193b308d01de6c12c40e1104f7704b97d02ecba3..06fb8aa3e981d661f8e4226a34307f03cd6d9a70 100644 --- a/src/pystencils/compound_op_mapping.py +++ b/src/pystencils/reduction_op_mapping.py @@ -11,7 +11,7 @@ _available_operator_interface: set[ReductionOp] = { } -def compound_op_to_expr(op: ReductionOp, op1, op2) -> PsExpression: +def reduction_op_to_expr(op: ReductionOp, op1, op2) -> PsExpression: if op in _available_operator_interface: match op: case ReductionOp.Add: @@ -24,7 +24,7 @@ def compound_op_to_expr(op: ReductionOp, op1, op2) -> PsExpression: operator = PsDiv case _: raise FreezeError( - f"Found unsupported operation type for compound assignments: {op}." + f"Found unsupported operation type for reduction assignments: {op}." ) return operator(op1, op2) else: @@ -35,5 +35,5 @@ def compound_op_to_expr(op: ReductionOp, op1, op2) -> PsExpression: return PsCall(PsMathFunction(MathFunctions.Max), [op1, op2]) case _: raise FreezeError( - f"Found unsupported operation type for compound assignments: {op}." + f"Found unsupported operation type for reduction assignments: {op}." )