State synthesis for quantum devices

include/cudaq/Optimizer/Dialect/CC/CCOps.td

@@ -898,6 +898,26 @@ def cc_AddressOfOp : CCOp<"address_of", [Pure,
   }];
 }
 
+def cc_GetStateOp : CCOp<"get_state", [Pure] > {
+  let summary = "Get state from kernel with the provided name.";
+  let description = [{
+    This operation is created by argument synthesis of state pointer arguments
+    for quantum devices. It takes a kernel name as ASCIIZ string literal value
+    and returns the kernel's quantum state. The operation is replaced by a call
+    to the kernel with the provided name in ReplaceStateByKernel pass.
+
+    ```mlir
+      %0 = cc.get_state "callee" : !cc.ptr<!cc.state>
+    ```
+  }];
+
+  let arguments = (ins StrAttr:$calleeName);
+  let results = (outs cc_PointerType:$result);
+  let assemblyFormat = [{
+     $calleeName `:` qualified(type(results)) attr-dict
+  }];
+}
+
 def cc_GlobalOp : CCOp<"global", [IsolatedFromAbove, Symbol]> {
   let summary = "Create a global constant or variable";
   let description = [{

include/cudaq/Optimizer/Transforms/Passes.td

@@ -799,6 +799,44 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
   }];
 }
 
+def ReplaceStateWithKernel : Pass<"replace-state-with-kernel", "mlir::func::FuncOp"> {
+  let summary =
+    "Replace `quake.init_state` instructions with call to the kernel generating the state";
+  let description = [{
+    Argument synthesis for state pointers for quantum devices substitutes state
+    argument by a new state created from `__nvqpp_cudaq_state_get` intrinsic, which
+    in turn accepts the name for the synthesized kernel that generated the state.
+
+    This optimization completes the replacement of `quake.init_state` instruction by:
+
+    - Replace `quake.init_state` by a call that `get_state` call refers to.
+    - Remove all unneeded instructions.
+
+    For example:
+
+    Before ReplaceStateWithKernel (replace-state-with-kernel):
+    ```
+    func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
+      %0 = cc.string_literal "__nvqpp__mlirgen__test_init_state.modified_0" : !cc.ptr<!cc.array<i8 x 45>>
+      %1 = cc.cast %0 : (!cc.ptr<!cc.array<i8 x 45>>) -> !cc.ptr<i8>
+      %2 = call @__nvqpp_cudaq_state_get(%1) : (!cc.ptr<i8>) -> !cc.ptr<!cc.state>
+      %3 = call @__nvqpp_cudaq_state_numberOfQubits(%2) : (!cc.ptr<!cc.state>) -> i64
+      %4 = quake.alloca !quake.veq<?>[%3 : i64]
+      %5 = quake.init_state %4, %2 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
+      return
+    }
+    ```
+
+    After ReplaceStateWithKernel (replace-state-with-kernel):
+    ```
+    func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
+      %5 = call @__nvqpp__mlirgen__test_init_state.modified_0() : () -> !quake.veq<?>
+      return
+    }
+    ```
+  }];
+}
+
 def StatePreparation : Pass<"state-prep", "mlir::ModuleOp"> {
   let summary =
     "Convert state vector data into gates";

lib/Optimizer/Transforms/CMakeLists.txt

@@ -51,6 +51,7 @@ add_cudaq_library(OptTransforms
   QuakeSynthesizer.cpp
   RefToVeqAlloc.cpp
   RegToMem.cpp
+  ReplaceStateWithKernel.cpp
   StatePreparation.cpp
   UnitarySynthesis.cpp
   WiresToWiresets.cpp

lib/Optimizer/Transforms/ReplaceStateWithKernel.cpp

@@ -0,0 +1,132 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2024 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 "PassDetails.h"
+#include "cudaq/Optimizer/Builder/Intrinsics.h"
+#include "cudaq/Optimizer/Dialect/CC/CCOps.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/Passes.h"
+#include <span>
+
+namespace cudaq::opt {
+#define GEN_PASS_DEF_REPLACESTATEWITHKERNEL
+#include "cudaq/Optimizer/Transforms/Passes.h.inc"
+} // namespace cudaq::opt
+
+#define DEBUG_TYPE "replace-state-with-kernel"
+
+using namespace mlir;
+
+namespace {
+
+static bool isCall(Operation *op, std::vector<const char *> &&names) {
+  if (op) {
+    if (auto callOp = dyn_cast<func::CallOp>(op)) {
+      if (auto calleeAttr = callOp.getCalleeAttr()) {
+        auto funcName = calleeAttr.getValue().str();
+        if (std::find(names.begin(), names.end(), funcName) != names.end())
+          return true;
+      }
+    }
+  }
+  return false;
+}
+
+static bool isNumberOfQubitsCall(Operation *op) {
+  return isCall(op, {cudaq::getNumQubitsFromCudaqState});
+}
+
+// clang-format off
+/// Replace `quake.init_state` by a call to a (modified) kernel that produced
+/// the state.
+///
+/// ```
+///  %0 = cc.get_state "__nvqpp__mlirgen__test_init_state.modified_0" : !cc.ptr<!cc.state>
+///  %1 = call @__nvqpp_cudaq_state_numberOfQubits(%0) : (!cc.ptr<!cc.state>) -> i64
+///  %2 = quake.alloca !quake.veq<?>[%1 : i64]
+///  %3 = quake.init_state %2, %0 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
+/// ───────────────────────────────────────────
+/// ...
+///  %5 = call @callee.modified_0() : () -> !quake.veq<?>
+/// ```
+// clang-format on
+class ReplaceStateWithKernelPattern
+    : public OpRewritePattern<quake::InitializeStateOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(quake::InitializeStateOp initState,
+                                PatternRewriter &rewriter) const override {
+    auto *alloca = initState.getOperand(0).getDefiningOp();
+    auto stateOp = initState.getOperand(1);
+
+    if (auto ptrTy = dyn_cast<cudaq::cc::PointerType>(stateOp.getType())) {
+      if (isa<cudaq::cc::StateType>(ptrTy.getElementType())) {
+        auto *numOfQubits = alloca->getOperand(0).getDefiningOp();
+
+        if (auto getState = stateOp.getDefiningOp<cudaq::cc::GetStateOp>()) {
+          auto calleeName = getState.getCalleeName();
+          rewriter.replaceOpWithNewOp<func::CallOp>(
+              initState, initState.getType(), calleeName, mlir::ValueRange{});
+
+          if (alloca->getUses().empty())
+            rewriter.eraseOp(alloca);
+          else {
+            alloca->emitError(
+                "Failed to remove `quake.alloca` in state synthesis");
+            return failure();
+          }
+          if (isNumberOfQubitsCall(numOfQubits)) {
+            if (numOfQubits->getUses().empty())
+              rewriter.eraseOp(numOfQubits);
+            else {
+              numOfQubits->emitError("Failed to remove runtime call to get "
+                                     "number of qubits in state synthesis");
+              return failure();
+            }
+          }
+          return success();
+        }
+        numOfQubits->emitError(
+            "Failed to replace `quake.init_state` in state synthesis");
+      }
+    }
+    return failure();
+  }
+};
+
+class ReplaceStateWithKernelPass
+    : public cudaq::opt::impl::ReplaceStateWithKernelBase<
+          ReplaceStateWithKernelPass> {
+public:
+  using ReplaceStateWithKernelBase::ReplaceStateWithKernelBase;
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    auto func = getOperation();
+    RewritePatternSet patterns(ctx);
+    patterns.insert<ReplaceStateWithKernelPattern>(ctx);
+
+    LLVM_DEBUG(llvm::dbgs()
+               << "Before replace state with kernel: " << func << '\n');
+
+    if (failed(applyPatternsAndFoldGreedily(func.getOperation(),
+                                            std::move(patterns))))
+      signalPassFailure();
+
+    LLVM_DEBUG(llvm::dbgs()
+               << "After replace state with kerenl: " << func << '\n');
+  }
+};
+} // namespace

python/runtime/cudaq/algorithms/py_state.cpp

@@ -96,8 +96,9 @@ class PyRemoteSimulationState : public RemoteSimulationState {
     }
   }
 
-  std::pair<std::string, std::vector<void *>> getKernelInfo() const override {
-    return {kernelName, argsData->getArgs()};
+  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 {

python/runtime/cudaq/platform/py_alt_launch_kernel.cpp

@@ -523,7 +523,7 @@ MlirModule synthesizeKernel(const std::string &name, MlirModule module,
   auto isLocalSimulator = platform.is_simulator() && !platform.is_emulated();
   auto isSimulator = isLocalSimulator || isRemoteSimulator;
 
-  cudaq::opt::ArgumentConverter argCon(name, unwrap(module), isSimulator);
+  cudaq::opt::ArgumentConverter argCon(name, unwrap(module));
   argCon.gen(runtimeArgs.getArgs());
   std::string kernName = cudaq::runtime::cudaqGenPrefixName + name;
   SmallVector<StringRef> kernels = {kernName};