Address PR comments

lib/Optimizer/Transforms/QuakeSynthesizer.cpp

@@ -455,11 +455,9 @@ class QuakeSynthesizer
     // Keep track of the stdVec sizes.
     std::vector<std::tuple<std::size_t, Type, std::uint64_t>> stdVecInfo;
 
-    for (auto iter : llvm::enumerate(arguments)) {
-      if (iter.index() < startingArgIdx)
-        continue;
-      auto argNum = iter.index();
-      auto argument = iter.value();
+    for (std::size_t argNum = startingArgIdx, end = arguments.size();
+         argNum < end; argNum++) {
+      auto argument = arguments[argNum];
       std::size_t offset = structLayout.second[argNum - startingArgIdx];
 
       // Get the argument type

python/runtime/cudaq/algorithms/py_vqe.cpp

@@ -63,14 +63,21 @@ bool isArgumentStdVec(MlirModule &module, const std::string &kernelName,
   return isa<cudaq::cc::StdvecType>(kernel.getArgument(argIdx).getType());
 }
 
-/// @brief Run `cudaq::observe` on the provided kernel and spin operator.
-observe_result pyObserve(py::object &kernel, spin_op &spin_operator,
-                         py::args args, const int shots,
-                         bool argMapperProvided = false) {
+/// @brief Return the kernel name and MLIR module for a kernel.
+static inline std::pair<std::string, MlirModule>
+getKernelNameAndModule(py::object &kernel) {
   if (py::hasattr(kernel, "compile"))
     kernel.attr("compile")();
   auto kernelName = kernel.attr("name").cast<std::string>();
   auto kernelMod = kernel.attr("module").cast<MlirModule>();
+  return std::make_pair(kernelName, kernelMod);
+}
+
+/// @brief Run `cudaq::observe` on the provided kernel and spin operator.
+observe_result pyObserve(py::object &kernel, spin_op &spin_operator,
+                         py::args args, const int shots,
+                         bool argMapperProvided = false) {
+  auto [kernelName, kernelMod] = getKernelNameAndModule(kernel);
   auto &platform = cudaq::get_platform();
   args = simplifiedValidateInputArguments(args);
   auto *argData = toOpaqueArgs(args, kernelMod, kernelName);
@@ -108,10 +115,7 @@ observe_result pyObserve(py::object &kernel, spin_op &spin_operator,
 /// implementation that requires the variation parameters to be the first
 /// argument in the kernel.
 static bool firstArgIsCompatibleWithRemoteVQE(py::object &kernel) {
-  if (py::hasattr(kernel, "compile"))
-    kernel.attr("compile")();
-  auto kernelName = kernel.attr("name").cast<std::string>();
-  auto kernelMod = kernel.attr("module").cast<MlirModule>();
+  auto [kernelName, kernelMod] = getKernelNameAndModule(kernel);
   auto kernelFunc = getKernelFuncOp(kernelMod, kernelName);
   if (kernelFunc.getNumArguments() < 1)
     return false;
@@ -132,11 +136,7 @@ pyVQE_remote_cpp(cudaq::quantum_platform &platform, py::object &kernel,
                  spin_op &hamiltonian, cudaq::optimizer &optimizer,
                  cudaq::gradient *gradient, py::function *argumentMapper,
                  const int n_params, const int shots) {
-
-  if (py::hasattr(kernel, "compile"))
-    kernel.attr("compile")();
-  auto kernelName = kernel.attr("name").cast<std::string>();
-  auto kernelMod = kernel.attr("module").cast<MlirModule>();
+  auto [kernelName, kernelMod] = getKernelNameAndModule(kernel);
   auto ctx = std::make_unique<ExecutionContext>("observe", /*shots=*/0);
   ctx->kernelName = kernelName;
   ctx->spin = &hamiltonian;

python/runtime/utils/PyRemoteSimulatorQPU.cpp

@@ -27,13 +27,7 @@ class PyRemoteSimulatorQPU : public cudaq::BaseRemoteSimulatorQPU {
                  cudaq::optimizer &optimizer, const int n_params,
                  const std::size_t shots) override {
     cudaq::ExecutionContext *executionContextPtr =
-        [&]() -> cudaq::ExecutionContext * {
-      std::scoped_lock<std::mutex> lock(m_contextMutex);
-      const auto iter = m_contexts.find(std::this_thread::get_id());
-      if (iter == m_contexts.end())
-        return nullptr;
-      return iter->second;
-    }();
+        getExecutionContextForMyThread();
 
     auto *wrapper = reinterpret_cast<const cudaq::ArgWrapper *>(kernelArgs);
     auto m_module = wrapper->mod;
@@ -70,13 +64,8 @@ class PyRemoteSimulatorQPU : public cudaq::BaseRemoteSimulatorQPU {
     auto *mlirContext = m_module->getContext();
 
     cudaq::ExecutionContext *executionContextPtr =
-        [&]() -> cudaq::ExecutionContext * {
-      std::scoped_lock<std::mutex> lock(m_contextMutex);
-      const auto iter = m_contexts.find(std::this_thread::get_id());
-      if (iter == m_contexts.end())
-        return nullptr;
-      return iter->second;
-    }();
+        getExecutionContextForMyThread();
+
     // Default context for a 'fire-and-ignore' kernel launch; i.e., no context
     // was set before launching the kernel. Use a static variable per thread to
     // set up a single-shot execution context for this case.
@@ -120,13 +109,8 @@ class PyNvcfSimulatorQPU : public cudaq::BaseNvcfSimulatorQPU {
     auto *mlirContext = m_module->getContext();
 
     cudaq::ExecutionContext *executionContextPtr =
-        [&]() -> cudaq::ExecutionContext * {
-      std::scoped_lock<std::mutex> lock(m_contextMutex);
-      const auto iter = m_contexts.find(std::this_thread::get_id());
-      if (iter == m_contexts.end())
-        return nullptr;
-      return iter->second;
-    }();
+        getExecutionContextForMyThread();
+
     // Default context for a 'fire-and-ignore' kernel launch; i.e., no context
     // was set before launching the kernel. Use a static variable per thread to
     // set up a single-shot execution context for this case.

runtime/common/BaseRemoteSimulatorQPU.h

@@ -53,6 +53,16 @@ class BaseRemoteSimulatorQPU : public cudaq::QPU {
   std::unique_ptr<mlir::MLIRContext> m_mlirContext;
   std::unique_ptr<cudaq::RemoteRuntimeClient> m_client;
 
+  /// @brief Return a pointer to the execution context for this thread. It will
+  /// return `nullptr` if it was not found in `m_contexts`.
+  cudaq::ExecutionContext *getExecutionContextForMyThread() {
+    std::scoped_lock<std::mutex> lock(m_contextMutex);
+    const auto iter = m_contexts.find(std::this_thread::get_id());
+    if (iter == m_contexts.end())
+      return nullptr;
+    return iter->second;
+  }
+
 public:
   BaseRemoteSimulatorQPU()
       : QPU(),
@@ -100,13 +110,7 @@ class BaseRemoteSimulatorQPU : public cudaq::QPU {
                  cudaq::optimizer &optimizer, const int n_params,
                  const std::size_t shots) override {
     cudaq::ExecutionContext *executionContextPtr =
-        [&]() -> cudaq::ExecutionContext * {
-      std::scoped_lock<std::mutex> lock(m_contextMutex);
-      const auto iter = m_contexts.find(std::this_thread::get_id());
-      if (iter == m_contexts.end())
-        return nullptr;
-      return iter->second;
-    }();
+        getExecutionContextForMyThread();
 
     if (executionContextPtr && executionContextPtr->name == "tracer")
       return;
@@ -135,13 +139,7 @@ class BaseRemoteSimulatorQPU : public cudaq::QPU {
         name, qpu_id, m_simName);
 
     cudaq::ExecutionContext *executionContextPtr =
-        [&]() -> cudaq::ExecutionContext * {
-      std::scoped_lock<std::mutex> lock(m_contextMutex);
-      const auto iter = m_contexts.find(std::this_thread::get_id());
-      if (iter == m_contexts.end())
-        return nullptr;
-      return iter->second;
-    }();
+        getExecutionContextForMyThread();
 
     if (executionContextPtr && executionContextPtr->name == "tracer") {
       return;
@@ -173,13 +171,7 @@ class BaseRemoteSimulatorQPU : public cudaq::QPU {
         name, qpu_id, m_simName);
 
     cudaq::ExecutionContext *executionContextPtr =
-        [&]() -> cudaq::ExecutionContext * {
-      std::scoped_lock<std::mutex> lock(m_contextMutex);
-      const auto iter = m_contexts.find(std::this_thread::get_id());
-      if (iter == m_contexts.end())
-        return nullptr;
-      return iter->second;
-    }();
+        getExecutionContextForMyThread();
 
     if (executionContextPtr && executionContextPtr->name == "tracer") {
       return;

runtime/cudaq/algorithms/vqe.h

@@ -15,6 +15,28 @@
 
 namespace cudaq {
 
+/// \brief This is an internal helper function to reduce duplicated code in the
+/// user-facing `vqe()` functions below. Users should not directly call this
+/// function.
+template <typename QuantumKernel, typename... Args,
+          typename = std::enable_if_t<
+              std::is_invocable_v<QuantumKernel, std::vector<double>, Args...>>>
+static inline optimization_result
+remote_vqe(cudaq::quantum_platform &platform, QuantumKernel &&kernel,
+           cudaq::spin_op &H, cudaq::optimizer &optimizer,
+           cudaq::gradient *gradient, const int n_params,
+           const std::size_t shots, Args &&...args) {
+  auto ctx = std::make_unique<ExecutionContext>("observe", shots);
+  ctx->kernelName = cudaq::getKernelName(kernel);
+  ctx->spin = &H;
+  platform.set_exec_ctx(ctx.get());
+  auto serializedArgsBuffer = serializeArgs(args...);
+  platform.launchVQE(ctx->kernelName, serializedArgsBuffer.data(), gradient, H,
+                     optimizer, n_params, shots);
+  platform.reset_exec_ctx();
+  return ctx->optResult.value_or(optimization_result{});
+}
+
 ///
 /// \brief Compute the minimal eigenvalue of \p H with VQE.
 ///
@@ -73,19 +95,9 @@ optimization_result vqe(QuantumKernel &&kernel, cudaq::spin_op H,
   }
 
   auto &platform = cudaq::get_platform();
-  if (platform.supports_remote_vqe()) {
-    auto ctx = std::make_unique<ExecutionContext>("observe", /*shots=*/0);
-    ctx->kernelName = cudaq::getKernelName(kernel);
-    ctx->spin = &H;
-    platform.set_exec_ctx(ctx.get());
-    auto serializedArgsBuffer = serializeArgs(args...);
-    platform.launchVQE(cudaq::getKernelName(kernel),
-                       /*kernelArgs=*/serializedArgsBuffer.data(),
-                       /*gradient=*/nullptr, H, optimizer, n_params,
-                       /*shots=*/0);
-    platform.reset_exec_ctx();
-    return ctx->optResult.value_or(optimization_result{});
-  }
+  if (platform.supports_remote_vqe())
+    return remote_vqe(platform, kernel, H, optimizer, /*gradient=*/nullptr,
+                      n_params, /*shots=*/0, args...);
 
   return optimizer.optimize(n_params, [&](const std::vector<double> &x,
                                           std::vector<double> &grad_vec) {
@@ -153,19 +165,9 @@ optimization_result vqe(std::size_t shots, QuantumKernel &&kernel,
   }
 
   auto &platform = cudaq::get_platform();
-  if (platform.supports_remote_vqe()) {
-    auto ctx = std::make_unique<ExecutionContext>("observe", /*shots=*/shots);
-    ctx->kernelName = cudaq::getKernelName(kernel);
-    ctx->spin = &H;
-    platform.set_exec_ctx(ctx.get());
-    auto serializedArgsBuffer = serializeArgs(args...);
-    platform.launchVQE(cudaq::getKernelName(kernel),
-                       serializedArgsBuffer.data(), /*gradient=*/nullptr, H,
-                       optimizer, n_params,
-                       /*shots=*/shots);
-    platform.reset_exec_ctx();
-    return ctx->optResult.value_or(optimization_result{});
-  }
+  if (platform.supports_remote_vqe())
+    return remote_vqe(platform, kernel, H, optimizer, /*gradient=*/nullptr,
+                      n_params, shots, args...);
 
   return optimizer.optimize(n_params, [&](const std::vector<double> &x,
                                           std::vector<double> &grad_vec) {
@@ -234,19 +236,9 @@ optimization_result vqe(QuantumKernel &&kernel, cudaq::gradient &gradient,
       "std::tuple<Args...>(std::vector<double>) ArgMapper function object.");
 
   auto &platform = cudaq::get_platform();
-  if (platform.supports_remote_vqe()) {
-    auto ctx = std::make_unique<ExecutionContext>("observe", /*shots=*/0);
-    ctx->kernelName = cudaq::getKernelName(kernel);
-    ctx->spin = &H;
-    platform.set_exec_ctx(ctx.get());
-    auto serializedArgsBuffer = serializeArgs(args...);
-    platform.launchVQE(cudaq::getKernelName(kernel),
-                       /*kernelArgs=*/serializedArgsBuffer.data(), &gradient, H,
-                       optimizer, n_params,
-                       /*shots=*/0);
-    platform.reset_exec_ctx();
-    return ctx->optResult.value_or(optimization_result{});
-  }
+  if (platform.supports_remote_vqe())
+    return remote_vqe(platform, kernel, H, optimizer, &gradient, n_params,
+                      /*shots=*/0, args...);
 
   auto requires_grad = optimizer.requiresGradients();
   // If there are additional arguments, we need to clone the gradient and