transformations: (memref-to-dsd) Support memref.load ops (#3338)

This pass translates memref allocs to csl array allocs with a get_dsd on
top of it, which will be used by various compute ops. This PR add
support for memref.load ops, which should continue to load from the
memref, not from a get_dsd op.

---------

Co-authored-by: n-io <n-io@users.noreply.github.com>

Author: n-io

tests/filecheck/transforms/memref-to-dsd.mlir

@@ -114,6 +114,12 @@ builtin.module {
 // CHECK-NEXT: %29 = arith.constant 510 : i16
 // CHECK-NEXT: %30 = "csl.get_mem_dsd"(%b, %29) : (memref<510xf32>, i16) -> !csl<dsd mem1d_dsd>
 
+%38 = memref.load %b[%28] : memref<510xf32>
+"test.op"(%38) : (f32) -> ()
+
+// CHECK-NEXT: %31 = memref.load %b[%13] : memref<510xf32>
+// CHECK-NEXT: "test.op"(%31) : (f32) -> ()
+
 }) {sym_name = "program"} :  () -> ()
 }
 // CHECK-NEXT: }) {"sym_name" = "program"} :  () -> ()

xdsl/transforms/memref_to_dsd.py

@@ -89,6 +89,24 @@ def match_and_rewrite(self, op: csl.GetMemDsdOp, rewriter: PatternRewriter, /):
                 raise ValueError("Failed to resolve GetMemDsdOp called on dsd type")
 
 
+class FixMemrefLoadOnGetDsd(RewritePattern):
+    """
+    Memref load ops should load from the underlying memref, not from the dsd.
+    """
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: memref.Load, rewriter: PatternRewriter, /):
+        if isinstance(op.memref.type, csl.DsdType):
+            if isinstance(op.memref, OpResult) and isinstance(
+                op.memref.op, csl.GetMemDsdOp
+            ):
+                rewriter.replace_matched_op(
+                    memref.Load.get(op.memref.op.base_addr, op.indices)
+                )
+            else:
+                raise ValueError("Failed to resolve memref.load called on dsd type")
+
+
 class LowerSubviewOpPass(RewritePattern):
     """Lowers memref.subview to dsd ops"""
 
@@ -355,12 +373,10 @@ def apply(self, ctx: MLContext, op: ModuleOp) -> None:
                     LowerAllocOpPass(),
                     DsdOpUpdateType(),
                     RetainAddressOfOpPass(),
+                    FixMemrefLoadOnGetDsd(),
+                    FixGetDsdOnGetDsd(),
                 ]
             ),
             apply_recursively=False,
         )
         forward_pass.rewrite_module(op)
-        cleanup_pass = PatternRewriteWalker(
-            FixGetDsdOnGetDsd(),
-        )
-        cleanup_pass.rewrite_module(op)