From d37093f87eaa753d5e2338243a3a17a150ba1e87 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 3 Sep 2026 11:59:07 -0700 Subject: [PATCH 1/2] [Caching] Fix per-construct split false fallback on ndarray-vs-non-ndarray mixed ptrs The whole-element/component recompute-safety guard (added to close a real alias_analysis blind spot) had a conservative catch-all that rejected ANY mixed ExternalPtr/MatrixPtr pair whose matrix side did not normalize to an external pointer. But a MatrixPtr over a local alloca, a global temp, or a field is a different memory space than an ndarray read, and alias_analysis already reports such a pair `different` (checked before this guard runs). Overriding that to "may overlap" disabled the split for any kernel that recomputes a whole-element ndarray read past a component write to a local/temp/field -- including qipc's giant graph_do_while `_step_kernel`, which regressed warm+1-edit time-to-first- step from ~11s back to ~27s (whole-kernel recompile). Fix: `as_ndarray_ptr` now follows the matrix-ptr chain to the underlying ndarray external pointer (handling higher-rank elements), and the guard only overrides alias_analysis when BOTH sides are ndarray-backed; otherwise it defers to alias_analysis's verdict. The genuine whole-element-vs-component ndarray hazard (same ndarray, mixed access) still falls back. Validated on a Blackwell node: new regression test fails before / passes after; the true-positive test still falls back; full test_per_offload_cache.py suite 78 passed / 3 skipped (cpu+cuda); qipc warm+1-edit 26.8s -> 11.6s; qipc test_abd_freefall.py 4 passed (split output correct). --- .../split_frontend_per_construct.cpp | 22 +++++++++++---- tests/python/test_per_offload_cache.py | 28 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/quadrants/transforms/split_frontend_per_construct.cpp b/quadrants/transforms/split_frontend_per_construct.cpp index 574761bc26..c7467d7b05 100644 --- a/quadrants/transforms/split_frontend_per_construct.cpp +++ b/quadrants/transforms/split_frontend_per_construct.cpp @@ -566,13 +566,19 @@ bool internal_func_is_memory_free(const std::string &name) { // The ndarray (external) access a load/store pointer resolves to, directly or through a MatrixPtr element; nullptr if // it isn't one. +// The ndarray external pointer this pointer ultimately indexes into, following the matrix-ptr chain (a component +// access `a[i][c]` is a MatrixPtrStmt over the element's ExternalPtrStmt; higher-rank elements nest further). Null when +// the base is not an ndarray -- a local alloca, a global temp, or a field -- which lives in a different memory space. ExternalPtrStmt *as_ndarray_ptr(Stmt *p) { - if (p == nullptr) + while (p != nullptr) { + if (auto *e = p->cast()) + return e; + if (auto *mp = p->cast()) { + p = mp->origin; + continue; + } return nullptr; - if (auto *e = p->cast()) - return e; - if (auto *mp = p->cast()) - return mp->origin != nullptr ? mp->origin->cast() : nullptr; + } return nullptr; } @@ -628,8 +634,12 @@ bool whole_element_read_may_overlap_component_write(Stmt *a, Stmt *b) { return false; ExternalPtrStmt *ea = as_ndarray_ptr(a); ExternalPtrStmt *eb = as_ndarray_ptr(b); + // Only override alias_analysis when both sides are ndarray-backed. A matrix ptr over a non-ndarray (a local alloca, a + // global temp, or a field) is a different memory space than the ndarray read, and alias_analysis already reported the + // pair `different` (checked before this helper runs), so that verdict is authoritative -- this is not the mixed + // whole-element/component ndarray blind spot this helper exists to close. if (ea == nullptr || eb == nullptr) - return true; // a matrix ptr with a non-external origin: cannot prove the element addresses disjoint + return false; return irpass::analysis::maybe_same_address(ea, eb); } diff --git a/tests/python/test_per_offload_cache.py b/tests/python/test_per_offload_cache.py index 0f93a36663..65d3f0bf20 100644 --- a/tests/python/test_per_offload_cache.py +++ b/tests/python/test_per_offload_cache.py @@ -814,6 +814,34 @@ def whole_vs_component(a: qd.types.NDArray[vec2, 1], out: qd.types.ndarray()) -> assert np.allclose(out.to_numpy(), 7.0, atol=1e-2), out.to_numpy() +@test_utils.test(arch=[qd.cpu, qd.cuda], offline_cache=False) +def test_per_construct_frontend_split_whole_element_vs_non_ndarray_component_ok() -> None: + # Companion to the test above: the intervening component write targets a *field* `f[i][0]` -- a MatrixPtrStmt whose + # origin is a GlobalPtrStmt, not an ndarray -- so it lives in a different memory space than the ndarray read `s[0]` + # and cannot alias it. alias_analysis already reports the pair different; the whole-element/component guard must + # defer to that and let the split fire, rather than rejecting every mixed ExternalPtr/MatrixPtr pair. The same + # blind spot (a MatrixPtr over a local alloca) disabled the split for qipc's giant `_step_kernel`, so this is the + # regression guard for that fix. + f = qd.Vector.field(2, qd.f32, shape=(_N,)) + + @qd.kernel + def whole_vs_field_component(s: qd.types.ndarray(), out: qd.types.ndarray()) -> None: + base = s[0] # whole-element ndarray read, recomputed into construct 2 + for i in range(_N): # construct 1: component write to a field (a different memory space than the ndarray) + f[i][0] = 2.0 + for i in range(out.shape[0]): # construct 2: reuse the snapshot + out[i] = base + + s = qd.ndarray(qd.f32, shape=(_N,)) + out = qd.ndarray(qd.f32, shape=(_N,)) + s.from_numpy(np.arange(_N, dtype=np.float32) + 7.0) + whole_vs_field_component(s, out) + + obs = whole_vs_field_component._primal.per_offload_cache_observations + assert obs.frontend_constructs_total >= 2, obs # split fires: a field write cannot alias the ndarray read + assert np.allclose(out.to_numpy(), 7.0, atol=1e-2), out.to_numpy() # base = s[0] = 7.0 + + @test_utils.test(arch=[qd.cpu, qd.cuda], offline_cache=False) def test_per_construct_frontend_split_fallback_carried_rmw_local() -> None: # Two constructs each read-modify-write the same local `s`, and the second also stores it. The second construct From 99ab47d5687fcaaa50bf9838e761ce55230e48d9 Mon Sep 17 00:00:00 2001 From: Hugh Perkins Date: Thu, 3 Sep 2026 14:13:38 -0700 Subject: [PATCH 2/2] Trim comments on the split alias-guard fix --- .../split_frontend_per_construct.cpp | 20 +++++-------------- tests/python/test_per_offload_cache.py | 14 +++++-------- 2 files changed, 10 insertions(+), 24 deletions(-) diff --git a/quadrants/transforms/split_frontend_per_construct.cpp b/quadrants/transforms/split_frontend_per_construct.cpp index c7467d7b05..d7d4fc5951 100644 --- a/quadrants/transforms/split_frontend_per_construct.cpp +++ b/quadrants/transforms/split_frontend_per_construct.cpp @@ -564,11 +564,7 @@ bool internal_func_is_memory_free(const std::string &name) { return kMemoryFree.count(name) > 0; } -// The ndarray (external) access a load/store pointer resolves to, directly or through a MatrixPtr element; nullptr if -// it isn't one. -// The ndarray external pointer this pointer ultimately indexes into, following the matrix-ptr chain (a component -// access `a[i][c]` is a MatrixPtrStmt over the element's ExternalPtrStmt; higher-rank elements nest further). Null when -// the base is not an ndarray -- a local alloca, a global temp, or a field -- which lives in a different memory space. +// The ndarray access a pointer resolves to, following the MatrixPtr chain. Null when the base is not an ndarray. ExternalPtrStmt *as_ndarray_ptr(Stmt *p) { while (p != nullptr) { if (auto *e = p->cast()) @@ -621,12 +617,9 @@ bool grad_companion_may_alias(Stmt *a, Stmt *b) { return arg_a->arg_id == arg_b->arg_id; } -// alias_analysis compares a whole-element ndarray read (`base = a[i]`, a bare ExternalPtrStmt) against a component -// write to the same element (`a[j][c] = ...`, a MatrixPtrStmt over an ExternalPtrStmt) as `different`, because only the -// component side carries a matrix origin. But a whole-element read covers every component, so it observes such a write -// whenever the two element addresses may coincide. Normalize both to their external origins and re-check: same-arg, -// possibly-same-index -> may-alias. Matrix-vs-matrix and external-vs-external pairs are already precise from the raw -// maybe_same_address check; cross-arg mixed pairs come back `different` here and stay the launch guard's concern. +// alias_analysis calls a whole-element ndarray read `a[i]` and a component write `a[j][c]` `different`, because only +// the write carries a matrix origin. But a whole-element read covers every component, so it can observe such a write +// when the element indices may coincide. Normalize both to their ndarray origins and re-check. bool whole_element_read_may_overlap_component_write(Stmt *a, Stmt *b) { const bool mixed = (a != nullptr && a->is() && b != nullptr && b->is()) || (a != nullptr && a->is() && b != nullptr && b->is()); @@ -634,10 +627,7 @@ bool whole_element_read_may_overlap_component_write(Stmt *a, Stmt *b) { return false; ExternalPtrStmt *ea = as_ndarray_ptr(a); ExternalPtrStmt *eb = as_ndarray_ptr(b); - // Only override alias_analysis when both sides are ndarray-backed. A matrix ptr over a non-ndarray (a local alloca, a - // global temp, or a field) is a different memory space than the ndarray read, and alias_analysis already reported the - // pair `different` (checked before this helper runs), so that verdict is authoritative -- this is not the mixed - // whole-element/component ndarray blind spot this helper exists to close. + // A non-ndarray pointer cannot alias an ndarray, so defer to alias_analysis instead of assuming overlap. if (ea == nullptr || eb == nullptr) return false; return irpass::analysis::maybe_same_address(ea, eb); diff --git a/tests/python/test_per_offload_cache.py b/tests/python/test_per_offload_cache.py index 65d3f0bf20..2ea2484716 100644 --- a/tests/python/test_per_offload_cache.py +++ b/tests/python/test_per_offload_cache.py @@ -816,18 +816,14 @@ def whole_vs_component(a: qd.types.NDArray[vec2, 1], out: qd.types.ndarray()) -> @test_utils.test(arch=[qd.cpu, qd.cuda], offline_cache=False) def test_per_construct_frontend_split_whole_element_vs_non_ndarray_component_ok() -> None: - # Companion to the test above: the intervening component write targets a *field* `f[i][0]` -- a MatrixPtrStmt whose - # origin is a GlobalPtrStmt, not an ndarray -- so it lives in a different memory space than the ndarray read `s[0]` - # and cannot alias it. alias_analysis already reports the pair different; the whole-element/component guard must - # defer to that and let the split fire, rather than rejecting every mixed ExternalPtr/MatrixPtr pair. The same - # blind spot (a MatrixPtr over a local alloca) disabled the split for qipc's giant `_step_kernel`, so this is the - # regression guard for that fix. + # A field write cannot alias an ndarray read, so the split must fire. This guards the fix for qipc's `_step_kernel`, + # where a mixed ndarray-read / matrix-ptr-write pair with a non-ndarray write wrongly forced the whole-kernel path. f = qd.Vector.field(2, qd.f32, shape=(_N,)) @qd.kernel def whole_vs_field_component(s: qd.types.ndarray(), out: qd.types.ndarray()) -> None: - base = s[0] # whole-element ndarray read, recomputed into construct 2 - for i in range(_N): # construct 1: component write to a field (a different memory space than the ndarray) + base = s[0] # recomputed into construct 2 + for i in range(_N): # construct 1: component write to a field f[i][0] = 2.0 for i in range(out.shape[0]): # construct 2: reuse the snapshot out[i] = base @@ -839,7 +835,7 @@ def whole_vs_field_component(s: qd.types.ndarray(), out: qd.types.ndarray()) -> obs = whole_vs_field_component._primal.per_offload_cache_observations assert obs.frontend_constructs_total >= 2, obs # split fires: a field write cannot alias the ndarray read - assert np.allclose(out.to_numpy(), 7.0, atol=1e-2), out.to_numpy() # base = s[0] = 7.0 + assert np.allclose(out.to_numpy(), 7.0, atol=1e-2), out.to_numpy() @test_utils.test(arch=[qd.cpu, qd.cuda], offline_cache=False)