[Perf] Cache adstack-sizer metadata per task across SPIR-V + LLVM-GPU; per-snode / DeviceAllocation invalidation (Part 2/2) (#620)

Author: duburcqa

docs/source/user_guide/autodiff.md

@@ -289,6 +289,8 @@ The evaluation happens in different places depending on the backend, but the res
 
 Either way, the per-thread stride and each adstack's offset / max-size land in a small buffer the main kernel reads on every push. The backing heap grows on demand to match the largest size any launch has needed so far, and is reused across subsequent launches - you do not need to reserve memory up front.
 
+The sized result is cached per task and reused while the loop bounds are unchanged. Although changing loop bounds at runtime is possible, it comes with some limitations for performance reasons. See [What can go wrong](#what-can-go-wrong) for details.
+
 The on-device sizer relies on two common hardware features (64-bit integer arithmetic and raw-pointer storage-buffer access). Every mainstream GPU from late 2018 onward supports both.
 
 #### Manual override
@@ -347,13 +349,11 @@ A large `ndrange` combined with several loop-carried variables multiplies quickl
 
 ## What can go wrong
 
-- **Adstack overflow at `qd.sync()` (sizer under-estimated the bound).** Surfaces asynchronously at the next `qd.sync()` as `QuadrantsAssertionError: Adstack overflow ...`. On unusually intricate nested loops - typically deeply nested `for i in range(arr[...])` with cumulative-index arithmetic - the sizer can compute a bound that is mathematically tighter than the actual push count. This is a bug; please file it with `QD_DUMP_IR=1` set so the kernel IR ships with the report. Workarounds, in order of convenience:
-  - shorten the innermost dynamic loop;
-  - precompute its worst-case trip into a scalar field the kernel only reads;
-  - split the inner section into its own `@qd.kernel`;
-  - pass `ad_stack_size=N` to `qd.init()` with `N` large enough to cover the real push count (bypasses the sizer).
+- **Adstack overflow.** Surfaces as `QuadrantsAssertionError: Adstack overflow ...` at the next Quadrants Python entry. The message names the offending kernel + offload task and the most likely cause:
+  - *Untracked tensor mutation between launches.* A tensor backing a data-dependent loop bound was written to outside Quadrants's tracking - typically a DLPack zero-copy mutation through a torch tensor sharing storage with a Quadrants ndarray, or a raw pointer write through a non-torch consumer. The cached adstack capacity was sized against the value before the mutation; if the mutation grew the bound, the next launch overflows. Fix: route the write through a Quadrants API (`Ndarray.write` / `Ndarray.fill` / a kernel that writes the value). Alternatively, catch the exception and re-launch - Quadrants invalidates the cached bound on raise, so the retry runs against the live state. Kernel state may be inconsistent after an overflow; do not retry the same step without restarting from a clean state.
+  - *Sizer under-estimated the bound (Quadrants bug).* On unusually intricate nested loops - typically deeply nested `for i in range(arr[...])` with cumulative-index arithmetic - the sizer can compute a bound that is mathematically tighter than the actual push count. To file a bug: clear `/tmp/ir/`, rerun your script with `QD_DUMP_IR=1` set in the environment so Quadrants dumps the kernel IR there, then open an issue on the Quadrants repo with the contents of `/tmp/ir/` attached as a zip. Workaround: pass a generous `ad_stack_size=N` to `qd.init()` with `N` large enough to cover the real push count (bypasses the sizer).
 - **Out-of-memory before the kernel even runs.** A reverse pass through many loop-carried variables at a large ndrange can ask the runtime for more adstack memory than the device can physically back, even when the sizer's number is correct. Surfaces as an allocator OOM at launch time. Remedies are the ones listed under *Avoiding OOM on GPU* above: fewer loop-carried variables, a smaller ndrange, manual checkpointing, or more device-memory headroom.
-- **Loop bounds backed by a mutated ndarray.** A reverse-mode kernel with `for i in range(n[j])` requires `n[j]` to hold the same value at the forward call and at `.grad()`. If anything writes to `n[j]` between those two points - the differentiable kernel itself, or any other kernel call - the computed gradient may come out wrong, sometimes as an `Adstack overflow` exception at `qd.sync()`, sometimes silently. The safe rule: populate loop-bound ndarrays before the forward call and leave them untouched until `.grad()` returns. The reason for that is Quadrants' adstack sizer design: it reads the loop bound separately at each dispatch, which includes forward and backward calls. Tape-based eager AD like [PyTorch's autograd](https://pytorch.org/docs/stable/notes/autograd.html) is not affected, since the trip count is recorded as the forward runs and reused at backward time.
+- **Loop bounds backed by a mutated ndarray.** A reverse-mode kernel with `for i in range(n[j])` requires `n[j]` to hold the same value at the forward call and at `.grad()`. If anything writes to `n[j]` between those two points - the differentiable kernel itself, or any other kernel call - the backward call will trigger an `Adstack overflow` exception or the computed gradient would come out silently wrong. The safe rule: populate loop-bound ndarrays before the forward call and leave them untouched until `.grad()` returns. The reason for that is Quadrants' adstack sizer design: it reads the loop bound separately at each dispatch, which includes forward and backward calls. Tape-based eager AD like [PyTorch's autograd](https://pytorch.org/docs/stable/notes/autograd.html) is not affected, since the trip count is recorded as the forward runs and reused at backward time.
 
 ## Performance characteristics
 

docs/source/user_guide/debug.md

@@ -38,7 +38,7 @@ a = x[-1]     # AssertionError in debug mode
 
 #### Adstack overflow
 
-`debug=True` also enables a deferred bounds check on the adstack used by reverse-mode autodiff. A push past the per-stack capacity (set via `qd.init(ad_stack_size=...)` or per-alloca by `determine_ad_stack_size`) raises `RuntimeError("[Aa]dstack overflow")` on the next `qd.sync()`. Without the check, an adstack overflow silently writes past the per-thread slab and produces a wrong gradient.
+The adstack overflow check on reverse-mode autodiff runs always, on every backend, regardless of `debug`. A push past the per-stack capacity raises `QuadrantsAssertionError("[Aa]dstack overflow")` at the next Quadrants Python entry that polls the overflow flag after the offending kernel has executed - kernel launch, host-side field / ndarray read, or `qd.sync()`. On CPU this is the entry of the offending launch itself; on GPU it can be one or more entries later, since the GPU may not have run the offending push by the time the poll at end-of-launch fires. The error message describes the cause (untracked tensor mutation between launches, or sizer under-estimate caused by a bug in Quadrants) and the recovery flow; see [Autodiff -> What can go wrong](autodiff.md) for the full description.
 
 ### Assertions in kernels
 

docs/source/user_guide/init_options.md

@@ -86,11 +86,12 @@ Default `False`. Turns on every available correctness check. Use while iterating
 
 Enables:
 - field-bounds check on tensor indexing (out-of-range index raises `RuntimeError`);
-- adstack-overflow check on reverse-mode autodiff (overflow raises `RuntimeError` on the next `qd.sync()`);
 - kernel `assert` statements;
 - integer-overflow guards on arithmetic;
 - IR verification after every compiler pass.
 
+The adstack-overflow check on reverse-mode autodiff runs unconditionally on every backend regardless of `debug`; see [Autodiff -> What can go wrong](autodiff.md) for the contract.
+
 **Cost.** Significant on both compile time (verifier walks the IR after every transform; extra runtime checks expand the emitted code; ~21s extra observed on adstack-heavy kernels) and runtime. For just the field-bounds check in a release build without the rest, use [`check_out_of_bound`](#check_out_of_bound) below.
 
 ### `check_out_of_bound`
@@ -101,23 +102,22 @@ Default `False`. Enables the field-bounds check on tensor indexing - an out-of-r
 
 Interaction with `debug`:
 
-| Flags | Field bounds | Adstack overflow | Other `debug` checks |
-|-------|--------------|------------------|----------------------|
-| neither | off | off | off |
-| `check_out_of_bound=True` only | on | off | off |
-| `debug=True` | on | on | on |
+| Flags | Field bounds | Other `debug` checks |
+|-------|--------------|----------------------|
+| neither | off | off |
+| `check_out_of_bound=True` only | on | off |
+| `debug=True` | on | on |
 
 - `debug=True` always implies `check_out_of_bound=True` (the field-bounds check fires whenever debug mode is on).
-- The adstack-overflow check on reverse-mode autodiff (a push past the per-stack capacity raises `RuntimeError("[Aa]dstack overflow")` on the next `qd.sync()`) is on its own gate, controlled by `debug` - it is not enabled by `check_out_of_bound` alone.
 
 Per-backend support:
 
-| Backend | Field bounds check | Adstack overflow check |
-|---------|--------------------|------------------------|
-| CPU | with `check_out_of_bound=True` or `debug=True` | with `debug=True` |
-| CUDA | with `check_out_of_bound=True` or `debug=True` | with `debug=True` |
-| AMDGPU | with `check_out_of_bound=True` or `debug=True` | with `debug=True` |
-| Metal | never (no in-kernel assertion mechanism) | with `debug=True` |
-| Vulkan | never (no in-kernel assertion mechanism) | with `debug=True` |
+| Backend | Field bounds check |
+|---------|--------------------|
+| CPU | with `check_out_of_bound=True` or `debug=True` |
+| CUDA | with `check_out_of_bound=True` or `debug=True` |
+| AMDGPU | with `check_out_of_bound=True` or `debug=True` |
+| Metal | never (no in-kernel assertion mechanism) |
+| Vulkan | never (no in-kernel assertion mechanism) |
 
-Metal and Vulkan lack the assertion extension that the field-bounds check relies on; `check_out_of_bound=True` is silently reset to `False` on those backends at `qd.init` time and a warning is logged. The adstack-overflow check is gated independently of the assertion extension, so `debug=True` activates it on every backend including Metal and Vulkan.
+Metal and Vulkan lack the assertion extension that the field-bounds check relies on; `check_out_of_bound=True` is silently reset to `False` on those backends at `qd.init` time and a warning is logged.

python/quadrants/lang/impl.py

@@ -24,6 +24,7 @@
     QuadrantsRuntimeError,
     QuadrantsSyntaxError,
     QuadrantsTypeError,
+    handle_exception_from_cpp,
 )
 from quadrants.lang.expr import Expr, make_expr_group
 from quadrants.lang.field import Field, ScalarField
@@ -576,7 +577,13 @@ def sync(self):
             return
         self.materialize()
         assert self._prog is not None
-        self._prog.synchronize()
+        try:
+            self._prog.synchronize()
+        except Exception as e:
+            wrapped = handle_exception_from_cpp(e)
+            if wrapped is e:
+                raise
+            raise wrapped from None
 
 
 pyquadrants = PyQuadrants()