[Caching] Fix per-construct split false fallback on ndarray-vs-non-ndarray mixed pointers

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<ExternalPtrStmt>())
+      return e;
+    if (auto *mp = p->cast<MatrixPtrStmt>()) {
+      p = mp->origin;
+      continue;
+    }
     return nullptr;
-  if (auto *e = p->cast<ExternalPtrStmt>())
-    return e;
-  if (auto *mp = p->cast<MatrixPtrStmt>())
-    return mp->origin != nullptr ? mp->origin->cast<ExternalPtrStmt>() : nullptr;
+  }
   return nullptr;
 }
 
@@ -615,21 +617,19 @@ 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<ExternalPtrStmt>() && b != nullptr && b->is<MatrixPtrStmt>()) ||
                      (a != nullptr && a->is<MatrixPtrStmt>() && b != nullptr && b->is<ExternalPtrStmt>());
   if (!mixed)
     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);
 }
 

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