dialects: (snitch_stream) make snitch_stream parameters go outwards-in

tests/filecheck/dialects/snitch_stream/convert_snitch_stream_to_snitch.mlir

@@ -8,8 +8,8 @@
 "snitch_stream.streaming_region"(%A, %B, %C) <{
     "stride_patterns" = [
         #snitch_stream.stride_pattern<ub = [2], strides = [8]>,
-        #snitch_stream.stride_pattern<ub = [2, 3], strides = [24, 8]>,
-        #snitch_stream.stride_pattern<ub = [2, 3, 4, 5], strides = [480, 160, 40, 8]>
+        #snitch_stream.stride_pattern<ub = [3, 2], strides = [8, 24]>,
+        #snitch_stream.stride_pattern<ub = [5, 4, 3, 2], strides = [8, 40, 160, 480]>
     ],
     "operandSegmentSizes" = array<i32: 2, 1>
 }> ({

tests/filecheck/transforms/convert_memref_stream_to_snitch.mlir

@@ -33,7 +33,7 @@ memref_stream.streaming_region {
 // CHECK-NEXT:  %{{.*}} = builtin.unrealized_conversion_cast %A : memref<2xf64> to !riscv.reg<>
 // CHECK-NEXT:  %{{.*}} = builtin.unrealized_conversion_cast %B : memref<3xf64> to !riscv.reg<>
 // CHECK-NEXT:  %{{.*}} = builtin.unrealized_conversion_cast %C : memref<3x2xf64> to !riscv.reg<>
-// CHECK-NEXT:    "snitch_stream.streaming_region"(%2, %3, %4) <{"stride_patterns" = [#snitch_stream.stride_pattern<ub = [2, 3], strides = [0, 8]>, #snitch_stream.stride_pattern<ub = [2, 3], strides = [8, 0]>, #snitch_stream.stride_pattern<ub = [6], strides = [8]>], "operandSegmentSizes" = array<i32: 2, 1>}> ({
+// CHECK-NEXT:    "snitch_stream.streaming_region"(%2, %3, %4) <{"stride_patterns" = [#snitch_stream.stride_pattern<ub = [3, 2], strides = [8, 0]>, #snitch_stream.stride_pattern<ub = [3, 2], strides = [0, 8]>, #snitch_stream.stride_pattern<ub = [6], strides = [8]>], "operandSegmentSizes" = array<i32: 2, 1>}> ({
 // CHECK-NEXT:  ^0(%a : !stream.readable<!riscv.freg<>>, %b : !stream.readable<!riscv.freg<>>, %c : !stream.writable<!riscv.freg<>>):
 // CHECK-NEXT:      %{{.*}} = builtin.unrealized_conversion_cast %a : !stream.readable<!riscv.freg<>> to !stream.readable<f64>
 // CHECK-NEXT:      %{{.*}} = builtin.unrealized_conversion_cast %b : !stream.readable<!riscv.freg<>> to !stream.readable<f64>
@@ -78,7 +78,7 @@ memref_stream.streaming_region {
 
 // CHECK-NEXT:    %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : memref<1x1x8x8xf64> to !riscv.reg<>
 // CHECK-NEXT:    %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : memref<1x1x3x3xf64> to !riscv.reg<>
-// CHECK-NEXT:    "snitch_stream.streaming_region"(%{{.*}}, %{{.*}}) <{"stride_patterns" = [#snitch_stream.stride_pattern<ub = [3, 3, 6, 6], strides = [8, 64, 8, 64]>, #snitch_stream.stride_pattern<ub = [3, 3, 36], strides = [8, 24, 0]>], "operandSegmentSizes" = array<i32: 2, 0>}> ({
+// CHECK-NEXT:    "snitch_stream.streaming_region"(%{{.*}}, %{{.*}}) <{"stride_patterns" = [#snitch_stream.stride_pattern<ub = [6, 6, 3, 3], strides = [64, 8, 64, 8]>, #snitch_stream.stride_pattern<ub = [36, 3, 3], strides = [0, 24, 8]>], "operandSegmentSizes" = array<i32: 2, 0>}> ({
 // CHECK-NEXT:    ^{{.*}}(%{{.*}} : !stream.readable<!riscv.freg<>>, %{{.*}} : !stream.readable<!riscv.freg<>>):
 // CHECK-NEXT:      %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : !stream.readable<!riscv.freg<>> to !stream.readable<f64>
 // CHECK-NEXT:      %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : !stream.readable<!riscv.freg<>> to !stream.readable<f64>
@@ -103,7 +103,7 @@ memref_stream.streaming_region {
 // CHECK-NEXT:    %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : memref<8x8xf64> to !riscv.reg<>
 // CHECK-NEXT:    %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : memref<8x8xf64> to !riscv.reg<>
 // CHECK-NEXT:    %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : memref<8x8xf64> to !riscv.reg<>
-// CHECK-NEXT:    "snitch_stream.streaming_region"(%{{.*}}, %{{.*}}, %{{.*}}) <{"stride_patterns" = [#snitch_stream.stride_pattern<ub = [8, 8, 8], strides = [8, 0, 64]>, #snitch_stream.stride_pattern<ub = [8, 8, 8], strides = [64, 8, 0]>, #snitch_stream.stride_pattern<ub = [64], strides = [8]>], "operandSegmentSizes" = array<i32: 3, 0>}> ({
+// CHECK-NEXT:    "snitch_stream.streaming_region"(%{{.*}}, %{{.*}}, %{{.*}}) <{"stride_patterns" = [#snitch_stream.stride_pattern<ub = [8, 8, 8], strides = [64, 0, 8]>, #snitch_stream.stride_pattern<ub = [8, 8, 8], strides = [0, 8, 64]>, #snitch_stream.stride_pattern<ub = [64], strides = [8]>], "operandSegmentSizes" = array<i32: 3, 0>}> ({
 // CHECK-NEXT:    ^{{.*}}(%{{.*}} : !stream.readable<!riscv.freg<>>, %{{.*}} : !stream.readable<!riscv.freg<>>, %{{.*}} : !stream.readable<!riscv.freg<>>):
 // CHECK-NEXT:      %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : !stream.readable<!riscv.freg<>> to !stream.readable<f64>
 // CHECK-NEXT:      %{{.*}} = builtin.unrealized_conversion_cast %{{.*}} : !stream.readable<!riscv.freg<>> to !stream.readable<f64>

tests/interpreters/test_snitch_stream_interpreter.py

@@ -16,8 +16,8 @@
     [
         ((6,), (1,), tuple(range(6))),
         ((6,), (2,), tuple(range(0, 12, 2))),
-        ((2, 3), (1, 2), tuple(range(6))),
-        ((2, 3, 4), (1, 2, 6), tuple(range(24))),
+        ((3, 2), (2, 1), tuple(range(6))),
+        ((4, 3, 2), (6, 2, 1), tuple(range(24))),
     ],
 )
 def test_stride_pattern_offsets(
@@ -33,11 +33,11 @@ def test_stride_pattern_offsets(
     "inputs, outputs",
     [
         (((24,), (1,)), ((24,), (1,))),
-        (((4, 3, 2), (1, 4, 12)), ((24,), (1,))),
-        (((2, 3), (8, 16)), ((6,), (8,))),
-        (((2, 3), (0, 8)), ((2, 3), (0, 8))),
-        (((2, 3), (8, 0)), ((2, 3), (8, 0))),
-        (((3, 3, 1, 6, 1, 1), (1, 2, 3, 4, 5, 6)), ((3, 3, 6), (1, 2, 4))),
+        (((2, 3, 4), (12, 4, 1)), ((24,), (1,))),
+        (((3, 2), (16, 8)), ((6,), (8,))),
+        (((3, 2), (8, 0)), ((3, 2), (8, 0))),
+        (((3, 2), (0, 8)), ((3, 2), (0, 8))),
+        (((1, 1, 6, 1, 3, 3), (6, 5, 4, 3, 2, 1)), ((6, 3, 3), (4, 2, 1))),
     ],
 )
 def test_simplify_stride_pattern(
@@ -92,7 +92,7 @@ def test_snitch_stream_interpreter():
                 ArrayAttr(
                     (
                         snitch_stream.StridePattern.from_bounds_and_strides(
-                            [3, 2], [8, 24]
+                            [2, 3], [24, 8]
                         ),
                     )
                 ),

xdsl/backend/riscv/lowering/convert_snitch_stream_to_snitch.py

@@ -21,6 +21,9 @@ def insert_stride_pattern_ops(
     strides: builtin.ArrayAttr[builtin.IntAttr],
     dm: builtin.IntAttr,
 ):
+    """
+    `ub` and `strides` must go from the outermost dimension inwards
+    """
     # reference implementation:
     # https://github.com/pulp-platform/snitch/blob/d026f47843f0ea6c269244c4e6851e0e09141ec3/sw/snRuntime/src/ssr.h#L73
     #
@@ -60,8 +63,8 @@ def insert_stride_pattern_ops(
 
     ints = tuple(builtin.IntAttr(i) for i in range(rank))
 
-    b_ops = tuple(riscv.LiOp(b.data) for b in ub)
-    s_ops = tuple(riscv.LiOp(s.data) for s in strides)
+    b_ops = tuple(riscv.LiOp(b.data) for b in reversed(ub.data))
+    s_ops = tuple(riscv.LiOp(s.data) for s in reversed(strides.data))
     new_b_ops = tuple(riscv.AddiOp(b_op.rd, -1) for b_op in b_ops)
     set_bound_ops = tuple(
         snitch.SsrSetDimensionBoundOp(new_b_op, dm, i)

xdsl/dialects/snitch_stream.py

@@ -60,12 +60,12 @@ class StridePattern(ParametrizedAttribute):
 
     ```
     // 2D access pattern
-    #pat = #snitch_stream.stride_pattern<ub = [8, 16], strides = [128, 8]>
+    #pat = #snitch_stream.stride_pattern<ub = [16, 8], strides = [8, 128]>
     // Corresponds to the following locations
     // for i in range(16):
     //   for j in range(8):
     //     yield i * 8 + j * 128
-    // Note that the upper bounds and strides go from the innermost loop outwards
+    // Note that the upper bounds and strides go from the outermost loop inwards
     ```
     """
 
@@ -125,13 +125,11 @@ def verify(self) -> None:
             )
 
     def offset_iter(self) -> Iterator[int]:
-        for indices in product(
-            *(range(bound.data) for bound in reversed(self.ub.data))
-        ):
+        for indices in product(*(range(bound.data) for bound in self.ub.data)):
             indices: tuple[int, ...] = indices
             yield sum(
                 index * stride.data
-                for (index, stride) in zip(indices, reversed(self.strides.data))
+                for (index, stride) in zip(indices, self.strides.data)
             )
 
     def offsets(self) -> tuple[int, ...]:
@@ -145,7 +143,7 @@ def simplified(self) -> StridePattern:
         e.g.
 
         ```
-        stride_pattern<ub = [4, 3, 2], strides = [1, 4, 12]>
+        stride_pattern<ub = [2, 3, 4], strides = [12, 4, 1]>
         ->
         stride_pattern<ub = [24], strides = [1]
         ```
@@ -161,11 +159,11 @@ def simplified(self) -> StridePattern:
         )
 
         # Outermost bound and stride
-        ub0, s0 = tuples[-1]
+        ub0, s0 = tuples[0]
 
         # Start with the second outermost loop bounds
-        second_outermost_dim = len(tuples) - 2
-        while 0 <= second_outermost_dim:
+        second_outermost_dim = 1
+        while second_outermost_dim < len(tuples):
             # Next bound and stride to fold into outermost
             ubd, sd = tuples[second_outermost_dim]
             if s0 == ubd * sd:
@@ -174,15 +172,15 @@ def simplified(self) -> StridePattern:
                 s0 = sd
                 # Decrement the index into tuples for what the new second outermost loop
                 # bound is
-                second_outermost_dim -= 1
+                second_outermost_dim += 1
             else:
                 # The second outermost loop does not match, do not try to further simplify
                 break
 
-        # ub and s include all the tuples up to and including the second outermost dim,
-        # followed by the new outermost bound and stride
-        ub = (*(bound for bound, _ in tuples[: second_outermost_dim + 1]), ub0)
-        s = (*(stride for _, stride in tuples[: second_outermost_dim + 1]), s0)
+        # ub and s include the new outermost bound and stride,
+        # followed by all the tuples up to and including the second outermost dim
+        ub = (ub0, *(bound for bound, _ in tuples[second_outermost_dim:]))
+        s = (s0, *(stride for _, stride in tuples[second_outermost_dim:]))
 
         return StridePattern.from_bounds_and_strides(ub, s)
 

xdsl/transforms/convert_memref_stream_to_snitch_stream.py

@@ -106,11 +106,9 @@ def match_and_rewrite(
             for map, shape in zip(op.indexing_maps, shapes, strict=True)
         )
         bounds = tuple(bound.data for bound in op.bounds)
-        # snitch_stream stride patterns are from innermost out
-        # TODO: make consistent
         stride_patterns = tuple(
             snitch_stream.StridePattern.from_bounds_and_strides(
-                bounds[: len(strides)][::-1], strides[::-1]
+                bounds[: len(strides)], strides
             ).simplified()
             for strides in all_strides
         )