Repository: exo-lang/exo

Issue:
bind_expr Bugs
1. bind_expr doesn't accept an argument that is a curosr to an expressions, but rather it must be given a list always if we are using a cursor. This is a bug in the arguments processor. 
Bug example:
```
def test_bind_expr_issue():
    @proc
    def bar(n: size, x: f32[n], y: f32[n], z: f32[n]):
        for i in seq(0, n):
            x[i] = y[i]
            x[i] = z[i]
    
    yi_cursor = bar.find("y[_]")
    zi_cursor = bar.find("z[_]")
    bar = bind_expr(bar, [yi_cursor], "y_read")     <----  Okay
    bar = bind_expr(bar, zi_cursor, "z_read")       <----- Error
````

2. Passing a list of more than one cursor to bind_expr binds only the first cursor unless `cse` is set to True. I find this very confusing because why else would I want to pass a list of many cursors:
```
def test_bind_expr_issue():
    @proc
    def bar(n: size, x: f32[n], y: f32[n], z: f32[n]):
        for i in seq(0, n):
            x[i] = y[i]
            z[i] = y[i]
    
    yi_cursors = bar.find("y[_]", many=True)
    print(bind_expr(bar, yi_cursors, "y_read"))
    print(bind_expr(bar, yi_cursors, "y_read", cse=True))
```
```
def bar(n: size, x: f32[n] @ DRAM, y: f32[n] @ DRAM, z: f32[n] @ DRAM):
    for i in seq(0, n):
        y_read: R
        y_read = y[i]
        x[i] = y_read
        z[i] = y[i]
def bar(n: size, x: f32[n] @ DRAM, y: f32[n] @ DRAM, z: f32[n] @ DRAM):
    for i in seq(0, n):
        y_read: R
        y_read = y[i]
        x[i] = y_read
        z[i] = y_read
```
Perhpas we should internally set cse=True whenever the user passes a list of cursors. I think generally it is unclear to me what cse implies in the context of cursors and in the contex of template arguments. We should at least document this behvaior if it is what is expected.
3. bind_expr currently assigns the type R to the new allocation. I think this is a bit confusing since when I bind something I would expect the default to be the type of the original expressions, otherwise you can easily get into a situation where you have cross-precision computation even though your original computation didn't have that. Changing the type of some computation shouldn't be something that is silently done by a scheduling operations, but only by explicit direction from the processor by using set_precision. Another reason why this is confusing is that stage_mem would use the precision of the staged buffer which makes it a bit inconsinstent for bind_expr to do something else.

Discussion:
SamirDroubi: It is expected that BLAS tests will fail.
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/exo-lang/exo/pull/542?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`295dd32`)](https://app.codecov.io/gh/exo-lang/exo/commit/295dd32a1dde0cb88499e255cdd86853081e67c7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) 87.05% compared to head [(`e22456b`)](https://app.codecov.io/gh/exo-lang/exo/pull/542?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) 86.77%.


: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     #542      +/-   ##
==========================================
- Coverage   87.05%   86.77%   -0.28%     
==========================================
  Files          81       80       -1     
  Lines       19313    19001     -312     
==========================================
- Hits        16813    16489     -324     
- Misses       2500     2512      +12     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/exo-lang/exo/pull/542?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/apps/aarch64/filter1D/filter1D_cursor_forwarding.py b/apps/aarch64/filter1D/filter1D_cursor_forwarding.py
index fbc8245..c5aab1c 100644
--- a/apps/aarch64/filter1D/filter1D_cursor_forwarding.py
+++ b/apps/aarch64/filter1D/filter1D_cursor_forwarding.py
@@ -42,7 +42,6 @@ filter1D = fission(filter1D, filter1D.find("xX4[_] = _").after())
 # set memories & precision
 filter1D = set_memory(filter1D, "sum #1", Neon)
 filter1D = set_memory(filter1D, "xX4", Neon)
-filter1D = set_precision(filter1D, "xX4", "f32")
 
 # replace
 filter1D = replace(filter1D, "for outXi in _:_", neon_zero_4xf32)
diff --git a/apps/aarch64/filter1D/filter1D_old.py b/apps/aarch64/filter1D/filter1D_old.py
index b90c04e..36fe3fc 100644
--- a/apps/aarch64/filter1D/filter1D_old.py
+++ b/apps/aarch64/filter1D/filter1D_old.py
@@ -34,7 +34,6 @@ filter1D = autofission(filter1D, filter1D.find("xX4[_] = _").after())
 # set memories & precision
 filter1D = set_memory(filter1D, "sum #1", Neon)
 filter1D = set_memory(filter1D, "xX4", Neon)
-filter1D = set_precision(filter1D, "xX4", "f32")
 
 # replace
 filter1D = replace(filter1D, "for outXi in _:_", neon_zero_4xf32)
diff --git a/src/exo/API_scheduling.py b/src/exo/API_scheduling.py
index d74a72b..ee28555 100644
--- a/src/exo/API_scheduling.py
+++ b/src/exo/API_scheduling.py
@@ -917,20 +917,19 @@ def rewrite_expr(proc, expr_cursor, new_expr):
     return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
 
 
-@sched_op([ExprCursorA(many=True), NameA, BoolA])
-def bind_expr(proc, expr_cursors, new_name, cse=False):
+@sched_op([ListOrElemA(ExprCursorA), NameA])
+def bind_expr(proc, expr_cursors, new_name):
     """
-    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. It will stop upon
-    encountering a read of any buffer that the expression depends on.
+    Bind some numeric/data-value type expression(s) into a new intermediate,
+    scalar-sized buffer. Attempts to perform common sub-expression
+    elimination while binding. It will stop upon encountering a read of any
+    buffer that the expression depends on. The precision of the new allocation
+    is that of the bound expression.
 
     args:
         expr_cursors    - a list of cursors to multiple instances of the
                           same expression
         new_name        - a string to name the new buffer
-        cse             - (bool) use common sub-expression elimination?
 
     rewrite:
         bind_expr(..., '32.0 * x[i]', 'b')
@@ -947,7 +946,7 @@ def bind_expr(proc, expr_cursors, new_name, cse=False):
             "can be bound by bind_expr()"
         )
 
-    ir, fwd = scheduling.DoBindExpr(new_name, exprs, cse)
+    ir, fwd = scheduling.DoBindExpr(new_name, exprs)
     return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
 
 
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index aab03d6..02d63e8 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -1151,12 +1151,9 @@ def DoRewriteExpr(expr_cursor, new_expr):
     return expr_cursor._replace(new_expr)
 
 
-def DoBindExpr(new_name, expr_cursors, cse=False):
+def DoBindExpr(new_name, expr_cursors):
     assert expr_cursors
 
-    if not cse:
-        expr_cursors = expr_cursors[0:1]
-
     expr = expr_cursors[0]._node
     assert isinstance(expr, LoopIR.expr)
     assert expr.type.is_numeric()
@@ -1168,8 +1165,10 @@ def DoBindExpr(new_name, expr_cursors, cse=False):
     init_c = get_enclosing_stmt_cursor(expr_cursors[0])
 
     new_name = Sym(new_name)
-    alloc_s = LoopIR.Alloc(new_name, T.R, None, None, expr.srcinfo)
-    assign_s = LoopIR.Assign(new_name, T.R, None, [], expr, None, expr.srcinfo)
+    alloc_s = LoopIR.Alloc(new_name, expr.type.basetype(), None, None, expr.srcinfo)
+    assign_s = LoopIR.Assign(
+        new_name, expr.type.basetype(), None, [], expr, None, expr.srcinfo
+    )
     ir, fwd = init_c.before()._insert([alloc_s, assign_s])
 
     new_read = LoopIR.Read(new_name, [], expr.type, expr.srcinfo)


Test patch:
diff --git a/tests/golden/test_schedules/test_bind_cursor_arg.txt b/tests/golden/test_schedules/test_bind_cursor_arg.txt
new file mode 100644
index 0000000..adb68dd
--- /dev/null
+++ b/tests/golden/test_schedules/test_bind_cursor_arg.txt
@@ -0,0 +1,4 @@
+def foo(a: R @ DRAM):
+    const: R
+    const = 1.0
+    a = const
\ No newline at end of file
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 2dfbd3c..6924f2e 100644
--- a/tests/golden/test_schedules/test_bind_expr_diff_indices.txt
+++ b/tests/golden/test_schedules/test_bind_expr_diff_indices.txt
@@ -2,7 +2,7 @@ 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
         x[i] = x[i] - y[i]
-        tmp: R
+        tmp: i8
         tmp = x[i] + y[i] + 1.0
         w[i + 1] = tmp
         x[i] = y[i]
diff --git a/tests/golden/test_schedules/test_bind_lhs.txt b/tests/golden/test_schedules/test_bind_lhs.txt
index cdc26b6..7ef3dcd 100644
--- a/tests/golden/test_schedules/test_bind_lhs.txt
+++ b/tests/golden/test_schedules/test_bind_lhs.txt
@@ -2,9 +2,9 @@ def myfunc_cpu(inp: i32[1, 1, 16] @ DRAM, out: i32[1, 1, 16] @ DRAM):
     for ii in seq(0, 1):
         for jj in seq(0, 1):
             for kk in seq(0, 16):
-                inp_ram: R
+                inp_ram: i32
                 inp_ram = inp[ii, jj, kk]
-                out_ram: R
+                out_ram: i32
                 out_ram = out[ii, jj, kk]
                 out[ii, jj, kk] += out_ram + inp_ram
                 out[ii, jj, kk] = out[ii, jj, kk] * inp_ram
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_simple_bind_expr.txt b/tests/golden/test_schedules/test_simple_bind_expr.txt
index 11194be..531e9d9 100644
--- a/tests/golden/test_schedules/test_simple_bind_expr.txt
+++ b/tests/golden/test_schedules/test_simple_bind_expr.txt
@@ -1,5 +1,5 @@
 def bar(n: size, x: i8[n] @ DRAM, y: i8[n] @ DRAM, z: i8[n] @ DRAM):
     for i in seq(0, n):
-        z_tmp: R
+        z_tmp: i8
         z_tmp = x[i] + y[i]
         z[i] = z_tmp
\ No newline at end of file
diff --git a/tests/test_neon.py b/tests/test_neon.py
index 92840ce..8283b7d 100644
--- a/tests/test_neon.py
+++ b/tests/test_neon.py
@@ -310,7 +310,7 @@ def simple_math_neon_sched():
 
         p = stage_mem(p, "for ii in _:_ #0", "x[4 * io : 4 * io + 4]", "xVec")
 
-        p = bind_expr(p, "y[_]", "yVec", cse=True)
+        p = bind_expr(p, p.find("y[_]", many=True), "yVec")
         p = autolift_alloc(p, "yVec: _", keep_dims=True)
         p = fission(p, p.find("yVec[_] = _").after())
 
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 4426128..f9e3b53 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -1091,11 +1091,20 @@ def test_bind_lhs(golden):
                     out[ii, jj, kk] += out[ii, jj, kk] + inp[ii, jj, kk]
                     out[ii, jj, kk] = out[ii, jj, kk] * inp[ii, jj, kk]
 
-    myfunc_cpu = bind_expr(myfunc_cpu, "inp[_]", "inp_ram", cse=True)
-    myfunc_cpu = bind_expr(myfunc_cpu, "out[_]", "out_ram", cse=True)
+    myfunc_cpu = bind_expr(myfunc_cpu, myfunc_cpu.find("inp[_]", many=True), "inp_ram")
+    myfunc_cpu = bind_expr(myfunc_cpu, myfunc_cpu.find("out[_]", many=True), "out_ram")
     assert str(myfunc_cpu) == golden
 
 
+def test_bind_cursor_arg(golden):
+    @proc
+    def foo(a: R):
+        a = 1.0
+
+    foo = bind_expr(foo, foo.find("1.0"), "const")
+    assert str(foo) == golden
+
+
 def test_divide_with_recompute(golden):
     @proc
     def foo(n: size, A: i8[n + 3]):
@@ -1640,7 +1649,7 @@ def test_bind_expr_diff_indices(golden):
             x[i] = y[i]
             w[i] = x[i] + y[i] + 1.0
 
-    bar = bind_expr(bar, "x[i]+y[i]+1.0", "tmp", cse=True)
+    bar = bind_expr(bar, bar.find("x[i]+y[i]+1.0", many=True), "tmp")
     assert str(bar) == golden
 
 
diff --git a/tests/test_x86.py b/tests/test_x86.py
index e1944b5..1d142f3 100644
--- a/tests/test_x86.py
+++ b/tests/test_x86.py
@@ -108,7 +108,7 @@ def simple_math_avx2_sched():
         p = replace(p, "for i0 in _:_ #0", mm256_loadu_ps)
         p = replace(p, "for i0 in _:_ #0", mm256_storeu_ps)
 
-        p = bind_expr(p, "y[_]", "yVec", cse=True)
+        p = bind_expr(p, p.find("y[_]", many=True), "yVec")
         p = autolift_alloc(p, "yVec: _", keep_dims=True)
         p = set_memory(p, "yVec", AVX2)
         p = old_fission_after(p, "yVec[_] = _")