[Fixup] Phase 1: enriched adstack overflow diagnostic - registry-id cmpxchg + kernel/task name in raise message

Author: duburcqa

quadrants/codegen/llvm/codegen_llvm.cpp

@@ -1940,6 +1940,16 @@ 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<OffloadedTask>(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) {
+    uint32_t id = prog->register_adstack_sizing_info(static_cast<const void *>(&current_task->ad_stack), kernel_name,
+                                                     task_codegen_id, /*allocated_max_sizes=*/{});
+    current_task->ad_stack.registry_id = id;
+  }
 
   for (auto &arg : func->args()) {
     kernel_args.push_back(&arg);
@@ -2010,6 +2020,20 @@ void TaskCodeGenLLVM::finalize_offloaded_task_function() {
       current_task->arr_writes.erase(std::unique(current_task->arr_writes.begin(), current_task->arr_writes.end()),
                                      current_task->arr_writes.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<int> 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<int>(a.max_size_compile_time));
+      }
+      uint32_t id = prog->register_adstack_sizing_info(static_cast<const void *>(&current_task->ad_stack), kernel_name,
+                                                       task_codegen_id, std::move(allocated_max_sizes));
+      current_task->ad_stack.registry_id = id;
+    }
   }
 
   // entry_block should jump to the body after all allocas are inserted
@@ -2464,10 +2488,12 @@ void TaskCodeGenLLVM::emit_ad_stack_row_claim_llvm() {
   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. The
-  // condition is hoisted to a structured if so the not-overflowing fast path skips the atomic entirely - one
-  // function call to fetch the flag pointer plus one CreateICmpUGT comparison, the same compare we already
-  // emitted for the clamp.
+  // 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);
@@ -2479,6 +2505,17 @@ void TaskCodeGenLLVM::emit_ad_stack_row_claim_llvm() {
     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<uint64_t>(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);
@@ -2742,7 +2779,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<uint64_t>(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);

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<StaticAdStackBoundExpr> 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,

quadrants/program/program.cpp

@@ -2,6 +2,8 @@
 
 #include "program.h"
 
+#include <mutex>
+
 #include "quadrants/program/adstack_size_expr_eval.h"
 #include "quadrants/ir/statements.h"
 #include "quadrants/ir/type_factory.h"
@@ -239,6 +241,76 @@ void Program::synchronize_and_assert() {
   program_impl_->synchronize_and_assert();
 }
 
+uint32_t Program::register_adstack_sizing_info(const void *ad_stack_ptr,
+                                               const std::string &kernel_name,
+                                               int task_id_in_kernel,
+                                               std::vector<int> allocated_max_sizes) {
+  std::lock_guard<std::mutex> lk(adstack_sizing_info_registry_mutex_);
+  // Idempotent re-registration: same `ad_stack_ptr` yields the same id across re-compiles. The pointer is
+  // task-stable for the kernel's lifetime; if a kernel is destroyed and a new one happens to hit the same
+  // address, the previous entry's metadata is overwritten - acceptable because the diagnostic message reads
+  // metadata at the same instant as the lookup, and the new entry's `kernel_name` / `allocated_max_sizes`
+  // describe the live kernel.
+  auto it = adstack_sizing_info_id_by_ptr_.find(ad_stack_ptr);
+  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);
+    return it->second;
+  }
+  uint32_t id = static_cast<uint32_t>(adstack_sizing_info_registry_.size());
+  AdStackSizingInfoEntry entry;
+  entry.ad_stack_ptr = ad_stack_ptr;
+  entry.kernel_name = kernel_name;
+  entry.task_id_in_kernel = task_id_in_kernel;
+  entry.allocated_max_sizes = std::move(allocated_max_sizes);
+  adstack_sizing_info_registry_.push_back(std::move(entry));
+  adstack_sizing_info_id_by_ptr_.emplace(ad_stack_ptr, id);
+  return id;
+}
+
+const Program::AdStackSizingInfoEntry *Program::lookup_adstack_sizing_info(uint32_t id) const {
+  std::lock_guard<std::mutex> lk(adstack_sizing_info_registry_mutex_);
+  if (id == 0 || id >= adstack_sizing_info_registry_.size()) {
+    return nullptr;
+  }
+  return &adstack_sizing_info_registry_[id];
+}
+
+std::string Program::diagnose_adstack_overflow_message(uint32_t task_id) const {
+  std::string identity_block;
+  if (task_id != 0) {
+    const auto *entry = lookup_adstack_sizing_info(task_id);
+    if (entry != nullptr) {
+      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";
+    }
+  }
+  return identity_block +
+         "Two possible causes:\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"
+         "Note: kernel state may be inconsistent post-overflow; do not retry the same step without "
+         "addressing the cause and restarting from a clean state.";
+}
+
 StreamSemaphore Program::flush() {
   return program_impl_->flush();
 }

quadrants/program/program.h

@@ -3,10 +3,13 @@
 #pragma once
 
 #include <functional>
+#include <mutex>
 #include <optional>
 #include <atomic>
 #include <stack>
 #include <shared_mutex>
+#include <string>
+#include <vector>
 
 #define QD_RUNTIME_HOST
 #include "quadrants/ir/frontend_ir.h"
@@ -214,6 +217,29 @@ class QD_DLL_EXPORT Program {
     return *adstack_cache_;
   }
 
+  // Identity registry for adstack-sizer info. 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.
+  struct AdStackSizingInfoEntry {
+    const void *ad_stack_ptr{nullptr};
+    std::string kernel_name;
+    int task_id_in_kernel{0};
+    std::vector<int> allocated_max_sizes;
+  };
+  uint32_t register_adstack_sizing_info(const void *ad_stack_ptr,
+                                        const std::string &kernel_name,
+                                        int task_id_in_kernel,
+                                        std::vector<int> allocated_max_sizes);
+  const AdStackSizingInfoEntry *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).
+  // Returns the dual-cause message body to embed in the `QuadrantsAssertionError` raised at the poll site.
+  std::string diagnose_adstack_overflow_message(uint32_t task_id) const;
+
   /**
    * Destroys a new SNode tree.
    *
@@ -352,6 +378,14 @@ class QD_DLL_EXPORT Program {
   // Adstack caching state (per-task metadata, bytecode, size-expr results, generation counters). All adstack-specific
   // surface lives in `program/adstack_size_expr_eval.{h,cpp}`; routed through `adstack_cache()` getter.
   std::unique_ptr<AdStackCache> adstack_cache_;
+
+  // Adstack-sizing-info identity registry. See `register_adstack_sizing_info`. 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". Reverse lookup map keyed by `ad_stack_ptr` keeps `register_adstack_sizing_info`
+  // idempotent across re-launches of the same kernel.
+  std::vector<AdStackSizingInfoEntry> adstack_sizing_info_registry_{AdStackSizingInfoEntry{}};
+  std::unordered_map<const void *, uint32_t> adstack_sizing_info_id_by_ptr_;
+  mutable std::mutex adstack_sizing_info_registry_mutex_;
   std::stack<int> free_snode_tree_ids_;
 
   std::vector<std::unique_ptr<Function>> functions_;

quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp

@@ -639,25 +639,29 @@ void LlvmRuntimeExecutor::check_adstack_overflow() {
     return;
   }
   int64_t flag = __atomic_exchange_n(adstack_overflow_flag_host_ptr_, (int64_t)0, __ATOMIC_RELAXED);
-  if (flag != 0) {
-    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. Two "
-        "possible causes:\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"
-        "Note: kernel state may be inconsistent post-overflow; do not retry the same step without "
-        "addressing the cause and restarting from a clean state.");
+  if (flag == 0) {
+    return;
+  }
+  // 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 = __atomic_exchange_n(adstack_overflow_task_id_host_ptr_, (int64_t)0, __ATOMIC_RELAXED);
+    task_id = static_cast<uint32_t>(recorded);
   }
+  Program *prog = (program_impl_ != nullptr) ? program_impl_->program : nullptr;
+  std::string diagnostic = (prog != nullptr) ? prog->diagnose_adstack_overflow_message(task_id)
+                                             : std::string(
+                                                   "Adstack overflow: a reverse-mode autodiff kernel pushed more "
+                                                   "elements than the adstack capacity 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,