Simplifying cudaq::state support.

This PR greatly simplifies the use of cudaq::get_state by limiting its
applicability strictly to local simulation execution contexts. Prior
attempts were made to make cudaq::get_state generally available regardless
of the target execution context. This level of support is being discontinued.

- Move the qvector ctor overloads to the cudaq::state class. This
  simplifies qvector by making it only possible to initialize a
  qvector with a cudaq::state object and dropping all the overloads.
- Fix tests.  Adapt the regression tests.
- Drop state with non-simulation targets.
- Add code to distinguish between floats and complex floats and get rid
  of a bunch of redundant code that wasn't needed. Also, update tests.
- Make all command line arguments be in the same order.
- Don't implicit handle density matrix as a flattened vector (#6)

Signed-off-by: Eric Schweitz <[email protected]>

Author: schweitzpgi

include/cudaq/Frontend/nvqpp/ASTBridge.h

@@ -288,6 +288,10 @@ class QuakeBridgeVisitor
   bool TraverseCXXConstructExpr(clang::CXXConstructExpr *x,
                                 DataRecursionQueue *q = nullptr);
   bool VisitCXXConstructExpr(clang::CXXConstructExpr *x);
+  bool TraverseCXXTemporaryObjectExpr(clang::CXXTemporaryObjectExpr *x,
+                                      DataRecursionQueue *q = nullptr) {
+    return TraverseCXXConstructExpr(x, q);
+  }
   bool VisitCXXOperatorCallExpr(clang::CXXOperatorCallExpr *x);
   bool VisitCXXParenListInitExpr(clang::CXXParenListInitExpr *x);
   bool WalkUpFromCXXOperatorCallExpr(clang::CXXOperatorCallExpr *x);

include/cudaq/Optimizer/Builder/Intrinsics.h

@@ -53,10 +53,14 @@ static constexpr const char getNumQubitsFromCudaqState[] =
     "__nvqpp_cudaq_state_numberOfQubits";
 
 // Create a new state from data.
-static constexpr const char createCudaqStateFromDataFP64[] =
-    "__nvqpp_cudaq_state_createFromData_fp64";
-static constexpr const char createCudaqStateFromDataFP32[] =
-    "__nvqpp_cudaq_state_createFromData_fp32";
+static constexpr const char createCudaqStateFromDataComplexF64[] =
+    "__nvqpp_cudaq_state_createFromData_complex_f64";
+static constexpr const char createCudaqStateFromDataComplexF32[] =
+    "__nvqpp_cudaq_state_createFromData_complex_f32";
+static constexpr const char createCudaqStateFromDataF64[] =
+    "__nvqpp_cudaq_state_createFromData_f64";
+static constexpr const char createCudaqStateFromDataF32[] =
+    "__nvqpp_cudaq_state_createFromData_f32";
 
 // Delete a state created by the runtime functions above.
 static constexpr const char deleteCudaqState[] = "__nvqpp_cudaq_state_delete";

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

@@ -121,7 +121,7 @@ def quake_InitializeStateOp : QuakeOp<"init_state",
   }];
 
   let arguments = (ins
-    VeqType:$targets,
+    NonStruqRefType:$targets,
     AnyStateInitType:$state
   );
   let results = (outs VeqType);

lib/Frontend/nvqpp/ConvertExpr.cpp

@@ -682,9 +682,12 @@ bool QuakeBridgeVisitor::VisitCastExpr(clang::CastExpr *x) {
         if (cxxExpr->getNumArgs() == 1)
           return true;
     }
-    if (isa<ComplexType>(castToTy) && isa<ComplexType>(peekValue().getType())) {
+    if (isa<ComplexType>(castToTy) && isa<ComplexType>(peekValue().getType()))
       return true;
-    }
+    if (isa<quake::StateType>(castToTy))
+      if (auto ptrTy = dyn_cast<cudaq::cc::PointerType>(peekValue().getType()))
+        if (isa<quake::StateType>(ptrTy.getElementType()))
+          return pushValue(builder.create<cudaq::cc::LoadOp>(loc, popValue()));
     if (auto funcTy = peelPointerFromFunction(castToTy))
       if (auto fromTy = dyn_cast<cc::CallableType>(peekValue().getType())) {
         auto inputs = funcTy.getInputs();
@@ -1003,8 +1006,8 @@ bool QuakeBridgeVisitor::VisitMaterializeTemporaryExpr(
   // The following cases are λ expressions, quantum data, or a std::vector view.
   // In those cases, there is nothing to materialize, so we can just pass the
   // Value on the top of the stack.
-  if (isa<cc::CallableType, quake::VeqType, quake::RefType, cc::SpanLikeType>(
-          ty))
+  if (isa<cc::CallableType, quake::VeqType, quake::RefType, cc::SpanLikeType,
+          quake::StateType>(ty))
     return true;
 
   // If not one of the above special cases, then materialize the value to a
@@ -2689,6 +2692,11 @@ bool QuakeBridgeVisitor::VisitInitListExpr(clang::InitListExpr *x) {
   }
 
   // List has 1 or more members.
+  if (size == 1 && isa<clang::MaterializeTemporaryExpr>(x->getInit(0)))
+    if (auto alloc = peekValue().getDefiningOp<cudaq::cc::AllocaOp>())
+      if (auto arrTy = dyn_cast<cudaq::cc::ArrayType>(initListTy))
+        if (alloc.getElementType() == arrTy.getElementType())
+          return true;
   auto last = lastValues(size);
   bool allRef = std::all_of(last.begin(), last.end(), [](auto v) {
     return isa<quake::RefType, quake::VeqType>(v.getType());
@@ -2916,6 +2924,32 @@ bool QuakeBridgeVisitor::VisitCXXConstructExpr(clang::CXXConstructExpr *x) {
             loc, quake::VeqType::getUnsized(builder.getContext()), sizeVal));
       }
 
+      if (ctorName == "state") {
+        // cudaq::state ctor can be materialized when using local simulators and
+        // converting raw data to state vectors. Use a runtime helper function
+        // to perform the conversion.
+        Value stdvec = popValue();
+        auto stateTy = cudaq::cc::PointerType::get(
+            quake::StateType::get(builder.getContext()));
+        if (auto stdvecTy = dyn_cast<cudaq::cc::StdvecType>(stdvec.getType())) {
+          auto dataTy = cudaq::cc::PointerType::get(stdvecTy.getElementType());
+          Value data =
+              builder.create<cudaq::cc::StdvecDataOp>(loc, dataTy, stdvec);
+          auto i64Ty = builder.getI64Type();
+          Value size =
+              builder.create<cudaq::cc::StdvecSizeOp>(loc, i64Ty, stdvec);
+          return pushValue(builder.create<quake::CreateStateOp>(
+              loc, stateTy, ValueRange{data, size}));
+        }
+        if (auto alloc = stdvec.getDefiningOp<cudaq::cc::AllocaOp>()) {
+          Value size = alloc.getSeqSize();
+          return pushValue(builder.create<quake::CreateStateOp>(
+              loc, stateTy, ValueRange{alloc, size}));
+        }
+        TODO_loc(loc, "unhandled state constructor");
+        return false;
+      }
+
       // lambda determines: is `t` a cudaq::state* ?
       auto isStateType = [&](Type t) {
         if (auto ptrTy = dyn_cast<cc::PointerType>(t))
@@ -2925,9 +2959,17 @@ bool QuakeBridgeVisitor::VisitCXXConstructExpr(clang::CXXConstructExpr *x) {
 
       if (ctorName == "qudit") {
         auto initials = popValue();
+        if (isa<quake::StateType>(initials.getType()))
+          if (auto load = initials.getDefiningOp<cudaq::cc::LoadOp>())
+            initials = load.getPtrvalue();
         if (isStateType(initials.getType())) {
-          TODO_x(loc, x, mangler, "qudit(state) ctor");
-          return false;
+          Value alloca = builder.create<quake::AllocaOp>(loc);
+          auto veq1Ty = quake::VeqType::get(builder.getContext(), 1);
+          Value initSt = builder.create<quake::InitializeStateOp>(
+              loc, veq1Ty, ValueRange{alloca, initials});
+          if (auto initOp = initials.getDefiningOp<quake::CreateStateOp>())
+            builder.create<quake::DeleteStateOp>(loc, initOp);
+          return pushValue(builder.create<quake::ExtractRefOp>(loc, initSt, 0));
         }
         bool ok = false;
         if (auto ptrTy = dyn_cast<cc::PointerType>(initials.getType()))
@@ -2953,57 +2995,26 @@ bool QuakeBridgeVisitor::VisitCXXConstructExpr(clang::CXXConstructExpr *x) {
           return pushValue(builder.create<quake::AllocaOp>(
               loc, quake::VeqType::getUnsized(ctx), initials));
         }
-        if (isa<quake::StateType>(initials.getType())) {
+        if (isa<quake::StateType>(initials.getType()))
           if (auto load = initials.getDefiningOp<cudaq::cc::LoadOp>())
             initials = load.getPtrvalue();
-        }
         if (isStateType(initials.getType())) {
           Value state = initials;
           auto i64Ty = builder.getI64Type();
           auto numQubits =
               builder.create<quake::GetNumberOfQubitsOp>(loc, i64Ty, state);
           auto veqTy = quake::VeqType::getUnsized(ctx);
           Value alloc = builder.create<quake::AllocaOp>(loc, veqTy, numQubits);
-          return pushValue(builder.create<quake::InitializeStateOp>(
-              loc, veqTy, alloc, state));
+          Value initSt = builder.create<quake::InitializeStateOp>(loc, veqTy,
+                                                                  alloc, state);
+          if (auto initOp = initials.getDefiningOp<quake::CreateStateOp>())
+            builder.create<quake::DeleteStateOp>(loc, initOp);
+          return pushValue(initSt);
         }
-        // Otherwise, it is the cudaq::qvector(std::vector<complex>) ctor.
-        Value numQubits;
-        Type initialsTy = initials.getType();
-        if (auto ptrTy = dyn_cast<cc::PointerType>(initialsTy)) {
-          if (auto arrTy = dyn_cast<cc::ArrayType>(ptrTy.getElementType())) {
-            if (arrTy.isUnknownSize()) {
-              if (auto allocOp = initials.getDefiningOp<cc::AllocaOp>())
-                if (auto size = allocOp.getSeqSize())
-                  numQubits =
-                      builder.create<math::CountTrailingZerosOp>(loc, size);
-            } else {
-              std::size_t arraySize = arrTy.getSize();
-              if (!std::has_single_bit(arraySize)) {
-                reportClangError(x, mangler,
-                                 "state vector must be a power of 2 in length");
-              }
-              numQubits = builder.create<arith::ConstantIntOp>(
-                  loc, std::countr_zero(arraySize), 64);
-            }
-          }
-        } else if (auto stdvecTy = dyn_cast<cc::StdvecType>(initialsTy)) {
-          Value vecLen = builder.create<cc::StdvecSizeOp>(
-              loc, builder.getI64Type(), initials);
-          numQubits = builder.create<math::CountTrailingZerosOp>(loc, vecLen);
-          auto ptrTy = cc::PointerType::get(stdvecTy.getElementType());
-          initials = builder.create<cc::StdvecDataOp>(loc, ptrTy, initials);
-        }
-        if (!numQubits) {
-          reportClangError(
-              x, mangler,
-              "internal error: could not determine the number of qubits");
-          return false;
-        }
-        auto veqTy = quake::VeqType::getUnsized(ctx);
-        auto alloc = builder.create<quake::AllocaOp>(loc, veqTy, numQubits);
-        return pushValue(builder.create<quake::InitializeStateOp>(
-            loc, veqTy, alloc, initials));
+        reportClangError(
+            x, mangler,
+            "internal error: could not determine the number of qubits");
+        return false;
       }
       if ((ctorName == "qspan" || ctorName == "qview") &&
           isa<quake::VeqType>(peekValue().getType())) {

lib/Optimizer/Builder/Intrinsics.cpp

@@ -304,11 +304,17 @@ static constexpr IntrinsicCode intrinsicTable[] = {
   }
 )#"},
 
-    {cudaq::createCudaqStateFromDataFP32, {}, R"#(
-  func.func private @__nvqpp_cudaq_state_createFromData_fp32(%p : !cc.ptr<i8>, %s : i64) -> !cc.ptr<!quake.state>
+    {cudaq::createCudaqStateFromDataComplexF32, {}, R"#(
+  func.func private @__nvqpp_cudaq_state_createFromData_complex_f32(%p : !cc.ptr<i8>, %s : i64) -> !cc.ptr<!quake.state>
 )#"},
-    {cudaq::createCudaqStateFromDataFP64, {}, R"#(
-  func.func private @__nvqpp_cudaq_state_createFromData_fp64(%p : !cc.ptr<i8>, %s : i64) -> !cc.ptr<!quake.state>
+    {cudaq::createCudaqStateFromDataComplexF64, {}, R"#(
+  func.func private @__nvqpp_cudaq_state_createFromData_complex_f64(%p : !cc.ptr<i8>, %s : i64) -> !cc.ptr<!quake.state>
+)#"},
+    {cudaq::createCudaqStateFromDataF32, {}, R"#(
+  func.func private @__nvqpp_cudaq_state_createFromData_f32(%p : !cc.ptr<i8>, %s : i64) -> !cc.ptr<!quake.state>
+)#"},
+    {cudaq::createCudaqStateFromDataF64, {}, R"#(
+  func.func private @__nvqpp_cudaq_state_createFromData_f64(%p : !cc.ptr<i8>, %s : i64) -> !cc.ptr<!quake.state>
 )#"},
 
     {cudaq::deleteCudaqState, {}, R"#(

lib/Optimizer/CodeGen/ConvertToQIRAPI.cpp

@@ -785,8 +785,8 @@ struct QmemRAIIOpRewrite : public OpConversionPattern<cudaq::codegen::RAIIOp> {
 
     Value sizeOperand;
     if (!adaptor.getAllocSize()) {
-      auto type = cast<quake::VeqType>(allocTy);
-      auto constantSize = type.getSize();
+      auto type = dyn_cast<quake::VeqType>(allocTy);
+      auto constantSize = type ? type.getSize() : 1;
       sizeOperand =
           rewriter.create<arith::ConstantIntOp>(loc, constantSize, 64);
     } else {

lib/Optimizer/CodeGen/QuakeToCodegen.cpp

@@ -80,15 +80,27 @@ class CreateStateOpPattern : public OpRewritePattern<quake::CreateStateOp> {
 
     auto bufferTy = buffer.getType();
     auto ptrTy = cast<cudaq::cc::PointerType>(bufferTy);
-    auto arrTy = cast<cudaq::cc::ArrayType>(ptrTy.getElementType());
-    auto eleTy = arrTy.getElementType();
+    auto arrTy = dyn_cast<cudaq::cc::ArrayType>(ptrTy.getElementType());
+    auto eleTy = arrTy ? arrTy.getElementType() : ptrTy.getElementType();
     auto is64Bit = isa<Float64Type>(eleTy);
+    bool isComplex = false;
 
-    if (auto cTy = dyn_cast<ComplexType>(eleTy))
+    if (auto cTy = dyn_cast<ComplexType>(eleTy)) {
       is64Bit = isa<Float64Type>(cTy.getElementType());
+      isComplex = true;
+    }
+
+    auto createStateFunc = [&]() {
+      if (isComplex) {
+        if (is64Bit)
+          return cudaq::createCudaqStateFromDataComplexF64;
+        return cudaq::createCudaqStateFromDataComplexF32;
+      }
+      if (is64Bit)
+        return cudaq::createCudaqStateFromDataF64;
+      return cudaq::createCudaqStateFromDataF32;
+    }();
 
-    auto createStateFunc = is64Bit ? cudaq::createCudaqStateFromDataFP64
-                                   : cudaq::createCudaqStateFromDataFP32;
     cudaq::IRBuilder irBuilder(ctx);
     auto result = irBuilder.loadIntrinsic(module, createStateFunc);
     assert(succeeded(result) && "loading intrinsic should never fail");
@@ -97,11 +109,9 @@ class CreateStateOpPattern : public OpRewritePattern<quake::CreateStateOp> {
     auto statePtrTy = cudaq::cc::PointerType::get(stateTy);
     auto i8PtrTy = cudaq::cc::PointerType::get(rewriter.getI8Type());
     auto cast = rewriter.create<cudaq::cc::CastOp>(loc, i8PtrTy, buffer);
-    auto one = rewriter.create<arith::ConstantIntOp>(loc, 1, size.getType());
-    auto powsz = rewriter.create<arith::ShLIOp>(loc, size.getType(), one, size);
 
     rewriter.replaceOpWithNewOp<func::CallOp>(
-        createStateOp, statePtrTy, createStateFunc, ValueRange{cast, powsz});
+        createStateOp, statePtrTy, createStateFunc, ValueRange{cast, size});
     return success();
   }
 };

lib/Optimizer/CodeGen/VerifyNVQIRCalls.cpp

@@ -46,8 +46,10 @@ struct VerifyNVQIRCallOpsPass
           cudaq::opt::QIRArrayQubitAllocateArrayWithStateComplex32,
           cudaq::opt::QIRArrayQubitAllocateArrayWithStateComplex64,
           cudaq::getNumQubitsFromCudaqState,
-          cudaq::createCudaqStateFromDataFP32,
-          cudaq::createCudaqStateFromDataFP64,
+          cudaq::createCudaqStateFromDataComplexF32,
+          cudaq::createCudaqStateFromDataComplexF64,
+          cudaq::createCudaqStateFromDataF32,
+          cudaq::createCudaqStateFromDataF64,
           cudaq::deleteCudaqState};
       // It must be either NVQIR extension functions or in the allowed list.
       return std::find(NVQIR_FUNCS.begin(), NVQIR_FUNCS.end(), functionName) !=

runtime/common/ExecutionContext.h

@@ -11,7 +11,6 @@
 #include "Future.h"
 #include "NoiseModel.h"
 #include "SampleResult.h"
-#include "SimulationState.h"
 #include "Trace.h"
 #include "cudaq/algorithms/optimizer.h"
 #include "cudaq/operators.h"
@@ -20,6 +19,8 @@
 
 namespace cudaq {
 
+class SimulationState;
+
 /// The ExecutionContext is an abstraction to indicate how a CUDA-Q kernel
 /// should be executed.
 class ExecutionContext {

runtime/common/SimulationState.h

@@ -8,7 +8,10 @@
 
 #pragma once
 
+#include "cudaq/utils/cudaq_utils.h"
+#include "cudaq/utils/matrix.h"
 #include <algorithm>
+#include <bitset>
 #include <complex>
 #include <memory>
 #include <optional>
@@ -28,10 +31,11 @@ using TensorStateData =
 /// @brief state_data is a variant type
 /// encoding different forms of user state vector data
 /// we support.
-using state_data = std::variant<
-    std::vector<std::complex<double>>, std::vector<std::complex<float>>,
-    std::pair<std::complex<double> *, std::size_t>,
-    std::pair<std::complex<float> *, std::size_t>, TensorStateData>;
+using state_data = std::variant<std::vector<std::complex<double>>,
+                                std::vector<std::complex<float>>,
+                                std::pair<std::complex<double> *, std::size_t>,
+                                std::pair<std::complex<float> *, std::size_t>,
+                                complex_matrix, TensorStateData>;
 
 /// @brief The `SimulationState` interface provides and extension point
 /// for concrete circuit simulation sub-types to describe their
@@ -69,20 +73,34 @@ class SimulationState {
   /// and extract the data pointer and size.
   template <typename ScalarType = double>
   auto getSizeAndPtr(const state_data &data) {
-    auto type = data.index();
-    std::tuple<std::size_t, void *> sizeAndPtr;
-    if (type == 0)
-      sizeAndPtr = getSizeAndPtrFromVec<double, ScalarType>(data);
-    else if (type == 1)
-      sizeAndPtr = getSizeAndPtrFromVec<float, ScalarType>(data);
-    else if (type == 2)
-      sizeAndPtr = getSizeAndPtrFromPair<double, ScalarType>(data);
-    else if (type == 3)
-      sizeAndPtr = getSizeAndPtrFromPair<float, ScalarType>(data);
-    else
-      throw std::runtime_error("unsupported data type for state.");
-
-    return sizeAndPtr;
+    if (std::holds_alternative<std::vector<std::complex<double>>>(data))
+      return getSizeAndPtrFromVec<double, ScalarType>(data);
+    if (std::holds_alternative<std::vector<std::complex<float>>>(data))
+      return getSizeAndPtrFromVec<float, ScalarType>(data);
+    if (std::holds_alternative<std::pair<std::complex<double> *, std::size_t>>(
+            data))
+      return getSizeAndPtrFromPair<double, ScalarType>(data);
+    if (std::holds_alternative<std::pair<std::complex<float> *, std::size_t>>(
+            data))
+      return getSizeAndPtrFromPair<float, ScalarType>(data);
+    if (std::holds_alternative<complex_matrix>(data)) {
+      // Complex matrix is double precision only
+      if constexpr (!std::is_same_v<double, ScalarType>)
+        throw std::runtime_error("[sim-state] invalid data precision.");
+      auto &cMat = std::get<complex_matrix>(data);
+      if (cMat.rows() != cMat.cols())
+        throw std::runtime_error(
+            "[sim-state] complex matrix must be square for density matrix.");
+      // Check that it must be a power of 2
+      if (std::bitset<64>(cMat.rows()).count() != 1)
+        throw std::runtime_error("[sim-state] complex matrix size must be a "
+                                 "power of 2 for density matrix.");
+      return std::make_tuple(
+          cMat.size(),
+          reinterpret_cast<void *>(const_cast<complex_matrix &>(cMat).get_data(
+              complex_matrix::order::row_major)));
+    }
+    throw std::runtime_error("unsupported data type for state vector.");
   }
 
   /// @brief Subclass-specific creator method for