introduced new operation cudaq::save_state and support

kvmto · kvmto · commit 0c43b7050356 · 2025-09-02T14:18:03.000Z
Signed-off-by: Kevin Mato &lt;kmato@nvidia.com&gt;
diff --git a/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h b/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h
@@ -122,6 +122,8 @@ static constexpr const char QISApplyKrausChannel[] =
 
 static constexpr const char QISTrap[] = "__quantum__qis__trap";
 
+static constexpr const char QISSaveState[] = "__quantum__qis__save_state";
+
 /// Since apply noise is actually a call back to `C++` code, the `QIR` data type
 /// `Array` of `Qubit*` must be converted into a `cudaq::qvector`, which is
 /// presently a `std::vector<cudaq::qubit>` but with an extremely restricted
diff --git a/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td b/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.td
@@ -530,6 +530,17 @@ def quake_ApplyNoiseOp : QuakeOp<"apply_noise", [AttrSizedOperandSegments]> {
   }];
 }
 
+def quake_SaveStateOp : QuakeOp<"save_state"> {
+  let summary = "Save quantum state representation compatible with the simulator.";
+  let description = [{
+    This operation provides support for the `cudaq::save_state`
+    function. This function is only valid in simulation contexts where the
+    simulator is part of the same process as the C++ host executable itself.
+  }];
+
+  // No arguments are needed.
+}
+
 //===----------------------------------------------------------------------===//
 // Memory and register conversion instructions: These operations are useful for
 // intermediate conversions between memory-SSA and value-SSA semantics and vice
diff --git a/lib/Frontend/nvqpp/ConvertExpr.cpp b/lib/Frontend/nvqpp/ConvertExpr.cpp
@@ -1634,6 +1634,13 @@ bool QuakeBridgeVisitor::VisitCallExpr(clang::CallExpr *x) {
       return false;
     }
 
+
+
+    if (funcName == "save_state") {
+      builder.create<quake::SaveStateOp>(loc, TypeRange{}, ValueRange{});
+      return true;
+    }
+
     if (funcName == "mx" || funcName == "my" || funcName == "mz") {
       // Measurements always return a bool or a std::vector<bool>.
       bool useStdvec =
diff --git a/lib/Optimizer/Builder/Intrinsics.cpp b/lib/Optimizer/Builder/Intrinsics.cpp
@@ -536,6 +536,7 @@ static constexpr IntrinsicCode intrinsicTable[] = {
   func.func private @__quantum__qis__convert_array_to_stdvector(!qir_array) -> !qir_array
   func.func private @__quantum__qis__free_converted_stdvector(!qir_array)
 
+  func.func private @__quantum__qis__save_state()
   llvm.func @generalizedInvokeWithRotationsControlsTargets(i64, i64, i64, i64, !qir_llvmptr, ...) attributes {sym_visibility = "private"}
   llvm.func @__quantum__qis__apply_kraus_channel_generalized(i64, i64, i64, i64, i64, ...) attributes {sym_visibility = "private"}
 )#"},
diff --git a/lib/Optimizer/CodeGen/ConvertToQIRAPI.cpp b/lib/Optimizer/CodeGen/ConvertToQIRAPI.cpp
@@ -1515,6 +1515,20 @@ struct AllocaOpPattern : public OpConversionPattern<cudaq::cc::AllocaOp> {
   }
 };
 
+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{});
+    return success();
+  }
+};
+
+
 /// Convert the quake types in `func::FuncOp` signatures.
 struct FuncSignaturePattern : public OpConversionPattern<func::FuncOp> {
   using OpConversionPattern::OpConversionPattern;
diff --git a/lib/Optimizer/Transforms/CMakeLists.txt b/lib/Optimizer/Transforms/CMakeLists.txt
@@ -29,6 +29,7 @@ add_cudaq_library(OptTransforms
   DistributedDeviceCall.cpp
   EraseNoise.cpp
   EraseNopCalls.cpp
+  EraseSaveState.cpp
   EraseVectorCopyCtor.cpp
   ExpandControlVeqs.cpp
   ExpandMeasurements.cpp
diff --git a/lib/Optimizer/Transforms/EraseSaveState.cpp b/lib/Optimizer/Transforms/EraseSaveState.cpp
@@ -0,0 +1,43 @@
+/*******************************************************************************
+ * Copyright (c) 2025 NVIDIA Corporation & Affiliates.                         *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include "PassDetails.h"
+#include "cudaq/Optimizer/Builder/Intrinsics.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#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"
+} // namespace cudaq::opt
+
+#define DEBUG_TYPE "erase-save-state"
+
+using namespace mlir;
+
+/// \file
+/// 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> {
+public:
+  using OpRewritePattern<Op>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(Op saveState,
+                                PatternRewriter &rewriter) const override {
+    rewriter.eraseOp(saveState);
+    return success();
+  }
+};
+
+} // namespace
diff --git a/python/cudaq/kernel/ast_bridge.py b/python/cudaq/kernel/ast_bridge.py
@@ -2626,6 +2626,11 @@ def bodyBuilder(iterVal):
                     quake.ApplyNoiseOp(params, [asVeq], key=key)
                     return
 
+
+                if node.func.attr == 'save_state':
+                    quake.SaveStateOp()
+                    return
+
                 if node.func.attr == 'compute_action':
                     # There can only be 2 arguments here.
                     action = None
diff --git a/python/tests/builder/test_save_state.py b/python/tests/builder/test_save_state.py
@@ -0,0 +1,39 @@
+# ============================================================================ #
+# Copyright (c) 2025 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+
+import os
+
+import pytest
+import numpy as np
+
+import cudaq
+
+@pytest.mark.parametrize('target', ['density-matrix-cpu', 'stim'])
+def test_save_state_builtin(target: str):
+    cudaq.set_target(target)
+
+    noise = cudaq.NoiseModel()
+
+    @cudaq.kernel
+    def bell_depol2(d: float, flag: bool):
+        q, r = cudaq.qubit(), cudaq.qubit()
+        h(q)
+        cudaq.save_state()
+       
+        x.ctrl(q, r) 
+        cudaq.save_state()
+
+        if flag:
+            cudaq.apply_noise(cudaq.Depolarization2, d, q, r)
+        else:
+            cudaq.apply_noise(cudaq.Depolarization2, [d], q, r)
+
+    counts = cudaq.sample(bell_depol2, 0.2, True, noise_model=noise)
+    assert len(counts) == 4
+    print(counts)
+
diff --git a/runtime/common/ExecutionContext.h b/runtime/common/ExecutionContext.h
@@ -142,5 +142,19 @@ class ExecutionContext {
   /// Note: Measurement Syndrome Matrix is defined in
   /// https://arxiv.org/pdf/2407.13826.
   std::optional<std::pair<std::size_t, std::size_t>> msm_dimensions;
+
+  /// @brief For each possible error, this is a "flips" vector of length "num
+  /// qubits", where "num qubits" is the number of qubits known to the simulator
+  /// at the time of the error mechanism. This is populated when using the "msm"
+  /// mode (i.e. this->name == "msm")
+  std::vector<std::vector<bool>> msm_x_flips; // msm_x_flips[error_id][qubit_id]
+  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]
+
+  /// @brief Save the current simulation state in the recorded states storage.
+  void save_state(const SimulationState *state);
 };
 } // namespace cudaq
diff --git a/runtime/nvqir/CircuitSimulator.h b/runtime/nvqir/CircuitSimulator.h
@@ -149,6 +149,15 @@ class CircuitSimulator {
   /// https://arxiv.org/pdf/2407.13826.
   virtual void generateMSM() {}
 
+  /// @brief Return the internal state representation. This
+  /// is meant for subtypes to override
+  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;
+
   /// @brief Apply exp(-i theta PauliTensorProd) to the underlying state.
   /// This must be provided by subclasses.
   virtual void applyExpPauli(double theta,
@@ -540,6 +549,13 @@ class CircuitSimulatorBase : public CircuitSimulator {
         "Simulation data not available for this simulator backend.");
   }
 
+  /// @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() {
+        throw std::runtime_error("Simulation data not available for this simulator backend.");
+  }
+
   /// @brief Handle basic sampling tasks by storing the qubit index for
   /// processing in resetExecutionContext. Return true to indicate this is
   /// sampling and to exit early. False otherwise.
diff --git a/runtime/nvqir/NVQIR.cpp b/runtime/nvqir/NVQIR.cpp
@@ -839,6 +839,18 @@ void __quantum__qis__apply_kraus_channel_generalized(
   va_end(args);
 }
 
+
+static void
+__quantum__qis__save_state() {
+
+  auto *ctx = nvqir::getCircuitSimulatorInternal()->getExecutionContext();
+  if (!ctx)
+    return;
+
+  std::unique_ptr<cudaq::SimulationState> state = nvqir::getCircuitSimulatorInternal()->getCurrentSimulationState();
+  ctx->save_state(state.get());
+}
+
 namespace details {
 struct FakeQubit {
   std::int8_t *id;
diff --git a/unittests/integration/save_state_tester.cpp b/unittests/integration/save_state_tester.cpp
@@ -0,0 +1,102 @@
+/*******************************************************************************
+ * Copyright (c) 2025 NVIDIA Corporation & Affiliates.                         *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include "CUDAQTestUtils.h"
+#include "cudaq/builder/kernels.h"
+#include <iostream>
+
+
+#if defined(CUDAQ_BACKEND_STIM)
+CUDAQ_TEST(SaveStateTester, checkStimState) {
+
+ struct multi_round_kernel {
+    void operator()(int num_qubits, int num_rounds,
+                    double noise_probability) __qpu__ {
+      cudaq::qvector q(num_qubits);
+      for (int round = 0; round < num_rounds; round++) {
+        h(q[0]);
+        for (int qi = 0; qi < num_qubits; qi++)
+          cudaq::apply_noise<cudaq::depolarization_channel>(noise_probability,
+                                                            q[qi]);
+        for (int qi = 1; qi < num_qubits; qi++)
+          x<cudaq::ctrl>(q[qi - 1], q[qi]);
+          
+		// as well here, cudaq::snap_state();
+        auto syndrome_z = syn_z(q);
+    	auto syndrome_x = syn_x(q);
+    	cudaq::save_state();
+    	/* decoding at each round */
+      }
+      mz(q);
+      for (int qi = 0; qi < num_qubits; qi++)
+        reset(q[qi]);
+    }
+  };
+
+  int num_qubits = 5;
+  int num_rounds = 3;
+  double noise_bf_prob = 0.0625;
+
+  cudaq::noise_model noise;
+  cudaq::bit_flip_channel bf(noise_bf_prob);
+  for (std::size_t i = 0; i < num_qubits; i++)
+    noise.add_channel("mz", {i}, bf);
+  cudaq::set_noise(noise);
+
+  cudaq::ExecutionContext ctx_msm_size("msm_size");
+  auto &platform = cudaq::get_platform();
+  platform.set_exec_ctx(&ctx_msm_size);
+  multi_round_kernel{}(num_qubits, num_rounds);
+  platform.reset_exec_ctx();
+
+  cudaq::ExecutionContext ctx_msm("msm");
+  ctx_msm.noiseModel = &noise;
+  ctx_msm.msm_dimensions = ctx_msm_size.msm_dimensions;
+  platform.set_exec_ctx(&ctx_msm);
+  multi_round_kernel{}(num_qubits, num_rounds);
+  platform.reset_exec_ctx();
+  
+  // accessing the execution context to get the recorded states
+  assert(ctx_msm.recorded_states.size() == num_rounds && "Expected 3 state snapshots");
+
+  for (std::size_t round = 0; round < ctx_msm.recorded_states.size(); ++round) {
+    const auto &state = ctx_msm.recorded_states[round];
+
+    assert(state.getNumQubits() == num_qubits && "Number of qubits must equal num_qubits");
+
+    for (std::size_t qi = 0; qi < state.getNumQubits(); ++qi) {
+        const auto &x_row = state.getTableau().x_output[qi];
+        const auto &z_row = state.getTableau().z_output[qi];
+
+        // Each row should have the same number of qubits.
+        assert(x_row.num_qubits == state.getNumQubits());
+        assert(z_row.num_qubits == state.getNumQubits());
+
+        // Check that each character in the row is a valid Pauli symbol.
+        for (char c : x_row.str()) {
+            assert((c == 'I' || c == 'X' || c == 'Y' || c == 'Z') &&
+                   "Invalid symbol in tableau X-output row");
+        }
+        for (char c : z_row.str()) {
+            assert((c == 'I' || c == 'X' || c == 'Y' || c == 'Z') &&
+                   "Invalid symbol in tableau Z-output row");
+        }
+    }
+    const auto &frame = state.getPauliFrame();
+
+    // Frame size should match number of qubits.
+    assert(frame.size() == state.getNumQubits());
+
+    // Each frame entry must be a valid Pauli operator.
+    for (char c : frame) {
+        assert((c == 'I' || c == 'X' || c == 'Y' || c == 'Z') &&
+               "Invalid symbol in Pauli frame");
+    }
+  }
+}
+#endif
