[python] Unify caching and key cache on CompileTarget hash (NVIDIA#4859)

Finally! Caching for Python done right 🙃 This builds on top of
@atgeller's work in NVIDIA#4607,
but using the CompileTarget now that it's available.

We store the hashes of the compilation options that affect the compiled
output within `CompiledModule`. We then determine whether the cached
module can be reused by just comparing those hashes to the ones obtained
for the new module.

We distinguish three cases when it comes to caching:

1. When `CompileTarget.fullySpecialize == true`, we disable caching.
This is the case where arguments get inlined into the IR, so keeping
caches here would require hashing the runtime arguments. We currently
don't support that. This mostly affects remote devices, so not having
caching here is reasonable (and the current behaviour already).
2. When `CompileTarget.fullySpecialize == false` and there are no
captured kernels, we key the cache on a hash of `CompileTarget`. This is
the 'best' case, where the compilation options within `CompileTarget`
fully determine the output of the compilation.
3. When `CompileTarget.fullySpecialize == false` and there are captured
kernels, we key the cache on a pair of hashes: the hash of
`CompileTarget` (as in 2) and the hash of the module IR itself. This
detects changes to the captured kernels, as the captured targets are
inlined into the IR before computing the hash.

This should also subsume the second level of caching we were performing
within the ExecutionContext. I'm going to run some benchmarks to make
sure this hasn't regressed.

---------

Signed-off-by: Luca Mondada <luca@mondada.net>

Author: lmondada

cudaq/include/cudaq/Target/CompileTarget.h

@@ -21,6 +21,15 @@ class CompileTarget {
   /// Hook to update the pass pipeline before compilation.
   virtual void updatePassPipeline(std::string &passPipeline) const {}
 
+  /// Return a hash of this target's configuration, used as a cache key to
+  /// decide whether a previously compiled module can be reused.
+  ///
+  /// A hash of 0 disables caching.
+  ///
+  /// TODO: CompileTarget should be made non-virtual. Once it is, this method
+  /// becomes unnecessary and callers can use std::hash<CompileTarget> directly.
+  virtual std::size_t hash() const;
+
   /// Whether to recompile the kernel in the presence of an AOT-compiled module.
   ///
   /// If this is `false` and an AOT-compiled kernel (in the form of a function
@@ -137,3 +146,8 @@ class CompileTarget {
 };
 
 } // namespace cudaq
+
+template <>
+struct std::hash<cudaq::CompileTarget> {
+  std::size_t operator()(const cudaq::CompileTarget &t) const noexcept;
+};

cudaq/lib/Target/CompileTarget.cpp

@@ -9,6 +9,7 @@
 #include "cudaq/Target/CompileTarget.h"
 #include "cudaq/runtime/logger/logger.h"
 #include <cctype>
+#include <functional>
 
 /// Replace `%KEY%` and `%KEY:default%` placeholders from runtime options.
 static void substitutePipelinePlaceholders(
@@ -97,3 +98,41 @@ cudaq::CompileTarget::CompileTarget(
     CUDAQ_INFO("{:<27} {}\n", "disable_qubit_mapping:", "true");
   }
 }
+
+template <typename T>
+inline void hash_combine(std::size_t &seed, const T &val) {
+  std::hash<T> hasher;
+  seed ^= hasher(val) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+}
+
+template <typename... Args>
+inline std::size_t hash_val(const Args &...args) {
+  std::size_t seed = 0;
+  (hash_combine(seed, args), ...);
+  return seed;
+}
+
+std::size_t std::hash<cudaq::CompileTarget>::operator()(
+    const cudaq::CompileTarget &t) const noexcept {
+  std::size_t seed = hash_val(
+      t.pipelineConfig.overridePassPipeline, t.pipelineConfig.highLevelPipeline,
+      t.pipelineConfig.midLevelPipeline, t.pipelineConfig.lowLevelPipeline,
+      t.pipelineConfig.codegenTranslation, t.pipelineConfig.postCodeGenPasses,
+      t.pipelineConfig.skipTargetLoweringPipeline,
+      t.pipelineConfig.disableQubitMapping,
+      t.pipelineConfig.replaceStateWithKernel, t.pipelineConfig.addMeasurements,
+      t.overrideAOTCompilation, t.emulate, t.warnNamedMeasurements,
+      t.supportConditionalsOnMeasureResults, t.supportDeviceCalls,
+      t.storeReorderIdx, t.emitResourceCounts, t.emitJit, t.emitTargetCode,
+      t.fullySpecialize, t.isLocalSimulator, t.argumentSynthChangeSemantics);
+
+  // Optional spin observable: include its string representation when present.
+  if (t.pauliTermSplitObservable)
+    hash_combine(seed, t.pauliTermSplitObservable->to_string());
+
+  return seed;
+}
+
+std::size_t cudaq::CompileTarget::hash() const {
+  return std::hash<cudaq::CompileTarget>{}(*this);
+}

python/cudaq/runtime/observe.py

@@ -38,16 +38,11 @@ def __broadcastObserve(kernel, spin_operator, *args, shots_count=0, qpu_id=0):
     ctx.kernelName = kernel_name
     has_vector_args = isa_kernel_decorator(kernel) and any(
         hasattr(a, 'shape') and len(a.shape) == 2 for a in args)
-    if has_vector_args:
-        ctx.allowJitEngineCaching = True
-        ctx.useParametricJit = True
     policy = cudaq_runtime.ObservePolicy(ctx, kernel_name, spin_operator)
     for i, a in enumerate(argSet):
         ctx.batchIteration = i
         results.append(
             cudaq_runtime.launch_observe(policy, ctx, lambda a=a: kernel(*a)))
-    if has_vector_args:
-        ctx.unset_jit_engine()
     return results
 
 
@@ -171,7 +166,6 @@ def __computeTermExpectation(term, observe_result):
         else:
             ctx = cudaq_runtime.ExecutionContext('observe', 0, qpu_id)
         ctx.setSpinOperator(localOp)
-        ctx.allowJitEngineCaching = True
         if num_trajectories is not None:
             if noise_model is None:
                 raise RuntimeError(
@@ -202,7 +196,6 @@ def __computeTermExpectation(term, observe_result):
             results.append(
                 cudaq_runtime.ObserveResult(exp_val, op,
                                             observeResult.counts()))
-        ctx.unset_jit_engine()
 
     if noise_model != None:
         cudaq_runtime.unset_noise()

python/cudaq/runtime/sample.py

@@ -166,7 +166,6 @@ def sample(kernel,
         ctx = cudaq_runtime.ExecutionContext("sample", shots_count)
         ctx.kernelName = kernel_name
         ctx.explicitMeasurements = explicit_measurements
-        ctx.allowJitEngineCaching = True
         policy = cudaq_runtime.SamplePolicy(ctx, kernel_name,
                                             explicit_measurements)
 
@@ -194,7 +193,6 @@ def sample(kernel,
                     "results when executed. Exiting shot loop to avoid infinite "
                     "loop.")
                 break
-        ctx.unset_jit_engine()
         return counts
     finally:
         if set_noise_for_call:

python/runtime/common/py_ExecutionContext.cpp

@@ -42,18 +42,8 @@ void bindExecutionContext(nanobind::module_ &mod) {
               &cudaq::ExecutionContext::numberTrajectories)
       .def_rw("explicitMeasurements",
               &cudaq::ExecutionContext::explicitMeasurements)
-      .def_rw("allowJitEngineCaching",
-              &cudaq::ExecutionContext::allowCompiledModuleCaching)
-      .def_rw("useParametricJit", &cudaq::ExecutionContext::useParametricJit)
       .def_ro("invocationResultBuffer",
               &cudaq::ExecutionContext::invocationResultBuffer)
-      .def("unset_jit_engine",
-           [&](cudaq::ExecutionContext &execCtx) {
-             if (execCtx.cachedCompiledModule) {
-               execCtx.cachedCompiledModule = std::nullopt;
-               execCtx.allowCompiledModuleCaching = false;
-             }
-           })
       .def("setSpinOperator",
            [](cudaq::ExecutionContext &ctx, cudaq::spin_op &spin) {
              ctx.spin = spin;

python/runtime/cudaq/algorithms/py_run.cpp

@@ -73,7 +73,7 @@ static detail::RunResultSpan
 pyRunTheKernel(const std::string &name, quantum_platform &platform,
                mlir::ModuleOp mod, CompiledModule *compiled,
                std::size_t shots_count, std::size_t qpu_id,
-               OpaqueArguments &opaques, bool allowCaching) {
+               OpaqueArguments &opaques) {
   if (!name.ends_with(".run"))
     throw std::runtime_error("`cudaq.run` only supports runnable kernels.");
   // Set the `run` attribute on the module to indicate this is a run context
@@ -101,7 +101,7 @@ pyRunTheKernel(const std::string &name, quantum_platform &platform,
         [[maybe_unused]] auto result =
             clean_launch_module(name, mod, opaques, compiled);
       },
-      platform, name, name, shots_count, layoutInfo, qpu_id, allowCaching);
+      platform, name, name, shots_count, layoutInfo, qpu_id);
 
   return results;
 }
@@ -138,7 +138,7 @@ run_impl(const std::string &shortName, MlirModule module,
   {
     nanobind::gil_scoped_release release;
     span = pyRunTheKernel(shortName, platform, mod, compiled, shots_count,
-                          qpu_id, opaques, true);
+                          qpu_id, opaques);
   }
   auto results = pyReadResults(span, mod, shots_count, shortName);
 
@@ -217,7 +217,7 @@ run_async_impl(const std::string &shortName, MlirModule module,
             platform.set_noise(&noise_model.value());
           try {
             auto span = pyRunTheKernel(name, platform, mod, nullptr,
-                                       shots_count, qpu_id, opaques, false);
+                                       shots_count, qpu_id, opaques);
             sp.set_value(span);
             ep.set_value("");
           } catch (std::runtime_error &e) {

python/runtime/cudaq/platform/py_alt_launch_kernel.cpp

@@ -40,6 +40,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
@@ -649,68 +650,39 @@ static void appendTheResultValue(ModuleOp module, const std::string &name,
   runtimeArgs.emplace_back(buf, [](void *ptr) { std::free(ptr); });
 }
 
-/// In a sample launch context, the (`JIT` compiled) CompiledModule may be
-/// cached so that it can be called many times in a loop without being
-/// recompiled. This exploits the fact that the arguments processed at the
-/// sample callsite are invariant by the definition of a `CUDA-Q` kernel.
-template <std::invocable F>
-  requires std::is_invocable_r_v<cudaq::CompiledModule, F>
-static cudaq::CompiledModule with_compiled_module_cache(F &&f) {
-  auto *currentExecCtx = cudaq::getExecutionContext();
-
-  auto getCache = [currentExecCtx]() -> std::optional<cudaq::CompiledModule> {
-    if (currentExecCtx && currentExecCtx->allowCompiledModuleCaching)
-      return currentExecCtx->cachedCompiledModule;
-    return std::nullopt;
-  };
-  auto saveCache = [currentExecCtx](cudaq::CompiledModule compiled) {
-    if (currentExecCtx && currentExecCtx->allowCompiledModuleCaching) {
-      if (!currentExecCtx->cachedCompiledModule)
-        currentExecCtx->cachedCompiledModule = compiled;
-    }
-  };
+/// Compute a hash of the IR given by the `ModuleOp`.
+static std::size_t hashModuleOp(ModuleOp mod) {
+  llvm::hash_code h{0};
+  mod.walk([&h](Operation *op) {
+    h = llvm::hash_combine(h, OperationEquivalence::computeHash(op));
+  });
+  return static_cast<std::size_t>(h);
+}
 
-  auto cachedModule = getCache();
-  if (cachedModule)
-    return *cachedModule;
-  auto compiled = f();
-  saveCache(compiled);
-  return compiled;
+/// Obtain a fresh `CompileTarget` for the current execution context.
+static std::unique_ptr<cudaq::CompileTarget> getCompileTargetImpl() {
+  auto *ctx = cudaq::getExecutionContext();
+  if (!ctx)
+    return cudaq::get_compile_target(cudaq::other_policies{});
+
+  return cudaq::policies::withPolicy(ctx->name, [&](auto policy) {
+    using Policy = std::decay_t<decltype(policy)>;
+    if constexpr (std::is_same_v<Policy, cudaq::observe_policy>) {
+      policy.spin = ctx->spin.value();
+    }
+    return cudaq::get_compile_target(policy);
+  });
 }
 
 static cudaq::CompiledModule
 compileModuleImpl(const std::string &name, ModuleOp mod,
-                  const std::vector<void *> &rawArgs, bool isEntryPoint) {
+                  const std::vector<void *> &rawArgs, bool isEntryPoint,
+                  std::unique_ptr<cudaq::CompileTarget> target = nullptr) {
+  if (!target)
+    target = getCompileTargetImpl();
   cudaq::SourceModule src{name, mod.getAsOpaquePointer()};
-
-  // Only cache on local simulators
-  auto cacheable =
-      cudaq::is_simulator_platform() && !cudaq::is_emulated_platform();
-
-  auto compile = [&]() {
-    cudaq::CompiledModule compiled;
-    auto *ctx = cudaq::getExecutionContext();
-    if (!ctx) {
-      auto target = cudaq::get_compile_target(cudaq::other_policies{});
-      return cudaq_internal::compiler::compileModule(std::move(target), src,
-                                                     {rawArgs}, isEntryPoint);
-    }
-
-    return cudaq::policies::withPolicy(ctx->name, [&](auto policy) {
-      using Policy = std::decay_t<decltype(policy)>;
-      if constexpr (std::is_same_v<Policy, cudaq::observe_policy>) {
-        policy.spin = ctx->spin.value();
-      }
-      auto target = cudaq::get_compile_target(policy);
-      return cudaq_internal::compiler::compileModule(std::move(target), src,
-                                                     {rawArgs}, isEntryPoint);
-    });
-  };
-
-  if (!cacheable) {
-    return compile();
-  }
-  return with_compiled_module_cache(compile);
+  return cudaq_internal::compiler::compileModule(std::move(target), src,
+                                                 {rawArgs}, isEntryPoint);
 }
 
 // Launching the module \p mod will modify its content, such as by argument
@@ -720,38 +692,48 @@ static cudaq::KernelThunkResultType
 pyLaunchModule(const std::string &name, ModuleOp mod,
                cudaq::CompiledModule *cachedModule,
                const std::vector<void *> &rawArgs) {
-  bool isCachable = [&]() {
-    // Must have a slot to read/write the cache from. Callers opt out of the
-    // cache by passing nullptr.
-    if (!cachedModule)
-      return false;
-    auto &platform = cudaq::get_platform();
-    // Must be local simulator
-    if (!platform.is_simulator() || platform.is_emulated())
-      return false;
+  auto target = getCompileTargetImpl();
+  auto targetHash = target->hash();
+
+  // We don't cache kernels that inline all arguments, as any change to the
+  // runtime arguments would invalidate the cache. Currently, synthesis is
+  // all-or-nothing, but if arg-by-arg synthesis is supported, then that will
+  // need to be detected.
+  bool cacheable = cachedModule && !target->fullySpecialize && targetHash != 0;
+
+  // Normally, we assume that the module IR is constant given the uniqued name.
+  // However, kernels that capture other kernels inline captured kernels into
+  // the module, so we need to handle this case specially.
+  bool hasCaptures = [&]() {
     auto func = cudaq::getKernelFuncOp(mod, name);
-    // TODO: currently, synthesis is all-or-nothing, but if arg-by-arg
-    // synthesis is supported, then that will need to be detected
-    if (cudaq::opt::factory::isFullySynthesized(func))
-      return false;
-    // Caching for kernels with lifted arguments is not currently supported.
     for (unsigned i = 0; i < func.getNumArguments(); ++i)
       if (func.getArgAttr(i, "quake.pylifted"))
-        return false;
-    return true;
+        return true;
+    return false;
   }();
-
-  // Cache hit only if the cached module's entry point matches this launch's.
-  // Notably, run has a different entry point so can't share a cache with
-  // other launch modes.
-  if (isCachable && cachedModule->getName() == name)
+  // Hash detects changes to callables as they have been merged into `mod`.
+  std::size_t moduleHash = (cacheable && hasCaptures) ? hashModuleOp(mod) : 0;
+
+  // Cache hit: same kernel, same target configuration, same module content.
+  if (cacheable && cachedModule->getName() == name &&
+      cachedModule->getMetadata().targetHash == targetHash &&
+      cachedModule->getMetadata().moduleHash == moduleHash) {
+    CUDAQ_INFO("Reusing cached module with name {} and hash ({}, {})", name,
+               targetHash, moduleHash);
     return cudaq::streamlinedLaunchModule(*cachedModule, rawArgs);
+  }
 
+  CUDAQ_INFO("Compiling module {}", name);
   mlir::OwningOpRef<ModuleOp> clone = mod.clone();
-  auto compiled = compileModuleImpl(name, clone.get(), rawArgs, true);
+  auto compiled =
+      compileModuleImpl(name, clone.get(), rawArgs, true, std::move(target));
   auto res = cudaq::streamlinedLaunchModule(compiled, rawArgs);
-  if (isCachable)
+  if (cacheable) {
+    CUDAQ_INFO("Caching module {} with hash ({}, {})", name, targetHash,
+               moduleHash);
+    compiled.setCacheKey(targetHash, moduleHash);
     *cachedModule = std::move(compiled);
+  }
   return res;
 }