Repository: exo-lang/exo

Issue:
inline_window bug
Here is a test case that replicates the issue:
```
def test_inline_window3(golden):
    @proc
    def inner_memset(x: [R][16] @ DRAM):
        for i in seq(0, 16):
            x[i] = 0.0
    @proc
    def memset(n: size, x: [R][n] @ DRAM):
        assert n % 16 == 0
        for io in seq(0, n / 16):
            x_1 = x[16 * io + 0:16 * io + 16]
            inner_memset(x_1)
    memset = inline_window(memset, "x_1 = _")
```

The issue is at the following line: https://github.com/exo-lang/exo/blob/c791819a7b9459f87d032bf65837ecdbe3bf8448/src/exo/LoopIR_scheduling.py#L1008

We assume that the accesses to a window will always have the same dimensions as the window definition, but that's not always the case.

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/exo-lang/exo/pull/595?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) Report
Attention: Patch coverage is `70.58824%` with `10 lines` in your changes are missing coverage. Please review.
> Project coverage is 84.37%. Comparing base [(`c791819`)](https://app.codecov.io/gh/exo-lang/exo/commit/c791819a7b9459f87d032bf65837ecdbe3bf8448?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) to head [(`43cbefb`)](https://app.codecov.io/gh/exo-lang/exo/pull/595?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang).

| [Files](https://app.codecov.io/gh/exo-lang/exo/pull/595?dropdown=coverage&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\_unification.py](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfdW5pZmljYXRpb24ucHk=) | 0.00% | [5 Missing :warning: ](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) |
| [src/exo/new\_eff.py](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9uZXdfZWZmLnB5) | 40.00% | [3 Missing :warning: ](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) |
| [src/exo/LoopIR.py](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVIucHk=) | 80.00% | [1 Missing :warning: ](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) |
| [src/exo/LoopIR\_dataflow.py](https://app.codecov.io/gh/exo-lang/exo/pull/595?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfZGF0YWZsb3cucHk=) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/exo-lang/exo/pull/595?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     #595      +/-   ##
==========================================
- Coverage   86.29%   84.37%   -1.93%     
==========================================
  Files          88       88              
  Lines       21160    21166       +6     
==========================================
- Hits        18260    17858     -402     
- Misses       2900     3308     +408     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/exo-lang/exo/pull/595?dropdown=coverage&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/LoopIR.py b/src/exo/LoopIR.py
index 300ad3a..3665a26 100644
--- a/src/exo/LoopIR.py
+++ b/src/exo/LoopIR.py
@@ -80,7 +80,7 @@ module LoopIR {
          | Alloc( sym name, type type, mem mem )
          | Free( sym name, type type, mem mem )
          | Call( proc f, expr* args )
-         | WindowStmt( sym lhs, expr rhs )
+         | WindowStmt( sym name, expr rhs )
          attributes( effect? eff, srcinfo srcinfo )
 
     loop_mode = Seq()
@@ -970,7 +970,7 @@ class LoopIR_Compare:
                 self.match_e(a1, a2) for a1, a2 in zip(s1.args, s2.args)
             )
         elif isinstance(s1, LoopIR.WindowStmt):
-            return self.match_name(s1.lhs, s2.lhs) and self.match_e(s1.rhs, s2.rhs)
+            return self.match_name(s1.name, s2.name) and self.match_e(s1.rhs, s2.rhs)
         else:
             assert False, f"bad case: {type(s1)}"
 
@@ -1119,7 +1119,7 @@ class FreeVars(LoopIR_Do):
             if s.name not in self.env:
                 self.fv.add(s.name)
         elif styp is LoopIR.WindowStmt:
-            self.env[s.lhs] = True
+            self.env[s.name] = True
         elif styp is LoopIR.If:
             self.do_e(s.cond)
             self.push()
@@ -1206,9 +1206,9 @@ class Alpha_Rename(LoopIR_Rewrite):
             return [((s2 and s2[0]) or s).update(name=new_name)]
         elif isinstance(s, LoopIR.WindowStmt):
             rhs = self.map_e(s.rhs) or s.rhs
-            lhs = s.lhs.copy()
-            self.env[s.lhs] = lhs
-            return [s.update(lhs=lhs, rhs=rhs)]
+            name = s.name.copy()
+            self.env[s.name] = name
+            return [s.update(name=name, rhs=rhs)]
         elif isinstance(s, LoopIR.If):
             self.push()
             stmts = super().map_s(s)
diff --git a/src/exo/LoopIR_compiler.py b/src/exo/LoopIR_compiler.py
index a06232b..ebdc4e6 100644
--- a/src/exo/LoopIR_compiler.py
+++ b/src/exo/LoopIR_compiler.py
@@ -858,8 +858,8 @@ class Compiler:
             rhs = self.comp_e(s.rhs)
             assert isinstance(s.rhs, LoopIR.WindowExpr)
             mem = self.mems[s.rhs.name]
-            lhs = self.new_varname(s.lhs, typ=s.rhs.type, mem=mem)
-            self.add_line(f"struct {win_struct} {lhs} = {rhs};")
+            name = self.new_varname(s.name, typ=s.rhs.type, mem=mem)
+            self.add_line(f"struct {win_struct} {name} = {rhs};")
         elif isinstance(s, LoopIR.If):
             cond = self.comp_e(s.cond)
             self.add_line(f"if ({cond}) {{")
diff --git a/src/exo/LoopIR_dataflow.py b/src/exo/LoopIR_dataflow.py
index 7c7aed5..e6e4e00 100644
--- a/src/exo/LoopIR_dataflow.py
+++ b/src/exo/LoopIR_dataflow.py
@@ -133,7 +133,7 @@ class LoopIR_Dependencies(LoopIR_Do):
             self._lhs = None
         elif isinstance(s, LoopIR.WindowStmt):
             rhs_buf = self._alias.get(s.rhs.name, s.rhs.name)
-            self._alias[s.lhs] = rhs_buf
+            self._alias[s.name] = rhs_buf
             self._lhs = rhs_buf
             self._depends[rhs_buf].add(rhs_buf)
             self.do_e(s.rhs)
diff --git a/src/exo/LoopIR_pprint.py b/src/exo/LoopIR_pprint.py
index 807a441..38c7e18 100644
--- a/src/exo/LoopIR_pprint.py
+++ b/src/exo/LoopIR_pprint.py
@@ -428,7 +428,7 @@ def _print_stmt(stmt, env: PrintEnv, indent: str) -> list[str]:
 
     elif isinstance(stmt, LoopIR.WindowStmt):
         rhs = _print_expr(stmt.rhs, env)
-        return [f"{indent}{env.get_name(stmt.lhs)} = {rhs}"]
+        return [f"{indent}{env.get_name(stmt.name)} = {rhs}"]
 
     elif isinstance(stmt, LoopIR.Alloc):
         mem = f" @{stmt.mem.name()}" if stmt.mem else ""
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index c4c29b6..fec0d5e 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -1056,10 +1056,14 @@ def DoInlineWindow(window_cursor):
             new_idxs = calc_idx(rd.idx)
             old_typ = window_s.rhs.type
             new_typ = old_typ.update(src_buf=buf_name, idx=new_idxs)
-            return {"name": buf_name, "idx": new_idxs, "type": new_typ}
+            return rd.update(name=buf_name, idx=new_idxs, type=new_typ)
         elif isinstance(rd, LoopIR.Read):
-            new_idxs = calc_idx(rd.idx)
-            return {"name": buf_name, "idx": new_idxs}
+            if rd.idx:
+                new_idxs = calc_idx(rd.idx)
+                return rd.update(name=buf_name, idx=new_idxs)
+            else:
+                assert isinstance(c.parent()._node, LoopIR.Call)
+                return window_s.rhs
 
     def mk_stride_expr(c):
         e = c._node
@@ -1073,11 +1077,13 @@ def DoInlineWindow(window_cursor):
         return {"name": window_s.rhs.name, "idx": idxs}
 
     for c in get_rest_of_block(window_cursor):
-        ir, fwd = _replace_reads(ir, fwd, c, window_s.lhs, mk_read)
+        ir, fwd = _replace_reads(
+            ir, fwd, c, window_s.name, mk_read, only_replace_attrs=False
+        )
         ir, fwd = _replace_pats(
-            ir, fwd, c, f"stride({repr(window_s.lhs)}, _)", mk_stride_expr
+            ir, fwd, c, f"stride({repr(window_s.name)}, _)", mk_stride_expr
         )
-        ir, fwd = _replace_writes(ir, fwd, c, window_s.lhs, mk_write)
+        ir, fwd = _replace_writes(ir, fwd, c, window_s.name, mk_write)
 
     return ir, fwd
 
diff --git a/src/exo/LoopIR_unification.py b/src/exo/LoopIR_unification.py
index 32a84f9..93078ef 100644
--- a/src/exo/LoopIR_unification.py
+++ b/src/exo/LoopIR_unification.py
@@ -59,7 +59,7 @@ def Get_Live_Variables(stmt_cursor):
 
     for s in reversed(covering_stmts):
         if isinstance(s, LoopIR.WindowStmt):
-            live_vars[s.lhs] = s.rhs.type
+            live_vars[s.name] = s.rhs.type
         elif isinstance(s, LoopIR.Alloc):
             live_vars[s.name] = s.type
         elif isinstance(s, LoopIR.If):
@@ -922,10 +922,10 @@ class Unification:
         elif isinstance(ps, LoopIR.WindowStmt):
             self.unify_e(ps.rhs, bs.rhs)
             # new name identification is similar to Alloc
-            pvar = BufVar(ps.lhs, ps.rhs.type.as_tensor, use_win=False)
-            pvar.set_buf_solution(bs.lhs)
-            self.buf_unknowns[ps.lhs] = pvar
-            self.bbuf_types[bs.lhs] = bs.rhs.type.as_tensor
+            pvar = BufVar(ps.name, ps.rhs.type.as_tensor, use_win=False)
+            pvar.set_buf_solution(bs.name)
+            self.buf_unknowns[ps.name] = pvar
+            self.bbuf_types[bs.name] = bs.rhs.type.as_tensor
 
         # tail recursion
         self.unify_stmts(proc_s, block_s)
diff --git a/src/exo/effectcheck.py b/src/exo/effectcheck.py
index 347f2c9..c178b07 100644
--- a/src/exo/effectcheck.py
+++ b/src/exo/effectcheck.py
@@ -123,7 +123,7 @@ class InferEffects:
         elif isinstance(stmt, LoopIR.Alloc):
             self._types[stmt.name] = stmt.type
         elif isinstance(stmt, LoopIR.WindowStmt):
-            self._types[stmt.lhs] = stmt.rhs.type
+            self._types[stmt.name] = stmt.rhs.type
         else:
             pass
 
@@ -249,7 +249,7 @@ class InferEffects:
             )
         elif isinstance(stmt, LoopIR.WindowStmt):
             return LoopIR.WindowStmt(
-                stmt.lhs, stmt.rhs, eff_null(stmt.srcinfo), stmt.srcinfo
+                stmt.name, stmt.rhs, eff_null(stmt.srcinfo), stmt.srcinfo
             )
 
         else:
@@ -967,7 +967,7 @@ class CheckEffects:
                 # src_shape   = stmt.rhs.type.src_type.shape()
                 w_idx = stmt.rhs.type.idx
                 src_buf = stmt.rhs.type.src_buf
-                dst_buf = stmt.lhs
+                dst_buf = stmt.name
 
                 src_dims = [
                     d for d, w in enumerate(w_idx) if isinstance(w, LoopIR.Interval)
diff --git a/src/exo/mem_analysis.py b/src/exo/mem_analysis.py
index f6d7831..33e5f4b 100644
--- a/src/exo/mem_analysis.py
+++ b/src/exo/mem_analysis.py
@@ -135,7 +135,7 @@ class MemoryAnalysis:
 
         elif styp is LoopIR.WindowStmt:
             mem = self.get_e_mem(s.rhs)
-            self.mem_env[s.lhs] = mem
+            self.mem_env[s.name] = mem
             return s
 
         elif styp is LoopIR.Call:
diff --git a/src/exo/new_eff.py b/src/exo/new_eff.py
index db14299..10da1d4 100644
--- a/src/exo/new_eff.py
+++ b/src/exo/new_eff.py
@@ -429,7 +429,7 @@ def possible_config_writes(stmts):
                 # accumulate windowing expressions
                 awin = lift_e(s.rhs)
                 assert isinstance(awin, AWin)
-                aenv = self.windows[-1] + AEnv(s.lhs, awin, addnames=False)
+                aenv = self.windows[-1] + AEnv(s.name, awin, addnames=False)
                 self.windows[-1] = aenv
 
             elif isinstance(s, LoopIR.Call):
@@ -467,7 +467,7 @@ def globenv(stmts):
             aenvs.append(AEnv(globname, rhs, addnames=True))
         elif isinstance(s, LoopIR.WindowStmt):
             win = lift_e(s.rhs)
-            aenvs.append(AEnv(s.lhs, win))
+            aenvs.append(AEnv(s.name, win))
         elif isinstance(s, LoopIR.Alloc):
             win = AWinAlloc(s.name, s.type.shape())
             aenvs.append(AEnv(s.name, win))
@@ -1254,7 +1254,7 @@ def get_changing_scalars(stmts, changeset=None, aliases=None):
             for nm in pchgs:
                 add_name(nm)
         elif isinstance(s, LoopIR.WindowStmt):
-            aliases[s.lhs] = s.rhs.name
+            aliases[s.name] = s.rhs.name
         else:
             pass
 
@@ -2271,7 +2271,7 @@ class _OverApproxEffects(LoopIR_Do):
                 self.add_name(name)
             return  # don't call do_e on all of the arguments
         elif isinstance(s, LoopIR.WindowStmt):
-            self._aliases[s.lhs] = s.rhs.name
+            self._aliases[s.name] = s.rhs.name
 
         super().do_s(s)
 
@@ -2344,7 +2344,7 @@ class _Check_Aliasing_Helper(LoopIR_Do):
             name = s.rhs.name
             while name in self._aliases:
                 name = self._aliases[name]
-            self._aliases[s.lhs] = name
+            self._aliases[s.name] = name
         else:
             super().do_s(s)
 
diff --git a/src/exo/pattern_match.py b/src/exo/pattern_match.py
index 3210df5..7dca84c 100644
--- a/src/exo/pattern_match.py
+++ b/src/exo/pattern_match.py
@@ -245,7 +245,7 @@ class PatternMatch:
         elif isinstance(stmt, LoopIR.WindowStmt):
             if isinstance(pat, PAST.Assign):
                 return (
-                    self.match_name(pat.name, stmt.lhs)
+                    self.match_name(pat.name, stmt.name)
                     and pat.idx == []
                     and self.match_e(pat.rhs, stmt.rhs)
                 )
diff --git a/src/exo/prec_analysis.py b/src/exo/prec_analysis.py
index 263634d..671ed97 100644
--- a/src/exo/prec_analysis.py
+++ b/src/exo/prec_analysis.py
@@ -132,7 +132,7 @@ class PrecisionAnalysis(LoopIR_Rewrite):
 
         elif isinstance(s, LoopIR.WindowStmt):
             # update the type binding for this symbol...
-            self.set_type(result[0].lhs, result[0].rhs.type)
+            self.set_type(result[0].name, result[0].rhs.type)
 
         elif isinstance(s, LoopIR.Alloc):
             typ = result[0].type


Test patch:
diff --git a/tests/golden/test_schedules/test_inline_window3.txt b/tests/golden/test_schedules/test_inline_window3.txt
new file mode 100644
index 0000000..49a7e18
--- /dev/null
+++ b/tests/golden/test_schedules/test_inline_window3.txt
@@ -0,0 +1,6 @@
+def memset(n: size, x: [R][n] @ DRAM):
+    assert n % 16 == 0
+    res: R @ DRAM
+    for io in seq(0, n / 16):
+        inner_memset(x[16 * io:16 + 16 * io])
+        res += x[16 * io]
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 9f8a153..74cea3d 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -2249,6 +2249,25 @@ def test_inline_window2(golden):
     assert str(simplify(foo)) == golden
 
 
+def test_inline_window3(golden):
+    @proc
+    def inner_memset(x: [R][16] @ DRAM):
+        for i in seq(0, 16):
+            x[i] = 0.0
+
+    @proc
+    def memset(n: size, x: [R][n] @ DRAM):
+        assert n % 16 == 0
+        res: R
+        for io in seq(0, n / 16):
+            x_1 = x[16 * io : 16 * io + 16]
+            inner_memset(x_1)
+            res += x_1[0]
+
+    memset = inline_window(memset, "x_1 = _")
+    assert str(simplify(memset)) == golden
+
+
 def test_lift_if_second_statement_in_then_error():
     @proc
     def foo(m: size, x: R[m]):