Repository: exo-lang/exo

Issue:
divide_loop should validate perfectness with a real solver
`divide_loop`'s `perfect` argument is an important example of a scheduling directive that performs some custom analyses to check for validity. It currently has ad-hoc checks for constant extents, or for static asserts of the trivial form `assert size%k == 0`. It does not employ a real solver, so it does not generalize even to simple extensions like `assert size1%k == 0 and size2%k == 0`.

Directives like this should be much more powerful in their ability to use static asserts and other conditions in their reasoning.

That could mean:

- Figuring out how `perfect=True` can be made to rely entirely on existing analysis checks
- Adding a simple lowering to an existing solver (likely Z3) for performing these types of queries

The latter is probably useful infrastructure that could be used by many scheduling primitives and higher-level automations, but `divide_loop` seems like a reasonable motivating example with which to start.

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/exo-lang/exo/pull/579?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`98e3255`)](https://app.codecov.io/gh/exo-lang/exo/commit/98e3255605df94ceabb6bc8956263022e51e317b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) 88.10% compared to head [(`19f83c9`)](https://app.codecov.io/gh/exo-lang/exo/pull/579?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) 88.13%.

| [Files](https://app.codecov.io/gh/exo-lang/exo/pull/579?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) | Patch % | Lines |
|---|---|---|
| [src/exo/LoopIR\_scheduling.py](https://app.codecov.io/gh/exo-lang/exo/pull/579?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfc2NoZWR1bGluZy5weQ==) | 88.88% | [1 Missing :warning: ](https://app.codecov.io/gh/exo-lang/exo/pull/579?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #579      +/-   ##
==========================================
+ Coverage   88.10%   88.13%   +0.02%     
==========================================
  Files          83       83              
  Lines       19703    19708       +5     
==========================================
+ Hits        17360    17369       +9     
+ Misses       2343     2339       -4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/exo-lang/exo/pull/579?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang).


Patch:
diff --git a/src/exo/API_scheduling.py b/src/exo/API_scheduling.py
index 9485324..558df09 100644
--- a/src/exo/API_scheduling.py
+++ b/src/exo/API_scheduling.py
@@ -1646,7 +1646,7 @@ def divide_loop(proc, loop_cursor, div_const, new_iters, tail="guard", perfect=F
 
     stmt = loop_cursor._impl
 
-    ir, fwd = scheduling.DoSplit(
+    ir, fwd = scheduling.DoDivideLoop(
         stmt,
         quot=div_const,
         outer_iter=new_iters[0],
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index 388c576..54d65c3 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -628,7 +628,9 @@ def DoDivideWithRecompute(
     return ir, fwd
 
 
-def DoSplit(loop_cursor, quot, outer_iter, inner_iter, tail="guard", perfect=False):
+def DoDivideLoop(
+    loop_cursor, quot, outer_iter, inner_iter, tail="guard", perfect=False
+):
     loop = loop_cursor._node
     N = loop.hi
     outer_i = Sym(outer_iter)
@@ -674,26 +676,15 @@ def DoSplit(loop_cursor, quot, outer_iter, inner_iter, tail="guard", perfect=Fal
         outer_hi = szop("/", N, inner_hi)  # floor div
     elif tail_strategy == "perfect":
         if not isinstance(N, LoopIR.Const):
-            is_N_divisible = False
-            for pred in loop_cursor.get_root().preds:
-                if (
-                    isinstance(pred, LoopIR.BinOp)
-                    and pred.op == "=="
-                    and isinstance(pred.rhs, LoopIR.Const)
-                    and pred.rhs.val == 0
-                    and isinstance(pred.lhs, LoopIR.BinOp)
-                    and pred.lhs.op == "%"
-                    and isinstance(pred.lhs.rhs, LoopIR.Const)
-                    and pred.lhs.rhs.val > 0
-                    and pred.lhs.rhs.val % quot == 0
-                    and isinstance(pred.lhs.lhs, LoopIR.Read)
-                    and pred.lhs.lhs.name == loop.hi.name
-                ):
-                    is_N_divisible = True
-
-            if not is_N_divisible:
-                raise SchedulingError(f"cannot perfectly split the '{loop.iter}' loop.")
-
+            hi_mod_quot = boolop("%", N, cnst(quot), T.index)
+            try:
+                ir = loop_cursor.get_root()
+                loop = loop_cursor._node
+                Check_CompareExprs(ir, [loop], hi_mod_quot, "==", cnst(0))
+            except SchedulingError:
+                raise SchedulingError(
+                    f"cannot perfectly split the '{loop.iter}' loop " f"by {quot}"
+                )
             outer_hi = boolop("/", N, cnst(quot), T.index)
         else:
             if N.val % quot != 0:
@@ -2800,7 +2791,7 @@ class DoExtractMethod(Cursor_Rewrite):
             self.pop()
 
             if body or orelse:
-                return [s.update(body=body or s.body, orelse=orelse or s.orlse)]
+                return [s.update(body=body or s.body, orelse=orelse or s.orelse)]
 
             return None
 
@@ -4148,7 +4139,7 @@ __all__ = [
     "DoCommuteExpr",
     "DoLeftReassociateExpr",
     "DoSpecialize",
-    "DoSplit",
+    "DoDivideLoop",
     "DoUnroll",
     "DoAddLoop",
     "DoCutLoop",
diff --git a/src/exo/new_eff.py b/src/exo/new_eff.py
index dcc43e5..db14299 100644
--- a/src/exo/new_eff.py
+++ b/src/exo/new_eff.py
@@ -2121,7 +2121,7 @@ def Check_ExprBound(proc, stmts, expr, op, value, exception=True):
         query = ADef(e < AInt(value))
         err_msg = f"greater than {value}"
     elif op == "==":
-        query = ADef(e=AInt(value))
+        query = ADef(AEq(e, AInt(value)))
         err_msg = f"equal to {value}"
     else:
         assert False, "Bad case"


Test patch:
diff --git a/tests/golden/test_schedules/test_divide_loop_perfect3.txt b/tests/golden/test_schedules/test_divide_loop_perfect3.txt
new file mode 100644
index 0000000..8ad5f20
--- /dev/null
+++ b/tests/golden/test_schedules/test_divide_loop_perfect3.txt
@@ -0,0 +1,7 @@
+def foo(m: size, n: size, A: R[m, n] @ DRAM):
+    assert n % 4 == 0 and m % 8 == 0
+    for io in seq(0, m / 8):
+        for ii in seq(0, 8):
+            for jo in seq(0, n / 4):
+                for ji in seq(0, 4):
+                    A[ii + 8 * io, ji + 4 * jo] = 0.2
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index cb62ea9..3a5f569 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -1265,6 +1265,19 @@ def test_divide_loop_perfect2(golden):
     assert str(simplify(foo)) == golden
 
 
+def test_divide_loop_perfect3(golden):
+    @proc
+    def foo(m: size, n: size, A: R[m, n]):
+        assert n % 4 == 0 and m % 8 == 0
+        for i in seq(0, m):
+            for j in seq(0, n):
+                A[i, j] = 0.2
+
+    foo = divide_loop(foo, "i", 8, ["io", "ii"], perfect=True)
+    foo = divide_loop(foo, "j", 4, ["jo", "ji"], perfect=True)
+    assert str(simplify(foo)) == golden
+
+
 def test_divide_loop_perfect_fail():
     @proc
     def foo(n: size, A: i8[n]):