[Perf] Cache adstack-sizer metadata per task across SPIR-V + LLVM-GPU; per-snode / DeviceAllocation invalidation

Author: duburcqa

quadrants/codegen/llvm/codegen_llvm.cpp

@@ -7,6 +7,8 @@
 #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/extension.h"
@@ -1975,6 +1977,39 @@ 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 ((static_cast<uint32_t>(kv.second) & static_cast<uint32_t>(irpass::ExternalPtrAccess::WRITE)) == 0) {
+          continue;
+        }
+        if (!kv.first.empty()) {
+          current_task->arr_writes.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());
+    }
   }
 
   // entry_block should jump to the body after all allocas are inserted

quadrants/codegen/llvm/llvm_compiled_data.h

@@ -95,12 +95,26 @@ 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<int> 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<int> arr_writes;
+
   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);
 };
 
 struct LLVMCompiledTask {

quadrants/codegen/spirv/kernel_utils.h

@@ -212,6 +212,15 @@ struct TaskAttributes {
   };
   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<int> snode_writes;
+
   static std::string buffers_name(BufferInfo b);
 
   std::string debug_string() const;
@@ -222,7 +231,8 @@ struct TaskAttributes {
             task_type,
             buffer_binds,
             range_for_attribs,
-            ad_stack);
+            ad_stack,
+            snode_writes);
 };
 
 /**

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"
 
@@ -198,6 +200,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_);