[Python] Qvector init from state

include/cudaq/Optimizer/Builder/Factory.h

@@ -119,6 +119,12 @@ inline mlir::LLVM::LLVMStructType stdVectorImplType(mlir::Type eleTy) {
   return mlir::LLVM::LLVMStructType::getLiteral(ctx, eleTys);
 }
 
+inline mlir::Type stateImplType(mlir::Type eleTy) {
+  auto *ctx = eleTy.getContext();
+  auto eTy = cudaq::opt::factory::getCharType(ctx);
+  return cudaq::opt::factory::getPointerType(eTy);
+}
+
 // Host side types for std::string and std::vector
 
 cudaq::cc::StructType stlStringType(mlir::MLIRContext *ctx);

include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h

@@ -55,6 +55,8 @@ static constexpr const char QIRArrayQubitAllocateArrayWithStateComplex32[] =
     "__quantum__rt__qubit_allocate_array_with_state_complex32";
 static constexpr const char QIRArrayQubitAllocateArrayWithStatePtr[] =
     "__quantum__rt__qubit_allocate_array_with_state_ptr";
+static constexpr const char QIRArrayQubitAllocateArrayWithCudaqStatePtr[] =
+    "__quantum__rt__qubit_allocate_array_with_cudaq_state_ptr";
 static constexpr const char QIRQubitAllocate[] =
     "__quantum__rt__qubit_allocate";
 static constexpr const char QIRArrayQubitReleaseArray[] =

lib/Optimizer/CodeGen/ConvertToQIR.cpp

@@ -143,13 +143,17 @@ class QmemRAIIOpRewrite
     StringRef functionName;
     if (Type eleTy = dyn_cast<LLVM::LLVMPointerType>(ccState.getType())
                          .getElementType()) {
+      if (auto elePtrTy = dyn_cast<LLVM::LLVMPointerType>(eleTy))
+        eleTy = elePtrTy.getElementType();
       if (auto arrayTy = dyn_cast<LLVM::LLVMArrayType>(eleTy))
         eleTy = arrayTy.getElementType();
       bool fromComplex = false;
       if (auto complexTy = dyn_cast<LLVM::LLVMStructType>(eleTy)) {
         fromComplex = true;
         eleTy = complexTy.getBody()[0];
       }
+      if (eleTy == rewriter.getI8Type())
+        functionName = cudaq::opt::QIRArrayQubitAllocateArrayWithCudaqStatePtr;
       if (eleTy == rewriter.getF64Type())
         functionName =
             fromComplex
@@ -187,6 +191,7 @@ class QmemRAIIOpRewrite
     // Create QIR allocation with initializer function.
     auto *ctx = rewriter.getContext();
     auto ptrTy = cudaq::opt::factory::getPointerType(ctx);
+
     FlatSymbolRefAttr raiiSymbolRef =
         cudaq::opt::factory::createLLVMFunctionSymbol(
             functionName, array_qbit_type, {i64Ty, ptrTy}, parentModule);
@@ -2011,6 +2016,8 @@ void cudaq::opt::initializeTypeConversions(LLVMTypeConverter &typeConverter) {
       [](quake::VeqType type) { return getArrayType(type.getContext()); });
   typeConverter.addConversion(
       [](quake::RefType type) { return getQubitType(type.getContext()); });
+  typeConverter.addConversion(
+      [](cc::StateType type) { return factory::stateImplType(type); });
   typeConverter.addConversion([](cc::CallableType type) {
     return lambdaAsPairOfPointers(type.getContext());
   });
@@ -2026,6 +2033,9 @@ void cudaq::opt::initializeTypeConversions(LLVMTypeConverter &typeConverter) {
     if (isa<NoneType>(eleTy))
       return factory::getPointerType(type.getContext());
     eleTy = typeConverter.convertType(eleTy);
+    if (isa<NoneType>(eleTy))
+      return factory::getPointerType(type.getContext());
+
     if (auto arrTy = dyn_cast<cc::ArrayType>(eleTy)) {
       // If array has a static size, it becomes an LLVMArrayType.
       assert(arrTy.isUnknownSize());

lib/Optimizer/Transforms/GenKernelExecution.cpp

@@ -434,8 +434,10 @@ class GenerateKernelExecution
         hasTrailingData = true;
         continue;
       }
-      if (isa<cudaq::cc::PointerType>(currEleTy))
+      if (isa<cudaq::cc::PointerType>(currEleTy) &&
+          !isStatePointerType(currEleTy)) {
         continue;
+      }
 
       // cast to the struct element type, void* -> TYPE *
       argPtr = builder.create<cudaq::cc::CastOp>(
@@ -933,6 +935,13 @@ class GenerateKernelExecution
     builder.create<cudaq::cc::StoreOp>(loc, endPtr, sret2);
   }
 
+  static bool isStatePointerType(mlir::Type ty) {
+    if (auto ptrTy = dyn_cast<cudaq::cc::PointerType>(ty)) {
+      return isa<cudaq::cc::StateType>(ptrTy.getElementType());
+    }
+    return false;
+  }
+
   static MutableArrayRef<BlockArgument>
   dropAnyHiddenArguments(MutableArrayRef<BlockArgument> args,
                          FunctionType funcTy, bool hasThisPointer) {
@@ -941,7 +950,8 @@ class GenerateKernelExecution
         cudaq::cc::numberOfHiddenArgs(hasThisPointer, hiddenSRet);
     if (count > 0 && args.size() >= count &&
         std::all_of(args.begin(), args.begin() + count, [](auto i) {
-          return isa<cudaq::cc::PointerType>(i.getType());
+          return isa<cudaq::cc::PointerType>(i.getType()) &&
+                 !isStatePointerType(i.getType());
         }))
       return args.drop_front(count);
     return args;
@@ -1207,8 +1217,9 @@ class GenerateKernelExecution
         hasTrailingData = true;
         continue;
       }
-      if (isa<cudaq::cc::PointerType>(inTy))
+      if (isa<cudaq::cc::PointerType>(inTy) && !isStatePointerType(inTy))
         continue;
+
       stVal = builder.create<cudaq::cc::InsertValueOp>(loc, stVal.getType(),
                                                        stVal, arg, idx);
     }

python/cudaq/kernel/ast_bridge.py

@@ -6,6 +6,7 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 import ast
+import hashlib
 import graphlib
 import sys, os
 from typing import Callable
@@ -19,6 +20,8 @@
 from ..mlir.dialects import builtin, func, arith, math, complex
 from ..mlir._mlir_libs._quakeDialects import cudaq_runtime, load_intrinsic, register_all_dialects
 
+State = cudaq_runtime.State
+
 # This file implements the CUDA-Q Python AST to MLIR conversion.
 # It provides a `PyASTBridge` class that implements the `ast.NodeVisitor` type
 # to walk the Python AST for a `cudaq.kernel` annotated function and generate
@@ -110,6 +113,8 @@ def __init__(self, **kwargs):
         symbol table, which maps variable names to constructed `mlir.Values`. 
         """
         self.valueStack = deque()
+        self.cudaqStateHashes = kwargs[
+            'cudaqStateHashes'] if 'cudaqStateHashes' in kwargs else None
         self.knownResultType = kwargs[
             'knownResultType'] if 'knownResultType' in kwargs else None
         if 'existingModule' in kwargs:
@@ -1842,19 +1847,27 @@ def bodyBuilder(iterVal):
                         return
 
                 if node.func.attr == 'qvector':
-                    value = self.ifPointerThenLoad(self.popValue())
-                    if (IntegerType.isinstance(value.type)):
+                    valueOrPtr = self.popValue()
+                    initializerTy = valueOrPtr.type
+
+                    if cc.PointerType.isinstance(initializerTy):
+                        initializerTy = cc.PointerType.getElementType(
+                            initializerTy)
+
+                    if (IntegerType.isinstance(initializerTy)):
                         # handle `cudaq.qvector(n)`
+                        value = self.ifPointerThenLoad(valueOrPtr)
                         ty = self.getVeqType()
                         qubits = quake.AllocaOp(ty, size=value).result
                         self.pushValue(qubits)
                         return
-                    if cc.StdvecType.isinstance(value.type):
+                    if cc.StdvecType.isinstance(initializerTy):
                         # handle `cudaq.qvector(initState)`
 
                         # Validate the length in case of a constant initializer:
                         # `cudaq.qvector([1., 0., ...])`
                         # `cudaq.qvector(np.array([1., 0., ...]))`
+                        value = self.ifPointerThenLoad(valueOrPtr)
                         listScalar = None
                         arrNode = node.args[0]
                         if isinstance(arrNode, ast.List):
@@ -1899,14 +1912,32 @@ def bodyBuilder(iterVal):
                         veqTy = quake.VeqType.get(self.ctx)
 
                         qubits = quake.AllocaOp(veqTy, size=numQubits).result
-                        initials = cc.StdvecDataOp(ptrTy, value).result
+                        data = cc.StdvecDataOp(ptrTy, value).result
                         init = quake.InitializeStateOp(veqTy, qubits,
-                                                       initials).result
+                                                       data).result
+                        self.pushValue(init)
+                        return
+
+                    if cc.StateType.isinstance(initializerTy):
+                        # handle `cudaq.qvector(state)`
+                        statePtr = self.ifNotPointerThenStore(valueOrPtr)
+
+                        symName = '__nvqpp_cudaq_state_numberOfQubits'
+                        load_intrinsic(self.module, symName)
+                        i64Ty = self.getIntegerType()
+                        numQubits = func.CallOp([i64Ty], symName,
+                                                [statePtr]).result
+
+                        veqTy = quake.VeqType.get(self.ctx)
+                        qubits = quake.AllocaOp(veqTy, size=numQubits).result
+                        init = quake.InitializeStateOp(veqTy, qubits,
+                                                       statePtr).result
+
                         self.pushValue(init)
                         return
 
                     self.emitFatalError(
-                        f"unsupported qvector argument type: {value.type} (unknown)",
+                        f"unsupported qvector argument type: {value.type}",
                         node)
                     return
 
@@ -3330,6 +3361,54 @@ def visit_BinOp(self, node):
         else:
             self.emitFatalError(f"unhandled binary operator - {node.op}", node)
 
+    def __store_state(self, value: State):
+        # Compute a unique hash string for the state data
+        hashValue = hashlib.sha1(value).hexdigest(
+        )[:10] + self.name.removeprefix('__nvqppBuilderKernel_')
+        print(f'state has value for {self.name}: {hashValue}')
+
+        stateTy = cc.StateType.get(self.ctx)
+        statePtrTy = cc.PointerType.get(self.ctx, stateTy)
+
+        # Generate a function that stores the state value in a global
+        globalTy = statePtrTy
+        globalName = f'nvqpp.cudaq.state.{hashValue}'
+        setStateName = f'nvqpp.set.cudaq.state.{hashValue}'
+        with InsertionPoint.at_block_begin(self.module.body):
+            cc.GlobalOp(TypeAttr.get(globalTy), globalName, external=True)
+            setStateFunc = func.FuncOp(setStateName,
+                                       FunctionType.get(inputs=[statePtrTy],
+                                                        results=[]),
+                                       loc=self.loc)
+            entry = setStateFunc.add_entry_block()
+            with InsertionPoint(entry):
+                zero = self.getConstantInt(0)
+                address = cc.AddressOfOp(cc.PointerType.get(self.ctx, globalTy),
+                                         FlatSymbolRefAttr.get(globalName))
+                ptr = cc.ComputePtrOp(
+                    cc.PointerType.get(self.ctx, statePtrTy), address, [zero],
+                    DenseI32ArrayAttr.get([kDynamicPtrIndex], context=self.ctx))
+
+                cc.StoreOp(entry.arguments[0], ptr)
+                func.ReturnOp([])
+
+        # Record the unique hash value
+        if hashValue not in self.cudaqStateHashes:
+            self.cudaqStateHashes.append(hashValue)
+
+        # Store the state into a global variable
+        cudaq_runtime.storePointerToCudaqState(self.name, hashValue, value)
+
+        # Return the pointer to stored state
+        zero = self.getConstantInt(0)
+        address = cc.AddressOfOp(cc.PointerType.get(self.ctx, globalTy),
+                                 FlatSymbolRefAttr.get(globalName)).result
+        ptr = cc.ComputePtrOp(
+            cc.PointerType.get(self.ctx, statePtrTy), address, [zero],
+            DenseI32ArrayAttr.get([kDynamicPtrIndex], context=self.ctx)).result
+        statePtr = cc.LoadOp(ptr).result
+        return statePtr
+
     def visit_Name(self, node):
         """
         Visit `ast.Name` nodes and extract the correct value from the symbol table.
@@ -3378,6 +3457,11 @@ def visit_Name(self, node):
             # Only support a small subset of types here
             complexType = type(1j)
             value = self.capturedVars[node.id]
+
+            if isinstance(value, State):
+                self.pushValue(self.__store_state(value))
+                return
+
             if isinstance(value, (list, np.ndarray)) and isinstance(
                     value[0], (int, bool, float, np.float32, np.float64,
                                complexType, np.complex64, np.complex128)):
@@ -3445,7 +3529,7 @@ def visit_Name(self, node):
                 errorType = f"{errorType}[{type(value[0]).__name__}]"
 
             self.emitFatalError(
-                f"Invalid type for variable ({node.id}) captured from parent scope (only int, bool, float, complex, and list/np.ndarray[int|bool|float|complex] accepted, type was {errorType}).",
+                f"Invalid type for variable ({node.id}) captured from parent scope (only int, bool, float, complex, cudaq.State, and list/np.ndarray[int|bool|float|complex] accepted, type was {errorType}).",
                 node)
 
         # Throw an exception for the case that the name is not
@@ -3474,9 +3558,11 @@ def compile_to_mlir(astModule, metadata, **kwargs):
     lineNumberOffset = kwargs['location'] if 'location' in kwargs else ('', 0)
     parentVariables = kwargs[
         'parentVariables'] if 'parentVariables' in kwargs else {}
+    cudaqStateHashes = kwargs['cudaqStateHashes']
 
     # Create the AST Bridge
     bridge = PyASTBridge(verbose=verbose,
+                         cudaqStateHashes=cudaqStateHashes,
                          knownResultType=returnType,
                          returnTypeIsFromPython=True,
                          locationOffset=lineNumberOffset,

python/cudaq/kernel/kernel_decorator.py

@@ -39,6 +39,7 @@ class PyKernelDecorator(object):
 
     def __init__(self, function, verbose=False, module=None, kernelName=None):
         self.kernelFunction = function
+        self.cudaqStateHashes = []
         self.module = None if module == None else module
         self.verbose = verbose
         self.name = kernelName if kernelName != None else self.kernelFunction.__name__
@@ -163,7 +164,8 @@ def compile(self):
             verbose=self.verbose,
             returnType=self.returnType,
             location=self.location,
-            parentVariables=self.globalScopedVars)
+            parentVariables=self.globalScopedVars,
+            cudaqStateHashes=self.cudaqStateHashes)
 
         # Grab the dependent capture variables, if any
         self.dependentCaptures = extraMetadata[
@@ -294,13 +296,16 @@ def __call__(self, *args):
                                             self.module,
                                             *processedArgs,
                                             callable_names=callableNames)
+            cudaq_runtime.deletePointersToCudaqState(self.cudaqStateHashes)
         else:
-            return cudaq_runtime.pyAltLaunchKernelR(
+            result = cudaq_runtime.pyAltLaunchKernelR(
                 self.name,
                 self.module,
                 mlirTypeFromPyType(self.returnType, self.module.context),
                 *processedArgs,
                 callable_names=callableNames)
+            cudaq_runtime.deletePointersToCudaqState(self.cudaqStateHashes)
+            return result
 
 
 def kernel(function=None, **kwargs):

python/cudaq/kernel/utils.py

@@ -15,6 +15,7 @@
 from typing import Callable, List
 import ast, sys, traceback
 
+State = cudaq_runtime.State
 qvector = cudaq_runtime.qvector
 qubit = cudaq_runtime.qubit
 pauli_word = cudaq_runtime.pauli_word
@@ -87,6 +88,8 @@ def emitFatalErrorOverride(msg):
         if annotation.value.id == 'cudaq':
             if annotation.attr in ['qview', 'qvector']:
                 return quake.VeqType.get(ctx)
+            if annotation.attr in ['State']:
+                return cc.PointerType.get(ctx, cc.StateType.get(ctx))
             if annotation.attr == 'qubit':
                 return quake.RefType.get(ctx)
             if annotation.attr == 'pauli_word':
@@ -193,6 +196,8 @@ def mlirTypeFromPyType(argType, ctx, **kwargs):
         return ComplexType.get(mlirTypeFromPyType(np.float32, ctx))
     if argType == pauli_word:
         return cc.CharspanType.get(ctx)
+    if argType == State:
+        return cc.PointerType.get(ctx, cc.StateType.get(ctx))
 
     if argType in [list, np.ndarray, List]:
         if 'argInstance' not in kwargs:

python/runtime/cudaq/algorithms/py_state.cpp

@@ -124,7 +124,6 @@ void bindPyState(py::module &mod) {
                                   dataTypeSize}, /* strides */
                                  true            /* readonly */
           );
-
         return py::buffer_info(dataPtr, dataTypeSize, /*itemsize */
                                desc, 1,               /* ndim */
                                {shape[0]},            /* shape */
@@ -153,9 +152,20 @@ void bindPyState(py::module &mod) {
                   reinterpret_cast<std::complex<float> *>(info.ptr),
                   info.size));
             }
-
-            return state::from_data(std::make_pair(
-                reinterpret_cast<std::complex<double> *>(info.ptr), info.size));
+            if (info.format ==
+                py::format_descriptor<std::complex<double>>::format()) {
+              return state::from_data(std::make_pair(
+                  reinterpret_cast<std::complex<double> *>(info.ptr),
+                  info.size));
+            }
+            throw std::runtime_error(
+                "A numpy array with only floating point elements passed to "
+                "state.from_data. input must be of complex float type, "
+                "please "
+                "add to your array creation `dtype=numpy.complex64` if "
+                "simulation is FP32 and `dtype=numpy.complex128` if "
+                "simulation if FP64, or dtype=cudaq.complex() for "
+                "precision-agnostic code");
           },
           "Return a state from data.")
       .def_static(