From 0189043b7182de1a666a59e0912cef8ebbb0757d Mon Sep 17 00:00:00 2001 From: Kostiantyn Dementiev <68417998+konstde00@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:58:24 +0200 Subject: [PATCH 1/2] Keep the stage_mem load unless the block overwrites the whole window Currently stage_mem skips the initial load whenever the staged block writes to the buffer, even if the block writes only part of the window, so the elements it doesn't write are stored back as garbage (#318). Now the load is skipped only when the block provably overwrites every in-bounds element of the window. A structural check covers loop nests which sweep the window with unit stride (including divide_loop guard tails), other shapes go through the existing location-set analysis, and if the solver can't decide - the load is kept. Added 4 tests: the #318 case, a full overwrite (no load), a conditional single-element write and a compile-and-run check of the #318 case. Fixes #318 --- src/exo/rewrite/LoopIR_scheduling.py | 13 +- src/exo/rewrite/new_eff.py | 172 ++++++++++++++++++ .../test_stage_mem_conditional_overwrite.txt | 7 + .../test_stage_mem_full_overwrite.txt | 6 + .../test_stage_mem_partial_overwrite.txt | 8 + tests/test_schedules.py | 53 ++++++ 6 files changed, 257 insertions(+), 2 deletions(-) create mode 100644 tests/golden/test_schedules/test_stage_mem_conditional_overwrite.txt create mode 100644 tests/golden/test_schedules/test_stage_mem_full_overwrite.txt create mode 100644 tests/golden/test_schedules/test_stage_mem_partial_overwrite.txt diff --git a/src/exo/rewrite/LoopIR_scheduling.py b/src/exo/rewrite/LoopIR_scheduling.py index 1b0950210..5974c5002 100644 --- a/src/exo/rewrite/LoopIR_scheduling.py +++ b/src/exo/rewrite/LoopIR_scheduling.py @@ -27,6 +27,7 @@ Check_BufferReduceOnly, Check_Bounds, Check_Access_In_Window, + Check_Window_Overwritten, Check_IsDeadAfter, Check_IsIdempotent, Check_ExprBound, @@ -4055,7 +4056,15 @@ def mk_write(c, block_cursor): ir, fwd, c, buf_name, partial(mk_write, block_cursor=fwd(block_cursor)) ) - if actualR and not WShadow: + load_needed = actualR and not WShadow + if actualW and not load_needed: + # Without a load, the store below copies the whole staging buffer back, + # so the block must overwrite every (in-bounds) element of the window. + load_needed = not Check_Window_Overwritten( + block_cursor.get_root(), block, buf_name, w_exprs, buf_typ.shape() + ) + + if load_needed: 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: @@ -4133,7 +4142,7 @@ def mk_write(c, block_cursor): # new alloc, load_nest + new_body + store_nest new_block_c = fwd(block_cursor[0]).as_block().expand(0, len(block_cursor) - 1) - if actualR and not WShadow: + if load_needed: new_block_c = new_block_c.expand(1, 0) if actualW: new_block_c = new_block_c.expand(0, 1) diff --git a/src/exo/rewrite/new_eff.py b/src/exo/rewrite/new_eff.py index 9dfa88c72..064c5aacd 100644 --- a/src/exo/rewrite/new_eff.py +++ b/src/exo/rewrite/new_eff.py @@ -2065,6 +2065,178 @@ def get_locset(w_idxs): ) +def _affine_index(e): + """ + Returns the affine form of an index expression as a dictionary from + variables to coefficients (None maps to the constant term), + or None if the expression is not affine + """ + if isinstance(e, LoopIR.Const): + return {None: e.val} + elif isinstance(e, LoopIR.Read) and len(e.idx) == 0: + return {e.name: 1} + elif isinstance(e, LoopIR.USub): + arg = _affine_index(e.arg) + return None if arg is None else {x: -c for x, c in arg.items()} + elif isinstance(e, LoopIR.BinOp) and e.op in ("+", "-", "*"): + lhs = _affine_index(e.lhs) + rhs = _affine_index(e.rhs) + if lhs is None or rhs is None: + return None + if e.op == "*": + if set(lhs) <= {None}: + lhs, rhs = rhs, lhs + if not set(rhs) <= {None}: + return None + return {x: c * rhs.get(None, 0) for x, c in lhs.items()} + sign = 1 if e.op == "+" else -1 + res = dict(lhs) + for x, c in rhs.items(): + res[x] = res.get(x, 0) + sign * c + return res + else: + return None + + +def _affine_eq(a, b): + return all(a.get(x, 0) == b.get(x, 0) for x in set(a) | set(b)) + + +def _is_bounds_guard(cond, idx, shape): + """ + Returns True if cond only excludes locations outside of the buffer, + i.e. it is a conjunction of idx[d] < shape[d] and 0 <= idx[d], + where idx and shape are affine forms + """ + if isinstance(cond, LoopIR.BinOp) and cond.op == "and": + return _is_bounds_guard(cond.lhs, idx, shape) and _is_bounds_guard( + cond.rhs, idx, shape + ) + if not isinstance(cond, LoopIR.BinOp) or cond.op not in ("<", ">", "<=", ">="): + return False + lhs, rhs = _affine_index(cond.lhs), _affine_index(cond.rhs) + if lhs is None or rhs is None: + return False + op = cond.op + if op in (">", ">="): + lhs, rhs, op = rhs, lhs, op.replace(">", "<") + for i, n in zip(idx, shape): + if op == "<" and n is not None and _affine_eq(lhs, i) and _affine_eq(rhs, n): + return True + if op == "<=" and _affine_eq(lhs, {None: 0}) and _affine_eq(rhs, i): + return True + return False + + +def _writes_whole_window(s, buf, w_exprs, shape, loops, guards): + """ + Returns True if executing s writes every in-bounds location of the window + w_exprs of buf with a single assignment, where each dimension of the window + is either a point or swept with unit stride by exactly one loop, and the + assignment is only guarded by checks against the buffer bounds. + loops and guards are the loops and conditions around s inside the block. + """ + if isinstance(s, LoopIR.For): + loop = (s.iter, _affine_index(s.lo), _affine_index(s.hi)) + return any( + _writes_whole_window(b, buf, w_exprs, shape, loops + [loop], guards) + for b in s.body + ) + if isinstance(s, LoopIR.If) and len(s.orelse) == 0: + return any( + _writes_whole_window(b, buf, w_exprs, shape, loops, guards + [s.cond]) + for b in s.body + ) + if not isinstance(s, LoopIR.Assign) or s.name != buf: + return False + + idxs = [_affine_index(i) for i in s.idx] + if any(idx is None for idx in idxs): + return False + shape = [_affine_index(n) for n in shape] + if not all(_is_bounds_guard(g, idxs, shape) for g in guards): + return False + + used = set() + for w, idx in zip(w_exprs, idxs): + swept = [l for l in loops if idx.get(l[0], 0) != 0] + if not isinstance(w, tuple): + # a point of the window must be written exactly + pt = _affine_index(w) + if swept or pt is None or not _affine_eq(idx, pt): + return False + continue + + w_lo, w_hi = _affine_index(w[0]), _affine_index(w[1]) + if w_lo is None or w_hi is None or len(swept) != 1: + return False + it, lo, hi = swept[0] + if idx[it] != 1 or it in used or lo is None or hi is None: + return False + used.add(it) + # idx = it + off, and it sweeps [lo, hi), so [lo + off, hi + off) is written + off = {x: c for x, c in idx.items() if x != it} + lo_off = {x: lo.get(x, 0) + off.get(x, 0) for x in set(lo) | set(off)} + hi_off = {x: hi.get(x, 0) + off.get(x, 0) for x in set(hi) | set(off)} + if not _affine_eq(lo_off, w_lo) or not _affine_eq(hi_off, w_hi): + return False + + # a loop that sweeps no dimension might run zero times + return len(used) == len(loops) + + +def Check_Window_Overwritten(proc, block, buf, w_exprs, shape): + """ + Returns True if the block always writes every in-bounds location of + the window w_exprs of buf (whose dimensions are given by shape) + Returns False otherwise, including when this cannot be proven + """ + assert len(block) > 0 + assert len(w_exprs) == len(shape) + + # fast path for the common shape: a loop nest that assigns to the window + if any(_writes_whole_window(s, buf, w_exprs, shape, [], []) for s in block): + return True + + ctxt = ContextExtraction(proc, block) + + p = ctxt.get_control_predicate() + G = ctxt.get_pre_globenv() + + slv = SMTSolver(verbose=False) + slv.push() + slv.assume(AMay(p)) + + # build a location set describing the window, + # restricted to the allocated region of the buffer + coords = [Sym(f"i{i}") for i, _ in enumerate(w_exprs)] + bounds = [] + for i, w, n in zip(coords, w_exprs, shape): + if isinstance(w, tuple): + w_lo, w_hi = w + bounds.append(AAnd(lift_e(w_lo) <= AInt(i), AInt(i) < lift_e(w_hi))) + else: + bounds.append(AEq(AInt(i), lift_e(w))) + bounds.append(AAnd(AInt(0) <= AInt(i), AInt(i) < lift_e(n))) + pt = LS.Point(buf, [AInt(i) for i in coords], T.index) + window_set = LFilter(AAnd(*bounds), pt) + for i in reversed(coords): + window_set = LBigUnion(i, window_set) + + a = G(stmts_effs(block)) + Wr = getsets([ES.WRITE_ALL], a)[0] + try: + is_overwritten = slv.verify(ADef(is_empty(LDiff(window_set, Wr)))) + except TypeError as e: + # the query quantifies over the window and the writes, so z3 may + # answer "unknown"; treat it as not proven + if "unknown result from z3" not in str(e): + raise + is_overwritten = False + slv.pop() + return is_overwritten + + def Check_Bounds(proc, alloc_stmt, block): if len(block) == 0: return diff --git a/tests/golden/test_schedules/test_stage_mem_conditional_overwrite.txt b/tests/golden/test_schedules/test_stage_mem_conditional_overwrite.txt new file mode 100644 index 000000000..3017122bd --- /dev/null +++ b/tests/golden/test_schedules/test_stage_mem_conditional_overwrite.txt @@ -0,0 +1,7 @@ +def foo(n: size, x: i8[n] @ DRAM): + assert n > 0 + xtmp: i8 @ DRAM + xtmp = x[0] + if n > 4: + xtmp = 0.0 + x[0] = xtmp \ No newline at end of file diff --git a/tests/golden/test_schedules/test_stage_mem_full_overwrite.txt b/tests/golden/test_schedules/test_stage_mem_full_overwrite.txt new file mode 100644 index 000000000..c89d3beaf --- /dev/null +++ b/tests/golden/test_schedules/test_stage_mem_full_overwrite.txt @@ -0,0 +1,6 @@ +def foo(x: i8[10] @ DRAM): + xtmp: i8[8 - 0] @ DRAM + for i in seq(0, 8): + xtmp[i - 0] = 0.0 + for i0 in seq(0, 8 - 0): + x[i0 + 0] = xtmp[i0] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_stage_mem_partial_overwrite.txt b/tests/golden/test_schedules/test_stage_mem_partial_overwrite.txt new file mode 100644 index 000000000..4fea598d4 --- /dev/null +++ b/tests/golden/test_schedules/test_stage_mem_partial_overwrite.txt @@ -0,0 +1,8 @@ +def foo(x: i8[10] @ DRAM): + xtmp: i8[9 - 0] @ DRAM + for i0 in seq(0, 9 - 0): + xtmp[i0] = x[i0 + 0] + for i in seq(0, 8): + xtmp[i - 0] = 0.0 + for i0 in seq(0, 9 - 0): + x[i0 + 0] = xtmp[i0] \ No newline at end of file diff --git a/tests/test_schedules.py b/tests/test_schedules.py index 26594ba57..459ba6c66 100644 --- a/tests/test_schedules.py +++ b/tests/test_schedules.py @@ -1,5 +1,6 @@ from __future__ import annotations +import numpy as np import pytest from exo import ParseFragmentError, proc, DRAM, Procedure, config @@ -3168,6 +3169,58 @@ def foo(n: size, m: size, x: f32[n], y: f32[n]): assert str(foo) == golden +def test_stage_mem_partial_overwrite(golden): + # The block only writes x[0:8], so the window x[0:9] must be loaded first, + # otherwise the store copies an uninitialized value back into x[8] + @proc + def foo(x: i8[10]): + for i in seq(0, 8): + x[i] = 0.0 + + foo = stage_mem(foo, foo.body(), "x[0:9]", "xtmp") + assert str(foo) == golden + + +def test_stage_mem_full_overwrite(golden): + # The block writes the whole window, so no load is needed + @proc + def foo(x: i8[10]): + for i in seq(0, 8): + x[i] = 0.0 + + foo = stage_mem(foo, foo.body(), "x[0:8]", "xtmp") + assert str(foo) == golden + + +def test_stage_mem_conditional_overwrite(golden): + # A write under a condition does not always overwrite the window + @proc + def foo(n: size, x: i8[n]): + assert n > 0 + if n > 4: + x[0] = 0.0 + + foo = stage_mem(foo, foo.body(), "x[0]", "xtmp") + assert str(foo) == golden + + +def test_stage_mem_partial_overwrite_runs(compiler): + @proc + def foo(x: i8[10]): + for i in seq(0, 8): + x[i] = 0.0 + + foo_staged = rename(stage_mem(foo, foo.body(), "x[0:9]", "xtmp"), "foo_staged") + lib = compiler.compile([foo, foo_staged]) + + expected = np.arange(1, 11, dtype=np.int8) + actual = expected.copy() + lib.foo(None, expected) + lib.foo_staged(None, actual) + + np.testing.assert_array_equal(actual, expected) + + def test_new_expr_multi_vars(golden): @proc def bar(n: size, arr: R[n] @ DRAM): From 2fb5ac05787fa7e168e992d6504374ddce6ebdba Mon Sep 17 00:00:00 2001 From: Kostiantyn Dementiev <68417998+konstde00@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:14:17 +0200 Subject: [PATCH 2/2] Add tests for guarded and reversed overwrite cases Two more stage_mem cases: a block under a divide_loop style bounds guard, which overwrites the whole window (no load), and a reversed write which covers only part of the window (the load is kept). Also a direct test of Check_Window_Overwritten for the case when the solver can't decide - the window is treated as not overwritten. --- .../test_stage_mem_guarded_overwrite.txt | 9 +++++++ ...t_stage_mem_reversed_partial_overwrite.txt | 8 +++++++ tests/test_new_eff.py | 20 ++++++++++++++++ tests/test_schedules.py | 24 +++++++++++++++++++ 4 files changed, 61 insertions(+) create mode 100644 tests/golden/test_schedules/test_stage_mem_guarded_overwrite.txt create mode 100644 tests/golden/test_schedules/test_stage_mem_reversed_partial_overwrite.txt diff --git a/tests/golden/test_schedules/test_stage_mem_guarded_overwrite.txt b/tests/golden/test_schedules/test_stage_mem_guarded_overwrite.txt new file mode 100644 index 000000000..804cc6770 --- /dev/null +++ b/tests/golden/test_schedules/test_stage_mem_guarded_overwrite.txt @@ -0,0 +1,9 @@ +def foo(n: size, x: i8[n] @ DRAM): + for io in seq(0, (n + 3) / 4): + xtmp: i8[4 * io + 4 - 4 * io] @ DRAM + for ii in seq(0, 4): + if n > ii + 4 * io and 0 <= ii + 4 * io: + xtmp[ii + 4 * io - 4 * io] = 0.0 + for i0 in seq(0, 4 * io + 4 - 4 * io): + if i0 + 4 * io < n: + x[i0 + 4 * io] = xtmp[i0] \ No newline at end of file diff --git a/tests/golden/test_schedules/test_stage_mem_reversed_partial_overwrite.txt b/tests/golden/test_schedules/test_stage_mem_reversed_partial_overwrite.txt new file mode 100644 index 000000000..6ed9866e9 --- /dev/null +++ b/tests/golden/test_schedules/test_stage_mem_reversed_partial_overwrite.txt @@ -0,0 +1,8 @@ +def foo(x: i8[16] @ DRAM): + xtmp: i8[16 - 0] @ DRAM + for i0 in seq(0, 16 - 0): + xtmp[i0] = x[i0 + 0] + for i in seq(0, 8): + xtmp[-i + 15 - 0] = 0.0 + for i0 in seq(0, 16 - 0): + x[i0 + 0] = xtmp[i0] \ No newline at end of file diff --git a/tests/test_new_eff.py b/tests/test_new_eff.py index 7f079e4a7..4f14f3bb3 100644 --- a/tests/test_new_eff.py +++ b/tests/test_new_eff.py @@ -37,6 +37,26 @@ def test_debug_let_and_mod(): slv.verify(F) +def test_window_overwritten_unknown_is_not_proven(monkeypatch): + @proc + def foo(x: i8[10]): + x[0] = 0.0 + x[1] = 0.0 + + ir = foo._loopir_proc + x = ir.args[0] + srcinfo = ir.body[0].srcinfo + window = [(LoopIR.Const(0, T.index, srcinfo), LoopIR.Const(2, T.index, srcinfo))] + assert Check_Window_Overwritten(ir, ir.body, x.name, window, x.type.shape()) + + # if the solver cannot decide, the window does not count as overwritten + def unknown(self, e): + raise TypeError("unknown result from z3") + + monkeypatch.setattr(SMTSolver, "verify", unknown) + assert not Check_Window_Overwritten(ir, ir.body, x.name, window, x.type.shape()) + + def test_reorder_stmts_fail(): @proc def foo(N: size, x: R[N]): diff --git a/tests/test_schedules.py b/tests/test_schedules.py index 459ba6c66..10c184e74 100644 --- a/tests/test_schedules.py +++ b/tests/test_schedules.py @@ -3204,6 +3204,30 @@ def foo(n: size, x: i8[n]): assert str(foo) == golden +def test_stage_mem_guarded_overwrite(golden): + # Guards that only skip out-of-bounds elements still overwrite the window + @proc + def foo(n: size, x: i8[n]): + for io in seq(0, (n + 3) / 4): + for ii in seq(0, 4): + if n > ii + 4 * io and 0 <= ii + 4 * io: + x[ii + 4 * io] = 0.0 + + foo = stage_mem(foo, foo.find_loop("ii"), "x[4 * io:4 * io + 4]", "xtmp") + assert str(foo) == golden + + +def test_stage_mem_reversed_partial_overwrite(golden): + # The block only writes x[8:16], so the window x[0:16] must be loaded first + @proc + def foo(x: i8[16]): + for i in seq(0, 8): + x[-i + 15] = 0.0 + + foo = stage_mem(foo, foo.body(), "x[0:16]", "xtmp") + assert str(foo) == golden + + def test_stage_mem_partial_overwrite_runs(compiler): @proc def foo(x: i8[10]):