transforms: (memref-to-dsd) Support 1d subview of nd memref

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

@@ -98,27 +98,17 @@ builtin.module {
 // CHECK-NEXT: %26 = "test.op"() : () -> !csl<dsd mem1d_dsd>
 // CHECK-NEXT: "csl.fadds"(%26, %26, %26) : (!csl<dsd mem1d_dsd>, !csl<dsd mem1d_dsd>, !csl<dsd mem1d_dsd>) -> ()
 
-%33 = "csl.variable"() : () -> !csl.var<memref<512xf32>>
-%34 = "csl.load_var"(%33) : (!csl.var<memref<512xf32>>) -> memref<512xf32>
-"csl.store_var"(%33, %34) : (!csl.var<memref<512xf32>>, memref<512xf32>) -> ()
-
-// CHECK-NEXT: %27 = "csl.variable"() : () -> !csl.var<!csl<dsd mem1d_dsd>>
-// CHECK-NEXT: %28 = "csl.load_var"(%27) : (!csl.var<!csl<dsd mem1d_dsd>>) -> !csl<dsd mem1d_dsd>
-// CHECK-NEXT: "csl.store_var"(%27, %28) : (!csl.var<!csl<dsd mem1d_dsd>>, !csl<dsd mem1d_dsd>) -> ()
-
-// ensure that pre-existing get_mem_dsd ops access the underlying buffer, not the get_mem_dsd created on top of it
-
-%36 = arith.constant 510 : i16
-%37 = "csl.get_mem_dsd"(%b, %36) : (memref<510xf32>, i16) -> !csl<dsd mem1d_dsd>
-
-// 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) -> ()
+%39 = memref.alloc() {"alignment" = 64 : i64} : memref<3x64xf32>
+%40 = "memref.subview"(%39, %0) <{"operandSegmentSizes" = array<i32: 1, 1, 0, 0>, "static_offsets" = array<i64: 0, -9223372036854775808>, "static_sizes" = array<i64: 1, 32>, "static_strides" = array<i64: 1, 1>}> : (memref<3x64xf32>, index) -> memref<32xf32, strided<[1], offset: ?>>
+
+// CHECK-NEXT: %27 = "csl.zeros"() : () -> memref<3x64xf32>
+// CHECK-NEXT: %28 = arith.constant 3 : i16
+// CHECK-NEXT: %29 = arith.constant 64 : i16
+// CHECK-NEXT: %30 = "csl.get_mem_dsd"(%27, %28, %29) : (memref<3x64xf32>, i16, i16) -> !csl<dsd mem4d_dsd>
+// CHECK-NEXT: %31 = arith.constant 32 : i16
+// CHECK-NEXT: %32 = "csl.get_mem_dsd"(%27, %31) <{"offsets" = [0 : i16, -9223372036854775808 : i64]}> : (memref<3x64xf32>, i16) -> !csl<dsd mem1d_dsd>
+// CHECK-NEXT: %33 = arith.index_cast %0 : index to si16
+// CHECK-NEXT: %34 = "csl.increment_dsd_offset"(%32, %33) <{"elem_type" = f32}> : (!csl<dsd mem1d_dsd>, si16) -> !csl<dsd mem1d_dsd>
 
 }) {sym_name = "program"} :  () -> ()
 }

xdsl/dialects/csl/csl.py

@@ -1166,9 +1166,9 @@ def verify_(self) -> None:
             raise VerifyException(
                 "DSD of type mem4d_dsd must have between 1 and 4 dimensions"
             )
-        if self.offsets is not None and len(self.offsets) != len(self.sizes):
+        if self.offsets is not None and len(self.offsets) < len(self.sizes):
             raise VerifyException(
-                "Dimensions of offsets must match dimensions of sizes"
+                "Dimensions of offsets must match (or exceed) dimensions of sizes"
             )
         if self.strides is not None and len(self.strides) != len(self.sizes):
             raise VerifyException(

xdsl/transforms/memref_to_dsd.py

@@ -1,3 +1,4 @@
+import collections
 from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import cast
@@ -120,6 +121,35 @@
     @op_type_rewrite_pattern
     def match_and_rewrite(self, op: memref.SubviewOp, rewriter: PatternRewriter, /):
         assert isa(op.source.type, MemRefType[Attribute])
+        assert isa(op.result.type, MemRefType[Attribute])
+
+        if len(op.result.type.get_shape()) == 1 and len(op.source.type.get_shape()) > 1:
+            # 1d subview onto a nd memref
+            sizes = op.static_sizes.get_values()
+            scounts = collections.Counter(sizes)
+            if 1 in scounts:
+                scounts.pop(1)
+            assert (
+                len(scounts) == 1
+            ), "1d access into nd memref must specify one size > 1"
+            size, counts = scounts.most_common()[0]
+            assert (
+                counts == 1
+            ), "1d access into nd memref can only specify one size > 1, which can occur only once"
+            size_op = arith.ConstantOp.from_int_and_width(size, 16)
+            offsets = [
+                IntegerAttr(o, 16 if o != memref.SubviewOp.DYNAMIC_INDEX else 64)
+                for o in op.static_offsets.get_values()
+            ]
+            dsd_op = csl.GetMemDsdOp(
+                operands=[op.source, [size_op]],
+                properties={"offsets": ArrayAttr(offsets)},
+                result_types=[csl.DsdType(csl.DsdKind.mem1d_dsd)],
+            )
+            offset_ops = self._update_offsets(op, dsd_op) if op.offsets else []
+            rewriter.replace_matched_op([size_op, dsd_op, *offset_ops])
+            return
+
         assert len(op.static_sizes) == 1, "not implemented"
         assert len(op.static_offsets) == 1, "not implemented"
         assert len(op.static_strides) == 1, "not implemented"
@@ -214,7 +244,7 @@
 
         static_offsets = cast(Sequence[int], subview.static_offsets.get_values())
 
-        if static_offsets[0] == memref.SubviewOp.DYNAMIC_INDEX:
+        if subview.offsets:
             ops.append(cast_op := arith.IndexCastOp(subview.offsets[0], csl.i16_value))
             ops.append(
                 csl.IncrementDsdOffsetOp.build(