transformations: Add new csl-wrapper-hoist-buffers pass

tests/filecheck/transforms/csl-wrapper-hoist-buffers.mlir

@@ -0,0 +1,39 @@
+// RUN: xdsl-opt %s -p "csl-wrapper-hoist-buffers" | filecheck %s
+
+builtin.module {
+  "csl_wrapper.module"() <{"width" = 1022 : i16, "height" = 510 : i16, "params" = [], "program_name" = "gauss_seidel_func"}> ({
+  ^0(%0 : i16, %1 : i16, %2 : i16, %3 : i16):
+    "csl_wrapper.yield"() <{"fields" = []}> : () -> ()
+  }, {
+  ^1(%4 : i16, %5 : i16):
+    %arg0 = memref.alloc() : memref<512xf32>
+    %6 = "csl.addressof"(%arg0) : (memref<512xf32>) -> !csl.ptr<f32, #csl<ptr_kind many>, #csl<ptr_const var>>
+    "csl.export"(%6) <{"var_name" = "arg0", "type" = !csl.ptr<f32, #csl<ptr_kind many>, #csl<ptr_const var>>}> : (!csl.ptr<f32, #csl<ptr_kind many>, #csl<ptr_const var>>) -> ()
+    "csl.export"() <{"var_name" = @gauss_seidel_func, "type" = () -> ()}> : () -> ()
+    csl.func @gauss_seidel_func() {
+      %7 = memref.alloc() {"alignment" = 64 : i64} : memref<510xf32>
+      "test.op"(%7) : (memref<510xf32>) -> ()
+      csl.return
+    }
+    "csl_wrapper.yield"() <{"fields" = []}> : () -> ()
+  }) : () -> ()
+}
+
+// CHECK-NEXT: builtin.module {
+// CHECK-NEXT:   "csl_wrapper.module"() <{"width" = 1022 : i16, "height" = 510 : i16, "params" = [], "program_name" = "gauss_seidel_func"}> ({
+// CHECK-NEXT:   ^0(%0 : i16, %1 : i16, %2 : i16, %3 : i16):
+// CHECK-NEXT:     "csl_wrapper.yield"() <{"fields" = []}> : () -> ()
+// CHECK-NEXT:   }, {
+// CHECK-NEXT:   ^1(%4 : i16, %5 : i16):
+// CHECK-NEXT:     %6 = memref.alloc() {"alignment" = 64 : i64} : memref<510xf32>
+// CHECK-NEXT:     %arg0 = memref.alloc() : memref<512xf32>
+// CHECK-NEXT:     %7 = "csl.addressof"(%arg0) : (memref<512xf32>) -> !csl.ptr<f32, #csl<ptr_kind many>, #csl<ptr_const var>>
+// CHECK-NEXT:     "csl.export"(%7) <{"var_name" = "arg0", "type" = !csl.ptr<f32, #csl<ptr_kind many>, #csl<ptr_const var>>}> : (!csl.ptr<f32, #csl<ptr_kind many>, #csl<ptr_const var>>) -> ()
+// CHECK-NEXT:     "csl.export"() <{"var_name" = @gauss_seidel_func, "type" = () -> ()}> : () -> ()
+// CHECK-NEXT:     csl.func @gauss_seidel_func() {
+// CHECK-NEXT:       "test.op"(%6) : (memref<510xf32>) -> ()
+// CHECK-NEXT:       csl.return
+// CHECK-NEXT:     }
+// CHECK-NEXT:     "csl_wrapper.yield"() <{"fields" = []}> : () -> ()
+// CHECK-NEXT:   }) : () -> ()
+// CHECK-NEXT: }

xdsl/tools/command_line_tool.py

@@ -111,6 +111,11 @@ def get_csl_wrapper_to_csl():
 
         return csl_wrapper_to_csl.CslWrapperToCslPass
 
+    def get_csl_wrapper_hoist_buffers():
+        from xdsl.transforms import csl_wrapper_hoist_buffers
+
+        return csl_wrapper_hoist_buffers.CslWrapperHoistBuffers
+
     def get_dce():
         from xdsl.transforms import dead_code_elimination
 
@@ -426,6 +431,7 @@ def get_stencil_bufferize():
         "csl-stencil-bufferize": get_csl_stencil_bufferize,
         "csl-stencil-to-csl-wrapper": get_csl_stencil_to_csl_wrapper,
         "csl-wrapper-to-csl": get_csl_wrapper_to_csl,
+        "csl-wrapper-hoist-buffers": get_csl_wrapper_hoist_buffers,
         "dce": get_dce,
         "distribute-stencil": get_distribute_stencil,
         "dmp-to-mpi": get_lower_halo_to_mpi,

xdsl/transforms/csl_wrapper_hoist_buffers.py

@@ -0,0 +1,52 @@
+from dataclasses import dataclass
+
+from xdsl.context import MLContext
+from xdsl.dialects import memref
+from xdsl.dialects.builtin import ModuleOp
+from xdsl.dialects.csl import csl_wrapper
+from xdsl.passes import ModulePass
+from xdsl.pattern_rewriter import (
+    PatternRewriter,
+    PatternRewriteWalker,
+    RewritePattern,
+    op_type_rewrite_pattern,
+)
+from xdsl.rewriter import InsertPoint
+
+
+@dataclass(frozen=True)
+class HoistBuffers(RewritePattern):
+    """
+    Hoists buffers to `csl_wrapper.program_module`-level.
+    """
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: memref.Alloc, rewriter: PatternRewriter, /):
+        wrapper = op.parent_op()
+        while wrapper and not isinstance(wrapper, csl_wrapper.ModuleOp):
+            wrapper = wrapper.parent_op()
+
+        # no action required if this op exists on module-level
+        if not wrapper or wrapper == op.parent_op():
+            return
+
+        assert len(op.dynamic_sizes) == 0, "not implemented"
+        assert len(op.symbol_operands) == 0, "not implemented"
+
+        rewriter.insert_op(
+            alloc := op.clone(), InsertPoint.at_start(wrapper.program_module.block)
+        )
+        rewriter.replace_matched_op([], new_results=[alloc.memref])
+
+
+@dataclass(frozen=True)
+class CslWrapperHoistBuffers(ModulePass):
+    """
+    Hoists buffers to the `csl_wrapper.program_module`-level.
+    """
+
+    name = "csl-wrapper-hoist-buffers"
+
+    def apply(self, ctx: MLContext, op: ModuleOp) -> None:
+        module_pass = PatternRewriteWalker(HoistBuffers())
+        module_pass.rewrite_module(op)