New launch kernel for python

docs/sphinx/api/languages/cpp_api.rst

@@ -91,6 +91,8 @@ Common
 
 .. doxygenclass:: cudaq::RemoteSimulationState
 
+.. doxygenclass:: cudaq::QPUState
+
 .. doxygenclass:: cudaq::registry::RegisteredType
     :members:
 

docs/sphinx/targets/cpp/pasqal.cpp

@@ -3,9 +3,9 @@
 // nvq++ --target pasqal pasqal.cpp -o out.x
 // ./out.x
 // ```
-// Assumes a valid set of credentials (PASQAL_AUTH_TOKEN, PASQAL_PROJECT_ID)
-// have been set. To set PASQAL_AUTH_TOKEN from Pasqal Cloud username and
-// password, use pasqal_auth.py in this folder.
+// Assumes a valid set of credentials (`PASQAL_AUTH_TOKEN`, `PASQAL_PROJECT_ID`)
+// have been set. To set `PASQAL_AUTH_TOKEN` from Pasqal Cloud username and
+// password, use `pasqal_auth.py` in this folder.
 
 #include "cudaq/algorithms/evolve.h"
 #include "cudaq/dynamics_integrators.h"
@@ -15,7 +15,7 @@
 #include <map>
 #include <vector>
 
-// This example illustrates how to use Pasqal's EMU_MPS emulator over Pasqal's
+// This example illustrates how to use Pasqal EMU_MPS emulator over Pasqal
 // cloud via CUDA-Q. Contact Pasqal at [email protected] or through
 // https://community.pasqal.com for assistance.
 

docs/sphinx/targets/cpp/quera_basic.cpp

@@ -16,7 +16,7 @@
 // Credentials must be set before running this program.
 // Amazon Braket costs apply.
 
-// This example illustrates how to use QuEra's Aquila device on Braket with
+// This example illustrates how to use QuEra Aquila device on Braket with
 // CUDA-Q. It is a CUDA-Q implementation of the getting started materials for
 // Braket available here:
 // https://docs.aws.amazon.com/braket/latest/developerguide/braket-get-started-hello-ahs.html

docs/sphinx/targets/cpp/quera_intro.cpp

@@ -13,7 +13,7 @@
 #include <map>
 #include <vector>
 
-// This example illustrates how to use QuEra's Aquila device on Braket with
+// This example illustrates how to use QuEra Aquila device on Braket with
 // CUDA-Q. It is a CUDA-Q implementation of the getting started materials for
 // Braket available here:
 // https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/analog_hamiltonian_simulation/01_Introduction_to_Aquila.ipynb

docs/sphinx/targets/python/quera_intro.py

@@ -2,7 +2,7 @@
 from cudaq.operator import *
 import numpy as np
 
-# This example illustrates how to use QuEra's Aquila device on Braket with CUDA-Q.
+# This example illustrates how to use QuEra Aquila device on Braket with CUDA-Q.
 # It is a CUDA-Q implementation of the getting started materials for Braket available here:
 # https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/analog_hamiltonian_simulation/01_Introduction_to_Aquila.ipynb
 

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

@@ -1595,7 +1595,7 @@ def quake_CreateStateOp : QuakeOp<"create_state", [Pure]> {
     cc_PointerType:$data,
     AnySignlessInteger:$length
   );
-  let results = (outs cc_PointerType:$result);
+  let results = (outs PointerOf<[quake_StateType]>:$result);
   let assemblyFormat = [{
       $data `,` $length `:` functional-type(operands, results) attr-dict
   }];
@@ -1613,7 +1613,7 @@ def QuakeOp_DeleteStateOp : QuakeOp<"delete_state", [] > {
     ```
   }];
 
-  let arguments = (ins cc_PointerType:$state);
+  let arguments = (ins PointerOf<[quake_StateType]>:$state);
   let results = (outs);
   let assemblyFormat = [{
       $state `:` type(operands) attr-dict
@@ -1633,30 +1633,43 @@ def quake_GetNumberOfQubitsOp : QuakeOp<"get_number_of_qubits", [Pure] > {
     ```
   }];
 
-  let arguments = (ins cc_PointerType:$state);
+  let arguments = (ins PointerOf<[quake_StateType]>:$state);
   let results = (outs AnySignlessInteger:$result);
   let assemblyFormat = [{
       $state `:` functional-type(operands, results) attr-dict
   }];
 }
 
-def QuakeOp_GetStateOp : QuakeOp<"get_state", [Pure] > {
+def QuakeOp_MaterializeStateOp : QuakeOp<"materialize_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 symbol references:
+      - @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`.
+
+    This operation will return the state of the original kernel called with
+    arguments passed to `cudaq::get_state`.
+
+    The operation may be replaced by calls to the @num_qubits and @init calls,
+    which will reproduce the specified state in the `ReplaceStateByKernel`
+    pass.
 
     ```mlir
-      %0 = quake.get_state "callee" : !cc.ptr<!quake.state>
+      %0 = quake.materialize_state @num_qubits, @init : !cc.ptr<!quake.state>
     ```
   }];
 
-  let arguments = (ins StrAttr:$calleeName);
-  let results = (outs cc_PointerType:$result);
+  let arguments = (ins
+    FlatSymbolRefAttr:$numQubitsFunc,
+    FlatSymbolRefAttr:$initFunc
+  );
+  let results = (outs PointerOf<[quake_StateType]>:$result);
   let assemblyFormat = [{
-     $calleeName `:` qualified(type(results)) attr-dict
+     $numQubitsFunc `,` $initFunc `:` qualified(type(results)) attr-dict
   }];
 }
 

include/cudaq/Optimizer/Transforms/Passes.td

@@ -814,8 +814,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<!quake.state>
-      %5 = cc.get_number_of_qubits %4 : (!cc.ptr<!quake.state>) -> i64
+      %4 = quake.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!quake.state>
+      %5 = quake.get_number_of_qubits %4 : (!cc.ptr<!quake.state>) -> i64
       %6 = quake.alloca !quake.veq<?>[%5 : i64]
       %7 = quake.init_state %6, %4 : (!quake.veq<?>, !cc.ptr<!quake.state>) -> !quake.veq<?>
 
@@ -841,7 +841,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<!quake.state>
+      %4 = quake.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!quake.state>
       call @__nvqpp__mlirgen__sub_kernel(%4) : (!cc.ptr<!quake.state>) -> ()
       return
     }
@@ -851,9 +851,79 @@ 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<!quake.state>
+      %4 = quake.create_state %3, %c8_i64  : (!cc.ptr<!cc.array<complex<f32> x 8>>, i64) -> !cc.ptr<!quake.state>
       call @__nvqpp__mlirgen__sub_kernel(%4) : (!cc.ptr<!quake.state>) -> ()
-      call @__nvqpp_cudaq_state_delete(%4) : (!cc.ptr<!quake.state>) -> ()
+      quake.delete_state %3 : !cc.ptr<!quake.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 = [{
+    This optimization replaces `quake.init_state`, `quake.get_number_of_qubits`,
+    and `quake.materialize_state` operations invoked on state pointers during
+    argument synthesis for quantum devices.
+
+    Before this optimization, argument synthesis for state pointers for quantum
+    devices substituted a state created from the `quake.materialize_state`
+    operation for the state argument.
+
+    The `quake.materialize_state` operation accepts symbols for the synthesized
+    kernels `@num_qubits` and `@init` that argument synthesis generated from
+    the original kernel call that generated the state, e.g.,
+    the `cudaq::get_state` call that refers to the result of a specific quantum
+    kernel being invoked with a set of parameters
+
+    For example, for the user code:
+      ```
+      state = cudaq::get_state(callee, args)
+      caller(state)
+      ```
+
+    The argument synthesis generated the following new kernels from the `callee`
+    and synthesized them to substitute their arguments with `args`:
+      ```
+      func.func @callee_num_qubits(arguments) -> i64
+      func.func @callee_init(qubits: !quake.veq<?>, arguments) -> !quake.veq<?>
+      ```
+
+    The argument synthesis also substituted the state argument in the `caller`
+    with:
+      ```
+      quake.materialize_state @callee_num_qubits @callee_init: !cc.ptr<!quake.state>
+      ```
+
+    This optimization performs the replacements for the the following operations 
+    that use a state produced by  `quake.materialize_state @num_qubits, @init`
+    operation:
+
+    - Replace `quake.get_number_of_qubits` operation by call to `@num_qubits`
+    - Replace `quake.init_state` operation by call to `@init`
+    - Clean up unused `quake.materialize_state` operation
+
+    For example:
+
+    Before ReplaceStateWithKernel (replace-state-with-kernel):
+    ```
+    func.func @foo() {
+      %0 = quake.materialize_state @callee.num_qubits_0, @callee.init_0: !cc.ptr<!quake.state>
+      %1 = quake.get_number_of_qubits %0 : (!cc.ptr<!quake.state>) -> i64
+      %2 = quake.alloca !quake.veq<?>[%1 : i64]
+      %3 = quake.init_state %2, %0 : (!quake.veq<?>, !cc.ptr<!quake.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/Dialect/CC/CCOps.cpp

@@ -980,6 +980,14 @@ struct FuseWithConstantArray
 
           return success();
         }
+        if (auto complexTy = dyn_cast<ComplexType>(extval.getType())) {
+          std::int32_t i = extval.getRawConstantIndices()[0];
+          auto cval = cast<ArrayAttr>(conarr.getConstantValues()[i]);
+          rewriter.replaceOpWithNewOp<complex::ConstantOp>(extval, complexTy,
+                                                           cval);
+
+          return success();
+        }
       }
     return failure();
   }

lib/Optimizer/Transforms/CMakeLists.txt

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

lib/Optimizer/Transforms/ReplaceStateWithKernel.cpp

@@ -0,0 +1,135 @@
+/*******************************************************************************
+ * 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"
+
+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 function
+/// that computes the number of qubits for a state.
+///
+/// ```mlir
+///  %0 = quake.materialize_state @callee.num_qubits_0, @callee.init_0 : !cc.ptr<!quake.state>
+///  %1 = quake.get_number_of_qubits %0 : (!cc.ptr<!quake.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.getState();
+    auto materializeState = stateOp.getDefiningOp<quake::MaterializeStateOp>();
+    if (!materializeState) {
+      LLVM_DEBUG(llvm::dbgs() << "ReplaceStateWithKernel: failed to replace "
+                                 "`quake.get_num_qubits`: "
+                              << stateOp << '\n');
+      return failure();
+    }
+
+    auto numQubitsFunc = materializeState.getNumQubitsFunc();
+    rewriter.setInsertionPoint(numQubits);
+    rewriter.replaceOpWithNewOp<func::CallOp>(
+        numQubits, numQubits.getType(), numQubitsFunc, mlir::ValueRange{});
+    return success();
+  }
+};
+
+// clang-format off
+/// Replace `quake.init_state` by a call to a (modified) kernel that produced
+/// the state.
+///
+/// ```mlir
+///  %0 = quake.materialize_state @callee.num_qubits_0, @callee.init_0 : !cc.ptr<!quake.state>
+///  %3 = quake.init_state %2, %0 : (!quake.veq<?>, !cc.ptr<!quake.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.getTargets();
+    auto stateOp = initState.getState();
+
+    if (auto ptrTy = dyn_cast<cudaq::cc::PointerType>(stateOp.getType())) {
+      if (isa<quake::StateType>(ptrTy.getElementType())) {
+        auto materializeState =
+            stateOp.getDefiningOp<quake::MaterializeStateOp>();
+        if (!materializeState) {
+          LLVM_DEBUG(llvm::dbgs() << "ReplaceStateWithKernel: failed to "
+                                     "replace `quake.init_state`: "
+                                  << stateOp << '\n');
+          return failure();
+        }
+
+        auto initName = materializeState.getInitFunc();
+        rewriter.setInsertionPoint(initState);
+        rewriter.replaceOpWithNewOp<func::CallOp>(initState,
+                                                  initState.getType(), initName,
+                                                  mlir::ValueRange{allocaOp});
+
+        return success();
+      }
+    }
+    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