Add concatenate_arrays (#717)

* Add merge_temporary_arrays transformation

* type annotate

* move offset computation to a separate loop

* check dependency_names before using map_instructions

* rename

* generate a new array name if not given

* expect a list of args instead of comma separated string of args

* fix flake8 errors

* make mypy's job easier

* Rename, clean up concatenate_arrays

Co-authored-by: Andreas Kloeckner <inform@tiker.net>

Author: isuruf

loopy/__init__.py

@@ -105,6 +105,7 @@
 from loopy.transform.precompute import precompute
 from loopy.transform.buffer import buffer_array
 from loopy.transform.fusion import fuse_kernels
+from loopy.transform.concatenate import concatenate_arrays
 
 from loopy.transform.arithmetic import (
         fold_constants,
@@ -232,6 +233,7 @@
 
         "precompute", "buffer_array",
         "fuse_kernels",
+        "concatenate_arrays",
 
         "fold_constants", "collect_common_factors_on_increment",
 

loopy/transform/concatenate.py

@@ -0,0 +1,130 @@
+__copyright__ = "Copyright (C) 2022 Isuru Fernando"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+__doc__ = """
+.. currentmodule:: loopy
+.. autofunction:: concatenate_arrays
+"""
+
+from typing import Sequence, Optional, List
+
+from loopy.kernel.data import ArrayArg, KernelArgument, TemporaryVariable, auto
+from loopy.symbolic import SubstitutionRuleMappingContext
+from loopy.kernel import LoopKernel
+from loopy.translation_unit import for_each_kernel
+
+import pymbolic.primitives as prim
+from pytools import all_equal
+
+
+@for_each_kernel
+def concatenate_arrays(
+        kernel: LoopKernel,
+        array_names: Sequence[str],
+        new_name: Optional[str] = None,
+        axis_nr: int = 0) -> LoopKernel:
+    """Merges arrays (temporaries or arguments) into one array along the axis
+    given by *axis_nr*.
+
+    :arg array_names: a list of names of temporary variables.
+
+    :arg axis_nr: the (zero-based) index of the axis of the arrays to be merged.
+
+    :arg new_name: new name for the merged temporary. If not given, a new name
+        is generated.
+    """
+    assert isinstance(kernel, LoopKernel)
+
+    var_name_gen = kernel.get_var_name_generator()
+    new_name = new_name or var_name_gen("concatenated_array")
+    new_aggregate = prim.Variable(new_name)
+
+    arrays = []
+    for array_name in array_names:
+        ary = kernel.get_var_descriptor(array_name)
+        if ary.shape is None or ary.shape is auto:
+            raise ValueError(f"Shape of temporary variable '{array_name}' is "
+                    "unknown. Cannot merge with unknown shapes")
+
+        assert isinstance(ary.shape, tuple)
+        shape = list(ary.shape)
+        # make the shape value at axis_nr a constant so that we can
+        # check that the rest of the attributes (except name) are equal.
+        shape[axis_nr] = 1
+        arrays.append(ary.copy(shape=tuple(shape), name=new_name))
+
+    if not all_equal(arrays):
+        raise ValueError("Arrays must be identical except for shape "
+                "(except for shape) in order to concatenate.")
+
+    offsets = {}
+    axis_length = 0
+    for array_name in array_names:
+        offsets[array_name] = axis_length
+        ary = kernel.temporary_variables[array_name]
+        assert isinstance(ary.shape, tuple)
+        axis_length += ary.shape[axis_nr]
+
+    new_ary = arrays[0]
+    new_shape = list(new_ary.shape)
+    new_shape[axis_nr] = axis_length
+    new_ary = new_ary.copy(shape=tuple(new_shape))
+
+    # {{{ rewrite subscripts
+
+    from loopy.transform.padding import SubscriptRewriter
+
+    def modify_array_access(expr):
+        idx = expr.index
+        if not isinstance(idx, tuple):
+            idx = (idx,)
+        idx = list(idx)
+        idx[axis_nr] += offsets[expr.aggregate.name]
+
+        return new_aggregate.index(tuple(idx))
+
+    rule_mapping_context = SubstitutionRuleMappingContext(
+            kernel.substitutions, var_name_gen)
+    aash = SubscriptRewriter(rule_mapping_context,
+            array_names, modify_array_access)
+    kernel = rule_mapping_context.finish_kernel(aash.map_kernel(kernel))
+
+    # }}}
+
+    if isinstance(new_ary, TemporaryVariable):
+        new_tvs = {name: tv for name, tv in kernel.temporary_variables.items()
+                if name not in array_names}
+        new_tvs[new_name] = new_ary
+        return kernel.copy(temporary_variables=new_tvs)
+    elif isinstance(new_ary, ArrayArg):
+        new_args: List[KernelArgument] = []
+        inserted = False
+        for arg in kernel.args:
+            if arg.name in array_names:
+                if not inserted:
+                    new_args.append(new_ary)
+                    inserted = True
+            else:
+                new_args.append(arg)
+        return kernel.copy(args=new_args)
+    else:
+        raise AssertionError()

loopy/transform/padding.py

@@ -31,7 +31,7 @@
 from loopy.diagnostic import LoopyError
 
 
-class ArrayAxisSplitHelper(RuleAwareIdentityMapper):
+class SubscriptRewriter(RuleAwareIdentityMapper):
     def __init__(self, rule_mapping_context, arg_names, handler):
         super().__init__(rule_mapping_context)
         self.arg_names = arg_names
@@ -43,6 +43,21 @@ def map_subscript(self, expr, expn_state):
         else:
             return super().map_subscript(expr, expn_state)
 
+    def map_kernel(self, kernel, within=lambda *args: True):
+        new_insns = [
+            # While subst rules are not allowed in assignees, the mapper
+            # may perform tasks entirely unrelated to subst rules, so
+            # we must map assignees, too.
+            insn if not kernel.substitutions and not within(kernel, insn, ()) \
+            and not any(name in self.arg_names for name in \
+                insn.dependency_names()) else
+            self.map_instruction(kernel,
+                insn.with_transformed_expressions(
+                    lambda expr: self(expr, kernel, insn)))  # noqa: B023
+            for insn in kernel.instructions]
+
+        return kernel.copy(instructions=new_insns)
+
 
 # {{{ split_array_dim (deprecated since June 2016)
 
@@ -236,7 +251,7 @@ def split_access_axis(expr):
 
     rule_mapping_context = SubstitutionRuleMappingContext(
             kernel.substitutions, var_name_gen)
-    aash = ArrayAxisSplitHelper(rule_mapping_context,
+    aash = SubscriptRewriter(rule_mapping_context,
             set(array_to_rest.keys()), split_access_axis)
     kernel = rule_mapping_context.finish_kernel(aash.map_kernel(kernel))
 
@@ -367,7 +382,7 @@ def split_access_axis(expr):
 
     rule_mapping_context = SubstitutionRuleMappingContext(
             kernel.substitutions, var_name_gen)
-    aash = ArrayAxisSplitHelper(rule_mapping_context,
+    aash = SubscriptRewriter(rule_mapping_context,
             {array_name}, split_access_axis)
     kernel = rule_mapping_context.finish_kernel(aash.map_kernel(kernel))
 

test/test_transform.py

@@ -1608,6 +1608,25 @@ def test_prefetch_to_same_temp_var(ctx_factory):
     lp.auto_test_vs_ref(ref_tunit, ctx, t_unit)
 
 
+def test_concatenate_arrays(ctx_factory):
+    ctx = ctx_factory()
+
+    t_unit = lp.make_kernel(
+        "{[i]: 0<=i<10}",
+        """
+        <> a[i] = x[i]    {id=init_a}
+        <> b[i] = y[i]    {id=init_b}
+        out[i] = a[i] + b[i] {id=insn,dep=init_a:init_b}
+        """)
+
+    t_unit = lp.add_dtypes(t_unit, {"x": "float64", "y": "float64"})
+    ref_t_unit = t_unit
+
+    t_unit = lp.concatenate_arrays(t_unit, ["a", "b"], "c")
+    assert t_unit.default_entrypoint.temporary_variables["c"].shape == (20,)
+    lp.auto_test_vs_ref(ref_t_unit, ctx, t_unit)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])