[Caching] Ref 7: CPU per-task artifact cache (ORC object serialize/load)

docs/source/user_guide/init_options.md

@@ -13,7 +13,7 @@ Setting `offline_cache=False` is intended to emulate cold-start, i.e. a fresh Py
 When `offline_cache=True`, compilation artifacts persist on disk under `offline_cache_file_path` (default `~/.cache/quadrants/qdcache`), so a later Python process reuses them instead of recompiling. Setting `offline_cache=False` (or `QD_OFFLINE_CACHE=0`) forces a cold start: Quadrants recompiles kernels and neither reads nor writes its own on-disk cache. (On CUDA the driver keeps its own separate cache of compiled GPU code at `~/.nv/ComputeCache` that this flag does not disable; `offline_cache=False` only stops that cache from serving results across runs. Set `CUDA_CACHE_DISABLE=1` to turn it off entirely.)
 
 The flag gates two disk tiers, not just the whole-kernel one:
-- On CUDA and AMDGPU it also enables a **per-task artifact cache**: a kernel is compiled as several *tasks* (roughly one per parallel loop), so editing one task reuses the *other* tasks' compiled code across processes instead of recompiling the whole kernel.
+- On CUDA, AMDGPU, and CPU it also enables a **per-task artifact cache**: a kernel is compiled as several *tasks* (roughly one per parallel loop), so editing one task reuses the *other* tasks' compiled code across processes instead of recompiling the whole kernel.
 - There is no eviction policy for these per-task artifacts yet, so this cache grows over time; wipe `offline_cache_file_path` occasionally if it gets large.
 
 The separate source-level cache used by [fastcache](./fastcache.md) kernels is controlled by `src_ll_cache` (on by default), not by `offline_cache`; with `offline_cache=False` it still writes its own bookkeeping files to disk, so set `src_ll_cache=False` as well to stop that too.

quadrants/codegen/codegen.cpp

@@ -72,11 +72,11 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
   const int n = (int)offloads.size();
   std::vector<std::unique_ptr<LLVMCompiledTask>> data(n);
 
-  // Cross-process per-task artifact cache: a hit skips a task's compilation and carries cached code plus launch
-  // metadata to the launcher. An empty dir disables the tier.
+  // Per-task artifact cache: a hit skips a task's compilation and carries cached code to the launcher. Empty dir = off.
   const std::string art_dir = pertask_artifact_dir_ref();
-  const bool artifact_tier =
-      (compile_config_.arch == Arch::cuda || compile_config_.arch == Arch::amdgpu) && !art_dir.empty();
+  const bool per_task_cache_enabled = (compile_config_.arch == Arch::cuda || compile_config_.arch == Arch::amdgpu ||
+                                       arch_is_cpu(compile_config_.arch)) &&
+                                      !art_dir.empty();
   const PerTaskArtifactCache artifact_cache(art_dir);
   const DeviceCapabilityConfig pertask_caps = prog->get_device_caps();
   std::vector<std::string> pertask_keys(n);
@@ -97,11 +97,12 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
       //  - external-func: the body has the so/bc path + name, not its contents, so a stale update keeps the key;
       //  - real-func (FuncCallStmt): the printer emits only the callee name, but codegen inlines its body;
       //  - mem_access_opt (BLS/read-only hints): serialized from an unordered_map, so its order varies by process.
-      bool artifact_eligible = artifact_tier && offload->as<OffloadedStmt>()->mem_access_opt.get_all().empty();
-      if (artifact_eligible) {
-        irpass::analysis::gather_statements(offload.get(), [&artifact_eligible](Stmt *s) {
+      bool eligible_for_per_task_cache =
+          per_task_cache_enabled && offload->as<OffloadedStmt>()->mem_access_opt.get_all().empty();
+      if (eligible_for_per_task_cache) {
+        irpass::analysis::gather_statements(offload.get(), [&eligible_for_per_task_cache](Stmt *s) {
           if (s->is<AdStackAllocaStmt>() || s->is<ExternalFuncCallStmt>() || s->is<FuncCallStmt>()) {
-            artifact_eligible = false;
+            eligible_for_per_task_cache = false;
           }
           return false;
         });
@@ -110,7 +111,7 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
       // Key on `#index` too: the module bakes the task's kernel-wide index into its symbol names, so two
       // byte-identical tasks at different indices must not alias (they would collide at link).
       std::string cache_key;
-      if (artifact_eligible) {
+      if (eligible_for_per_task_cache) {
         cache_key = get_hashed_per_task_cache_key(compile_config_, pertask_caps, offload->as<OffloadedStmt>(), kernel) +
                     "#" + std::to_string(i);
         // Under kernel profiling, drop the name-free cross-kernel aliasing: the artifact carries OffloadedTask::name,
@@ -141,13 +142,12 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
   }
   worker.flush();
 
-  // Build one self-contained artifact per task before the whole-module link consumes `data`.
-  //
-  // AMDGPU links each task with a separate `ld.lld`, so the per-task path only pays off when the tier is on; with it
-  // off it stays on the single whole-module link. CUDA's per-task load is cheap, so it always takes the per-task path.
+  // Build one artifact per task before the whole-module link consumes `data`. CUDA always takes this path; AMDGPU and
+  // CPU only when offline_cache is enabled, since their per-task link/serialize only pays off then.
   std::vector<PerConstructArtifact> per_construct_artifacts;
   const bool build_per_construct_artifacts =
-      compile_config_.arch == Arch::cuda || (compile_config_.arch == Arch::amdgpu && artifact_tier);
+      compile_config_.arch == Arch::cuda ||
+      ((compile_config_.arch == Arch::amdgpu || arch_is_cpu(compile_config_.arch)) && per_task_cache_enabled);
   if (build_per_construct_artifacts) {
     for (int i = 0; i < n; i++) {
       if (!data[i])
@@ -176,8 +176,8 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
     }
   }
 
-  // A cross-process hit leaves a task with no module, so skip the whole-kernel link; the launcher assembles the
-  // backend module from the per-task artifacts. Still concatenate every task's metadata into `tasks`.
+  // A cache hit leaves a task with no module, so skip the whole-kernel link and let the launcher assemble from the
+  // per-task artifacts. Still gather every task's metadata into `tasks`.
   const bool code_only_tasks = std::any_of(data.begin(), data.end(), [](const auto &d) { return d && !d->module; });
   LLVMCompiledKernel llvm_compiled_kernel;
   if (code_only_tasks) {
@@ -207,7 +207,7 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
   }
   // Record per-task reuse counts for PerOffloadCacheObservations.tasks_* (read back by the compilation manager). Only
   // when the tier ran, so non-CUDA / cache-off compiles keep the -1 sentinel instead of a misleading 0.
-  if (artifact_tier && prog != nullptr) {
+  if (per_task_cache_enabled && prog != nullptr) {
     auto &cc = prog->per_construct_cache();
     std::lock_guard<std::mutex> g(cc.mu);
     cc.last_task_stats[kernel->get_name()] = {n, n_hit.load(), n_recompiled.load()};

quadrants/codegen/llvm/per_task_artifact_cache.h

@@ -85,6 +85,16 @@ class PerTaskArtifactCache {
     }
   }
 
+  // Drop a cached record so a later process refills it. try_load rejects framing-level corruption, but a malformed
+  // payload is only detectable by the backend loader.
+  void erase(const std::string &ir_key) const {
+    if (dir_.empty()) {
+      return;
+    }
+    std::error_code ec;
+    std::filesystem::remove(path_for(ir_key), ec);
+  }
+
  private:
   // The IR key is a hex digest plus a `#<index>` suffix; '#' is legal but shell-awkward, so map it (and '/') to '_'.
   // The `.qdb` extension is mandatory -- the binary serializer refuses any other suffix.

quadrants/runtime/cpu/jit_cpu.cpp

@@ -32,6 +32,7 @@
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Scalar.h"
@@ -49,6 +50,7 @@
 #include "quadrants/jit/jit_session.h"
 #include "quadrants/util/file_sequence_writer.h"
 #include "quadrants/runtime/llvm/llvm_context.h"
+#include "quadrants/codegen/llvm/per_task_artifact_cache.h"
 
 namespace quadrants::lang {
 
@@ -80,10 +82,14 @@ class JITSessionCPU;
 class JITModuleCPU : public JITModule {
  private:
   JITSessionCPU *session_;
-  JITDylib *dylib_;
+  // One dylib on the whole-module path, one per task on the per-task path. Lookups search them all; task entry names
+  // are unique across dylibs.
+  std::vector<JITDylib *> dylibs_;
 
  public:
-  JITModuleCPU(JITSessionCPU *session, JITDylib *dylib) : session_(session), dylib_(dylib) {
+  JITModuleCPU(JITSessionCPU *session, JITDylib *dylib) : session_(session), dylibs_{dylib} {
+  }
+  JITModuleCPU(JITSessionCPU *session, std::vector<JITDylib *> dylibs) : session_(session), dylibs_(std::move(dylibs)) {
   }
 
   void *lookup_function(const std::string &name) override;
@@ -104,6 +110,8 @@ class JITSessionCPU : public JITSession {
   std::vector<llvm::orc::JITDylib *> all_libs_;
   int module_counter_;
   SectionMemoryManager *memory_manager_;
+  // Built on first per-task cache miss, reused for the rest of the session.
+  std::unique_ptr<llvm::TargetMachine> pertask_target_machine_;
 
  public:
   JITSessionCPU(QuadrantsLLVMContext *tlctx,
@@ -175,6 +183,75 @@ class JITSessionCPU : public JITSession {
     return new_module_raw_ptr;
   }
 
+  JITModule *add_module_per_task(std::vector<PerConstructArtifact> artifacts, int max_reg) override {
+    QD_ASSERT(max_reg == 0);  // No need to specify max_reg on CPUs
+    std::lock_guard<std::mutex> _(mut_);
+
+    const PerTaskArtifactCache artifact_cache(pertask_artifact_dir_ref());
+
+    // One dylib per task keeps each task's object in its own symbol namespace, so shared helper symbols never collide.
+    std::vector<llvm::orc::JITDylib *> dylibs;
+    dylibs.reserve(artifacts.size());
+    for (auto &art : artifacts) {
+      // Advance up front: an error below throws with the dylib already created, so this id must not be reused.
+      const int mod_id = module_counter_++;
+      auto dylib_expect = es_.createJITDylib(fmt::format("pertask_{}", mod_id));
+      QD_ASSERT(dylib_expect);
+      auto &dylib = dylib_expect.get();
+      dylib.addGenerator(
+          cantFail(llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(dl_.getGlobalPrefix())));
+
+      std::unique_ptr<llvm::MemoryBuffer> obj;
+      const bool from_cache = !art.code.empty();
+      if (from_cache) {
+        obj = llvm::MemoryBuffer::getMemBufferCopy(llvm::StringRef(art.code.data(), art.code.size()),
+                                                   fmt::format("pertask_{}", mod_id));
+      } else {
+        QD_ASSERT(art.module);
+        obj = compile_module_to_object(*art.module);
+        if (!art.key.empty()) {
+          PerTaskArtifact rec;
+          rec.tasks = art.tasks;
+          rec.used_tree_ids = art.used_tree_ids;
+          rec.struct_for_tls_sizes = art.struct_for_tls_sizes;
+          rec.code.assign(obj->getBufferStart(), obj->getBufferEnd());
+          artifact_cache.store(art.key, rec);
+        }
+      }
+      // A corrupt cached object must not abort the process, so don't cantFail. Drop the record so a later run refills
+      // it and raise a catchable error.
+      if (auto err = object_layer_.add(dylib, std::move(obj))) {
+        if (!art.key.empty()) {
+          artifact_cache.erase(art.key);
+        }
+        QD_ERROR("Failed to load per-task CPU object into the JIT (offline cache may be corrupt): {}",
+                 llvm::toString(std::move(err)));
+      }
+      // add() links lazily, so link errors (e.g. a corrupt relocation) only surface at lookup, by which point the key
+      // is gone. Force materialization now to catch them while we can still erase the record.
+      for (const auto &task : art.tasks) {
+#ifdef __APPLE__
+        auto sym = es_.lookup({&dylib}, mangle_(task.name));
+#else
+        auto sym = es_.lookup({&dylib}, es_.intern(task.name));
+#endif
+        if (!sym) {
+          if (!art.key.empty()) {
+            artifact_cache.erase(art.key);
+          }
+          QD_ERROR("Failed to materialize per-task CPU object for \"{}\" (offline cache may be corrupt): {}", task.name,
+                   llvm::toString(sym.takeError()));
+        }
+      }
+      dylibs.push_back(&dylib);
+    }
+
+    auto new_module = std::make_unique<JITModuleCPU>(this, std::move(dylibs));
+    auto *new_module_raw_ptr = new_module.get();
+    modules.push_back(std::move(new_module));
+    return new_module_raw_ptr;
+  }
+
   void *lookup(const std::string Name) override {
     std::lock_guard<std::mutex> _(mut_);
 #ifdef __APPLE__
@@ -187,21 +264,50 @@ class JITSessionCPU : public JITSession {
     return symbol->getAddress().toPtr<void *>();
   }
 
-  void *lookup_in_module(JITDylib *lib, const std::string Name) {
+  void *lookup_in_modules(const std::vector<JITDylib *> &libs, const std::string Name) {
     std::lock_guard<std::mutex> _(mut_);
 #ifdef __APPLE__
-    auto symbol = es_.lookup({lib}, mangle_(Name));
+    auto symbol = es_.lookup(libs, mangle_(Name));
 #else
-    auto symbol = es_.lookup({lib}, es_.intern(Name));
+    auto symbol = es_.lookup(libs, es_.intern(Name));
 #endif
     if (!symbol)
       QD_ERROR("Function \"{}\" not found", Name);
     return symbol->getAddress().toPtr<void *>();
   }
+
+ private:
+  // Compile a per-task module to a host object file for the on-disk cache.
+  std::unique_ptr<llvm::MemoryBuffer> compile_module_to_object(llvm::Module &M) {
+    if (!pertask_target_machine_) {
+      auto expected_jtmb = llvm::orc::JITTargetMachineBuilder::detectHost();
+      if (!expected_jtmb) {
+        QD_ERROR("LLVM TargetMachineBuilder has failed.");
+      }
+      // detectHost() carries the running core's explicit feature vector, matching the whole-module JIT compiler; a bare
+      // CPU name with empty features would enable that model's default superset and emit illegal instructions.
+      //
+      // Large code model: each per-task object gets its own SectionMemoryManager, so under memory pressure RTDyld can
+      // place a task's .text and .rodata more than 2GB apart, overflowing the small model's 32-bit RIP-relative refs
+      // and reading constants from garbage at launch. 64-bit addressing makes the section distance irrelevant.
+      auto jtmb = std::move(*expected_jtmb);
+      jtmb.setCodeModel(llvm::CodeModel::Large);
+      auto expected_tm = jtmb.createTargetMachine();
+      QD_ERROR_UNLESS(expected_tm, "Could not allocate target machine!");
+      pertask_target_machine_ = std::move(*expected_tm);
+    }
+    M.setDataLayout(pertask_target_machine_->createDataLayout());
+    llvm::orc::SimpleCompiler compiler(*pertask_target_machine_);
+    auto obj = compiler(M);
+    if (!obj) {
+      QD_ERROR("Per-task CPU object compilation failed");
+    }
+    return std::move(*obj);
+  }
 };
 
 void *JITModuleCPU::lookup_function(const std::string &name) {
-  return session_->lookup_in_module(dylib_, name);
+  return session_->lookup_in_modules(dylibs_, name);
 }
 
 std::unique_ptr<JITSession> create_llvm_jit_session_cpu(QuadrantsLLVMContext *tlctx,

quadrants/runtime/cpu/kernel_launcher.cpp

@@ -339,7 +339,13 @@ KernelLauncher::Handle KernelLauncher::register_llvm_kernel(const LLVM::Compiled
     auto *executor = get_runtime_executor();
 
     auto data = compiled.get_internal_data().compiled_data.clone();
-    auto *jit_module = executor->create_jit_module(std::move(data.module));
+    JITModule *jit_module = nullptr;
+    if (!data.per_construct_artifacts.empty()) {
+      // Per-task cache hit path: assemble the kernel from per-task objects instead of a whole-kernel module.
+      jit_module = executor->create_jit_module_per_task(std::move(data.per_construct_artifacts));
+    } else {
+      jit_module = executor->create_jit_module(std::move(data.module));
+    }
 
     std::vector<TaskFunc> task_funcs;
     std::vector<int32_t> checkpoint_id_per_task;

quadrants/runtime/program_impls/llvm/llvm_program.cpp

@@ -1,6 +1,8 @@
 #include "quadrants/runtime/program_impls/llvm/llvm_program.h"
 
 #include "llvm/IR/Module.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/TargetParser/Host.h"
 
 #include "quadrants/codegen/cpu/codegen_cpu.h"
 #include "quadrants/codegen/llvm/llvm_compiled_data.h"
@@ -50,6 +52,19 @@ LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_, KernelProfilerBase *pro
       pertask_dir += "_" + AMDGPUContext::get_instance().get_mcpu();
     }
 #endif
+    if (arch_is_cpu(config_.arch)) {
+      // Host objects are CPU-specific, so scope the dir by host triple + CPU; a shared path must never serve an object
+      // built for an incompatible CPU.
+      auto jtmb = llvm::orc::JITTargetMachineBuilder::detectHost();
+      std::string tag =
+          (jtmb ? jtmb->getTargetTriple().str() : std::string("unknown")) + "_" + llvm::sys::getHostCPUName().str();
+      for (char &c : tag) {
+        if (c == '/' || c == ':' || c == ' ' || c == '\\') {
+          c = '_';
+        }
+      }
+      pertask_dir += "_cpu_" + tag;
+    }
   }
   pertask_artifact_dir_ref() = pertask_dir;
 }