[Fixup] Erase ndarray_data_gen on Ndarray destroy + parametrize cache-invalidation test over qd.ndarray and field

Author: duburcqa

quadrants/program/adstack_size_expr_eval.h

@@ -126,6 +126,13 @@ class AdStackCache {
   void bump_ndarray_data_gen(void *devalloc_ptr) {
     ++ndarray_data_gen_[devalloc_ptr];
   }
+  // Drop a per-DeviceAllocation entry. Called from `Ndarray::~Ndarray()` so the holder address can be reused by a
+  // future allocation without inheriting the destroyed ndarray's stale generation. Leftover snapshots in
+  // `per_task_ad_stack_cache_` / `llvm_per_task_ad_stack_cache_` referencing the dropped key fall back to gen=0
+  // on the next lookup (their stored snapshot will not match), which forces a fresh sizer dispatch and self-heals.
+  void erase_ndarray_data_gen(void *devalloc_ptr) {
+    ndarray_data_gen_.erase(devalloc_ptr);
+  }
 
  private:
   std::unordered_map<const SerializedSizeExpr *, SizeExprCacheEntry> size_expr_cache_;

quadrants/program/ndarray.cpp

@@ -98,6 +98,10 @@ Ndarray::Ndarray(DeviceAllocation &devalloc,
 
 Ndarray::~Ndarray() {
   if (prog_) {
+    // Drop any cached `ndarray_data_gen` entry keyed by `&ndarray_alloc_` before the holder address can be reused
+    // by a future Ndarray allocation. Without this, a new Ndarray that happens to occupy the same heap address
+    // would pick up the destroyed ndarray's last generation counter and could falsely match a cache snapshot.
+    prog_->adstack_cache().erase_ndarray_data_gen(const_cast<DeviceAllocation *>(&ndarray_alloc_));
     ndarray_alloc_.device->dealloc_memory(ndarray_alloc_);
   }
 }

tests/python/test_adstack.py

@@ -2089,65 +2089,101 @@ def compute(a: qd.types.ndarray(dtype=qd.i32, ndim=1)):
         assert x.grad[i] == pytest.approx(expected, rel=1e-5)
 
 
+@pytest.mark.parametrize(
+    "trip_count_source",
+    ["qd_ndarray", "field"],
+    ids=["qd_ndarray", "field"],
+)
 @test_utils.test(require=qd.extension.adstack)
-def test_adstack_metadata_cache_invalidates_on_qd_ndarray_host_mutation():
-    # Pins per-task adstack metadata cache invalidation against host-side ndarray writes. A reverse-mode kernel
-    # whose inner trip count is `range(n[i])` over a `qd.ndarray` populates the cache on the first launch with
-    # `max_size = n[i]` evaluated against the current ndarray contents. A subsequent host-side mutation via
-    # `Ndarray::write` (no kernel involvement) must bump `ndarray_data_gen` so the next launch evicts the cache
-    # entry and re-runs the sizer against the new values.
+def test_adstack_metadata_cache_invalidates_on_host_mutation(trip_count_source):
+    # Pins per-task adstack metadata cache invalidation against host-side mutation of the structure that supplies
+    # the inner trip count. A reverse-mode kernel whose inner range is `range(n[i])` populates the cache on the
+    # first launch with `max_size = n[i]` evaluated against the current contents. A subsequent host-side mutation
+    # via `Ndarray::write` (qd.ndarray case) or via the `SNodeRwAccessorsBank` writer kernel (field case) must
+    # bump the matching generation counter so the next launch evicts the entry and re-runs the sizer.
     #
-    # Internal details: the cache key on every backend is `(AdStackSizingInfo *, snode_write_gen, ndarray_data_gen)`.
-    # Without the host-mutation bump, the second launch sees the same key, returns the stale `max_size` for the
-    # previous ndarray contents, and the reverse pass walks the wrong number of inner iters per outer iteration.
-    # On Metal the symptom is heap reads from out-of-bounds slots that produce garbage gradients (e.g. `x.grad[k]`
-    # in the dozens instead of the analytical `0.8`); on CPU the host-eval path replays observed reads against the
-    # mutated ndarray and recovers without the explicit gen bump, so the test asserts gradient values rather than
-    # an overflow trap to catch the bug on every backend uniformly.
+    # Internal details: the cache key on every backend is `(AdStackSizingInfo *, snode_write_gen[snode_ids],
+    # ndarray_data_gen[devalloc])`. The qd.ndarray path goes through `ndarray_data_gen` keyed by the
+    # `DeviceAllocation` holder address; the field path goes through `snode_write_gen` keyed by `SNode::id`.
+    # Without the bump on either side the second launch sees the same key, returns the stale `max_size` for the
+    # previous contents, and the reverse pass walks the wrong number of inner iters per outer iteration. On Metal
+    # the symptom is heap reads from out-of-bounds slots that produce garbage gradients (e.g. `x.grad[k]` in the
+    # dozens instead of the analytical `0.8`); on CPU the host-eval path replays observed reads and recovers
+    # without the explicit gen bump, so the test asserts gradient values rather than an overflow trap to catch the
+    # bug on every backend uniformly.
     N = 4
     N_X = 16
 
     x = qd.field(qd.f32, shape=(N_X,), needs_grad=True)
     loss = qd.field(qd.f32, shape=(), needs_grad=True)
 
-    @qd.kernel
-    def compute(n: qd.types.ndarray(dtype=qd.i32, ndim=1)):
-        for i_e in range(n.shape[0]):
-            accum = 0.0
-            for j in range(n[i_e]):
-                accum = accum + x[j] * x[j]
-            loss[None] += accum
+    if trip_count_source == "qd_ndarray":
+        n_obj = qd.ndarray(qd.i32, shape=(N,))
 
-    n_arr = qd.ndarray(qd.i32, shape=(N,))
-    for i in range(N):
-        n_arr[i] = 8
+        @qd.kernel
+        def compute(n: qd.types.ndarray(dtype=qd.i32, ndim=1)):
+            for i_e in range(n.shape[0]):
+                accum = 0.0
+                for j in range(n[i_e]):
+                    accum = accum + x[j] * x[j]
+                loss[None] += accum
+
+        def set_n(val):
+            for i in range(N):
+                n_obj[i] = val
+
+        def call_compute():
+            compute(n_obj)
 
+        def call_compute_grad():
+            compute.grad(n_obj)
+
+    else:
+        n_obj = qd.field(qd.i32, shape=(N,))
+
+        @qd.kernel
+        def compute():
+            for i_e in range(N):
+                accum = 0.0
+                for j in range(n_obj[i_e]):
+                    accum = accum + x[j] * x[j]
+                loss[None] += accum
+
+        def set_n(val):
+            for i in range(N):
+                n_obj[i] = val
+
+        def call_compute():
+            compute()
+
+        def call_compute_grad():
+            compute.grad()
+
+    set_n(8)
     for i in range(N_X):
         x[i] = 0.1
     loss[None] = 0.0
-    compute(n_arr)
+    call_compute()
     loss.grad[None] = 1.0
     for i in range(N_X):
         x.grad[i] = 0.0
-    compute.grad(n_arr)
+    call_compute_grad()
     qd.sync()
     assert loss[None] == pytest.approx(N * 8 * 0.01, rel=1e-5)
     for k in range(8):
         assert x.grad[k] == pytest.approx(N * 2 * 0.1, rel=1e-5)
     for k in range(8, N_X):
         assert x.grad[k] == pytest.approx(0.0, abs=1e-7)
 
-    for i in range(N):
-        n_arr[i] = 16
-
+    set_n(16)
     for i in range(N_X):
         x[i] = 0.1
     loss[None] = 0.0
-    compute(n_arr)
+    call_compute()
     loss.grad[None] = 1.0
     for i in range(N_X):
         x.grad[i] = 0.0
-    compute.grad(n_arr)
+    call_compute_grad()
     qd.sync()
     assert loss[None] == pytest.approx(N * 16 * 0.01, rel=1e-5)
     for k in range(N_X):