[Fixup] Bump ndarray_data_gen on host-side ndarray writes + fill so cached adstack-sizer metadata is evicted

Author: duburcqa

quadrants/program/ndarray.cpp

@@ -1,5 +1,6 @@
 #include <numeric>
 
+#include "quadrants/program/adstack_size_expr_eval.h"
 #include "quadrants/program/ndarray.h"
 #include "quadrants/program/program.h"
 #include "fp16.h"
@@ -185,6 +186,10 @@ void Ndarray::write(const std::vector<int> &I, TypedConstant val) const {
   staging_buf_->device->memcpy_internal(this->ndarray_alloc_.get_ptr(index * size_), staging_buf_->get_ptr(), size_);
 
   prog_->synchronize_and_assert();
+  // Host-side mutation of the ndarray contents: bump the per-DeviceAllocation generation so any cached
+  // adstack-sizer metadata that depended on `ExternalTensorRead` of this ndarray is evicted on next launch.
+  // Keyed by the same `&ndarray_alloc_` the kernel launchers use in `bump_writes_for_kernel_*`.
+  prog_->adstack_cache().bump_ndarray_data_gen(const_cast<DeviceAllocation *>(&ndarray_alloc_));
 }
 
 int64 Ndarray::read_int(const std::vector<int> &i) {

quadrants/program/program.cpp

@@ -442,6 +442,10 @@ void Program::fill_ndarray_fast_u32(Ndarray *ndarray, uint32_t val) {
   // This is a temporary solution to bypass device api. Should be moved to CommandList once available in CUDA.
   program_impl_->fill_ndarray(ndarray->ndarray_alloc_,
                               ndarray->get_nelement() * ndarray->get_element_size() / sizeof(uint32_t), val);
+  // Host-issued device fill mutates the ndarray contents without going through a quadrants kernel, so
+  // `bump_writes_for_kernel_*` will not cover it. Bump explicitly using the same `&ndarray_alloc_` key the
+  // kernel launchers use, so any cached adstack-sizer metadata depending on this ndarray is evicted.
+  adstack_cache_->bump_ndarray_data_gen(&ndarray->ndarray_alloc_);
 }
 
 std::pair<const StructType *, size_t> Program::get_struct_type_with_data_layout(const StructType *old_ty,

tests/python/test_adstack.py

@@ -2089,6 +2089,71 @@ def compute(a: qd.types.ndarray(dtype=qd.i32, ndim=1)):
         assert x.grad[i] == pytest.approx(expected, rel=1e-5)
 
 
+@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.
+    #
+    # 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.
+    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
+
+    n_arr = qd.ndarray(qd.i32, shape=(N,))
+    for i in range(N):
+        n_arr[i] = 8
+
+    for i in range(N_X):
+        x[i] = 0.1
+    loss[None] = 0.0
+    compute(n_arr)
+    loss.grad[None] = 1.0
+    for i in range(N_X):
+        x.grad[i] = 0.0
+    compute.grad(n_arr)
+    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
+
+    for i in range(N_X):
+        x[i] = 0.1
+    loss[None] = 0.0
+    compute(n_arr)
+    loss.grad[None] = 1.0
+    for i in range(N_X):
+        x.grad[i] = 0.0
+    compute.grad(n_arr)
+    qd.sync()
+    assert loss[None] == pytest.approx(N * 16 * 0.01, rel=1e-5)
+    for k in range(N_X):
+        assert x.grad[k] == pytest.approx(N * 2 * 0.1, rel=1e-5)
+
+
 @test_utils.test(require=qd.extension.adstack)
 def test_adstack_inner_range_bounded_by_multidim_ndarray_read():
     # Pins multi-axis stride handling in the `ExternalTensorRead` evaluator. The sizer routes a reverse-mode inner trip