[custatevec] getStateData to return data consistent with qpp and cudaq::from_state

runtime/cudaq/builder/kernels.cpp

@@ -48,7 +48,7 @@ std::vector<std::size_t> getControlIndices(std::size_t grayRank) {
   return ctrlIds;
 }
 
-std::size_t mEntry(std::size_t row, std::size_t col) {
+int mEntry(std::size_t row, std::size_t col) {
   auto b_and_g = row & ((col >> 1) ^ col);
   std::size_t sum_of_ones = 0;
   while (b_and_g > 0) {

runtime/nvqir/custatevec/CuStateVecCircuitSimulator.cu

@@ -520,9 +520,45 @@ public:
   }
 
   cudaq::State getStateData() override {
+    // Handle empty state (e.g., no qubit allocation)
+    if (stateDimension == 0)
+      return cudaq::State{{stateDimension}, {}};
+
     std::vector<std::complex<ScalarType>> tmp(stateDimension);
-    cudaMemcpy(tmp.data(), deviceStateVector,
-               stateDimension * sizeof(CudaDataType), cudaMemcpyDeviceToHost);
+    // Use custatevec accessor to retrieve the view
+    custatevecAccessorDescriptor_t accessor;
+    const uint32_t nIndexBits = std::log2(stateDimension);
+    // Note: we use MSB bit ordering when reporting the state vector
+    // hence, bit ordering vector = [N-1, N-2, ..., 0]
+    std::vector<int32_t> bitOrdering(nIndexBits);
+    std::iota(std::rbegin(bitOrdering), std::rend(bitOrdering), 0);
+    std::size_t extraWorkspaceSizeInBytes = 0;
+    // create accessor view
+    HANDLE_ERROR(custatevecAccessorCreateView(
+        handle, deviceStateVector, cuStateVecCudaDataType, nIndexBits,
+        &accessor, bitOrdering.data(), bitOrdering.size(),
+        /*maskBitString*/ nullptr, /*maskOrdering*/ nullptr,
+        /*maskLen*/ 0, &extraWorkspaceSizeInBytes));
+    // allocate external workspace if necessary
+    void *extraWorkspace = nullptr;
+    if (extraWorkspaceSizeInBytes > 0)
+      HANDLE_CUDA_ERROR(cudaMalloc(&extraWorkspace, extraWorkspaceSizeInBytes));
+
+    // set external workspace
+    HANDLE_ERROR(custatevecAccessorSetExtraWorkspace(
+        handle, accessor, extraWorkspace, extraWorkspaceSizeInBytes));
+
+    // get all state vector components: [0, stateDimension)
+    HANDLE_ERROR(custatevecAccessorGet(handle, accessor, tmp.data(),
+                                       /*begin*/ 0,
+                                       /*end*/
+                                       stateDimension));
+    // destroy descriptor
+    HANDLE_ERROR(custatevecAccessorDestroy(accessor));
+    // free extra workspace if allocated
+    if (extraWorkspaceSizeInBytes > 0)
+      HANDLE_CUDA_ERROR(cudaFree(extraWorkspace));
+
     if constexpr (std::is_same_v<ScalarType, float>) {
       std::vector<std::complex<double>> data;
       std::transform(tmp.begin(), tmp.end(), std::back_inserter(data),

unittests/integration/kernels_tester.cpp

@@ -151,6 +151,33 @@ CUDAQ_TEST(KernelsTester, checkFromState) {
     auto counts = cudaq::sample(*kernel);
     counts.dump();
   }
+
+  {
+    // Random unitary state
+    const std::size_t numQubits = 2;
+    auto randHam = cudaq::spin_op::random(numQubits, numQubits * numQubits);
+    auto eigenVectors = randHam.to_matrix().eigenvectors();
+    // Map the ground state to a cudaq::state
+    std::vector<std::complex<double>> expectedData(eigenVectors.rows());
+    for (std::size_t i = 0; i < eigenVectors.rows(); i++)
+      expectedData[i] = eigenVectors(i, 0);
+    auto kernel = cudaq::make_kernel();
+    auto qubits = kernel.qalloc(numQubits);
+    cudaq::from_state(kernel, qubits, expectedData);
+    std::cout << kernel << "\n";
+    auto ss = cudaq::get_state(kernel);
+    const std::complex<double> globalPhase = [&]() {
+      // find the first non-zero element to compute the global phase factor
+      for (std::size_t i = 0; i < (1u << numQubits); i++)
+        if (std::abs(ss[i]) > 1e-3)
+          return expectedData[i] / ss[i];
+      // Something wrong (state vector all zeros!)
+      return std::complex<double>(0.0, 0.0);
+    }();
+    // Check the state (accounted for global phase)
+    for (std::size_t i = 0; i < (1u << numQubits); i++)
+      EXPECT_NEAR(std::abs(globalPhase * ss[i] - expectedData[i]), 0.0, 1e-6);
+  }
 }
 
 #endif