Add C++ kernel_builder::qalloc(state&) (#1722)

* Add C++ kernel_builder::qalloc(state&) * Simplify implementation after merge * Further simplify the code
NVIDIA · May 31, 2024 · 8cb7ec9 · 8cb7ec9
1 parent a6ca2dd
commit 8cb7ec9
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/runtime/cudaq/builder/kernel_builder.cpp b/runtime/cudaq/builder/kernel_builder.cpp
@@ -619,6 +619,22 @@ QuakeValue qalloc(ImplicitLocOpBuilder &builder,
   return QuakeValue(builder, qubits);
 }
 
+QuakeValue qalloc(mlir::ImplicitLocOpBuilder &builder, cudaq::state *state,
+                  StateVectorStorage &stateVectorStorage) {
+  auto *context = builder.getContext();
+  auto statePtrTy = cudaq::cc::PointerType::get(cc::StateType::get(context));
+  auto statePtr = builder.create<cc::CastOp>(
+      builder.getLoc(), statePtrTy,
+      builder.create<arith::ConstantIntOp>(
+          reinterpret_cast<std::intptr_t>(state), 64));
+  // Add the initialize state op
+  Value qubits = builder.create<quake::AllocaOp>(
+      quake::VeqType::get(context, state->get_num_qubits()));
+  qubits = builder.create<quake::InitializeStateOp>(qubits.getType(), qubits,
+                                                    statePtr);
+  return QuakeValue(builder, qubits);
+}
+
 QuakeValue constantVal(ImplicitLocOpBuilder &builder, double val) {
   llvm::APFloat d(val);
   Value constant =

diff --git a/runtime/cudaq/builder/kernel_builder.h b/runtime/cudaq/builder/kernel_builder.h
@@ -183,6 +183,10 @@ QuakeValue qalloc(mlir::ImplicitLocOpBuilder &builder,
                   StateVectorStorage &stateVectorData,
                   StateVectorVariant &&state, simulation_precision precision);
 
+/// @brief Allocate a `qvector` from a user provided state.
+QuakeValue qalloc(mlir::ImplicitLocOpBuilder &builder, cudaq::state *state,
+                  StateVectorStorage &stateVectorData);
+
 /// @brief Create a QuakeValue representing a constant floating-point number
 QuakeValue constantVal(mlir::ImplicitLocOpBuilder &builder, double val);
 
@@ -478,6 +482,12 @@ class kernel_builder : public details::kernel_builder_base {
                            simulation_precision::fp32);
   }
 
+  // Overload for `cudaq::state`
+  QuakeValue qalloc(const cudaq::state &state) {
+    return details::qalloc(*opBuilder.get(), const_cast<cudaq::state *>(&state),
+                           stateVectorStorage);
+  }
+
   /// @brief Return a `QuakeValue` representing the constant floating-point
   /// value.
   QuakeValue constantVal(double val) {

diff --git a/unittests/integration/builder_tester.cpp b/unittests/integration/builder_tester.cpp
@@ -1269,6 +1269,26 @@ TEST(BuilderTester, checkFromStateVector) {
     EXPECT_EQ(counts.size(), 1);
     EXPECT_EQ(counts.count("110"), 1000);
   }
+
+  {
+    // qalloc with state
+    auto kernel = cudaq::make_kernel();
+    auto qubits = kernel.qalloc(2);
+    kernel.h(qubits[0]);
+    auto state = cudaq::get_state(kernel);
+    // Send on the state to the next kernel
+    auto anotherOne = cudaq::make_kernel();
+    auto newQubits = anotherOne.qalloc(state);
+    anotherOne.x<cudaq::ctrl>(newQubits[0], newQubits[1]);
+    std::cout << anotherOne << "\n";
+    auto counts = cudaq::sample(anotherOne);
+    std::size_t counter = 0;
+    for (auto &[k, v] : counts) {
+      counter += v;
+      EXPECT_TRUE(k == "00" || k == "11");
+    }
+    EXPECT_EQ(counter, 1000);
+  }
 }
 
 CUDAQ_TEST(BuilderTester, checkCanProgressivelyBuild) {