diff --git a/docs/source/user_guide/init_options.md b/docs/source/user_guide/init_options.md index 20c42dfdd4..a2d923bdb3 100644 --- a/docs/source/user_guide/init_options.md +++ b/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. diff --git a/quadrants/codegen/codegen.cpp b/quadrants/codegen/codegen.cpp index bcff2f13b6..a419153a18 100644 --- a/quadrants/codegen/codegen.cpp +++ b/quadrants/codegen/codegen.cpp @@ -72,11 +72,11 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() { const int n = (int)offloads.size(); std::vector> 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 cache. + // 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_cache_on = - (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 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_cache_on && offload->as()->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()->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() || s->is() || s->is()) { - 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(), 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 cache 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 per_construct_artifacts; const bool build_per_construct_artifacts = - compile_config_.arch == Arch::cuda || (compile_config_.arch == Arch::amdgpu && artifact_cache_on); + 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) { @@ -206,8 +206,8 @@ LLVMCompiledKernel KernelCodeGen::compile_kernel_to_module() { } } // Record per-task reuse counts for PerOffloadCacheObservations.tasks_* (read back by the compilation manager). Only - // when the artifact cache ran, so non-CUDA / cache-off compiles keep the -1 sentinel instead of a misleading 0. - if (artifact_cache_on && prog != nullptr) { + // when the per-task cache ran, so a compile with it off keeps the -1 sentinel instead of a misleading 0. + if (per_task_cache_enabled && prog != nullptr) { auto &cc = prog->per_construct_cache(); std::lock_guard g(cc.mu); cc.last_task_stats[kernel->get_name()] = {n, n_hit.load(), n_recompiled.load()}; diff --git a/quadrants/codegen/llvm/per_task_artifact_cache.h b/quadrants/codegen/llvm/per_task_artifact_cache.h index b61fb94fba..7dc8f6e96e 100644 --- a/quadrants/codegen/llvm/per_task_artifact_cache.h +++ b/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 `#` suffix; '#' is legal but shell-awkward, so map it (and '/') to '_'. // The `.qdb` extension is mandatory -- the binary serializer refuses any other suffix. diff --git a/quadrants/runtime/cpu/jit_cpu.cpp b/quadrants/runtime/cpu/jit_cpu.cpp index a3f750205d..7a2bd8e6c9 100644 --- a/quadrants/runtime/cpu/jit_cpu.cpp +++ b/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 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 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 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 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 artifacts, int max_reg) override { + QD_ASSERT(max_reg == 0); // No need to specify max_reg on CPUs + std::lock_guard _(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 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 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 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(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 _(mut_); #ifdef __APPLE__ @@ -187,21 +264,49 @@ class JITSessionCPU : public JITSession { return symbol->getAddress().toPtr(); } - void *lookup_in_module(JITDylib *lib, const std::string Name) { + void *lookup_in_modules(const std::vector &libs, const std::string Name) { std::lock_guard _(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(); } + + private: + // Compile a per-task module to a host object file for the on-disk cache. + std::unique_ptr 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."); + } + // Use detectHost()'s explicit features; a bare CPU name would enable features the running core may lack. + auto jtmb = std::move(*expected_jtmb); + // On x86 RTDyld can place a per-task object's .text and .rodata over 2GB apart, past the reach of 32-bit + // RIP-relative refs; the large code model's 64-bit addressing avoids that. + if (jtmb.getTargetTriple().isX86()) { + 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 create_llvm_jit_session_cpu(QuadrantsLLVMContext *tlctx, diff --git a/quadrants/runtime/cpu/kernel_launcher.cpp b/quadrants/runtime/cpu/kernel_launcher.cpp index 7d7fcfde86..f7e27adb21 100644 --- a/quadrants/runtime/cpu/kernel_launcher.cpp +++ b/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 task_funcs; std::vector checkpoint_id_per_task; diff --git a/quadrants/runtime/program_impls/llvm/CMakeLists.txt b/quadrants/runtime/program_impls/llvm/CMakeLists.txt index 922c047b70..46b99a033d 100644 --- a/quadrants/runtime/program_impls/llvm/CMakeLists.txt +++ b/quadrants/runtime/program_impls/llvm/CMakeLists.txt @@ -11,6 +11,7 @@ target_include_directories(llvm_program_impl ${PROJECT_SOURCE_DIR} ${PROJECT_SOURCE_DIR}/external/eigen ${PROJECT_SOURCE_DIR}/external/spdlog/include + ${PROJECT_SOURCE_DIR}/external/PicoSHA2 ${LLVM_INCLUDE_DIRS} ) diff --git a/quadrants/runtime/program_impls/llvm/llvm_program.cpp b/quadrants/runtime/program_impls/llvm/llvm_program.cpp index bbded19886..c273e5785b 100644 --- a/quadrants/runtime/program_impls/llvm/llvm_program.cpp +++ b/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" @@ -27,6 +29,10 @@ #include "quadrants/codegen/llvm/compiled_kernel_data.h" #include "quadrants/codegen/llvm/per_task_artifact_cache.h" +#include "picosha2.h" + +#include + namespace quadrants::lang { LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_, KernelProfilerBase *profiler) : ProgramImpl(config_), compilation_workers("compile", config_.print_ir ? 1 : config_.num_compile_threads) { @@ -50,6 +56,31 @@ LlvmProgramImpl::LlvmProgramImpl(CompileConfig &config_, KernelProfilerBase *pro pertask_dir += "_" + AMDGPUContext::get_instance().get_mcpu(); } #endif + if (arch_is_cpu(config_.arch)) { + // Host objects match detectHost()'s triple, CPU and features, so a cache path shared between hosts is scoped by + // all three. + auto jtmb = llvm::orc::JITTargetMachineBuilder::detectHost(); + std::string tag = + (jtmb ? jtmb->getTargetTriple().str() : std::string("unknown")) + "_" + llvm::sys::getHostCPUName().str(); + if (jtmb) { + // Sort so the hash ignores feature order; the full vector is too long to use directly. + std::vector features = jtmb->getFeatures().getFeatures(); + std::sort(features.begin(), features.end()); + std::string joined; + for (const auto &f : features) { + joined += f + ","; + } + std::string feat_hex; + picosha2::hash256_hex_string(joined, feat_hex); + tag += "_" + feat_hex.substr(0, 16); + } + for (char &c : tag) { + if (c == '/' || c == ':' || c == ' ' || c == '\\') { + c = '_'; + } + } + pertask_dir += "_cpu_" + tag; + } } pertask_artifact_dir_ref() = pertask_dir; } diff --git a/tests/python/test_per_offload_cache.py b/tests/python/test_per_offload_cache.py index 0f93a36663..ee0445c6d3 100644 --- a/tests/python/test_per_offload_cache.py +++ b/tests/python/test_per_offload_cache.py @@ -1289,12 +1289,12 @@ def kernel_struct(out1: qd.types.ndarray(), out2: qd.types.ndarray()) -> None: assert np.allclose(out2.to_numpy(), _C[0], atol=1.0), out2.to_numpy() -# --- Cross-process per-task artifact cache (CUDA + AMDGPU backend reuse) ---------------------------------------------- +# --- Cross-process per-task artifact cache (CUDA + AMDGPU + CPU) ------------------------------------------------------ # # The per-task artifact cache stores each offloaded task's fully compiled code + launch metadata on disk, keyed by the -# task's own IR (name-free), so a later process reuses an unchanged task instead of recompiling it. CUDA and AMDGPU -# fill it, and it is gated on `offline_cache`. Reuse is reported on `PerOffloadCacheObservations.tasks_*` (-1 when the -# artifact cache did not run). +# task's own IR (name-free), so a later process reuses an unchanged task instead of recompiling it. CUDA and AMDGPU fill +# it with GPU code and CPU with a host object; it is gated on `offline_cache`. Reuse is reported on +# `PerOffloadCacheObservations.tasks_*` (-1 when offline_cache is disabled). @test_utils.test(arch=[qd.cuda, qd.amdgpu], offline_cache=False) @@ -1372,3 +1372,74 @@ def k_second(x: qd.types.ndarray(qd.f32, ndim=1)) -> None: finally: qd.reset() shutil.rmtree(cache_dir, ignore_errors=True) + + +@test_utils.test(arch=qd.cpu, offline_cache=False) +def test_per_task_artifact_cache_disabled_without_offline_cache_cpu() -> None: + # With offline_cache off the per-task cache never runs, so its counts stay at the -1 sentinel; the frontend split + # still fires. + @qd.kernel + def kernel_two_loops(x: qd.types.ndarray(qd.f32, ndim=1)) -> None: + for i in x: + x[i] = x[i] * 2.0 + 1.0 + for i in x: + x[i] = x[i] - 3.0 + + arr = qd.ndarray(qd.f32, shape=(_N,)) + arr.from_numpy(np.arange(_N, dtype=np.float32)) + kernel_two_loops(arr) + + obs = kernel_two_loops._primal.per_offload_cache_observations + assert obs.frontend_constructs_total == 2, obs + assert obs.tasks_total == -1, obs + assert obs.tasks_cache_hit == -1, obs + assert obs.tasks_recompiled == -1, obs + assert np.allclose(arr.to_numpy(), np.arange(_N) * 2.0 + 1.0 - 3.0), arr.to_numpy() + + +def test_per_task_artifact_cache_reuses_shared_task_cross_process_cpu() -> None: + # A fresh process with a warm disk cache loads an unchanged task instead of recompiling it. Re-init with the same + # cache path emulates the second process. + if qd.cpu not in test_utils.expected_archs(): + pytest.skip("this variant exercises the CPU per-task artifact cache") + + cache_dir = tempfile.mkdtemp() + try: + qd.init(arch=qd.cpu, offline_cache=True, offline_cache_file_path=cache_dir) + + @qd.kernel + def k_first(x: qd.types.ndarray(qd.f32, ndim=1)) -> None: + for i in x: + x[i] = x[i] * 2.0 + 1.0 + for i in x: + x[i] = x[i] - 3.0 + + a = qd.ndarray(qd.f32, shape=(_N,)) + a.from_numpy(np.arange(_N, dtype=np.float32)) + k_first(a) + obs1 = k_first._primal.per_offload_cache_observations + assert obs1.tasks_total >= 2, obs1 + assert obs1.tasks_cache_hit == 0, obs1 + + # Second "process": fresh runtime, same disk. k_second's first loop matches k_first's, so that task is a hit. + qd.init(arch=qd.cpu, offline_cache=True, offline_cache_file_path=cache_dir) + + @qd.kernel + def k_second(x: qd.types.ndarray(qd.f32, ndim=1)) -> None: + for i in x: + x[i] = x[i] * 2.0 + 1.0 + for i in x: + x[i] = x[i] + 7.0 + + b = qd.ndarray(qd.f32, shape=(_N,)) + b.from_numpy(np.arange(_N, dtype=np.float32)) + k_second(b) + obs2 = k_second._primal.per_offload_cache_observations + assert obs2.tasks_cache_hit > 0, obs2 + assert obs2.tasks_recompiled >= 1, obs2 + assert obs2.tasks_cache_hit + obs2.tasks_recompiled == obs2.tasks_total, obs2 + + assert np.allclose(b.to_numpy(), np.arange(_N) * 2.0 + 1.0 + 7.0), b.to_numpy() + finally: + qd.reset() + shutil.rmtree(cache_dir, ignore_errors=True)