Repository: exo-lang/exo

Issue:
reorder_loop() issue: errors out when preceded with divide_loop()
```
ret_val = rename(gemm, "gemm_custom")
ret_val = divide_loop(ret_val, 'j', 32, ['jo', 'ji'], perfect=True)
ret_val = divide_loop(ret_val, 'i', 32, ['io', 'ii'], perfect=True)
ret_val = reorder_loops(ret_val, 'ii jo')
print(ret_val)
```

Discussion:

Patch:
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index 94ed869..3459d79 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -458,8 +458,8 @@ def DoSplit(loop_cursor, quot, hi, lo, tail="guard", perfect=False):
         return op("+", op("*", cnst(quot), rd(hi_i)), rd(lo_i))
 
     # short-hands for sanity
-    def boolop(op, lhs, rhs):
-        return LoopIR.BinOp(op, lhs, rhs, T.bool, srcinfo)
+    def boolop(op, lhs, rhs, typ):
+        return LoopIR.BinOp(op, lhs, rhs, typ, srcinfo)
 
     def szop(op, lhs, rhs):
         return LoopIR.BinOp(op, lhs, rhs, lhs.type, srcinfo)
@@ -484,7 +484,7 @@ def DoSplit(loop_cursor, quot, hi, lo, tail="guard", perfect=False):
         idx_sub = substitute(srcinfo)
 
         def guard_wrapper(body):
-            cond = boolop("<", idx_sub, N)
+            cond = boolop("<", idx_sub, N, T.bool)
             return LoopIR.If(cond, body, [], None, srcinfo)
 
         ir, fwd_wrap = fwd(loop_cursor).body()._wrap(guard_wrapper, "body")
@@ -520,7 +520,7 @@ def DoSplit(loop_cursor, quot, hi, lo, tail="guard", perfect=False):
                     f"cannot perfectly split the '{split_loop.iter}' loop."
                 )
 
-            hi_rng = boolop("/", N, cnst(quot))
+            hi_rng = boolop("/", N, cnst(quot), T.index)
         else:
             if N.val % quot != 0:
                 raise SchedulingError(
@@ -573,7 +573,7 @@ def DoSplit(loop_cursor, quot, hi, lo, tail="guard", perfect=False):
             cut_i, LoopIR.Const(0, T.index, srcinfo), Ntail, cut_body, None, srcinfo
         )
         if tail_strategy == "cut_and_guard":
-            cond = boolop(">", Ntail, LoopIR.Const(0, T.int, srcinfo))
+            cond = boolop(">", Ntail, LoopIR.Const(0, T.int, srcinfo), T.bool)
             cut_s = LoopIR.If(cond, [cut_s], [], None, srcinfo)
 
         ir, fwd_ins = fwd(loop_cursor).after()._insert([cut_s])


Test patch:
diff --git a/tests/golden/test_schedules/test_divide_loop_perfect2.txt b/tests/golden/test_schedules/test_divide_loop_perfect2.txt
new file mode 100644
index 0000000..fef1251
--- /dev/null
+++ b/tests/golden/test_schedules/test_divide_loop_perfect2.txt
@@ -0,0 +1,8 @@
+def foo(n: size, A: i8[n] @ DRAM):
+    assert n % 4 == 0
+    for io in seq(0, n / 4):
+        tile: i8[4] @ DRAM
+        for ii in seq(0, 4):
+            tile[ii] = 0.2
+        for i0 in seq(0, 4):
+            A[i0 + 4 * io] = tile[i0]
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 3fa18a1..1b9f393 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -968,6 +968,18 @@ def test_divide_loop_perfect(golden):
     assert str(foo) == golden
 
 
+def test_divide_loop_perfect2(golden):
+    @proc
+    def foo(n: size, A: i8[n]):
+        assert n % 4 == 0
+        for i in seq(0, n):
+            A[i] = 0.2
+
+    foo = divide_loop(foo, "i", 4, ["io", "ii"], perfect=True)
+    foo = stage_mem(foo, "for ii in _:_", "A[io*4:io*4+4]", "tile")
+    assert str(simplify(foo)) == golden
+
+
 def test_divide_loop_perfect_fail():
     @proc
     def foo(n: size, A: i8[n]):