wip, save_state lowered. need for library mode and llvm translate

Signed-off-by: Kevin Mato <[email protected]>

Author: kvmto

lib/Frontend/nvqpp/ConvertExpr.cpp

@@ -1634,8 +1634,6 @@ bool QuakeBridgeVisitor::VisitCallExpr(clang::CallExpr *x) {
       return false;
     }
 
-
-
     if (funcName == "save_state") {
       builder.create<quake::SaveStateOp>(loc, TypeRange{}, ValueRange{});
       return true;

lib/Optimizer/CodeGen/ConvertToQIRAPI.cpp

@@ -1516,19 +1516,18 @@ struct AllocaOpPattern : public OpConversionPattern<cudaq::cc::AllocaOp> {
   }
 };
 
-struct SaveStateOpRewrite
-    : public OpConversionPattern<quake::SaveStateOp> {
+struct SaveStateOpRewrite : public OpConversionPattern<quake::SaveStateOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(quake::SaveStateOp saveState, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<func::CallOp>(saveState, TypeRange{}, cudaq::opt::QISSaveState, ValueRange{});
+    rewriter.replaceOpWithNewOp<func::CallOp>(
+        saveState, TypeRange{}, cudaq::opt::QISSaveState, ValueRange{});
     return success();
   }
 };
 
-
 /// Convert the quake types in `func::FuncOp` signatures.
 struct FuncSignaturePattern : public OpConversionPattern<func::FuncOp> {
   using OpConversionPattern::OpConversionPattern;

lib/Optimizer/Transforms/EraseSaveState.cpp

@@ -12,7 +12,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
- 
+
 namespace cudaq::opt {
 #define GEN_PASS_DEF_ERASESAVESTATE
 #include "cudaq/Optimizer/Transforms/Passes.h.inc"
@@ -26,7 +26,6 @@ using namespace mlir;
 /// This pass exists simply to remove all the quake.save_state (and related)
 /// Ops from the IR.
 
-
 namespace {
 template <typename Op>
 class EraseSaveStatePattern : public OpRewritePattern<Op> {
@@ -40,4 +39,4 @@ class EraseSaveStatePattern : public OpRewritePattern<Op> {
   }
 };
 
-} // namespace
+} // namespace

runtime/common/ExecutionContext.h

@@ -15,93 +15,101 @@
 #include "Trace.h"
 #include "cudaq/algorithms/optimizer.h"
 #include "cudaq/operators.h"
+#include <iostream>
 #include <optional>
 #include <string_view>
-#include <iostream>
 
 #include "nvqir/stim/StimState.h"
 
 namespace cudaq {
 
-struct RecordStorage{
+struct RecordStorage {
 
   size_t memory_limit;
   size_t current_memory;
   RecordStorage(size_t limit = 1e9) : memory_limit(limit), current_memory(0) {}
 
   std::vector<std::unique_ptr<SimulationState>> recordedStates;
 
-    void save_state(SimulationState *state) {
-        recordedStates.push_back(clone_state(state));
-    }
-    const std::vector<std::unique_ptr<SimulationState>>& get_recorded_states() const { return recordedStates; }
-    
-    void clear() { recordedStates.clear(); }
-    void dump_recorded_states() const {
-      for (std::size_t i = 0; i < recordedStates.size(); i++) {
-        recordedStates[i]->dump(std::cout);
-      }
+  void save_state(SimulationState *state) {
+    recordedStates.push_back(clone_state(state));
+  }
+  const std::vector<std::unique_ptr<SimulationState>> &
+  get_recorded_states() const {
+    return recordedStates;
+  }
+
+  void clear() { recordedStates.clear(); }
+  void dump_recorded_states() const {
+    for (std::size_t i = 0; i < recordedStates.size(); i++) {
+      recordedStates[i]->dump(std::cout);
     }
-    
+  }
+
 private:
-    std::unique_ptr<SimulationState> clone_state(SimulationState* state) {
-        if (state->isArrayLike()) {
-            // Handle array-like states (CusvState, etc.)
-            return clone_array_like_state(state);
-        } else {
-            // Handle specialized states (CuDensityMatState, StimState, etc.)
-            return clone_specialized_state(state);
-        }
+  std::unique_ptr<SimulationState> clone_state(SimulationState *state) {
+    if (state->isArrayLike()) {
+      // Handle array-like states (CusvState, etc.)
+      return clone_array_like_state(state);
+    } else {
+      // Handle specialized states (CuDensityMatState, StimState, etc.)
+      return clone_specialized_state(state);
+    }
+  }
+
+  std::unique_ptr<SimulationState>
+  clone_array_like_state(SimulationState *state) {
+    auto numQubits = state->getNumQubits();
+    if (numQubits > 20) { // Prevent exponential explosion
+      throw std::runtime_error("State too large to clone via amplitudes");
+    }
+
+    // Generate all basis states
+    auto totalStates = 1ULL << numQubits;
+    std::vector<std::vector<int>> basisStates;
+    for (size_t i = 0; i < totalStates; ++i) {
+      std::vector<int> basis(numQubits);
+      for (size_t j = 0; j < numQubits; ++j) {
+        basis[j] = (i >> j) & 1;
+      }
+      basisStates.push_back(basis);
     }
-    
-    std::unique_ptr<SimulationState> clone_array_like_state(SimulationState* state) {
-        auto numQubits = state->getNumQubits();
-        if (numQubits > 20) {  // Prevent exponential explosion
-            throw std::runtime_error("State too large to clone via amplitudes");
-        }
-        
-        // Generate all basis states
-        auto totalStates = 1ULL << numQubits;
-        std::vector<std::vector<int>> basisStates;
-        for (size_t i = 0; i < totalStates; ++i) {
-            std::vector<int> basis(numQubits);
-            for (size_t j = 0; j < numQubits; ++j) {
-                basis[j] = (i >> j) & 1;
-            }
-            basisStates.push_back(basis);
-        }
-        
-        auto amplitudes = state->getAmplitudes(basisStates);
-        
-        // Create new state with appropriate precision
-        if (state->getPrecision() == SimulationState::precision::fp32) {
-            std::vector<std::complex<float>> floatAmps;
-            for (const auto& amp : amplitudes) {
-                floatAmps.emplace_back(static_cast<float>(amp.real()), 
-                                     static_cast<float>(amp.imag()));
-            }
-            return state->createFromData(floatAmps);
-        } else {
-            return state->createFromData(amplitudes);
-        }
+
+    auto amplitudes = state->getAmplitudes(basisStates);
+
+    // Create new state with appropriate precision
+    if (state->getPrecision() == SimulationState::precision::fp32) {
+      std::vector<std::complex<float>> floatAmps;
+      for (const auto &amp : amplitudes) {
+        floatAmps.emplace_back(static_cast<float>(amp.real()),
+                               static_cast<float>(amp.imag()));
+      }
+      return state->createFromData(floatAmps);
+    } else {
+      return state->createFromData(amplitudes);
     }
-    
-    std::unique_ptr<SimulationState> clone_specialized_state(SimulationState* state) {
-        // Try dynamic_cast to known types that have clone methods
-        // this triggerd fatal error: library_types.h: No such file or directory
-        //if (auto* densityState = dynamic_cast<const CuDensityMatState*>(state)) {
-        //    return CuDensityMatState::clone(*densityState);
-        //}
-        
-        if (auto* stimState = dynamic_cast<const StimState*>(state)) {
-            return stimState->clone();
-        }
-        
-        // For unknown specialized types, try createFromSizeAndPtr as fallback
-        // This might work for some specialized states
-        // auto tensor = state->getTensor(0); 
-        // return state->createFromSizeAndPtr(tensor.get_num_elements(), tensor.data, 1);
+  }
+
+  std::unique_ptr<SimulationState>
+  clone_specialized_state(SimulationState *state) {
+    // Try dynamic_cast to known types that have clone methods
+    // this triggerd fatal error: library_types.h: No such file or directory
+    // if (auto* densityState = dynamic_cast<const CuDensityMatState*>(state)) {
+    //    return CuDensityMatState::clone(*densityState);
+    //}
+
+    if (auto *cloneable = dynamic_cast<ClonableState *>(state)) {
+      return cloneable->clone();
     }
+
+    // Fallback for non-cloneable specialized states
+    throw std::runtime_error("Specialized state type does not support cloning");
+    // For unknown specialized types, try createFromSizeAndPtr as fallback
+    // This might work for some specialized states
+    // auto tensor = state->getTensor(0);
+    // return state->createFromSizeAndPtr(tensor.get_num_elements(),
+    // tensor.data, 1);
+  }
 };
 
 /// The ExecutionContext is an abstraction to indicate how a CUDA-Q kernel
@@ -239,27 +247,24 @@ class ExecutionContext {
   std::vector<std::vector<bool>> msm_z_flips; // msm_z_flips[error_id][qubit_id]
 
   /// @brief For each shot, this is a vector of error IDs.
-  /// This is populated when using the "sample" mode (i.e. this->name == "sample")
-  std::vector<std::vector<std::size_t>> errors_per_shot; // errors_per_shot[shot][error_id]
+  /// This is populated when using the "sample" mode (i.e. this->name ==
+  /// "sample")
+  std::vector<std::vector<std::size_t>>
+      errors_per_shot; // errors_per_shot[shot][error_id]
 
   /// @brief Save the current simulation state in the recorded states storage.
-  void save_state(SimulationState *state){
-    recordStorage.save_state(state);
-  }
+  void save_state(SimulationState *state) { recordStorage.save_state(state); }
 
   /// @brief Get the recorded states saved during execution.
-  const std::vector<std::unique_ptr<SimulationState>>& get_recorded_states() const {
+  const std::vector<std::unique_ptr<SimulationState>> &
+  get_recorded_states() const {
     return recordStorage.get_recorded_states();
   }
 
   /// @brief Clear the recorded states saved during execution.
-  void clear_recorded_states(){
-    recordStorage.clear();
-  }
+  void clear_recorded_states() { recordStorage.clear(); }
 
   /// @brief Dump the recorded states saved during execution.
-  void dump_recorded_states() const{
-    recordStorage.dump_recorded_states();
-  }
+  void dump_recorded_states() const { recordStorage.dump_recorded_states(); }
 };
 } // namespace cudaq

runtime/common/SimulationState.h

@@ -17,20 +17,21 @@
 
 namespace cudaq {
 class SimulationState;
+class ClonableState;
 
 /// Enum to specify the initial quantum state.
 enum class InitialState { ZERO, UNIFORM };
 
-
 /// @brief StimData now stores a list of (pointer, size) pairs
 /// according to convention:
 /// 0: pointer to num_qubits, size = 1
 /// 1: pointer to msm_err_count, size = 1
 /// 2: pointer to num_stabilizers, size = 1
-/// 3: x_output array, size = x_output_size (X stabilizers + destabilisers + phase bits)
-/// 4: z_output array, size = z_output_size (Z stabilizers + destabilisers + phase bits)
-/// 5: frame array (Pauli frame), size = 2*num_qubits
-using StimData = std::vector<std::pair<void*, std::size_t>>;
+/// 3: x_output array, size = x_output_size (X stabilizers + destabilisers +
+/// phase bits) 4: z_output array, size = z_output_size (Z stabilizers +
+/// destabilisers + phase bits) 5: frame array (Pauli frame), size =
+/// 2*num_qubits
+using StimData = std::vector<std::pair<void *, std::size_t>>;
 
 /// @brief Encapsulates a list of tensors (data pointer and dimensions).
 // Note: tensor data is expected in column-major.
@@ -42,8 +43,7 @@ using 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>, TensorStateData,
-    StimData>;
+    std::pair<std::complex<float> *, std::size_t>, TensorStateData, StimData>;
 
 /// @brief The `SimulationState` interface provides and extension point
 /// for concrete circuit simulation sub-types to describe their
@@ -149,7 +149,8 @@ class SimulationState {
             "Cannot initialize state vector/density matrix state by stim "
             "data. Please use stabilizer simulator backends.");
       auto &dataCasted = std::get<StimData>(data);
-      return createFromSizeAndPtr(dataCasted.size(), const_cast<StimData*>(&dataCasted), data.index());
+      return createFromSizeAndPtr(
+          dataCasted.size(), const_cast<StimData *>(&dataCasted), data.index());
     }
     // Flat array state data
     // Check the precision first. Get the size and
@@ -269,4 +270,12 @@ class SimulationState {
   /// @brief Destructor
   virtual ~SimulationState() {}
 };
+
+/// @brief Interface for SimulationState subtypes that support cloning.
+class ClonableState {
+public:
+  virtual ~ClonableState() = default;
+  virtual std::unique_ptr<SimulationState> clone() const = 0;
+};
+
 } // namespace cudaq

runtime/cudaq/builder/kernels.h

@@ -59,9 +59,9 @@ std::vector<double> getAlphaY(const std::span<double> data,
 /// to its internal representation. This implementation follows the algorithm
 /// defined in `https://arxiv.org/pdf/quant-ph/0407010.pdf`.
 template <typename Kernel>
-void from_state(Kernel &&kernel, QuakeValue &qubits,
-                const std::span<std::complex<double>> data,
-                std::size_t inNumQubits = 0) {
+inline void from_state(Kernel &&kernel, QuakeValue &qubits,
+                       const std::span<std::complex<double>> data,
+                       std::size_t inNumQubits = 0) {
   std::make_signed_t<std::size_t> numQubits =
       qubits.constantSize().value_or(inNumQubits);
   if (numQubits <= 0)
@@ -113,7 +113,7 @@ void from_state(Kernel &&kernel, QuakeValue &qubits,
 /// @brief Construct a CUDA-Q kernel that produces the
 /// given state. This overload will return the `kernel_builder` as a
 /// `unique_ptr`.
-auto from_state(const std::span<std::complex<double>> data) {
+inline auto from_state(const std::span<std::complex<double>> data) {
   auto numQubits = std::log2(data.size());
   std::vector<details::KernelBuilderType> empty;
   auto kernel = std::make_unique<kernel_builder<>>(empty);

runtime/cudaq/qis/qubit_qis.h

@@ -1386,6 +1386,8 @@ void apply_noise(Args &&...args) {
       details::tuple_slice_last<qubit_arity>(std::forward_as_tuple(args...)));
 }
 
+inline void save_state() { return; }
+
 } // namespace cudaq
 
 #define __qop__ __attribute__((annotate("user_custom_quantum_operation")))

runtime/nvqir/CircuitSimulator.h

@@ -151,12 +151,13 @@ class CircuitSimulator {
 
   /// @brief Return the internal state representation. This
   /// is meant for subtypes to override
-  virtual std::unique_ptr<cudaq::SimulationState> getSimulationState()  = 0;
+  virtual std::unique_ptr<cudaq::SimulationState> getSimulationState() = 0;
 
   /// @brief Get the current simulation state.
   /// The method returns the current state of teh simulation without flushing
   /// the gate queue.
-  virtual std::unique_ptr<cudaq::SimulationState> getCurrentSimulationState()  = 0;
+  virtual std::unique_ptr<cudaq::SimulationState>
+  getCurrentSimulationState() = 0;
 
   /// @brief Apply exp(-i theta PauliTensorProd) to the underlying state.
   /// This must be provided by subclasses.
@@ -553,7 +554,8 @@ class CircuitSimulatorBase : public CircuitSimulator {
   /// The method returns the current state of teh simulation without flushing
   /// the gate queue.
   virtual std::unique_ptr<cudaq::SimulationState> getCurrentSimulationState() {
-        throw std::runtime_error("Simulation data not available for this simulator backend.");
+    throw std::runtime_error(
+        "Simulation data not available for this simulator backend.");
   }
 
   /// @brief Handle basic sampling tasks by storing the qubit index for

runtime/nvqir/NVQIR.cpp

@@ -839,15 +839,14 @@ void __quantum__qis__apply_kraus_channel_generalized(
   va_end(args);
 }
 
-
-static void
-__quantum__qis__save_state() {
+static void __quantum__qis__save_state() {
 
   auto *ctx = nvqir::getCircuitSimulatorInternal()->getExecutionContext();
   if (!ctx)
     return;
 
-  std::unique_ptr<cudaq::SimulationState> state = nvqir::getCircuitSimulatorInternal()->getCurrentSimulationState();
+  std::unique_ptr<cudaq::SimulationState> state =
+      nvqir::getCircuitSimulatorInternal()->getCurrentSimulationState();
   ctx->save_state(state.get());
 }