[Perf] Drop redundant per-launch update_adstack_sizing_info_size_exprs + use find() over operator[] in arr_reads kNone bump

Author: duburcqa

docs/source/user_guide/autodiff.md

@@ -289,7 +289,7 @@ 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 formula's inputs 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 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.
 

quadrants/program/adstack_size_expr_eval.cpp

@@ -1869,10 +1869,9 @@ void bump_writes_for_kernel_llvm(Program *prog,
     // 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) {
-      if (static_cast<size_t>(arg_id) >= ctx->device_allocation_type.size()) {
-        continue;
-      }
-      if (ctx->device_allocation_type[arg_id] != LaunchContextBuilder::DevAllocType::kNone) {
+      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);
@@ -1910,10 +1909,9 @@ void bump_writes_for_kernel_llvm(Program *prog,
   // 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) {
-      if (static_cast<size_t>(arg_id) >= ctx->device_allocation_type.size()) {
-        continue;
-      }
-      if (ctx->device_allocation_type[arg_id] != LaunchContextBuilder::DevAllocType::kNone) {
+      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);

quadrants/runtime/llvm/llvm_adstack_lazy_claim.cpp

@@ -691,17 +691,6 @@ std::size_t LlvmRuntimeExecutor::publish_adstack_metadata(const AdStackSizingInf
   if (n_stacks == 0 || num_threads == 0) {
     return 0;
   }
-  // Refresh the registry entry's `size_exprs` snapshot from the live (post-push_back, stable)
-  // ad_stack. The codegen-time registration filled `size_exprs` from `&current_task->ad_stack` but
-  // that pointer's underlying `unique_ptr<OffloadedTask>` heap was freed by the launcher's
-  // `current_task = nullptr` after `offloaded_tasks.push_back(*current_task)`; the registry's
-  // snapshot is now backed by its own heap-owned copy, so this refresh just keeps it in sync with
-  // any post-codegen mutations a backend might apply to `ad_stack.size_exprs`. Idempotent and
-  // bounded-cost: copies one `std::vector<SerializedSizeExpr>` per launch per task with adstacks.
-  if (program_impl_ != nullptr && program_impl_->program != nullptr && ad_stack.registry_id != 0) {
-    program_impl_->program->adstack_cache().update_adstack_sizing_info_size_exprs(ad_stack.registry_id,
-                                                                                  ad_stack.size_exprs);
-  }
   auto align_up_8 = [](std::size_t n) -> std::size_t { return (n + 7u) & ~std::size_t{7u}; };
   // Allocate / grow the two device-side metadata arrays. Capacity is in u64 entries, kept at or above n_stacks.
   // On GPU these buffers are written exclusively by the device-side sizer kernel (`runtime_eval_adstack_size_expr`);