Storage reduction pass

src/gt4py/analysis/passes.py

@@ -1297,6 +1297,94 @@ def apply(cls, transform_data: TransformData) -> None:
         cls.DemoteSymbols.apply(transform_data.implementation_ir, demotables)
 
 
+class ReduceTemporaryStoragesPass(TransformPass):
+    """Demote temporary symbols used within a single multi-stage to 2D fields if the following
+    constraints are satisfied:
+
+    1. The multi-stage iteration order is not parallel.
+    2. There are no offsets in the k-direction.
+    3. All of the accesses are in the same k-interval.
+    """
+
+    class ReducibleFieldsCollector(gt_ir.IRNodeVisitor):
+        @classmethod
+        def apply(cls, node: gt_ir.StencilImplementation) -> Set[str]:
+            collector = cls()
+            return collector(node)
+
+        def __call__(self, node: gt_ir.StencilImplementation) -> Set[str]:
+            assert isinstance(node, gt_ir.StencilImplementation)
+            self.interval: gt_ir.AxisInterval = None
+            self.iteration_order: gt_ir.IterationOrder = None
+            self.multi_stage: str = ""
+            self.reduced_fields: Dict[str, Dict[str, Union[str, gt_ir.AxisInterval]]] = {
+                temp_field: dict(multi_stage="", interval=None)
+                for temp_field in node.temporary_fields
+            }
+            self.visit(node)
+            return set(self.reduced_fields.keys())
+
+        def visit_MultiStage(self, node: gt_ir.MultiStage) -> None:
+            self.iteration_order = node.iteration_order
+            self.multi_stage = node.name
+            self.generic_visit(node)
+
+        def visit_ApplyBlock(self, node: gt_ir.ApplyBlock) -> None:
+            self.interval = node.interval
+            self.generic_visit(node)
+
+        def visit_FieldRef(self, node: gt_ir.FieldRef) -> None:
+            field_name = node.name
+            if field_name in self.reduced_fields:
+                if self.iteration_order == gt_ir.IterationOrder.PARALLEL:
+                    self.reduced_fields.pop(field_name)
+                else:
+                    offsets: List[int] = list(node.offset.values())
+                    if offsets[-1] != 0:
+                        self.reduced_fields.pop(field_name)
+                    else:
+                        interval = self.reduced_fields[field_name]["interval"]
+                        multi_stage = self.reduced_fields[field_name]["multi_stage"]
+                        if (interval is not None and interval != self.interval) or (
+                            multi_stage != "" and multi_stage != self.multi_stage
+                        ):
+                            self.reduced_fields.pop(field_name)
+                        else:
+                            self.reduced_fields[field_name]["interval"] = self.interval
+                            self.reduced_fields[field_name]["multi_stage"] = self.multi_stage
+
+    class StorageReducer(gt_ir.IRNodeVisitor):
+        @classmethod
+        def apply(cls, node: gt_ir.StencilImplementation, reduced_fields: Set[str]) -> None:
+            instance = cls(reduced_fields)
+            return instance(node)
+
+        def __init__(self, reduced_fields: Set[str]):
+            self.reduced_fields = reduced_fields
+
+        def __call__(self, node: gt_ir.StencilImplementation) -> gt_ir.StencilImplementation:
+            assert isinstance(node, gt_ir.StencilImplementation)
+            return self.visit(node)
+
+        def visit_StencilImplementation(self, node: gt_ir.StencilImplementation) -> None:
+            self.iir = node
+            for field_name in self.reduced_fields:
+                assert field_name in node.temporary_fields, "Tried to reduce API field to 2D."
+                node.fields[field_name].axes.pop()
+            self.generic_visit(node)
+
+        def visit_FieldRef(self, node: gt_ir.FieldRef) -> None:
+            if node.name in self.reduced_fields:
+                field_decl = self.iir.fields[node.name]
+                node.offset = {axis: node.offset[axis] for axis in field_decl.axes}
+
+    @classmethod
+    def apply(cls, transform_data: TransformData) -> None:
+        reduced_fields = cls.ReducibleFieldsCollector.apply(transform_data.implementation_ir)
+        if len(reduced_fields) > 0:
+            cls.StorageReducer.apply(transform_data.implementation_ir, reduced_fields)
+
+
 class HousekeepingPass(TransformPass):
     class WarnIfNoEffect(gt_ir.IRNodeVisitor):
         """Warn if StencilImplementation has no effect."""
@@ -1317,10 +1405,7 @@ def __call__(self, stencil_name: str, node: gt_ir.StencilImplementation) -> None
         def visit_StencilImplementation(self, node: gt_ir.StencilImplementation):
             # Emit warning if stencil has no effect, i.e. does not read or write to any api fields
             if not node.has_effect:
-                warnings.warn(
-                    f"Stencil `{self.stencil_name}` has no effect.",
-                    RuntimeWarning,
-                )
+                warnings.warn(f"Stencil `{self.stencil_name}` has no effect.", RuntimeWarning)
 
     class PruneEmptyNodes(gt_ir.IRNodeMapper):
         """Removes empty multi-stages, stage groups, and stages."""

src/gt4py/analysis/transformer.py

@@ -32,6 +32,7 @@
     InitInfoPass,
     MergeBlocksPass,
     NormalizeBlocksPass,
+    ReduceTemporaryStoragesPass,
 )
 
 
@@ -111,6 +112,10 @@ def __call__(self, definition_ir, options):
         # into local scalars
         DemoteLocalTemporariesToVariablesPass.apply(self.transform_data)
 
+        # turn temporary fields that are only written and read within the horizontal plane
+        # into 2D i-j fields
+        ReduceTemporaryStoragesPass.apply(self.transform_data)
+
         # prune some stages that don't have effect
         HousekeepingPass.apply(self.transform_data)
 

src/gt4py/backend/debug_backend.py

@@ -56,16 +56,24 @@ def _make_regional_computation(self, iteration_order, interval_definition):
         else:
             range_args = [loop_bounds[1] + " -1", loop_bounds[0] + " -1", "-1"]
 
-        range_expr = "range({args})".format(args=", ".join(a for a in range_args))
-        seq_axis = self.impl_node.domain.sequential_axis.name
-        source_lines.append("for {ax} in {range_expr}:".format(ax=seq_axis, range_expr=range_expr))
+        if range_args != self.range_args:
+            self.range_args = range_args
+            range_expr = "range({args})".format(args=", ".join(a for a in range_args))
+            seq_axis = self.impl_node.domain.sequential_axis.name
+            source_lines.append(
+                "for {ax} in {range_expr}:".format(ax=seq_axis, range_expr=range_expr)
+            )
 
         return source_lines
 
     def make_temporary_field(
-        self, name: str, dtype: gt_ir.DataType, extent: gt_definitions.Extent
+        self,
+        name: str,
+        dtype: gt_ir.DataType,
+        extent: gt_definitions.Extent,
+        axes: list = gt_definitions.CartesianSpace.names,
     ):
-        source_lines = super().make_temporary_field(name, dtype, extent)
+        source_lines = super().make_temporary_field(name, dtype, extent, axes)
         source_lines.extend(self._make_field_accessor(name, extent.to_boundary().lower_indices))
 
         return source_lines
@@ -112,10 +120,13 @@ def make_stage_source(self, iteration_order: gt_ir.IterationOrder, regions: list
     def visit_FieldRef(self, node: gt_ir.FieldRef):
         assert node.name in self.block_info.accessors
         index = []
-        for ax in self.domain.axes_names:
+        for ax in node.offset.keys():
             offset = "{:+d}".format(node.offset[ax]) if ax in node.offset else ""
             index.append("{ax}{offset}".format(ax=ax, offset=offset))
 
+        # Extend index with zeros...
+        index.extend(["0"] * (len(self.block_info.extent) - len(index)))
+
         source = "{name}{marker}[{index}]".format(
             marker=self.origin_marker, name=node.name, index=", ".join(index)
         )

src/gt4py/backend/gt_backends.py

@@ -452,14 +452,23 @@ def visit_StencilImplementation(
         arg_fields = []
         tmp_fields = []
         storage_ids = []
+        used_axes = dict()
+        all_axes = [axis.lower() for axis in gt_definitions.CartesianSpace.names]
+
         max_ndim = 0
         for name, field_decl in node.fields.items():
             if name not in node.unreferenced:
                 max_ndim = max(max_ndim, len(field_decl.axes))
+                axes = "".join(field_decl.axes).lower()
+                selector = ["1" if axis in axes else "0" for axis in all_axes]
+                used_axes[axes] = dict(name=axes.lower(), selector=", ".join(selector))
+
                 field_attributes = {
                     "name": field_decl.name,
                     "dtype": self._make_cpp_type(field_decl.data_type),
+                    "axes": axes,
                 }
+
                 if field_decl.is_api:
                     if field_decl.layout_id not in storage_ids:
                         storage_ids.append(field_decl.layout_id)
@@ -497,6 +506,7 @@ def visit_StencilImplementation(
             stage_functors=stage_functors,
             stencil_unique_name=self.class_name,
             tmp_fields=tmp_fields,
+            used_axes=used_axes.values(),
         )
 
         sources: Dict[str, Dict[str, str]] = {"computation": {}, "bindings": {}}
@@ -543,8 +553,7 @@ def generate(self) -> Type["StencilObject"]:
 
         # Generate and return the Python wrapper class
         return self.make_module(
-            pyext_module_name=pyext_module_name,
-            pyext_file_path=pyext_file_path,
+            pyext_module_name=pyext_module_name, pyext_file_path=pyext_file_path
         )
 
     def generate_computation(self) -> Dict[str, Union[str, Dict]]:

src/gt4py/backend/numpy_backend.py

@@ -88,13 +88,17 @@ def _make_regional_computation(
             range_args = [loop_bounds[1] + " -1", loop_bounds[0] + " -1", "-1"]
 
         if iteration_order != gt_ir.IterationOrder.PARALLEL:
-            range_expr = "range({args})".format(args=", ".join(a for a in range_args))
-            seq_axis = self.impl_node.domain.sequential_axis.name
-            source_lines.append(
-                "for {ax} in {range_expr}:".format(ax=seq_axis, range_expr=range_expr)
-            )
+            if self.range_args != range_args:
+                self.range_args = range_args
+                range_expr = "range({args})".format(args=", ".join(a for a in range_args))
+                seq_axis = self.impl_node.domain.sequential_axis.name
+                source_lines.append(
+                    "for {ax} in {range_expr}:".format(ax=seq_axis, range_expr=range_expr)
+                )
             source_lines.extend(" " * self.indent_size + line for line in body_sources)
         else:
+            # Clear range args on parallel intervals
+            self.range_args.clear()
             source_lines.append(
                 "{interval_k_start_name} = {lb}".format(
                     interval_k_start_name=self.interval_k_start_name, lb=loop_bounds[0]
@@ -110,9 +114,13 @@ def _make_regional_computation(
         return source_lines
 
     def make_temporary_field(
-        self, name: str, dtype: gt_ir.DataType, extent: gt_definitions.Extent
-    ) -> List[str]:
-        source_lines = super().make_temporary_field(name, dtype, extent)
+        self,
+        name: str,
+        dtype: gt_ir.DataType,
+        extent: gt_definitions.Extent,
+        axes: list = gt_definitions.CartesianSpace.names,
+    ):
+        source_lines = super().make_temporary_field(name, dtype, extent, axes)
         source_lines.extend(self._make_field_origin(name, extent.to_boundary().lower_indices))
 
         return source_lines
@@ -138,7 +146,7 @@ def visit_FieldRef(self, node: gt_ir.FieldRef) -> str:
         lower_extent = list(extent.lower_indices)
         upper_extent = list(extent.upper_indices)
 
-        for d, ax in enumerate(self.domain.axes_names):
+        for d, ax in enumerate(node.offset.keys()):
             idx = node.offset.get(ax, 0)
             if idx:
                 lower_extent[d] += idx
@@ -160,28 +168,31 @@ def visit_FieldRef(self, node: gt_ir.FieldRef) -> str:
             )
 
         k_ax = self.domain.sequential_axis.name
-        k_offset = node.offset.get(k_ax, 0)
-        if is_parallel:
-            start_expr = self.interval_k_start_name
-            start_expr += " {:+d}".format(k_offset) if k_offset else ""
-            end_expr = self.interval_k_end_name
-            end_expr += " {:+d}".format(k_offset) if k_offset else ""
-            index.append(
-                "{name}{marker}[2] + {start}:{name}{marker}[2] + {stop}".format(
-                    name=node.name, start=start_expr, marker=self.origin_marker, stop=end_expr
+        if k_ax in node.offset:
+            k_offset = node.offset.get(k_ax, 0)
+            if is_parallel:
+                start_expr = self.interval_k_start_name
+                start_expr += " {:+d}".format(k_offset) if k_offset else ""
+                end_expr = self.interval_k_end_name
+                end_expr += " {:+d}".format(k_offset) if k_offset else ""
+                index.append(
+                    "{name}{marker}[2] + {start}:{name}{marker}[2] + {stop}".format(
+                        name=node.name, start=start_expr, marker=self.origin_marker, stop=end_expr
+                    )
                 )
-            )
-        else:
-            idx = "{:+d}".format(k_offset) if k_offset else ""
-            index.append(
-                "{name}{marker}[{d}] + {ax}{idx}".format(
-                    name=node.name,
-                    marker=self.origin_marker,
-                    d=len(self.domain.parallel_axes),
-                    ax=k_ax,
-                    idx=idx,
+            else:
+                idx = "{:+d}".format(k_offset) if k_offset else ""
+                index.append(
+                    "{name}{marker}[{d}] + {ax}{idx}".format(
+                        name=node.name,
+                        marker=self.origin_marker,
+                        d=len(self.domain.parallel_axes),
+                        ax=k_ax,
+                        idx=idx,
+                    )
                 )
-            )
+        else:
+            index.append("0:1" if is_parallel else "0")
 
         source = "{name}[{index}]".format(name=node.name, index=", ".join(index))
 

src/gt4py/backend/python_generator.py

@@ -78,6 +78,7 @@ def __init__(
         self.param_names = None
 
         self.var_refs_defined = set()
+        self.range_args = []
 
     def __call__(self, impl_node: gt_ir.Node, sources: gt_text.TextBlock):
         assert isinstance(impl_node, gt_ir.StencilImplementation)
@@ -110,7 +111,7 @@ def __call__(self, impl_node: gt_ir.Node, sources: gt_text.TextBlock):
         return self.sources
 
     def make_temporary_field(
-        self, name: str, data_type: gt_ir.DataType, extent: gt_definitions.Extent
+        self, name: str, data_type: gt_ir.DataType, extent: gt_definitions.Extent, axes: list
     ):
         source_lines = []
         boundary = extent.to_boundary()
@@ -119,7 +120,10 @@ def make_temporary_field(
                 domain=self.domain_arg_name, d=d, size=" {:+d}".format(size) if size > 0 else ""
             )
             for d, size in enumerate(boundary.frame_size)
+            if d < len(axes)
         )
+        for i in range(len(extent) - len(axes)):
+            shape += ", 1"
         source_lines.append(
             "{name} = {np_prefix}.empty(({shape}), dtype={np_prefix}.{dtype})".format(
                 name=name, np_prefix=self.numpy_prefix, shape=shape, dtype=data_type.dtype.name
@@ -293,7 +297,7 @@ def visit_StencilImplementation(self, node: gt_ir.StencilImplementation):
                 field = node.fields[name]
                 self.sources.extend(
                     self.make_temporary_field(
-                        field.name, field.data_type, node.fields_extents[field.name]
+                        field.name, field.data_type, node.fields_extents[field.name], field.axes
                     )
                 )
             self.sources.empty_line()