Repository: exo-lang/exo

Issue:
Check_BufferRW in stage_mem is totally broken

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/exo-lang/exo/pull/614?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 `93.22034%` with `4 lines` in your changes are missing coverage. Please review.
> Project coverage is 86.45%. Comparing base [(`775f81a`)](https://app.codecov.io/gh/exo-lang/exo/commit/775f81a129560c4dbb75f84c97551298f866ba28?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) to head [(`5240d5d`)](https://app.codecov.io/gh/exo-lang/exo/pull/614?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/614?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\_scheduling.py](https://app.codecov.io/gh/exo-lang/exo/pull/614?src=pr&el=tree&filepath=src%2Fexo%2FLoopIR_scheduling.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfc2NoZWR1bGluZy5weQ==) | 92.30% | [4 Missing :warning: ](https://app.codecov.io/gh/exo-lang/exo/pull/614?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     #614      +/-   ##
==========================================
- Coverage   86.80%   86.45%   -0.36%     
==========================================
  Files          92       92              
  Lines       21474    21520      +46     
==========================================
- Hits        18641    18605      -36     
- Misses       2833     2915      +82     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/exo-lang/exo/pull/614?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).

yamaguchi1024: This PR unsurprising breaks the BLAS repo. @skeqiqevian Can you review the PR? @SamirDroubi Can we go through the changes to BLAS next week?
skeqiqevian: The following should fail. In general, we should fail (instead of just "not replacing") when an existing window expr in the staged scope is not entirely a subset or not entirely disjoint with our staged window expr.

```
def test_counterexample():
    @proc
    def foo(N: size, A: i8[N, N]):
        sum: i8
        for i in seq(0, N):
            for jo in seq(0, N / 4):
                for ji in seq(0, 4):
                    A[i, 4*jo + ji] = 0.0
                    y = A[i:i+1, 4*jo+ji]
                    sum += y

    foo = stage_mem(foo, "for ji in _:_", "A[i, 4*jo:4*jo+4]", "tile")
    print(foo)
```
yamaguchi1024: I'm giving up, assigning @skeqiqevian 🙏 
skeqiqevian: As currently implemented, this will not close #318. 
skeqiqevian: @yamaguchi1024 can you review when you get a chance? 

I removed the changes about writes shadowing reads since that implementation wasn't correct. In general, you have to reason about the entire block's read and write sets, which seems hard to do. I think doing such reasoning might be able to address #318 if we check that the block's write set is the same as the staged window. 
yamaguchi1024: @skeqiqevian Thanks Kevin! I agree that we have to reason about the entire block's read and write sets to make it correct, but I'm not sure why you had to remove it. The incorrect write-shadowing-read behaviour is not introduced in this PR. That's why BLAS CI is failing.

w/o this PR:
```
void exo_dasum_stride_1( void *ctxt, int_fast32_t n, struct exo_win_1f64c x, double* result ) { 
// assert stride(x, 0) == 1
double result_;
result_ = 0.0;
__m256d var10;
var10 = _mm256_setzero_pd();
__m256d var15[8];
var15[0] = _mm256_setzero_pd();
```

w/ this PR:
```
5: void exo_dasum_stride_1( void *ctxt, int_fast32_t n, struct exo_win_1f64c x, double* result ) {         
5: // assert stride(x, 0) == 1                    
5: double result_;                                                                                         
5: result_ = *result;                                                                                      
5: result_ = 0.0;                                                                                          
5: __m256d var10;                                                                                                                                                                                                     
5: var10 = _mm256_setzero_pd();                                                                            
5: __m256d var15[8];                                                                                       
5: var15[0] = _mm256_setzero_pd();   
```
We have this extra `result_ = *result;`.

Even if reasoning about complete write-shadowing-read is out-of-scope of this PR, can't we at least make a conservative but sound shadowing effect? I think what I was trying to do was incorrect because it didn't look at the buffer indices of read and writes.
skeqiqevian: I couldn't think of a tractable way to solve your problem - you need to check that every read is shadowed by a write in order to make that change. And even if there is a simple, but conservative check, I think it makes more sense to include that in separate PR. 
yamaguchi1024: Okay I think that's fine. I want to approve this PR but I can't because I started it. Do you want to approve it?

Also @SamirDroubi heads up this will change BLAS codegen slightly, but I think it's a negligible overhead
SamirDroubi: > Okay I think that's fine. I want to approve this PR but I can't because I started it. Do you want to approve it?
> 
> Also @SamirDroubi heads up this will change BLAS codegen slightly, but I think it's a negligible overhead

I would like first to see what changed and make sure it is not a big deal. I will build BLAS with this tomorrow and see the diff.

Patch:
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index a368974..c2cc2b9 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -24,9 +24,9 @@ from .new_eff import (
     Check_DeleteConfigWrite,
     Check_ExtendEqv,
     Check_ExprEqvInContext,
-    Check_BufferRW,
     Check_BufferReduceOnly,
     Check_Bounds,
+    Check_Access_In_Window,
     Check_IsDeadAfter,
     Check_IsIdempotent,
     Check_ExprBound,
@@ -42,6 +42,8 @@ import exo.API as api
 from .pattern_match import match_pattern
 from .memory import DRAM
 
+from functools import partial
+
 # --------------------------------------------------------------------------- #
 # --------------------------------------------------------------------------- #
 # Wrapper for LoopIR_Rewrite for scheduling directives which takes procedure cursor
@@ -196,45 +198,61 @@ def _replace_pats(ir, fwd, c, pat, repl, only_replace_attrs=True, use_sym_id=Tru
     # TODO: consider the implications of composing O(n) forwarding functions.
     #   will we need a special data structure? A chunkier operation for
     #   multi-way replacement?
-    cur_fwd = lambda x: x
     c = fwd(c)
+    todos = []
     for rd in match_pattern(c, pat, use_sym_id=use_sym_id):
+        if c_repl := repl(rd):
+            todos.append((rd, c_repl))
+
+    cur_fwd = lambda x: x
+    for (rd, c_repl) in todos:
         rd = cur_fwd(rd)
-        if not (c_repl := repl(rd)):
-            continue
         ir, fwd_rd = _replace_helper(rd, c_repl, only_replace_attrs)
         cur_fwd = _compose(fwd_rd, cur_fwd)
     return ir, _compose(cur_fwd, fwd)
 
 
 def _replace_reads(ir, fwd, c, sym, repl, only_replace_attrs=True):
-    cur_fwd = lambda x: x
     c = fwd(c)
+    todos = []
     for rd in match_pattern(c, f"{repr(sym)}[_]", use_sym_id=True):
         # Need [_] to pattern match against window expressions
+        if c_repl := repl(rd):
+            todos.append((rd, c_repl))
+
+    cur_fwd = lambda x: x
+    for (rd, c_repl) in todos:
         rd = cur_fwd(rd)
-        if not (c_repl := repl(rd)):
-            continue
         ir, fwd_rd = _replace_helper(rd, c_repl, only_replace_attrs)
         cur_fwd = _compose(fwd_rd, cur_fwd)
     return ir, _compose(cur_fwd, fwd)
 
 
-def _replace_writes(ir, fwd, c, sym, repl, only_replace_attrs=True):
-    cur_fwd = lambda x: x
+def _replace_writes(
+    ir, fwd, c, sym, repl, only_replace_attrs=True, match_assign=True, match_reduce=True
+):
     c = fwd(c)
 
     # TODO: Consider optimizing to just one call of [match_pattern]
-    matches = match_pattern(c, f"{repr(sym)} = _", use_sym_id=True) + match_pattern(
-        c, f"{repr(sym)} += _", use_sym_id=True
-    )
+    matches = []
+    if match_assign:
+        matches = match_pattern(c, f"{repr(sym)} = _", use_sym_id=True)
+    if match_reduce:
+        matches = matches + match_pattern(c, f"{repr(sym)} += _", use_sym_id=True)
+
+    todos = []
     for block in matches:
         assert len(block) == 1  # match_pattern on stmts return blocks
-        s = cur_fwd(block[0])
-        if not (c_repl := repl(s)):
-            continue
+        s = block[0]
+        if c_repl := repl(s):
+            todos.append((s, c_repl))
+
+    cur_fwd = lambda x: x
+    for (s, c_repl) in todos:
+        s = cur_fwd(s)
         ir, fwd_s = _replace_helper(s, c_repl, only_replace_attrs)
         cur_fwd = _compose(fwd_s, cur_fwd)
+
     return ir, _compose(cur_fwd, fwd)
 
 
@@ -3818,7 +3836,6 @@ def DoFoldBuffer(alloc_cursor, dim_idx, new_size):
 
 
 def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False):
-    proc = block_cursor.get_root()
     new_name = Sym(new_name)
 
     def get_typ_mem():
@@ -3869,7 +3886,9 @@ def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False):
                 pt = LoopIR.BinOp("-", w.pt, off, T.index, w.srcinfo)
                 return LoopIR.Point(pt, w.srcinfo)
 
-        return [off_w(w_i, w_e[0]) for w_i, w_e in zip(w_idx, w_exprs)]
+        w_los = [w_e[0] if isinstance(w_e, tuple) else w_e for w_e in w_exprs]
+
+        return [off_w(w_i, w_e) for w_i, w_e in zip(w_idx, w_los)]
 
     ir = block_cursor.get_root()
     block = [s._node for s in block_cursor]
@@ -3938,8 +3957,68 @@ def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False):
 
         return ir, fwd
 
-    isR, isW = Check_BufferRW(ir, block, buf_name, n_dims)
-    if isR:
+    def idx_contained_by_window(idx, block_cursor):
+        """
+        Returns True if idx always lies in staged window range.
+        Returns False if idx never lies in staged window range.
+        Otherwise, will raise a SchedulingError.
+        """
+        p = idx.get_root()
+        return Check_Access_In_Window(p, idx, w_exprs, block_cursor)
+
+    actualR = actualW = False
+    WShadow = False
+    # Conservatively, shadowing logic only works for single element staging windows.
+    w_is_pt = all(not isinstance(w, tuple) for w in w_exprs)
+
+    def mk_read(c, block_cursor):
+        nonlocal actualR
+        rd = c._node
+
+        if isinstance(rd, LoopIR.Read):
+            if idx_contained_by_window(c, block_cursor):
+                _idx = rewrite_idx(rd.idx)
+                actualR = True
+                return {"name": new_name, "idx": _idx}
+        elif isinstance(rd, LoopIR.WindowExpr):
+            if any(
+                isinstance(w, LoopIR.Interval) and not isinstance(w_e, tuple)
+                for w, w_e in zip(rd.idx, w_exprs)
+            ):
+                raise SchedulingError(
+                    f"Existing WindowExpr {rd} has a widnowed dimension which is not windowed in the new staged window."
+                )
+
+            if idx_contained_by_window(c, block_cursor):
+                _idx = rewrite_win(rd.idx)
+                _typ = T.Window(new_typ, rd.type.as_tensor, new_name, _idx)
+                actualR = True
+                return {"name": new_name, "idx": _idx, "type": _typ}
+
+    def mk_write(c, block_cursor):
+        nonlocal actualR
+        nonlocal actualW
+        nonlocal WShadow
+        s = c._node
+        if isinstance(s, (LoopIR.Assign, LoopIR.Reduce)):
+            if idx_contained_by_window(c, block_cursor):
+                actualW = True
+                if isinstance(s, LoopIR.Reduce):
+                    actualR = True
+                if not actualR and w_is_pt:
+                    WShadow = True
+                return {"name": new_name, "idx": rewrite_idx(s.idx)}
+
+    for c in block_cursor:
+        ir, fwd = _replace_reads(
+            ir, fwd, c, buf_name, partial(mk_read, block_cursor=fwd(block_cursor))
+        )
+
+        ir, fwd = _replace_writes(
+            ir, fwd, c, buf_name, partial(mk_write, block_cursor=fwd(block_cursor))
+        )
+
+    if actualR and not WShadow:
         load_iter = [Sym(f"i{i}") for i, _ in enumerate(shape)]
         load_widx = [LoopIR.Read(s, [], T.index, srcinfo) for s in load_iter]
         if use_accum_zero:
@@ -3980,7 +4059,8 @@ def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False):
             ir, fwd = insert_safety_guards(
                 ir, fwd, get_inner_stmt(load_nest_c), load_rhs, buf_typ
             )
-    if isW:
+
+    if actualW:
         store_iter = [Sym(f"i{i}") for i, _ in enumerate(shape)]
         store_ridx = [LoopIR.Read(s, [], T.index, srcinfo) for s in store_iter]
         cp_store_ridx = store_ridx.copy()
@@ -4020,38 +4100,19 @@ def DoStageMem(block_cursor, buf_name, w_exprs, new_name, use_accum_zero=False):
             ir, fwd, store_stmt_c, store_stmt_c._node, buf_typ
         )
 
-    def mk_read(c):
-        rd = c._node
-        if isinstance(rd, LoopIR.Read):
-            return {
-                "name": new_name,
-                "idx": rewrite_idx(rd.idx),
-                "type": rd.type,  # non-ideal, but easiest for now
-            }
-        elif isinstance(rd, LoopIR.WindowExpr):
-            w_idx = rewrite_win(rd.idx)
-            return {
-                "name": new_name,
-                "idx": w_idx,
-                "type": T.Window(new_typ, rd.type.as_tensor, new_name, w_idx),
-            }
-
-    def mk_write(c):
-        s = c._node
-        return {"name": new_name, "idx": rewrite_idx(s.idx)}
-
-    for c in block_cursor:
-        ir, fwd = _replace_reads(ir, fwd, c, buf_name, mk_read)
-        ir, fwd = _replace_writes(ir, fwd, c, buf_name, mk_write)
-
     # new alloc, load_nest + new_body + store_nest
     new_block_c = fwd(block_cursor[0]).as_block().expand(0, len(block_cursor) - 1)
-    if isR:
+    if actualR and not WShadow:
         new_block_c = new_block_c.expand(1, 0)
-    if isW:
+    if actualW:
         new_block_c = new_block_c.expand(0, 1)
-    alloc_c = new_block_c[0].prev()
-    Check_Bounds(ir, alloc_c._node, [c._node for c in new_block_c])
+    if not actualR and not actualW:
+        raise SchedulingError(
+            f"Cannot stage '{buf_name}' with the given window shape. Wrong window shape, or '{buf_name}' not accessed in the given scope?"
+        )
+
+    Check_Bounds(ir, new_alloc[0], [c._node for c in new_block_c])
+
     return ir, fwd
 
 
diff --git a/src/exo/internal_cursors.py b/src/exo/internal_cursors.py
index 182c9e9..2694d05 100644
--- a/src/exo/internal_cursors.py
+++ b/src/exo/internal_cursors.py
@@ -210,6 +210,9 @@ class Block(Cursor):
     def parent(self) -> Node:
         return self._anchor
 
+    def depth(self) -> int:
+        return self._anchor.depth()
+
     def before(self) -> Gap:
         return self[0].before()
 
@@ -611,6 +614,9 @@ class Node(Cursor):
             raise InvalidCursorError("cursor does not have a parent")
         return Node(self._root, self._path[:-1])
 
+    def depth(self) -> int:
+        return len(self._path)
+
     def before(self) -> Gap:
         return Gap(self._root, self, GapType.Before)
 
@@ -772,6 +778,9 @@ class Gap(Cursor):
             return self._anchor
         return self._anchor.parent()
 
+    def depth(self) -> int:
+        return self._anchor.depth()
+
     def anchor(self) -> Node:
         return self._anchor
 
diff --git a/src/exo/new_eff.py b/src/exo/new_eff.py
index 10da1d4..f2a0de4 100644
--- a/src/exo/new_eff.py
+++ b/src/exo/new_eff.py
@@ -731,7 +731,6 @@ def _lsstr(ls, prec=0):
         typ = f":{ls.type}" if ls.type else ""
         return f"{{({ls.name},{coords}{typ})}}"
     elif isinstance(ls, LS.WholeBuf):
-        coords = ",".join([str(a) for a in ls.coords])
         return f"{{{ls.name}|{ls.ndim}}}"
     elif isinstance(ls, (LS.Union, LS.Isct, LS.Diff)):
         if isinstance(ls, LS.Union):
@@ -1104,6 +1103,27 @@ def get_changing_globset(env):
 # Extraction of Effects from programs
 
 
+def window_effs(e):
+    eff_access = []
+    syms = {}
+    for i, w in enumerate(e.idx):
+        if isinstance(w, LoopIR.Interval):
+            syms[i] = Sym(f"EXO_EFFECTS_WINDOW_TEMP_INDEX_{i}")
+            eff_access.append(lift_e(syms[i]))
+        else:
+            eff_access.append(lift_e(w.pt))
+
+    eff = [E.Read(e.name, [idx for idx in eff_access])]
+
+    for i, w in enumerate(e.idx):
+        if isinstance(w, LoopIR.Interval):
+            sym = syms[i]
+            bds = AAnd(lift_e(w.lo) <= AInt(sym), AInt(sym) < lift_e(w.hi))
+            eff = E.Loop(syms[i][E.Guard(bds, eff)])
+
+    return eff
+
+
 def expr_effs(e):
     if isinstance(e, LoopIR.Read):
         if e.type.is_numeric():
@@ -1963,7 +1983,7 @@ def Check_ExprEqvInContext(proc, expr0, stmts0, expr1, stmts1=None):
         raise SchedulingError(f"Expressions are not equivalent:\n{expr0}\nvs.\n{expr1}")
 
 
-def Check_BufferRW(proc, stmts, buf, ndim):
+def Check_BufferReduceOnly(proc, stmts, buf, ndim):
     assert len(stmts) > 0
     ctxt = ContextExtraction(proc, stmts)
 
@@ -1976,40 +1996,101 @@ def Check_BufferRW(proc, stmts, buf, ndim):
 
     wholebuf = LS.WholeBuf(buf, ndim)
     a = G(stmts_effs(stmts))
-    Mod, Rd, Red = getsets([ES.MODIFY, ES.READ_H, ES.REDUCE], a)
-    write = LIsct(wholebuf, Mod)
-    read = LIsct(wholebuf, LUnion(Rd, Red))
+    RW = getsets([ES.READ_WRITE], a)[0]
+    readwrite = LIsct(wholebuf, RW)
 
-    no_read = slv.verify(ADef(is_empty(read)))
-    no_write = slv.verify(ADef(is_empty(write)))
+    no_rw = slv.verify(ADef(is_empty(readwrite)))
     slv.pop()
+    if not no_rw:
+        raise SchedulingError(
+            f"The buffer {buf} is accessed in a way other than "
+            f"simply reducing into it"
+        )
 
-    return (not no_read), (not no_write)
 
+# TODO: I think idxs should be passed as either a read, window, or write (assign/reduce)
+def Check_Access_In_Window(proc, access_cursor, w_exprs, block_cursor):
+    """
+    Returns True if idxs always lies within w_exprs
+    Returns False if idxs never lies within w_exprs
+    Raises a SchedulingError otherwise
 
-def Check_BufferReduceOnly(proc, stmts, buf, ndim):
-    assert len(stmts) > 0
-    ctxt = ContextExtraction(proc, stmts)
+    block_cursor is the context in which to interpret the access in.
+    """
 
+    access = access_cursor._node
+    block = [x._node for x in block_cursor]
+    idxs = access.idx
+    assert len(idxs) == len(w_exprs)
+
+    ctxt = ContextExtraction(proc, block)
     p = ctxt.get_control_predicate()
-    G = ctxt.get_pre_globenv()
 
     slv = SMTSolver(verbose=False)
     slv.push()
     slv.assume(AMay(p))
 
-    wholebuf = LS.WholeBuf(buf, ndim)
-    a = G(stmts_effs(stmts))
-    RW = getsets([ES.READ_WRITE], a)[0]
-    readwrite = LIsct(wholebuf, RW)
+    # build a location set describing the allocated region of the buffer
+    name = access.name
+
+    def get_locset(w_idxs):
+        coords = [Sym(f"i{i}") for i, _ in enumerate(w_idxs)]
+        bounds = []
+        for i, w_access in zip(coords, w_idxs):
+            # Need four cases because accesses and w_exprs have different representations
+            if isinstance(w_access, tuple):
+                (w_lo, w_hi) = w_access
+                bounds.append(AAnd(lift_e(w_lo) <= AInt(i), AInt(i) < lift_e(w_hi)))
+            elif isinstance(w_access, LoopIR.Interval):
+                w_lo = w_access.lo
+                w_hi = w_access.hi
+                bounds.append(AAnd(lift_e(w_lo) <= AInt(i), AInt(i) < lift_e(w_hi)))
+            elif isinstance(w_access, LoopIR.Point):
+                bounds.append(AEq(AInt(i), lift_e(w_access.pt)))
+            else:
+                bounds.append(AEq(AInt(i), lift_e(w_access)))
+        bounds = AAnd(*bounds)
 
-    no_rw = slv.verify(ADef(is_empty(readwrite)))
-    slv.pop()
-    if not no_rw:
-        raise SchedulingError(
-            f"The buffer {buf} is accessed in a way other than "
-            f"simply reducing into it"
-        )
+        pt = LS.Point(name, [AInt(i) for i in coords], T.index)
+        locset = LFilter(bounds, pt)
+        for i in reversed(coords):
+            locset = LBigUnion(i, locset)
+        return locset
+
+    window_locset = get_locset(w_exprs)
+
+    access_locset = get_locset(idxs)
+
+    # Surround access_locset with the appropriate bounds based on the context block_cursor.
+    cursor = access_cursor
+    while cursor.depth() != block_cursor.depth():
+        if isinstance(cursor._node, LoopIR.For):
+            loop = cursor._node
+            bounds = AAnd(
+                lift_e(loop.lo) <= AInt(loop.iter), AInt(loop.iter) < lift_e(loop.hi)
+            )
+            access_locset = LFilter(bounds, access_locset)
+        elif isinstance(cursor._node, LoopIR.If):
+            cond = lift_e(cursor._node.cond)
+            if cursor._path[-1][0] == "orelse":
+                cond = ANot(lift_e(cursor._node.cond))
+            access_locset = LFilter(cond, access_locset)
+
+        cursor = cursor.parent()
+
+    if slv.verify(ADef(is_empty(LIsct(access_locset, window_locset)))):
+        # access_locset is disjoint from window_set
+        slv.pop()
+        return False
+
+    if slv.verify(ADef(is_empty(LDiff(access_locset, window_locset)))):
+        # access_locset is a subset of window_set
+        slv.pop()
+        return True
+
+    raise SchedulingError(
+        f"Buffer has accesses which are neither fully within nor disjoint from the window"
+    )
 
 
 def Check_Bounds(proc, alloc_stmt, block):


Test patch:
diff --git a/tests/golden/test_schedules/test_stage_mem_assign.txt b/tests/golden/test_schedules/test_stage_mem_assign.txt
new file mode 100644
index 0000000..45f67a5
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_assign.txt
@@ -0,0 +1,6 @@
+def foo(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+    result = 0.0
+    for i in seq(0, n):
+        tile: f32 @ DRAM
+        tile = x[i]
+        result = tile + tile
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_stage_mem_assign2.txt b/tests/golden/test_schedules/test_stage_mem_assign2.txt
new file mode 100644
index 0000000..b134638
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_assign2.txt
@@ -0,0 +1,7 @@
+def foo(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+    result = 0.0
+    for i in seq(0, n):
+        tile: f32 @ DRAM
+        tile = x[i]
+        result = tile
+        result = tile
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_stage_mem_asum.txt b/tests/golden/test_schedules/test_stage_mem_asum.txt
new file mode 100644
index 0000000..36dde00
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_asum.txt
@@ -0,0 +1,6 @@
+def asum(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+    result = 0.0
+    for i in seq(0, n):
+        tile: f32 @ DRAM
+        tile = x[i]
+        result += select(0.0, tile, tile, -tile)
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_stage_mem_okay.txt b/tests/golden/test_schedules/test_stage_mem_okay.txt
new file mode 100644
index 0000000..b08ea0f
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_okay.txt
@@ -0,0 +1,12 @@
+def foo(x: i8[10, 10, 10] @ DRAM):
+    y: i8 @ DRAM
+    x_tmp: i8[10, 10] @ DRAM
+    for i0 in seq(0, 10):
+        for i1 in seq(0, 10):
+            x_tmp[i0, i1] = x[i0, 0, i1]
+    for i in seq(0, 10):
+        x_tmp[i, i] = 1.0
+        y = x_tmp[2, 3]
+    for i0 in seq(0, 10):
+        for i1 in seq(0, 10):
+            x[i0, 0, i1] = x_tmp[i0, i1]
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_stage_mem_recursive.txt b/tests/golden/test_schedules/test_stage_mem_recursive.txt
new file mode 100644
index 0000000..0c631a4
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_recursive.txt
@@ -0,0 +1,17 @@
+def recursive(n: size, y: R[n] @ DRAM, x: R[n] @ DRAM):
+    assert n > 2
+    assert (-2 + n) % 4 == 0
+    for io in seq(0, (-2 + n) / 4):
+        y_tmp: R[4] @ DRAM
+        y_tmp[0] = y[1 + 4 * io] + y[4 * io] + x[4 * io]
+        y_tmp[1] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[
+            1 + 4 * io]
+        y_tmp[2] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[
+            1 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io]) + x[2 +
+                                                                      4 * io]
+        y_tmp[3] = y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[
+            1 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io]) + x[
+                2 + 4 * io] + (y[1 + 4 * io] + y[4 * io] + x[4 * io] +
+                               y[1 + 4 * io] + x[1 + 4 * io]) + x[3 + 4 * io]
+        for i0 in seq(0, 4):
+            y[2 + i0 + 4 * io] = y_tmp[i0]
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_stage_mem_reduce.txt b/tests/golden/test_schedules/test_stage_mem_reduce.txt
new file mode 100644
index 0000000..5304c32
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_reduce.txt
@@ -0,0 +1,6 @@
+def foo(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+    result = 0.0
+    for i in seq(0, n):
+        tile: f32 @ DRAM
+        tile = x[i]
+        result += tile + tile
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_stage_mem_reduce2.txt b/tests/golden/test_schedules/test_stage_mem_reduce2.txt
new file mode 100644
index 0000000..0fb0601
--- /dev/null
+++ b/tests/golden/test_schedules/test_stage_mem_reduce2.txt
@@ -0,0 +1,6 @@
+def foo(x: f32[30] @ DRAM, result: f32 @ DRAM):
+    tmp: f32 @ DRAM
+    tmp = 0.0
+    for i in seq(0, 30):
+        x[i] = tmp
+    result = tmp
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index c1308dc..eddc1a5 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -2847,10 +2847,55 @@ def test_fail_stage_mem():
                                     * B[4 * k + kk, 4 * j + jj]
                                 )
 
-    with pytest.raises(SchedulingError, match="accessed out-of-bounds"):
+    with pytest.raises(
+        SchedulingError,
+        match="Buffer has accesses which are neither fully within nor disjoint from the window",
+    ):
         sqmat = stage_mem(sqmat, "for ii in _: _", "B[4*i:4*i+4, 4*k:4*k+4]", "Btile")
 
 
+def test_stage_mem_recursive(golden):
+    @proc
+    def recursive(n: size, y: R[n] @ DRAM, x: R[n] @ DRAM):
+        assert n > 2
+        assert (-2 + n) % 4 == 0
+        for io in seq(0, (-2 + n) / 4):
+            y[2 + 4 * io] = y[1 + 4 * io] + y[4 * io] + x[4 * io]
+            y[3 + 4 * io] = (
+                y[1 + 4 * io] + y[4 * io] + x[4 * io] + y[1 + 4 * io] + x[1 + 4 * io]
+            )
+            y[4 + 4 * io] = (
+                y[1 + 4 * io]
+                + y[4 * io]
+                + x[4 * io]
+                + y[1 + 4 * io]
+                + x[1 + 4 * io]
+                + (y[1 + 4 * io] + y[4 * io] + x[4 * io])
+                + x[2 + 4 * io]
+            )
+            y[5 + 4 * io] = (
+                y[1 + 4 * io]
+                + y[4 * io]
+                + x[4 * io]
+                + y[1 + 4 * io]
+                + x[1 + 4 * io]
+                + (y[1 + 4 * io] + y[4 * io] + x[4 * io])
+                + x[2 + 4 * io]
+                + (
+                    y[1 + 4 * io]
+                    + y[4 * io]
+                    + x[4 * io]
+                    + y[1 + 4 * io]
+                    + x[1 + 4 * io]
+                )
+                + x[3 + 4 * io]
+            )
+
+    block = recursive.find_loop("io").body()
+    recursive = stage_mem(recursive, block, "y[2+4*io : 6+4*io]", "y_tmp")
+    assert str(simplify(recursive)) == golden
+
+
 def test_stage_mem_twice(golden):
     # This test now finds a way to stage the buffer B twice
     @proc
@@ -4421,3 +4466,130 @@ def test_parallelize_loop(golden):
 
     foo = parallelize_loop(foo, foo.find_loop("i"))
     assert str(foo) == golden
+
+
+def test_stage_mem_should_fail():
+    @proc
+    def foo(x: i8[10, 10, 10]):
+        for i in seq(0, 10):
+            x[i, i, i] = 1.0
+
+    with pytest.raises(SchedulingError, match="Buffer has accesses"):
+        foo = stage_mem(foo, foo.find_loop("i"), "x[0:10, 0:2, 0:10]", "x_tmp")
+
+
+def test_stage_mem_should_fail2():
+    @proc
+    def foo(x: i8[10, 10, 10]):
+        y: i8
+        for i in seq(0, 10):
+            y = x[i, i, i]
+
+    with pytest.raises(SchedulingError, match="Buffer has accesses"):
+        foo = stage_mem(foo, foo.find_loop("i"), "x[0:10, 0, 0:10]", "x_tmp")
+
+
+def test_stage_mem_should_fail3():
+    @proc
+    def foo(x: i8[10, 10, 10], y: i8[10]):
+        for i in seq(0, 10):
+            x[i, i, i] = 0.0
+
+    with pytest.raises(SchedulingError, match="Cannot stage"):
+        foo = stage_mem(foo, foo.find_loop("i"), "y[0:10]", "y_tmp")
+
+
+def test_stage_mem_okay(golden):
+    @proc
+    def foo(x: i8[10, 10, 10]):
+        y: i8
+        for i in seq(0, 10):
+            x[i, 0, i] = 1.0
+            y = x[2, 0, 3]
+
+    foo = stage_mem(foo, foo.find_loop("i"), "x[0:10, 0, 0:10]", "x_tmp")
+    assert str(simplify(foo)) == golden
+
+
+def test_stage_mem_should_fail4():
+    @proc
+    def foo(N: size, A: i8[N, N]):
+        sum_: i8
+        for i in seq(0, N):
+            for jo in seq(0, N / 4):
+                for ji in seq(0, 4):
+                    A[i, 4 * jo + ji] = 0.0
+                    y = A[i : i + 1, 4 * jo + ji]
+                    sum_ += y[0]
+
+    with pytest.raises(SchedulingError, match="Existing WindowExpr"):
+        foo = stage_mem(foo, "for ji in _:_", "A[i, 4*jo:4*jo+4]", "tile")
+
+
+def test_stage_mem_should_fail5():
+    @proc
+    def foo(N: size, A: i8[N, N]):
+        for i in seq(0, N):
+            for jo in seq(0, N / 4):
+                for ji in seq(0, 4):
+                    A[i, 4 * jo + ji] = 0.0
+                    y = A[i : i + 4, 4 * jo + ji]
+
+    with pytest.raises(SchedulingError, match="Buffer has accesses"):
+        foo = stage_mem(foo, "for ji in _:_", "A[i-1:i+1, 4*jo:4*jo+4]", "tile")
+
+
+def test_stage_mem_asum(golden):
+    @proc
+    def asum(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+        result = 0.0
+        for i in seq(0, n):
+            result += select(0.0, x[i], x[i], -x[i])
+
+    asum = stage_mem(asum, "result += _", "x[i]", "tile")
+    assert str(simplify(asum)) == golden
+
+
+def test_stage_mem_reduce(golden):
+    @proc
+    def foo(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+        result = 0.0
+        for i in seq(0, n):
+            result += x[i] + x[i]
+
+    foo = stage_mem(foo, "result += _", "x[i]", "tile")
+    assert str(simplify(foo)) == golden
+
+
+def test_stage_mem_assign(golden):
+    @proc
+    def foo(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+        result = 0.0
+        for i in seq(0, n):
+            result = x[i] + x[i]
+
+    foo = stage_mem(foo, "result = _ #1", "x[i]", "tile")
+    assert str(simplify(foo)) == golden
+
+
+def test_stage_mem_assign2(golden):
+    @proc
+    def foo(n: size, x: [f32][n] @ DRAM, result: f32 @ DRAM):
+        result = 0.0
+        for i in seq(0, n):
+            result = x[i]
+            result = x[i]
+
+    foo = stage_mem(foo, foo.find("result = _ #1").expand(0, 1), "x[i]", "tile")
+    assert str(simplify(foo)) == golden
+
+
+def test_stage_mem_reduce2(golden):
+    @proc
+    def foo(x: f32[30], result: f32):
+        result = 0.0
+        for i in seq(0, 30):
+            x[i] = result
+
+    foo = stage_mem(foo, foo.body(), "result", "tmp")
+    assert str(simplify(foo)) == golden