New launch kernel for python

python/runtime/cudaq/algorithms/py_state.cpp

@@ -75,6 +75,7 @@ class PyRemoteSimulationState : public RemoteSimulationState {
                           std::size_t size, std::size_t returnOffset)
       : argsData(argsDataToOwn), kernelMod(args.mod) {
     this->kernelName = in_kernelName;
+    this->args = argsData->getArgs();
   }
 
   void execute() const override {
@@ -98,11 +99,6 @@ class PyRemoteSimulationState : public RemoteSimulationState {
     }
   }
 
-  std::optional<std::pair<std::string, std::vector<void *>>>
-  getKernelInfo() const override {
-    return std::make_pair(kernelName, argsData->getArgs());
-  }
-
   std::complex<double> overlap(const cudaq::SimulationState &other) override {
     const auto &otherState =
         dynamic_cast<const PyRemoteSimulationState &>(other);
@@ -121,7 +117,7 @@ class PyRemoteSimulationState : public RemoteSimulationState {
     return context.overlapResult.value();
   }
 
-  ~PyRemoteSimulationState() { delete argsData; }
+  virtual ~PyRemoteSimulationState() override { delete argsData; }
 };
 
 /// @brief Run `cudaq::get_state` for remote execution targets on the provided
@@ -140,6 +136,39 @@ state pyGetStateRemote(py::object kernel, py::args args) {
                                            size, returnOffset));
 }
 
+/// @brief Python implementation of the `QPUState`.
+// Note: Python kernel arguments are wrapped hence need to be unwrapped
+// accordingly.
+class PyQPUState : public QPUState {
+  // Holder of args data for clean-up.
+  cudaq::OpaqueArguments *argsData;
+
+public:
+  PyQPUState(const std::string &in_kernelName,
+             cudaq::OpaqueArguments *argsDataToOwn)
+      : argsData(argsDataToOwn) {
+    this->kernelName = in_kernelName;
+    this->args = argsData->getArgs();
+  }
+
+  virtual ~PyQPUState() override { delete argsData; }
+};
+
+/// @brief Run `cudaq::get_state` for qpu targets on the provided
+/// kernel and args
+state pyGetStateQPU(py::object kernel, py::args args) {
+  if (py::hasattr(kernel, "compile"))
+    kernel.attr("compile")();
+
+  auto kernelName = kernel.attr("name").cast<std::string>();
+  args = simplifiedValidateInputArguments(args);
+  auto kernelMod = kernel.attr("module").cast<MlirModule>();
+  auto *argData = toOpaqueArgs(args, kernelMod, kernelName);
+  auto [argWrapper, size, returnOffset] =
+      pyCreateNativeKernel(kernelName, kernelMod, *argData);
+  return state(new PyQPUState(kernelName, argData));
+}
+
 state pyGetStateLibraryMode(py::object kernel, py::args args) {
   return details::extractState([&]() mutable {
     if (0 == args.size())
@@ -671,6 +700,8 @@ index pair.
           return pyGetStateRemote(kernel, args);
         if (holder.getTarget().name == "orca-photonics")
           return pyGetStateLibraryMode(kernel, args);
+        if (holder.getTarget().is_remote() || holder.getTarget().is_emulated())
+          return pyGetStateQPU(kernel, args);
         return pyGetState(kernel, args);
       },
       R"#(Return the :class:`State` of the system after execution of the provided `kernel`.

python/runtime/cudaq/platform/py_alt_launch_kernel.cpp

@@ -337,15 +337,31 @@ pyAltLaunchKernelBase(const std::string &name, MlirModule module,
   if (launch) {
     auto &platform = cudaq::get_platform();
     auto uReturnOffset = static_cast<std::uint64_t>(returnOffset);
-    if (platform.is_remote() || platform.is_emulated()) {
+    auto isRemoteSimulator =
+        platform.get_remote_capabilities().isRemoteSimulator;
+    auto isQuantumDevice =
+        !isRemoteSimulator && (platform.is_remote() || platform.is_emulated());
+
+    if (isRemoteSimulator) {
+      // Remote simulator - use altLaunchKernel to support returning values.
+      // TODO: after cudaq::run support this should be merged with the quantum
+      // device case.
       auto *wrapper = new cudaq::ArgWrapper{mod, names, rawArgs};
       auto dynamicResult = cudaq::altLaunchKernel(
           name.c_str(), thunk, reinterpret_cast<void *>(wrapper), size,
           uReturnOffset);
       if (dynamicResult.data_buffer || dynamicResult.size)
         throw std::runtime_error("not implemented: support dynamic results");
       delete wrapper;
+    } else if (isQuantumDevice) {
+      // Quantum devices or their emulation - we can use streamlinedLaunchKernel
+      // as quantum platform do not support direct returns.
+      auto dynamicResult =
+          cudaq::streamlinedLaunchKernel(name.c_str(), runtimeArgs.getArgs());
+      if (dynamicResult.data_buffer || dynamicResult.size)
+        throw std::runtime_error("not implemented: support dynamic results");
     } else {
+      // Local simulator - use altLaunchKernel with the thunk function.
       auto dynamicResult = cudaq::altLaunchKernel(name.c_str(), thunk, rawArgs,
                                                   size, uReturnOffset);
       if (dynamicResult.data_buffer || dynamicResult.size)

python/runtime/utils/PyFermioniqRESTQPU.cpp

@@ -7,6 +7,7 @@
  ******************************************************************************/
 
 #include "common/ArgumentWrapper.h"
+#include "cudaq/Optimizer/InitAllDialects.h"
 #include "cudaq/platform/fermioniq/FermioniqBaseQPU.h"
 
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
@@ -27,41 +28,69 @@ void registerLLVMDialectTranslation(MLIRContext *context);
 namespace cudaq {
 
 class PyFermioniqRESTQPU : public cudaq::FermioniqBaseQPU {
+private:
+  /// Creates new context without mlir initialization.
+  MLIRContext *createContext() {
+    DialectRegistry registry;
+    cudaq::opt::registerCodeGenDialect(registry);
+    cudaq::registerAllDialects(registry);
+    auto context = new MLIRContext(registry);
+    context->loadAllAvailableDialects();
+    registerLLVMDialectTranslation(*context);
+    return context;
+  }
+
 protected:
   std::tuple<ModuleOp, MLIRContext *, void *>
   extractQuakeCodeAndContext(const std::string &kernelName,
                              void *data) override {
+    auto [mod, ctx] = extractQuakeCodeAndContextImpl(kernelName);
+    void *updatedArgs = nullptr;
+    if (data) {
+      auto *wrapper = reinterpret_cast<cudaq::ArgWrapper *>(data);
+      updatedArgs = wrapper->rawArgs;
+    }
+    return {mod, ctx, updatedArgs};
+  }
+
+  std::tuple<ModuleOp, MLIRContext *>
+  extractQuakeCodeAndContextImpl(const std::string &kernelName) {
 
-    auto *wrapper = reinterpret_cast<cudaq::ArgWrapper *>(data);
-    auto m_module = wrapper->mod;
-    auto callableNames = wrapper->callableNames;
+    cudaq::info("extract quake code\n");
+
+    MLIRContext *context = createContext();
 
-    auto *context = m_module->getContext();
     static bool initOnce = [&] {
       registerToQIRTranslation();
       registerToOpenQASMTranslation();
       registerToIQMJsonTranslation();
-      registerLLVMDialectTranslation(*context);
       return true;
     }();
     (void)initOnce;
 
+    // Get the quake representation of the kernel
+    auto quakeCode = cudaq::get_quake_by_name(kernelName);
+    auto m_module = parseSourceString<ModuleOp>(quakeCode, context);
+    if (!m_module)
+      throw std::runtime_error("module cannot be parsed");
+
     // Here we have an opportunity to run any passes that are
     // specific to python before the rest of the RemoteRESTQPU workflow
-    auto cloned = m_module.clone();
+    auto cloned = m_module->clone();
     PassManager pm(cloned.getContext());
-    pm.addNestedPass<func::FuncOp>(cudaq::opt::createPySynthCallableBlockArgs(
-        SmallVector<StringRef>(callableNames.begin(), callableNames.end())));
+
+    pm.addPass(cudaq::opt::createLambdaLiftingPass());
     cudaq::opt::addAggressiveEarlyInlining(pm);
-    pm.addPass(mlir::createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(cudaq::opt::createClassicalMemToReg());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
     pm.addNestedPass<mlir::func::FuncOp>(
         cudaq::opt::createUnwindLoweringPass());
-    pm.addPass(mlir::createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
     pm.addPass(cudaq::opt::createApplyOpSpecializationPass());
     pm.addPass(createInlinerPass());
     pm.addPass(cudaq::opt::createExpandMeasurementsPass());
-    pm.addPass(createCanonicalizerPass());
-    pm.addPass(createCSEPass());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(createCSEPass());
     if (failed(pm.run(cloned)))
       throw std::runtime_error(
           "Failure to synthesize callable block arguments in PyRemoteRESTQPU ");
@@ -75,8 +104,10 @@ class PyFermioniqRESTQPU : public cudaq::FermioniqBaseQPU {
     // The remote rest qpu workflow will need the module string in
     // the internal registry.
     __cudaq_deviceCodeHolderAdd(kernelName.c_str(), moduleStr.c_str());
-    return std::make_tuple(cloned, context, wrapper->rawArgs);
+    return std::make_tuple(cloned, context);
   }
+
+  void cleanupContext(MLIRContext *context) override { delete context; }
 };
 } // namespace cudaq
 

python/runtime/utils/PyRemoteRESTQPU.cpp

@@ -9,6 +9,7 @@
 #include "common/ArgumentWrapper.h"
 #include "common/BaseRemoteRESTQPU.h"
 #include "common/RuntimeMLIRCommonImpl.h"
+#include "cudaq/Optimizer/InitAllDialects.h"
 
 // [RFC]:
 // The RemoteRESTQPU implementation that is now split across several files needs
@@ -57,40 +58,65 @@ TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
 // implement some core functionality here in PyRemoteRESTQPU so we don't load
 // twice
 class PyRemoteRESTQPU : public cudaq::BaseRemoteRESTQPU {
+private:
+  /// Creates new context without mlir initialization.
+  MLIRContext *createContext() {
+    DialectRegistry registry;
+    cudaq::opt::registerCodeGenDialect(registry);
+    cudaq::registerAllDialects(registry);
+    auto context = new MLIRContext(registry);
+    context->loadAllAvailableDialects();
+    registerLLVMDialectTranslation(*context);
+    return context;
+  }
+
 protected:
   std::tuple<ModuleOp, MLIRContext *, void *>
   extractQuakeCodeAndContext(const std::string &kernelName,
                              void *data) override {
+    auto [mod, ctx] = extractQuakeCodeAndContextImpl(kernelName);
+    void *updatedArgs = nullptr;
+    if (data) {
+      auto *wrapper = reinterpret_cast<cudaq::ArgWrapper *>(data);
+      updatedArgs = wrapper->rawArgs;
+    }
+    return {mod, ctx, updatedArgs};
+  }
 
-    auto *wrapper = reinterpret_cast<cudaq::ArgWrapper *>(data);
-    auto m_module = wrapper->mod;
-    auto callableNames = wrapper->callableNames;
+  std::tuple<ModuleOp, MLIRContext *>
+  extractQuakeCodeAndContextImpl(const std::string &kernelName) {
+
+    MLIRContext *context = createContext();
 
-    auto *context = m_module->getContext();
     static bool initOnce = [&] {
       registerToQIRTranslation();
       registerToOpenQASMTranslation();
       registerToIQMJsonTranslation();
-      registerLLVMDialectTranslation(*context);
       return true;
     }();
     (void)initOnce;
 
+    // Get the quake representation of the kernel
+    auto quakeCode = cudaq::get_quake_by_name(kernelName);
+    auto m_module = parseSourceString<ModuleOp>(quakeCode, context);
+    if (!m_module)
+      throw std::runtime_error("module cannot be parsed");
+
     // Here we have an opportunity to run any passes that are
     // specific to python before the rest of the RemoteRESTQPU workflow
-    auto cloned = m_module.clone();
+    auto cloned = m_module->clone();
     PassManager pm(cloned.getContext());
-    pm.addNestedPass<func::FuncOp>(cudaq::opt::createPySynthCallableBlockArgs(
-        SmallVector<StringRef>(callableNames.begin(), callableNames.end())));
+
+    pm.addPass(cudaq::opt::createLambdaLiftingPass());
     cudaq::opt::addAggressiveEarlyInlining(pm);
-    pm.addPass(mlir::createCanonicalizerPass());
-    pm.addNestedPass<mlir::func::FuncOp>(
-        cudaq::opt::createUnwindLoweringPass());
-    pm.addPass(mlir::createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(cudaq::opt::createClassicalMemToReg());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindLoweringPass());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
     pm.addPass(cudaq::opt::createApplyOpSpecializationPass());
     pm.addPass(createInlinerPass());
-    pm.addPass(createCanonicalizerPass());
-    pm.addPass(createCSEPass());
+    pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+    pm.addNestedPass<func::FuncOp>(createCSEPass());
     if (failed(pm.run(cloned)))
       throw std::runtime_error(
           "Failure to synthesize callable block arguments in PyRemoteRESTQPU ");
@@ -103,8 +129,10 @@ class PyRemoteRESTQPU : public cudaq::BaseRemoteRESTQPU {
     // The remote rest qpu workflow will need the module string in
     // the internal registry.
     __cudaq_deviceCodeHolderAdd(kernelName.c_str(), moduleStr.c_str());
-    return std::make_tuple(cloned, context, wrapper->rawArgs);
+    return std::make_tuple(cloned, context);
   }
+
+  void cleanupContext(MLIRContext *context) override { delete context; }
 };
 } // namespace cudaq
 

python/tests/backends/test_IQM.py

@@ -201,6 +201,86 @@ def test_IQM_state_preparation_builder():
     assert assert_close(counts["11"], 0., 2)
 
 
+def test_IQM_state_synthesis_from_simulator():
+
+    @cudaq.kernel
+    def kernel(state: cudaq.State):
+        qubits = cudaq.qvector(state)
+
+    state = cudaq.State.from_data(
+        np.array([1. / np.sqrt(2.), 1. / np.sqrt(2.), 0., 0.], dtype=complex))
+
+    counts = cudaq.sample(kernel, state)
+    print(counts)
+    assert "00" in counts
+    assert "10" in counts
+    assert assert_close(counts["01"], 0., 2)
+    assert assert_close(counts["11"], 0., 2)
+
+    synthesized = cudaq.synthesize(kernel, state)
+    counts = cudaq.sample(synthesized)
+    assert '00' in counts
+    assert '10' in counts
+    assert assert_close(counts["01"], 0., 2)
+    assert assert_close(counts["11"], 0., 2)
+
+
+def test_IQM_state_synthesis_from_simulator_builder():
+
+    kernel, state = cudaq.make_kernel(cudaq.State)
+    qubits = kernel.qalloc(state)
+
+    state = cudaq.State.from_data(
+        np.array([1. / np.sqrt(2.), 1. / np.sqrt(2.), 0., 0.], dtype=complex))
+
+    counts = cudaq.sample(kernel, state)
+    assert "00" in counts
+    assert "10" in counts
+    assert assert_close(counts["01"], 0., 2)
+    assert assert_close(counts["11"], 0., 2)
+
+
+def test_IQM_state_synthesis():
+
+    @cudaq.kernel
+    def init(n: int):
+        q = cudaq.qvector(n)
+        x(q[0])
+
+    @cudaq.kernel
+    def kernel(s: cudaq.State):
+        q = cudaq.qvector(s)
+        x(q[1])
+
+    s = cudaq.get_state(init, 2)
+    s = cudaq.get_state(kernel, s)
+    counts = cudaq.sample(kernel, s)
+    assert '10' in counts
+    assert assert_close(counts["00"], 0., 2)
+    assert assert_close(counts["01"], 0., 2)
+    assert assert_close(counts["11"], 0., 2)
+
+
+def test_IQM_state_synthesis_builder():
+
+    init, n = cudaq.make_kernel(int)
+    qubits = init.qalloc(n)
+    init.x(qubits[0])
+
+    s = cudaq.get_state(init, 2)
+
+    kernel, state = cudaq.make_kernel(cudaq.State)
+    qubits = kernel.qalloc(state)
+    kernel.x(qubits[1])
+
+    s = cudaq.get_state(kernel, s)
+    counts = cudaq.sample(kernel, s)
+    assert '10' in counts
+    assert assert_close(counts["00"], 0., 2)
+    assert assert_close(counts["01"], 0., 2)
+    assert assert_close(counts["11"], 0., 2)
+
+
 def test_exp_pauli():
 
     @cudaq.kernel