transformations: New stencil-shape-minimize pass

tests/filecheck/transforms/stencil-shape-minimize.mlir

@@ -0,0 +1,27 @@
+// RUN: xdsl-opt -p stencil-shape-minimize --split-input-file %s | filecheck %s
+
+builtin.module {
+  func.func @different_input_offsets(%out : !stencil.field<[-4,68]xf64>, %in : !stencil.field<[-4,68]xf64>) {
+    %tin = stencil.load %in : !stencil.field<[-4,68]xf64> -> !stencil.temp<[-1,65]xf64>
+    %tout = stencil.apply(%arg = %tin : !stencil.temp<[-1,65]xf64>) -> (!stencil.temp<[0,64]xf64>) {
+      %x = stencil.access %arg[-1] : !stencil.temp<[-1,65]xf64>
+      %y = stencil.access %arg[1] : !stencil.temp<[-1,65]xf64>
+      %o = arith.addf %x, %y : f64
+      stencil.return %o : f64
+    }
+    stencil.store %tout to %out(<[0], [64]>) : !stencil.temp<[0,64]xf64> to !stencil.field<[-4,68]xf64>
+    func.return
+  }
+}
+
+// CHECK:       func.func @different_input_offsets(%out : !stencil.field<[-1,65]xf64>, %in : !stencil.field<[-1,65]xf64>) {
+// CHECK-NEXT:    %tin = stencil.load %in : !stencil.field<[-1,65]xf64> -> !stencil.temp<[-1,65]xf64>
+// CHECK-NEXT:    %tout = stencil.apply(%arg = %tin : !stencil.temp<[-1,65]xf64>) -> (!stencil.temp<[0,64]xf64>) {
+// CHECK-NEXT:      %x = stencil.access %arg[-1] : !stencil.temp<[-1,65]xf64>
+// CHECK-NEXT:      %y = stencil.access %arg[1] : !stencil.temp<[-1,65]xf64>
+// CHECK-NEXT:      %o = arith.addf %x, %y : f64
+// CHECK-NEXT:      stencil.return %o : f64
+// CHECK-NEXT:    }
+// CHECK-NEXT:    stencil.store %tout to %out(<[0], [64]>) : !stencil.temp<[0,64]xf64> to !stencil.field<[-1,65]xf64>
+// CHECK-NEXT:    func.return
+// CHECK-NEXT:  }

xdsl/tools/command_line_tool.py

@@ -415,6 +415,11 @@ def get_stencil_bufferize():
 
         return stencil_bufferize.StencilBufferize
 
+    def get_stencil_shape_minimize():
+        from xdsl.transforms import stencil_shape_minimize
+
+        return stencil_shape_minimize.StencilShapeMinimize
+
     return {
         "arith-add-fastmath": get_arith_add_fastmath,
         "loop-hoist-memref": get_loop_hoist_memref,
@@ -494,6 +499,7 @@ def get_stencil_bufferize():
         "stencil-tensorize-z-dimension": get_stencil_tensorize_z_dimension,
         "stencil-unroll": get_stencil_unroll,
         "stencil-bufferize": get_stencil_bufferize,
+        "stencil-shape-minimize": get_stencil_shape_minimize,
         "test-lower-linalg-to-snitch": get_test_lower_linalg_to_snitch,
     }
 

xdsl/transforms/stencil_shape_minimize.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass, field
+
+from xdsl.context import MLContext
+from xdsl.dialects import builtin, func, stencil
+from xdsl.ir import Attribute
+from xdsl.passes import ModulePass
+from xdsl.pattern_rewriter import (
+    GreedyRewritePatternApplier,
+    PatternRewriter,
+    PatternRewriteWalker,
+    RewritePattern,
+    TypeConversionPattern,
+    attr_type_rewrite_pattern,
+    op_type_rewrite_pattern,
+)
+
+
+@dataclass
+class ShapeAnalysis(TypeConversionPattern):
+    seen: set[stencil.TempType[Attribute]] = field(default_factory=set)
+
+    @attr_type_rewrite_pattern
+    def convert_type(self, typ: stencil.TempType[Attribute], /) -> Attribute | None:
+        self.seen.add(typ)
+
+
+@dataclass
+class ShapeMinimisation(TypeConversionPattern):
+    shape: stencil.StencilBoundsAttr | None = None
+
+    @attr_type_rewrite_pattern
+    def convert_type(self, typ: stencil.FieldType[Attribute], /) -> Attribute | None:
+        if typ.bounds != self.shape and self.shape:
+            return stencil.FieldType(self.shape, typ.element_type)
+
+
+@dataclass(frozen=True)
+class FuncOpShapeUpdate(RewritePattern):
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /):
+        op.update_function_type()
+
+
+@dataclass(frozen=True)
+class StencilShapeMinimize(ModulePass):
+    """
+    Minimises the shapes of `stencil.field` types that have been over-allocated and are larger than necessary.
+    """
+
+    name = "stencil-shape-minimize"
+
+    def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
+        analysis = ShapeAnalysis(seen=set())
+        PatternRewriteWalker(analysis).rewrite_module(op)
+        bounds = set(
+            t.bounds
+            for t in analysis.seen
+            if isinstance(t.bounds, stencil.StencilBoundsAttr)
+        )
+        if not bounds:
+            return
+        shape: stencil.StencilBoundsAttr = bounds.pop()
+        for b in bounds:
+            shape = shape | b
+
+        PatternRewriteWalker(
+            GreedyRewritePatternApplier(
+                [
+                    ShapeMinimisation(shape=shape),
+                    FuncOpShapeUpdate(),
+                ]
+            )
+        ).rewrite_module(op)