diff --git a/quadrants/transforms/split_frontend_per_construct.cpp b/quadrants/transforms/split_frontend_per_construct.cpp index 574761bc26..d7d4fc5951 100644 --- a/quadrants/transforms/split_frontend_per_construct.cpp +++ b/quadrants/transforms/split_frontend_per_construct.cpp @@ -564,15 +564,17 @@ 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 access a pointer resolves to, following the MatrixPtr chain. Null when the base is not an ndarray. 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; } @@ -615,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()); @@ -628,8 +627,9 @@ 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); + // A non-ndarray pointer cannot alias an ndarray, so defer to alias_analysis instead of assuming overlap. 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..2ea2484716 100644 --- a/tests/python/test_per_offload_cache.py +++ b/tests/python/test_per_offload_cache.py @@ -814,6 +814,30 @@ 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: + # 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] # 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 + + 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() + + @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