Repository: exo-lang/exo

Issue:
partition_loop() should work on loops with a variable upper bound
Title. Right now, it only works on loops that have a constant upper bound such as 
`for i in par(0, 4)`

It would assist with the implementation of TRSM and TRMM if it could be changed to work on loops like
`for i in par(0, N)`

Discussion:
alexreinking: This will conflict with the changes I'm making to convert the primitive directives to use cursors.

Patch:
diff --git a/src/exo/API_scheduling.py b/src/exo/API_scheduling.py
index 9217447..9f4668d 100644
--- a/src/exo/API_scheduling.py
+++ b/src/exo/API_scheduling.py
@@ -1540,7 +1540,9 @@ def cut_loop(proc, loop_cursor, cut_point):
     """
     stmt = loop_cursor._impl
     proc_c = ic.Cursor.root(proc)
-    loopir = scheduling.DoPartitionLoop(stmt, cut_point).apply_proc(proc_c._node())
+    loopir = scheduling.DoPartitionLoop(proc_c._node(), stmt, cut_point).apply_proc(
+        proc_c._node()
+    )
 
     return Procedure(loopir, _provenance_eq_Procedure=proc)
 
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index bf10837..c093e6c 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -247,35 +247,42 @@ def DoReorderStmt(f_cursor, s_cursor):
 
 
 class DoPartitionLoop(LoopIR_Rewrite):
-    def __init__(self, loop_cursor, num):
+    def __init__(self, proc, loop_cursor, num):
+        assert num > 0
         self.stmt = loop_cursor._node()
         self.partition_by = num
+        self.proc = proc
 
     def map_s(self, s):
         if s is self.stmt:
             assert isinstance(s, LoopIR.Seq)
 
-            if not isinstance(s.hi, LoopIR.Const):
-                raise SchedulingError("expected loop bound to be a literal")
-
-            if s.hi.val <= self.partition_by:
+            part_by = LoopIR.Const(self.partition_by, T.int, s.srcinfo)
+            new_hi = LoopIR.BinOp("-", s.hi, part_by, T.int, s.srcinfo)
+            try:
+                Check_IsPositiveExpr(
+                    self.proc,
+                    [s],
+                    LoopIR.BinOp(
+                        "+", new_hi, LoopIR.Const(1, T.int, s.srcinfo), T.int, s.srcinfo
+                    ),
+                )
+            except SchedulingError:
                 raise SchedulingError(
-                    "expected loop bound to be larger than partitioning value"
+                    f"expected the new loop bound {new_hi} to be always non-negative"
                 )
 
-            part_by = LoopIR.Const(self.partition_by, T.int, s.srcinfo)
             loop1 = s.update(hi=part_by, eff=None)
 
             # all uses of the loop iteration in the second body need
             # to be offset by the partition value
             iter2 = s.iter.copy()
-            hi2 = LoopIR.Const(s.hi.val - part_by.val, T.int, s.srcinfo)
             iter2_node = LoopIR.Read(iter2, [], T.index, s.srcinfo)
             iter_off = LoopIR.BinOp("+", iter2_node, part_by, T.index, s.srcinfo)
             env = {s.iter: iter_off}
 
             body2 = SubstArgs(s.body, env).result()
-            loop2 = s.update(iter=iter2, hi=hi2, body=body2, eff=None)
+            loop2 = s.update(iter=iter2, hi=new_hi, body=body2, eff=None)
 
             return [loop1, loop2]
 


Test patch:
diff --git a/tests/golden/test_schedules/test_cut_loop2.txt b/tests/golden/test_schedules/test_cut_loop2.txt
new file mode 100644
index 0000000..ac2fddb
--- /dev/null
+++ b/tests/golden/test_schedules/test_cut_loop2.txt
@@ -0,0 +1,8 @@
+def foo(n: size):
+    assert n > 3
+    for i in seq(0, 3):
+        x: R @ DRAM
+        x = 0.0
+    for i in seq(0, -3 + n):
+        x: R @ DRAM
+        x = 0.0
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_syrk_cut_loop.txt b/tests/golden/test_schedules/test_syrk_cut_loop.txt
new file mode 100644
index 0000000..46f7192
--- /dev/null
+++ b/tests/golden/test_schedules/test_syrk_cut_loop.txt
@@ -0,0 +1,15 @@
+def SYRK(M: size, K: size, A: [f32][M, K] @ DRAM, A_t: [f32][M, K] @ DRAM,
+         C: [f32][M, M] @ DRAM):
+    assert M >= 1
+    assert K >= 1
+    assert stride(A, 1) == 1
+    assert stride(A_t, 1) == 1
+    assert stride(C, 1) == 1
+    for io in seq(0, M / 4):
+        for ii in seq(0, 4):
+            for j in seq(0, 1):
+                for k in seq(0, K):
+                    C[ii + 4 * io, j] += A[ii + 4 * io, k] * A_t[j, k]
+            for j in seq(0, ii + 4 * io):
+                for k in seq(0, K):
+                    C[ii + 4 * io, 1 + j] += A[ii + 4 * io, k] * A_t[1 + j, k]
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index a53166e..35c43b5 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -1756,3 +1756,61 @@ def test_formatted_expr_errors_2():
         expand_dim(
             bar, "tmp : _", "n", FormattedExprStr("_", i_type_R_read_cursor)
         )  # should be error
+
+
+def test_syrk_cut_loop(golden):
+    @proc
+    def SYRK(
+        M: size,
+        K: size,
+        A: [f32][M, K] @ DRAM,
+        A_t: [f32][M, K] @ DRAM,
+        C: [f32][M, M] @ DRAM,
+    ):
+        assert M >= 1
+        assert K >= 1
+        assert stride(A, 1) == 1
+        assert stride(A_t, 1) == 1
+        assert stride(C, 1) == 1
+        for io in seq(0, M / 4):
+            for ii in seq(0, 4):
+                for j in seq(0, 4 * io + ii + 1):
+                    for k in seq(0, K):
+                        C[4 * io + ii, j] += A[4 * io + ii, k] * A_t[j, k]
+
+    SYRK = cut_loop(SYRK, "for j in _:_", 1)
+    assert str(simplify(SYRK)) == golden
+
+
+def test_cut_loop1():
+    @proc
+    def foo(n: size):
+        for i in seq(0, n):
+            x: R
+            x = 0.0
+
+    with pytest.raises(SchedulingError, match="expected the new loop bound"):
+        foo = cut_loop(foo, "for i in _:_", 3)
+
+
+def test_cut_loop2(golden):
+    @proc
+    def foo(n: size):
+        assert n > 3
+        for i in seq(0, n):
+            x: R
+            x = 0.0
+
+    foo = cut_loop(foo, "for i in _:_", 3)
+    assert str(simplify(foo)) == golden
+
+
+def test_cut_loop3():
+    @proc
+    def foo(n: size):
+        for i in seq(0, n):
+            x: R
+            x = 0.0
+
+    with pytest.raises(TypeError, match="cut_loop: expected a positive integer"):
+        foo = cut_loop(foo, "for i in _:_", -3)