added tests and cleaned loop blocking logic

loki/transformations/loop_blocking.py

@@ -4,118 +4,207 @@
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
-
+from pymbolic.primitives import Expression
 
 from loki import Subroutine, ir, FindNodes, Assignment, Sum, Product, Loop, LoopRange, IntLiteral, \
     Quotient, InlineCall, DeferredTypeSymbol, Transformer, Array, Variable, RangeIndex, Scalar, \
-    SubstituteExpressions, FindVariables
+    SubstituteExpressions, FindVariables, parse_expr, fgen
+
+
+
+def ceil_division(iexpr1: Expression, iexpr2: Expression) -> Expression:
+    """
+    Returns ceiled division expression of two integer expressions iexpr1/iexpr2.
+    """
+    return Sum(children=(Quotient(numerator=Sum(children=(iexpr1, IntLiteral(-1))),
+                                  denominator=iexpr2), IntLiteral(1)))
+
+
+
+def negate(i: Expression) -> Expression:
+    return Product((IntLiteral(-1), i))
+
+
+def num_iterations(loop_range: LoopRange) -> Expression:
+    """
+    Returns total number of iterations of a loop.
+
+    Given a loop, this returns an expression that computes the total number of
+    iterations of the loop, i.e.
+    (start,stop,step) -> ceil(stop-start/step)
+    """
+    start = loop_range.start
+    stop = loop_range.stop
+    step = loop_range.step
+    if step is None:
+        return stop if isinstance(start, IntLiteral) and start.value == 1 else Sum(
+            (stop, negate(start), IntLiteral(1)))
+    else:
+        return Sum((Quotient(Sum((stop, negate(start))), step), IntLiteral(1)))
+
+
+def normalized_loop_range(loop_range: LoopRange) -> LoopRange:
+    """
+    Returns the normalized LoopRange of a given LoopRange.
+
+    Returns the normalized LoopRange which corresponds to a loop with the same
+    number of iterations but starts at 1 and has stride 1, i.e.
+    (start,stop,step) -> (1,num_iter,1)
+    """
+    return LoopRange((1, num_iterations(loop_range)))
+
+
+def iteration_number(iter_idx: Variable, loop_range: LoopRange) -> Expression:
+    """
+    Returns the normalized iteration number of the iteration variable, i.e.
+    (iter_idx - start + step)/step = (iter_idx-start)/step + 1
+    """
+    if loop_range.step is None:
+        return Sum((Sum((iter_idx, negate(loop_range.start))), IntLiteral(1)))
+    else:
+        return Sum((Quotient(Sum((iter_idx, negate(loop_range.start))), loop_range.step),
+                   IntLiteral(1)))
+
+
+def iteration_index(iter_num: Variable, loop_range: LoopRange) -> Expression:
+    """
+    Returns the iteration index of the loop based on the iteration number, i.e.
+    iter_idx = (iter_num-1)*step+start
+    """
+    if loop_range.step is None:
+        return Sum((iter_num, loop_range.start))
+    else:
+        return Sum(
+            (Product((Sum((iter_num, IntLiteral(-1))), loop_range.step)), loop_range.start))
 
 
 def split_loop(routine: Subroutine, loop: ir.Loop, block_size: int):
     """
-    Assumes that loop is contiguous in blocking variable
-    into a blocked loop in those variables. Does not block
-    the loop data
+    Blocks a loop by splitting it into an outer loop and inner loop of size `block_size`.
+
+    Parameters
+    ----------
+    routine: :any:`Subroutine`
+        Subroutine object containing the loop. New variables introduced in the
+        loop splitting will be declared in the body of routine.
+    loop: :any:`Loop`
+        Loop to be split.
+    block_size: int
+        inner loop size (size of blocking blocks)
     """
     # blocking variable declarations
-    decls = FindNodes(ir.VariableDeclaration).visit(routine.spec)
-    block_size_var = loop.variable.clone(name=loop.variable.name + '_loop_block_size')
-    num_blocks_var = loop.variable.clone(name=loop.variable.name + '_loop_num_blocks')
-    outer_loop_var = loop.variable.clone(name=loop.variable.name + '_loop_outer_idx')
-    inner_loop_var = loop.variable.clone(name=loop.variable.name + '_loop_local')
+    block_size_var = loop.variable.clone(name=loop.variable.name + "_loop_block_size")
+    num_blocks_var = loop.variable.clone(name=loop.variable.name + "_loop_num_blocks")
+    block_idx = loop.variable.clone(name=loop.variable.name + "_loop_block_idx")
+    inner_loop_var = loop.variable.clone(name=loop.variable.name + "_loop_local")
+    iter_num_var = loop.variable.clone(name=loop.variable.name + "_loop_iter_num")
     global_loop_var = loop.variable
-    block_start = loop.variable.clone(name=loop.variable.name + '_loop_block_start')
-    block_end = loop.variable.clone(name=loop.variable.name + '_loop_block_end')
+    block_start = loop.variable.clone(name=loop.variable.name + "_loop_block_start")
+    block_end = loop.variable.clone(name=loop.variable.name + "_loop_block_end")
 
-    routine.variables += (block_size_var, num_blocks_var, outer_loop_var, block_start, block_end)
-
-    # Inner loop
-    inner_body = Assignment(global_loop_var, Sum(children=(
-        Product(children=(outer_loop_var, block_size_var)), inner_loop_var)))
-    inner_loop = Loop(variable=inner_loop_var, body=(inner_body,) + loop.body,
-                      bounds=LoopRange((IntLiteral(1), block_size_var)))
-
-    # blocking variables outside loop initalization
-    block_loop_inits = [ir.Assignment(block_size_var, IntLiteral(block_size))]
-    q = Quotient(numerator=Sum(children=(loop.bounds.upper, block_size_var, IntLiteral(-1))),
-                 denominator=block_size_var)
-    block_loop_inits.append(ir.Assignment(num_blocks_var, q))
-    change_map = {loop: block_loop_inits}
-    #  Outer loop bounds + body
+    routine.variables += (
+        block_size_var, num_blocks_var, block_idx, block_start, block_end, iter_num_var)
 
-    outer_bounds = LoopRange((IntLiteral(1), num_blocks_var))
-    outer_body = (
+    # block index calculations
+    blocking_body = (
         Assignment(block_start,
-                   Sum(children=(
-                       Product(children=(
-                           Sum(children=(
-                               outer_loop_var, IntLiteral(-1)
-                           )),
-                           block_size_var
-                       )),
-                       IntLiteral(1)
-                   ))
+                   parse_expr(f"({block_idx} - 1) * {block_size_var} + 1")
                    ),
         Assignment(block_end,
                    InlineCall(DeferredTypeSymbol('MIN', scope=routine),
-                              parameters=(Product(children=(outer_loop_var, block_size_var)),
+                              parameters=(Product(children=(block_idx, block_size_var)),
                                           loop.bounds.upper))
                    ))
-    outer_loop = Loop(variable=outer_loop_var, body=outer_body + (inner_loop,),
-                      bounds=outer_bounds)
 
-    change_map[loop].append(outer_loop)
+    # Outer loop blocking variable assignments
+    loop_range = loop.bounds
+    block_loop_inits = (ir.Assignment(block_size_var, IntLiteral(block_size)),
+                        ir.Assignment(num_blocks_var,
+                                      ceil_division(num_iterations(loop_range), block_size_var)))
+
+    # Inner loop
+    iteration_nums = (
+        Assignment(iter_num_var, parse_expr(f"{block_start}+{inner_loop_var}-1")),
+        Assignment(global_loop_var, iteration_index(iter_num_var, loop_range))
+    )
+    inner_loop = Loop(variable=inner_loop_var, body=iteration_nums + loop.body,
+                      bounds=LoopRange((IntLiteral(1), block_size_var)))
+
+    #  Outer loop bounds + body
+    outer_loop = Loop(variable=block_idx, body=blocking_body + (inner_loop,),
+                      bounds=LoopRange((IntLiteral(1), num_blocks_var)))
 
+    change_map = {loop: block_loop_inits + (outer_loop,)}
     Transformer(change_map, inplace=True).visit(routine.ir)
     return loop.variable, inner_loop, outer_loop
 
 
-def blocked_shape(a: Array, blocking_var: Variable, loop: ir.Loop):
-    dims = tuple(RangeIndex(
-        (1, loop.bounds.upper)) if isinstance(dim, Scalar) and blocking_var in dim else dim for dim
-                 in a.dimensions)
+def blocked_shape(a: Array, blocking_vars: list[Variable], loop: ir.Loop):
+    """
+    calculates the dimensions for a blocked version of the array.
+    """
+    dims = tuple(
+        loop.bounds.upper if isinstance(dim, Scalar) and any(
+            blocking_var in dim for blocking_var in blocking_vars) else dim for dim
+        in a.dimensions)
     return dims
 
 
 def blocked_type(a: Array):
     return a.type.clone(intent=None)
 
 
-def block_loop_dims(a: Array, blocking_var: Variable, loop: ir.Loop):
+def block_loop_dims(a: Array, blocking_vars: list[Variable], loop: ir.Loop):
+    """
+    calculates the dimension index of the blocked array inside the loop.
+    """
     dims = tuple(
-        loop.variable if isinstance(dim, Scalar) and blocking_var in dim else dim for dim
+        loop.variable if isinstance(dim, Scalar) and any(
+            blocking_var in dim for blocking_var in blocking_vars) else dim for dim
         in a.dimensions)
     return dims
 
 
-def insert_block_loop_arrays(routine: Subroutine, loop_var: Variable, inner_loop: ir.Loop,
+def insert_block_loop_arrays(routine: Subroutine, blocking_indices: list[Variable], inner_loop: ir.Loop,
                              outer_loop: ir.Loop, block_size: int):
     """
+    Replaces arrays inside the inner loop with blocked counterparts.
+
     This routine declares array variables to hold the blocks of the arrays used inside
-    the loop and replaces array variables inside the loop with their blocked counterparts
+    the loop and replaces array variables inside the loop with their blocked counterparts.
+    An array is blocked with the leading dimensions
+
+    Parameters
+    ----------
+    routine : Subroutine
+
+    blocking_indices: list of  :any:`Variable`
+        list of the index variables that arrays inside the loop should be blocked by.
+    inner_loop: :any:`Loop`
+        inner loop after loop splitting
+    outer_loop: :any:`Loop`
+        outer loop after loop splitting
+    block_size: int
+        size of blocked arrays
     """
     # Declare Blocked arrays
-    arrays = [var for var in FindVariables().visit(inner_loop.body) if
-              isinstance(var, Array) and loop_var in var]
+    arrays = tuple(var for var in FindVariables().visit(inner_loop.body) if
+                   isinstance(var, Array) and any(
+                       bv in var for bv in blocking_indices))
     name_map = {a.name: a.name + '_block' for a in arrays}
     block_arrays = tuple(
-        a.clone(name=name_map[a.name], dimensions=blocked_shape(a, loop_var, inner_loop),
+        a.clone(name=name_map[a.name], dimensions=blocked_shape(a, blocking_indices, inner_loop),
                 type=blocked_type(a)) for a in arrays)
     routine.variables += block_arrays
-
+    # Replace arrays in loop with blocked arrays and update idx
     block_array_expr = [
-        a.clone(name=name_map[a.name], dimensions=block_loop_dims(a, loop_var, inner_loop)) for a in
+        a.clone(name=name_map[a.name], dimensions=block_loop_dims(a, blocking_indices, inner_loop)) for a
+        in
         arrays]
-    # Replace arrays in loop with blocked arrays and update idx
     body = SubstituteExpressions(dict(zip(arrays, block_array_expr))).visit(inner_loop.body)
     inner_loop._update(body=body)
 
 
 def handle_device_transfers(routine: Subroutine, loop_var: Variable, inner_loop: ir.Loop,
                             outer_loop: ir.Loop, block_size: int):
     pass
-
-
-def block_loop(routine: Subroutine, loop: ir.Loop, block_size: int):
-    loop_var, inner_loop, outer_loop = split_loop(routine, loop, block_size)
-    insert_block_loop_arrays(routine, loop_var, inner_loop, outer_loop, block_size)