State synthesis for quantum devices

include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td

@@ -1467,19 +1467,29 @@ def QuakeOp_GetStateOp : QuakeOp<"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.
+    for quantum devices.
+
+    It takes two kernel names as ASCIIZ string literals:
+      - "num_qubits" for determining the size of the allocation to initialize
+      - "init" for initializing the state the same way as the original kernel
+        passed to `cudaq::get_state`) as ASCIIZ string literal
+
+    And returns the quantum state of the original kernel passed to
+    `cudaq::get_state`. The operation is replaced by calls to the kernels with
+    the provided names in `ReplaceStateByKernel` pass.
 
     ```mlir
-      %0 = quake.get_state "callee" : !cc.ptr<!cc.state>
+      %0 = quake.get_state "num_qubits" "init" : !cc.ptr<!cc.state>
     ```
   }];
 
-  let arguments = (ins StrAttr:$calleeName);
+  let arguments = (ins
+    StrAttr:$numQubitsFuncName,
+    StrAttr:$initFuncName
+  );
   let results = (outs cc_PointerType:$result);
   let assemblyFormat = [{
-     $calleeName `:` qualified(type(results)) attr-dict
+     $numQubitsFuncName $initFuncName `:` qualified(type(results)) attr-dict
   }];
 }
 

include/cudaq/Optimizer/Transforms/Passes.h

@@ -58,6 +58,13 @@ std::unique_ptr<mlir::Pass>
 createArgumentSynthesisPass(mlir::ArrayRef<mlir::StringRef> funcNames,
                             mlir::ArrayRef<mlir::StringRef> substitutions);
 
+/// Helper function to build an argument synthesis pass. The names of the
+/// functions and the substitutions text can be built as an unzipped pair of
+/// lists.
+std::unique_ptr<mlir::Pass>
+createArgumentSynthesisPass(const std::vector<std::string> &funcNames,
+                            const std::vector<std::string> &substitutions);
+
 // declarative passes
 #define GEN_PASS_DECL
 #define GEN_PASS_REGISTRATION

include/cudaq/Optimizer/Transforms/Passes.td

@@ -778,8 +778,8 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
     func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
       %c8_i64 = arith.constant 8 : i64
       %0 = cc.address_of @foo.rodata_synth_0 : !cc.ptr<!cc.array<complex<f32> x 8>>
-      %4 = cc.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
-      %5 = cc.get_number_of_qubits %4 : (!cc.ptr<!cc.state>) -> i64
+      %4 = quake.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
+      %5 = quake.get_number_of_qubits %4 : (!cc.ptr<!cc.state>) -> i64
       %6 = quake.alloca !quake.veq<?>[%5 : i64]
       %7 = quake.init_state %6, %4 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
 
@@ -805,7 +805,7 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
     Before DeleteStates (delete-states):
     ```  func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
       ...
-      %4 = call @__nvqpp_cudaq_state_createFromData_fp32(%3, %c8_i64) : (!cc.ptr<i8>, i64) -> !cc.ptr<!cc.state>
+      %4 = quake.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
       call @__nvqpp__mlirgen__sub_kernel(%4) : (!cc.ptr<!cc.state>) -> ()
       return
     }
@@ -815,9 +815,47 @@ def DeleteStates : Pass<"delete-states", "mlir::ModuleOp"> {
     ```
     func.func @foo() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} {
       ...
-      %4 = call @__nvqpp_cudaq_state_createFromData_fp32(%3, %c8_i64) : (!cc.ptr<i8>, i64) -> !cc.ptr<!cc.state>
+      %4 = quake.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!cc.state>
       call @__nvqpp__mlirgen__sub_kernel(%4) : (!cc.ptr<!cc.state>) -> ()
-      call @__nvqpp_cudaq_state_delete(%4) : (!cc.ptr<!cc.state>) -> ()
+      quake.delete_state %3 : !cc.ptr<!cc.state>
+      return
+    }
+    ```
+  }];
+}
+
+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() {
+      %0 = quake.get_state "callee.num_qubits_0" "callee.init_0": !cc.ptr<!cc.state>
+      %1 = quake.get_number_of_qubits %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<?>
+      return
+    }
+    ```
+
+    After ReplaceStateWithKernel (replace-state-with-kernel):
+    ```
+    func.func @foo()  {
+      %1 = call @callee.num_qubits_0() : () -> i64
+      %2 = quake.alloca !quake.veq<?>[%1 : i64]
+      %3 = call @callee.init_0(%2) : (!quake.veq<?>) -> !quake.veq<?>
       return
     }
     ```

lib/Optimizer/Transforms/ArgumentSynthesis.cpp

@@ -163,3 +163,12 @@ cudaq::opt::createArgumentSynthesisPass(ArrayRef<StringRef> funcNames,
   return std::make_unique<ArgumentSynthesisPass>(
       ArgumentSynthesisOptions{pairs});
 }
+
+std::unique_ptr<mlir::Pass> cudaq::opt::createArgumentSynthesisPass(
+    const std::vector<std::string> &funcNames,
+    const std::vector<std::string> &substitutions) {
+  return cudaq::opt::createArgumentSynthesisPass(
+      mlir::SmallVector<mlir::StringRef>{funcNames.begin(), funcNames.end()},
+      mlir::SmallVector<mlir::StringRef>{substitutions.begin(),
+                                         substitutions.end()});
+}

lib/Optimizer/Transforms/CMakeLists.txt

@@ -52,6 +52,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 - 2025 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 {
+// clang-format off
+/// Replace `quake.get_number_of_qubits` by a call to a a function
+/// that computes the number of qubits for a state.
+///
+/// ```
+///  %0 = quake.get_state "callee.num_qubits_0" "callee.init_0" : !cc.ptr<!cc.state>
+///  %1 = quake.get_number_of_qubits %0 : (!cc.ptr<!cc.state>) -> i64
+/// ───────────────────────────────────────────
+/// ...
+///  %1 = call @callee.num_qubits_0() : () -> i64
+/// ```
+// clang-format on
+class ReplaceGetNumQubitsPattern
+    : public OpRewritePattern<quake::GetNumberOfQubitsOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(quake::GetNumberOfQubitsOp numQubits,
+                                PatternRewriter &rewriter) const override {
+
+    auto stateOp = numQubits.getOperand();
+    if (auto getState = stateOp.getDefiningOp<quake::GetStateOp>()) {
+      auto numQubitsName = getState.getNumQubitsFuncName();
+
+      rewriter.setInsertionPoint(numQubits);
+      rewriter.replaceOpWithNewOp<func::CallOp>(
+          numQubits, numQubits.getType(), numQubitsName, mlir::ValueRange{});
+      return success();
+    }
+    return numQubits->emitError(
+        "ReplaceStateWithKernel: failed to replace `quake.get_num_qubits`");
+  }
+};
+
+// clang-format off
+/// Replace `quake.init_state` by a call to a (modified) kernel that produced
+/// the state.
+///
+/// ```
+///  %0 = quake.get_state "callee.num_qubits_0" "callee.init_0" : !cc.ptr<!cc.state>
+///  %3 = quake.init_state %2, %0 : (!quake.veq<?>, !cc.ptr<!cc.state>) -> !quake.veq<?>
+/// ───────────────────────────────────────────
+/// ...
+/// %3 = call @callee.init_0(%2): (!quake.veq<?>) -> !quake.veq<?>
+/// ```
+// clang-format on
+class ReplaceInitStatePattern
+    : public OpRewritePattern<quake::InitializeStateOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(quake::InitializeStateOp initState,
+                                PatternRewriter &rewriter) const override {
+    auto allocaOp = initState.getOperand(0);
+    auto stateOp = initState.getOperand(1);
+
+    if (auto ptrTy = dyn_cast<cudaq::cc::PointerType>(stateOp.getType())) {
+      if (isa<cudaq::cc::StateType>(ptrTy.getElementType())) {
+        if (auto getState = stateOp.getDefiningOp<quake::GetStateOp>()) {
+          auto initName = getState.getInitFuncName();
+
+          rewriter.setInsertionPoint(initState);
+          rewriter.replaceOpWithNewOp<func::CallOp>(
+              initState, initState.getType(), initName,
+              mlir::ValueRange{allocaOp});
+
+          return success();
+        }
+
+        return initState->emitError(
+            "ReplaceStateWithKernel: failed to replace `quake.init_state`");
+      }
+    }
+    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<ReplaceGetNumQubitsPattern, ReplaceInitStatePattern>(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 kernel: " << func << '\n');
+  }
+};
+} // namespace

python/runtime/cudaq/algorithms/py_state.cpp

@@ -98,8 +98,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

@@ -535,7 +535,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};