QPU.cpp

/*******************************************************************************
 * Copyright (c) 2025 - 2026 NVIDIA Corporation & Affiliates.                  *
 * All rights reserved.                                                        *
 *                                                                             *
 * This source code and the accompanying materials are made available under    *
 * the terms of the Apache License 2.0 which accompanies this distribution.    *
 ******************************************************************************/

#include "QPU.h"
#include "common/ArgumentConversion.h"
#include "common/Environment.h"
#include "common/ExecutionContext.h"
#include "common/JIT.h"
#include "common/RuntimeMLIR.h"
#include "cudaq/Optimizer/Builder/Intrinsics.h"
#include "cudaq/Optimizer/Builder/Runtime.h"
#include "cudaq/Optimizer/CodeGen/OpenQASMEmitter.h"
#include "cudaq/Optimizer/CodeGen/Passes.h"
#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
#include "cudaq/Optimizer/Transforms/AddMetadata.h"
#include "cudaq/Optimizer/Transforms/Passes.h"
#include "mlir/ExecutionEngine/ExecutionEngine.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Export.h"
#include "mlir/Transforms/Passes.h"

using namespace mlir;

static void specializeKernel(const std::string &name, ModuleOp module,
                             const std::vector<void *> &rawArgs,
                             Type resultTy = {},
                             bool enablePythonCodegenDump = false,
                             bool isEntryPoint = true) {
  PassManager pm(module.getContext());
  cudaq::opt::ArgumentConverter argCon(name, module);
  argCon.gen(name, module, rawArgs);
  SmallVector<std::string> kernels;
  SmallVector<std::string> substs;
  for (auto *kInfo : argCon.getKernelSubstitutions()) {
    std::string kernName =
        cudaq::runtime::cudaqGenPrefixName + kInfo->getKernelName().str();
    kernels.emplace_back(kernName);
    std::string substBuff;
    llvm::raw_string_ostream ss(substBuff);
    ss << kInfo->getSubstitutionModule();
    substs.emplace_back(substBuff);
  }

  // Collect references for the argument synthesis.
  SmallVector<StringRef> kernelRefs{kernels.begin(), kernels.end()};
  SmallVector<StringRef> substRefs{substs.begin(), substs.end()};

  // Run a pass manager to specialize & optimize the kernel to be launched.
  pm.addPass(cudaq::opt::createArgumentSynthesisPass(
      kernelRefs, substRefs, /*changeSemantics=*/false));
  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
  pm.addPass(cudaq::opt::createLambdaLifting({.constantPropagation = true}));
  // We must inline these lambda calls before apply specialization as it does
  // not perform control/adjoint specialization across function call boundary.
  cudaq::opt::addAggressiveInlining(pm);
  pm.addPass(
      cudaq::opt::createApplySpecialization({.constantPropagation = true}));
  cudaq::opt::addAggressiveInlining(pm);
  pm.addPass(cudaq::opt::createDistributedDeviceCall());
  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
  if (resultTy && isEntryPoint) {
    // If we're expecting a result, then we want to call the .thunk function so
    // that the result is properly marshaled. Add the GKE pass to generate the
    // .thunk. At this point, the kernel should have been specialized so it has
    // an arity of 0.
    auto nullary = true;
    for (auto arg : rawArgs)
      if (!arg) {
        nullary = false;
        break;
      }
    pm.addPass(
        cudaq::opt::createGenerateKernelExecution({.positNullary = nullary}));
  }
  pm.addPass(createSymbolDCEPass());
  if (enablePythonCodegenDump) {
    module.getContext()->disableMultithreading();
    pm.enableIRPrinting();
  }
  if (failed(pm.run(module)))
    throw std::runtime_error("Could not successfully apply argument synth.");
}

/// Lowers \p module to LLVM code. The LLVM code will use "full QIR" as the
/// transport layer. If \p kernelName and \p args are provided, they will
/// specialize the selected entry-point kernel.
std::string cudaq::detail::lower_to_qir_llvm(const std::string &name,
                                             ModuleOp module,
                                             OpaqueArguments &args,
                                             const std::string &format) {
  ScopedTraceWithContext(cudaq::TIMING_JIT, "getQIR", name);
  // Translate the module to QIR transport layer (as LLVM code).
  cudaq::detail::mergeAllCallableClosures(module, name, args.getArgs());
  specializeKernel(name, module, args.getArgs());
  PassManager pm(module.getContext());
  cudaq::opt::addAggressiveInlining(pm);
  cudaq::opt::createTargetFinalizePipeline(pm);
  cudaq::opt::addAOTPipelineConvertToQIR(pm, format);
  if (failed(pm.run(module)))
    throw std::runtime_error("Conversion to " + format + " failed.");
  llvm::LLVMContext llvmContext;
  llvmContext.setOpaquePointers(false);
  std::unique_ptr<llvm::Module> llvmModule =
      translateModuleToLLVMIR(module, llvmContext);
  if (!llvmModule)
    return "{translation failed}";
  std::string result;
  llvm::raw_string_ostream os(result);
  llvmModule->print(os, nullptr);
  os.flush();
  return result;
}

/// Lowers \p module to `Open QASM 2`. The output will be a string of `Open
/// QASM` code. \p kernelName and \p args should be provided, as they will
/// specialize the selected entry-point kernel.
std::string cudaq::detail::lower_to_openqasm(const std::string &name,
                                             ModuleOp module,
                                             OpaqueArguments &args) {
  ScopedTraceWithContext(cudaq::TIMING_JIT, "getASM", name);
  // Translate module to OpenQASM2 transport layer.
  cudaq::detail::mergeAllCallableClosures(module, name, args.getArgs());
  specializeKernel(name, module, args.getArgs());
  auto *ctx = module.getContext();
  PassManager pm(ctx);
  cudaq::opt::createTargetFinalizePipeline(pm);
  cudaq::opt::createPipelineTransformsForPythonToOpenQASM(pm);
  cudaq::opt::addPipelineTranslateToOpenQASM(pm);
  const bool enablePrintMLIRBeforeAndAfterEachPass =
      cudaq::getEnvBool("CUDAQ_MLIR_PRINT_EACH_PASS", false);
  if (enablePrintMLIRBeforeAndAfterEachPass) {
    ctx->disableMultithreading();
    pm.enableIRPrinting();
  }
  if (failed(pm.run(module)))
    throw std::runtime_error("Conversion to OpenQASM failed.");
  std::string result;
  llvm::raw_string_ostream os(result);
  if (failed(cudaq::translateToOpenQASM(module, os)))
    return "{translation failed}";
  os.flush();
  return result;
}

/// Scan \p module and set flags in the current platform context accordingly.
static void updateExecutionContext(ModuleOp module) {
  auto *currentExecCtx = cudaq::getExecutionContext();
  if (!currentExecCtx)
    return;

  for (auto &artifact : module) {
    quake::detail::QuakeFunctionAnalysis analysis{&artifact};
    auto info = analysis.getAnalysisInfo();
    if (info.empty())
      continue;
    auto result = info[&artifact];
    if (result.hasConditionalsOnMeasure) {
      currentExecCtx->hasConditionalsOnMeasureResults = true;
      break;
    }
  }
}

static std::optional<cudaq::JitEngine> alreadyBuiltJITCode() {
  auto *currentExecCtx = cudaq::getExecutionContext();
  if (!currentExecCtx || !currentExecCtx->allowJitEngineCaching)
    return std::nullopt;
  return currentExecCtx->jitEng;
}

/// In a sample launch context, the (`JIT` compiled) execution engine 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.
static void cacheJITForPerformance(cudaq::JitEngine jit) {
  auto *currentExecCtx = cudaq::getExecutionContext();
  if (currentExecCtx && currentExecCtx->allowJitEngineCaching) {
    if (!currentExecCtx->jitEng)
      currentExecCtx->jitEng = jit;
  }
}

namespace {
struct PythonLauncher : public cudaq::ModuleLauncher {
  cudaq::CompiledKernel compileModule(const std::string &name, ModuleOp module,
                                      const std::vector<void *> &rawArgs,
                                      Type resultTy,
                                      bool isEntryPoint) override {
    // Check the ExecutionContext JIT cache first (used by cudaq.sample to
    // avoid recompiling on every shot).
    if (auto jit = alreadyBuiltJITCode())
      return cudaq::createCompiledKernel(*jit, name, /*hasResult=*/!!resultTy);

    ScopedTraceWithContext(cudaq::TIMING_LAUNCH,
                           "PythonLauncher::compileModule");
    const bool enablePythonCodegenDump =
        cudaq::getEnvBool("CUDAQ_PYTHON_CODEGEN_DUMP", false);

    std::string fullName = cudaq::runtime::cudaqGenPrefixName + name;

    // 1. Check that this call is sane.
    if (enablePythonCodegenDump)
      module.dump();
    auto funcOp = module.lookupSymbol<func::FuncOp>(fullName);
    if (!funcOp)
      throw std::runtime_error("no kernel named " + name + " found in module");

    // 2. Merge other modules (e.g., if there are device kernel calls).
    cudaq::detail::mergeAllCallableClosures(module, name, rawArgs);

    // Mark all newly merged kernels private.
    for (auto &op : module)
      if (auto f = dyn_cast<func::FuncOp>(op))
        if (f != funcOp)
          f.setPrivate();

    updateExecutionContext(module);

    // 3. Specialize the kernel (argument synthesis, optimization).
    CUDAQ_INFO("Run Argument Synth.\n");
    if (enablePythonCodegenDump)
      module.dump();
    specializeKernel(name, module, rawArgs, resultTy, enablePythonCodegenDump,
                     isEntryPoint);

    // 4. Lower to QIR and JIT compile.
    auto jit = cudaq::createQIRJITEngine(module, "qir:");
    cacheJITForPerformance(jit);

    return cudaq::createCompiledKernel(jit, name,
                                       /*hasResult=*/!!resultTy &&
                                           isEntryPoint);
  }
};
} // namespace

CUDAQ_REGISTER_TYPE(cudaq::ModuleLauncher, PythonLauncher, default)