Allow divide_dim to divide to non-literal expressions

src/exo/API_scheduling.py

@@ -1372,8 +1372,6 @@ def divide_dim(proc, alloc_cursor, dim_idx, quotient):
         `x : R[n, 3, 4, m]`
         `x[i, j / 4, j % 4, k] = ...`
     """
-    if quotient == 1:
-        raise ValueError("why are you trying to divide by 1?")
     stmt = alloc_cursor._impl
     if not (0 <= dim_idx < len(stmt._node.type.shape())):
         raise ValueError(f"Cannot divide out-of-bounds dimension index {dim_idx}")

src/exo/LoopIR_scheduling.py

@@ -264,6 +264,24 @@ def Check_CompareExprs(proc, stmts, lhs, op, rhs):
     Check_ExprBound(proc, stmts, expr, op, 0)
 
 
+def Check_IsDivisible(proc, stmts, expr, quot):
+    failed = False
+    if not isinstance(expr, LoopIR.Const):
+        try:
+            quot = LoopIR.Const(quot, T.int, null_srcinfo())
+            expr_mod_quot = LoopIR.BinOp("%", expr, quot, T.index, null_srcinfo())
+            zero = LoopIR.Const(0, T.int, null_srcinfo())
+            Check_CompareExprs(proc, stmts, expr_mod_quot, "==", zero)
+        except SchedulingError:
+            failed = True
+    else:
+        # Fast path
+        failed = expr.val % quot != 0
+
+    if failed:
+        raise SchedulingError(f"cannot perfectly divide '{expr}' by {quot}")
+
+
 def extract_env(c: ic.Cursor) -> List[Tuple[Sym, ic.Cursor]]:
     """
     Extract the environment of live variables at `c`.
@@ -303,6 +321,28 @@ def move_back(c):
         return c.prev()
 
 
+# --------------------------------------------------------------------------- #
+# --------------------------------------------------------------------------- #
+# IR Building Helpers
+
+
+def divide_expr(e, quot):
+    assert isinstance(e, LoopIR.expr)
+    if isinstance(quot, int):
+        quot_int = quot
+        quot_ir = LoopIR.Const(quot, e.type, e.srcinfo)
+    elif isinstance(quot, LoopIR.expr):
+        quot_int = quot.val
+        quot_ir = quot
+    else:
+        assert False, f"Bad case {type(quot)}"
+    if isinstance(e, LoopIR.Const) and e.val % quot == 0:
+        div = LoopIR.Const(e.val // quot_int, e.type, e.srcinfo)
+    else:
+        div = LoopIR.BinOp("/", e, quot_ir, e.type, e.srcinfo)
+    return div
+
+
 # --------------------------------------------------------------------------- #
 # --------------------------------------------------------------------------- #
 # Scheduling directives
@@ -728,25 +768,10 @@ def ceildiv(lhs, rhs):
     elif tail_strategy in ["cut", "cut_and_guard"]:
         outer_hi = szop("/", N, inner_hi)  # floor div
     elif tail_strategy == "perfect":
-        if not isinstance(N, LoopIR.Const):
-            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:
-                raise SchedulingError(
-                    f"cannot perfectly split the '{loop.iter}' loop "
-                    f"because {quot} does not evenly divide "
-                    f"{N.val}"
-                )
-            outer_hi = cnst(N.val // quot)
+        ir = loop_cursor.get_root()
+        loop = loop_cursor._node
+        Check_IsDivisible(ir, [loop], N, quot)
+        outer_hi = divide_expr(N, quot)
     else:
         assert False, f"bad tail strategy: {tail_strategy}"
 
@@ -1728,17 +1753,13 @@ def DoDivideDim(alloc_cursor, dim_idx, quotient):
     old_typ = alloc_s.type
     old_shp = old_typ.shape()
     dim = old_shp[dim_idx]
-    if not isinstance(dim, LoopIR.Const):
-        raise SchedulingError(f"Cannot divide non-literal dimension: {dim}")
-    if not dim.val % quotient == 0:
-        raise SchedulingError(f"Cannot divide {dim.val} evenly by {quotient}")
-    denom = quotient
-    numer = dim.val // denom
+    Check_IsDivisible(alloc_cursor.get_root(), [alloc_s], dim, quotient)
+    numer = divide_expr(dim, quotient)
     new_shp = (
         old_shp[:dim_idx]
         + [
-            LoopIR.Const(numer, T.int, dim.srcinfo),
-            LoopIR.Const(denom, T.int, dim.srcinfo),
+            numer,
+            LoopIR.Const(quotient, T.int, dim.srcinfo),
         ]
         + old_shp[dim_idx + 1 :]
     )

tests/golden/test_schedules/test_divide_dim_2.txt

@@ -0,0 +1,6 @@
+def foo(n: size, m: size, A: R[n + m + 12] @ DRAM):
+    x: R[n, 3 * m, 4, m] @ DRAM
+    for i in seq(0, n):
+        for j in seq(0, 12):
+            for k in seq(0, m):
+                x[i, j / 4, j % 4, k] = A[i + j + k]

tests/golden/test_schedules/test_divide_dim_3.txt

@@ -0,0 +1,6 @@
+def foo(n: size, m: size):
+    x: R[n, 1, (7 + m) / 8 * 8 / 8, 8, 1, m, 1] @ DRAM
+    for i in seq(0, n):
+        for j in seq(0, m):
+            for k in seq(0, m):
+                x[i, 0, j / 8, j % 8, 0, k, 0] = 2.0

tests/test_schedules.py

@@ -1012,6 +1012,34 @@ def foo(n: size, m: size, A: R[n + m + 12]):
     assert str(foo) == golden
 
 
+def test_divide_dim_2(golden):
+    @proc
+    def foo(n: size, m: size, A: R[n + m + 12]):
+        x: R[n, 12 * m, m]
+        for i in seq(0, n):
+            for j in seq(0, 12):
+                for k in seq(0, m):
+                    x[i, j, k] = A[i + j + k]
+
+    foo = simplify(divide_dim(foo, "x", 1, 4))
+    assert str(foo) == golden
+
+
+def test_divide_dim_3(golden):
+    @proc
+    def foo(n: size, m: size):
+        x: R[n, ((m + 7) / 8) * 8, m]
+        for i in seq(0, n):
+            for j in seq(0, m):
+                for k in seq(0, m):
+                    x[i, j, k] = 2.0
+
+    for i in range(2, -1, -1):
+        foo = divide_dim(foo, "x", i, 1)
+    foo = simplify(divide_dim(foo, "x", 2, 8))
+    assert str(foo) == golden
+
+
 def test_divide_dim_fail_1():
     @proc
     def foo(n: size, m: size, A: R[n + m + 12]):
@@ -1024,10 +1052,20 @@ def foo(n: size, m: size, A: R[n + m + 12]):
     with pytest.raises(ValueError, match="out-of-bounds"):
         divide_dim(foo, "x", 3, 4)
 
-    with pytest.raises(SchedulingError, match="Cannot divide 12 evenly"):
+    with pytest.raises(SchedulingError, match="cannot perfectly divide"):
         divide_dim(foo, "x", 1, 5)
 
 
+def test_divide_dim_fail_2():
+    @proc
+    def foo(n: size, m: size):
+        x: R[n, 3 * m, m]
+
+    with pytest.raises(SchedulingError, match="cannot perfectly divide"):
+        for i in range(3):
+            divide_dim(foo, "x", i, 15)
+
+
 def test_mult_dim_1(golden):
     @proc
     def foo(n: size, m: size, A: R[n + m + 12]):
@@ -1516,7 +1554,7 @@ def foo(n: size, A: i8[n]):
         for i in seq(0, n):
             A[i] = 1.0
 
-    with pytest.raises(SchedulingError, match="cannot perfectly split"):
+    with pytest.raises(SchedulingError, match="cannot perfectly divide"):
         foo = divide_loop(foo, "i", 4, ["io", "ii"], perfect=True)