diff --git a/docs/source/user_guide/autodiff.md b/docs/source/user_guide/autodiff.md index cf120cca8b..e9885bfb6d 100644 --- a/docs/source/user_guide/autodiff.md +++ b/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 diff --git a/docs/source/user_guide/debug.md b/docs/source/user_guide/debug.md index 044a15c13f..94dfa94407 100644 --- a/docs/source/user_guide/debug.md +++ b/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 diff --git a/docs/source/user_guide/init_options.md b/docs/source/user_guide/init_options.md index d7f9e33310..8ef71245fc 100644 --- a/docs/source/user_guide/init_options.md +++ b/docs/source/user_guide/init_options.md @@ -74,11 +74,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` @@ -89,23 +90,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. diff --git a/python/quadrants/lang/impl.py b/python/quadrants/lang/impl.py index ab41296afc..bea9b65867 100644 --- a/python/quadrants/lang/impl.py +++ b/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() diff --git a/quadrants/codegen/llvm/codegen_llvm.cpp b/quadrants/codegen/llvm/codegen_llvm.cpp index c36e4c5d67..6bf252a2d1 100644 --- a/quadrants/codegen/llvm/codegen_llvm.cpp +++ b/quadrants/codegen/llvm/codegen_llvm.cpp @@ -7,8 +7,11 @@ #include "llvm/IR/Module.h" #include "llvm/Linker/Linker.h" #include "quadrants/analysis/offline_cache_util.h" +#include "quadrants/ir/analysis.h" +#include "quadrants/ir/snode.h" #include "quadrants/ir/statements.h" #include "quadrants/ir/transforms.h" +#include "quadrants/program/adstack_size_expr_eval.h" #include "quadrants/program/extension.h" #include "quadrants/runtime/program_impls/llvm/llvm_program.h" #include "quadrants/codegen/llvm/struct_llvm.h" @@ -1938,6 +1941,21 @@ std::string TaskCodeGenLLVM::init_offloaded_task_function(OffloadedStmt *stmt, s func = llvm::Function::Create(task_function_type, llvm::Function::ExternalLinkage, task_kernel_name, module.get()); current_task = std::make_unique(task_kernel_name); + // Pre-register the per-task AdStackSizingInfo so the registry id is assigned BEFORE codegen visits any + // `AdStackPushStmt`. Metadata (allocated_max_sizes) is filled in at `finalize_offloaded_task_function` + // time after the alloca scan completes; the registry call is idempotent on the same `ad_stack_ptr` so + // the second call updates the entry in place. Skipping registration when `prog == nullptr` (C++-only + // tests) leaves `registry_id == 0`, which the codegen-emitted cmpxchg short-circuits. + if (prog != nullptr) { + // Reserve a registry id at task START so codegen can bake the immediate before any push site is + // visited. Metadata + size_exprs are filled in at `finalize_offloaded_task_function` time below + // (idempotent re-registration on the same identity_key updates the entry in place). Identity + // key here is the raw `¤t_task->ad_stack` address; the registry never derefs it. + uint32_t id = prog->adstack_cache().register_adstack_sizing_info( + static_cast(¤t_task->ad_stack), kernel_name, task_codegen_id, /*allocated_max_sizes=*/{}, + /*size_exprs=*/{}); + current_task->ad_stack.registry_id = id; + } for (auto &arg : func->args()) { kernel_args.push_back(&arg); @@ -1975,6 +1993,65 @@ void TaskCodeGenLLVM::finalize_offloaded_task_function() { current_task->ad_stack.allocas = ad_stack_allocas_info_; current_task->ad_stack.size_exprs = ad_stack_size_exprs_; current_task->ad_stack.bound_expr = ad_stack_static_bound_expr_; + // Snodes the task body mutates. Persisted on `OffloadedTask::snode_writes` so the LLVM + // launcher can invalidate the per-task adstack metadata cache when a kernel that runs in + // between mutated a SNode an enclosing `size_expr::FieldLoad` reads. Mirrors the SPIR-V + // analogue in `spirv_codegen.cpp`. Sorted + deduplicated for stable serialisation. + if (current_offload != nullptr) { + auto snode_rw = irpass::analysis::gather_snode_read_writes(current_offload); + current_task->snode_writes.reserve(snode_rw.second.size()); + for (auto *s : snode_rw.second) { + if (s != nullptr) { + current_task->snode_writes.push_back(s->id); + } + } + std::sort(current_task->snode_writes.begin(), current_task->snode_writes.end()); + current_task->snode_writes.erase( + std::unique(current_task->snode_writes.begin(), current_task->snode_writes.end()), + current_task->snode_writes.end()); + // Ndarray args this task writes to. Same role as `snode_writes` but for ndarray data; + // covers `size_expr::ExternalTensorRead` invalidation. The first element of each + // `arg_id_path` key is the kernel-arg slot, which is what `Program::ndarray_data_gen_` + // is keyed by (via the bound DeviceAllocation). + auto arr_access = irpass::detect_external_ptr_access_in_task(current_offload); + for (const auto &kv : arr_access) { + if (kv.first.empty()) { + continue; + } + const uint32_t access_bits = static_cast(kv.second); + if ((access_bits & static_cast(irpass::ExternalPtrAccess::WRITE)) != 0) { + current_task->arr_writes.push_back(kv.first.front()); + } + if ((access_bits & static_cast(irpass::ExternalPtrAccess::READ)) != 0) { + current_task->arr_reads.push_back(kv.first.front()); + } + } + std::sort(current_task->arr_writes.begin(), current_task->arr_writes.end()); + current_task->arr_writes.erase(std::unique(current_task->arr_writes.begin(), current_task->arr_writes.end()), + current_task->arr_writes.end()); + std::sort(current_task->arr_reads.begin(), current_task->arr_reads.end()); + current_task->arr_reads.erase(std::unique(current_task->arr_reads.begin(), current_task->arr_reads.end()), + current_task->arr_reads.end()); + } + // Register the per-task AdStackSizingInfo with the Program-side identity registry. The id is baked + // into the lazy-claim overflow path's `cmpxchg(0, id)` so the host raise site can name the offending + // kernel + task in its diagnostic message. Empty alloca list = no adstack pushes in this task; skip + // registration to keep the registry compact. + if (!current_task->ad_stack.allocas.empty() && prog != nullptr) { + std::vector allocated_max_sizes; + allocated_max_sizes.reserve(current_task->ad_stack.allocas.size()); + for (const auto &a : current_task->ad_stack.allocas) { + allocated_max_sizes.push_back(static_cast(a.max_size_compile_time)); + } + // Update the entry with the live metadata + per-alloca size_exprs. The size_exprs are copied + // into the registry so the diagnose path can walk them without dereferencing the launcher's + // unstable `OffloadedTask::ad_stack` pointer (freed by `current_task = nullptr` after + // by-value `offloaded_tasks.push_back(*current_task)`). + uint32_t id = prog->adstack_cache().register_adstack_sizing_info( + static_cast(¤t_task->ad_stack), kernel_name, task_codegen_id, + std::move(allocated_max_sizes), current_task->ad_stack.size_exprs); + current_task->ad_stack.registry_id = id; + } } // entry_block should jump to the body after all allocas are inserted @@ -2408,10 +2485,12 @@ void TaskCodeGenLLVM::emit_ad_stack_row_claim_llvm() { // Per-task capacity slot for the defense-in-depth bounds check: clamp the claimed row at `capacity - 1` so any // overshoot stays in-bounds. For tasks without a captured `bound_expr` the launcher writes UINT32_MAX into this slot - // so the clamp is inert. The divergence-overflow signal that the SPIR-V codegen emits via OpAtomicUMax is not yet - // wired on the LLVM side - it requires a `__atomic_or_n` against `runtime->adstack_overflow_flag` and a matching - // runtime-side getter; in its absence we still get the in-bounds clamp, so the kernel cannot silently corrupt the - // heap end. Surfacing the divergence is future work. + // so the clamp is inert. On overshoot (`claimed_row > capacity - 1`) the codegen also OR-1's the host-visible + // adstack overflow flag (`runtime->adstack_overflow_flag_dev_ptr`, which the host allocated as pinned UVA-mapped + // memory in `LlvmRuntimeExecutor::materialize_runtime`) so the host poll surfaces the divergence at the next + // Quadrants Python entry. The atomic crosses the host/device boundary cleanly because the slot is in + // pinned host memory; required hardware envelope is the same Pascal+ / GFX9+ that the existing pinned-host + // H2D-async pattern already requires. llvm::Value *capacities_base = call("LLVMRuntime_get_adstack_bound_row_capacities", get_runtime()); llvm::Value *capacity_slot_ptr = builder->CreateGEP(i32ty, capacities_base, task_id_i64); llvm::Value *capacity = builder->CreateLoad(i32ty, capacity_slot_ptr); @@ -2426,6 +2505,38 @@ void TaskCodeGenLLVM::emit_ad_stack_row_claim_llvm() { llvm::Value *cmp = builder->CreateICmpUGT(claimed_row, clamp_upper); llvm::Value *clamped_row = builder->CreateSelect(cmp, clamp_upper, claimed_row); builder->CreateStore(clamped_row, row_id_var); + + // Overflow signal: on `claimed_row > clamp_upper`, atomically OR 1 into the pinned-host overflow flag and + // record the offending task identity in the companion `adstack_overflow_task_id_dev_ptr` slot via a + // `cmpxchg(0, registry_id)`. Only the FIRST overflowing thread's id sticks; subsequent threads observe + // a non-zero value and their cmpxchg fails harmlessly. The condition is hoisted to a structured if so + // the not-overflowing fast path skips both atomics entirely - one function call to fetch the pointers + // plus one CreateICmpUGT comparison (the same compare we already emitted for the clamp). + auto *current_function = builder->GetInsertBlock()->getParent(); + auto *overflow_then_block = llvm::BasicBlock::Create(*llvm_context, "adstack_overflow_signal", current_function); + auto *overflow_merge_block = llvm::BasicBlock::Create(*llvm_context, "adstack_overflow_merge", current_function); + builder->CreateCondBr(cmp, overflow_then_block, overflow_merge_block); + builder->SetInsertPoint(overflow_then_block); + { + auto *i64ty_local = llvm::Type::getInt64Ty(*llvm_context); + llvm::Value *flag_ptr = call("LLVMRuntime_get_adstack_overflow_flag_dev_ptr", get_runtime()); + llvm::Value *one_i64 = llvm::ConstantInt::get(i64ty_local, 1); + builder->CreateAtomicRMW(llvm::AtomicRMWInst::Or, flag_ptr, one_i64, llvm::MaybeAlign(), + llvm::AtomicOrdering::Monotonic); + // Record the registry id (0 means "not registered"; skip the cmpxchg in that case so the slot stays + // zero and the host raise site falls through to the generic dual-cause message). Each offload task + // emits its own lazy-claim block, so the immediate is task-local at codegen time. + if (current_task != nullptr && current_task->ad_stack.registry_id != 0) { + llvm::Value *task_id_ptr = call("LLVMRuntime_get_adstack_overflow_task_id_dev_ptr", get_runtime()); + llvm::Value *expected_zero = llvm::ConstantInt::get(i64ty_local, 0); + llvm::Value *new_id = + llvm::ConstantInt::get(i64ty_local, static_cast(current_task->ad_stack.registry_id)); + builder->CreateAtomicCmpXchg(task_id_ptr, expected_zero, new_id, llvm::MaybeAlign(), + llvm::AtomicOrdering::Monotonic, llvm::AtomicOrdering::Monotonic); + } + builder->CreateBr(overflow_merge_block); + } + builder->SetInsertPoint(overflow_merge_block); } // Return (creating on first call) the per-stack `alloca i64` that holds the live push count for this stack on the @@ -2686,7 +2797,11 @@ void TaskCodeGenLLVM::visit(AdStackPushStmt *stmt) { llvm::Value *max_size_addr = builder->CreateGEP(i64ty, ad_stack_max_sizes_ptr_llvm_, stack_id_i64); llvm::Value *max_size = builder->CreateLoad(i64ty, max_size_addr); llvm::Value *stack_base = get_ad_stack_base_llvm(stack); - call("stack_push", get_runtime(), stack_base, max_size, tlctx->get_constant(stack->element_size_in_bytes())); + auto *i64ty_local = llvm::Type::getInt64Ty(*llvm_context); + llvm::Value *registry_id_const = llvm::ConstantInt::get( + i64ty_local, current_task != nullptr ? static_cast(current_task->ad_stack.registry_id) : 0u); + call("stack_push", get_runtime(), stack_base, max_size, tlctx->get_constant(stack->element_size_in_bytes()), + registry_id_const); auto primal_ptr = call("stack_top_primal", stack_base, tlctx->get_constant(stack->element_size_in_bytes())); primal_ptr = builder->CreateBitCast(primal_ptr, llvm::PointerType::get(tlctx->get_data_type(stmt->ret_type), 0)); builder->CreateStore(llvm_val[stmt->v], primal_ptr); diff --git a/quadrants/codegen/llvm/llvm_compiled_data.h b/quadrants/codegen/llvm/llvm_compiled_data.h index 1606dd31f8..820b8c3476 100644 --- a/quadrants/codegen/llvm/llvm_compiled_data.h +++ b/quadrants/codegen/llvm/llvm_compiled_data.h @@ -73,6 +73,14 @@ struct AdStackSizingInfo { // the actual gate-passing thread count; `nullopt` falls through to dispatched-threads worst-case sizing (no behavior // change versus a kernel without this metadata). std::optional bound_expr; + // Identity in `Program::adstack_sizing_info_registry_`. Assigned at `finalize_offloaded_task_function` + // time after the registry idempotently maps `&this` to a u32 id. Baked as an immediate into the + // codegen-emitted lazy-claim `cmpxchg(0, registry_id)` so the host raise site can name the offending + // kernel + task in its diagnostic message. `0` means "not registered" - the lazy-claim emit short- + // circuits the cmpxchg in that case (no information to record). NOT serialised to the offline cache: + // ids are assigned per `Program` lifetime, not per-kernel-content; a deserialised task re-registers + // itself at the next launch. + uint32_t registry_id{0}; QD_IO_DEF(per_thread_stride, per_thread_stride_float, per_thread_stride_int, @@ -95,12 +103,32 @@ class OffloadedTask { int dynamic_shared_array_bytes{0}; AdStackSizingInfo ad_stack{}; + // Snode IDs this task writes to (read-modify-write counts as a write). Computed at codegen time + // by walking the offloaded IR with `gather_snode_read_writes`. Consumed at launch time: each id + // here bumps `Program::snode_write_gen_[id]` so the per-task adstack metadata cache invalidates + // whenever a kernel that ran since the cache was recorded mutated a SNode a downstream + // `size_expr::FieldLoad` may read. Mirrors the SPIR-V `TaskAttributes::snode_writes` field. + std::vector snode_writes; + // Argument arg_ids this task writes to (WRITE bit set in `irpass::detect_external_ptr_access_in_task`). + // Consumed at launch time to bump `Program::ndarray_data_gen_` for the bound DeviceAllocation so + // the per-task adstack metadata cache invalidates when a kernel that ran since the cache was + // recorded mutated an ndarray a downstream `size_expr::ExternalTensorRead` reads. Mirrors the + // SPIR-V `KernelContextAttributes::arr_access` WRITE-bit set, but stored per-task here because + // LLVM codegen does not aggregate `arr_access` to the kernel level. + std::vector arr_writes; + // Argument arg_ids this task reads (READ bit set in `irpass::detect_external_ptr_access_in_task`). Consumed at + // launch time only on backends with an H2D blit per launch (LLVM-GPU CUDA / AMDGPU `DevAllocType::kNone`) and on + // CPU LLVM with `DevAllocType::kNone` host args: the data pointer is stable across launches, so the cache key by + // data ptr cannot detect content mutations the user performed outside Quadrants's tracking. Mirrors the SPIR-V + // `KernelContextAttributes::arr_access` READ-bit set. + std::vector arr_reads; + explicit OffloadedTask(const std::string &name = "", int block_dim = 0, int grid_dim = 0, int dynamic_shared_array_bytes = 0) : name(name), block_dim(block_dim), grid_dim(grid_dim), dynamic_shared_array_bytes(dynamic_shared_array_bytes) {}; - QD_IO_DEF(name, block_dim, grid_dim, dynamic_shared_array_bytes, ad_stack); + QD_IO_DEF(name, block_dim, grid_dim, dynamic_shared_array_bytes, ad_stack, snode_writes, arr_writes, arr_reads); }; struct LLVMCompiledTask { diff --git a/quadrants/codegen/spirv/detail/spirv_codegen.h b/quadrants/codegen/spirv/detail/spirv_codegen.h index 3dfc9b508d..b0beb165b2 100644 --- a/quadrants/codegen/spirv/detail/spirv_codegen.h +++ b/quadrants/codegen/spirv/detail/spirv_codegen.h @@ -307,6 +307,24 @@ class TaskCodegen : public IRVisitor { // via `load_top`, so the bootstrap value is dead memory and writing it through a possibly-unclaimed `row_id_var` // would corrupt arbitrary heap rows). Only the `count_var` increment is kept so push and pop stay balanced. std::unordered_set ad_stack_bootstrap_pushes_; + + // Function-scope u32 alloca holding the maximum-overflow-signal-seen-this-thread across every adstack push site + // in the task. Each `AdStackPushStmt` visitor updates it with `signal = (count >= max) ? stack_id + 1 : 0` via + // OpUMax (register-only, no atomic, no global memory access on the not-overflowing fast path). Right before the + // kernel's `OpReturn`, an `if any_overflow_signal_var > 0 then atomic_fetch_max(host_visible_buf[0], + // any_overflow_signal_var)` emits the host-visible signal exactly once per task per kernel. The host polls the + // buffer (Apple Silicon: shared memory; Vulkan: HOST_VISIBLE | HOST_COHERENT) without any DtoH or sync drain. + // Lazy-allocated by the first `AdStackPushStmt` visitor; nullptr when the task has no adstack pushes (no + // task-end emit needed). Reset to {} per task in `generate_*_kernel`. + spirv::Value any_overflow_signal_var_; + // Set to true when the task body contains at least one `AdStackPushStmt` visit. Read by + // `emit_adstack_task_end_overflow_check` to decide whether to emit the task-end overflow check and + // its AdStackOverflow / AdStackTaskRegistryId buffer accesses. Forward-only tasks never set this and + // therefore never request the AdStack-side buffer binds, which keeps the SPIR-V launcher's bind + // path from null-binding `AdStackTaskRegistryId` on a Program that has not allocated it + // (allocation only fires inside `publish_adstack_metadata_spirv` for kernels with at least one + // adstack-bearing task; null-binding would crash Metal's `rw_buffer` device-equality assertion). + bool task_has_adstack_push_{false}; // Function-scope OpVariable initialized to UINT32_MAX at task entry; overwritten with the atomically claimed row // index when codegen visits `ad_stack_lca_block_float_`. `get_ad_stack_heap_thread_base_float()` loads this variable // and multiplies against the runtime float stride to produce the per-thread heap base, replacing the prior `invoc_id @@ -340,6 +358,12 @@ class TaskCodegen : public IRVisitor { spirv::Value get_ad_stack_metadata_stride_float(); spirv::Value get_ad_stack_metadata_stride_int(); void ensure_ad_stack_metadata_loaded(AdStackSpirv &info); + // Lazy-allocate the per-task overflow-signal accumulator at the function entry block. See the member + // doc on `any_overflow_signal_var_` for the design rationale. + spirv::Value ensure_any_overflow_signal_var(); + // Emit the task-end overflow check at the current insertion point. Must be called from each + // `generate_*_kernel` just before the closing `OpReturn`. No-op when the task has no adstack push sites. + void emit_adstack_task_end_overflow_check(); // Routes to the correct backing-typed pointer (`*f32` for `heap_float`, `*i32` for `heap_int`) based on // `info.heap_kind`. See comment on the implementation for the bool<->i32 conversion contract. spirv::Value ad_stack_slot_ptr(AdStackSpirv &info, spirv::Value idx, bool primal); diff --git a/quadrants/codegen/spirv/kernel_utils.cpp b/quadrants/codegen/spirv/kernel_utils.cpp index 408af81db1..793dcfa021 100644 --- a/quadrants/codegen/spirv/kernel_utils.cpp +++ b/quadrants/codegen/spirv/kernel_utils.cpp @@ -39,6 +39,9 @@ std::string TaskAttributes::buffers_name(BufferInfo b) { if (b.type == BufferType::AdStackBoundRowCapacity) { return "AdStackBoundRowCapacity"; } + if (b.type == BufferType::AdStackTaskRegistryId) { + return "AdStackTaskRegistryId"; + } if (b.type == BufferType::AdStackHeapFloat) { return "AdStackHeapFloat"; } diff --git a/quadrants/codegen/spirv/kernel_utils.h b/quadrants/codegen/spirv/kernel_utils.h index bf00e6530b..7c12196b0a 100644 --- a/quadrants/codegen/spirv/kernel_utils.h +++ b/quadrants/codegen/spirv/kernel_utils.h @@ -59,6 +59,16 @@ struct TaskAttributes { // internal bug, not user-recoverable), and `synchronize()` surfaces it as a clear actionable error rather than // letting it silently corrupt gradients via OOB writes. AdStackBoundRowCapacity, + // Per-kernel StorageBuffer holding the `Program::adstack_sizing_info_registry_` id per task + // (`uint[num_tasks_in_kernel]`). Populated by the SPIR-V launcher in + // `GfxRuntime::publish_adstack_metadata_spirv` immediately after registering each adstack-bearing + // task with the Program-side identity registry; slot `task_id_in_kernel` carries the registry id + // for that task (0 for tasks without adstacks). The codegen task-end emit reads slot + // `task_id_in_kernel` and `OpAtomicCompareExchange`'s it into `AdStackOverflow[1]` on overflow, + // recording the FIRST overflowing task's registry id for the host raise site to look up + // kernel name + offload task index in its diagnostic message. Allocated and grown on demand + // following the same pattern as `AdStackBoundRowCapacity`. + AdStackTaskRegistryId, }; struct BufferInfo { @@ -208,10 +218,27 @@ struct TaskAttributes { uint32_t per_thread_stride_int_compile_time{0}; std::vector allocas; std::optional bound_expr; + // Identity in `Program::adstack_sizing_info_registry_`. Assigned at SPIR-V codegen time after the + // Program registry idempotently maps `&this` to a u32 id. Baked as an immediate into the codegen- + // emitted task-end overflow path's `cmpxchg(0, registry_id)` against slot 1 of the AdStackOverflow + // buffer so the host raise site can name the offending kernel + task in its diagnostic message. `0` + // means "not registered" - the codegen short-circuits the cmpxchg in that case. NOT serialised to the + // offline cache: ids are assigned per `Program` lifetime; a deserialised task re-registers itself at + // the next launch. + uint32_t registry_id{0}; QD_IO_DEF(per_thread_stride_float_compile_time, per_thread_stride_int_compile_time, allocas, bound_expr); }; AdStackSizingAttribs ad_stack; + // Snode IDs this task writes to (read-modify-write counts as a write). Computed at SPIR-V codegen time + // by walking the offloaded IR with `gather_snode_read_writes`. Consumed by the SPIR-V launcher on every + // `launch_kernel` call: each id here bumps `Program::snode_write_gen_[id]` so the per-task adstack + // metadata cache invalidates whenever a kernel that ran since the cache was recorded mutated a SNode + // a downstream `size_expr::FieldLoad` may read. Stored as raw IDs (not `SNode *`) so the field + // survives offline-cache load-store; the runtime resolves the pointer on demand via + // `Program::get_snode_by_id` only if it ever needs to call into snode-specific APIs. + std::vector snode_writes; + static std::string buffers_name(BufferInfo b); std::string debug_string() const; @@ -222,7 +249,8 @@ struct TaskAttributes { task_type, buffer_binds, range_for_attribs, - ad_stack); + ad_stack, + snode_writes); }; /** diff --git a/quadrants/codegen/spirv/spirv_codegen.cpp b/quadrants/codegen/spirv/spirv_codegen.cpp index 2d6b601e37..d3be1a9688 100644 --- a/quadrants/codegen/spirv/spirv_codegen.cpp +++ b/quadrants/codegen/spirv/spirv_codegen.cpp @@ -17,7 +17,9 @@ #include "quadrants/codegen/spirv/spirv_ir_builder.h" #include "quadrants/codegen/spirv/detail/spirv_codegen.h" #include "quadrants/codegen/spirv/spirv_shared_array_retyping.h" +#include "quadrants/ir/analysis.h" #include "quadrants/ir/transforms.h" +#include "quadrants/ir/snode.h" #include "quadrants/math/arithmetic.h" #include "quadrants/codegen/ir_dump.h" @@ -36,6 +38,7 @@ constexpr char kExtArrBufferName[] = "ext_arr_buffer"; constexpr char kAdStackOverflowBufferName[] = "adstack_overflow_buffer"; constexpr char kAdStackRowCounterBufferName[] = "adstack_row_counter_buffer"; constexpr char kAdStackBoundRowCapacityBufferName[] = "adstack_bound_row_capacity_buffer"; +constexpr char kAdStackTaskRegistryIdBufferName[] = "adstack_task_registry_id_buffer"; constexpr char kAdStackHeapFloatBufferName[] = "adstack_heap_float_buffer"; constexpr char kAdStackHeapIntBufferName[] = "adstack_heap_int_buffer"; constexpr char kAdStackMetadataBufferName[] = "adstack_metadata_buffer"; @@ -66,6 +69,8 @@ std::string buffer_instance_name(BufferInfo b) { return kAdStackRowCounterBufferName; case BufferType::AdStackBoundRowCapacity: return kAdStackBoundRowCapacityBufferName; + case BufferType::AdStackTaskRegistryId: + return kAdStackTaskRegistryIdBufferName; case BufferType::AdStackHeapFloat: return kAdStackHeapFloatBufferName; case BufferType::AdStackHeapInt: @@ -198,6 +203,23 @@ TaskCodegen::Result TaskCodegen::run() { task_attribs_.ad_stack.per_thread_stride_float_compile_time = ad_stack_heap_per_thread_stride_float_; task_attribs_.ad_stack.per_thread_stride_int_compile_time = ad_stack_heap_per_thread_stride_int_; + // Snodes the task body mutates (any `GlobalStore` or `AtomicOp` whose dest resolves to a + // `GlobalPtrStmt`). Persisted on `task_attribs_.snode_writes` so the SPIR-V launcher can bump + // `Program::snode_write_gen_` for each id on every `launch_kernel` call - that is the precise signal + // the per-task adstack metadata cache uses to invalidate when a prior kernel may have changed a + // value an enclosing `size_expr::FieldLoad` reads. Stored as raw int ids (not `SNode *`) so the + // field round-trips through the offline cache without resolving the pointer at serialise time. + auto snode_rw = irpass::analysis::gather_snode_read_writes(task_ir_); + task_attribs_.snode_writes.reserve(snode_rw.second.size()); + for (auto *s : snode_rw.second) { + if (s != nullptr) { + task_attribs_.snode_writes.push_back(s->id); + } + } + std::sort(task_attribs_.snode_writes.begin(), task_attribs_.snode_writes.end()); + task_attribs_.snode_writes.erase(std::unique(task_attribs_.snode_writes.begin(), task_attribs_.snode_writes.end()), + task_attribs_.snode_writes.end()); + Result res; res.spirv_code = ir_->finalize(); res.task_attribs = std::move(task_attribs_); @@ -1899,6 +1921,9 @@ void TaskCodegen::generate_serial_kernel(OffloadedStmt *stmt) { // The computation for a single work is wrapped inside a function, so that // we can do grid-strided loop. ir_->start_function(kernel_function_); + // Initialise the adstack overflow-signal accumulator before any user code so the zero-init dominates + // every push site and the task-end read. See `ensure_any_overflow_signal_var` doc for details. + ensure_any_overflow_signal_var(); spirv::Value cond = ir_->eq(ir_->get_global_invocation_id(0), ir_->uint_immediate_number(ir_->u32_type(), 0)); // if (gl_GlobalInvocationID.x > 0) spirv::Label then_label = ir_->new_label(); @@ -1914,6 +1939,7 @@ void TaskCodegen::generate_serial_kernel(OffloadedStmt *stmt) { ir_->make_inst(spv::OpBranch, merge_label); ir_->start_label(merge_label); + emit_adstack_task_end_overflow_check(); ir_->make_inst(spv::OpReturn); // return; ir_->make_inst(spv::OpFunctionEnd); // } Close kernel @@ -1949,6 +1975,7 @@ void TaskCodegen::generate_range_for_kernel(OffloadedStmt *stmt) { range_for_attribs.end = (stmt->const_end ? stmt->end_value : stmt->end_offset); ir_->start_function(kernel_function_); + ensure_any_overflow_signal_var(); const std::string total_elems_name("total_elems"); spirv::Value total_elems; spirv::Value begin_expr_value; @@ -2115,6 +2142,7 @@ void TaskCodegen::generate_range_for_kernel(OffloadedStmt *stmt) { // loop merge ir_->start_label(merge_label); + emit_adstack_task_end_overflow_check(); ir_->make_inst(spv::OpReturn); ir_->make_inst(spv::OpFunctionEnd); @@ -2130,6 +2158,7 @@ void TaskCodegen::generate_struct_for_kernel(OffloadedStmt *stmt) { // The computation for a single work is wrapped inside a function, so that // we can do grid-strided loop. ir_->start_function(kernel_function_); + ensure_any_overflow_signal_var(); auto listgen_buffer = get_buffer_value(BufferType::ListGen, PrimitiveType::u32); auto listgen_count_ptr = ir_->struct_array_access(ir_->u32_type(), listgen_buffer, ir_->const_i32_zero_); @@ -2173,6 +2202,7 @@ void TaskCodegen::generate_struct_for_kernel(OffloadedStmt *stmt) { } ir_->start_label(loop_merge); + emit_adstack_task_end_overflow_check(); ir_->make_inst(spv::OpReturn); // return; ir_->make_inst(spv::OpFunctionEnd); // } Close kernel @@ -2631,6 +2661,71 @@ spirv::Value TaskCodegen::ad_stack_heap_int_ptr(spirv::Value slot_offset, spirv: // caches the buffer + stride eagerly so it dominates every sibling body). The adjoint offset for f32 adstacks // is a derived `OpIAdd` rather than an extra buffer load - mirrors the host launcher's `offset + max_size` // prefix-sum layout for the primal/adjoint pair. +// Allocate the per-task overflow-signal accumulator at the function entry block and zero-initialize it. +// Holds the maximum `stack_id + 1` value seen across all overflowing push sites in this thread; 0 means no +// overflow. Read + conditionally atomic-max'd to the host-visible AdStackOverflow buffer at task-end. +// Called eagerly from each `generate_*_kernel` at task start so the init-store dominates every push site +// and every task-end read; lazy alloc + init at the first push site would put the store inside whatever +// conditional block the first push happened to live in, leaving the var undefined on paths that bypass it. +spirv::Value TaskCodegen::ensure_any_overflow_signal_var() { + if (any_overflow_signal_var_.id != 0) { + return any_overflow_signal_var_; + } + any_overflow_signal_var_ = ir_->alloca_variable(ir_->u32_type()); + ir_->store_variable(any_overflow_signal_var_, ir_->uint_immediate_number(ir_->u32_type(), 0)); + return any_overflow_signal_var_; +} + +// Emit the task-end overflow check at the current insertion point, just before the kernel's `OpReturn`. +// Reads `any_overflow_signal_var_`; if non-zero, atomic-max'es it into slot 0 of the host-visible +// AdStackOverflow buffer. The buffer is host-readable (Apple Silicon shared memory; Vulkan +// HOST_VISIBLE | HOST_COHERENT), so the host polls slot 0 directly without any DtoH or sync drain. No-op +// when the task has no adstack push sites (the var is unallocated). +void TaskCodegen::emit_adstack_task_end_overflow_check() { + // Skip the entire emit (including the AdStackOverflow / AdStackTaskRegistryId buffer accesses) when + // the task body never visited an `AdStackPushStmt`. Forward-only tasks would otherwise force the + // launcher's bind path to wire AdStackTaskRegistryId for kernels where + // `publish_adstack_metadata_spirv` never allocated the buffer (no task in the kernel has adstacks), + // crashing Metal's `rw_buffer` device-equality assertion on the kDeviceNullAllocation fallback. + if (!task_has_adstack_push_ || any_overflow_signal_var_.id == 0) { + return; + } + spirv::Value zero = ir_->uint_immediate_number(ir_->u32_type(), 0); + spirv::Value cur = ir_->load_variable(any_overflow_signal_var_, ir_->u32_type()); + spirv::Value has_overflow = ir_->ne(cur, zero); + spirv::Label then_label = ir_->new_label(); + spirv::Label merge_label = ir_->new_label(); + ir_->make_inst(spv::OpSelectionMerge, merge_label, spv::SelectionControlMaskNone); + ir_->make_inst(spv::OpBranchConditional, has_overflow, then_label, merge_label); + ir_->start_label(then_label); + { + spirv::Value overflow_buffer = get_buffer_value(BufferType::AdStackOverflow, PrimitiveType::u32); + spirv::Value overflow_signal_ptr = + ir_->struct_array_access(ir_->u32_type(), overflow_buffer, ir_->uint_immediate_number(ir_->i32_type(), 0)); + ir_->make_value(spv::OpAtomicUMax, ir_->u32_type(), overflow_signal_ptr, + /*scope=*/ir_->const_i32_one_, + /*semantics=*/ir_->const_i32_zero_, cur); + // Record the offending task's `Program::adstack_sizing_info_registry_` id into slot 1 via + // `cmpxchg(0, registry_id)`. The launcher pre-writes the registry id into + // `AdStackTaskRegistryId[task_id_in_kernel]` per task; the codegen reads that slot at the + // task-end emit and atomically swaps it into AdStackOverflow[1] when the latter is still 0 + // (i.e. this is the FIRST overflowing thread across all tasks in the dispatch). The host + // raise site reads slot 1 and routes through `Program::diagnose_adstack_overflow_message` to + // produce a kernel-name + offload-task-index identity block. + spirv::Value task_registry_buffer = get_buffer_value(BufferType::AdStackTaskRegistryId, PrimitiveType::u32); + spirv::Value task_registry_ptr = ir_->struct_array_access( + ir_->u32_type(), task_registry_buffer, ir_->uint_immediate_number(ir_->i32_type(), task_id_in_kernel_)); + spirv::Value registry_id = ir_->load_variable(task_registry_ptr, ir_->u32_type()); + spirv::Value overflow_task_id_ptr = + ir_->struct_array_access(ir_->u32_type(), overflow_buffer, ir_->uint_immediate_number(ir_->i32_type(), 1)); + ir_->make_value(spv::OpAtomicCompareExchange, ir_->u32_type(), overflow_task_id_ptr, + /*scope=*/ir_->const_i32_one_, /*sem_eq=*/ir_->const_i32_zero_, + /*sem_neq=*/ir_->const_i32_zero_, registry_id, /*comparator=*/zero); + ir_->make_inst(spv::OpBranch, merge_label); + } + ir_->start_label(merge_label); +} + void TaskCodegen::ensure_ad_stack_metadata_loaded(AdStackSpirv &info) { if (info.offset_val.id != 0) { return; @@ -2758,6 +2853,7 @@ spirv::SType TaskCodegen::ad_stack_backing_type(const AdStackSpirv &info) const } void TaskCodegen::visit(AdStackPushStmt *stmt) { + task_has_adstack_push_ = true; auto &info = ad_stacks_.at(stmt->stack); ensure_ad_stack_metadata_loaded(info); spirv::Value count = ir_->load_variable(info.count_var, ir_->u32_type()); @@ -2781,48 +2877,12 @@ void TaskCodegen::visit(AdStackPushStmt *stmt) { return; } - if (compile_config_ && compile_config_->debug) { - // Debug build: map an OOB push to the last valid slot via a `GLSLstd450UMin` clamp, issue the primal/adjoint store, - // and publish `signal = (count >= max_size) ? stack_id + 1 : 0` to the host-visible AdStackOverflow buffer via - // `OpAtomicUMax`. The atomic-max with 0 cannot raise the host-visible value, so the runtime only sees the flag set - // on an actual overflow; concurrent threads that all witness the same overflow on the same stack publish the same - // value deterministically. The clamp + OpSelect formulation collapses what would otherwise be a per-push structured - // if-then-else region into straight-line code, which spirv-cross emits as straight-line MSL - critical for - // reverse-grad kernels with hundreds of adstacks pushed inside an inner loop on Apple's MSL-compiler-service - // shader-size threshold. `max_val` is the runtime-published AdStackMetadata bound cached on `info.max_size_val` by - // `ensure_ad_stack_metadata_loaded`, not a compile-time immediate. The gate is `debug` (not `check_out_of_bound`) - // so the field bounds check and the adstack overflow check stay on independent flags - Metal / Vulkan force-disable - // `check_out_of_bound` because they lack `Extension::assertion`, but `debug` reaches this codepath unaffected. - spirv::Value max_val = info.max_size_val; - spirv::Value max_minus_one = ir_->sub(max_val, one); - spirv::Value clamped_idx = ir_->call_glsl450(ir_->u32_type(), GLSLstd450UMin, count, max_minus_one); - spirv::Value primal_ptr = ad_stack_slot_ptr(info, clamped_idx, /*primal=*/true); - ir_->store_variable(primal_ptr, val); - if (info.heap_kind != AdStackHeapKind::heap_int) { - spirv::Value adjoint_ptr = ad_stack_slot_ptr(info, clamped_idx, /*primal=*/false); - ir_->store_variable(adjoint_ptr, ir_->get_zero(backing_type)); - } - ir_->store_variable(info.count_var, ir_->add(count, one)); - spirv::Value overflow_signal = - ir_->select(ir_->ge(count, max_val), ir_->uint_immediate_number(ir_->u32_type(), info.stack_id + 1), - ir_->uint_immediate_number(ir_->u32_type(), 0)); - spirv::Value overflow_buffer = get_buffer_value(BufferType::AdStackOverflow, PrimitiveType::u32); - spirv::Value overflow_ptr = - ir_->struct_array_access(ir_->u32_type(), overflow_buffer, ir_->uint_immediate_number(ir_->i32_type(), 0)); - ir_->make_value(spv::OpAtomicUMax, ir_->u32_type(), overflow_ptr, - /*scope=*/ir_->const_i32_one_, - /*semantics=*/ir_->const_i32_zero_, overflow_signal); - return; - } - - // Release build: trust the bound from `determine_ad_stack_size` (or the host-evaluated `size_expr` when the static - // pre-pass leaves a runtime-bound stack). Skip the clamp and the `OpAtomicUMax` overflow signal. Mirrors LLVM's - // release-build push, which also drops the runtime overflow guard - any unresolved stack is a hard pre-pass error - // there, and on the SPIR-V heap-backing path the host launcher is responsible for publishing a `max_size` that fits - // the kernel's actual push depth (and growing the heap if it does not, see the runtime-capacity grow-and-retry path - // in `runtime/gfx/runtime.cpp`). Cuts the per-push MSL emit down to the slot stores plus the `count++`, which is - // the dominant lever for keeping cross-compiled MSL below Apple's MSL-compiler-service shader-size limits on - // arches that hit them. + // Plain store + count increment. Always-on overflow detection runs alongside the store via a + // thread-local OpUMax accumulator (`any_overflow_signal_var_`); the per-push cost is two register + // operations (compare + max-update). One conditional `OpAtomicUMax` to the host-visible AdStackOverflow + // buffer is emitted ONCE per task at task-end (see `emit_adstack_task_end_overflow_check`), so push + // sites do not touch global memory for overflow signaling. + spirv::Value max_val = info.max_size_val; spirv::Value primal_ptr = ad_stack_slot_ptr(info, count, /*primal=*/true); ir_->store_variable(primal_ptr, val); if (info.heap_kind != AdStackHeapKind::heap_int) { @@ -2830,6 +2890,15 @@ void TaskCodegen::visit(AdStackPushStmt *stmt) { ir_->store_variable(adjoint_ptr, ir_->get_zero(backing_type)); } ir_->store_variable(info.count_var, ir_->add(count, one)); + // Update the per-task overflow-signal accumulator. `signal = (count >= max) ? stack_id + 1 : 0`; running max + // across all push sites in this thread. No global memory access. + spirv::Value any_overflow_var = ensure_any_overflow_signal_var(); + spirv::Value overflow_signal = + ir_->select(ir_->ge(count, max_val), ir_->uint_immediate_number(ir_->u32_type(), info.stack_id + 1), + ir_->uint_immediate_number(ir_->u32_type(), 0)); + spirv::Value prev = ir_->load_variable(any_overflow_var, ir_->u32_type()); + spirv::Value updated = ir_->call_glsl450(ir_->u32_type(), GLSLstd450UMax, prev, overflow_signal); + ir_->store_variable(any_overflow_var, updated); } void TaskCodegen::visit(AdStackPopStmt *stmt) { diff --git a/quadrants/program/adstack_size_expr_eval.cpp b/quadrants/program/adstack_size_expr_eval.cpp index 6f68390442..068331d515 100644 --- a/quadrants/program/adstack_size_expr_eval.cpp +++ b/quadrants/program/adstack_size_expr_eval.cpp @@ -25,15 +25,24 @@ namespace quadrants::lang { namespace { +using ReadSink = std::vector; + +// Forward-declared, defined further down. Reads SNode `snode_id` at `indices` via the per-launch read +// cache (when active) so multiple size-expr trees evaluated within the same outer launch share a single +// reader-kernel dispatch per `(snode_id, indices)` pair. +int64_t read_field_with_launch_cache(int snode_id, const std::vector &indices, Program *prog); + int64_t evaluate_node(const SerializedSizeExpr &expr, int32_t node_idx, - const std::unordered_map &bound_vars, + std::unordered_map &bound_vars, Program *prog, - LaunchContextBuilder *ctx); + LaunchContextBuilder *ctx, + ReadSink *reads); int64_t evaluate_field_load(const SerializedSizeExprNode &node, - const std::unordered_map &bound_vars, - Program *prog) { + std::unordered_map &bound_vars, + Program *prog, + ReadSink *reads) { QD_ASSERT_INFO(node.snode_id >= 0, "SerializedSizeExpr FieldLoad with no snode_id"); SNode *snode = prog->get_snode_by_id(node.snode_id); QD_ASSERT_INFO(snode != nullptr, @@ -54,13 +63,28 @@ int64_t evaluate_field_load(const SerializedSizeExprNode &node, indices.push_back(static_cast(it->second)); } } - auto accessors = prog->get_snode_rw_accessors_bank().get(snode); - return accessors.read_int(indices); + int64_t v = read_field_with_launch_cache(node.snode_id, indices, prog); + if (reads != nullptr) { + AdStackCache::SizeExprReadObservation obs; + obs.kind = AdStackCache::SizeExprReadObservation::FieldLoadObs; + obs.snode_id = node.snode_id; + obs.indices = std::move(indices); + obs.arg_shape_axis = 0; + obs.prim_dt = 0; + obs.observed_value = v; + // Snapshot the SNode's write gen so the next replay can fast-skip when no kernel has written this SNode + // since record time (the dominant case for a steady-state reverse-mode loop with stable bounds). + obs.observed_gen = prog->adstack_cache().snode_write_gen(node.snode_id); + reads->push_back(std::move(obs)); + } + return v; } int64_t evaluate_external_tensor_read(const SerializedSizeExprNode &node, - const std::unordered_map &bound_vars, - LaunchContextBuilder *ctx) { + std::unordered_map &bound_vars, + Program *prog, + LaunchContextBuilder *ctx, + ReadSink *reads) { QD_ASSERT_INFO(ctx != nullptr, "SerializedSizeExpr ExternalTensorRead evaluated with no LaunchContextBuilder; the launcher " "must pass the current launch's context in"); @@ -104,30 +128,60 @@ int64_t evaluate_external_tensor_read(const SerializedSizeExprNode &node, } } auto prim_dt = static_cast(node.const_value); + int64_t v; switch (prim_dt) { case PrimitiveTypeID::i32: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; case PrimitiveTypeID::i64: - return static_cast(data_ptr)[linear]; + v = static_cast(data_ptr)[linear]; + break; case PrimitiveTypeID::u32: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; case PrimitiveTypeID::u64: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; case PrimitiveTypeID::i16: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; case PrimitiveTypeID::u16: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; case PrimitiveTypeID::i8: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; case PrimitiveTypeID::u8: - return static_cast(static_cast(data_ptr)[linear]); + v = static_cast(static_cast(data_ptr)[linear]); + break; default: QD_ERROR("SerializedSizeExpr ExternalTensorRead: unsupported element type {}", node.const_value); + v = 0; } - return 0; + if (reads != nullptr) { + AdStackCache::SizeExprReadObservation obs; + obs.kind = AdStackCache::SizeExprReadObservation::ExternalReadObs; + obs.snode_id = 0; + obs.indices.reserve(resolved.size()); + for (auto r : resolved) + obs.indices.push_back(static_cast(r)); + obs.arg_id_path = node.arg_id_path; + obs.arg_shape_axis = 0; + obs.prim_dt = static_cast(prim_dt); + obs.observed_value = v; + obs.observed_devalloc = data_ptr; + if (prog != nullptr) { + // Snapshot the ndarray's data gen so the next replay can fast-skip when no kernel / Ndarray API write + // has touched the underlying buffer since record time. Mirrors the FieldLoad fast-skip; covers the same + // steady-state hot path for ndarray-bounded reverse-mode loops. + obs.observed_gen = prog->adstack_cache().ndarray_data_gen(data_ptr); + } + reads->push_back(std::move(obs)); + } + return v; } -int64_t evaluate_external_tensor_shape(const SerializedSizeExprNode &node, LaunchContextBuilder *ctx) { +int64_t evaluate_external_tensor_shape(const SerializedSizeExprNode &node, LaunchContextBuilder *ctx, ReadSink *reads) { QD_ASSERT_INFO(ctx != nullptr, "SerializedSizeExpr ExternalTensorShape evaluated with no LaunchContextBuilder; the launcher " "must pass the current launch's context into the evaluator to resolve ndarray shapes"); @@ -141,14 +195,26 @@ int64_t evaluate_external_tensor_shape(const SerializedSizeExprNode &node, Launc // garbage - often zero when the adjacent field is zero-initialised - and the containing tree collapses to // zero. The adstack max_size is clamped to 1 on a zero tree result, which under-bounds real push counts and // trips an overflow assertion at the next `qd.sync()`. - return static_cast(ctx->get_struct_arg_host(arg_indices)); + int64_t v = static_cast(ctx->get_struct_arg_host(arg_indices)); + if (reads != nullptr) { + AdStackCache::SizeExprReadObservation obs; + obs.kind = AdStackCache::SizeExprReadObservation::ExternalShapeObs; + obs.snode_id = 0; + obs.arg_id_path = node.arg_id_path; + obs.arg_shape_axis = node.arg_shape_axis; + obs.prim_dt = 0; + obs.observed_value = v; + reads->push_back(std::move(obs)); + } + return v; } int64_t evaluate_node(const SerializedSizeExpr &expr, int32_t node_idx, - const std::unordered_map &bound_vars, + std::unordered_map &bound_vars, Program *prog, - LaunchContextBuilder *ctx) { + LaunchContextBuilder *ctx, + ReadSink *reads) { QD_ASSERT_INFO(node_idx >= 0 && static_cast(node_idx) < expr.nodes.size(), "SerializedSizeExpr node_idx {} out of bounds (size={})", node_idx, expr.nodes.size()); const auto &node = expr.nodes[node_idx]; @@ -156,23 +222,23 @@ int64_t evaluate_node(const SerializedSizeExpr &expr, case SizeExpr::Kind::Const: return node.const_value; case SizeExpr::Kind::FieldLoad: - return evaluate_field_load(node, bound_vars, prog); + return evaluate_field_load(node, bound_vars, prog, reads); case SizeExpr::Kind::Add: - return evaluate_node(expr, node.operand_a, bound_vars, prog, ctx) + - evaluate_node(expr, node.operand_b, bound_vars, prog, ctx); + return evaluate_node(expr, node.operand_a, bound_vars, prog, ctx, reads) + + evaluate_node(expr, node.operand_b, bound_vars, prog, ctx, reads); case SizeExpr::Kind::Sub: - return std::max(evaluate_node(expr, node.operand_a, bound_vars, prog, ctx) - - evaluate_node(expr, node.operand_b, bound_vars, prog, ctx), + return std::max(evaluate_node(expr, node.operand_a, bound_vars, prog, ctx, reads) - + evaluate_node(expr, node.operand_b, bound_vars, prog, ctx, reads), 0); case SizeExpr::Kind::Mul: - return evaluate_node(expr, node.operand_a, bound_vars, prog, ctx) * - evaluate_node(expr, node.operand_b, bound_vars, prog, ctx); + return evaluate_node(expr, node.operand_a, bound_vars, prog, ctx, reads) * + evaluate_node(expr, node.operand_b, bound_vars, prog, ctx, reads); case SizeExpr::Kind::Max: - return std::max(evaluate_node(expr, node.operand_a, bound_vars, prog, ctx), - evaluate_node(expr, node.operand_b, bound_vars, prog, ctx)); + return std::max(evaluate_node(expr, node.operand_a, bound_vars, prog, ctx, reads), + evaluate_node(expr, node.operand_b, bound_vars, prog, ctx, reads)); case SizeExpr::Kind::MaxOverRange: { - int64_t begin = evaluate_node(expr, node.operand_a, bound_vars, prog, ctx); - int64_t end = evaluate_node(expr, node.operand_b, bound_vars, prog, ctx); + int64_t begin = evaluate_node(expr, node.operand_a, bound_vars, prog, ctx, reads); + int64_t end = evaluate_node(expr, node.operand_b, bound_vars, prog, ctx, reads); // Guard against pathological trip counts. The evaluator walks `[begin, end)` linearly and re-evaluates the // body at every i; a range of several million would stall the launch hot path for seconds. Real reverse-mode // trip counts sit well below this cap (a few hundred to a few thousand in practice); anything above is @@ -183,14 +249,24 @@ int64_t evaluate_node(const SerializedSizeExpr &expr, "Shrink the enclosing reverse-mode loop or restructure the `SizeExpr` source kernel.", end - begin, kMaxOverRangeIterations); int64_t result = 0; - auto extended = bound_vars; + // Bind `var_id` in `bound_vars` for the duration of the loop and restore the outer-scope value (or erase, if + // there was none) before returning, so nested `MaxOverRange` bindings of the same `var_id` stay correct without + // cloning the entire map per iteration. + auto prev_it = bound_vars.find(node.var_id); + bool had_prev = prev_it != bound_vars.end(); + int64_t prev_val = had_prev ? prev_it->second : 0; for (int64_t i = begin; i < end; ++i) { - extended[node.var_id] = i; - int64_t v = evaluate_node(expr, node.body_node_idx, extended, prog, ctx); + bound_vars[node.var_id] = i; + int64_t v = evaluate_node(expr, node.body_node_idx, bound_vars, prog, ctx, reads); if (v > result) { result = v; } } + if (had_prev) { + bound_vars[node.var_id] = prev_val; + } else { + bound_vars.erase(node.var_id); + } return result; } case SizeExpr::Kind::BoundVariable: { @@ -201,9 +277,9 @@ int64_t evaluate_node(const SerializedSizeExpr &expr, return it->second; } case SizeExpr::Kind::ExternalTensorShape: - return evaluate_external_tensor_shape(node, ctx); + return evaluate_external_tensor_shape(node, ctx, reads); case SizeExpr::Kind::ExternalTensorRead: - return evaluate_external_tensor_read(node, bound_vars, ctx); + return evaluate_external_tensor_read(node, bound_vars, prog, ctx, reads); } QD_ERROR("unreachable SerializedSizeExpr kind {}", node.kind); return 0; @@ -414,7 +490,8 @@ int32_t encode_subtree(const SerializedSizeExpr &src, LaunchContextBuilder *ctx, const FieldLoadDeviceEmitter &fl_emitter, std::vector &out_nodes, - std::vector &out_indices) { + std::vector &out_indices, + ReadSink *reads) { QD_ASSERT_INFO(src_idx >= 0 && static_cast(src_idx) < src.nodes.size(), "encode_subtree: src_idx {} out of bounds (size={})", src_idx, src.nodes.size()); const bool subtree_needs_device = contains_device_leaf[src_idx]; @@ -426,7 +503,7 @@ int32_t encode_subtree(const SerializedSizeExpr &src, // `FieldLoad` / `ExternalTensorShape` leaves - the device interpreter does not know how to walk SNodes or // index into `args_type`. std::unordered_map empty_bound; - int64_t val = evaluate_node(src, src_idx, empty_bound, prog, ctx); + int64_t val = evaluate_node(src, src_idx, empty_bound, prog, ctx, reads); AdStackSizeExprDeviceNode dn = make_empty_device_node(static_cast(AdStackSizeExprDeviceKind::kConst)); dn.const_value = val; out_nodes.push_back(dn); @@ -454,9 +531,9 @@ int32_t encode_subtree(const SerializedSizeExpr &src, case SizeExpr::Kind::Mul: case SizeExpr::Kind::Max: { int32_t a = encode_subtree(src, node.operand_a, contains_device_leaf, free_vars, var_id_remap, prog, ctx, - fl_emitter, out_nodes, out_indices); + fl_emitter, out_nodes, out_indices, reads); int32_t b = encode_subtree(src, node.operand_b, contains_device_leaf, free_vars, var_id_remap, prog, ctx, - fl_emitter, out_nodes, out_indices); + fl_emitter, out_nodes, out_indices, reads); AdStackSizeExprDeviceKind dk = AdStackSizeExprDeviceKind::kAdd; if (kind == SizeExpr::Kind::Sub) dk = AdStackSizeExprDeviceKind::kSub; @@ -472,11 +549,11 @@ int32_t encode_subtree(const SerializedSizeExpr &src, } case SizeExpr::Kind::MaxOverRange: { int32_t a = encode_subtree(src, node.operand_a, contains_device_leaf, free_vars, var_id_remap, prog, ctx, - fl_emitter, out_nodes, out_indices); + fl_emitter, out_nodes, out_indices, reads); int32_t b = encode_subtree(src, node.operand_b, contains_device_leaf, free_vars, var_id_remap, prog, ctx, - fl_emitter, out_nodes, out_indices); + fl_emitter, out_nodes, out_indices, reads); int32_t body = encode_subtree(src, node.body_node_idx, contains_device_leaf, free_vars, var_id_remap, prog, ctx, - fl_emitter, out_nodes, out_indices); + fl_emitter, out_nodes, out_indices, reads); AdStackSizeExprDeviceNode dn = make_empty_device_node(static_cast(AdStackSizeExprDeviceKind::kMaxOverRange)); dn.operand_a = a; @@ -610,12 +687,840 @@ int32_t encode_subtree(const SerializedSizeExpr &src, } // namespace +namespace { + +// Per-launch cache of `FieldLoad` re-reads, keyed by `(snode_id, indices)`. Within one host-side eval root +// call the SNode field values are pinned (no other kernel runs concurrently), so deduping repeats across +// the size-expr trees evaluated in that window is correctness-safe. +struct LaunchScopedReadCache { + struct Key { + int snode_id; + std::vector indices; + bool operator==(const Key &o) const noexcept { + return snode_id == o.snode_id && indices == o.indices; + } + }; + struct KeyHash { + std::size_t operator()(const Key &k) const noexcept { + std::size_t h = std::hash{}(k.snode_id); + for (int v : k.indices) { + h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ull + (h << 6) + (h >> 2); + } + return h; + } + }; + std::unordered_map map; +}; +thread_local LaunchScopedReadCache *t_launch_read_cache = nullptr; + +int64_t read_field_with_launch_cache(int snode_id, const std::vector &indices, Program *prog) { + SNode *snode = prog->get_snode_by_id(snode_id); + if (snode == nullptr) { + return std::numeric_limits::min(); + } + if (t_launch_read_cache != nullptr) { + LaunchScopedReadCache::Key key{snode_id, indices}; + auto it = t_launch_read_cache->map.find(key); + if (it != t_launch_read_cache->map.end()) { + return it->second; + } + int64_t v = prog->get_snode_rw_accessors_bank().get(snode).read_int(indices); + t_launch_read_cache->map.emplace(std::move(key), v); + return v; + } + return prog->get_snode_rw_accessors_bank().get(snode).read_int(indices); +} + +// Read the input that `obs` describes against the live state and `ctx`. Caller compares the result to +// `obs.observed_value` to decide whether the cached `SizeExprCacheEntry` is still valid. Each `obs.kind` +// mirrors the corresponding leaf in `evaluate_field_load` / `evaluate_external_tensor_shape` / +// `evaluate_external_tensor_read`. +int64_t replay_one_observation(const AdStackCache::SizeExprReadObservation &obs, + Program *prog, + LaunchContextBuilder *ctx) { + using Obs = AdStackCache::SizeExprReadObservation; + switch (obs.kind) { + case Obs::FieldLoadObs: { + // Gen-counter fast skip: when no kernel has bumped this SNode's write generation since record time, + // the underlying field value cannot have changed and we can return the recorded `observed_value` + // without dispatching a reader kernel. The dispatch is the dominant per-launch cost on the hot path + // for steady-state reverse-mode loops with stable bounds. + if (prog != nullptr && prog->adstack_cache().snode_write_gen(obs.snode_id) == obs.observed_gen) { + return obs.observed_value; + } + int64_t v = read_field_with_launch_cache(obs.snode_id, obs.indices, prog); + if (v == std::numeric_limits::min()) { + return obs.observed_value + 1; // force a mismatch if SNode disappeared + } + return v; + } + case Obs::ExternalShapeObs: { + if (ctx == nullptr) { + return obs.observed_value + 1; + } + std::vector arg_indices(obs.arg_id_path.begin(), obs.arg_id_path.end()); + arg_indices.push_back(TypeFactory::SHAPE_POS_IN_NDARRAY); + arg_indices.push_back(obs.arg_shape_axis); + return static_cast(ctx->get_struct_arg_host(arg_indices)); + } + case Obs::ExternalReadObs: { + if (ctx == nullptr || obs.arg_id_path.empty()) { + return obs.observed_value + 1; + } + int arg_id = obs.arg_id_path[0]; + ArgArrayPtrKey key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto it = ctx->array_ptrs.find(key); + if (it == ctx->array_ptrs.end()) { + return obs.observed_value + 1; + } + void *data_ptr = it->second; + // Gen-counter fast skip: when the data pointer is the same `DeviceAllocation *` we observed at record + // time AND its data generation has not been bumped since (no kernel write, no host-side `Ndarray.write` + // / `fill`), the underlying scalar cannot have changed and we can return the recorded value without + // dereferencing the device pointer (which on GPU would be a DtoH copy, on CPU a host load). + if (prog != nullptr && data_ptr == obs.observed_devalloc && + prog->adstack_cache().ndarray_data_gen(data_ptr) == obs.observed_gen) { + return obs.observed_value; + } + int64_t linear = 0; + int64_t stride = 1; + for (std::size_t i = obs.indices.size(); i > 0; --i) { + linear += static_cast(obs.indices[i - 1]) * stride; + if (i - 1 > 0) { + std::vector sh_idx(obs.arg_id_path.begin(), obs.arg_id_path.end()); + sh_idx.push_back(TypeFactory::SHAPE_POS_IN_NDARRAY); + sh_idx.push_back(static_cast(i - 1)); + stride *= static_cast(ctx->get_struct_arg_host(sh_idx)); + } + } + switch (static_cast(obs.prim_dt)) { + case PrimitiveTypeID::i32: + return static_cast(static_cast(data_ptr)[linear]); + case PrimitiveTypeID::i64: + return static_cast(data_ptr)[linear]; + case PrimitiveTypeID::u32: + return static_cast(static_cast(data_ptr)[linear]); + case PrimitiveTypeID::u64: + return static_cast(static_cast(data_ptr)[linear]); + case PrimitiveTypeID::i16: + return static_cast(static_cast(data_ptr)[linear]); + case PrimitiveTypeID::u16: + return static_cast(static_cast(data_ptr)[linear]); + case PrimitiveTypeID::i8: + return static_cast(static_cast(data_ptr)[linear]); + case PrimitiveTypeID::u8: + return static_cast(static_cast(data_ptr)[linear]); + default: + return obs.observed_value + 1; + } + } + } + return obs.observed_value + 1; +} +} // namespace + +bool AdStackCache::try_size_expr_cache_hit(Program *prog, + const SerializedSizeExpr *expr_key, + LaunchContextBuilder *ctx, + int64_t &out_result) { + auto it = size_expr_cache_.find(expr_key); + if (it == size_expr_cache_.end()) { + return false; + } + const auto &entry = it->second; + for (const auto &obs : entry.reads) { + int64_t now = replay_one_observation(obs, prog, ctx); + if (now != obs.observed_value) { + size_expr_cache_.erase(it); + return false; + } + } + out_result = entry.result; + return true; +} + +void AdStackCache::record_size_expr_eval(const SerializedSizeExpr *expr_key, + int64_t result, + std::vector reads) { + size_expr_cache_[expr_key] = SizeExprCacheEntry{result, std::move(reads)}; +} + +bool AdStackCache::try_spirv_bytecode_cache_hit(Program *prog, + const void *attribs_key, + LaunchContextBuilder *ctx, + std::vector &out_bytecode) { + auto it = spirv_bytecode_cache_.find(attribs_key); + if (it == spirv_bytecode_cache_.end()) { + return false; + } + const auto &entry = it->second; + for (const auto &obs : entry.reads) { + int64_t now = replay_one_observation(obs, prog, ctx); + if (now != obs.observed_value) { + spirv_bytecode_cache_.erase(it); + return false; + } + } + out_bytecode = entry.bytecode; + return true; +} + +void AdStackCache::record_spirv_bytecode_eval(const void *attribs_key, + std::vector bytecode, + std::vector reads) { + spirv_bytecode_cache_[attribs_key] = SpirvBytecodeCacheEntry{std::move(bytecode), std::move(reads)}; +} + +void AdStackCache::record_per_task_ad_stack(const void *attribs_key, + std::vector metadata, + uint32_t stride_float, + uint32_t stride_int, + std::vector> snode_gens, + std::vector> arg_gens) { + per_task_ad_stack_cache_[attribs_key] = PerTaskAdStackCacheEntry{std::move(metadata), stride_float, stride_int, + std::move(snode_gens), std::move(arg_gens)}; +} + +bool AdStackCache::try_per_task_ad_stack_cache_hit(const void *attribs_key, + LaunchContextBuilder *ctx, + PerTaskAdStackCacheEntry &out) { + auto it = per_task_ad_stack_cache_.find(attribs_key); + if (it == per_task_ad_stack_cache_.end()) { + return false; + } + const auto &entry = it->second; + for (const auto &snode_pair : entry.snode_gens) { + if (snode_write_gen(snode_pair.first) != snode_pair.second) { + per_task_ad_stack_cache_.erase(it); + return false; + } + } + for (const auto &arg_tuple : entry.arg_gens) { + int arg_id = std::get<0>(arg_tuple); + void *recorded_devalloc = std::get<1>(arg_tuple); + uint64_t recorded_gen = std::get<2>(arg_tuple); + void *current_devalloc = nullptr; + if (ctx != nullptr) { + ArgArrayPtrKey key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto ap_it = ctx->array_ptrs.find(key); + if (ap_it != ctx->array_ptrs.end()) { + current_devalloc = ap_it->second; + } + } + if (current_devalloc != recorded_devalloc) { + per_task_ad_stack_cache_.erase(it); + return false; + } + if (ndarray_data_gen(recorded_devalloc) != recorded_gen) { + per_task_ad_stack_cache_.erase(it); + return false; + } + } + out = entry; + return true; +} + +void AdStackCache::record_llvm_per_task_ad_stack(const void *attribs_key, + std::vector offsets, + std::vector max_sizes, + uint64_t stride_combined, + uint64_t stride_float, + uint64_t stride_int, + std::vector> snode_gens, + std::vector> arg_gens) { + llvm_per_task_ad_stack_cache_[attribs_key] = + LlvmPerTaskAdStackCacheEntry{std::move(offsets), std::move(max_sizes), stride_combined, stride_float, + stride_int, std::move(snode_gens), std::move(arg_gens)}; +} + +bool AdStackCache::try_llvm_per_task_ad_stack_cache_hit(const void *attribs_key, + LaunchContextBuilder *ctx, + LlvmPerTaskAdStackCacheEntry &out) { + auto it = llvm_per_task_ad_stack_cache_.find(attribs_key); + if (it == llvm_per_task_ad_stack_cache_.end()) { + return false; + } + const auto &entry = it->second; + for (const auto &snode_pair : entry.snode_gens) { + if (snode_write_gen(snode_pair.first) != snode_pair.second) { + llvm_per_task_ad_stack_cache_.erase(it); + return false; + } + } + for (const auto &arg_tuple : entry.arg_gens) { + int arg_id = std::get<0>(arg_tuple); + void *recorded_devalloc = std::get<1>(arg_tuple); + uint64_t recorded_gen = std::get<2>(arg_tuple); + void *current_devalloc = nullptr; + if (ctx != nullptr) { + ArgArrayPtrKey key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto ap_it = ctx->array_ptrs.find(key); + if (ap_it != ctx->array_ptrs.end()) { + current_devalloc = ap_it->second; + } + } + if (current_devalloc != recorded_devalloc) { + llvm_per_task_ad_stack_cache_.erase(it); + return false; + } + if (ndarray_data_gen(recorded_devalloc) != recorded_gen) { + llvm_per_task_ad_stack_cache_.erase(it); + return false; + } + } + out = entry; + return true; +} + +// Per-thread backing for `SizeExprLaunchScope`. The outer scope on each thread points `t_launch_read_cache` here +// after clearing the map; nested scopes are no-ops. +thread_local LaunchScopedReadCache t_launch_read_cache_storage{}; + +SizeExprLaunchScope::SizeExprLaunchScope() : owns_(t_launch_read_cache == nullptr) { + if (owns_) { + t_launch_read_cache_storage.map.clear(); + t_launch_read_cache = &t_launch_read_cache_storage; + } +} +SizeExprLaunchScope::~SizeExprLaunchScope() { + if (owns_) { + t_launch_read_cache = nullptr; + } +} + int64_t evaluate_adstack_size_expr(const SerializedSizeExpr &expr, Program *prog, LaunchContextBuilder *ctx) { if (expr.nodes.empty()) { return -1; } + // Open a `SizeExprLaunchScope` if no enclosing one is active, so repeated reads within this eval share + // the launch read cache. Callers that issue several `evaluate_adstack_size_expr` calls back-to-back + // should open their own scope to span all of them. + SizeExprLaunchScope local_scope; + + // Cache fast path: replay the recorded reads against the live state and reuse the cached result if + // every input still matches. The full walk runs only on cache miss. + if (prog != nullptr) { + int64_t cached; + if (prog->adstack_cache().try_size_expr_cache_hit(prog, &expr, ctx, cached)) { + return cached; + } + } std::unordered_map empty_bound_vars; - return evaluate_node(expr, static_cast(expr.nodes.size() - 1), empty_bound_vars, prog, ctx); + std::vector reads; + int64_t result = + evaluate_node(expr, static_cast(expr.nodes.size() - 1), empty_bound_vars, prog, ctx, &reads); + if (prog != nullptr) { + prog->adstack_cache().record_size_expr_eval(&expr, result, std::move(reads)); + } + return result; +} + +namespace { + +// Diagnose-time leaf reader: resolves an `ExternalTensorRead` against the captured +// `AdStackCache::DiagnoseLaunchSnapshot` and the program's `Device::map` interface. Returns -1 on any failure +// (missing arg in snapshot, unrecognised primitive type, mapping failure) so the caller can substitute the +// `?` placeholder for that stack while keeping the rest of the message intact. +// +// Single-scalar staging-buffer pattern (mirrors `Ndarray::read` in `program/ndarray.cpp`): allocate a tiny +// `host_read=true` staging buffer, `memcpy_internal` the one element from the ndarray's device buffer into +// staging, then map staging to read the value host-side. This works on every backend because every +// `Device` implementation supports `host_read=true` allocations + `map` + `memcpy_internal`. For `kNone` +// numpy passthrough the captured pointer is already host-readable; we read it directly. +int64_t read_diagnose_external_tensor(const SerializedSizeExprNode &node, + const std::vector &resolved_indices, + Program *prog, + const AdStackCache::DiagnoseLaunchSnapshot &snapshot) { + if (node.arg_id_path.empty()) { + return -1; + } + int arg_id = node.arg_id_path[0]; + auto ptr_it = snapshot.data_ptrs.find(arg_id); + if (ptr_it == snapshot.data_ptrs.end() || ptr_it->second == nullptr) { + return -1; + } + auto type_it = snapshot.dev_alloc_types.find(arg_id); + if (type_it == snapshot.dev_alloc_types.end()) { + return -1; + } + auto shape_it = snapshot.shapes.find(arg_id); + if (shape_it == snapshot.shapes.end()) { + return -1; + } + const std::vector &shape = shape_it->second; + // Compose C-order linear offset across resolved indices (mirrors `evaluate_external_tensor_read`'s stride + // math; we cannot share the helper because that one routes through `LaunchContextBuilder::get_struct_arg_host` + // which is not available here). + if (resolved_indices.size() > shape.size() && !shape.empty()) { + // More indices than rank - the size_expr was lowered against a different shape; skip. + return -1; + } + int64_t linear = 0; + int64_t stride = 1; + for (std::size_t i = resolved_indices.size(); i > 0; --i) { + linear += resolved_indices[i - 1] * stride; + if (i - 1 > 0 && i - 1 < shape.size()) { + stride *= static_cast(shape[i - 1]); + } + } + auto prim_dt = static_cast(node.const_value); + std::size_t elem_size = 0; + switch (prim_dt) { + case PrimitiveTypeID::i8: + case PrimitiveTypeID::u8: + elem_size = 1; + break; + case PrimitiveTypeID::i16: + case PrimitiveTypeID::u16: + elem_size = 2; + break; + case PrimitiveTypeID::i32: + case PrimitiveTypeID::u32: + elem_size = 4; + break; + case PrimitiveTypeID::i64: + case PrimitiveTypeID::u64: + elem_size = 8; + break; + default: + return -1; + } + std::size_t byte_offset = static_cast(linear) * elem_size; + // Decode the captured pointer to host bytes. + std::vector staging_bytes(elem_size); + if (type_it->second == LaunchContextBuilder::DevAllocType::kNone) { + // Numpy passthrough: ptr is already a raw host pointer. + const uint8_t *src = static_cast(ptr_it->second) + byte_offset; + std::memcpy(staging_bytes.data(), src, elem_size); + } else if (type_it->second == LaunchContextBuilder::DevAllocType::kNdarray) { + if (prog == nullptr) { + return -1; + } + auto *alloc = static_cast(ptr_it->second); + if (alloc == nullptr || alloc->device == nullptr) { + return -1; + } + Device::AllocParams params; + params.host_write = false; + params.host_read = true; + params.size = elem_size; + params.usage = AllocUsage::Storage; + auto [staging, alloc_res] = alloc->device->allocate_memory_unique(params); + if (alloc_res != RhiResult::success || !staging) { + return -1; + } + alloc->device->memcpy_internal(staging->get_ptr(), alloc->get_ptr(byte_offset), elem_size); + void *mapped = nullptr; + if (alloc->device->map(*staging, &mapped) != RhiResult::success || mapped == nullptr) { + return -1; + } + std::memcpy(staging_bytes.data(), mapped, elem_size); + alloc->device->unmap(*staging); + } else { + return -1; + } + // Sign- / zero-extend to int64 according to the captured primitive type. + switch (prim_dt) { + case PrimitiveTypeID::i8: + return static_cast(*reinterpret_cast(staging_bytes.data())); + case PrimitiveTypeID::u8: + return static_cast(*reinterpret_cast(staging_bytes.data())); + case PrimitiveTypeID::i16: + return static_cast(*reinterpret_cast(staging_bytes.data())); + case PrimitiveTypeID::u16: + return static_cast(*reinterpret_cast(staging_bytes.data())); + case PrimitiveTypeID::i32: + return static_cast(*reinterpret_cast(staging_bytes.data())); + case PrimitiveTypeID::u32: + return static_cast(*reinterpret_cast(staging_bytes.data())); + case PrimitiveTypeID::i64: + return *reinterpret_cast(staging_bytes.data()); + case PrimitiveTypeID::u64: + return static_cast(*reinterpret_cast(staging_bytes.data())); + default: + return -1; + } +} + +// Mirror of `evaluate_node` for diagnose-time evaluation. Same tree-walk semantics; differs only in the leaf +// case for `ExternalTensorRead` / `ExternalTensorShape`, which route through the snapshot + `Device::map` path +// instead of `LaunchContextBuilder`. Returns -1 on any leaf-resolution failure to short-circuit the rest of +// the walk and let the caller fall back to the static dual-cause body. +int64_t evaluate_node_for_diagnose(const SerializedSizeExpr &expr, + int32_t node_idx, + std::unordered_map &bound_vars, + Program *prog, + const AdStackCache::DiagnoseLaunchSnapshot &snapshot) { + if (node_idx < 0 || static_cast(node_idx) >= expr.nodes.size()) { + return -1; + } + const auto &node = expr.nodes[node_idx]; + switch (static_cast(node.kind)) { + case SizeExpr::Kind::Const: + return node.const_value; + case SizeExpr::Kind::FieldLoad: { + // Field reads stay on the existing host path - they do not depend on `LaunchContextBuilder` and the + // SNode reader-kernel dispatch is host-driven. We pass `nullptr` ReadSink so the recorded observations + // do not leak into the cache from a diagnose-only walk. + return evaluate_field_load(node, bound_vars, prog, /*reads=*/nullptr); + } + case SizeExpr::Kind::Add: { + int64_t a = evaluate_node_for_diagnose(expr, node.operand_a, bound_vars, prog, snapshot); + int64_t b = evaluate_node_for_diagnose(expr, node.operand_b, bound_vars, prog, snapshot); + if (a < 0 || b < 0) { + return -1; + } + return a + b; + } + case SizeExpr::Kind::Sub: { + int64_t a = evaluate_node_for_diagnose(expr, node.operand_a, bound_vars, prog, snapshot); + int64_t b = evaluate_node_for_diagnose(expr, node.operand_b, bound_vars, prog, snapshot); + if (a < 0 || b < 0) { + return -1; + } + return std::max(a - b, 0); + } + case SizeExpr::Kind::Mul: { + int64_t a = evaluate_node_for_diagnose(expr, node.operand_a, bound_vars, prog, snapshot); + int64_t b = evaluate_node_for_diagnose(expr, node.operand_b, bound_vars, prog, snapshot); + if (a < 0 || b < 0) { + return -1; + } + return a * b; + } + case SizeExpr::Kind::Max: { + int64_t a = evaluate_node_for_diagnose(expr, node.operand_a, bound_vars, prog, snapshot); + int64_t b = evaluate_node_for_diagnose(expr, node.operand_b, bound_vars, prog, snapshot); + if (a < 0 || b < 0) { + return -1; + } + return std::max(a, b); + } + case SizeExpr::Kind::MaxOverRange: { + int64_t begin = evaluate_node_for_diagnose(expr, node.operand_a, bound_vars, prog, snapshot); + int64_t end = evaluate_node_for_diagnose(expr, node.operand_b, bound_vars, prog, snapshot); + if (begin < 0 || end < 0) { + return -1; + } + // Same iteration cap as the live evaluator; refusing to enumerate prevents diagnose from stalling + // the error path on a pathological trip count. + constexpr int64_t kMaxOverRangeIterations = int64_t{1} << 24; + if (end > begin && end - begin > kMaxOverRangeIterations) { + return -1; + } + int64_t result = 0; + auto prev_it = bound_vars.find(node.var_id); + bool had_prev = prev_it != bound_vars.end(); + int64_t prev_val = had_prev ? prev_it->second : 0; + for (int64_t i = begin; i < end; ++i) { + bound_vars[node.var_id] = i; + int64_t v = evaluate_node_for_diagnose(expr, node.body_node_idx, bound_vars, prog, snapshot); + if (v < 0) { + if (had_prev) { + bound_vars[node.var_id] = prev_val; + } else { + bound_vars.erase(node.var_id); + } + return -1; + } + if (v > result) { + result = v; + } + } + if (had_prev) { + bound_vars[node.var_id] = prev_val; + } else { + bound_vars.erase(node.var_id); + } + return result; + } + case SizeExpr::Kind::BoundVariable: { + auto it = bound_vars.find(node.var_id); + if (it == bound_vars.end()) { + return -1; + } + return it->second; + } + case SizeExpr::Kind::ExternalTensorShape: { + if (node.arg_id_path.empty()) { + return -1; + } + int arg_id = node.arg_id_path[0]; + auto shape_it = snapshot.shapes.find(arg_id); + if (shape_it == snapshot.shapes.end()) { + return -1; + } + if (node.arg_shape_axis < 0 || static_cast(node.arg_shape_axis) >= shape_it->second.size()) { + return -1; + } + return static_cast(shape_it->second[node.arg_shape_axis]); + } + case SizeExpr::Kind::ExternalTensorRead: { + // Resolve indices from bound_vars first, then dispatch to the snapshot-aware reader. + std::vector resolved(node.indices.size()); + for (std::size_t i = 0; i < node.indices.size(); ++i) { + int32_t raw = node.indices[i]; + if (raw >= 0) { + resolved[i] = raw; + } else { + int32_t var_id = -(raw + 1); + auto bv = bound_vars.find(var_id); + if (bv == bound_vars.end()) { + return -1; + } + resolved[i] = bv->second; + } + } + return read_diagnose_external_tensor(node, resolved, prog, snapshot); + } + } + return -1; +} + +} // namespace + +int64_t evaluate_adstack_size_expr_for_diagnose(const SerializedSizeExpr &expr, Program *prog) { + if (expr.nodes.empty() || prog == nullptr) { + return -1; + } + const AdStackCache::DiagnoseLaunchSnapshot *snapshot = prog->adstack_cache().get_diagnose_snapshot(); + if (snapshot == nullptr) { + return -1; + } + std::unordered_map bound_vars; + return evaluate_node_for_diagnose(expr, static_cast(expr.nodes.size() - 1), bound_vars, prog, *snapshot); +} + +uint32_t AdStackCache::register_adstack_sizing_info(const void *identity_key, + const std::string &kernel_name, + int task_id_in_kernel, + std::vector allocated_max_sizes, + std::vector size_exprs) { + std::lock_guard lk(adstack_sizing_info_registry_mutex_); + // Idempotent re-registration: same `identity_key` yields the same id across re-compiles and updates the + // entry's metadata + size_exprs in place. The key is just an opaque dedup token - the registry never + // dereferences it; all data needed by the diagnose path is copied into the entry below. + auto it = adstack_sizing_info_id_by_ptr_.find(identity_key); + if (it != adstack_sizing_info_id_by_ptr_.end()) { + auto &entry = adstack_sizing_info_registry_[it->second]; + entry.kernel_name = kernel_name; + entry.task_id_in_kernel = task_id_in_kernel; + entry.allocated_max_sizes = std::move(allocated_max_sizes); + entry.size_exprs = std::move(size_exprs); + return it->second; + } + uint32_t id = static_cast(adstack_sizing_info_registry_.size()); + AdStackSizingInfoEntry entry; + entry.identity_key = identity_key; + entry.kernel_name = kernel_name; + entry.task_id_in_kernel = task_id_in_kernel; + entry.allocated_max_sizes = std::move(allocated_max_sizes); + entry.size_exprs = std::move(size_exprs); + adstack_sizing_info_registry_.push_back(std::move(entry)); + adstack_sizing_info_id_by_ptr_.emplace(identity_key, id); + return id; +} + +void AdStackCache::update_adstack_sizing_info_size_exprs(uint32_t id, std::vector size_exprs) { + std::lock_guard lk(adstack_sizing_info_registry_mutex_); + if (id == 0 || id >= adstack_sizing_info_registry_.size()) { + return; + } + adstack_sizing_info_registry_[id].size_exprs = std::move(size_exprs); +} + +std::optional AdStackCache::lookup_adstack_sizing_info(uint32_t id) const { + std::lock_guard lk(adstack_sizing_info_registry_mutex_); + if (id == 0 || id >= adstack_sizing_info_registry_.size()) { + return std::nullopt; + } + return adstack_sizing_info_registry_[id]; +} + +std::string AdStackCache::diagnose_adstack_overflow_message(uint32_t task_id) const { + return diagnose_adstack_overflow(task_id).message; +} + +AdStackCache::AdStackOverflowDiagnosis AdStackCache::diagnose_adstack_overflow(uint32_t task_id) const { + // Lazy LLVM capture: if the launcher stashed a pending ctx pointer for this launch (LLVM defers eager + // capture to avoid the per-launch snapshot cost), capture now before walking size_exprs. SPIR-V already + // captured eagerly at launch, so `pending_launch_ctx_` is null there. + if (pending_launch_ctx_ != nullptr) { + const_cast(this)->capture_diagnose_snapshot(*pending_launch_ctx_); + } + std::string identity_block; + std::string disambiguation_block; + // Cause classifier: when the synchronous re-run produces required > allocated for ANY stack, the most likely + // cause is an untracked tensor mutation (DLPack-bypass etc.). When all required <= allocated, the pre-pass + // undersized the bound (Quadrants bug). When we cannot re-evaluate (e.g. no captured launch snapshot, or a + // leaf type the diagnose evaluator does not support) we fall through to the static dual-cause body. + enum class Cause { Unknown, DLPackBypass, QuadrantsBug }; + Cause cause = Cause::Unknown; + + if (task_id != 0) { + auto entry_opt = lookup_adstack_sizing_info(task_id); + if (entry_opt.has_value()) { + const auto &entry = *entry_opt; + identity_block = " Offending task: kernel `" + entry.kernel_name + "` offload task #" + + std::to_string(entry.task_id_in_kernel) + "; per-stack allocated max_size = ["; + for (size_t i = 0; i < entry.allocated_max_sizes.size(); ++i) { + if (i != 0) { + identity_block += ", "; + } + identity_block += std::to_string(entry.allocated_max_sizes[i]); + } + identity_block += "].\n"; + + // Synchronous sizer rerun: walk each stack's `SerializedSizeExpr` and evaluate against the live host / + // SNode state. Stacks whose tree contains an `ExternalTensorShape` or `ExternalTensorRead` leaf go + // through the snapshot-based `evaluate_adstack_size_expr_for_diagnose` (see its declaration for the + // `Device::map` design rationale). Pure host-resolvable trees go through the standard host evaluator. + // The disambiguation is best-effort: if every stack's tree resolves we get a precise classification; + // otherwise we report what we have and fall back to the static dual-cause hint. + if (!entry.size_exprs.empty()) { + std::vector required_sizes; + std::vector required_known; + size_t any_grew = 0; + size_t any_unknown = 0; + size_t total = std::min(entry.size_exprs.size(), entry.allocated_max_sizes.size()); + for (size_t i = 0; i < total; ++i) { + const auto &expr = entry.size_exprs[i]; + bool host_resolvable = true; + for (const auto &node : expr.nodes) { + auto k = static_cast(node.kind); + if (k == SizeExpr::Kind::ExternalTensorShape || k == SizeExpr::Kind::ExternalTensorRead) { + host_resolvable = false; + break; + } + } + int64_t v = -1; + if (host_resolvable && !expr.nodes.empty()) { + // Pure host-resolvable: SNode field loads, constants, arithmetic. `ctx == nullptr` is safe because + // every leaf we kept is host-resolvable; ETS / ETR are the only kinds that touch ctx and we + // filtered them out. + SizeExprLaunchScope scope; + v = evaluate_adstack_size_expr(expr, prog_, nullptr); + } else if (!expr.nodes.empty()) { + // Tree contains ETR / ETS leaves. The diagnose evaluator resolves them through the captured launch + // snapshot (`Device::map`-based ndarray reads). On failure (no snapshot, allocation cannot be + // mapped, unsupported dtype) the helper returns -1 and we fall through to the `?` placeholder. + int64_t diag = evaluate_adstack_size_expr_for_diagnose(expr, prog_); + if (diag >= 0) { + v = diag; + } + } + required_sizes.push_back(v); + required_known.push_back(!expr.nodes.empty() && v >= 0); + if (required_known.back() && static_cast(v) > entry.allocated_max_sizes[i]) { + ++any_grew; + } + if (!required_known.back()) { + ++any_unknown; + } + } + if (any_grew > 0) { + cause = Cause::DLPackBypass; + } else if (any_unknown == 0 && total > 0) { + cause = Cause::QuadrantsBug; + } + // Only print the rerun line when at least one stack's bound resolves to a real value. With every leaf + // unresolved the line would be `required = [?, ?, ...]` which adds zero signal beyond the dual-cause + // body that follows; the omission keeps the message focused on actionable content. + if (any_unknown < total) { + disambiguation_block = " Synchronous sizer rerun: required max_size = ["; + for (size_t i = 0; i < required_sizes.size(); ++i) { + if (i != 0) { + disambiguation_block += ", "; + } + if (required_known[i]) { + disambiguation_block += std::to_string(required_sizes[i]); + } else { + disambiguation_block += "?"; + } + } + disambiguation_block += "]."; + if (any_unknown > 0) { + disambiguation_block += + " (`?` = sizer rerun could not resolve this stack's bound against the captured " + "launch state)."; + } + disambiguation_block += "\n"; + } + } + } + } + + std::string body; + if (cause == Cause::DLPackBypass) { + body = + "Cause (sync sizer rerun): a tensor backing a data-dependent loop bound was mutated 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 DLPack consumer. " + "The cached adstack capacity was sized against the value before the mutation. Recovery: route " + "the mutation through Quadrants APIs (`Ndarray.write` / `fill` / kernel writes) so the cache " + "invalidates correctly, OR set a generous initial cap if a workload-change milestone genuinely " + "grew capacity. Restart the iteration / training loop from a clean state.\n"; + } else if (cause == Cause::QuadrantsBug) { + body = + "Cause (sync sizer rerun): the freshly-computed required size does not exceed the allocated " + "size for any stack - this is a Quadrants bug. The pre-pass resolved the alloca to a bound " + "tighter than the actual runtime push count: either the enclosing loop shape is outside the " + "current `SizeExpr` grammar, or the Bellman-Ford analyzer undercounted the forward-pass " + "accumulation. Please file with the kernel IR (`QD_DUMP_IR=1`).\n"; + } else { + body = + "Two possible causes (synchronous sizer rerun was not conclusive - some `SizeExpr` trees " + "depend on ndarray contents that are not host-resolvable without a per-launch context, or the " + "task-id slot was empty so the registry pointer could not be confirmed live):\n" + " 1. A tensor backing a data-dependent loop bound was mutated 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 DLPack consumer). The cached " + "adstack capacity was sized against the value before the mutation. Recovery: route the " + "mutation through Quadrants APIs (`Ndarray.write` / `fill` / kernel writes) so the cache " + "invalidates correctly, OR set a generous initial cap if a workload-change milestone " + "genuinely grew capacity. Restart the iteration / training loop from a clean state.\n" + " 2. (Quadrants bug) the pre-pass resolved the alloca to a bound tighter than the actual " + "runtime push count - the enclosing loop shape is outside the current `SizeExpr` grammar, or " + "the Bellman-Ford analyzer undercounted the forward-pass accumulation. Please file with the " + "kernel IR (`QD_DUMP_IR=1`).\n"; + } + AdStackOverflowDiagnosis result; + result.message = identity_block + disambiguation_block + body + + "Note: kernel state may be inconsistent post-overflow; do not retry the same " + "step without addressing the cause and restarting from a clean state."; + // Flag the cache as confirmed-invalid only when the sync rerun positively identified DLPack-bypass (`required + // > allocated` for at least one stack with every leaf resolved against the live snapshot). Unknown is a rare + // fallback now that the snapshot-based evaluator handles ndarray-bound leaves; treating it as + // confirmed-bypass would silently retry against a possibly-broken cache. Quadrants-bug is excluded for the + // same reason - the next launch would re-run the same wrong sizer and produce the same wrong bound. + result.confirmed_invalid_cache = (cause == Cause::DLPackBypass); + return result; +} + +void AdStackCache::capture_diagnose_snapshot(const LaunchContextBuilder &ctx) { + diagnose_snapshot_.data_ptrs.clear(); + diagnose_snapshot_.dev_alloc_types.clear(); + diagnose_snapshot_.shapes.clear(); + // Pull just the data-pointer slot for each arg; the grad-pointer slot is irrelevant to size_expr leaves. + for (const auto &kv : ctx.array_ptrs) { + if (kv.first.ptr_type == TypeFactory::DATA_PTR_POS_IN_NDARRAY) { + diagnose_snapshot_.data_ptrs[kv.first.arg_id] = kv.second; + } + } + diagnose_snapshot_.dev_alloc_types = ctx.device_allocation_type; + // Mirror the per-arg shape vectors `LaunchContextBuilder` populated alongside the args-buffer writes. Going + // through this side map rather than `args_type->get_element_offset` avoids the spurious "Cannot treat as + // TensorType" diagnostics emitted when an axis lookup overruns the actual rank, and keeps the diagnose path + // independent of `args_type` lifetime. + for (const auto &kv : ctx.ndarray_shapes) { + std::vector shape32(kv.second.begin(), kv.second.end()); + diagnose_snapshot_.shapes[kv.first] = std::move(shape32); + } + diagnose_snapshot_.valid = true; +} + +const AdStackCache::DiagnoseLaunchSnapshot *AdStackCache::get_diagnose_snapshot() const { + return diagnose_snapshot_.valid ? &diagnose_snapshot_ : nullptr; } void clip_effective_rows_by_loop_trip_count(std::size_t &effective_rows, @@ -662,7 +1567,8 @@ std::vector encode_bytecode_common(std::vector encode_bytecode_common(std::vector(root_src_idx), contains_device_leaf, free_vars, - var_id_remap, prog, ctx, fl_emitter, nodes, indices); + var_id_remap, prog, ctx, fl_emitter, nodes, indices, reads); if (max_nodes_per_stack > 0) { const std::size_t per_stack = nodes.size() - nodes_before; QD_ERROR_IF(per_stack > static_cast(max_nodes_per_stack), @@ -950,8 +1856,129 @@ std::vector encode_adstack_size_expr_device_bytecode_for_spirv( *out_base_psb = root_psb + static_cast(place_byte_offset); return true; }; - return encode_bytecode_common(std::move(stack_headers), exprs, prog, ctx, fl_emitter, - spirv::kAdStackSizerMaxNodesPerStack); + // Bytecode fast path: replay the recorded host-fold reads against the live state and reuse the cached + // bytecode if every input still matches. The full encode runs only on cache miss. + if (prog != nullptr) { + std::vector cached; + if (prog->adstack_cache().try_spirv_bytecode_cache_hit(prog, static_cast(&ad_stack), ctx, cached)) { + return cached; + } + } + std::vector reads; + std::vector bytecode = encode_bytecode_common(std::move(stack_headers), exprs, prog, ctx, fl_emitter, + spirv::kAdStackSizerMaxNodesPerStack, &reads); + if (prog != nullptr) { + prog->adstack_cache().record_spirv_bytecode_eval(static_cast(&ad_stack), bytecode, std::move(reads)); + } + return bytecode; +} + +void bump_writes_for_kernel_llvm(Program *prog, + LaunchContextBuilder *ctx, + const std::vector &offloaded_tasks) { + if (prog == nullptr) { + return; + } + auto bump_data_ptr = [&](int arg_id) { + ArgArrayPtrKey data_key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto it = ctx->array_ptrs.find(data_key); + if (it != ctx->array_ptrs.end() && it->second != nullptr) { + prog->adstack_cache().bump_ndarray_data_gen(it->second); + } + }; + for (const auto &task : offloaded_tasks) { + for (int snode_id : task.snode_writes) { + prog->adstack_cache().bump_snode_write_gen(snode_id); + } + for (int arg_id : task.arr_writes) { + bump_data_ptr(arg_id); + } + // Read-only `DevAllocType::kNone` args also need a bump: the user's host array is either H2D-blitted to a + // temporary device buffer (CUDA / AMDGPU) or read directly (CPU), and in both cases the data pointer used as + // the cache key is stable across launches, so a content mutation the user performed outside Quadrants's + // tracking is invisible to the metadata cache without an explicit bump. Mirrors the SPIR-V `kone_h2d_blit` + // rule in `bump_writes_for_kernel_spirv`. + for (int arg_id : task.arr_reads) { + auto type_it = ctx->device_allocation_type.find(arg_id); + if (type_it == ctx->device_allocation_type.end() || + type_it->second != LaunchContextBuilder::DevAllocType::kNone) { + continue; + } + bump_data_ptr(arg_id); + } + } +} + +void bump_writes_for_kernel_llvm(Program *prog, + LaunchContextBuilder *ctx, + const std::vector> &snode_writes_per_task, + const std::vector> &arr_writes_per_task, + const std::vector> &arr_reads_per_task) { + if (prog == nullptr) { + return; + } + auto bump_data_ptr = [&](int arg_id) { + ArgArrayPtrKey data_key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto it = ctx->array_ptrs.find(data_key); + if (it != ctx->array_ptrs.end() && it->second != nullptr) { + prog->adstack_cache().bump_ndarray_data_gen(it->second); + } + }; + for (const auto &task_snodes : snode_writes_per_task) { + for (int snode_id : task_snodes) { + prog->adstack_cache().bump_snode_write_gen(snode_id); + } + } + for (const auto &task_args : arr_writes_per_task) { + for (int arg_id : task_args) { + bump_data_ptr(arg_id); + } + } + // Read-only `DevAllocType::kNone` args: see the comment in the CUDA / AMDGPU overload for why CPU LLVM also + // needs the bump. Empty `arr_reads_per_task` is the legal cache-miss path (offline-cache load that did not + // capture per-task arr_reads); skip the loop without raising. + for (const auto &task_args : arr_reads_per_task) { + for (int arg_id : task_args) { + auto type_it = ctx->device_allocation_type.find(arg_id); + if (type_it == ctx->device_allocation_type.end() || + type_it->second != LaunchContextBuilder::DevAllocType::kNone) { + continue; + } + bump_data_ptr(arg_id); + } + } +} + +void bump_writes_for_kernel_spirv( + Program *prog, + LaunchContextBuilder *ctx, + const std::vector &task_attribs, + const std::vector, irpass::ExternalPtrAccess>> &arr_access) { + if (prog == nullptr) { + return; + } + for (const auto &task : task_attribs) { + for (int snode_id : task.snode_writes) { + prog->adstack_cache().bump_snode_write_gen(snode_id); + } + } + for (const auto &kv : arr_access) { + const std::vector &indices = kv.first; + uint32_t access = uint32_t(kv.second); + QD_ASSERT(indices.size() == 1); + int arg_id = indices[0]; + bool kernel_writes = (access & uint32_t(irpass::ExternalPtrAccess::WRITE)) != 0; + bool kone_h2d_blit = (access & uint32_t(irpass::ExternalPtrAccess::READ)) != 0 && + ctx->device_allocation_type[arg_id] == LaunchContextBuilder::DevAllocType::kNone; + if (!kernel_writes && !kone_h2d_blit) { + continue; + } + ArgArrayPtrKey data_key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto it = ctx->array_ptrs.find(data_key); + if (it != ctx->array_ptrs.end()) { + prog->adstack_cache().bump_ndarray_data_gen(it->second); + } + } } } // namespace quadrants::lang diff --git a/quadrants/program/adstack_size_expr_eval.h b/quadrants/program/adstack_size_expr_eval.h index cedbedd1ea..ec8f777203 100644 --- a/quadrants/program/adstack_size_expr_eval.h +++ b/quadrants/program/adstack_size_expr_eval.h @@ -2,17 +2,295 @@ #include #include +#include +#include +#include +#include #include #include "quadrants/codegen/llvm/llvm_compiled_data.h" #include "quadrants/codegen/spirv/kernel_utils.h" #include "quadrants/ir/adstack_size_expr.h" +#include "quadrants/program/program.h" #include "quadrants/transforms/static_adstack_analysis.h" namespace quadrants::lang { -class Program; class LaunchContextBuilder; +class Program; + +// Adstack-specific state owned by `Program` and routed through `program->adstack_cache().method(...)`. Holds two +// orthogonal pieces: +// 1. The per-task adstack-sizer metadata caches (SPIR-V + LLVM-GPU), the encoded SPIR-V bytecode cache, the +// per-launch SizeExpr-eval result cache, and the per-snode / per-DeviceAllocation generation counters that +// drive precise invalidation. +// 2. The adstack-overflow identity registry + diagnostic classifier that the codegen-emitted overflow path +// reads through (`Program::launch_kernel` populates `DiagnoseLaunchSnapshot`; the registry maps task ids +// to kernel + offload-task identities + per-stack capacities, and `diagnose_adstack_overflow` runs the +// synchronous sizer rerun against the captured snapshot to classify the failure mode). +// Both pieces are adstack-internal and lived in `Program` historically; consolidating them here keeps the +// `Program` surface focused on cross-feature program state. +class AdStackCache { + public: + // Back-reference to `Program` is used by the diagnose path to reach `evaluate_adstack_size_expr` / + // `evaluate_adstack_size_expr_for_diagnose` (free functions that take `Program *`) and by the registry methods + // to access `get_compute_device()` for `Device::map`-based ndarray reads. Stored as a raw pointer because + // `AdStackCache` is owned by `Program` and shares its lifetime - the back-ref cannot dangle. + explicit AdStackCache(Program *prog) : prog_(prog) { + } + + // One input read observed during a `evaluate_adstack_size_expr` walk. The cache entry records these so a subsequent + // lookup re-reads the same inputs and compares to `observed_value`; a single mismatch forces a full re-walk. + // `observed_gen` snapshots `snode_write_gen` (FieldLoadObs) or `ndarray_data_gen` (ExternalReadObs) at record + // time. The replay walk uses it as a fast-path short-circuit: if the gen counter has not advanced, the value + // cannot have changed and the dispatch (reader kernel for SNode reads, device-pointer deref for ndarray reads) + // is skipped. ExternalShapeObs reads the args buffer per launch (cheap host memory access), so it does not need + // a gen and leaves this field at 0. + struct SizeExprReadObservation { + enum Kind : uint8_t { FieldLoadObs, ExternalShapeObs, ExternalReadObs }; + Kind kind; + int snode_id; + std::vector indices; + std::vector arg_id_path; + int arg_shape_axis; + int prim_dt; + int64_t observed_value; + uint64_t observed_gen{0}; + void *observed_devalloc{nullptr}; + }; + struct SizeExprCacheEntry { + int64_t result; + std::vector reads; + }; + bool try_size_expr_cache_hit(Program *prog, + const SerializedSizeExpr *expr_key, + LaunchContextBuilder *ctx, + int64_t &out_result); + void record_size_expr_eval(const SerializedSizeExpr *expr_key, + int64_t result, + std::vector reads); + void invalidate_size_expr() { + size_expr_cache_.clear(); + } + + // Cache for encoded SPIR-V adstack-sizer bytecode. Same dep-tracking contract as `try_size_expr_cache_hit` but the + // cached payload is the encoded bytes rather than an integer. + struct SpirvBytecodeCacheEntry { + std::vector bytecode; + std::vector reads; + }; + bool try_spirv_bytecode_cache_hit(Program *prog, + const void *attribs_key, + LaunchContextBuilder *ctx, + std::vector &out_bytecode); + void record_spirv_bytecode_eval(const void *attribs_key, + std::vector bytecode, + std::vector reads); + void invalidate_spirv_bytecode() { + spirv_bytecode_cache_.clear(); + } + + // Per-task adstack metadata output cache for the SPIR-V on-device sizer. + struct PerTaskAdStackCacheEntry { + std::vector metadata; + uint32_t stride_float{0}; + uint32_t stride_int{0}; + std::vector> snode_gens; + std::vector> arg_gens; + }; + bool try_per_task_ad_stack_cache_hit(const void *attribs_key, + LaunchContextBuilder *ctx, + PerTaskAdStackCacheEntry &out); + void record_per_task_ad_stack(const void *attribs_key, + std::vector metadata, + uint32_t stride_float, + uint32_t stride_int, + std::vector> snode_gens, + std::vector> arg_gens); + void invalidate_per_task_ad_stack() { + per_task_ad_stack_cache_.clear(); + } + + // Per-task adstack metadata output cache for the LLVM-GPU on-device sizer (CUDA + AMDGPU). + struct LlvmPerTaskAdStackCacheEntry { + std::vector offsets; + std::vector max_sizes; + uint64_t stride_combined{0}; + uint64_t stride_float{0}; + uint64_t stride_int{0}; + std::vector> snode_gens; + std::vector> arg_gens; + }; + bool try_llvm_per_task_ad_stack_cache_hit(const void *attribs_key, + LaunchContextBuilder *ctx, + LlvmPerTaskAdStackCacheEntry &out); + void record_llvm_per_task_ad_stack(const void *attribs_key, + std::vector offsets, + std::vector max_sizes, + uint64_t stride_combined, + uint64_t stride_float, + uint64_t stride_int, + std::vector> snode_gens, + std::vector> arg_gens); + void invalidate_llvm_per_task_ad_stack() { + llvm_per_task_ad_stack_cache_.clear(); + } + + // Bulk-invalidate just the per-task adstack metadata caches on the overflow raise path. The + // `size_expr_cache_` and `spirv_bytecode_cache_` are intentionally NOT cleared: they self-validate via per-read + // observation walks on the next lookup, so a DLPack-bypass mutation surfaces there as a normal observation + // mismatch and triggers a fresh evaluation without explicit eviction. The per-task metadata caches need a + // force-drop because their gen-counter snapshots match when the user's mutation bypassed our tracking. + // Invalidation is bulk (every task) rather than targeted (just the offender) because a single shared DLPack / + // torch view can back multiple tasks in the same kernel queue: targeted invalidation would let the next launch + // hit a stale entry on a different task that reads the same now-mutated tensor and overflow again. + void invalidate_all_per_task() { + invalidate_per_task_ad_stack(); + invalidate_llvm_per_task_ad_stack(); + } + + uint64_t snode_write_gen(int snode_id) const { + auto it = snode_write_gen_.find(snode_id); + return it == snode_write_gen_.end() ? 0u : it->second; + } + void bump_snode_write_gen(int snode_id) { + ++snode_write_gen_[snode_id]; + } + uint64_t ndarray_data_gen(void *devalloc_ptr) const { + auto it = ndarray_data_gen_.find(devalloc_ptr); + return it == ndarray_data_gen_.end() ? 0u : it->second; + } + 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); + } + + // ----------------------------------------------------------------------------------------------------------- + // Adstack-overflow identity registry + diagnostic classifier + // ----------------------------------------------------------------------------------------------------------- + // Codegen registers each `OffloadedTask::ad_stack` once per kernel compilation and bakes the assigned id as + // an immediate into the lazy-claim overflow path; on overflow the codegen emits `cmpxchg(0, id)` against the + // pinned-host task-id slot. The host raise site reads the slot and routes through + // `diagnose_adstack_overflow_message(id)` to look up the kernel name, task index, and per-stack metadata for + // an enriched error message. Pointer ownership stays with `OffloadedTask`; entries are added but not removed + // - the registry size is bounded by the number of adstack-bearing tasks compiled in the program's lifetime, + // typically dozens. The diagnose path NEVER dereferences `identity_key`; all size-expression data is stored + // inline (`size_exprs`) so the entry is self-contained and immune to lifetime issues from the underlying + // `AdStackSizingInfo` (LLVM) / `AdStackSizingAttribs` (SPIR-V) struct moves. + struct AdStackSizingInfoEntry { + const void *identity_key{nullptr}; + std::string kernel_name; + int task_id_in_kernel{0}; + std::vector allocated_max_sizes; + std::vector size_exprs; + }; + uint32_t register_adstack_sizing_info(const void *identity_key, + const std::string &kernel_name, + int task_id_in_kernel, + std::vector allocated_max_sizes, + std::vector size_exprs); + // Refresh just the `size_exprs` snapshot in an existing registry entry. Used by the LLVM launcher on the first + // launch of a task whose codegen-time registration could not capture size_exprs (the codegen-time + // `current_task->ad_stack` had not yet been finalized). No-op for `id == 0` and ids outside the registry range. + void update_adstack_sizing_info_size_exprs(uint32_t id, std::vector size_exprs); + // Returns a *copy* of the registry entry (not a pointer into the underlying vector) so the caller can safely + // hold the data across operations that might trigger another `register_adstack_sizing_info` and grow / reallocate + // the registry vector (e.g. `evaluate_adstack_size_expr` dispatching a reader kernel that compiles a fresh + // task). Returns `std::nullopt` for the sentinel id `0` and for out-of-range ids. + std::optional lookup_adstack_sizing_info(uint32_t id) const; + // Format a diagnostic message for an overflow signal. `task_id` is the value read from the pinned-host task-id + // slot (0 if no thread overflowed; otherwise the registry id of the first overflowing task). The `message` + // field is embedded into the `QuadrantsAssertionError` raised at the poll site. The `confirmed_invalid_cache` + // field is true only when the synchronous sizer rerun classified the failure as a stale-cache / + // DLPack-bypass case (`required > allocated` for at least one stack with every leaf resolved against the + // captured launch snapshot); the caller (LLVM `check_adstack_overflow` / SPIR-V `GfxRuntime::synchronize`) + // uses it to decide whether to bulk-invalidate the per-task metadata caches so the next launch auto-recovers. + // We deliberately do NOT invalidate on Unknown / Quadrants-bug because invalidating would mask sizer bugs and + // could let a never-confirmed cause silently retry against a possibly-broken cache. + struct AdStackOverflowDiagnosis { + std::string message; + bool confirmed_invalid_cache{false}; + }; + AdStackOverflowDiagnosis diagnose_adstack_overflow(uint32_t task_id) const; + // Convenience wrapper that returns just the message string; production code uses `diagnose_adstack_overflow` + // to also act on the confirmed-cause signal. + std::string diagnose_adstack_overflow_message(uint32_t task_id) const; + + // Snapshot of the most recent launch's context fields needed by `diagnose_adstack_overflow` to resolve + // ndarray-bound `SizeExpr` leaves (`ExternalTensorRead` / `ExternalTensorShape`) at error time, when the + // original `LaunchContextBuilder` is gone. Captured at the top of `Program::launch_kernel` BEFORE the + // launcher rewrites `array_ptrs` (the CPU launcher's `set_host_accessible_ndarray_ptrs` overwrites the + // `DeviceAllocation *` entry with a raw host pointer; capturing earlier keeps the original handle so the + // diagnose path can use the unified `Device::map` API instead of trusting backend-specific semantics). + // + // Design choice (vs. re-dispatching the on-device sizer at diagnose time): `Device::map` is virtual on + // every backend (CPU / CUDA / AMDGPU / Vulkan / Metal), so this snapshot-plus-map approach gets backend + // parity for free without re-entering the launcher's pipeline-setup machinery (compute pipelines / + // descriptor sets / command buffers / sync fences). The diagnose path stays out of the launch lifecycle. + struct DiagnoseLaunchSnapshot { + bool valid{false}; + // arg_id -> ctx->array_ptrs[(arg_id, DATA_PTR_POS_IN_NDARRAY)]. For `kNone` numpy passthrough this is a + // raw host pointer. For `kNdarray` (qd.ndarray) this is a `DeviceAllocation *` handle the diagnose path + // dereferences via `Device::map`. Captured before the CPU launcher's `set_host_accessible_ndarray_ptrs` + // overwrite so the handle is uniform across backends. + std::unordered_map data_ptrs; + std::unordered_map dev_alloc_types; + // Pre-extracted ndarray shapes (`ctx->get_struct_arg_host({arg_id, SHAPE_POS, axis})`) so the + // diagnose evaluator does not need a live `LaunchContextBuilder` to resolve `ExternalTensorShape` or + // multi-axis `ExternalTensorRead` strides. + std::unordered_map> shapes; + }; + // Capture the per-launch fields the diagnose evaluator needs (see `DiagnoseLaunchSnapshot`'s definition for + // the design rationale and field-by-field semantics). Called eagerly from `Program::launch_kernel` only on + // backends where the launch ctx is gone by the time overflow is detected (SPIR-V at `synchronize`); on LLVM + // backends the per-launch overflow poll runs while ctx is still in scope, so we stash the ctx pointer with + // `set_pending_launch_ctx` and let `diagnose_adstack_overflow` capture lazily on the (rare) overflow path. + void capture_diagnose_snapshot(const LaunchContextBuilder &ctx); + // Lazy-capture handoff: `Program::launch_kernel` on LLVM backends sets this to the in-scope ctx before + // forwarding into the launcher and clears it after the per-launch overflow poll returns. If the poll fires, + // `diagnose_adstack_overflow` reads the pointer and captures the snapshot just in time. Stored as a raw + // pointer because it is transient per-launch and never outlives the call frame that set it. + void set_pending_launch_ctx(const LaunchContextBuilder *ctx) { + pending_launch_ctx_ = ctx; + } + // Read-only accessor for the latest snapshot, used by `diagnose_adstack_overflow` to resolve ndarray-bound + // size_expr leaves. Returns `nullptr` when no launch has happened yet (e.g. a freshly constructed `Program` + // hits `synchronize` during teardown without a prior kernel launch). + const DiagnoseLaunchSnapshot *get_diagnose_snapshot() const; + + private: + Program *prog_{nullptr}; + std::unordered_map size_expr_cache_; + std::unordered_map spirv_bytecode_cache_; + std::unordered_map per_task_ad_stack_cache_; + std::unordered_map llvm_per_task_ad_stack_cache_; + std::unordered_map snode_write_gen_; + std::unordered_map ndarray_data_gen_; + + // Adstack-overflow identity registry storage. Index 0 is reserved as the "no overflow" sentinel so the + // codegen-emitted `cmpxchg(0, id)` cleanly distinguishes "task id recorded" from "slot still clean". The + // reverse lookup map (keyed by `identity_key`) keeps `register_adstack_sizing_info` idempotent across + // re-launches of the same kernel. + std::vector adstack_sizing_info_registry_{AdStackSizingInfoEntry{}}; + std::unordered_map adstack_sizing_info_id_by_ptr_; + mutable std::mutex adstack_sizing_info_registry_mutex_; + // Latest captured launch context snapshot for the diagnose path's ndarray-bound leaf resolution. See + // `DiagnoseLaunchSnapshot`'s comment above for why we capture in `Program::launch_kernel` before the launcher + // forwards. + // Single-threaded by construction: `capture_diagnose_snapshot` runs from `Program::launch_kernel` (Python + // launcher thread) and `get_diagnose_snapshot` runs from `diagnose_adstack_overflow` on the same thread; no + // mutex needed. The codegen-time identity registry above keeps its mutex because it is hit from compilation + // worker threads. + DiagnoseLaunchSnapshot diagnose_snapshot_; + // Transient ctx handoff for the lazy LLVM capture path. See `set_pending_launch_ctx`. + const LaunchContextBuilder *pending_launch_ctx_{nullptr}; +}; // Evaluates a compile-time captured `SerializedSizeExpr` against the current field state of `prog` and the // per-launch argument values in `ctx`, returning the concrete adstack capacity for this launch. Scalar i32/i64 @@ -22,6 +300,30 @@ class LaunchContextBuilder; // expression is empty (no symbolic bound captured), signalling to the caller to use the compile-time fallback. int64_t evaluate_adstack_size_expr(const SerializedSizeExpr &expr, Program *prog, LaunchContextBuilder *ctx); +// Diagnose-time variant that evaluates the same `SerializedSizeExpr` against the captured +// `AdStackCache::DiagnoseLaunchSnapshot` rather than a live `LaunchContextBuilder`. Used by +// `AdStackCache::diagnose_adstack_overflow` to resolve `ExternalTensorRead` / `ExternalTensorShape` leaves at +// error time against the live (potentially mutated) ndarray contents, without needing the launch ctx that is +// gone by sync time on async backends. The cross-backend `Device::map(*allocation, &host_ptr)` path is the +// design pivot - see `AdStackCache::DiagnoseLaunchSnapshot`'s comment for the rationale (vs. re-dispatching +// the on-device sizer). Returns -1 if any leaf cannot be resolved (e.g. an arg_id missing from the snapshot, +// or an allocation whose `Device::map` fails); callers fall back to the static dual-cause body in that case. +int64_t evaluate_adstack_size_expr_for_diagnose(const SerializedSizeExpr &expr, Program *prog); + +// RAII guard opening a thread-local read-cache scope. Every nested `evaluate_adstack_size_expr` running inside the +// scope shares one cache, so repeated `(snode_id, indices)` reads share a single reader-kernel dispatch. Place around +// any block that calls `evaluate_adstack_size_expr` more than once back-to-back. +class SizeExprLaunchScope { + public: + SizeExprLaunchScope(); + ~SizeExprLaunchScope(); + SizeExprLaunchScope(const SizeExprLaunchScope &) = delete; + SizeExprLaunchScope &operator=(const SizeExprLaunchScope &) = delete; + + private: + bool owns_; +}; + // Flattens every alloca's `SerializedSizeExpr` tree into the device-readable bytecode defined in // `quadrants/ir/adstack_size_expr_device.h` and returns the raw bytes ready to upload to a device scratch buffer. // Two transforms happen at encoding time: @@ -81,4 +383,35 @@ void clip_effective_rows_by_loop_trip_count(std::size_t &effective_rows, Program *prog, LaunchContextBuilder *ctx); +// Adstack-cache invalidation bump. Called from each backend's kernel launcher BEFORE the per-task +// `publish_adstack_metadata` loop runs, so the per-task metadata cache (`Program::*PerTaskAdStackCacheEntry`) snapshots +// the latest counters at record time and the next lookup detects any drift. Two sources contribute: +// +// - SNode writes: every task in the kernel lists its compile-time `snode_writes` set (computed at codegen via +// `irpass::analysis::gather_snode_read_writes`), bumped per id; covers `SizeExpr::FieldLoad` cache invalidation. +// - ndarray data writes: every arg slot the kernel writes to (`OffloadedTask::arr_writes` on LLVM-GPU, the kernel- +// level `ctx_attribs.arr_access` WRITE bits on SPIR-V) bumps the bound `DeviceAllocation`'s data generation. +// SPIR-V also bumps on the `kNone` READ branch to catch host-driven mutations of raw numpy / torch buffers blitted +// between launches; covers `SizeExpr::ExternalTensorRead` invalidation. +// +// The two helpers share the same Program-level effect; their signatures differ only because the codegen-time write +// sets are stored in different per-backend structs. Forward-only kernels (no adstack tasks) still call these to keep +// counters monotone, which is cheap (one map insert per snode_id at most). +void bump_writes_for_kernel_llvm(Program *prog, + LaunchContextBuilder *ctx, + const std::vector &offloaded_tasks); +// CPU launcher overload: per-task snode_writes / arr_writes / arr_reads are stored as separate parallel vectors on +// the launcher `Context` rather than as `OffloadedTask` clones, for legacy reasons documented in the CPU `Context` +// struct. +void bump_writes_for_kernel_llvm(Program *prog, + LaunchContextBuilder *ctx, + const std::vector> &snode_writes_per_task, + const std::vector> &arr_writes_per_task, + const std::vector> &arr_reads_per_task); +void bump_writes_for_kernel_spirv( + Program *prog, + LaunchContextBuilder *ctx, + const std::vector &task_attribs, + const std::vector, irpass::ExternalPtrAccess>> &arr_access); + } // namespace quadrants::lang diff --git a/quadrants/program/launch_context_builder.cpp b/quadrants/program/launch_context_builder.cpp index aec353eb04..c77d35228d 100644 --- a/quadrants/program/launch_context_builder.cpp +++ b/quadrants/program/launch_context_builder.cpp @@ -39,6 +39,7 @@ void LaunchContextBuilder::copy(const LaunchContextBuilder &other) { array_runtime_sizes = other.array_runtime_sizes; device_allocation_type = other.device_allocation_type; array_ptrs = other.array_ptrs; + ndarray_shapes = other.ndarray_shapes; } void LaunchContextBuilder::set_arg_float(int arg_id, float64 d) { @@ -301,6 +302,8 @@ void LaunchContextBuilder::set_arg_external_array_with_shape(int arg_id, } set_array_runtime_size(arg_id, size); set_array_device_allocation_type(arg_id, DevAllocType::kNone); + std::vector shape_int(shape.begin(), shape.end()); + ndarray_shapes[arg_id] = shape_int; for (int i = 0; i < shape.size(); i++) { set_struct_arg(std::array{arg_id, 0, i}, (int32)shape[i]); } @@ -329,6 +332,7 @@ void LaunchContextBuilder::set_args_ndarray(const std::vector &args_id, con array_ptrs[{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}] = (void *)ptr; set_array_device_allocation_type(arg_id, DevAllocType::kNdarray); const std::vector &shape = arr.shape; + ndarray_shapes[arg_id] = shape; size_t total_size = 1; for (int32 i = 0; i < shape.size(); i++) { set_struct_arg(std::array{arg_id, 0, i}, (int32)shape[i]); @@ -369,6 +373,7 @@ void LaunchContextBuilder::set_arg_ndarray_impl(int arg_id, // Set device allocation type and runtime size set_array_device_allocation_type(arg_id, DevAllocType::kNdarray); QD_ASSERT(shape.size() <= quadrants_max_num_indices); + ndarray_shapes[arg_id] = shape; size_t total_size = 1; for (int i = 0; i < shape.size(); i++) { set_struct_arg(std::array{arg_id, 0, i}, (int32)shape[i]); diff --git a/quadrants/program/launch_context_builder.h b/quadrants/program/launch_context_builder.h index d969637c51..b9385a19cc 100644 --- a/quadrants/program/launch_context_builder.h +++ b/quadrants/program/launch_context_builder.h @@ -173,6 +173,12 @@ class LaunchContextBuilder { std::unordered_map device_allocation_type; std::unordered_map array_ptrs; + // Per-arg ndarray shape, populated by `set_arg_external_array_with_shape` / `set_arg_ndarray*`. Mirrors what + // is already encoded into `arg_buffer_` via `set_struct_arg(std::array{arg_id, 0, axis}, shape[axis])`, but + // exposed as a flat vector here so the diagnose path (`Program::capture_diagnose_snapshot`) does not have to + // walk `args_type` element offsets. The args-type walk would emit spurious "Cannot treat as TensorType" + // diagnostics on out-of-rank axis lookups; mirroring shapes once at set-time is both cheaper and quieter. + std::unordered_map> ndarray_shapes; }; } // namespace quadrants::lang diff --git a/quadrants/program/ndarray.cpp b/quadrants/program/ndarray.cpp index 5d97865b4e..5bc78fd99f 100644 --- a/quadrants/program/ndarray.cpp +++ b/quadrants/program/ndarray.cpp @@ -1,5 +1,6 @@ #include +#include "quadrants/program/adstack_size_expr_eval.h" #include "quadrants/program/ndarray.h" #include "quadrants/program/program.h" #include "fp16.h" @@ -97,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(&ndarray_alloc_)); ndarray_alloc_.device->dealloc_memory(ndarray_alloc_); } } @@ -132,6 +137,10 @@ std::size_t Ndarray::get_nelement() const { } TypedConstant Ndarray::read(const std::vector &I) const { + // Surface any pending adstack overflow at this Quadrants Python entry. The internal `synchronize()` + // below drains the queue but does NOT raise; the explicit poll catches DLPack-bypass overflows from a + // previous launch within one entry of the offending kernel even when the user never calls `qd.sync()`. + prog_->check_adstack_overflow_and_assert(); prog_->synchronize(); size_t index = flatten_index(total_shape_, I); size_t size = data_type_size(get_element_data_type()); @@ -185,6 +194,10 @@ void Ndarray::write(const std::vector &I, TypedConstant val) const { staging_buf_->device->memcpy_internal(this->ndarray_alloc_.get_ptr(index * size_), staging_buf_->get_ptr(), size_); prog_->synchronize(); + // 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(&ndarray_alloc_)); } int64 Ndarray::read_int(const std::vector &i) { diff --git a/quadrants/program/program.cpp b/quadrants/program/program.cpp index 693d66dad0..d994cc2fac 100644 --- a/quadrants/program/program.cpp +++ b/quadrants/program/program.cpp @@ -2,8 +2,14 @@ #include "program.h" +#include + +#include "quadrants/ir/adstack_size_expr.h" +#include "quadrants/program/adstack_size_expr_eval.h" #include "quadrants/ir/statements.h" +#include "quadrants/ir/type_factory.h" #include "quadrants/program/extension.h" +#include "quadrants/program/launch_context_builder.h" #include "quadrants/codegen/cpu/codegen_cpu.h" #include "quadrants/struct/struct.h" #include "quadrants/runtime/program_impls/metal/metal_program.h" @@ -37,12 +43,12 @@ namespace quadrants::lang { std::atomic Program::num_instances_; -Program::Program(Arch desired_arch) : snode_rw_accessors_bank_(this) { +Program::Program(Arch desired_arch) + : snode_rw_accessors_bank_(this), adstack_cache_(std::make_unique(this)) { QD_TRACE("Program initializing..."); - // For performance considerations and correctness of QuantFloatType - // operations, we force floating-point operations to flush to zero on all - // backends (including CPUs). + // For performance considerations and correctness of QuantFloatType operations, we force floating-point operations to + // flush to zero on all backends (including CPUs). #if defined(_M_X64) || defined(__x86_64) _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON); #endif // defined(_M_X64) || defined(__x86_64) @@ -159,11 +165,38 @@ CompileResult Program::compile_kernel(const CompileConfig &compile_config, } void Program::launch_kernel(const CompiledKernelData &compiled_kernel_data, LaunchContextBuilder &ctx) { + // Diagnose-snapshot capture strategy depends on when the overflow check fires relative to ctx lifetime: + // - SPIR-V backends poll the overflow flag at `synchronize()` time, by which point the launch ctx is + // long gone; capture eagerly here, before the CPU launcher's `set_host_accessible_ndarray_ptrs` + // overwrites `array_ptrs[(arg_id, DATA_PTR_POS)]` from the original `DeviceAllocation *` handle. + // - LLVM backends poll the overflow flag in `check_adstack_overflow_and_assert` below, while ctx is + // still in scope. Stash the pointer; the diagnose path captures lazily on the rare overflow case. + const bool defer_to_overflow_path = arch_uses_llvm(compiled_kernel_data.arch()); + if (defer_to_overflow_path) { + adstack_cache_->set_pending_launch_ctx(&ctx); + } else { + adstack_cache_->capture_diagnose_snapshot(ctx); + } num_offloaded_tasks_on_last_call_ = compiled_kernel_data.num_tasks(); program_impl_->get_kernel_launcher().launch_kernel(compiled_kernel_data, ctx); if (compile_config().debug && arch_uses_llvm(compiled_kernel_data.arch())) { program_impl_->check_runtime_error(result_buffer); } + // Free per-launch poll on the pinned-host adstack overflow flag. Catches DLPack-bypass mutations and + // pre-pass undersizing within one Quadrants Python entry of the offending launch, including in async + // release loops that never call `qd.sync()`. SPIR-V backends' poll stays in `synchronize_and_assert()` + // because their overflow buffer needs `wait_idle()` to be coherent. + try { + program_impl_->check_adstack_overflow_and_assert(); + } catch (...) { + if (defer_to_overflow_path) { + adstack_cache_->set_pending_launch_ctx(nullptr); + } + throw; + } + if (defer_to_overflow_path) { + adstack_cache_->set_pending_launch_ctx(nullptr); + } } void Program::materialize_runtime() { @@ -183,15 +216,12 @@ static void remove_rw_accessor_cache(SNode *parent_snode, SNodeRwAccessorsBank * void Program::destroy_snode_tree(SNodeTree *snode_tree) { QD_ASSERT(arch_uses_llvm(compile_config().arch) || compile_config().arch == Arch::vulkan); - // When accessing a ti.field at Python scope, SNodeRwAccessorsBank creates - // a Quadrants Kernel to read/write the field in a JIT manner, which caches - // the compiled JIT Kernel so as to avoid recompilation when accessing the - // same field. + // When accessing a ti.field at Python scope, SNodeRwAccessorsBank creates a Quadrants Kernel to read/write the field + // in a JIT manner, which caches the compiled JIT Kernel so as to avoid recompilation when accessing the same field. - // This cache uses the place-SNode's address (SNode*) as the key, - // which becomes unsafe once the SNodeTree gets destroyed and that - // place-SNode's address gets reused by another SNode. We have to remove all - // cached kernels upon SNodeTree destruction. + // This cache uses the place-SNode's address (SNode*) as the key, which becomes unsafe once the SNodeTree gets + // destroyed and that place-SNode's address gets reused by another SNode. We have to remove all cached kernels upon + // SNodeTree destruction. SNode *root = snode_tree->root(); // Traverse SNodeTree to remove all cached RWAccessor kernels @@ -235,10 +265,8 @@ void Program::synchronize_and_assert() { program_impl_->synchronize_and_assert(); } -void Program::record_size_expr_eval(const SerializedSizeExpr *expr_key, - int64_t result, - std::vector reads) { - size_expr_cache_[expr_key] = SizeExprCacheEntry{result, std::move(reads)}; +void Program::check_adstack_overflow_and_assert() { + program_impl_->check_adstack_overflow_and_assert(); } StreamSemaphore Program::flush() { @@ -356,8 +384,8 @@ void Program::finalize() { } // Notify the backend that teardown has started before the two teardown syncs below. On LLVM this flips - // `LlvmProgramImpl::finalizing_` so `check_adstack_overflow()` short-circuits: otherwise a pending overflow - // flag from a kernel the user never synced explicitly would throw into the Program destructor path. + // `LlvmProgramImpl::finalizing_` so `check_adstack_overflow()` short-circuits: otherwise a pending overflow flag from + // a kernel the user never synced explicitly would throw into the Program destructor path. program_impl_->pre_finalize(); synchronize(); @@ -422,18 +450,13 @@ Ndarray *Program::create_ndarray(const DataType type, } void Program::delete_ndarray(Ndarray *ndarray) { - // [Note] Ndarray memory deallocation - // Ndarray's memory allocation is managed by Quadrants and Python can control - // this via Quadrants indirectly. For example, when an ndarray is GC-ed in - // Python, it signals Quadrants to free its memory allocation. But Quadrants - // will make sure **no pending kernels to be executed needs the ndarray** - // before it actually frees the memory. When `ti.reset()` is called, all - // ndarrays allocated in this program should be gone and no longer valid in - // Python. This isn't the best implementation, ndarrays should be managed by - // quadrants runtime instead of this giant program and it should be freed - // when: - // - Python GC signals quadrants that it's no longer useful - // - All kernels using it are executed. + // [Note] Ndarray memory deallocation Ndarray's memory allocation is managed by Quadrants and Python can control this + // via Quadrants indirectly. For example, when an ndarray is GC-ed in Python, it signals Quadrants to free its memory + // allocation. But Quadrants will make sure **no pending kernels to be executed needs the ndarray** before it actually + // frees the memory. When `ti.reset()` is called, all ndarrays allocated in this program should be gone and no longer + // valid in Python. This isn't the best implementation, ndarrays should be managed by quadrants runtime instead of + // this giant program and it should be freed when: - Python GC signals quadrants that it's no longer useful - All + // kernels using it are executed. if (ndarrays_.count(ndarray) && !program_impl_->used_in_kernel(ndarray->ndarray_alloc_.alloc_id)) { ndarrays_.erase(ndarray); } @@ -451,10 +474,13 @@ intptr_t Program::get_ndarray_data_ptr_as_int(const Ndarray *ndarray) { } 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. + // 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 Program::get_struct_type_with_data_layout(const StructType *old_ty, @@ -467,9 +493,7 @@ Program::~Program() { } DeviceCapabilityConfig translate_devcaps(const std::vector &device_caps) { - // Each device capability assignment is named like this: - // - `spirv_version=1.3` - // - `spirv_has_int8` + // Each device capability assignment is named like this: - `spirv_version=1.3` - `spirv_has_int8` DeviceCapabilityConfig cfg{}; for (const std::string &cap : device_caps) { std::string_view key; diff --git a/quadrants/program/program.h b/quadrants/program/program.h index d8fa88cd04..3d8b7b7425 100644 --- a/quadrants/program/program.h +++ b/quadrants/program/program.h @@ -1,14 +1,18 @@ -// Program - Quadrants program execution context +// Program - Quadrants program execution context #pragma once #include +#include #include #include #include #include +#include +#include #define QD_RUNTIME_HOST +#include "quadrants/ir/adstack_size_expr.h" #include "quadrants/ir/frontend_ir.h" #include "quadrants/ir/ir.h" #include "quadrants/ir/type_factory.h" @@ -29,6 +33,7 @@ namespace quadrants::lang { +class AdStackCache; class StructCompiler; /** @@ -49,15 +54,14 @@ class QD_DLL_EXPORT Program { using Kernel = quadrants::lang::Kernel; uint64 *result_buffer{nullptr}; // Note that this result_buffer is used - // only for runtime JIT functions (e.g. - // `runtime_memory_allocate_aligned`) + // only for runtime JIT functions (e.g. `runtime_memory_allocate_aligned`) std::vector> kernels; std::unique_ptr profiler{nullptr}; - // Note: for now we let all Programs share a single TypeFactory for smooth - // migration. In the future each program should have its own copy. + // Note: for now we let all Programs share a single TypeFactory for smooth migration. In the future each program + // should have its own copy. static TypeFactory &get_type_factory(); Program() : Program(default_compile_config.arch) { @@ -106,6 +110,13 @@ class QD_DLL_EXPORT Program { // Drain the queue and raise on any pending user-visible assert (e.g. adstack overflow). Bound to `qd.sync()`. void synchronize_and_assert(); + // Per-Quadrants-Python-entry poll for any pending adstack overflow signal. Unlike `synchronize_and_assert` + // this does NOT drain the queue: it only reads the pinned-host overflow flag (cheap host atomic load) and + // raises if set. Wired at every host-read entry point (`Ndarray::read`, `SNodeRwAccessorsBank` reads via + // `Program::launch_kernel`'s built-in poll) so a DLPack-bypass overflow surfaces within one entry of the + // offending launch even when the user never calls `qd.sync()`. + void check_adstack_overflow_and_assert(); + StreamSemaphore flush(); /** @@ -187,8 +198,7 @@ class QD_DLL_EXPORT Program { static int default_block_dim(const CompileConfig &config); - // Note this method is specific to LlvmProgramImpl, but we keep it here since - // it's exposed to python. + // Note this method is specific to LlvmProgramImpl, but we keep it here since it's exposed to python. void print_memory_profiler_info(); // Returns zero if the SNode is statically allocated @@ -208,34 +218,16 @@ class QD_DLL_EXPORT Program { // practice. SNode *get_snode_by_id(int snode_id); - // One input read observed during a `evaluate_adstack_size_expr` walk. The cache entry records these so - // a subsequent lookup re-reads the same inputs and compares to `observed_value`; a single mismatch - // forces a full re-walk. - struct SizeExprReadObservation { - enum Kind : uint8_t { FieldLoadObs, ExternalShapeObs, ExternalReadObs }; - Kind kind; - int snode_id; // FieldLoad - std::vector indices; // FieldLoad / ExternalReadObs: resolved indices - std::vector arg_id_path; // External*: arg_id_path - int arg_shape_axis; // ExternalShapeObs - int prim_dt; // ExternalReadObs: PrimitiveTypeID - int64_t observed_value; - }; - struct SizeExprCacheEntry { - int64_t result; - std::vector reads; - }; - // Replay the recorded reads against the live state and `ctx`; on full match write `result` to - // `out_result` and return true. Any mismatch erases the entry and returns false, leaving the caller - // to run a full eval and repopulate the cache via `record_size_expr_eval`. - bool try_size_expr_cache_hit(const SerializedSizeExpr *expr_key, LaunchContextBuilder *ctx, int64_t &out_result); - void record_size_expr_eval(const SerializedSizeExpr *expr_key, - int64_t result, - std::vector reads); - void invalidate_size_expr_cache() { - size_expr_cache_.clear(); + // Adstack-specific caching: per-task adstack-sizer metadata caches (SPIR-V + LLVM-GPU), encoded SPIR-V bytecode + // cache, per-launch SizeExpr-eval result cache, and per-snode / per-DeviceAllocation generation counters that drive + // precise invalidation. Defined in `program/adstack_size_expr_eval.h`. Lifecycle matches `Program`. + AdStackCache &adstack_cache() { + return *adstack_cache_; } + // Adstack-overflow identity registry, diagnostic classifier, and per-launch snapshot all live on + // `AdStackCache`. Callers route through `prog->adstack_cache().method(...)`. + /** * Destroys a new SNode tree. * @@ -351,12 +343,11 @@ class QD_DLL_EXPORT Program { return ndarrays_.size(); } - // TODO(zhanlue): Move these members and corresponding interfaces to - // ProgramImpl Ideally, Program should serve as a pure interface class and all - // the implementations should fall inside ProgramImpl + // TODO(zhanlue): Move these members and corresponding interfaces to ProgramImpl Ideally, Program should serve as a + // pure interface class and all the implementations should fall inside ProgramImpl // - // Once we migrated these implementations to ProgramImpl, lower-level objects - // could store ProgramImpl rather than Program. + // Once we migrated these implementations to ProgramImpl, lower-level objects could store ProgramImpl rather than + // Program. private: CompileConfig compile_config_; @@ -372,8 +363,10 @@ class QD_DLL_EXPORT Program { std::vector> snode_trees_; // Lazy cache for `get_snode_by_id`. Invalidated by `add_snode_tree` and `destroy_snode_tree`. std::unordered_map snode_id_cache_; - // Cache backing `try_size_expr_cache_hit` / `record_size_expr_eval`. - std::unordered_map size_expr_cache_; + // Adstack-specific state (per-task metadata caches, bytecode cache, size-expr results, generation counters, + // identity registry, diagnose-time launch snapshot). All adstack-specific surface lives in + // `program/adstack_size_expr_eval.{h,cpp}`; routed through `adstack_cache()` getter. + std::unique_ptr adstack_cache_; std::stack free_snode_tree_ids_; std::vector> functions_; diff --git a/quadrants/program/program_impl.h b/quadrants/program/program_impl.h index 1289b6cb5a..91fa344462 100644 --- a/quadrants/program/program_impl.h +++ b/quadrants/program/program_impl.h @@ -78,6 +78,16 @@ class ProgramImpl { synchronize(); } + /** + * Per-launch poll for any user-visible async error (currently: adstack overflow). Default no-op. LLVM + * backends override to read the pinned-host overflow flag without sync drain - the cost is one host atomic + * load. SPIR-V's overflow buffer needs `wait_idle()` to be coherent so its check stays in + * `synchronize_and_assert()`; promoting it to per-launch would tank the queue throughput on Apple Metal / + * Vulkan. + */ + virtual void check_adstack_overflow_and_assert() { + } + virtual StreamSemaphore flush() { synchronize(); return nullptr; diff --git a/quadrants/program/snode_rw_accessors_bank.cpp b/quadrants/program/snode_rw_accessors_bank.cpp index 72c236a552..59b136f5da 100644 --- a/quadrants/program/snode_rw_accessors_bank.cpp +++ b/quadrants/program/snode_rw_accessors_bank.cpp @@ -1,5 +1,6 @@ #include "quadrants/program/snode_rw_accessors_bank.h" +#include "quadrants/program/adstack_size_expr_eval.h" #include "quadrants/program/program.h" namespace quadrants::lang { @@ -45,6 +46,11 @@ void SNodeRwAccessorsBank::Accessors::write_float(const std::vector &I, flo prog_->synchronize(); const auto &compiled_kernel_data = get_or_compile(kernels_.writer_compiled, prog_, *writer_); prog_->launch_kernel(compiled_kernel_data, launch_ctx); + // Drives invalidation of the SPIR-V per-task adstack metadata cache: a runtime adstack bound that + // reads this snode (`SizeExpr::FieldLoad`) must see the GPU sizer re-run on the next launch after + // any host-side mutation. Bumped per snode_id so the cache only evicts entries whose `size_expr` + // actually depends on this specific snode. + prog_->adstack_cache().bump_snode_write_gen(snode_->id); } float64 SNodeRwAccessorsBank::Accessors::read_float(const std::vector &I) { @@ -65,6 +71,7 @@ void SNodeRwAccessorsBank::Accessors::write_int(const std::vector &I, int64 prog_->synchronize(); const auto &compiled_kernel_data = get_or_compile(kernels_.writer_compiled, prog_, *writer_); prog_->launch_kernel(compiled_kernel_data, launch_ctx); + prog_->adstack_cache().bump_snode_write_gen(snode_->id); } // for int32 and int64 @@ -75,6 +82,7 @@ void SNodeRwAccessorsBank::Accessors::write_uint(const std::vector &I, uint prog_->synchronize(); const auto &compiled_kernel_data = get_or_compile(kernels_.writer_compiled, prog_, *writer_); prog_->launch_kernel(compiled_kernel_data, launch_ctx); + prog_->adstack_cache().bump_snode_write_gen(snode_->id); } int64 SNodeRwAccessorsBank::Accessors::read_int(const std::vector &I) { diff --git a/quadrants/runtime/amdgpu/kernel_launcher.cpp b/quadrants/runtime/amdgpu/kernel_launcher.cpp index 9aa7265900..fdb633164b 100644 --- a/quadrants/runtime/amdgpu/kernel_launcher.cpp +++ b/quadrants/runtime/amdgpu/kernel_launcher.cpp @@ -1,6 +1,8 @@ #include "quadrants/runtime/amdgpu/kernel_launcher.h" #include "quadrants/rhi/amdgpu/amdgpu_context.h" +#include "quadrants/program/adstack_size_expr_eval.h" #include "quadrants/program/launch_context_builder.h" +#include "quadrants/program/program.h" #include "quadrants/runtime/llvm/llvm_runtime_executor.h" namespace quadrants::lang { @@ -278,6 +280,9 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx AMDGPUDriver::get_instance().memcpy_host_to_device_async(context_pointer, &ctx.get_context(), sizeof(RuntimeContext), nullptr); + // Adstack-cache invalidation bump - see `bump_writes_for_kernel_llvm` in `program/adstack_size_expr_eval.{h,cpp}`. + bump_writes_for_kernel_llvm(executor->get_program(), &ctx, offloaded_tasks); + if (ctx.graph_do_while_arg_id >= 0) { QD_ASSERT(ctx.graph_do_while_flag_dev_ptr); launch_offloaded_tasks_with_do_while(ctx, amdgpu_module, offloaded_tasks, context_pointer, arg_size); diff --git a/quadrants/runtime/cpu/kernel_launcher.cpp b/quadrants/runtime/cpu/kernel_launcher.cpp index ef059fff0d..d2c3a97900 100644 --- a/quadrants/runtime/cpu/kernel_launcher.cpp +++ b/quadrants/runtime/cpu/kernel_launcher.cpp @@ -1,4 +1,6 @@ #include "quadrants/runtime/cpu/kernel_launcher.h" +#include "quadrants/program/adstack_size_expr_eval.h" +#include "quadrants/program/program.h" #include "quadrants/rhi/arch.h" #include "quadrants/runtime/llvm/llvm_runtime_executor.h" @@ -25,6 +27,8 @@ void KernelLauncher::launch_offloaded_tasks(LaunchContextBuilder &ctx, // reducer below tightens specific slots. executor->publish_adstack_lazy_claim_buffers(task_funcs.size()); } + // Span every task's `publish_adstack_metadata` call below with one shared read cache. + SizeExprLaunchScope launch_scope; for (size_t i = 0; i < task_funcs.size(); ++i) { if (!ad_stacks[i].allocas.empty()) { executor->publish_adstack_metadata(ad_stacks[i], num_threads_per_task[i], &ctx); @@ -98,36 +102,33 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx auto *executor = get_runtime_executor(); ctx.get_context().runtime = executor->get_llvm_runtime(); - // For quadrants ndarrays, context.array_ptrs saves pointer to its - // |DeviceAllocation|, CPU backend actually want to use the raw ptr here. - const auto ¶meters = *launcher_ctx.parameters; - for (int i = 0; i < (int)parameters.size(); i++) { - const auto &kv = parameters[i]; - const auto &arg_id = kv.first; - const auto ¶meter = kv.second; - if (parameter.is_array) { - void *data_ptr = ctx.array_ptrs[{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}]; - void *grad_ptr = ctx.array_ptrs[{arg_id, TypeFactory::GRAD_PTR_POS_IN_NDARRAY}]; - - if (ctx.device_allocation_type[arg_id] == LaunchContextBuilder::DevAllocType::kNone) { - ctx.set_host_accessible_ndarray_ptrs(arg_id, (uint64)data_ptr, (uint64)grad_ptr); - if (arg_id == ctx.graph_do_while_arg_id) { - ctx.graph_do_while_flag_dev_ptr = data_ptr; - } - } else if (ctx.array_runtime_sizes[arg_id] > 0) { - uint64 host_ptr = (uint64)executor->get_device_alloc_info_ptr(*static_cast(data_ptr)); - ctx.set_array_device_allocation_type(arg_id, LaunchContextBuilder::DevAllocType::kNone); - uint64 host_ptr_grad = - grad_ptr == nullptr - ? 0 - : (uint64)executor->get_device_alloc_info_ptr(*static_cast(grad_ptr)); - ctx.set_host_accessible_ndarray_ptrs(arg_id, host_ptr, host_ptr_grad); - if (arg_id == ctx.graph_do_while_arg_id) { - ctx.graph_do_while_flag_dev_ptr = (void *)host_ptr; - } + // For quadrants ndarrays, context.array_ptrs saves pointer to its |DeviceAllocation|; the CPU backend wants the raw + // ptr here. Iterate only the precomputed array-typed `arg_id`s. + for (int arg_id : launcher_ctx.array_arg_ids) { + void *data_ptr = ctx.array_ptrs[{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}]; + void *grad_ptr = ctx.array_ptrs[{arg_id, TypeFactory::GRAD_PTR_POS_IN_NDARRAY}]; + + if (ctx.device_allocation_type[arg_id] == LaunchContextBuilder::DevAllocType::kNone) { + ctx.set_host_accessible_ndarray_ptrs(arg_id, (uint64)data_ptr, (uint64)grad_ptr); + if (arg_id == ctx.graph_do_while_arg_id) { + ctx.graph_do_while_flag_dev_ptr = data_ptr; + } + } else if (ctx.array_runtime_sizes[arg_id] > 0) { + uint64 host_ptr = (uint64)executor->get_device_alloc_info_ptr(*static_cast(data_ptr)); + ctx.set_array_device_allocation_type(arg_id, LaunchContextBuilder::DevAllocType::kNone); + uint64 host_ptr_grad = + grad_ptr == nullptr ? 0 + : (uint64)executor->get_device_alloc_info_ptr(*static_cast(grad_ptr)); + ctx.set_host_accessible_ndarray_ptrs(arg_id, host_ptr, host_ptr_grad); + if (arg_id == ctx.graph_do_while_arg_id) { + ctx.graph_do_while_flag_dev_ptr = (void *)host_ptr; } } } + // Adstack-cache invalidation bump - see `bump_writes_for_kernel_llvm` in `program/adstack_size_expr_eval.{h,cpp}`. + bump_writes_for_kernel_llvm(executor->get_program(), &ctx, launcher_ctx.snode_writes_per_task, + launcher_ctx.arr_writes_per_task, launcher_ctx.arr_reads_per_task); + if (ctx.graph_do_while_arg_id >= 0) { QD_ASSERT(ctx.graph_do_while_flag_dev_ptr); launch_offloaded_tasks_with_do_while(ctx, launcher_ctx.task_funcs, launcher_ctx.ad_stacks, @@ -154,9 +155,15 @@ KernelLauncher::Handle KernelLauncher::register_llvm_kernel(const LLVM::Compiled std::vector task_funcs; std::vector ad_stacks; std::vector num_threads_per_task; + std::vector> snode_writes_per_task; + std::vector> arr_writes_per_task; + std::vector> arr_reads_per_task; task_funcs.reserve(data.tasks.size()); ad_stacks.reserve(data.tasks.size()); num_threads_per_task.reserve(data.tasks.size()); + snode_writes_per_task.reserve(data.tasks.size()); + arr_writes_per_task.reserve(data.tasks.size()); + arr_reads_per_task.reserve(data.tasks.size()); for (auto &task : data.tasks) { auto *func_ptr = jit_module->lookup_function(task.name); QD_ASSERT_INFO(func_ptr, "Offloaded datum function {} not found", task.name); @@ -167,6 +174,9 @@ KernelLauncher::Handle KernelLauncher::register_llvm_kernel(const LLVM::Compiled // time. ad_stacks.push_back(task.ad_stack); num_threads_per_task.push_back(task.ad_stack.static_num_threads); + snode_writes_per_task.push_back(task.snode_writes); + arr_writes_per_task.push_back(task.arr_writes); + arr_reads_per_task.push_back(task.arr_reads); } // Populate ctx @@ -174,6 +184,18 @@ KernelLauncher::Handle KernelLauncher::register_llvm_kernel(const LLVM::Compiled ctx.task_funcs = std::move(task_funcs); ctx.ad_stacks = std::move(ad_stacks); ctx.num_threads_per_task = std::move(num_threads_per_task); + ctx.snode_writes_per_task = std::move(snode_writes_per_task); + ctx.arr_writes_per_task = std::move(arr_writes_per_task); + ctx.arr_reads_per_task = std::move(arr_reads_per_task); + + // Precompute the array-typed parameter `arg_id`s so `launch_llvm_kernel` does not have to walk the + // full parameters list and re-check `is_array` on every invocation. + ctx.array_arg_ids.clear(); + for (const auto &kv : *ctx.parameters) { + if (kv.second.is_array) { + ctx.array_arg_ids.push_back(kv.first); + } + } compiled.set_handle(handle); } diff --git a/quadrants/runtime/cpu/kernel_launcher.h b/quadrants/runtime/cpu/kernel_launcher.h index c14452c62c..059d491963 100644 --- a/quadrants/runtime/cpu/kernel_launcher.h +++ b/quadrants/runtime/cpu/kernel_launcher.h @@ -21,7 +21,17 @@ class KernelLauncher : public LLVM::KernelLauncher { // serial), so no launch-time gtmp resolution is needed on this backend. std::vector ad_stacks; std::vector num_threads_per_task; + // Per-task snode-write set / arg-write set / arg-read set, copied off `OffloadedTask::snode_writes` / + // `arr_writes` / `arr_reads` at register time. Used by `launch_llvm_kernel` to bump + // `Program::snode_write_gen_` / `ndarray_data_gen_` before each launch so the per-task adstack metadata cache + // invalidates when this kernel mutates a SNode / ndarray a downstream `size_expr` reads, plus the read-only + // `kNone` host-array case where the data pointer stays stable across launches but the user's content can change. + std::vector> snode_writes_per_task; + std::vector> arr_writes_per_task; + std::vector> arr_reads_per_task; const std::vector> *parameters; + // arg_ids of the array-typed entries in `parameters`, precomputed at register time. + std::vector array_arg_ids; }; public: diff --git a/quadrants/runtime/cuda/kernel_launcher.cpp b/quadrants/runtime/cuda/kernel_launcher.cpp index 412f244051..b695e8180b 100644 --- a/quadrants/runtime/cuda/kernel_launcher.cpp +++ b/quadrants/runtime/cuda/kernel_launcher.cpp @@ -2,6 +2,8 @@ #include "quadrants/runtime/cuda/cuda_utils.h" #include "quadrants/rhi/cuda/cuda_context.h" #include "quadrants/runtime/llvm/llvm_runtime_executor.h" +#include "quadrants/program/adstack_size_expr_eval.h" +#include "quadrants/program/program.h" #include @@ -347,6 +349,9 @@ void KernelLauncher::launch_llvm_kernel(Handle handle, LaunchContextBuilder &ctx sizeof(RuntimeContext), nullptr); } + // Adstack-cache invalidation bump - see `bump_writes_for_kernel_llvm` in `program/adstack_size_expr_eval.{h,cpp}`. + bump_writes_for_kernel_llvm(executor->get_program(), &ctx, offloaded_tasks); + if (ctx.graph_do_while_arg_id >= 0) { QD_ASSERT(ctx.graph_do_while_flag_dev_ptr); launch_offloaded_tasks_with_do_while(ctx, cuda_module, offloaded_tasks, device_context_ptr); diff --git a/quadrants/runtime/gfx/adstack_sizer_launch.cpp b/quadrants/runtime/gfx/adstack_sizer_launch.cpp index c9bb4fd0d8..ba88c8bca4 100644 --- a/quadrants/runtime/gfx/adstack_sizer_launch.cpp +++ b/quadrants/runtime/gfx/adstack_sizer_launch.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include "quadrants/codegen/spirv/adstack_sizer_shader.h" @@ -36,6 +37,34 @@ namespace gfx { namespace { +// Walk a `SerializedSizeExpr` tree structurally and append every leaf the GPU sizer will read at +// launch time into the dependency sets: `snode_ids` for `FieldLoad` leaves, `arg_ids` for +// `ExternalTensorShape` and `ExternalTensorRead` leaves. The walk reads no live values (no SNode +// access, no buffer dereference, no nested kernel launch) - it is pure tree inspection - so it is +// safe to call from anywhere in the launcher's pre-publish phase. Used by the per-task metadata cache +// to know which generation counters to snapshot at record time and re-check at lookup time. +void collect_size_expr_dep_keys(const SerializedSizeExpr &expr, + std::unordered_set &snode_ids, + std::unordered_set &arg_ids) { + for (const auto &node : expr.nodes) { + switch (static_cast(node.kind)) { + case SizeExpr::Kind::FieldLoad: + if (node.snode_id >= 0) { + snode_ids.insert(node.snode_id); + } + break; + case SizeExpr::Kind::ExternalTensorShape: + case SizeExpr::Kind::ExternalTensorRead: + if (!node.arg_id_path.empty()) { + arg_ids.insert(node.arg_id_path.front()); + } + break; + default: + break; + } + } +} + // True iff every SizeExpr across every adstack alloca in every task is host-resolvable WITHOUT triggering a // nested kernel launch. On SPIR-V backends both `ExternalTensorRead` and `FieldLoad` would require a device // round-trip from the host: ExternalTensorRead dereferences ndarray data that lives in GPU-private memory, @@ -76,6 +105,8 @@ void eval_per_task_metadata_on_host(const std::vector &adstack_task_indi LaunchContextBuilder &host_ctx, std::vector &per_task_ad_stack) { using HeapKind = spirv::TaskAttributes::AdStackAllocaAttribs::HeapKind; + // Span the per-task `evaluate_adstack_size_expr` calls below with one shared read cache. + SizeExprLaunchScope launch_scope; for (size_t ti : adstack_task_indices) { const auto &allocas = task_attribs[ti].ad_stack.allocas; auto &rt = per_task_ad_stack[ti]; @@ -147,6 +178,76 @@ std::vector GfxRuntime::publish_adstack_metadata_spirv( "encode AdStack SizeExpr bytecode. Ensure GfxProgramImpl passes `program_impl = this` " "into `GfxRuntime::Params`."); + // Register each adstack-bearing task with the Program-side identity registry so the host raise site + // can name the offending kernel + task in its diagnostic message. Idempotent: re-registration of the + // same `&task_attribs[ti].ad_stack` returns the same id and just refreshes the metadata. The + // `task_attribs` vector lives inside the cached compiled kernel handle, so the address is stable + // across launches. The id is written into `task_attribs[ti].ad_stack.registry_id` so the + // synchronous sizer rerun in `Program::diagnose_adstack_overflow_message` can find the live + // pointer; `set_adstack_sizing_info_pointer` flips the pointer-live sentinel that gates the deref. + // The `task_attribs` parameter is passed `const` here so we cast away constness for the in-place + // id stash; the cached compiled kernel owns the storage and outlives the launch. + for (size_t k = 0; k < adstack_task_indices.size(); ++k) { + size_t ti = adstack_task_indices[k]; + auto &mutable_attribs = + const_cast(task_attribs[ti].ad_stack); + // Skip re-registration once the registry has the entry. `task_attribs` lives inside the cached compiled-kernel + // handle and its allocas / size_exprs / max_size_compile_time are codegen-time immutable; a fresh compile would + // use a new `mutable_attribs` address with `registry_id == 0`, so this is the natural one-shot guard. + // Re-registering on every launch was costing ~2% of wallclock on the Metal rigid-step bench (vector copies for + // allocated_max_sizes + size_exprs, plus a deep copy of the kernel_name string into the entry). + if (mutable_attribs.registry_id != 0) { + continue; + } + std::vector allocated_max_sizes; + std::vector size_exprs; + allocated_max_sizes.reserve(mutable_attribs.allocas.size()); + size_exprs.reserve(mutable_attribs.allocas.size()); + for (const auto &a : mutable_attribs.allocas) { + allocated_max_sizes.push_back(static_cast(a.max_size_compile_time)); + size_exprs.push_back(a.size_expr); + } + uint32_t id = program_impl_->program->adstack_cache().register_adstack_sizing_info( + static_cast(&mutable_attribs), kernel_name, static_cast(ti), std::move(allocated_max_sizes), + std::move(size_exprs)); + mutable_attribs.registry_id = id; + } + + // Populate the per-task `BufferType::AdStackTaskRegistryId` buffer with the registry id assigned in + // the loop above. The codegen task-end overflow check reads slot `task_id_in_kernel_` and + // `OpAtomicCompareExchange`'s it into `AdStackOverflow[1]` when the latter is still 0 (the FIRST + // overflowing task across the dispatch records its identity). Slots for forward-only tasks default + // to 0; the codegen short-circuits the cmpxchg when the loaded id is 0 anyway, but the explicit fill + // keeps the buffer deterministic across launches and survives the offline-cache reload path. + // Allocation policy mirrors `adstack_bound_row_capacity_buffer_` in + // `adstack_bound_reducer_launch.cpp`: `host_write=true` SSBO, grow on amortised doubling, displaced + // buffer parked in `ctx_buffers_` for in-flight cmdlist safety. + { + const size_t needed_bytes = std::max(task_attribs.size(), 1) * sizeof(uint32_t); + if (!adstack_task_registry_id_buffer_ || adstack_task_registry_id_buffer_size_ < needed_bytes) { + size_t new_size = std::max(needed_bytes, 2 * adstack_task_registry_id_buffer_size_); + auto [buf, res] = device_->allocate_memory_unique({new_size, + /*host_write=*/true, + /*host_read=*/false, + /*export_sharing=*/false, AllocUsage::Storage}); + QD_ASSERT_INFO(res == RhiResult::success, "Failed to allocate adstack task registry id buffer (size={})", + new_size); + if (adstack_task_registry_id_buffer_) { + ctx_buffers_.push_back(std::move(adstack_task_registry_id_buffer_)); + } + adstack_task_registry_id_buffer_ = std::move(buf); + adstack_task_registry_id_buffer_size_ = new_size; + } + void *mapped = nullptr; + RhiResult map_res = device_->map_range(adstack_task_registry_id_buffer_->get_ptr(0), needed_bytes, &mapped); + QD_ASSERT_INFO(map_res == RhiResult::success, "Failed to map adstack task registry id buffer"); + uint32_t *slots = reinterpret_cast(mapped); + for (size_t ti = 0; ti < task_attribs.size(); ++ti) { + slots[ti] = task_attribs[ti].ad_stack.registry_id; + } + device_->unmap(*adstack_task_registry_id_buffer_); + } + // Fast path: when no SizeExpr in any adstack-bearing task contains an `ExternalTensorRead` leaf, every // capacity bound is host-resolvable through `evaluate_adstack_size_expr`, and the entire GPU sizer pipeline // (sizer-bytecode upload, per-task metadata-buffer alloc, `flush()` + `device_->wait_idle()` to force PSB @@ -161,6 +262,47 @@ std::vector GfxRuntime::publish_adstack_metadata_spirv( per_task_ad_stack); return per_task_ad_stack; } + + // Per-task metadata cache fast path. Each adstack-bearing task is keyed by the stable address of its + // `AdStackSizingAttribs` struct (lifetime of the compiled kernel); cached payload is the metadata + // bytes the sizer wrote back plus per-source generation snapshots tagged at record time. Hit means + // the entire sizer pipeline (`flush + wait_idle`, per-task metadata buffer alloc, cmdlist record, + // `submit_synced`, readback) drops out, which on Metal is the dominant per-launch cost. Partial hits + // fall through to a full pipeline run rather than threading a "skip task k" path through the cmdlist + // record loop; mixed hit / miss only happens immediately after an invalidation so the simpler code + // path is the right tradeoff. Soundness comes from `Program::try_per_task_ad_stack_cache_hit` + // re-checking every counter the recorded entry tagged: per-snode `snode_write_gen_` covers + // `FieldLoad` reads, per-DeviceAllocation `ndarray_data_gen_` covers `ExternalTensorRead` reads, and + // the cached `arg_id -> devalloc` map catches a different tensor at the same arg slot + // (`ExternalTensorShape` invariance is per-tensor, not across tensors). + { + bool all_hit = true; + for (size_t k = 0; k < adstack_task_indices.size() && all_hit; ++k) { + size_t ti = adstack_task_indices[k]; + AdStackCache::PerTaskAdStackCacheEntry entry; + if (program_impl_->program->adstack_cache().try_per_task_ad_stack_cache_hit( + static_cast(&task_attribs[ti].ad_stack), &host_ctx, entry)) { + auto &rt = per_task_ad_stack[ti]; + rt.metadata = std::move(entry.metadata); + rt.stride_float = entry.stride_float; + rt.stride_int = entry.stride_int; + } else { + all_hit = false; + } + } + if (all_hit) { + return per_task_ad_stack; + } + // Reset per-task strides clobbered by partial hits above; the GPU sizer pipeline below repopulates + // them from scratch. + for (size_t k = 0; k < adstack_task_indices.size(); ++k) { + size_t ti = adstack_task_indices[k]; + per_task_ad_stack[ti].metadata.clear(); + per_task_ad_stack[ti].stride_float = task_attribs[ti].ad_stack.per_thread_stride_float_compile_time; + per_task_ad_stack[ti].stride_int = task_attribs[ti].ad_stack.per_thread_stride_int_compile_time; + } + } + QD_ERROR_IF(!device_->get_caps().get(DeviceCapability::spirv_has_physical_storage_buffer) || !device_->get_caps().get(DeviceCapability::spirv_has_int64) || !device_->get_caps().get(DeviceCapability::spirv_has_int8) || @@ -219,6 +361,8 @@ std::vector GfxRuntime::publish_adstack_metadata_spirv( std::vector per_task_bytecode_offsets(adstack_task_indices.size()); std::vector per_task_metadata_bytes(adstack_task_indices.size()); size_t total_bytecode_bytes = 0; + // Span the per-task bytecode encoding below with one shared read cache. + SizeExprLaunchScope launch_scope; for (size_t k = 0; k < adstack_task_indices.size(); ++k) { size_t ti = adstack_task_indices[k]; per_task_bytecodes[k] = encode_adstack_size_expr_device_bytecode_for_spirv(task_attribs[ti].ad_stack, @@ -381,6 +525,37 @@ std::vector GfxRuntime::publish_adstack_metadata_spirv( ctx_buffers_.push_back(std::move(per_task_metadata_allocs[k])); } + // Record cache entries. Per task we walk every alloca's `size_expr` to build the dependency set + // (snode_ids referenced by `FieldLoad`, arg_ids referenced by `ExternalTensorShape` / + // `ExternalTensorRead`), snapshot the corresponding generation counters and the bound + // DeviceAllocation for each arg_id, and stash everything alongside the metadata bytes. The + // structural walk reads no live values, so it never re-enters `launch_kernel`. + for (size_t k = 0; k < adstack_task_indices.size(); ++k) { + size_t ti = adstack_task_indices[k]; + auto &rt = per_task_ad_stack[ti]; + std::unordered_set snode_ids; + std::unordered_set arg_ids; + for (const auto &alloca : task_attribs[ti].ad_stack.allocas) { + collect_size_expr_dep_keys(alloca.size_expr, snode_ids, arg_ids); + } + std::vector> snode_gens; + snode_gens.reserve(snode_ids.size()); + for (int snode_id : snode_ids) { + snode_gens.emplace_back(snode_id, program_impl_->program->adstack_cache().snode_write_gen(snode_id)); + } + std::vector> arg_gens; + arg_gens.reserve(arg_ids.size()); + for (int arg_id : arg_ids) { + ArgArrayPtrKey data_key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto ap_it = host_ctx.array_ptrs.find(data_key); + void *devalloc = (ap_it == host_ctx.array_ptrs.end()) ? nullptr : ap_it->second; + arg_gens.emplace_back(arg_id, devalloc, program_impl_->program->adstack_cache().ndarray_data_gen(devalloc)); + } + program_impl_->program->adstack_cache().record_per_task_ad_stack( + static_cast(&task_attribs[ti].ad_stack), rt.metadata, rt.stride_float, rt.stride_int, + std::move(snode_gens), std::move(arg_gens)); + } + return per_task_ad_stack; } diff --git a/quadrants/runtime/gfx/runtime.cpp b/quadrants/runtime/gfx/runtime.cpp index 0bf08c2c42..6d9f09b303 100644 --- a/quadrants/runtime/gfx/runtime.cpp +++ b/quadrants/runtime/gfx/runtime.cpp @@ -9,6 +9,7 @@ #include "quadrants/program/program.h" #include "quadrants/program/launch_context_builder.h" #include "quadrants/ir/type_factory.h" +#include "quadrants/common/exceptions.h" #include "quadrants/common/filesystem.hpp" #include @@ -533,6 +534,12 @@ void GfxRuntime::launch_kernel(KernelHandle handle, LaunchContextBuilder &host_c // Record commands const auto &task_attribs = ti_kernel->ti_kernel_attribs().tasks_attribs; + // Adstack-cache invalidation bump - see `bump_writes_for_kernel_spirv` in `program/adstack_size_expr_eval.{h,cpp}`. + if (program_impl_ != nullptr) { + bump_writes_for_kernel_spirv(program_impl_->program, &host_ctx, task_attribs, + ti_kernel->ti_kernel_attribs().ctx_attribs.arr_access); + } + // Device-side adstack SizeExpr evaluation: every task with adstack allocas has its per-alloca `max_size` / // `offset` metadata resolved by a dedicated compute shader (see `quadrants/runtime/gfx/adstack_sizer_launch.cpp` // for the full mechanism). The helper internally early-returns (after seeding the per-task vector with @@ -618,12 +625,17 @@ void GfxRuntime::launch_kernel(KernelHandle handle, LaunchContextBuilder &host_c auto it = src.find(bind.buffer.root_id); bindings->rw_buffer(bind.binding, it != src.end() ? it->second : kDeviceNullAllocation); } else if (bind.buffer.type == BufferType::AdStackOverflow) { - // SPIR-V codegen writes a non-zero sentinel into this single-u32 buffer whenever an AdStackPushStmt hits the - // overflow branch. Allocate it lazily on first use and reuse across launches; synchronize() reads it, raises on - // non-zero, and zeros it for the next window. + // Two-slot u32 buffer: [overflow_signal, task_registry_id]. The SPIR-V codegen task-end emit + // writes a non-zero `stack_id + 1` into slot 0 via OpAtomicUMax whenever any push site overflowed + // its `max_size`, and `cmpxchg(0, registry_id)` into slot 1 to record the offending + // `AdStackSizingInfo`'s `Program::adstack_sizing_info_registry_` id (only the FIRST overflowing + // task's id sticks; subsequent threads' cmpxchg fails harmlessly). Host `synchronize()` reads + // both slots, raises with `Program::diagnose_adstack_overflow_message(task_id)`, and zeros both + // for the next window. if (!adstack_overflow_buffer_) { - auto [buf, res] = device_->allocate_memory_unique({sizeof(uint32_t), /*host_write=*/true, /*host_read=*/true, - /*export_sharing=*/false, AllocUsage::Storage}); + auto [buf, res] = + device_->allocate_memory_unique({2 * sizeof(uint32_t), /*host_write=*/true, /*host_read=*/true, + /*export_sharing=*/false, AllocUsage::Storage}); QD_ASSERT_INFO(res == RhiResult::success, "Failed to allocate adstack overflow buffer"); adstack_overflow_buffer_ = std::move(buf); current_cmdlist_->buffer_fill(adstack_overflow_buffer_->get_ptr(0), kBufferSizeEntireSize, /*data=*/0); @@ -669,6 +681,21 @@ void GfxRuntime::launch_kernel(KernelHandle handle, LaunchContextBuilder &host_c } else { bindings->rw_buffer(bind.binding, kDeviceNullAllocation); } + } else if (bind.buffer.type == BufferType::AdStackTaskRegistryId) { + // Per-task `Program::adstack_sizing_info_registry_` ids written by the SPIR-V launcher in + // `publish_adstack_metadata_spirv` immediately after registering each adstack-bearing task; slot + // `ti` holds the registry id for that task (0 for tasks without adstacks). The codegen-emitted + // task-end overflow check reads slot `task_id_in_kernel_` and `OpAtomicCompareExchange`'s the + // value into `AdStackOverflow[1]` when the latter is still 0 - recording the FIRST overflowing + // task's identity. The codegen-side gate on `task_has_adstack_push_` ensures forward-only tasks + // never request this binding, so the buffer is always allocated by `publish_adstack_metadata_ + // spirv` whenever it is requested. The defensive null bind keeps the assertion path intact if + // a future codegen path requests the binding without the launcher having allocated. + if (adstack_task_registry_id_buffer_) { + bindings->rw_buffer(bind.binding, *adstack_task_registry_id_buffer_); + } else { + bindings->rw_buffer(bind.binding, kDeviceNullAllocation); + } } else if (bind.buffer.type == BufferType::AdStackHeapFloat) { // SPIR-V adstack primal/adjoint storage for f32 adstacks. Sized for `effective_rows`: the count of threads the // static-IR-bound reducer pre-counted as passing the captured gate, when the task has a captured `bound_expr` @@ -981,11 +1008,15 @@ void GfxRuntime::synchronize() { // against pending GPU writes. if (adstack_overflow_buffer_ && !finalizing_) { uint32_t flag_val = 0; + uint32_t task_id_val = 0; void *mapped = nullptr; QD_ASSERT(device_->map(*adstack_overflow_buffer_, &mapped) == RhiResult::success); - flag_val = *reinterpret_cast(mapped); + auto *slots = reinterpret_cast(mapped); + flag_val = slots[0]; + task_id_val = slots[1]; if (flag_val != 0) { - *reinterpret_cast(mapped) = 0; + slots[0] = 0; + slots[1] = 0; } device_->unmap(*adstack_overflow_buffer_); // UINT32_MAX is the dedicated sentinel the codegen-emitted defense-in-depth bounds check at the float Lowest @@ -1001,14 +1032,24 @@ void GfxRuntime::synchronize() { "LCA-block claim count. The bound is supposed to be exact by construction; reaching this signal " "means the reducer and the main pass observed different threads passing the captured gating " "predicate. File a bug with the kernel IR via `QD_DUMP_IR=1` and a minimal repro."); - QD_ERROR_IF(flag_val != 0, - "Adstack overflow (offending stack_id={}): a reverse-mode autodiff kernel pushed more elements " - "than the adstack capacity allows. Raised at the next qd.sync() rather than at the offending " - "kernel launch. The pre-pass resolved this alloca to a bound tighter than the actual runtime " - "push count - either the enclosing loop shape is outside the current `SizeExpr` grammar " - "(rewrite it, or extend the grammar), or the Bellman-Ford analyzer undercounted the " - "forward-pass accumulation on this stack (file a bug with the kernel IR via `QD_DUMP_IR=1`).", - flag_val - 1); + if (flag_val != 0) { + Program *prog = (program_impl_ != nullptr) ? program_impl_->program : nullptr; + std::string diagnostic; + if (prog != nullptr) { + auto diag = prog->adstack_cache().diagnose_adstack_overflow(task_id_val); + diagnostic = std::move(diag.message); + // See `LlvmRuntimeExecutor::check_adstack_overflow` for the rationale; only invalidate when the sizer rerun + // confirmed a stale cache (DLPack-bypass) so a Quadrants pre-pass bug is not silently masked. + if (diag.confirmed_invalid_cache) { + prog->adstack_cache().invalidate_all_per_task(); + } + } + throw QuadrantsAssertionError( + fmt::format("Adstack overflow: a reverse-mode autodiff kernel pushed more elements than the adstack " + "capacity allows. Raised at the next Quadrants Python entry rather than at the offending " + "kernel launch. Offending adstack index within the task: {}.\n{}", + flag_val - 1, diagnostic)); + } } fflush(stdout); } diff --git a/quadrants/runtime/gfx/runtime.h b/quadrants/runtime/gfx/runtime.h index 0733306034..dd804fa842 100644 --- a/quadrants/runtime/gfx/runtime.h +++ b/quadrants/runtime/gfx/runtime.h @@ -279,6 +279,16 @@ class QD_DLL_EXPORT GfxRuntime { std::unique_ptr adstack_bound_row_capacity_buffer_; size_t adstack_bound_row_capacity_buffer_size_{0}; + // Per-kernel `BufferType::AdStackTaskRegistryId` (`uint[num_tasks_in_kernel]`). Written by + // `publish_adstack_metadata_spirv` immediately after registering each adstack-bearing task with the + // Program-side identity registry: slot `ti` holds that task's registry id (0 for tasks without + // adstacks). The codegen task-end overflow check reads `slot[task_id_in_kernel_]` and + // `OpAtomicCompareExchange`'s it into `AdStackOverflow[1]` on overflow so the host raise site can + // name the offending kernel + task. Allocated and grown lazily on demand following the same + // pattern as `adstack_bound_row_capacity_buffer_`. + std::unique_ptr adstack_task_registry_id_buffer_; + size_t adstack_task_registry_id_buffer_size_{0}; + // Owning `ProgramImpl` back-reference; propagated from `Params::program_impl`. See the comment on // `Params::program_impl` for the contract. ProgramImpl *program_impl_{nullptr}; diff --git a/quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp b/quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp index f7de6d5780..914e64d5ba 100644 --- a/quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp +++ b/quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp @@ -28,11 +28,14 @@ #include "quadrants/runtime/llvm/llvm_runtime_executor.h" #include "quadrants/program/adstack_size_expr_eval.h" +#include "quadrants/program/program.h" +#include #include #include #include #include +#include #include #include "quadrants/ir/static_adstack_bound_reducer_device.h" @@ -456,29 +459,13 @@ void LlvmRuntimeExecutor::ensure_per_task_float_heap_post_reducer(std::size_t ta clip_effective_rows_by_loop_trip_count(effective_rows, *ad_stack.bound_expr, std::numeric_limits::max(), prog, ctx); } - // Read back the per-thread float stride (in bytes) that `publish_adstack_metadata` published into - // `runtime->adstack_per_thread_stride_float`. `AdStackSizingInfo::per_thread_stride_float` from the analysis pre-pass - // is in entry-count units (`2 * max_size`), not bytes, and would massively undersize the heap. - uint64_t stride_float_bytes_u64 = 0; - if (runtime_adstack_stride_float_field_ptr_ != nullptr) { - if (config_.arch == Arch::cuda) { -#if defined(QD_WITH_CUDA) - CUDADriver::get_instance().memcpy_device_to_host(&stride_float_bytes_u64, runtime_adstack_stride_float_field_ptr_, - sizeof(uint64_t)); -#else - QD_NOT_IMPLEMENTED; -#endif - } else if (config_.arch == Arch::amdgpu) { -#if defined(QD_WITH_AMDGPU) - AMDGPUDriver::get_instance().memcpy_device_to_host(&stride_float_bytes_u64, - runtime_adstack_stride_float_field_ptr_, sizeof(uint64_t)); -#else - QD_NOT_IMPLEMENTED; -#endif - } else { - stride_float_bytes_u64 = *reinterpret_cast(runtime_adstack_stride_float_field_ptr_); - } - } + // The per-thread float stride (in bytes) was just published into `runtime->adstack_per_thread_stride_float` by the + // matching `publish_adstack_metadata` call earlier in this task's per-task block. We stash the value host-side so + // we can read it directly here instead of paying a sync DtoH on every bound_expr task. The launcher pairs publish + // + reducer + post-reducer per task with no intervening publish for another task, so the stash is accurate at this + // call site. `AdStackSizingInfo::per_thread_stride_float` from the analysis pre-pass is in entry-count units + // (`2 * max_size`), not bytes, and would massively undersize the heap. + uint64_t stride_float_bytes_u64 = static_cast(last_published_stride_float_bytes_); const std::size_t needed_bytes = effective_rows * static_cast(stride_float_bytes_u64); // `QD_DEBUG_ADSTACK=1` opt-in diagnostic. Persistent so memory regressions can be debugged without re-instrumenting. if (std::getenv("QD_DEBUG_ADSTACK")) { @@ -641,36 +628,59 @@ void LlvmRuntimeExecutor::ensure_adstack_heap_float(std::size_t needed_bytes) { } void LlvmRuntimeExecutor::check_adstack_overflow() { - // Called from `synchronize()` on every sync so adstack overflow surfaces as a Python exception regardless of - // `compile_config.debug`. The runtime / result buffer may not exist yet (e.g. a C++ test that constructs Program - // without materializing the runtime and then triggers Program::finalize -> synchronize), so no-op in that case. - if (llvm_runtime_ == nullptr || result_buffer_cache_ == nullptr) { + // Called from `synchronize()` on every sync, plus other Quadrants Python entry points wired in + // `Program::check_adstack_overflow_and_raise`. The flag lives in pinned host memory (allocated at + // `materialize_runtime`); polling is a relaxed atomic exchange on the cached host pointer via + // `std::atomic` reinterpret_cast - no DtoH, no JIT call, no sync drain. Available on all backends because + // the pinned-host memory is in the host process address space regardless of where the kernel that wrote it ran. + // The reinterpret_cast is portable because `std::atomic` is layout-compatible with `int64_t` on every + // target (verified by the static_assert below); see also Itanium ABI / MSVC ABI lock-free guarantees. + // + // Returns early when the slot has not been allocated yet (e.g. a C++ test that constructs Program without + // materializing the runtime and then triggers `Program::finalize -> synchronize`). + static_assert(std::atomic::is_always_lock_free, + "std::atomic must be lock-free for the reinterpret_cast pattern below to be portable"); + if (adstack_overflow_flag_host_ptr_ == nullptr) { + return; + } + int64_t flag = + reinterpret_cast *>(adstack_overflow_flag_host_ptr_)->exchange(0, std::memory_order_relaxed); + if (flag == 0) { return; } - // `lookup_function` returns a host-callable pointer only on direct-dispatch backends (CPU LLVM JIT). CUDA / AMDGPU - // return device-side pointers and require the arg-marshalling `JITModule::call` path; the cached fast path therefore - // only fires when `direct_dispatch()` is true. - auto *runtime_jit_module = get_runtime_jit_module(); - if (runtime_jit_module->direct_dispatch()) { - if (adstack_overflow_retriever_ == nullptr) { - adstack_overflow_retriever_ = reinterpret_cast( - runtime_jit_module->lookup_function("runtime_retrieve_and_reset_adstack_overflow")); - QD_ASSERT(adstack_overflow_retriever_ != nullptr); + // Drain the companion task-id slot in the same poll. Both slots cleared so the next overflow records a fresh + // identity. `task_id == 0` means the kernel that overflowed pre-dates the registry wiring or its + // `ad_stack.registry_id` was unset for any reason (e.g. a deserialised offline-cache task that has not yet been + // re-registered); the diagnose helper falls through to the generic dual-cause message in that case. + uint32_t task_id = 0; + if (adstack_overflow_task_id_host_ptr_ != nullptr) { + int64_t recorded = reinterpret_cast *>(adstack_overflow_task_id_host_ptr_) + ->exchange(0, std::memory_order_relaxed); + task_id = static_cast(recorded); + } + Program *prog = (program_impl_ != nullptr) ? program_impl_->program : nullptr; + std::string diagnostic; + if (prog != nullptr) { + auto diag = prog->adstack_cache().diagnose_adstack_overflow(task_id); + diagnostic = std::move(diag.message); + // Auto-invalidate the per-task metadata caches when the synchronous sizer rerun confirmed the cache is stale + // (DLPack-bypass cause). The current run is corrupted (we are about to raise), but the next launch's sizer + // reruns from scratch against the live (mutated) state and the kernel runs to completion without further + // user intervention. Unknown / Quadrants-bug cases skip the invalidation so a real sizer bug is not masked + // by silent recompute. + if (diag.confirmed_invalid_cache) { + prog->adstack_cache().invalidate_all_per_task(); } - adstack_overflow_retriever_(llvm_runtime_); } else { - runtime_jit_module->call("runtime_retrieve_and_reset_adstack_overflow", llvm_runtime_); - } - auto flag = fetch_result(quadrants_result_buffer_error_id, result_buffer_cache_); - if (flag != 0) { - throw QuadrantsAssertionError( + diagnostic = "Adstack overflow: a reverse-mode autodiff kernel pushed more elements than the adstack capacity " - "allows. Raised at the next qd.sync() rather than at the offending kernel launch. The pre-pass " - "resolved this alloca to a bound tighter than the actual runtime push count - either the enclosing " - "loop shape is outside the current `SizeExpr` grammar (rewrite it, or extend the grammar), or the " - "Bellman-Ford analyzer undercounted the forward-pass accumulation on this stack (file a bug with " - "the kernel IR via `QD_DUMP_IR=1`)."); + "allows."; } + throw QuadrantsAssertionError( + "Adstack overflow: a reverse-mode autodiff kernel pushed more elements " + "than the adstack capacity allows. Raised at the next Quadrants Python " + "entry rather than at the offending kernel launch.\n" + + diagnostic); } std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInfo &ad_stack, @@ -762,12 +772,148 @@ std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInf runtime_adstack_stride_int_field_ptr_ = quadrants_union_cast_with_different_sizes( fetch_result_uint64(quadrants_result_buffer_ret_value_id + 4, result_buffer_cache_)); } - copy_h2d(runtime_adstack_offsets_field_ptr_, &offsets_dev_ptr, sizeof(void *)); - copy_h2d(runtime_adstack_max_sizes_field_ptr_, &max_sizes_dev_ptr, sizeof(void *)); + // The pointed-to scratch allocations are stable across launches (only `grow_to` swaps them). Skip the per-launch + // h2d that publishes the pointer values whenever they have not changed since the last call. On HIP / CUDA each + // skipped pointer-publish is one queue round-trip the launcher would otherwise pay; on a typical reverse-mode + // sweep this fires thousands of times. + if (offsets_dev_ptr != adstack_offsets_dev_ptr_published_) { + copy_h2d(runtime_adstack_offsets_field_ptr_, &offsets_dev_ptr, sizeof(void *)); + adstack_offsets_dev_ptr_published_ = offsets_dev_ptr; + } + if (max_sizes_dev_ptr != adstack_max_sizes_dev_ptr_published_) { + copy_h2d(runtime_adstack_max_sizes_field_ptr_, &max_sizes_dev_ptr, sizeof(void *)); + adstack_max_sizes_dev_ptr_published_ = max_sizes_dev_ptr; + } std::size_t stride = 0; const bool is_gpu_llvm = (config_.arch == Arch::cuda || config_.arch == Arch::amdgpu); + // Shared GPU async publish helper: pack `[stride_combined, stride_float, stride_int, offsets[n_stacks], + // max_sizes[n_stacks]]` into the pinned-host scratch (grow on demand, double-amortised), then issue 5 async H2Ds + // on the active stream and record the completion event. Used by both the host-eval branch (CUDA / AMDGPU + // resolvable size_exprs) and the on-device-sizer cache-hit branch. The driver's H2D DMA reads from the pinned + // bytes at execution time, so a `wait_pending()` at the top of the next call defends against an unusual + // interleaving where the GPU queue is backlogged and the next launch enters before the previous launch's last + // copy has been consumed. Only callable when `is_gpu_llvm` is true. + auto publish_metadata_pinned_async = [&](const uint64_t *offsets_src, const uint64_t *max_sizes_src, + uint64_t stride_combined_u64, uint64_t stride_float_u64, + uint64_t stride_int_u64) { + const std::size_t header_bytes = 3 * sizeof(uint64_t); + const std::size_t array_bytes = n_stacks * sizeof(uint64_t); + const std::size_t total_bytes = header_bytes + 2 * array_bytes; + auto wait_pending = [this]() { + if (!pinned_metadata_event_pending_) { + return; + } +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + CUDADriver::get_instance().event_synchronize(pinned_metadata_event_); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + AMDGPUDriver::get_instance().event_synchronize(pinned_metadata_event_); + } +#endif + pinned_metadata_event_pending_ = false; + }; + if (total_bytes > pinned_metadata_scratch_capacity_) { + wait_pending(); + if (pinned_metadata_scratch_ != nullptr) { +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + CUDADriver::get_instance().mem_free_host(pinned_metadata_scratch_); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + AMDGPUDriver::get_instance().mem_free_host(pinned_metadata_scratch_); + } +#endif + pinned_metadata_scratch_ = nullptr; + } + const std::size_t new_capacity = std::max(total_bytes, 2 * pinned_metadata_scratch_capacity_); +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + CUDADriver::get_instance().mem_alloc_host(&pinned_metadata_scratch_, new_capacity); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + // `hipHostMallocDefault == 0`. Coherent / portable / write-combined flags are intentionally not set; the + // workload is small payloads written linearly by the host and DMA-read by the GPU once. + AMDGPUDriver::get_instance().mem_alloc_host(&pinned_metadata_scratch_, new_capacity, 0u); + } +#endif + pinned_metadata_scratch_capacity_ = new_capacity; + } + if (pinned_metadata_event_ == nullptr) { + // `cuEventCreate` flag `0` (CU_EVENT_DEFAULT) means timing-enabled, which the driver costs us nothing to set + // up here and lets future profilers attach without re-creating the event. `hipEventCreateWithFlags` takes + // the same encoding. +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + CUDADriver::get_instance().event_create(&pinned_metadata_event_, 0u); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + AMDGPUDriver::get_instance().event_create(&pinned_metadata_event_, 0u); + } +#endif + } + wait_pending(); + auto *pinned = static_cast(pinned_metadata_scratch_); + pinned[0] = stride_combined_u64; + pinned[1] = stride_float_u64; + pinned[2] = stride_int_u64; + std::memcpy(pinned + 3, offsets_src, array_bytes); + std::memcpy(pinned + 3 + n_stacks, max_sizes_src, array_bytes); + // Queue the metadata copies on the stream the subsequent main-kernel dispatch will run on, so the GPU + // stream-orders the copies before the kernel reads `adstack_max_sizes` etc. CUDA: `CUDAContext::get_stream()` + // (configurable via `set_stream`, defaults to the null stream); AMDGPU: always the default stream because + // `AMDGPUContext::launch` passes `nullptr` to `hipLaunchKernel`. +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + void *active_stream = CUDAContext::get_instance().get_stream(); + CUDADriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_field_ptr_, pinned, + sizeof(uint64_t), active_stream); + if (runtime_adstack_stride_float_field_ptr_ != nullptr) { + CUDADriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_float_field_ptr_, pinned + 1, + sizeof(uint64_t), active_stream); + } + if (runtime_adstack_stride_int_field_ptr_ != nullptr) { + CUDADriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_int_field_ptr_, pinned + 2, + sizeof(uint64_t), active_stream); + } + CUDADriver::get_instance().memcpy_host_to_device_async(offsets_dev_ptr, pinned + 3, array_bytes, active_stream); + CUDADriver::get_instance().memcpy_host_to_device_async(max_sizes_dev_ptr, pinned + 3 + n_stacks, array_bytes, + active_stream); + CUDADriver::get_instance().event_record(pinned_metadata_event_, active_stream); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + void *active_stream = nullptr; + AMDGPUDriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_field_ptr_, pinned, + sizeof(uint64_t), active_stream); + if (runtime_adstack_stride_float_field_ptr_ != nullptr) { + AMDGPUDriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_float_field_ptr_, pinned + 1, + sizeof(uint64_t), active_stream); + } + if (runtime_adstack_stride_int_field_ptr_ != nullptr) { + AMDGPUDriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_int_field_ptr_, pinned + 2, + sizeof(uint64_t), active_stream); + } + AMDGPUDriver::get_instance().memcpy_host_to_device_async(offsets_dev_ptr, pinned + 3, array_bytes, active_stream); + AMDGPUDriver::get_instance().memcpy_host_to_device_async(max_sizes_dev_ptr, pinned + 3 + n_stacks, array_bytes, + active_stream); + AMDGPUDriver::get_instance().event_record(pinned_metadata_event_, active_stream); + } +#endif + pinned_metadata_event_pending_ = true; + }; + // Host-eval fast path. The on-device sizer kernel exists to handle one specific leaf, `ExternalTensorRead`, // whose ndarray data lives in GPU-private memory (`cudaMalloc` / `hipMalloc`, no UVA fallback) and thus // cannot be touched from the host. Every other SizeExpr leaf - `Const`, `BoundVariable`, @@ -809,6 +955,8 @@ std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInf // Guard `program_impl_->program` lookups against the C++-only-tests setup where `program_impl_` itself is null; // the on-device branch below already does this and falls back to `max_size_compile_time`. Program *prog = (program_impl_ != nullptr) ? program_impl_->program : nullptr; + // Span the per-stack `evaluate_adstack_size_expr` calls below with one shared read cache. + SizeExprLaunchScope launch_scope; std::vector host_max_sizes(n_stacks); for (std::size_t i = 0; i < n_stacks; ++i) { const SerializedSizeExpr *expr = (i < ad_stack.size_exprs.size()) ? &ad_stack.size_exprs[i] : nullptr; @@ -860,139 +1008,8 @@ std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInf copy_h2d(runtime_adstack_stride_int_field_ptr_, &stride_int_u64, sizeof(uint64_t)); } } else { - // Five-block payload packed into the pinned-host scratch as `[stride_combined, stride_float, stride_int, - // offsets[n_stacks], max_sizes[n_stacks]]`. Five async DMAs land on the matching device addresses; the driver's - // H2D DMA engine reads from the pinned bytes at execution time, so we must not overwrite the scratch before all - // copies have completed - hence the per-launch `event_record` after the last copy and the `event_synchronize` at - // the top of the next launch. - const std::size_t header_bytes = 3 * sizeof(uint64_t); - const std::size_t array_bytes = n_stacks * sizeof(uint64_t); - const std::size_t total_bytes = header_bytes + 2 * array_bytes; - - auto wait_pending = [this]() { - if (!pinned_metadata_event_pending_) { - return; - } -#if defined(QD_WITH_CUDA) - if (config_.arch == Arch::cuda) { - CUDADriver::get_instance().event_synchronize(pinned_metadata_event_); - } -#endif -#if defined(QD_WITH_AMDGPU) - if (config_.arch == Arch::amdgpu) { - AMDGPUDriver::get_instance().event_synchronize(pinned_metadata_event_); - } -#endif - pinned_metadata_event_pending_ = false; - }; - - // Grow / first-allocate the pinned host scratch and the per-launch completion event. Doubling growth - // means the pinned alloc / free traffic is amortised to O(log peak_total_bytes) across a run. - if (total_bytes > pinned_metadata_scratch_capacity_) { - wait_pending(); - if (pinned_metadata_scratch_ != nullptr) { -#if defined(QD_WITH_CUDA) - if (config_.arch == Arch::cuda) { - CUDADriver::get_instance().mem_free_host(pinned_metadata_scratch_); - } -#endif -#if defined(QD_WITH_AMDGPU) - if (config_.arch == Arch::amdgpu) { - AMDGPUDriver::get_instance().mem_free_host(pinned_metadata_scratch_); - } -#endif - pinned_metadata_scratch_ = nullptr; - } - std::size_t new_capacity = std::max(total_bytes, 2 * pinned_metadata_scratch_capacity_); -#if defined(QD_WITH_CUDA) - if (config_.arch == Arch::cuda) { - CUDADriver::get_instance().mem_alloc_host(&pinned_metadata_scratch_, new_capacity); - } -#endif -#if defined(QD_WITH_AMDGPU) - if (config_.arch == Arch::amdgpu) { - // `hipHostMallocDefault == 0`. Coherent / portable / write-combined flags are intentionally not set; - // the workload is small payloads written linearly by the host and DMA-read by the GPU once. - AMDGPUDriver::get_instance().mem_alloc_host(&pinned_metadata_scratch_, new_capacity, 0u); - } -#endif - pinned_metadata_scratch_capacity_ = new_capacity; - } - if (pinned_metadata_event_ == nullptr) { - // `cuEventCreate` flag `0` (CU_EVENT_DEFAULT) means timing-enabled, which the driver costs us nothing - // to set up here and lets future profilers attach without re-creating the event. `hipEventCreateWithFlags` - // takes the same encoding. -#if defined(QD_WITH_CUDA) - if (config_.arch == Arch::cuda) { - CUDADriver::get_instance().event_create(&pinned_metadata_event_, 0u); - } -#endif -#if defined(QD_WITH_AMDGPU) - if (config_.arch == Arch::amdgpu) { - AMDGPUDriver::get_instance().event_create(&pinned_metadata_event_, 0u); - } -#endif - } - // Block until any in-flight copies from the previous launch have finished pulling from the pinned scratch - // before we overwrite it. In steady state this is a no-op because the small DMAs finish well before the - // host loops back here; the wait exists only to defend against an unusual interleaving where the GPU - // queue is backlogged and the next launch enters this function before the previous launch's last copy - // has been consumed. - wait_pending(); - - auto *pinned = static_cast(pinned_metadata_scratch_); - pinned[0] = stride_combined_u64; - pinned[1] = stride_float_u64; - pinned[2] = stride_int_u64; - std::memcpy(pinned + 3, host_offsets.data(), array_bytes); - std::memcpy(pinned + 3 + n_stacks, host_max_sizes.data(), array_bytes); - - // Queue the metadata copies on the same stream the subsequent main-kernel dispatch will run on, so the - // GPU stream-orders the copies before the kernel reads `adstack_max_sizes` etc. On CUDA the active - // stream is `CUDAContext::get_instance().get_stream()` - configurable via `set_stream`, defaults to the - // null stream - and `CUDAContext::launch` dispatches kernels on the same handle. AMDGPU has no - // public stream-selection API: `AMDGPUContext::launch` always passes `nullptr` to `hipLaunchKernel` - // (i.e. the default stream), so the copies match that. -#if defined(QD_WITH_CUDA) - if (config_.arch == Arch::cuda) { - void *active_stream = CUDAContext::get_instance().get_stream(); - CUDADriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_field_ptr_, pinned, - sizeof(uint64_t), active_stream); - if (runtime_adstack_stride_float_field_ptr_ != nullptr) { - CUDADriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_float_field_ptr_, pinned + 1, - sizeof(uint64_t), active_stream); - } - if (runtime_adstack_stride_int_field_ptr_ != nullptr) { - CUDADriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_int_field_ptr_, pinned + 2, - sizeof(uint64_t), active_stream); - } - CUDADriver::get_instance().memcpy_host_to_device_async(offsets_dev_ptr, pinned + 3, array_bytes, active_stream); - CUDADriver::get_instance().memcpy_host_to_device_async(max_sizes_dev_ptr, pinned + 3 + n_stacks, array_bytes, - active_stream); - CUDADriver::get_instance().event_record(pinned_metadata_event_, active_stream); - } -#endif -#if defined(QD_WITH_AMDGPU) - if (config_.arch == Arch::amdgpu) { - void *active_stream = nullptr; // AMDGPUContext::launch always uses the default stream. - AMDGPUDriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_field_ptr_, pinned, - sizeof(uint64_t), active_stream); - if (runtime_adstack_stride_float_field_ptr_ != nullptr) { - AMDGPUDriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_float_field_ptr_, pinned + 1, - sizeof(uint64_t), active_stream); - } - if (runtime_adstack_stride_int_field_ptr_ != nullptr) { - AMDGPUDriver::get_instance().memcpy_host_to_device_async(runtime_adstack_stride_int_field_ptr_, pinned + 2, - sizeof(uint64_t), active_stream); - } - AMDGPUDriver::get_instance().memcpy_host_to_device_async(offsets_dev_ptr, pinned + 3, array_bytes, - active_stream); - AMDGPUDriver::get_instance().memcpy_host_to_device_async(max_sizes_dev_ptr, pinned + 3 + n_stacks, array_bytes, - active_stream); - AMDGPUDriver::get_instance().event_record(pinned_metadata_event_, active_stream); - } -#endif - pinned_metadata_event_pending_ = true; + publish_metadata_pinned_async(host_offsets.data(), host_max_sizes.data(), stride_combined_u64, stride_float_u64, + stride_int_u64); } } else { // GPU (CUDA / AMDGPU): encode the SizeExpr trees into device bytecode, upload, launch the sizer runtime @@ -1005,73 +1022,173 @@ std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInf // deref reads garbage. Moving the interpreter on-device keeps the pointer semantics intact - it reads the // data pointer out of `ctx->arg_buffer` (which the kernel will read too) and dereferences it where the // memory lives, with no migration / readback of the ndarray payload itself. - std::vector bytecode; - if (program_impl_ != nullptr && program_impl_->program != nullptr) { - bytecode = encode_adstack_size_expr_device_bytecode(ad_stack, program_impl_->program, ctx); - } else { - // No program attached (rare: C++-only tests that construct Program without a full runtime). Fall through - // to compile-time bounds by emitting an empty-tree bytecode - the device interpreter sees - // `root_node_idx == -1` for every stack and routes to `max_size_compile_time`. - bytecode = encode_adstack_size_expr_device_bytecode(ad_stack, nullptr, ctx); - } - // Grow the scratch buffer if the bytecode outgrew the cached capacity. Amortised doubling keeps the - // allocation traffic O(log max_bytecode_bytes) across a run. - const std::size_t bytecode_bytes = bytecode.size(); - if (bytecode_bytes > adstack_sizer_bytecode_capacity_) { - std::size_t new_cap = std::max(bytecode_bytes, 2 * adstack_sizer_bytecode_capacity_); - Device::AllocParams params{}; - params.size = new_cap; - params.host_read = false; - params.host_write = false; - params.export_sharing = false; - params.usage = AllocUsage::Storage; - DeviceAllocation new_alloc; - RhiResult res = llvm_device()->allocate_memory(params, &new_alloc); - QD_ERROR_IF(res != RhiResult::success, - "Failed to allocate {} bytes for the adstack sizer bytecode scratch buffer (err: {})", params.size, - int(res)); - adstack_sizer_bytecode_alloc_ = std::make_unique(std::move(new_alloc)); - adstack_sizer_bytecode_capacity_ = new_cap; + // + // Per-task metadata cache fast path: the sizer kernel's output (offsets / max_sizes / strides) is a + // deterministic function of (a) the per-task `AdStackSizingInfo *` (compile-time bytecode shape, stable + // for the kernel's lifetime), (b) every SNode value a `FieldLoad` leaf reads, and (c) every ndarray + // value an `ExternalTensorRead` leaf reads. Each launcher (cpu / cuda / amdgpu) bumps + // `Program::snode_write_gen_` / `ndarray_data_gen_` for everything this kernel may mutate before + // calling here, so the per-source generation snapshots stored alongside the cached payload catch any + // input change between launches and force a fresh sizer dispatch when needed. On hit, the cached + // offsets / max_sizes / strides are republished into the runtime struct via the same `copy_h2d` paths + // the host-eval branch above uses, and the entire bytecode-encode + h2d + sizer-kernel launch + + // 3x DtoH-stride pipeline is skipped. The cost of the sizer dispatch + DtoH stalls is small per + // launch on CUDA / AMDGPU, but a long sequence of reverse-mode launches over the same kernel + // pays it once per launch; the cache amortises that to once per generation-bump. + Program *prog = (program_impl_ != nullptr) ? program_impl_->program : nullptr; + bool llvm_metadata_cache_hit = false; + if (prog != nullptr) { + AdStackCache::LlvmPerTaskAdStackCacheEntry entry; + if (prog->adstack_cache().try_llvm_per_task_ad_stack_cache_hit(static_cast(&ad_stack), ctx, + entry)) { + QD_ASSERT(entry.offsets.size() == n_stacks && entry.max_sizes.size() == n_stacks); + // Publish the cached payload through the pinned-host async pipeline shared with the host-eval + // branch above: one pinned-scratch pack + five `memcpy_host_to_device_async` issued on the same + // stream the main kernel will dispatch on, ordered behind the previous launch's + // `pinned_metadata_event_pending_` wait. Packing the same `[stride_combined, stride_float, + // stride_int, offsets[n_stacks], max_sizes[n_stacks]]` shape keeps both branches' DMA pattern + // identical and removes the per-launch sync round-trips a `copy_h2d` would otherwise impose; on + // CPU `copy_h2d` is `memcpy` already so we keep the direct path there. + if (!is_gpu_llvm) { + copy_h2d(offsets_dev_ptr, entry.offsets.data(), n_stacks * sizeof(uint64_t)); + copy_h2d(max_sizes_dev_ptr, entry.max_sizes.data(), n_stacks * sizeof(uint64_t)); + copy_h2d(runtime_adstack_stride_field_ptr_, &entry.stride_combined, sizeof(uint64_t)); + if (runtime_adstack_stride_float_field_ptr_ != nullptr) { + copy_h2d(runtime_adstack_stride_float_field_ptr_, &entry.stride_float, sizeof(uint64_t)); + } + if (runtime_adstack_stride_int_field_ptr_ != nullptr) { + copy_h2d(runtime_adstack_stride_int_field_ptr_, &entry.stride_int, sizeof(uint64_t)); + } + } else { + publish_metadata_pinned_async(entry.offsets.data(), entry.max_sizes.data(), entry.stride_combined, + entry.stride_float, entry.stride_int); + } + stride = static_cast(entry.stride_combined); + stride_float_bytes = static_cast(entry.stride_float); + stride_int_bytes = static_cast(entry.stride_int); + llvm_metadata_cache_hit = true; + } } - void *bytecode_dev_ptr = get_device_alloc_info_ptr(*adstack_sizer_bytecode_alloc_); - copy_h2d(bytecode_dev_ptr, bytecode.data(), bytecode_bytes); + if (!llvm_metadata_cache_hit) { + std::vector bytecode; + if (program_impl_ != nullptr && program_impl_->program != nullptr) { + bytecode = encode_adstack_size_expr_device_bytecode(ad_stack, program_impl_->program, ctx); + } else { + // No program attached (rare: C++-only tests that construct Program without a full runtime). Fall through + // to compile-time bounds by emitting an empty-tree bytecode - the device interpreter sees + // `root_node_idx == -1` for every stack and routes to `max_size_compile_time`. + bytecode = encode_adstack_size_expr_device_bytecode(ad_stack, nullptr, ctx); + } + // Grow the scratch buffer if the bytecode outgrew the cached capacity. Amortised doubling keeps the + // allocation traffic O(log max_bytecode_bytes) across a run. + const std::size_t bytecode_bytes = bytecode.size(); + if (bytecode_bytes > adstack_sizer_bytecode_capacity_) { + std::size_t new_cap = std::max(bytecode_bytes, 2 * adstack_sizer_bytecode_capacity_); + Device::AllocParams params{}; + params.size = new_cap; + params.host_read = false; + params.host_write = false; + params.export_sharing = false; + params.usage = AllocUsage::Storage; + DeviceAllocation new_alloc; + RhiResult res = llvm_device()->allocate_memory(params, &new_alloc); + QD_ERROR_IF(res != RhiResult::success, + "Failed to allocate {} bytes for the adstack sizer bytecode scratch buffer (err: {})", params.size, + int(res)); + adstack_sizer_bytecode_alloc_ = std::make_unique(std::move(new_alloc)); + adstack_sizer_bytecode_capacity_ = new_cap; + } + void *bytecode_dev_ptr = get_device_alloc_info_ptr(*adstack_sizer_bytecode_alloc_); + copy_h2d(bytecode_dev_ptr, bytecode.data(), bytecode_bytes); - // Invoke the device interpreter. On CUDA / AMDGPU `JITModule::call` launches this as a single-thread kernel - // on the default stream and stream-orders it before the subsequent main-kernel dispatch, so the writes we - // do here are visible by the time the user's kernel reads `adstack_max_sizes` etc. - // - // The sizer kernel dereferences `ctx->arg_buffer` on device (that's how it resolves `ExternalTensorRead` leaves - // against ndarray pointers the caller packed into the arg buffer). AMDGPU always stages a device-side copy of - // `RuntimeContext` because HIP has no UVA fallback and the host pointer faults with `hipErrorIllegalAddress`. CUDA - // stages the device copy only when the driver + kernel do not expose HMM / system-allocated memory (queried via - // `CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS`): CUDA UVA covers pinned / CUDA-managed memory only, not the plain - // `std::make_unique()` backing, so a host pointer works on HMM-capable setups but faults otherwise - // (Turing without HMM, Windows, pre-535 Linux drivers) as `CUDA_ERROR_ILLEGAL_ADDRESS` at the next DtoH sync - // `illegal memory access ... while calling memcpy_device_to_host`. When the caller passes `nullptr` (HMM-capable - // CUDA) we fall back to the host pointer; the launcher gates the allocation so HMM-equipped setups pay no staging - // cost. - auto *const runtime_jit = get_runtime_jit_module(); - void *runtime_context_ptr_for_sizer = - device_runtime_context_ptr != nullptr ? device_runtime_context_ptr : static_cast(&ctx->get_context()); - runtime_jit->call("runtime_eval_adstack_size_expr", llvm_runtime_, - runtime_context_ptr_for_sizer, bytecode_dev_ptr); + // Invoke the device interpreter. On CUDA / AMDGPU `JITModule::call` launches this as a single-thread kernel + // on the default stream and stream-orders it before the subsequent main-kernel dispatch, so the writes we + // do here are visible by the time the user's kernel reads `adstack_max_sizes` etc. + // + // The sizer kernel dereferences `ctx->arg_buffer` on device (that's how it resolves `ExternalTensorRead` leaves + // against ndarray pointers the caller packed into the arg buffer). AMDGPU always stages a device-side copy of + // `RuntimeContext` because HIP has no UVA fallback and the host pointer faults with `hipErrorIllegalAddress`. + // CUDA stages the device copy only when the driver + kernel do not expose HMM / system-allocated memory (queried + // via `CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS`): CUDA UVA covers pinned / CUDA-managed memory only, not the + // plain `std::make_unique()` backing, so a host pointer works on HMM-capable setups but faults + // otherwise (Turing without HMM, Windows, pre-535 Linux drivers) as `CUDA_ERROR_ILLEGAL_ADDRESS` at the next DtoH + // sync `illegal memory access ... while calling memcpy_device_to_host`. When the caller passes `nullptr` + // (HMM-capable CUDA) we fall back to the host pointer; the launcher gates the allocation so HMM-equipped setups + // pay no staging cost. + auto *const runtime_jit = get_runtime_jit_module(); + void *runtime_context_ptr_for_sizer = + device_runtime_context_ptr != nullptr ? device_runtime_context_ptr : static_cast(&ctx->get_context()); + runtime_jit->call("runtime_eval_adstack_size_expr", llvm_runtime_, + runtime_context_ptr_for_sizer, bytecode_dev_ptr); - // Read back the per-kind strides published by `runtime_eval_adstack_size_expr` so we can size the float and int - // heaps independently host-side. The combined stride is unused by the split-heap codegen but kept around for - // legacy-kernel backward compatibility (mirrors `stride_int` in the unconditional-split layout). - uint64_t stride_combined_readback = 0; - uint64_t stride_float_readback = 0; - uint64_t stride_int_readback = 0; - copy_d2h(&stride_combined_readback, runtime_adstack_stride_field_ptr_, sizeof(uint64_t)); - if (runtime_adstack_stride_float_field_ptr_ != nullptr) { - copy_d2h(&stride_float_readback, runtime_adstack_stride_float_field_ptr_, sizeof(uint64_t)); - } - if (runtime_adstack_stride_int_field_ptr_ != nullptr) { - copy_d2h(&stride_int_readback, runtime_adstack_stride_int_field_ptr_, sizeof(uint64_t)); - } - stride = static_cast(stride_combined_readback); - stride_float_bytes = static_cast(stride_float_readback); - stride_int_bytes = static_cast(stride_int_readback); + // Read back the per-kind strides published by `runtime_eval_adstack_size_expr` so we can size the float and int + // heaps independently host-side. The combined stride is unused by the split-heap codegen but kept around for + // legacy-kernel backward compatibility (mirrors `stride_int` in the unconditional-split layout). + uint64_t stride_combined_readback = 0; + uint64_t stride_float_readback = 0; + uint64_t stride_int_readback = 0; + copy_d2h(&stride_combined_readback, runtime_adstack_stride_field_ptr_, sizeof(uint64_t)); + if (runtime_adstack_stride_float_field_ptr_ != nullptr) { + copy_d2h(&stride_float_readback, runtime_adstack_stride_float_field_ptr_, sizeof(uint64_t)); + } + if (runtime_adstack_stride_int_field_ptr_ != nullptr) { + copy_d2h(&stride_int_readback, runtime_adstack_stride_int_field_ptr_, sizeof(uint64_t)); + } + stride = static_cast(stride_combined_readback); + stride_float_bytes = static_cast(stride_float_readback); + stride_int_bytes = static_cast(stride_int_readback); + + // Record the cache entry so the next launch on this kernel can skip the sizer pipeline. We also + // need to read back the offsets / max_sizes arrays the sizer wrote to the device buffers - the + // cache hit path above republishes them, so we must store host copies here. n_stacks is small + // (a few dozen at most for any reasonable kernel) so the extra DtoH cost is negligible + // compared to the dispatch + sizer-kernel launch we are about to amortise away. + if (prog != nullptr) { + std::vector offsets_readback(n_stacks); + std::vector max_sizes_readback(n_stacks); + copy_d2h(offsets_readback.data(), offsets_dev_ptr, n_stacks * sizeof(uint64_t)); + copy_d2h(max_sizes_readback.data(), max_sizes_dev_ptr, n_stacks * sizeof(uint64_t)); + // Walk size_exprs structurally to gather the dependency keys (snode_ids referenced via + // FieldLoad, arg_ids referenced via ExternalTensorShape / ExternalTensorRead). Pure tree + // inspection - no live value reads, no nested kernel launches. Mirrors the SPIR-V analogue. + std::unordered_set snode_ids; + std::unordered_set arg_ids; + for (const auto &expr : ad_stack.size_exprs) { + for (const auto &node : expr.nodes) { + switch (static_cast(node.kind)) { + case SizeExpr::Kind::FieldLoad: + if (node.snode_id >= 0) + snode_ids.insert(node.snode_id); + break; + case SizeExpr::Kind::ExternalTensorShape: + case SizeExpr::Kind::ExternalTensorRead: + if (!node.arg_id_path.empty()) + arg_ids.insert(node.arg_id_path.front()); + break; + default: + break; + } + } + } + std::vector> snode_gens; + snode_gens.reserve(snode_ids.size()); + for (int snode_id : snode_ids) { + snode_gens.emplace_back(snode_id, prog->adstack_cache().snode_write_gen(snode_id)); + } + std::vector> arg_gens; + arg_gens.reserve(arg_ids.size()); + for (int arg_id : arg_ids) { + ArgArrayPtrKey data_key{arg_id, TypeFactory::DATA_PTR_POS_IN_NDARRAY}; + auto ap_it = ctx->array_ptrs.find(data_key); + void *devalloc = (ap_it == ctx->array_ptrs.end()) ? nullptr : ap_it->second; + arg_gens.emplace_back(arg_id, devalloc, prog->adstack_cache().ndarray_data_gen(devalloc)); + } + prog->adstack_cache().record_llvm_per_task_ad_stack( + static_cast(&ad_stack), std::move(offsets_readback), std::move(max_sizes_readback), + stride_combined_readback, stride_float_readback, stride_int_readback, std::move(snode_gens), + std::move(arg_gens)); + } + } // end if (!llvm_metadata_cache_hit) } // Legacy combined heap: not allocated. The unconditional-split codegen reads `heap_float` for f32 allocas and @@ -1108,6 +1225,7 @@ std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInf } ensure_adstack_heap_float(float_bytes); } + last_published_stride_float_bytes_ = stride_float_bytes; return needed_bytes; } diff --git a/quadrants/runtime/llvm/llvm_runtime_executor.cpp b/quadrants/runtime/llvm/llvm_runtime_executor.cpp index 658c139c0f..663e8ffbb0 100644 --- a/quadrants/runtime/llvm/llvm_runtime_executor.cpp +++ b/quadrants/runtime/llvm/llvm_runtime_executor.cpp @@ -17,6 +17,7 @@ #include "quadrants/rhi/cuda/cuda_driver.h" #include "quadrants/rhi/llvm/device_memory_pool.h" #include "quadrants/program/program_impl.h" +#include "quadrants/program/program.h" #if defined(QD_WITH_CUDA) #include "quadrants/rhi/cuda/cuda_context.h" @@ -193,6 +194,10 @@ void LlvmRuntimeExecutor::print_list_manager_info(void *list_manager, uint64 *re elements_per_chunk, element_size, size_MB); } +Program *LlvmRuntimeExecutor::get_program() const { + return program_impl_ != nullptr ? program_impl_->program : nullptr; +} + void LlvmRuntimeExecutor::synchronize() { if (config_.arch == Arch::cuda) { #if defined(QD_WITH_CUDA) @@ -544,6 +549,44 @@ void LlvmRuntimeExecutor::finalize() { pinned_metadata_scratch_ = nullptr; pinned_metadata_scratch_capacity_ = 0; } + // Release the pinned host slot used for the adstack overflow flag. Mirrors the pinned_metadata_scratch + // release above. The `LLVMRuntime::adstack_overflow_flag_dev_ptr` field that referenced this slot is in the + // runtime struct, which is destroyed below; the dangling reference is harmless because the kernel JIT + // module is also being torn down. + if (adstack_overflow_flag_host_ptr_ != nullptr) { +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + CUDADriver::get_instance().mem_free_host(adstack_overflow_flag_host_ptr_); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + AMDGPUDriver::get_instance().mem_free_host(adstack_overflow_flag_host_ptr_); + } +#endif + if (config_.arch != Arch::cuda && config_.arch != Arch::amdgpu) { + std::free(adstack_overflow_flag_host_ptr_); + } + adstack_overflow_flag_host_ptr_ = nullptr; + adstack_overflow_flag_dev_ptr_ = nullptr; + } + if (adstack_overflow_task_id_host_ptr_ != nullptr) { +#if defined(QD_WITH_CUDA) + if (config_.arch == Arch::cuda) { + CUDADriver::get_instance().mem_free_host(adstack_overflow_task_id_host_ptr_); + } +#endif +#if defined(QD_WITH_AMDGPU) + if (config_.arch == Arch::amdgpu) { + AMDGPUDriver::get_instance().mem_free_host(adstack_overflow_task_id_host_ptr_); + } +#endif + if (config_.arch != Arch::cuda && config_.arch != Arch::amdgpu) { + std::free(adstack_overflow_task_id_host_ptr_); + } + adstack_overflow_task_id_host_ptr_ = nullptr; + adstack_overflow_task_id_dev_ptr_ = nullptr; + } if (config_.arch == Arch::cuda || config_.arch == Arch::amdgpu) { preallocated_runtime_objects_allocs_.reset(); preallocated_runtime_memory_allocs_.reset(); @@ -731,6 +774,68 @@ void LlvmRuntimeExecutor::materialize_runtime(KernelProfilerBase *profiler, uint runtime_jit->call("LLVMRuntime_set_profiler_stop", llvm_runtime_, (void *)&KernelProfilerBase::profiler_stop); } + + // Allocate the pinned host slot for the adstack overflow flag and publish its device-mapped address into the + // runtime. The kernel-side `stack_push` writes the overflow signal here via a system-wide atomic; the host polls + // the same memory directly via `adstack_overflow_flag_host_ptr_`. CUDA / AMDGPU pinned host memory is already + // UVA-mapped so the same pointer is valid from both sides; on CPU the runtime is host-resident and the same + // pointer is used unchanged. Required hardware: NVIDIA Compute Capability 6.0+ / AMD GFX9+, the same envelope + // the existing pinned-host H2D-async pattern in `llvm_adstack_lazy_claim.cpp` already requires. + { + void *host_slot = nullptr; + if (config_.arch == Arch::cuda) { +#if defined(QD_WITH_CUDA) + CUDADriver::get_instance().mem_alloc_host(&host_slot, sizeof(int64_t)); +#else + QD_NOT_IMPLEMENTED; +#endif + } else if (config_.arch == Arch::amdgpu) { +#if defined(QD_WITH_AMDGPU) + AMDGPUDriver::get_instance().mem_alloc_host(&host_slot, sizeof(int64_t), 0u); +#else + QD_NOT_IMPLEMENTED; +#endif + } else { + host_slot = std::malloc(sizeof(int64_t)); + } + QD_ASSERT(host_slot != nullptr); + adstack_overflow_flag_host_ptr_ = static_cast(host_slot); + *adstack_overflow_flag_host_ptr_ = 0; + // CUDA `cuMemAllocHost_v2` and HIP `hipHostMalloc` with default flags both return UVA-mapped memory; the + // host pointer is also a valid device pointer on Pascal+ / GFX9+ hardware. On CPU the runtime is in host + // memory and the kernel runs as a function call, so the same pointer applies. + adstack_overflow_flag_dev_ptr_ = host_slot; + runtime_jit->call("runtime_set_adstack_overflow_flag_dev_ptr", llvm_runtime_, + adstack_overflow_flag_dev_ptr_); + } + // Companion task-id slot. Same allocation strategy as the flag above; placed in a separate page so the + // codegen-emitted `cmpxchg` does not contend with the flag's `atomic OR` on a shared cache line. Codegen + // writes the Program-assigned `adstack_sizing_info_id` here on the first overflowing thread; host reads + // the slot during the raise to look up the offending kernel / task in the Program-side registry. + { + void *host_slot = nullptr; + if (config_.arch == Arch::cuda) { +#if defined(QD_WITH_CUDA) + CUDADriver::get_instance().mem_alloc_host(&host_slot, sizeof(int64_t)); +#else + QD_NOT_IMPLEMENTED; +#endif + } else if (config_.arch == Arch::amdgpu) { +#if defined(QD_WITH_AMDGPU) + AMDGPUDriver::get_instance().mem_alloc_host(&host_slot, sizeof(int64_t), 0u); +#else + QD_NOT_IMPLEMENTED; +#endif + } else { + host_slot = std::malloc(sizeof(int64_t)); + } + QD_ASSERT(host_slot != nullptr); + adstack_overflow_task_id_host_ptr_ = static_cast(host_slot); + *adstack_overflow_task_id_host_ptr_ = 0; + adstack_overflow_task_id_dev_ptr_ = host_slot; + runtime_jit->call("runtime_set_adstack_overflow_task_id_dev_ptr", llvm_runtime_, + adstack_overflow_task_id_dev_ptr_); + } } void LlvmRuntimeExecutor::destroy_snode_tree(SNodeTree *snode_tree) { diff --git a/quadrants/runtime/llvm/llvm_runtime_executor.h b/quadrants/runtime/llvm/llvm_runtime_executor.h index 2404f3f7c2..623f711f79 100644 --- a/quadrants/runtime/llvm/llvm_runtime_executor.h +++ b/quadrants/runtime/llvm/llvm_runtime_executor.h @@ -107,6 +107,14 @@ class LlvmRuntimeExecutor { LaunchContextBuilder *ctx, void *device_runtime_context_ptr = nullptr); + // Accessor for the per-program `Program *` used by the launchers to bump per-snode and per- + // DeviceAllocation generation counters in `Program::snode_write_gen_` / `ndarray_data_gen_`. Used + // exclusively by the LLVM-GPU per-task adstack metadata cache (`record_llvm_per_task_ad_stack`) + // so a kernel write to a SNode or ndarray a downstream `size_expr::FieldLoad` / + // `ExternalTensorRead` reads invalidates cached metadata on the next launch. Returns nullptr in + // C++-only test contexts where `program_impl_` was constructed without a parent Program. + Program *get_program() const; + // Allocate-on-demand and clear the per-kernel lazy-claim arrays: // `adstack_row_counters[num_tasks]` = 0 (codegen-emitted LCA-block atomic-rmw target; each task counts its own // LCA-block-reaching threads in slot `task_codegen_id`) @@ -240,10 +248,24 @@ class LlvmRuntimeExecutor { // this cache, so avoid that. uint64 *result_buffer_cache_{nullptr}; - // Memoised host-callable pointer to `runtime_retrieve_and_reset_adstack_overflow`. Populated only on backends - // whose JIT module supports direct dispatch (CPU LLVM); other backends leave it null and route through - // `JITModule::call`. Valid for the lifetime of the runtime JIT module. - void (*adstack_overflow_retriever_)(void *){nullptr}; + // Pinned host slot for the adstack overflow flag. The kernel-side `stack_push` writes through the runtime's + // device-mapped address (`adstack_overflow_flag_dev_ptr_` published into `LLVMRuntime` at materialise time); + // the host polls `adstack_overflow_flag_host_ptr_` directly with a relaxed atomic exchange - no DtoH, no JIT call, + // no sync drain. Backends: + // - CPU LLVM: plain malloc; host and device addresses are identical. + // - CUDA: `cuMemAllocHost_v2` plus `cuMemHostGetDevicePointer` for the device-mapped address. + // - AMDGPU: `hipHostMalloc(0)` plus the HIP equivalent. + // Required hardware envelope (Pascal+ / GFX9+) matches what the existing pinned-host H2D-async pattern in + // `llvm_adstack_lazy_claim.cpp` already needs. + int64_t *adstack_overflow_flag_host_ptr_{nullptr}; + void *adstack_overflow_flag_dev_ptr_{nullptr}; + + // Companion task-id slot. Codegen emits `cmpxchg(0, id)` from the lazy-claim overflow path; the first + // overflowing thread's `Program::adstack_sizing_info_registry_` id sticks. Host reads the slot at the + // raise site for diagnostic message generation. Same allocation strategy as the flag above; on a fresh + // page so neither write contends with the flag's atomic OR. + int64_t *adstack_overflow_task_id_host_ptr_{nullptr}; + void *adstack_overflow_task_id_dev_ptr_{nullptr}; std::unique_ptr thread_pool_{nullptr}; std::shared_ptr device_{nullptr}; @@ -321,6 +343,17 @@ class LlvmRuntimeExecutor { DeviceAllocationUnique adstack_offsets_alloc_ = nullptr; DeviceAllocationUnique adstack_max_sizes_alloc_ = nullptr; std::size_t adstack_metadata_capacity_{0}; + // Last device-pointer values published into `runtime->adstack_offsets` / `adstack_max_sizes`. The pointed-to scratch + // allocations stay stable across launches (only grown via `grow_to`), so we skip the per-launch h2d when the pointer + // we would write matches what we already wrote on a prior launch. + void *adstack_offsets_dev_ptr_published_{nullptr}; + void *adstack_max_sizes_dev_ptr_published_{nullptr}; + // The float stride (in bytes) the most recent `publish_adstack_metadata` wrote to + // `runtime->adstack_per_thread_stride_float`. Read by `ensure_per_task_float_heap_post_reducer` for the matching + // task to size the float heap; we keep the value host-side so we do not need to DtoH-readback the stride field on + // every bound_expr task. The launcher pairs publish + reducer + post-reducer per task with no intervening publish, + // so the value remains accurate at the consumer call site. + std::size_t last_published_stride_float_bytes_{0}; // Per-launch scratch buffer used on GPU arches (CUDA / AMDGPU) to ship the `LlvmAdStackBoundReducerDeviceParams` blob // into for `runtime_eval_static_bound_count`. Allocated on demand on the first bound_expr task in a kernel, reused diff --git a/quadrants/runtime/llvm/runtime_module/internal_functions.h b/quadrants/runtime/llvm/runtime_module/internal_functions.h index 9cac852c32..6c9ded4dcf 100644 --- a/quadrants/runtime/llvm/runtime_module/internal_functions.h +++ b/quadrants/runtime/llvm/runtime_module/internal_functions.h @@ -45,17 +45,26 @@ i32 test_internal_func_args(RuntimeContext *context, float32 i, float32 j, int32 i32 test_stack(RuntimeContext *context) { auto *runtime = context->runtime; // Header u64 `n` + max_num_elements * 2 * element_size for primal+adjoint slot pairs. Allocate generously for - // the guard-case subtests below. - auto stack = new u8[8 + 16 * 2 * 4]; + // the guard-case subtests below. Stack-allocated rather than `new u8[...]` to keep the JIT bitcode free of + // `operator new[]` / `operator delete[]` references that some Linux JIT linker configurations cannot resolve. + u8 stack_storage[8 + 16 * 2 * 4]; + u8 *stack = stack_storage; stack_init(stack); + // Stash any prior overflow-flag pointer the host has installed and point the runtime at a local slot for + // the duration of this test. The local slot is host memory in this in-process test fixture, mirroring the + // pinned-host slot installed by `LlvmRuntimeExecutor::materialize_runtime` in production. + i64 *prev_flag_dev_ptr = runtime->adstack_overflow_flag_dev_ptr; + i64 local_flag = 0; + runtime->adstack_overflow_flag_dev_ptr = &local_flag; + // Basic push/pop accounting. - stack_push(runtime, stack, 16, 4); - stack_push(runtime, stack, 16, 4); - stack_push(runtime, stack, 16, 4); - stack_push(runtime, stack, 16, 4); + stack_push(runtime, stack, 16, 4, 0); + stack_push(runtime, stack, 16, 4, 0); + stack_push(runtime, stack, 16, 4, 0); + stack_push(runtime, stack, 16, 4, 0); QD_TEST_CHECK(*(u64 *)stack == 4, runtime); - QD_TEST_CHECK(runtime->adstack_overflow_flag == 0, runtime); + QD_TEST_CHECK(local_flag == 0, runtime); // stack_top_primal must point at slot (n - 1) (here: slot 3) when n > 0. QD_TEST_CHECK(stack_top_primal(stack, 4) == stack + sizeof(u64) + 3 * 2 * 4, runtime); @@ -76,19 +85,18 @@ i32 test_stack(RuntimeContext *context) { // * 2 * element_size, which would point into header territory and crash on read). QD_TEST_CHECK(stack_top_primal(stack, 4) == stack + sizeof(u64), runtime); - // Push past capacity: `n` stops at max_num_elements and `adstack_overflow_flag` flips to 1. + // Push past capacity: `n` stops at max_num_elements and the overflow flag flips to 1. for (int i = 0; i < 16; i++) { - stack_push(runtime, stack, 16, 4); + stack_push(runtime, stack, 16, 4, 0); } QD_TEST_CHECK(*(u64 *)stack == 16, runtime); - QD_TEST_CHECK(runtime->adstack_overflow_flag == 0, runtime); - stack_push(runtime, stack, 16, 4); // overflow push + QD_TEST_CHECK(local_flag == 0, runtime); + stack_push(runtime, stack, 16, 4, 0); // overflow push QD_TEST_CHECK(*(u64 *)stack == 16, runtime); - QD_TEST_CHECK(runtime->adstack_overflow_flag == 1, runtime); - // Reset the flag so subsequent tests in the same fixture are not poisoned. - runtime->adstack_overflow_flag = 0; + QD_TEST_CHECK(local_flag == 1, runtime); - delete[] stack; + // Restore the prior flag pointer so subsequent tests in the same fixture are not poisoned by our local slot. + runtime->adstack_overflow_flag_dev_ptr = prev_flag_dev_ptr; return 0; } diff --git a/quadrants/runtime/llvm/runtime_module/runtime.cpp b/quadrants/runtime/llvm/runtime_module/runtime.cpp index 88aa512542..3063022ade 100644 --- a/quadrants/runtime/llvm/runtime_module/runtime.cpp +++ b/quadrants/runtime/llvm/runtime_module/runtime.cpp @@ -580,10 +580,25 @@ struct LLVMRuntime { uint64 error_message_arguments[quadrants_error_message_max_num_arguments]; i32 error_message_lock = 0; i64 error_code = 0; - // Dedicated flag for adstack-overflow-specific errors. Separate from `error_code` so assertions (which set - // error_code=1 and are only surfaced when `compile_config.debug` is on) do not leak through the always-on poll - // that Program::synchronize runs. - i64 adstack_overflow_flag = 0; + // Dedicated overflow signal. Pointer to a 64-bit slot in pinned host memory (CUDA `cuMemAllocHost_v2`, + // HIP `hipHostMalloc`; CPU plain malloc; on this struct stored as the device-mapped address obtained via + // `cuMemHostGetDevicePointer` / HIP equivalent). The kernel-side `stack_push` writes via a system-wide + // atomic OR through this pointer; the host polls the corresponding host-side pointer (cached separately on + // `LlvmRuntimeExecutor`) without any DtoH or sync drain. Required hardware capability is system-scope + // atomics on host-mapped memory: NVIDIA Compute Capability 6.0+ (Pascal+, 2016) and AMD GFX9+ (Vega+, 2017), + // matching the existing pinned-host-scratch H2D-async pattern in `llvm_adstack_lazy_claim.cpp`. Separate from + // `error_code` so assertions (which set error_code=1, gated on `compile_config.debug`) do not leak through the + // always-on overflow poll. nullptr until `materialize_runtime` initialises it; nullptr-guarded in `stack_push`. + i64 *adstack_overflow_flag_dev_ptr = nullptr; + + // Pinned-host slot recording the Program-assigned u32 identity of the FIRST adstack-sizing-info whose + // overflow path fired this Quadrants entry window. Codegen emits a `cmpxchg(0, id)` immediately after the + // OR-1 on `adstack_overflow_flag_dev_ptr`; only the first overflowing thread's id sticks, subsequent + // threads' cmpxchg fails harmlessly. Host reads the slot during the raise to produce a diagnostic that + // names the offending kernel / task / stack via `Program::adstack_sizing_info_registry_`. nullptr until + // `materialize_runtime` allocates the slot. Lives on the same UVA-mapped pinned host page as the flag + // above so a single `mem_alloc_host` call covers both. + i64 *adstack_overflow_task_id_dev_ptr = nullptr; // Combined-heap fields. The codegen single-heap path reads these directly; the split-heap path leaves them untouched // and uses the per-kind fields below. Kept for backward compatibility with kernels that have not yet migrated to the @@ -678,6 +693,8 @@ STRUCT_FIELD(LLVMRuntime, adstack_offsets); STRUCT_FIELD(LLVMRuntime, adstack_max_sizes); STRUCT_FIELD(LLVMRuntime, adstack_row_counters); STRUCT_FIELD(LLVMRuntime, adstack_bound_row_capacities); +STRUCT_FIELD(LLVMRuntime, adstack_overflow_flag_dev_ptr); +STRUCT_FIELD(LLVMRuntime, adstack_overflow_task_id_dev_ptr); // NodeManager of node S (hash, pointer) managers the memory allocation of S_ch // It makes use of three ListManagers. @@ -1239,12 +1256,22 @@ void runtime_retrieve_and_reset_error_code(LLVMRuntime *runtime) { runtime->error_code = 0; } -void runtime_retrieve_and_reset_adstack_overflow(LLVMRuntime *runtime) { - // Paired with the relaxed atomic write in `stack_push`. The host calls this only after the thread pool has - // joined, so strictly no synchronization is required here, but use `__atomic_exchange_n` anyway to keep the - // read/reset symmetric with the write and to avoid annotating the single shared field as half-atomic. - i64 flag = __atomic_exchange_n(&runtime->adstack_overflow_flag, (i64)0, __ATOMIC_RELAXED); - runtime->set_result(quadrants_result_buffer_error_id, flag); +// Publish the device-mapped address of the pinned host slot the host allocated for the adstack overflow flag. +// Called once at materialise_runtime time after the host allocates the slot via `cuMemAllocHost_v2` / `hipHostMalloc` +// / plain malloc and obtains the device-mapped address (CUDA `cuMemHostGetDevicePointer` / HIP equivalent / identity +// on CPU). Subsequent kernel-side `stack_push` reads this pointer to write the overflow signal; the host polls the +// host-side address directly without involving any JIT helper. +void runtime_set_adstack_overflow_flag_dev_ptr(LLVMRuntime *runtime, void *dev_ptr) { + runtime->adstack_overflow_flag_dev_ptr = (i64 *)dev_ptr; +} + +// Companion to `runtime_set_adstack_overflow_flag_dev_ptr`. Called once at materialise_runtime alongside the +// flag setter. The task-id slot lives on the same pinned host page so a single allocation backs both. Codegen +// emits a `cmpxchg(0, baked_id)` against this pointer at the lazy-claim overflow path; only the first +// overflowing thread's id sticks. Host reads the slot during the raise to look up the offending kernel / +// task in `Program::adstack_sizing_info_registry_`. +void runtime_set_adstack_overflow_task_id_dev_ptr(LLVMRuntime *runtime, void *dev_ptr) { + runtime->adstack_overflow_task_id_dev_ptr = (i64 *)dev_ptr; } void runtime_retrieve_error_message(LLVMRuntime *runtime, int i) { @@ -1487,7 +1514,8 @@ void runtime_initialize(Ptr result_buffer, runtime->adstack_row_counters_capacity = 0; runtime->adstack_bound_row_capacities = nullptr; runtime->adstack_bound_row_capacities_capacity = 0; - runtime->adstack_overflow_flag = 0; + runtime->adstack_overflow_flag_dev_ptr = nullptr; + runtime->adstack_overflow_task_id_dev_ptr = nullptr; runtime->temporaries = (Ptr)runtime->allocate_aligned(runtime->runtime_objects_chunk, quadrants_global_tmp_buffer_size, quadrants_page_size); @@ -2569,12 +2597,12 @@ void quadrants_printf(LLVMRuntime *runtime, const char *format, Args &&...args) extern "C" { // local stack operations // The stack index `n` is clamped on read so that overflow (push past capacity) does not let subsequent pops and -// top-accesses underflow it and index far out of bounds. The corresponding stack_push sets -// `runtime->adstack_overflow_flag` and skips the increment instead of trapping, so the host-side launcher -// surfaces the failure as a Python exception rather than killing the process via __builtin_trap. When n == 0 -// (pop-after-overflow underflow path) we return a pointer to slot 0 - an uninitialized-but-in-bounds slot. The -// caller will read garbage from it, but the host raises on `runtime->adstack_overflow_flag` before any such -// value reaches user code. +// top-accesses underflow it and index far out of bounds. The corresponding stack_push writes through +// `runtime->adstack_overflow_flag_dev_ptr` (the device-mapped address of a pinned host slot) and skips the +// increment instead of trapping, so the host-side launcher surfaces the failure as a Python exception rather +// than killing the process via __builtin_trap. When n == 0 (pop-after-overflow underflow path) we return a +// pointer to slot 0 - an uninitialized-but-in-bounds slot. The caller will read garbage from it, but the host +// polls the pinned slot at every Quadrants Python entry and raises before any such value reaches user code. Ptr stack_top_primal(Ptr stack, std::size_t element_size) { auto n = *(u64 *)stack; std::size_t idx = n > 0 ? n - 1 : 0; @@ -2596,20 +2624,45 @@ void stack_pop(Ptr stack) { } } -void stack_push(LLVMRuntime *runtime, Ptr stack, size_t max_num_elements, std::size_t element_size) { +void stack_push(LLVMRuntime *runtime, + Ptr stack, + size_t max_num_elements, + std::size_t element_size, + i64 task_registry_id) { u64 &n = *(u64 *)stack; if (n + 1 > max_num_elements) { // Overflow: the loop has more iterations than the adstack capacity. Skip the push and flip the dedicated - // overflow flag so the host launcher throws at sync. Multiple CPU threads can hit this branch concurrently - // (thread pool dispatch over a multi-element field), so write the sentinel through `__atomic_store_n` with - // relaxed ordering: on x86-64/ARM64 this compiles to a regular naturally-aligned store, but it satisfies the - // C++11 memory model (plain non-atomic writes from multiple threads to the same object are a data race, even - // when every writer stores the same value). The host only reads the flag from `check_adstack_overflow()` - // after the thread pool has joined, so no ordering beyond "happens eventually" is required. - // `locked_task` was avoided because the AMDGPU JIT cannot retarget its host-side machinery - // (`hipErrorNoBinaryForGpu`). Using a separate field (not `error_code`) keeps this check distinct from - // assertion machinery, which is debug-gated. - __atomic_store_n(&runtime->adstack_overflow_flag, (i64)1, __ATOMIC_RELAXED); + // overflow flag in pinned host memory. The host polls the pinned slot at every Quadrants Python entry + // and raises a `QuadrantsAssertionError` with a diagnosis routed through a synchronous sizer that + // distinguishes a Quadrants bug (pre-pass undercount of the bound) from a user-side mutation that bypassed + // tracking (DLPack zero-copy is the typical case; the sizer's freshly-computed required size will exceed + // the cached allocated size in that case). + // + // Relaxed atomic ordering: multiple threads can hit this branch concurrently (CPU thread pool, GPU warp + // divergence) and they all store the same sentinel value, so no inter-thread ordering is required. On + // CPU this compiles to a naturally-aligned store; on CUDA/AMDGPU device kernels with the pointer aimed + // at pinned host memory (`cuMemAllocHost_v2` / `hipHostMalloc` with the device-mapped address obtained + // via `cuMemHostGetDevicePointer` / HIP equivalent) the store is a system-wide atomic on UVA host memory. + // Available on Compute Capability 6.0+ / GFX9+, the same hardware envelope the existing pinned-host + // H2D-async pattern in `llvm_adstack_lazy_claim.cpp` already requires. + // + // Nullptr-guard: `adstack_overflow_flag_dev_ptr` is nullptr until `materialize_runtime` initialises it. + // A kernel running before that (a stale cached kernel, a C++-only test) silently no-ops the overflow + // signal. The runtime cannot raise from device code; this is the safest behavior. + i64 *flag_ptr = runtime->adstack_overflow_flag_dev_ptr; + if (flag_ptr != nullptr) { + __atomic_store_n(flag_ptr, (i64)1, __ATOMIC_RELAXED); + } + // Record task identity in the companion pinned-host slot via cmpxchg(0, registry_id) so the host raise + // site can name the offending kernel + task in the diagnostic message. `task_registry_id == 0` means + // "not registered" (e.g. a deserialised offline-cache task that has not yet been re-registered); skip + // the cmpxchg so the slot stays zero and the host falls through to the generic dual-cause message. + i64 *task_id_ptr = runtime->adstack_overflow_task_id_dev_ptr; + if (task_id_ptr != nullptr && task_registry_id != 0) { + i64 expected = 0; + __atomic_compare_exchange_n(task_id_ptr, &expected, task_registry_id, /*weak=*/false, __ATOMIC_RELAXED, + __ATOMIC_RELAXED); + } return; } n += 1; diff --git a/quadrants/runtime/program_impls/llvm/llvm_program.h b/quadrants/runtime/program_impls/llvm/llvm_program.h index aa4bd8cf05..13848c52e5 100644 --- a/quadrants/runtime/program_impls/llvm/llvm_program.h +++ b/quadrants/runtime/program_impls/llvm/llvm_program.h @@ -169,6 +169,12 @@ class LlvmProgramImpl : public ProgramImpl { } } + void check_adstack_overflow_and_assert() override { + if (!finalizing_) { + runtime_exec_->check_adstack_overflow(); + } + } + LLVMRuntime *get_llvm_runtime() { return runtime_exec_->get_llvm_runtime(); } diff --git a/tests/cpp/program/test_diagnose_adstack_overflow.cpp b/tests/cpp/program/test_diagnose_adstack_overflow.cpp new file mode 100644 index 0000000000..150286c3ac --- /dev/null +++ b/tests/cpp/program/test_diagnose_adstack_overflow.cpp @@ -0,0 +1,119 @@ +#include "gtest/gtest.h" + +#include "quadrants/program/adstack_size_expr_eval.h" +#include "quadrants/program/program.h" + +namespace quadrants::lang { + +// `AdStackCache::register_adstack_sizing_info` + `lookup_adstack_sizing_info` + +// `diagnose_adstack_overflow_message` exercised in isolation. These are pure host-side data-structure +// operations, so the test does not need +// `materialize_runtime` or any device backend. +TEST(DiagnoseAdstackOverflow, RegistryAndLookup) { + Program prog(host_arch()); + + // Two distinct sizing infos. The pointer is just an identity key for `Program::register_*` - the test does not + // deref. + int dummy_a = 0; + int dummy_b = 0; + uint32_t id_a = + prog.adstack_cache().register_adstack_sizing_info(static_cast(&dummy_a), "kernel_a", /*task=*/0, + /*allocated_max_sizes=*/{16, 32}, /*size_exprs=*/{}); + uint32_t id_b = + prog.adstack_cache().register_adstack_sizing_info(static_cast(&dummy_b), "kernel_b", /*task=*/3, + /*allocated_max_sizes=*/{100}, /*size_exprs=*/{}); + EXPECT_NE(id_a, 0u); + EXPECT_NE(id_b, 0u); + EXPECT_NE(id_a, id_b); + + // Idempotent re-registration: same pointer returns the same id (and updates metadata in place). + uint32_t id_a_redo = + prog.adstack_cache().register_adstack_sizing_info(static_cast(&dummy_a), "kernel_a_v2", + /*task=*/2, + /*allocated_max_sizes=*/{8}, /*size_exprs=*/{}); + EXPECT_EQ(id_a, id_a_redo); + auto entry = prog.adstack_cache().lookup_adstack_sizing_info(id_a); + ASSERT_TRUE(entry.has_value()); + EXPECT_EQ(entry->kernel_name, "kernel_a_v2"); + EXPECT_EQ(entry->task_id_in_kernel, 2); + EXPECT_EQ(entry->allocated_max_sizes, std::vector({8})); + + // Lookup with id 0 (sentinel) returns nullopt. + EXPECT_FALSE(prog.adstack_cache().lookup_adstack_sizing_info(0).has_value()); + // Lookup with out-of-range id returns nullopt. + EXPECT_FALSE(prog.adstack_cache().lookup_adstack_sizing_info(static_cast(0xfffffffful)).has_value()); +} + +// Diagnose with an unknown / sentinel id falls back to the generic dual-cause body without crashing. +// The body must mention BOTH causes (DLPack bypass and Quadrants bug) so the user has actionable +// recovery information regardless of whether the runtime captured the offending task identity. +TEST(DiagnoseAdstackOverflow, GenericFallbackOnUnknownId) { + Program prog(host_arch()); + std::string msg = prog.adstack_cache().diagnose_adstack_overflow_message(/*task_id=*/0); + EXPECT_NE(msg.find("DLPack"), std::string::npos); + EXPECT_NE(msg.find("Quadrants bug"), std::string::npos); + // No identity-block prefix when there is nothing to look up. + EXPECT_EQ(msg.find("Offending task"), std::string::npos); +} + +// Diagnose with a registered id includes the kernel name + offload task index + per-stack allocated +// max_sizes in the identity block prefix. +TEST(DiagnoseAdstackOverflow, EnrichedMessageWithIdentity) { + Program prog(host_arch()); + int dummy = 0; + uint32_t id = + prog.adstack_cache().register_adstack_sizing_info(static_cast(&dummy), "compute_grad_diag", + /*task=*/4, + /*allocated_max_sizes=*/{32, 64}, /*size_exprs=*/{}); + std::string msg = prog.adstack_cache().diagnose_adstack_overflow_message(id); + EXPECT_NE(msg.find("Offending task"), std::string::npos); + EXPECT_NE(msg.find("compute_grad_diag"), std::string::npos); + EXPECT_NE(msg.find("offload task #4"), std::string::npos); + EXPECT_NE(msg.find("[32, 64]"), std::string::npos); + // Dual-cause body still present. + EXPECT_NE(msg.find("DLPack"), std::string::npos); + EXPECT_NE(msg.find("Quadrants bug"), std::string::npos); +} + +// The registry copies size_exprs into the entry, so the diagnose path is backend-independent: it +// walks the entry's heap-owned `size_exprs` regardless of whether the original kernel data lives in +// LLVM's `AdStackSizingInfo` or SPIR-V's `AdStackSizingAttribs`. Pin that uniform behaviour by +// registering with one empty SerializedSizeExpr and verifying the rerun reports `?`. +TEST(DiagnoseAdstackOverflow, BackendUniformSizeExprWalk) { + Program prog(host_arch()); + int identity = 0; + std::vector size_exprs(1); // one empty tree + uint32_t id = prog.adstack_cache().register_adstack_sizing_info(static_cast(&identity), "any_kernel", + /*task=*/0, + /*allocated_max_sizes=*/{16}, std::move(size_exprs)); + std::string msg = prog.adstack_cache().diagnose_adstack_overflow_message(id); + EXPECT_NE(msg.find("Offending task"), std::string::npos); + EXPECT_NE(msg.find("any_kernel"), std::string::npos); + // Empty `nodes` triggers the rerun-loop skip; with no resolvable stack, the rerun line is omitted and the + // dual-cause body fronts the diagnostic instead. + EXPECT_EQ(msg.find("Synchronous sizer rerun"), std::string::npos); + EXPECT_NE(msg.find("DLPack"), std::string::npos); + EXPECT_NE(msg.find("Quadrants bug"), std::string::npos); +} + +// `update_adstack_sizing_info_size_exprs` overwrites just the size_exprs without disturbing the rest +// of the entry. Used by the LLVM launcher on every launch to keep the registry in sync with the live +// `OffloadedTask::ad_stack`. +TEST(DiagnoseAdstackOverflow, UpdateSizeExprsRefreshesEntry) { + Program prog(host_arch()); + int identity = 0; + uint32_t id = prog.adstack_cache().register_adstack_sizing_info(static_cast(&identity), "k", + /*task=*/0, /*allocated_max_sizes=*/{8}, + /*size_exprs=*/{}); + // Empty registry size_exprs => no rerun line. + EXPECT_EQ(prog.adstack_cache().diagnose_adstack_overflow_message(id).find("Synchronous sizer rerun"), + std::string::npos); + prog.adstack_cache().update_adstack_sizing_info_size_exprs(id, std::vector(1)); + // After refresh, the rerun walks the new size_exprs but every leaf is unresolvable (empty tree), so the line + // is still omitted - same suppression as the BackendUniformSizeExprWalk case above. The `update` change is + // observable via the dual-cause body still fronting the message; we already pinned that path elsewhere. + EXPECT_EQ(prog.adstack_cache().diagnose_adstack_overflow_message(id).find("Synchronous sizer rerun"), + std::string::npos); +} + +} // namespace quadrants::lang diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 86d8ea33ec..ba9fa5fa4a 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -5,11 +5,13 @@ import subprocess import sys import textwrap +from contextlib import nullcontext import numpy as np import pytest import quadrants as qd +from quadrants.lang.exception import QuadrantsAssertionError from quadrants.lang.misc import is_extension_supported from tests import test_utils @@ -1069,7 +1071,7 @@ def test_adstack_overflow_raises(): # On LLVM the runtime raises QuadrantsAssertionError (subclass of AssertionError) from # check_adstack_overflow; on SPIR-V the gfx runtime raises RuntimeError via QD_ERROR. We accept either, # matching only the message prefix. - with pytest.raises((AssertionError, RuntimeError), match=r"[Aa]dstack overflow"): + with pytest.raises(QuadrantsAssertionError, match=r"[Aa]dstack overflow"): compute.grad() qd.sync() @@ -1081,7 +1083,7 @@ def test_adstack_overflow_flag_resets_after_catch(): # stale overflow exception every time they sync after the first one, which makes diagnosis and recovery # impossible. `debug=True` keeps the per-push bounds check live. compute, _, _ = _overflowing_compute() - with pytest.raises((AssertionError, RuntimeError), match=r"[Aa]dstack overflow"): + with pytest.raises(QuadrantsAssertionError, match=r"[Aa]dstack overflow"): compute.grad() qd.sync() # No new grad launch here - the flag must already be back to zero. @@ -1174,40 +1176,46 @@ def test_adstack_overflow_multithreaded(): # keeps the per-push bounds check live (release-build codegen elides it - see `test_adstack_overflow_raises` # for the rationale). compute, _, _ = _overflowing_compute(n_elements=16) - with pytest.raises((AssertionError, RuntimeError), match=r"[Aa]dstack overflow"): + with pytest.raises(QuadrantsAssertionError, match=r"[Aa]dstack overflow"): compute.grad() qd.sync() -def test_adstack_overflow_during_teardown_does_not_abort(tmp_path): +@pytest.mark.parametrize("force_sync", [False, True]) +def test_adstack_overflow_caught_then_clean_teardown(tmp_path, force_sync): # This test runs the kernel in a child process (not via `@test_utils.test`, which iterates arches), so it # cannot rely on the decorator's `require=qd.extension.adstack` skip. Guard manually: skip if the CPU backend # was not built with the adstack extension, matching what the sibling overflow tests get from the decorator. if not is_extension_supported(qd.cpu, qd.extension.adstack): pytest.skip("adstack extension not available on cpu") - # If a user launches an overflowing grad kernel and never calls `qd.sync()` before the process exits, the - # adstack-overflow flag is still set when Python interpreter teardown invokes `Program::finalize()`. The two - # teardown syncs inside `Program::finalize()` must not re-raise a `QuadrantsAssertionError` into the - # destructor path - doing so would terminate the process with `std::terminate()` instead of returning a clean - # exit code. A subprocess runs the overflowing-grad kernel without calling `qd.sync()` at all and exits; this - # test asserts that the child returns with exit code 0 rather than SIGABRT (-6) or any other non-zero code. + # Pins the per-launch-raise + clean-teardown contract for `Program::finalize()`. The per-launch + # `check_adstack_overflow_and_assert()` poll wired into `Program::launch_kernel` surfaces an overflow at + # the very next kernel-launch entry on synchronous backends (CPU). On async backends (CUDA / AMDGPU / + # Metal / Vulkan) the kernel may still be in flight when `launch_kernel` returns, so the post-launch poll + # reads a not-yet-set flag - the overflow is then surfaced at the next `qd.sync()` via the post-drain + # check in `LlvmProgramImpl::synchronize_and_assert` (or the host-mapped readback in + # `GfxRuntime::synchronize`). The `force_sync` parametrisation toggles whether the user issues an + # explicit `qd.sync()`. Either way the teardown contract holds: the two teardown `synchronize()` calls + # inside `Program::finalize()` re-enter the LLVM `check_adstack_overflow_and_assert` path, and + # `LlvmProgramImpl::pre_finalize()` must have set `finalizing_ = true` early enough that the per-launch + # poll AND the `synchronize_and_assert` poll BOTH short-circuit during the destructor. If either path + # re-raised, the destructor would `std::terminate()` instead of returning a clean exit code (-6 / SIGABRT + # on macOS, std::terminate's _Exit on linux). The subprocess asserts the child returns with exit code 0. # - # Internal details: `Program::finalize()` invokes `program_impl_->pre_finalize()` before the two teardown - # `synchronize()` calls. `LlvmProgramImpl::pre_finalize()` sets `finalizing_ = true` so - # `LlvmProgramImpl::synchronize()` short-circuits `check_adstack_overflow()`. Note the flag must be set - # *before* those syncs run - setting it only inside `LlvmProgramImpl::finalize()` (which is dispatched after - # them) is too late. The subprocess is launched from a temp file because `python -c ""` breaks - # Quadrants' kernel source-inspect (`getsourcelines` cannot find the source of an inlined `-c` string); the - # grad call is deliberately left unsynced so this is the teardown path, not the user-catch path. - # `debug=True` is required for the same reason as `test_adstack_overflow_raises`: release-build LLVM codegen - # elides the per-push bounds check on the premise the sizer's bound is tight, so a manually-misconfigured - # `ad_stack_size=32` kernel only flips the runtime overflow flag in debug mode. Without the flag set there - # is no flag for the teardown guard to swallow, and the bug this test pins (re-raise during destructor path - # leading to `std::terminate()`) cannot trigger. + # Internal details: the subprocess is launched from a temp file because `python -c ""` breaks + # Quadrants' kernel source-inspect (`getsourcelines` cannot find the source of an inlined `-c` string). + # `debug=True` is required because release-build LLVM codegen elides the per-push bounds check on the + # premise the sizer's bound is tight, so a manually-misconfigured `ad_stack_size=32` kernel only flips + # the runtime overflow flag in debug mode. Without the flag set there is no flag for the teardown + # guard to swallow, and the bug this test pins cannot trigger. child_script = textwrap.dedent( - """ + f""" + from contextlib import nullcontext + + import pytest import quadrants as qd + from quadrants.lang.exception import QuadrantsAssertionError qd.init(arch=qd.cpu, ad_stack_experimental_enabled=True, ad_stack_size=32, debug=True) @@ -1229,9 +1237,25 @@ def compute(): compute() y.grad[None] = 1.0 x.grad[0] = 0.0 - compute.grad() - # Intentionally no qd.sync() and no try/except here: the adstack-overflow flag is left set when the - # process exits, so teardown must swallow it via the `finalizing_` guard rather than re-raising. + + # CPU is the only arch the child runs on for now; synchronous-backend semantics apply. + # The per-launch poll wired into `Program::launch_kernel` surfaces the overflow at `compute.grad()`. + # On a future async-arch parametrisation, flip `is_sync_backend` and the qd.sync() branch becomes + # the raising path while compute.grad() drops to nullcontext. + is_sync_backend = True + force_sync = {force_sync} + + def raises_overflow(): + return pytest.raises(QuadrantsAssertionError, match=r"[Aa]dstack overflow") + + with raises_overflow() if is_sync_backend else nullcontext(): + compute.grad() + if force_sync: + with raises_overflow() if not is_sync_backend else nullcontext(): + qd.sync() + # Process exits without `qd.sync()` (when force_sync=False). Teardown's two `synchronize()` calls + # plus their per-launch polls must short-circuit on `finalizing_`; otherwise the destructor + # double-raises and the process exits non-zero / SIGABRTs. """ ) script_path = tmp_path / "overflow_teardown_child.py" @@ -1248,6 +1272,148 @@ def compute(): ) +@pytest.mark.needs_torch +@test_utils.test(require=qd.extension.adstack, debug=True) +def test_adstack_overflow_diagnostic_and_auto_recovery(): + import torch + + # Cross-backend regression for the always-on overflow detection + diagnostic + auto-recovery + # contract shipped on this branch. The kernel's inner trip count is bounded by `n[0]`, an `int32` + # ndarray. The per-task adstack metadata cache invalidation tracks `Ndarray.write` / + # `Ndarray.fill` via `Program::ndarray_data_gen_` - mutations that route through Quadrants APIs + # invalidate cleanly. DLPack zero-copy mutations (`.to_torch(copy=False)`) bypass that tracking, + # so the cache holds a stale `max_size`; the next reverse launch overflows. + # + # The contract pinned here: + # 1. The first launch with `n[0] = 2` populates the cache with `max_size = 2`. + # 2. A DLPack-backed torch view writes `n[0] = 64`. Quadrants's gen counter is NOT bumped. + # 3. The next reverse launch reads cached `max_size = 2`, pushes 64, overflows. The host poll + # raises with the enriched diagnostic naming kernel + offload-task index. The raise site + # ALSO bulk-invalidates the adstack-sizer caches on its way out. + # 4. The user catches the exception. They do NOT need to manually adjust `ad_stack_size`. On the NEXT reverse + # launch, the sizer reruns from scratch (cache invalidated), reads the mutated `n[0] = 64`, sizes capacity + # to 64, and the kernel runs to completion with the correct gradient. + # + # This auto-recovery contract is what lets the user's training-loop code recover from a + # transient cache-staleness window without per-iteration retries: the offending data has + # already been mutated in place; once the cache reflects it, every subsequent run just works. + n = qd.ndarray(qd.i32, shape=(1,)) + x = qd.field(qd.f32, shape=(1,), needs_grad=True) + y = qd.field(qd.f32, shape=(), needs_grad=True) + + @qd.kernel + def compute(n_arr: qd.types.ndarray(dtype=qd.i32, ndim=1)): + for i in x: + v = x[i] + for _ in range(n_arr[0]): + y[None] += qd.sin(v) + v = v + 1.0 + + # Step 1: small `n[0]`, kernel runs cleanly, cache populated with `max_size = 2`. + n[0] = 2 + x[0] = 0.1 + y[None] = 0.0 + compute(n) + y.grad[None] = 1.0 + x.grad[0] = 0.0 + compute.grad(n) + qd.sync() + + # Step 1.5: Python-side mutation through `Ndarray.write` (`n[0] = 8`). The setter routes through Quadrants's + # tracking and bumps `ndarray_data_gen_` for the bound DeviceAllocation, so the per-task adstack metadata + # cache invalidates cleanly: the next launch reruns the sizer with `n[0] = 8`, sizes capacity to 8, and the + # kernel runs to completion without raising. This pins the clean-invalidation contract on every backend + # (no DLPack involvement, no overflow, no recovery exception). + n[0] = 8 + y[None] = 0.0 + compute(n) + y.grad[None] = 1.0 + x.grad[0] = 0.0 + compute.grad(n) + qd.sync() + expected_after_clean_grow = sum(math.cos(0.1 + k) for k in range(8)) + assert x.grad[0] == pytest.approx(expected_after_clean_grow, rel=1e-4) + + # Reset state for the DLPack-bypass scenario: bring `n[0]` back down to 2 through the Quadrants-tracked + # setter so the next cached `max_size` is the small value the bypass mutation will outgrow. + n[0] = 2 + y[None] = 0.0 + compute(n) + y.grad[None] = 1.0 + x.grad[0] = 0.0 + compute.grad(n) + qd.sync() + + # The DLPack-bypass scenario below requires `to_torch(copy=False)` which is unsupported on + # Vulkan because Quadrants and torch do not currently share a command queue there + # (`_can_zerocopy_field` returns false on i32 ndarrays and the export raises + # `Zero-copy not available for arch=vulkan, dtype=i32`). Steps 1 + 4 above already verified + # the cleanly-running path; bail out before the bypass-mutation portion on Vulkan. + if qd.lang.impl.get_runtime().prog.config().arch == qd.vulkan: + return + + # Step 2: DLPack-bypass mutation. `Ndarray.write` would have bumped `ndarray_data_gen_` and + # invalidated the cache cleanly; `to_torch(copy=False)` shares storage with no Quadrants hook, + # so the cache sees no change. On Metal `to_torch(copy=False)` returns an `mps:0` tensor and + # writes through it dispatch asynchronously through Metal Performance Shaders; an explicit + # `torch.mps.synchronize()` is required to flush those writes to the shared buffer the + # Quadrants device kernel reads from. Without it the next Quadrants launch sees the stale + # `n[0] = 2` and the overflow detection misses entirely. CPU / CUDA / AMDGPU paths do not + # need the equivalent on this code path because their `to_torch` returns a tensor on a + # device where writes are coherent without an additional sync. + n_view = n.to_torch(copy=False) + n_view[0] = 64 + if qd.lang.impl.get_runtime().prog.config().arch == qd.metal: + torch.mps.synchronize() + qd.sync() + + # Step 3: next reverse launch may overflow. On backends with a stale-cache shortcut (LLVM-GPU + # `try_llvm_per_task_ad_stack_cache_hit`, SPIR-V `try_per_task_ad_stack_cache_hit`) the cached + # `max_size = 2` is reused because `ndarray_data_gen` has not been bumped, the kernel pushes 64, + # and the host poll raises with the enriched diagnostic naming kernel + offload-task index. On + # LLVM-CPU the host-eval branch always re-evaluates the size expression per launch via + # `try_size_expr_cache_hit`, which observes the live ndarray read and self-invalidates on + # mismatch - that path never raises here, so the second compute.grad call returns cleanly. The + # `pytest.raises if shortcut else nullcontext` pattern handles both paths uniformly without arch-narrowing the + # test. + y[None] = 0.0 + compute(n) + y.grad[None] = 1.0 + x.grad[0] = 0.0 + backend_uses_per_task_cache_shortcut = qd.lang.impl.get_runtime().prog.config().arch != qd.cpu + raises_overflow = pytest.raises(QuadrantsAssertionError, match=r"[Aa]dstack overflow") + with raises_overflow if backend_uses_per_task_cache_shortcut else nullcontext() as exc_info: + compute.grad(n) + qd.sync() + if exc_info is not None: + msg = str(exc_info.value) + # When the inner range is bounded by an ndarray read, the user sees the actual mutated size in the + # error (e.g. allocated=[1], required=[64]) and a recovery flow pointing at the tensor mutation + # performed outside Quadrants's tracking. The generic "this might also be a Quadrants bug" + # alternative only appears when the diagnostic cannot pin the cause down. + assert "DLPack" in msg, f"missing DLPack-bypass cause hint in: {msg}" + assert "Restart" in msg, f"missing recovery flow in: {msg}" + assert "Offending task" in msg, f"missing identity block in: {msg}" + assert "compute" in msg, f"missing kernel name in: {msg}" + assert "Synchronous sizer rerun: required max_size = [" in msg, f"missing sync-sizer-rerun line in: {msg}" + + # Step 4: auto-recovery. If the previous launch overflowed, the raise site bulk-invalidated the + # adstack-sizer caches when the synchronous sizer rerun confirmed a stale-cache cause. The next + # reverse launch reruns the sizer from scratch, reads `n[0] = 64`, sizes capacity to 64, and + # the kernel runs cleanly with no second overflow. Either way (stale-cache backend that + # recovered or auto-invalidating CPU backend that never overflowed), the closed-form gradient + # below is the contract for every backend. + y[None] = 0.0 + compute(n) + y.grad[None] = 1.0 + x.grad[0] = 0.0 + compute.grad(n) + qd.sync() + # Closed-form gradient sanity: y = sum_{k=0..n-1} sin(x + k), so dy/dx = sum cos(x + k). + expected = sum(math.cos(0.1 + k) for k in range(64)) + assert x.grad[0] == pytest.approx(expected, rel=1e-4) + + @pytest.mark.parametrize("n_iter", [30, 100]) @test_utils.test(require=qd.extension.adstack) def test_adstack_near_capacity(n_iter): @@ -2151,6 +2317,107 @@ 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_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[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) + + if trip_count_source == "qd_ndarray": + n_obj = qd.ndarray(qd.i32, shape=(N,)) + + @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 + call_compute() + loss.grad[None] = 1.0 + for i in range(N_X): + x.grad[i] = 0.0 + 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) + + set_n(16) + for i in range(N_X): + x[i] = 0.1 + loss[None] = 0.0 + call_compute() + loss.grad[None] = 1.0 + for i in range(N_X): + x.grad[i] = 0.0 + call_compute_grad() + 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