Follow-up to #2394 - handle remote REST QPU also (#2417)

Author: bmhowe23

runtime/common/BaseRemoteRESTQPU.h

@@ -657,6 +657,7 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
       localShots = executionContext->shots;
 
     executor->setShots(localShots);
+    bool isObserve = executionContext && executionContext->name == "observe";
 
     // If emulation requested, then just grab the function
     // and invoke it with the simulator
@@ -669,7 +670,7 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
       // Launch the execution of the simulated jobs asynchronously
       future = cudaq::details::future(std::async(
           std::launch::async,
-          [&, codes, localShots, kernelName, seed,
+          [&, codes, localShots, kernelName, seed, isObserve,
            reorderIdx = executionContext->reorderIdx,
            localJIT = std::move(jitEngines)]() mutable -> cudaq::sample_result {
             std::vector<cudaq::ExecutionResult> results;
@@ -680,7 +681,7 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
 
             bool hasConditionals =
                 cudaq::kernelHasConditionalFeedback(kernelName);
-            if (hasConditionals && codes.size() > 1)
+            if (hasConditionals && isObserve)
               throw std::runtime_error("error: spin_ops not yet supported with "
                                        "kernels containing conditionals");
             if (hasConditionals) {
@@ -710,18 +711,15 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
               }
             }
 
-            bool isObserve =
-                executionContext && executionContext->name == "observe";
             for (std::size_t i = 0; i < codes.size(); i++) {
               cudaq::ExecutionContext context("sample", localShots);
               context.reorderIdx = reorderIdx;
               cudaq::getExecutionManager()->setExecutionContext(&context);
               invokeJITKernelAndRelease(localJIT[i], kernelName);
               cudaq::getExecutionManager()->resetExecutionContext();
 
-              // If there are multiple codes, this is likely a spin_op.
-              // If so, use the code name instead of the global register.
-              if (isObserve || (codes.size() > 1)) {
+              if (isObserve) {
+                // Use the code name instead of the global register.
                 results.emplace_back(context.result.to_map(), codes[i].name);
                 results.back().sequentialData =
                     context.result.sequential_data();
@@ -743,7 +741,7 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
       // Allow developer to disable remote sending (useful for debugging IR)
       if (getEnvBool("DISABLE_REMOTE_SEND", false))
         return;
-      future = executor->execute(codes);
+      future = executor->execute(codes, isObserve);
     }
 
     // Keep this asynchronous if requested

runtime/common/BraketExecutor.h

@@ -78,8 +78,8 @@ class BraketExecutor : public Executor {
   ~BraketExecutor() = default;
 
   /// @brief Execute the provided Braket task
-  details::future
-  execute(std::vector<KernelExecution> &codesToExecute) override;
+  details::future execute(std::vector<KernelExecution> &codesToExecute,
+                          bool isObserve) override;
 
   /// @brief Set the server helper
   void setServerHelper(ServerHelper *helper) override;

runtime/common/Executor.cpp

@@ -10,8 +10,8 @@
 #include "common/Logger.h"
 
 namespace cudaq {
-details::future
-Executor::execute(std::vector<KernelExecution> &codesToExecute) {
+details::future Executor::execute(std::vector<KernelExecution> &codesToExecute,
+                                  bool isObserve) {
 
   serverHelper->setShots(shots);
 
@@ -53,7 +53,7 @@ Executor::execute(std::vector<KernelExecution> &codesToExecute) {
 
   config.insert({"shots", std::to_string(shots)});
   std::string name = serverHelper->name();
-  return details::future(ids, name, config);
+  return details::future(ids, name, config, isObserve);
 }
 } // namespace cudaq
 

runtime/common/Executor.h

@@ -41,7 +41,8 @@ class Executor : public registry::RegisteredType<Executor> {
 
   /// @brief Execute the provided quantum codes and return a future object
   /// The caller can make this synchronous by just immediately calling .get().
-  virtual details::future execute(std::vector<KernelExecution> &codesToExecute);
+  virtual details::future execute(std::vector<KernelExecution> &codesToExecute,
+                                  bool isObserve = false);
 };
 
 } // namespace cudaq

runtime/common/Future.cpp

@@ -41,9 +41,8 @@ sample_result future::get() {
     }
     auto c = serverHelper->processResults(resultResponse, id.first);
 
-    // If there are multiple jobs, this is likely a spin_op.
-    // If so, use the job name instead of the global register.
-    if (jobs.size() > 1) {
+    if (isObserve) {
+      // Use the job name instead of the global register.
       results.emplace_back(c.to_map(), id.second);
       results.back().sequentialData = c.sequential_data();
     } else {
@@ -75,6 +74,7 @@ future &future::operator=(future &other) {
   jobs = other.jobs;
   qpuName = other.qpuName;
   serverConfig = other.serverConfig;
+  isObserve = other.isObserve;
   if (other.wrapsFutureSampling) {
     wrapsFutureSampling = true;
     inFuture = std::move(other.inFuture);
@@ -86,6 +86,7 @@ future &future::operator=(future &&other) {
   jobs = other.jobs;
   qpuName = other.qpuName;
   serverConfig = other.serverConfig;
+  isObserve = other.isObserve;
   if (other.wrapsFutureSampling) {
     wrapsFutureSampling = true;
     inFuture = std::move(other.inFuture);
@@ -102,6 +103,7 @@ std::ostream &operator<<(std::ostream &os, future &f) {
   j["jobs"] = f.jobs;
   j["qpu"] = f.qpuName;
   j["config"] = f.serverConfig;
+  j["isObserve"] = f.isObserve;
   os << j.dump(4);
   return os;
 }
@@ -117,6 +119,7 @@ std::istream &operator>>(std::istream &is, future &f) {
   f.jobs = j["jobs"].get<std::vector<future::Job>>();
   f.qpuName = j["qpu"].get<std::string>();
   f.serverConfig = j["config"].get<std::map<std::string, std::string>>();
+  f.isObserve = j["isObserve"].get<bool>();
   return is;
 }
 

runtime/common/Future.h

@@ -49,6 +49,9 @@ class future {
   std::future<sample_result> inFuture;
   bool wrapsFutureSampling = false;
 
+  /// @brief Whether or not this is in support of an "observe" call
+  bool isObserve = false;
+
 public:
   /// @brief The constructor
   future() = default;
@@ -69,6 +72,11 @@ class future {
          std::map<std::string, std::string> &config)
       : jobs(_jobs), qpuName(qpuNameIn), serverConfig(config) {}
 
+  future(std::vector<Job> &_jobs, std::string &qpuNameIn,
+         std::map<std::string, std::string> &config, bool isObserve)
+      : jobs(_jobs), qpuName(qpuNameIn), serverConfig(config),
+        isObserve(isObserve) {}
+
   future &operator=(future &other);
   future &operator=(future &&other);
 

runtime/cudaq/platform/default/rest/helpers/braket/BraketExecutor.cpp

@@ -173,7 +173,8 @@ void BraketExecutor::setOutputNames(const KernelExecution &codeToExecute,
 }
 
 details::future
-BraketExecutor::execute(std::vector<KernelExecution> &codesToExecute) {
+BraketExecutor::execute(std::vector<KernelExecution> &codesToExecute,
+                        bool isObserve) {
   auto [dummy1, dummy2, messages] = checkHelperAndCreateJob(codesToExecute);
 
   std::string const defaultBucket = defaultBucketFuture.get();

runtime/cudaq/platform/orca/OrcaExecutor.h

@@ -17,8 +17,8 @@ namespace cudaq {
 /// API.
 class OrcaExecutor : public Executor {
 public:
-  details::future
-  execute(std::vector<KernelExecution> &codesToExecute) override {
+  details::future execute(std::vector<KernelExecution> &codesToExecute,
+                          bool isObserve) override {
     throw std::runtime_error(
         "ORCA backend does not support executing arbitrary kernels");
   }