[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 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 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. (CUDA caches each task's PTX; CPU caches a host relocatable object file. The cache is scoped per target -- by GPU compute capability on CUDA, by host CPU on CPU -- so a shared `offline_cache_file_path` never serves code built for a different device.)
 - 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,10 +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 entire compilation and carries cached PTX + launch
-  // metadata to the launcher. CUDA-only; an empty dir means offline cache is off (tier disabled).
+  // Cross-process per-task artifact cache: a hit skips a task's entire compilation and carries cached backend code +
+  // launch metadata to the launcher. CUDA + CPU; an empty dir means offline cache is off (tier disabled).
   const std::string art_dir = pertask_artifact_dir_ref();
-  const bool artifact_tier = compile_config_.arch == Arch::cuda && !art_dir.empty();
+  const bool artifact_tier =
+      (compile_config_.arch == Arch::cuda || 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);
@@ -140,10 +141,12 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() {
   }
   worker.flush();
 
-  // Build one self-contained artifact per task BEFORE the whole-module link consumes `data`: a hit carries cached PTX
-  // (null module), a miss builds+optimizes a module for the JIT to compile and store.
+  // Build one self-contained artifact per task BEFORE the whole-module link consumes `data`: a hit carries cached
+  // backend code (null module), a miss builds+optimizes a module for the JIT to compile and store. CUDA always takes
+  // this path (the Option B composite module); CPU only when the artifact tier is on (otherwise it keeps its
+  // whole-kernel module), so gate CPU on `artifact_tier`.
   std::vector<PerConstructArtifact> per_construct_artifacts;
-  if (compile_config_.arch == Arch::cuda) {
+  if (compile_config_.arch == Arch::cuda || artifact_tier) {
     for (int i = 0; i < n; i++) {
       if (!data[i])
         continue;
@@ -172,7 +175,7 @@ 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
-  // CUmodule from the per-task artifacts); still concatenate every task's metadata into `tasks`, which it runs off.
+  // kernel from the per-task artifacts); still concatenate every task's metadata into `tasks`, which it runs off.
   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) {

quadrants/codegen/llvm/per_task_artifact_cache.h

@@ -85,6 +85,17 @@ class PerTaskArtifactCache {
     }
   }
 
+  // Drop a cached record so a later process recompiles and refills it. Used when a backend rejects the payload bytes as
+  // corrupt: try_load only catches framing-level corruption (truncation / bad length header), not a decodable record
+  // whose `code` is a malformed backend object / PTX, which only the backend loader can detect.
+  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,15 @@ class JITSessionCPU;
 class JITModuleCPU : public JITModule {
  private:
   JITSessionCPU *session_;
-  JITDylib *dylib_;
+  // The whole-module path resolves in one dylib; the per-task path (add_module_per_task) resolves across N, one per
+  // task. A task's entry symbol lives in exactly one of them, so a by-name lookup that searches them all is
+  // unambiguous (only the unique task entry names are ever looked up; duplicated helper symbols never are).
+  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 +111,9 @@ class JITSessionCPU : public JITSession {
   std::vector<llvm::orc::JITDylib *> all_libs_;
   int module_counter_;
   SectionMemoryManager *memory_manager_;
+  // Lazily built host PIC target machine for the per-task object serialize path (add_module_per_task); reused across
+  // all misses in a session.
+  std::unique_ptr<llvm::TargetMachine> pertask_target_machine_;
 
  public:
   JITSessionCPU(QuadrantsLLVMContext *tlctx,
@@ -175,6 +185,69 @@ 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_);
+
+    // Cross-process fill site (mirror of runtime/cuda/jit_cuda.cpp): a hit already carries the cached host object in
+    // `code`, used verbatim; a miss compiles the per-task module to a host object and stores it -- object bytes plus
+    // the launch metadata that must travel with it -- under the task's IR key so a later process skips this
+    // compilation. No-ops when the dir is empty (offline cache off).
+    const PerTaskArtifactCache artifact_cache(pertask_artifact_dir_ref());
+
+    // One JITDylib per task (a task's entry symbol lives in exactly one), so per-task objects never collide on a
+    // shared helper / global symbol -- the CPU analog of CUDA's one-CUmodule-per-task, and why CPU needs no relink.
+    std::vector<llvm::orc::JITDylib *> dylibs;
+    dylibs.reserve(artifacts.size());
+    for (auto &art : artifacts) {
+      // Advance the counter up front: the error path below throws mid-iteration and the created dylib lingers in the
+      // session, so reusing this id on a later call would collide on the dylib name (createJITDylib would fail).
+      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) {
+        // Hit: the cached bytes are a relocatable host object; load them straight into the object layer.
+        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);
+        }
+      }
+      // `object_layer_.add` parses the object eagerly, so malformed bytes surface here -- offline-cache corruption on
+      // the hit path, or (defensively) a bad freshly-compiled object we just stored on the miss path. Don't `cantFail`
+      // -- that aborts the whole process on every launch through a corrupt cache. Drop the on-disk entry (either path
+      // may have written one) so a later process recompiles and refills it, and raise a catchable error instead of
+      // terminating (mirrors the CUDA per-task load, which QD_ERRORs on a bad module).
+      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)));
+      }
+      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 +260,61 @@ 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:
+  // Serialize a self-contained per-task module to a host relocatable object -- the "serialize" half of the per-task
+  // disk tier, called on a cache miss. Build the target machine the same way as KernelCodeGenCPU::optimize_module
+  // (host CPU, PIC so the object is loadable by the ORC object layer) so the emitted object matches the module that
+  // optimize_module already produced.
+  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.");
+      }
+      auto triple = expected_jtmb->getTargetTriple();
+      std::string err_str;
+      const llvm::Target *target = llvm::TargetRegistry::lookupTarget(triple.str(), err_str);
+      QD_ERROR_UNLESS(target, err_str);
+      llvm::TargetOptions options;
+      if (config_.fast_math) {
+        options.AllowFPOpFusion = llvm::FPOpFusion::Fast;
+        options.NoInfsFPMath = 1;
+        options.NoNaNsFPMath = 1;
+      } else {
+        options.AllowFPOpFusion = llvm::FPOpFusion::Strict;
+        options.NoInfsFPMath = 0;
+        options.NoNaNsFPMath = 0;
+      }
+      llvm::StringRef mcpu = llvm::sys::getHostCPUName();
+      pertask_target_machine_.reset(target->createTargetMachine(triple, mcpu.str(), "", options, llvm::Reloc::PIC_,
+                                                                llvm::CodeModel::Small,
+                                                                llvm::CodeGenOptLevel::Aggressive));
+      QD_ERROR_UNLESS(pertask_target_machine_.get(), "Could not allocate target machine!");
+    }
+    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,14 @@ 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 path: assemble the kernel from per-task objects (the cross-process artifact cache) instead of the
+      // whole-kernel module. Mirrors runtime/cuda/kernel_launcher.cpp.
+      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"
@@ -43,6 +45,20 @@ LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_, KernelProfilerBase *pro
       pertask_dir += "_sm_" + std::to_string(CUDAContext::get_instance().get_compute_capability());
     }
 #endif
+    if (arch_is_cpu(config_.arch)) {
+      // Artifacts hold host object code tied to the target machine, so scope the dir by host triple + CPU (as the
+      // CUDA branch scopes by sm); an NFS-shared cache path could otherwise serve an object built for one host to an
+      // incompatible CPU. Matches the target machine built in jit_cpu.cpp / KernelCodeGenCPU::optimize_module.
+      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;
 }