Repository: exo-lang/exo

Issue:
bind_expr with CSE is broken
Per the discussion in https://github.com/exo-lang/exo/pull/237 the CSE option for bind_expr incorrectly factors out right-hand-sides and generally breaks well-formedness of the code. This needs to be fixed.

Discussion:

Patch:
diff --git a/src/exo/API_scheduling.py b/src/exo/API_scheduling.py
index dcdb473..6ba0939 100644
--- a/src/exo/API_scheduling.py
+++ b/src/exo/API_scheduling.py
@@ -823,7 +823,8 @@ def bind_expr(proc, expr_cursors, new_name, cse=False):
     Bind some numeric/data-value type expression into a new intermediate,
     scalar-sized buffer.  If `cse=True` and more than one expression is
     pointed to, then this operation will attempt to perform
-    common sub-expression elimination while binding.
+    common sub-expression elimination while binding. It will stop upon
+    encountering a read of any buffer that the expression depends on.
 
     args:
         expr_cursors    - a list of cursors to multiple instances of the
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index b5cdeec..1ba35d9 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -1236,10 +1236,12 @@ class _BindExpr(LoopIR_Rewrite):
         self.orig_proc = proc
         self.new_name = Sym(new_name)
         self.exprs = exprs if cse else [exprs[0]]
+        self.expr_reads = set(sum([get_reads(e) for e in self.exprs], []))
         self.use_cse = cse
         self.found_expr = None
         self.placed_alloc = False
         self.sub_done = False
+        self.found_write = False
 
         super().__init__(proc)
 
@@ -1292,6 +1294,9 @@ class _BindExpr(LoopIR_Rewrite):
         return None
 
     def map_s(self, s):
+        if self.found_write:
+            return None
+
         if self.sub_done:
             return super().map_s(s)
 
@@ -1315,28 +1320,17 @@ class _BindExpr(LoopIR_Rewrite):
 
         if isinstance(s, (LoopIR.Assign, LoopIR.Reduce)):
             e = self.exprs[0]
-            # bind LHS when self.use_cse == True
-            if (
-                self.use_cse
-                and isinstance(e, LoopIR.Read)
-                and e.name == s.name
-                and e.type == s.type
-            ):
-                try:
-                    for i, j in zip(e.idx, s.idx):
-                        Check_ExprEqvInContext(self.orig_proc, [s], i, j)
-
-                    return [
-                        s.update(
-                            name=self.new_name,
-                            cast=None,
-                            idx=[],
-                            rhs=self.apply_e(s.rhs),
-                            eff=None,
-                        )
-                    ]
-                except SchedulingError:
-                    pass
+            new_rhs = self.map_e(s.rhs)
+
+            # terminate CSE if the expression is written to
+            if self.found_expr and self.use_cse:
+                for (name, type) in self.expr_reads:
+                    if s.name == name and s.type == type:
+                        self.found_write = True
+
+            if new_rhs is not None:
+                return [s.update(rhs=new_rhs)]
+            return None
 
         return super().map_s(s)
 


Test patch:
diff --git a/tests/golden/test_schedules/test_bind_expr_diff_indices.txt b/tests/golden/test_schedules/test_bind_expr_diff_indices.txt
index e5062bb..8adb304 100644
--- a/tests/golden/test_schedules/test_bind_expr_diff_indices.txt
+++ b/tests/golden/test_schedules/test_bind_expr_diff_indices.txt
@@ -1,8 +1,9 @@
 def bar(n: size, x: i8[n] @ DRAM, y: i8[n] @ DRAM, z: i8[n] @ DRAM):
     for i in seq(0, n - 1):
         w: i8[n] @ DRAM
-        w[i + 1] = x[i] + y[i]
-        w_tmp = x[i] + y[i]
-        w_tmp: R
-        w_tmp = w[i]
-        z[i] = w_tmp
+        x[i] = x[i] - y[i]
+        tmp: R
+        tmp = x[i] + y[i] + 1.0
+        w[i + 1] = tmp
+        x[i] = y[i]
+        w[i] = x[i] + y[i] + 1.0
diff --git a/tests/golden/test_schedules/test_bind_lhs.txt b/tests/golden/test_schedules/test_bind_lhs.txt
index 50e332a..eb81a88 100644
--- a/tests/golden/test_schedules/test_bind_lhs.txt
+++ b/tests/golden/test_schedules/test_bind_lhs.txt
@@ -6,5 +6,5 @@ def myfunc_cpu(inp: i32[1, 1, 16] @ DRAM, out: i32[1, 1, 16] @ DRAM):
                 inp_ram = inp[ii, jj, kk]
                 out_ram: R
                 out_ram = out[ii, jj, kk]
-                out_ram = out_ram + inp_ram
-                out_ram = out_ram * inp_ram
+                out[ii, jj, kk] = out_ram + inp_ram
+                out[ii, jj, kk] = out[ii, jj, kk] * inp_ram
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 6cac342..365c07a 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -824,11 +824,12 @@ def test_bind_expr_diff_indices(golden):
     def bar(n: size, x: i8[n], y: i8[n], z: i8[n]):
         for i in seq(0, n - 1):
             w: i8[n]
-            w[i + 1] = x[i] + y[i]
-            w[i] = x[i] + y[i]
-            z[i] = w[i]
+            x[i] = x[i] - y[i]
+            w[i + 1] = x[i] + y[i] + 1.0
+            x[i] = y[i]
+            w[i] = x[i] + y[i] + 1.0
 
-    bar = bind_expr(bar, "w[i]", "w_tmp", cse=True)
+    bar = bind_expr(bar, "x[i]+y[i]+1.0", "tmp", cse=True)
     assert str(bar) == golden