GridTools Stage Extents

src/gt4py/backend/gt_backends.py

@@ -470,10 +470,18 @@ def visit_StencilImplementation(self, node: gt_ir.StencilImplementation) -> Dict
             if name not in node.unreferenced
         ]
 
+        stage_extents = {}
         stage_functors = {}
         for multi_stage in node.multi_stages:
             for group in multi_stage.groups:
                 for stage in group.stages:
+                    compute_extent = stage.compute_extent
+                    extents: List[int] = []
+                    for i in range(compute_extent.ndims - 1):
+                        extents.extend(
+                            (compute_extent.lower_indices[i], compute_extent.upper_indices[i])
+                        )
+                    stage_extents[stage.name] = ", ".join([str(extent) for extent in extents])
                     stage_functors[stage.name] = self.visit(stage)
 
         multi_stages = []
@@ -491,6 +499,7 @@ def visit_StencilImplementation(self, node: gt_ir.StencilImplementation) -> Dict
             multi_stages=multi_stages,
             parameters=parameters,
             stage_functors=stage_functors,
+            stage_extents=stage_extents,
             stencil_unique_name=self.class_name,
             tmp_fields=tmp_fields,
         )
@@ -653,7 +662,7 @@ def generate_post_run(self) -> str:
         output_field_names = [
             name
             for name, info in self.args_data["field_info"].items()
-            if info.access == gt_definitions.AccessKind.READ_WRITE
+            if info and info.access == gt_definitions.AccessKind.READ_WRITE
         ]
 
         return "\n".join([f + "._set_device_modified()" for f in output_field_names])

src/gt4py/backend/templates/computation.src.in

@@ -42,6 +42,7 @@
             ]
         }
     }
+    - stage_extents: { str: str }
     - stencil_unique_name: str
     - tmp_fields: [{ "name": str, "dtype": str }]
 #}
@@ -259,7 +260,7 @@ void run(const std::array<gt::uint_t, 3>& domain,
         {%- for stage in step %}
             {%- filter indent(width=extra_indent) %}
             {{- stage_comma() }}
-            gt::make_stage<{{ stage }}_func>(
+            gt::make_stage_with_extent<{{ stage }}_func, gt::extent<{{ stage_extents[stage] }}>>(
                 p_{{ stage_functors[stage].args|map(attribute="name")|join("(), p_")}}()
             )
             {%- endfilter %}

tests/test_integration/stencil_definitions.py

@@ -236,3 +236,10 @@ def set_inner_as_kord(a4_1: Field3D, a4_2: Field3D, a4_3: Field3D, extm: Field3D
             a4_3 = a4_1
         else:
             diff_23 = a4_2 - a4_3
+
+
+@register
+def write_after_read(field: Field3D):
+    with computation(PARALLEL), interval(...):
+        tmp = field
+        field = tmp[-1, 0, 0]