Clean up QPU launches and remove MLIR header deps

cudaq/include/cudaq/Target/CompileTarget.h

@@ -21,6 +21,17 @@ class CompileTarget {
   /// Hook to update the pass pipeline before compilation.
   virtual void updatePassPipeline(std::string &passPipeline) 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
+  /// pointer) is provided, then compilation will be skipped and all other
+  /// options in this class will be ignored.
+  ///
+  /// If this is `true`, the AOT-compiled module (if it exists) will be
+  /// discarded and compilation will start from scratch, according to the
+  /// options in this class.
+  bool overrideAOTCompilation = false;
+
   /// Resolved MLIR pass-pipeline and `codegen` settings.
   struct PipelineConfig {
     /// If set, override compilation pipeline with this string.

cudaq/tools/nvqpp/nvq++.in

@@ -724,6 +724,8 @@ NVQIR_LIBS="${NVQIR_LIBS}${NVQIR_SIMULATION_BACKEND}"
 # Add the MLIR-related libraries if requested. 
 if ${ENABLE_MLIR_LIB_LINKING}; then 
     LINKLIBS="${LINKLIBS} -lcudaq-mlir-runtime -lcudaq-builder"
+else
+	PREPROCESSOR_DEFINES="${PREPROCESSOR_DEFINES} -D CUDAQ_DISABLE_JIT_COMPILER"
 fi
 
 if ${ENABLE_REALTIME_LOWERING}; then

python/runtime/interop/CMakeLists.txt

@@ -15,8 +15,7 @@ target_include_directories(cudaq-python-interop PRIVATE
     ${Python3_INCLUDE_DIRS}
 )
 target_link_libraries(cudaq-python-interop
-    PRIVATE nanobind-static Python3::Module cudaq
-    PUBLIC cudaq-mlir-runtime-headers)
+    PRIVATE nanobind-static Python3::Module cudaq)
 install (FILES PythonCppInterop.h PythonCppInteropDecls.h DESTINATION include/cudaq/python/)
 
 install(TARGETS cudaq-python-interop EXPORT cudaq-python-interop-targets DESTINATION lib)

runtime/common/BaseRemoteRESTQPU.h

@@ -14,7 +14,6 @@
 #include "common/KernelExecution.h"
 #include "common/Resources.h"
 #include "common/ServerHelper.h"
-#include "cudaq_internal/compiler/Compiler.h"
 #include "nvqir/AnalysisScope.h"
 #include "nvqir/resourcecounter/ResourceCounterScope.h"
 #include "cudaq/Target/TargetConfig.h"
@@ -56,8 +55,6 @@ inline observe_result observeResultFromCounts(const observe_policy &policy,
 
 class BaseRemoteRESTQPU : public QPU {
 protected:
-  using Compiler = cudaq_internal::compiler::Compiler;
-
   /// The number of shots
   std::optional<int> nShots;
 
@@ -280,47 +277,6 @@ class BaseRemoteRESTQPU : public QPU {
     return target;
   }
 
-  /// @brief Build the list of kernel executions for the given module under
-  /// a specific sampling policy. Source modules are lowered through the
-  /// configured pass pipeline; pre-compiled modules are emitted directly.
-  /// The resolved kernel name is returned via @p kernelName.
-  template <typename Policy>
-  std::pair<std::string, std::vector<cudaq::KernelExecution>>
-  compileKernelExecutions(Policy &policy, const AnyModule &module,
-                          KernelArgs args) {
-    Compiler compiler(getCompileTarget(policy));
-    std::string kernelName;
-    std::optional<CompiledModule> compiled;
-    if (std::holds_alternative<SourceModule>(module)) {
-      const auto &src = std::get<SourceModule>(module);
-      kernelName = src.getName();
-      CUDAQ_INFO("launching remote rest kernel ({})", kernelName);
-
-      auto [moduleOp, context] = Compiler::loadQuakeCodeByName(kernelName);
-
-      compiled = compiler.runPassPipeline(kernelName, moduleOp, args, true,
-                                          std::move(context));
-      if constexpr (std::is_same_v<Policy, sample_policy>) {
-        if (compiler.hasWarnedNamedMeasurements())
-          policy.warnedNamedMeasurements = true;
-      }
-    } else {
-      compiled = std::get<CompiledModule>(module);
-      kernelName = compiled->getName();
-      CUDAQ_INFO("launching remote rest kernel via module ({})", kernelName);
-    }
-
-    auto codes = compiler.emitKernelExecutions(*compiled);
-
-    // Propagate metadata from the compiled artifact to the execution context.
-    if (auto ctx = getExecutionContext()) {
-      ctx->hasConditionalsOnMeasureResults =
-          compiled->getMetadata().hasConditionalsOnMeasureResults;
-    }
-
-    return {kernelName, codes};
-  }
-
   void completeLaunchKernel(const std::string &kernelName,
                             std::vector<cudaq::KernelExecution> &&codes) {
     auto executionContext = cudaq::getExecutionContext();
@@ -514,7 +470,7 @@ class BaseRemoteRESTQPU : public QPU {
   }
 
   async_observe_result
-  completeLaunchKernel(async_observe_policy &policy,
+  completeLaunchKernel(const async_observe_policy &policy,
                        const std::string &kernelName,
                        std::vector<cudaq::KernelExecution> &&codes) {
     std::size_t localShots = 1000;

runtime/common/CompiledModule.h

@@ -263,34 +263,13 @@ class FatQuakeModule {
   ArtifactsStore artifacts;
 };
 
-/// @brief A compiled MLIR module, ready for execution or code generation.
-///
-/// Contains any number of named compilation artifacts (we currently support
-/// JIT binaries, optimized MLIR modules, and pre-computed resource metrics)
-/// that result from the compilation of a Quake MLIR module.
-///
-/// This type does not depend on MLIR/LLVM — it only keeps type-erased / opaque
-/// pointers. Build instances with
-/// `cudaq_internal::compiler::CompiledModuleHelper`.
-class CompiledModule : public FatQuakeModule {
-public:
-  using CompilationMetadata = FatQuakeModule::CompilationMetadata;
-
-private:
-  friend class cudaq_internal::compiler::CompiledModuleHelper;
-
-  CompiledModule(std::string kernelName)
-      : FatQuakeModule(std::move(kernelName)) {}
-
-public:
-  CompiledModule() : FatQuakeModule(std::string{}) {}
-};
-
 /// Bundle of artifacts that define a CUDA-Q kernel to be compiled and executed.
 ///
 /// Contains either a `nvq++`-compiled function pointer or an MLIR module,
 /// depending on the provenance of the kernel.
 class SourceModule : public FatQuakeModule {
+  friend class cudaq_internal::compiler::CompiledModuleHelper;
+
 public:
   SourceModule(std::string kernelName)
       : FatQuakeModule(std::move(kernelName)) {}
@@ -309,10 +288,35 @@ class SourceModule : public FatQuakeModule {
   const void *getMlirOpaqueModulePtr() const;
 };
 
-// TODO: remove once C++ launch can be cleanly split into compilation + launch.
-// Used by unifiedLaunchModule to compile kernels if they have not been compiled
-// before. In the future, unifiedLaunchModule should only accept compiled
-// modules.
+/// @brief A compiled MLIR module, ready for execution or code generation.
+///
+/// Contains any number of named compilation artifacts (we currently support
+/// JIT binaries, optimized MLIR modules, and pre-computed resource metrics)
+/// that result from the compilation of a Quake MLIR module.
+///
+/// This type does not depend on MLIR/LLVM — it only keeps type-erased / opaque
+/// pointers. Build instances with
+/// `cudaq_internal::compiler::CompiledModuleHelper`.
+class CompiledModule : public FatQuakeModule {
+public:
+  using CompilationMetadata = FatQuakeModule::CompilationMetadata;
+
+private:
+  friend class cudaq_internal::compiler::CompiledModuleHelper;
+
+  CompiledModule(std::string kernelName)
+      : FatQuakeModule(std::move(kernelName)) {}
+
+public:
+  // The choice of constructors is intentionally limited to:
+  //  - empty compiled modules for default construction
+  //  - compiled module from a source module to explicitly bypass the compiler
+  // For any other use case, you should go through the factory methods in
+  // `CompiledModuleHelper`.
+  CompiledModule() : FatQuakeModule(std::string{}) {}
+  explicit CompiledModule(SourceModule src) : FatQuakeModule(std::move(src)) {}
+};
+
 using AnyModule = std::variant<SourceModule, CompiledModule>;
 
 } // namespace cudaq

runtime/common/ExecutionContext.h

@@ -205,8 +205,8 @@ void resetExecutionContext();
 /// @brief Execute the given function within the given policy and execution
 /// context.
 template <typename Policy, typename Callable, typename... Args>
-auto with_policy_and_ctx(Policy &policy, ExecutionContext &ctx, Callable &&f,
-                         Args &&...args)
+auto with_policy_and_ctx(const Policy &policy, ExecutionContext &ctx,
+                         Callable &&f, Args &&...args)
     -> std::invoke_result_t<Callable, Args...> {
 
   // Save the outer execution context (if any) so we can restore it after.

runtime/cudaq/CMakeLists.txt

@@ -66,7 +66,6 @@ if (CUDA_FOUND)
       cudaq-common
       cudaq-nlopt
       cudaq-ensmallen
-      cudaq-mlir-runtime-headers # Public until compilation is moved out of the QPUs
       cudaq-logger
       CUDAQTargetConfig
     PRIVATE
@@ -94,7 +93,6 @@ else()
       cudaq-common
       cudaq-nlopt
       cudaq-ensmallen
-      cudaq-mlir-runtime-headers # Public until compilation is moved out of the QPUs
       cudaq-logger
       CUDAQTargetConfig
     PRIVATE

runtime/cudaq/algorithms/launch.h

@@ -9,6 +9,7 @@
 
 #pragma once
 
+#include "common/CompiledModule.h"
 #include "common/ExecutionContext.h"
 #include "common/KernelArgs.h"
 #include "cudaq/platform.h"
@@ -17,13 +18,34 @@
 #include "cudaq/runtime/logger/logger.h"
 #include "cudaq/utils/cudaq_utils.h"
 #include <stdexcept>
+#include <variant>
+
+namespace cudaq_internal::compiler {
+template <typename Policy>
+cudaq::CompiledModule
+compileModule(const Policy &policy,
+              std::unique_ptr<cudaq::CompileTarget> target,
+              const cudaq::SourceModule &src, cudaq::KernelArgs args,
+              bool isEntryPoint = true);
+} // namespace cudaq_internal::compiler
+
+// If JIT compilation is disabled, make compilation a no-op. QPUs may throw an
+// error if they expect a JIT-compiled module.
+#ifdef CUDAQ_DISABLE_JIT_COMPILER
+template <typename Policy>
+cudaq::CompiledModule cudaq_internal::compiler::compileModule(
+    const Policy &policy, std::unique_ptr<cudaq::CompileTarget> target,
+    const cudaq::SourceModule &src, cudaq::KernelArgs args, bool isEntryPoint) {
+  return cudaq::CompiledModule{src};
+}
+#endif
 
 namespace cudaq {
 namespace detail {
 
 /// @brief Execute the given function within the given execution context.
 template <typename Policy, typename Callable, typename... Args>
-auto launch(Policy &policy, std::size_t qpu_id, ExecutionContext &ctx,
+auto launch(const Policy &policy, std::size_t qpu_id, ExecutionContext &ctx,
             quantum_platform &platform, Callable &&f, Args &&...args)
     -> Policy::result_type {
 
@@ -51,7 +73,20 @@ auto launch(Policy &policy, std::size_t qpu_id, ExecutionContext &ctx,
   auto &qpu = platform.getQPU(qpu_id);
   ctx.executeKernelApi = [&qpu, &result, &policy](const AnyModule &module,
                                                   const KernelArgs &args) {
-    result = qpu.launchKernel(policy, module, args);
+    CompiledModule compiled;
+    if (const auto *source = std::get_if<SourceModule>(&module)) {
+      std::unique_ptr<cudaq::CompileTarget> target;
+      if constexpr (requires { policy.inner; }) {
+        target = cudaq::get_compile_target(policy.inner);
+      } else {
+        target = cudaq::get_compile_target(policy);
+      }
+      compiled = cudaq_internal::compiler::compileModule(
+          policy, std::move(target), *source, args);
+    } else {
+      compiled = std::get<CompiledModule>(module);
+    }
+    result = qpu.launchKernel(policy, compiled, args);
   };
 
   if constexpr (requires { policy.inner; })

runtime/cudaq/algorithms/sample/policy.h

@@ -37,7 +37,7 @@ struct sample_policy {
 
   /// @brief Flag to indicate that a warning about named measurement registers
   /// in sampling context has already been emitted.
-  bool warnedNamedMeasurements = false;
+  mutable bool warnedNamedMeasurements = false;
 
   /// @brief A vector containing information about how to reorder the global
   /// register after execution. Empty means no reordering.

runtime/cudaq/platform.h

@@ -47,14 +47,18 @@ std::unique_ptr<cudaq::CompileTarget> get_compile_target(const Policy &policy) {
   return getQuantumPlatformInternal()->getCompileTarget(policy);
 }
 
-/// Get the default compile target configuration used when JITing for Python.
+/// Get the default compile target configuration
+///
+/// This is suitable for local simulators, i.e. it will use
+/// AOT-compiled modules as-is if they exist, and otherwise JIT-compile the
+/// module as appropriate for a Python kernel.
 std::unique_ptr<cudaq::CompileTarget>
-getDefaultPythonCompileTarget(const sample_policy &policy);
+getDefaultCompileTarget(const sample_policy &policy);
 std::unique_ptr<cudaq::CompileTarget>
-getDefaultPythonCompileTarget(const observe_policy &policy);
+getDefaultCompileTarget(const observe_policy &policy);
 std::unique_ptr<cudaq::CompileTarget>
-getDefaultPythonCompileTarget(const other_policies &policy,
-                              ExecutionContext *context);
+getDefaultCompileTarget(const other_policies &policy,
+                        ExecutionContext *context);
 
 // Declare this function, implemented elsewhere
 std::string getQIR(const std::string &);