Repository: exo-lang/exo

Issue:
Fix lift_alloc to emit error when n_lifts is too large

Discussion:

Patch:
diff --git a/apps/gemmini/src/exo/conv.py b/apps/gemmini/src/exo/conv.py
index fe239eb..8693bda 100644
--- a/apps/gemmini/src/exo/conv.py
+++ b/apps/gemmini/src/exo/conv.py
@@ -182,9 +182,13 @@ def schedule_conv_3():
     conv = expand_dim(
         conv, "i_s: i8[_] #1", "30", "krow + orow_i", unsafe_disable_checks=True
     )
-    conv = old_lift_alloc(conv, "i_s : _", n_lifts=5, keep_dims=False)
-    conv = old_lift_alloc(conv, "w_s : _", n_lifts=4, keep_dims=False)
-    conv = old_lift_alloc(conv, "res : _", n_lifts=4, keep_dims=False)
+    # TODO: We should definitely use repeat for this!
+    conv = old_lift_alloc(conv, "i_s : _ #0", n_lifts=5, keep_dims=False)
+    conv = old_lift_alloc(conv, "i_s : _ #1", n_lifts=4, keep_dims=False)
+    conv = old_lift_alloc(conv, "w_s : _ #0", n_lifts=4, keep_dims=False)
+    conv = old_lift_alloc(conv, "w_s : _ #1", n_lifts=3, keep_dims=False)
+    conv = old_lift_alloc(conv, "res : _ #0", n_lifts=4, keep_dims=False)
+    conv = old_lift_alloc(conv, "res : _ #1", n_lifts=3, keep_dims=False)
 
     conv = old_fission_after(conv, "for kch_o in _:_ #0", n_lifts=6)
     conv = old_fission_after(conv, "for kch_o in _:_ #2", n_lifts=5)
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index 1ba35d9..30eb63e 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -1807,7 +1807,8 @@ class _DoLiftAllocSimple(LoopIR_Rewrite):
 # --------------------------------------------------------------------------- #
 # Lift Allocation scheduling directive
 
-
+# TODO: Implement autolift_alloc's logic using high-level scheduling metaprogramming and
+#       delete this code
 class _LiftAlloc(LoopIR_Rewrite):
     def __init__(self, proc, alloc_stmt, n_lifts, mode, size, keep_dims):
         assert isinstance(alloc_stmt, LoopIR.Alloc)
@@ -1848,12 +1849,16 @@ class _LiftAlloc(LoopIR_Rewrite):
 
     def map_s(self, s):
         if s is self.alloc_stmt:
-            # mark the point we want to lift this alloc-stmt to
-            n_up = min(self.n_lifts, len(self.ctrl_ctxt))
-            self.lift_site = self.ctrl_ctxt[-n_up]
+            if self.n_lifts > len(self.ctrl_ctxt):
+                raise SchedulingError(
+                    f"specified lift level {self.n_lifts} "
+                    "is higher than the number of loop "
+                    f"{len(self.ctrl_ctxt)}"
+                )
+            self.lift_site = self.ctrl_ctxt[-self.n_lifts]
 
             # extract the ranges and variables of enclosing loops
-            idxs, rngs = self.get_ctxt_itrs_and_rngs(n_up)
+            idxs, rngs = self.get_ctxt_itrs_and_rngs(self.n_lifts)
 
             # compute the lifted allocation buffer type, and
             # the new allocation statement


Test patch:
diff --git a/tests/pldi22/test_gemmini_conv_ae.py b/tests/pldi22/test_gemmini_conv_ae.py
index c8d8ce5..91f5167 100644
--- a/tests/pldi22/test_gemmini_conv_ae.py
+++ b/tests/pldi22/test_gemmini_conv_ae.py
@@ -125,10 +125,13 @@ def test_conv_ae(golden):
     gemmini = expand_dim(
         gemmini, "i_s: i8[_]", "30", "krow + orow_i", unsafe_disable_checks=True
     )
-    gemmini = old_lift_alloc(gemmini, "i_s : _", n_lifts=5, keep_dims=False)
-    gemmini = old_lift_alloc(gemmini, "w_s : _", n_lifts=4, keep_dims=False)
-
-    gemmini = old_lift_alloc(gemmini, "res : _", n_lifts=4, keep_dims=False)
+    # TODO: Use cursor + repeat to improve this!
+    gemmini = old_lift_alloc(gemmini, "i_s : _ #0", n_lifts=5, keep_dims=False)
+    gemmini = old_lift_alloc(gemmini, "i_s : _ #1", n_lifts=4, keep_dims=False)
+    gemmini = old_lift_alloc(gemmini, "w_s : _ #0", n_lifts=4, keep_dims=False)
+    gemmini = old_lift_alloc(gemmini, "w_s : _ #1", n_lifts=3, keep_dims=False)
+    gemmini = old_lift_alloc(gemmini, "res : _ #0", n_lifts=4, keep_dims=False)
+    gemmini = old_lift_alloc(gemmini, "res : _ #1", n_lifts=3, keep_dims=False)
 
     gemmini = old_fission_after(gemmini, "for kch_o in _:_ #0", n_lifts=6)
     gemmini = old_fission_after(gemmini, "for kch_o in _:_ #2", n_lifts=5)
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 365c07a..a7b6443 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -396,6 +396,18 @@ def test_lift_alloc_simple_error():
         lift_alloc(bar, "tmp_a : _", n_lifts=3)
 
 
+def test_autolift_alloc_error():
+    @proc
+    def bar(n: size, A: i8[n]):
+        for i in seq(0, n):
+            for j in seq(0, n):
+                tmp_a: i8
+                tmp_a = A[i]
+
+    with pytest.raises(SchedulingError, match="specified lift level 3 is"):
+        autolift_alloc(bar, "tmp_a : _", n_lifts=3)
+
+
 def test_expand_dim(golden):
     @proc
     def foo(n: size, m: size, x: i8):