Reduction Support

This MR introduces reductions to pystencils for scalar data types and thus covers #55.

User interface

• Adds reduction assignment classes to sympyextensions module: AddReductionAssignment, SubReductionAssignment, MulReductionAssignment, MinReductionAssignment, MaxReductionAssignment

These can be used as follows:

import pystencils as ps

r = ps.TypedSymbol("r", "double")
x, y = ps.fields(f"x, y: double[3D]", layout="fzyx")

assign_dot_prod = ps.AddReductionAssignment(r, x.center() * y.center())

• Alternatívely, you can also make use of the reduction_assignment or reduction_assignment_from_str functions:

from pystencils.sympyextensions import reduction_assignment, reduction_assignment_from_str
from pystencils.sympyextensions.reduction import ReductionOp

assign_dot_prod = reduction_assignment(r, ReductionOp.Add, x.center() * y.center())

assign_dot_prod = reduction_assignment_from_str(r, "+", x.center() * y.center())

Supported Backends

Generic CPUs

• Add reduction support for OpenMP

SIMD: SSE3, AVX2, AVX512

• Include a generated header file with horizontal operations performing a binary operation between a scalar variable and a SIMD vector. The SIMD vector is transformed to a scalar variable via reduction, and then the binary operation is applied to the other operand

CUDA

• Employ atomic reduction operations in all threads when the block size does not align with the warp size
• Optimization for alignment with warp size: perform warp-level reductions and only perform atomic operation on first thread of warp
• Include a header file with manual implementations for atomic operations that are not directly supported for floating point numbers: atomicMul, atomicMax, and atomicMin. These functions make use of a CAS mechanism.

Internal Changes

• Freeze handling for newly introduced ReductionAssignment nodes
• Add PsVecHorizontal vectorization node for conducting a binary operation between a scalar symbol and an extraction of a vector value (obtained by performing a reduction within a vector lane)
• Add dataclass ReductionInfo holding essential information about a reduction (i.e. reduction operation, initial value and the write-back pointer for exporting the reduction result) and create corresponding lookup table for symbols in KernelCreationContext
• Introduce NumericLimitsFunctions for initializing neutral elements for reductions making use of min/max operations
• Adapt Platform.select_function such that it either returns an PsExpression that replaces the function call or returns a PsExpression | tuple[tuple[PsStructuralNode, ...], PsAstNode] holding a PsAstNode that replaces the function call and tuple of structural nodes that are added before the replacement. The structural nodes allow adding preparatory code for the replacement, as needed for the warp-level reductions for GPU platforms
• Add ReductionFunctions.WriteBackToPtr function that is replaced with platform-dependent code in Platform.select_function
• Slightly adapt CPU/GPU Jit modules to support the handling of write-back pointers used for reductions

src/pystencils/backend/platforms/cuda.py

 from __future__ import annotations
+
+import math
+import operator
 from abc import ABC, abstractmethod
+from functools import reduce
+
+from pystencils.types import PsBoolType
 
 from ..ast import PsAstNode
+from ..constants import PsConstant
+from ...compound_op_mapping import compound_op_to_expr
 from ...sympyextensions.reduction import ReductionOp
 from ...types import constify, deconstify, PsPointerType, PsScalarType, PsCustomType
+from ...types.quick import UInt, SInt
 from ..exceptions import MaterializationError
 from .generic_gpu import GenericGpu
 
@@ -17,14 +26,14 @@ from ..kernelcreation import (
 )
 
 from ..kernelcreation.context import KernelCreationContext
-from ..ast.structural import PsBlock, PsConditional, PsDeclaration, PsStatement
+from ..ast.structural import PsBlock, PsConditional, PsDeclaration, PsStatement, PsAssignment
 from ..ast.expressions import (
     PsExpression,
     PsLiteralExpr,
     PsCast,
     PsCall,
     PsLookup,
-    PsBufferAcc, PsSymbolExpr
+    PsBufferAcc, PsSymbolExpr, PsConstantExpr, PsAdd, PsRem, PsEq
 )
 from ..ast.expressions import PsLt, PsAnd
 from ...types import PsSignedIntegerType, PsIeeeFloatType
@@ -292,28 +301,104 @@ class CudaPlatform(GenericGpu):
             case ReductionFunctions.WriteBackToPtr:
                 ptr_expr, symbol_expr = call.args
                 op = call.function.reduction_op
+                stype = symbol_expr.dtype
+                ptrtype = ptr_expr.dtype
+
+                warp_size = 32   # TODO: get from platform/user config
+
+                assert isinstance(ptr_expr, PsSymbolExpr) and isinstance(ptrtype, PsPointerType)
+                assert isinstance(symbol_expr, PsSymbolExpr) and isinstance(stype, PsScalarType)
+
+                if not isinstance(stype, PsIeeeFloatType) or stype.width not in (32, 64):
+                    NotImplementedError("atomic operations are only available for float32/64 datatypes")
 
-                assert isinstance(ptr_expr, PsSymbolExpr) and isinstance(ptr_expr.dtype, PsPointerType)
-                assert isinstance(symbol_expr, PsSymbolExpr) and isinstance(symbol_expr.dtype, PsScalarType)
+                # set up mask symbol for active threads in warp
+                mask = PsSymbol("__shfl_mask", UInt(32))
+                self._ctx.add_symbol(mask)
 
+                # workaround for subtractions -> use additions for reducing intermediate results
+                # similar to OpenMP reductions: local copies (negative sign) are added at the end
                 match op:
                     case ReductionOp.Sub:
-                        # workaround for unsupported atomicSub: use atomic add
-                        # similar to OpenMP reductions: local copies (negative sign) are added at the end
-                        call.function = CFunction("atomicAdd", [ptr_expr.dtype, symbol_expr.dtype],
-                                                  PsCustomType("void"))
-                        call.args = (ptr_expr, symbol_expr)
+                        actual_op = ReductionOp.Add
                     case _:
-                        call.function = CFunction(f"atomic{op.name}", [ptr_expr.dtype, symbol_expr.dtype],
-                                                  PsCustomType("void"))
-                        call.args = (ptr_expr, symbol_expr)
+                        actual_op = op
+
+                # perform local warp reductions
+                def gen_shuffle_instr(offset: int):
+                    return PsCall(CFunction("__shfl_xor_sync", [UInt(32), stype, SInt(32)], stype),
+                                  [PsSymbolExpr(mask), symbol_expr, PsConstantExpr(PsConstant(offset, SInt(32)))])
+
+                num_shuffles = math.frexp(warp_size)[1]
+                shuffles = [PsAssignment(symbol_expr,
+                                         compound_op_to_expr(actual_op, symbol_expr, gen_shuffle_instr(pow(2, i - 1))))
+                            for i in reversed(range(1, num_shuffles))]
+
+                # find first thread in warp
+                ispace = self._ctx.get_iteration_space()  # TODO: receive as argument in unfold_function?
+                is_valid_thread = self._get_condition_for_translation(ispace)
+                thread_indices_per_dim = [
+                    idx * PsConstantExpr(PsConstant(reduce(operator.mul, BLOCK_DIM[:i], 1), SInt(32)))
+                    for i, idx in enumerate(THREAD_IDX[:ispace.rank])
+                ]
+                tid: PsExpression = thread_indices_per_dim[0]
+                for t in thread_indices_per_dim[1:]:
+                    tid = PsAdd(tid, t)
+                first_thread_in_warp = PsEq(PsRem(tid, PsConstantExpr(PsConstant(warp_size, SInt(32)))),
+                                            PsConstantExpr(PsConstant(0, SInt(32))))
+                cond = PsAnd(is_valid_thread, first_thread_in_warp) if is_valid_thread else first_thread_in_warp
+
+                full_mask = PsLiteralExpr(PsLiteral("0xffffffff", UInt(32)))
+                ballot_instr = PsCall(CFunction("__ballot_sync", [UInt(32), SInt(32)], SInt(32)),
+                                                 [full_mask, is_valid_thread])
+                decl_mask = PsDeclaration(PsSymbolExpr(mask), ballot_instr if is_valid_thread else full_mask)
+
+                # use atomic operation on first thread of warp to sync
+                call.function = CFunction(f"atomic{actual_op.name}", [ptrtype, stype], PsCustomType("void"))
+                call.args = (ptr_expr, symbol_expr)
+
+                # assemble warp reduction
+                return PsConditional(is_valid_thread if is_valid_thread else PsConstantExpr(PsLiteral("true", PsBoolType)),
+                    PsBlock(
+                    [decl_mask]
+                    + shuffles
+                    + [PsConditional(cond, PsBlock([PsStatement(call)]))]))
 
-                if not isinstance(symbol_expr.dtype, PsIeeeFloatType) or symbol_expr.dtype.width not in (32, 64):
-                    NotImplementedError("atomicMul is only available for float32/64 datatypes")
+    #   Internals
 
-                return PsStatement(call)
+    # TODO: SYCL platform has very similar code for fetching conditionals -> move to GenericGPU?
 
-    #   Internals
+    def _get_condition_for_translation(
+            self, ispace: IterationSpace):
+
+        if self._omit_range_check:
+            return None
+
+        match ispace:
+            case FullIterationSpace():
+
+                dimensions = ispace.dimensions_in_loop_order()
+
+                conds = []
+                for dim in dimensions:
+                    ctr_expr = PsExpression.make(dim.counter)
+                    conds.append(PsLt(ctr_expr, dim.stop))
+
+                    if conds:
+                        condition: PsExpression = conds[0]
+                        for cond in conds[1:]:
+                            condition = PsAnd(condition, cond)
+                        return condition
+                    else:
+                        return None
+
+            case SparseIterationSpace():
+                sparse_ctr_expr = PsExpression.make(ispace.sparse_counter)
+                stop = PsExpression.make(ispace.index_list.shape[0])
+
+                return PsLt(sparse_ctr_expr.clone(), stop)
+            case _:
+                assert False, "Unknown iteration space"
 
     def _prepend_dense_translation(
         self, body: PsBlock, ispace: FullIterationSpace
@@ -321,7 +406,7 @@ class CudaPlatform(GenericGpu):
         ctr_mapping = self._thread_mapping(ispace)
 
         indexing_decls = []
-        conds = []
+        cond = self._get_condition_for_translation(ispace)
 
         dimensions = ispace.dimensions_in_loop_order()
 
@@ -335,14 +420,9 @@ class CudaPlatform(GenericGpu):
             indexing_decls.append(
                 self._typify(PsDeclaration(ctr_expr, ctr_mapping[dim.counter]))
             )
-            if not self._omit_range_check:
-                conds.append(PsLt(ctr_expr, dim.stop))
-
-        if conds:
-            condition: PsExpression = conds[0]
-            for cond in conds[1:]:
-                condition = PsAnd(condition, cond)
-            ast = PsBlock(indexing_decls + [PsConditional(condition, body)])
+
+        if cond:
+            ast = PsBlock(indexing_decls + [PsConditional(cond, body)])
         else:
             body.statements = indexing_decls + body.statements
             ast = body
@@ -355,6 +435,8 @@ class CudaPlatform(GenericGpu):
         factory = AstFactory(self._ctx)
         ispace.sparse_counter.dtype = constify(ispace.sparse_counter.get_dtype())
 
+        cond = self._get_condition_for_translation(ispace)
+
         sparse_ctr_expr = PsExpression.make(ispace.sparse_counter)
         ctr_mapping = self._thread_mapping(ispace)
 
@@ -377,10 +459,8 @@ class CudaPlatform(GenericGpu):
         ]
         body.statements = mappings + body.statements
 
-        if not self._omit_range_check:
-            stop = PsExpression.make(ispace.index_list.shape[0])
-            condition = PsLt(sparse_ctr_expr.clone(), stop)
-            ast = PsBlock([sparse_idx_decl, PsConditional(condition, body)])
+        if cond:
+            ast = PsBlock([sparse_idx_decl, PsConditional(cond, body)])
         else:
             body.statements = [sparse_idx_decl] + body.statements
             ast = body