Fix inline_window to correctly handle stride dimension

src/exo/LoopIR_scheduling.py

@@ -1074,6 +1074,23 @@ def calc_idx(self, idxs):
 
         return new_idxs
 
+    def calc_dim(self, dim):
+        assert dim < len(
+            [w for w in self.win_stmt.rhs.idx if isinstance(w, LoopIR.Interval)]
+        )
+
+        # Because our goal here is to offset `dim` in the original
+        # call argument to the point indexing to the windowing expression,
+        # new_dim should essencially be:
+        # `dim` + "number of LoopIR.Points in the windowing expression before the `dim` number of LoopIR.Interval"
+        new_dim = 0
+        for w in self.win_stmt.rhs.idx:
+            if isinstance(w, LoopIR.Interval):
+                dim -= 1
+            if dim == -1:
+                return new_dim
+            new_dim += 1
+
     def map_s(self, sc):
         s = sc._node()
         if s is self.win_stmt:
@@ -1150,7 +1167,7 @@ def map_e(self, e):
         elif etyp is LoopIR.StrideExpr:
             if self.win_stmt.lhs == e.name:
                 return LoopIR.StrideExpr(
-                    self.win_stmt.rhs.name, e.dim, e.type, e.srcinfo
+                    self.win_stmt.rhs.name, self.calc_dim(e.dim), e.type, e.srcinfo
                 )
 
         return super().map_e(e)

src/exo/platforms/gemmini.py

@@ -544,15 +544,15 @@ def make_ld_i8_block(name, ld_id, p=ld_i8_block, stride_val=None):
     "ld_i8_block_id1_v2", 1, stride_val="stride(src,0)"
 )
 ld_i8_block_id1_s2_v2 = make_ld_i8_block(
-    "ld_i8_block_id1_s2_v2", 1, stride_val="stride(src,2)"
+    "ld_i8_block_id1_s2_v2", 1, stride_val="stride(src,0)"
 )
 
 ld_i8_block_id2 = make_ld_i8_block("ld_i8_block_id2", 2)
 ld_i8_block_id2_v2 = make_ld_i8_block(
     "ld_i8_block_id2_v2", 2, stride_val="stride(src,0)"
 )
 ld_i8_block_id2_s2_v2 = make_ld_i8_block(
-    "ld_i8_block_id2_s2_v2", 2, stride_val="stride(src,2)"
+    "ld_i8_block_id2_s2_v2", 2, stride_val="stride(src,0)"
 )
 
 
@@ -626,11 +626,11 @@ def make_ld_i8(name, ld_id=0, p=ld_i8, sval=None):
 
 ld_i8_id1 = make_ld_i8("ld_i8_id1", 1)
 ld_i8_id1_v2 = make_ld_i8("ld_i8_id1_v2", 1, sval="stride(src, 0)")
-ld_i8_id1_s2_v2 = make_ld_i8("ld_i8_id1_s2_v2", 1, sval="stride(src, 2)")
+ld_i8_id1_s2_v2 = make_ld_i8("ld_i8_id1_s2_v2", 1, sval="stride(src, 0)")
 
 ld_i8_id2 = make_ld_i8("ld_i8_id2", 2)
 ld_i8_id2_v2 = make_ld_i8("ld_i8_id2_v2", 2, sval="stride(src, 0)")
-ld_i8_id2_s2_v2 = make_ld_i8("ld_i8_id2_s2_v2", 2, sval="stride(src, 2)")
+ld_i8_id2_s2_v2 = make_ld_i8("ld_i8_id2_s2_v2", 2, sval="stride(src, 0)")
 
 ld_i8_v2 = make_ld_i8("ld_i8_v2", 0, sval="stride(src,0)")
 
@@ -1025,7 +1025,7 @@ def make_st_acc_i8_s2_v2(p=st_acc_i8):
         p.find("ConfigStore.scale = _").after(),
         ConfigStore,
         "dst_stride",
-        "stride(dst, 2)",
+        "stride(dst, 0)",
     )
     # p = p.configwrite_after('ConfigStore.scale = _', ConfigStore, 'dst_stride', 'stride(dst, 2)')
     p = bind_config(p, "act", ConfigStore, "act")

tests/golden/test_schedules/test_inline_window2.txt

@@ -0,0 +1,3 @@
+def foo(n: size, m: size, k: size, x: R[n, m, k, 10] @ DRAM):
+    x[0, 0, 0, 0] = 0.0
+    bar(stride(x, 2))

tests/test_schedules.py

@@ -1079,6 +1079,22 @@ def foo(n: size, m: size, x: R[n, m]):
     assert str(foo) == golden
 
 
+def test_inline_window2(golden):
+    @proc
+    def bar(s: stride):
+        x: R
+        x = 0.0
+
+    @proc
+    def foo(n: size, m: size, k: size, x: R[n, m, k, 10]):
+        y = x[0, :, :, 0]
+        y[0, 0] = 0.0
+        bar(stride(y, 1))
+
+    foo = inline_window(foo, "y = _")
+    assert str(simplify(foo)) == golden
+
+
 def test_lift_if_second_statement_in_then_error():
     @proc
     def foo(m: size, x: R[m]):