Cherry pick PR #672 (22772f7 ) back, excluding decoding-server integration

Signed-off-by: Angela Burton <angelab@nvidia.com>

Author: anjbur

docs/sphinx/components/qec/introduction.rst

@@ -83,6 +83,12 @@ The code base class provides:
        using stabilizer_round = cudaq::qkernel<std::vector<cudaq::measure_result>(
            patch, const std::vector<std::size_t>&, const std::vector<std::size_t>&)>;
 
+   The two vector arguments of :code:`stabilizer_round` are the flattened X and
+   Z stabilizer *schedule* matrices, which can encode an optimized gate order
+   on top of the parity-check support. See
+   :cpp:func:`cudaq::qec::code::get_stabilizer_schedule_x` for the encoding
+   and the default (the plain parity matrices).
+
 4. **Protected Members**:
 
    - :code:`operation_encodings`: Maps operations to their quantum kernel implementations. The key is the ``operation`` enum and the value is a variant on the above kernel type aliases.
@@ -136,6 +142,20 @@ To implement a new quantum error correcting code:
            // Implement stabilizer measurements
        }
 
+   .. note::
+
+      The two vector arguments passed to the :code:`stabilizer_round` kernel
+      are the flattened X and Z stabilizer *schedule* matrices returned by
+      :cpp:func:`cudaq::qec::code::get_stabilizer_schedule_x` and
+      :cpp:func:`cudaq::qec::code::get_stabilizer_schedule_z`. By default
+      these equal the plain parity-check matrices (every entry 0 or 1), but a
+      code can override the :code:`get_stabilizer_schedule_*` methods to
+      encode a gate order, in which case entry :code:`k >= 1` means the
+      interaction executes at timestep :code:`k` (the built-in
+      :code:`surface_code` does this to avoid hook errors). A kernel that only
+      needs the support pattern should therefore test entries for
+      :code:`!= 0` rather than :code:`== 1`.
+
 4. **Register Operations**:
 
    In the constructor, register quantum kernels for each operation:
@@ -267,7 +287,18 @@ to prototype and develop new codes.
    .. note::
 
       The kernel registered for :code:`stabilizer_round` must be annotated to
-      return :code:`list[cudaq.measure_handle]`. 
+      return :code:`list[cudaq.measure_handle]`.
+
+   .. note::
+
+      As in C++, the two list arguments passed to the
+      :code:`stabilizer_round` kernel are the flattened X and Z stabilizer
+      schedule matrices, which default to the parity-check matrices. A Python
+      code can optionally define :code:`get_stabilizer_schedule_x` /
+      :code:`get_stabilizer_schedule_z` methods returning a 2D array with the
+      same shape and support pattern as the corresponding parity-check
+      matrix, where entry :code:`k >= 1` schedules that interaction at
+      timestep :code:`k`.
 
 3. **Implement the Code Class**:
 
@@ -482,6 +513,26 @@ print the layout. **Python:** :class:`cudaq_qec.stabilizer_grid` — see
 :ref:`qec_stabilizer_grid_cpp`. The header :file:`cudaq/qec/codes/surface_code.h`
 contains the full declaration.
 
+**Stabilizer measurement schedule**
+
+The surface code's :code:`stabilizer_round` kernel executes one depth-4
+extraction round: the X- and Z-check CNOTs are interleaved over four shared
+timesteps. Within each plaquette the CNOT order follows the standard zigzag
+schedule for the rotated surface code (`Tomita & Svore
+<https://arxiv.org/abs/1404.3747>`__): the X and Z plaquettes traverse their
+corners in transposed orders, selected per orientation so that mid-round
+ancilla faults ("hook errors", `Dennis et al.
+<https://arxiv.org/abs/quant-ph/0110143>`__) propagate onto data-qubit pairs
+perpendicular to the same-type logical operator. This preserves the full code
+distance :math:`d` under circuit-level noise; a naive schedule (both plaquette
+types in ascending qubit-index order) halves the effective distance of one
+memory basis. The schedule is available from the :code:`stabilizer_grid`
+helper via
+:cpp:func:`~cudaq::qec::surface_code::stabilizer_grid::get_cnot_schedule_x` /
+:code:`get_cnot_schedule_z` (matrix form) and
+:code:`get_cnot_schedule_pairs_x` / :code:`get_cnot_schedule_pairs_z` (flat
+pair-list form).
+
 Usage:
 
 .. tab:: Python

libs/qec/include/cudaq/qec/code.h

@@ -103,7 +103,10 @@ class code : public cudaqx::extension_point<code, const heterogeneous_map &> {
   /// @brief Type alias for two qubit quantum kernels
   using two_qubit_encoding = cudaq::qkernel<void(patch, patch)>;
 
-  /// @brief Type alias for stabilizer measurement kernels
+  /// @brief Type alias for stabilizer measurement kernels. The two vector
+  /// arguments are the flattened X and Z stabilizer schedule matrices (see
+  /// get_stabilizer_schedule_x/z()): entry 0 = no support, entry k >= 1 =
+  /// interaction at timestep k.
   using stabilizer_round = cudaq::qkernel<std::vector<cudaq::measure_result>(
       patch, const std::vector<std::size_t> &,
       const std::vector<std::size_t> &)>;
@@ -189,6 +192,26 @@ class code : public cudaqx::extension_point<code, const heterogeneous_map &> {
   /// @return Tensor representing Hz
   cudaqx::tensor<uint8_t> get_parity_z() const;
 
+  /// @brief Get the X-stabilizer schedule matrix passed to the code's
+  /// stabilizer_round kernel. Same shape and support pattern as
+  /// get_parity_x(); entry 0 means no support, entry k >= 1 means the
+  /// ancilla-data interaction executes at timestep k, so a code can encode an
+  /// optimized (e.g. hook-error-aware) gate order.
+  /// @return Tensor of scheduled interactions; the default implementation
+  /// returns get_parity_x() unchanged (every interaction at timestep 1, i.e.
+  /// ascending qubit-index order).
+  virtual cudaqx::tensor<uint8_t> get_stabilizer_schedule_x() const {
+    return get_parity_x();
+  }
+
+  /// @brief Get the Z-stabilizer schedule matrix passed to the code's
+  /// stabilizer_round kernel. See get_stabilizer_schedule_x().
+  /// @return Tensor of scheduled interactions; the default implementation
+  /// returns get_parity_z() unchanged.
+  virtual cudaqx::tensor<uint8_t> get_stabilizer_schedule_z() const {
+    return get_parity_z();
+  }
+
   /// @brief Get Lx stacked on Lz
   /// @return Tensor representing pauli observables
   cudaqx::tensor<uint8_t> get_pauli_observables_matrix() const;

libs/qec/include/cudaq/qec/codes/surface_code.h

@@ -196,6 +196,41 @@ class stabilizer_grid {
   /// @brief Get the stabilizers as a vector of cudaq::spin_op_terms
   std::vector<cudaq::spin_op_term> get_spin_op_stabilizers() const;
 
+  /// @brief Get the CNOT schedule matrix for the X stabilizers.
+  ///
+  /// @return Tensor with the same shape and support pattern as the X block of
+  /// the parity check matrix (num_x_stabilizers x distance^2). Entry 0 means
+  /// the ancilla does not touch that data qubit; entry k in [1, 4] means the
+  /// ancilla-data CNOT executes at timestep k of the stabilizer round.
+  ///
+  /// @note The per-plaquette CNOT order is chosen per the grid's orientation
+  /// so that mid-round ancilla faults ("hook errors"), which propagate onto
+  /// the data qubits of the remaining CNOTs, land perpendicular to the
+  /// same-type logical operator instead of along it, following the standard
+  /// zigzag schedule of https://arxiv.org/abs/1404.3747. Rows are ordered to
+  /// match the rows of to_parity_matrix()/code::get_parity_x().
+  cudaqx::tensor<uint8_t> get_cnot_schedule_x() const;
+
+  /// @brief Get the CNOT schedule matrix for the Z stabilizers.
+  ///
+  /// @return Tensor with the same shape and support pattern as the Z block of
+  /// the parity check matrix (num_z_stabilizers x distance^2). Entry 0 means
+  /// the ancilla does not touch that data qubit; entry k in [1, 4] means the
+  /// data-ancilla CNOT executes at timestep k of the stabilizer round.
+  ///
+  /// @note See get_cnot_schedule_x() for the hook-error rationale and row
+  /// ordering.
+  cudaqx::tensor<uint8_t> get_cnot_schedule_z() const;
+
+  /// @brief Get the X-stabilizer CNOT schedule as a flat list of
+  /// (stabilizer index, data index) pairs, ordered by timestep within each
+  /// stabilizer — the replay format for kernels that take an explicit CNOT
+  /// pair list. Stabilizer indices match the rows of get_cnot_schedule_x().
+  std::vector<std::size_t> get_cnot_schedule_pairs_x() const;
+
+  /// @brief Z-stabilizer counterpart of get_cnot_schedule_pairs_x().
+  std::vector<std::size_t> get_cnot_schedule_pairs_z() const;
+
   /// @brief Get the observables as a vector of cudaq::spin_op_terms
   ///
   /// @return The X logical observable first, followed by the Z logical
@@ -264,12 +299,15 @@ __qpu__ void prepm(patch p);
 ///
 /// @brief Perform stabilizer measurements on a surface_code patch
 /// @param p The patch to measure
-/// @param x_stabilizers Indices of X stabilizers to measure
-/// @param z_stabilizers Indices of Z stabilizers to measure
+/// @param x_stabilizer_schedule Flattened X-stabilizer CNOT schedule matrix
+/// (see stabilizer_grid::get_cnot_schedule_x): entry 0 = no support, entry
+/// k >= 1 = CNOT at timestep k
+/// @param z_stabilizer_schedule Flattened Z-stabilizer CNOT schedule matrix
+/// (see stabilizer_grid::get_cnot_schedule_z)
 /// @return Vector of measurement results
 __qpu__ std::vector<cudaq::measure_result>
-stabilizer(patch p, const std::vector<std::size_t> &x_stabilizers,
-           const std::vector<std::size_t> &z_stabilizers);
+stabilizer(patch p, const std::vector<std::size_t> &x_stabilizer_schedule,
+           const std::vector<std::size_t> &z_stabilizer_schedule);
 
 /// @brief surface_code implementation
 class surface_code : public cudaq::qec::code {
@@ -306,6 +344,14 @@ class surface_code : public cudaq::qec::code {
   surface_code(const heterogeneous_map &);
   // Grid constructor would be useful
 
+  /// @brief Get the hook-error-aware X-stabilizer CNOT schedule matrix.
+  /// See stabilizer_grid::get_cnot_schedule_x().
+  cudaqx::tensor<uint8_t> get_stabilizer_schedule_x() const override;
+
+  /// @brief Get the hook-error-aware Z-stabilizer CNOT schedule matrix.
+  /// See stabilizer_grid::get_cnot_schedule_z().
+  cudaqx::tensor<uint8_t> get_stabilizer_schedule_z() const override;
+
   /// @brief Extension creator function for the surface_code
   CUDAQ_EXTENSION_CUSTOM_CREATOR_FUNCTION(
       surface_code, static std::unique_ptr<cudaq::qec::code> create(

libs/qec/lib/codes/surface_code.cpp

@@ -386,6 +386,126 @@ stabilizer_grid::get_spin_op_stabilizers() const {
   return spin_op_stabs;
 }
 
+namespace {
+
+/// Corner offsets from a stabilizer grid coordinate to its four data qubits,
+/// listed in CNOT execution order (timesteps 1..4). Row-major ("Z"-shaped)
+/// visits NW, NE, SW, SE so the last two CNOTs touch a horizontally adjacent
+/// data pair; column-major ("N"-shaped) visits NW, SW, NE, SE so the last two
+/// touch a vertically adjacent pair. Assigning one shape to the X plaquettes
+/// and the transposed shape to the Z plaquettes is the standard zigzag
+/// schedule pair for the rotated surface code (Tomita & Svore,
+/// https://arxiv.org/abs/1404.3747; hook errors per Dennis et al.,
+/// https://arxiv.org/abs/quant-ph/0110143): it steers hook errors against
+/// the logical operators and stays conflict-free if the two timestep
+/// sequences are ever interleaved.
+constexpr int row_major_order[4][2] = {{-1, -1}, {-1, 0}, {0, -1}, {0, 0}};
+constexpr int col_major_order[4][2] = {{-1, -1}, {0, -1}, {-1, 0}, {0, 0}};
+
+cudaqx::tensor<uint8_t>
+build_cnot_schedule(const std::vector<vec2d> &stab_coords,
+                    const std::map<vec2d, size_t> &data_indices,
+                    std::size_t num_data, bool use_col_major_order) {
+  const auto &order = use_col_major_order ? col_major_order : row_major_order;
+  std::vector<std::vector<uint8_t>> rows;
+  rows.reserve(stab_coords.size());
+  for (const auto &sc : stab_coords) {
+    std::vector<uint8_t> row(num_data, 0);
+    for (int step = 0; step < 4; ++step) {
+      vec2d dq(sc.row + order[step][0], sc.col + order[step][1]);
+      auto it = data_indices.find(dq);
+      if (it != data_indices.end())
+        row[it->second] = step + 1;
+    }
+    rows.push_back(std::move(row));
+  }
+
+  // Match the row order of to_parity_matrix(): within one stabilizer type the
+  // sort comparator orders rows by the index of the first supported data
+  // qubit. That order is total only because no two same-type surface-code
+  // stabilizers share a first-support index — enforce it rather than assume
+  // it, since a violation would silently pair schedule rows with the wrong
+  // ancilla/parity rows.
+  std::vector<std::pair<std::size_t, std::size_t>>
+      keyed; // (first support, row)
+  keyed.reserve(rows.size());
+  for (std::size_t r = 0; r < rows.size(); ++r) {
+    auto it = std::find_if(rows[r].begin(), rows[r].end(),
+                           [](uint8_t v) { return v != 0; });
+    keyed.emplace_back(it - rows[r].begin(), r);
+  }
+  std::sort(keyed.begin(), keyed.end());
+  for (std::size_t r = 1; r < keyed.size(); ++r)
+    if (keyed[r].first == keyed[r - 1].first)
+      throw std::runtime_error(
+          "surface-code CNOT schedule: two same-type stabilizers share their "
+          "first supported data qubit, so the schedule rows cannot be aligned "
+          "with the parity-matrix rows.");
+
+  cudaqx::tensor<uint8_t> t({rows.size(), num_data});
+  for (std::size_t r = 0; r < rows.size(); ++r)
+    for (std::size_t c = 0; c < num_data; ++c)
+      t.at({r, c}) = rows[keyed[r].second][c];
+  return t;
+}
+
+/// Flatten a schedule matrix into (stabilizer, data) index pairs ordered by
+/// timestep within each stabilizer — the replay format used by kernels that
+/// take an explicit CNOT pair list.
+std::vector<std::size_t>
+schedule_to_pairs(const cudaqx::tensor<uint8_t> &sched) {
+  std::vector<std::size_t> pairs;
+  const auto num_stabs = sched.shape()[0];
+  const auto num_data = sched.shape()[1];
+  std::vector<std::pair<uint8_t, std::size_t>> row; // (step, data)
+  for (std::size_t s = 0; s < num_stabs; ++s) {
+    row.clear();
+    for (std::size_t d = 0; d < num_data; ++d)
+      if (sched.at({s, d}) != 0)
+        row.emplace_back(sched.at({s, d}), d);
+    std::sort(row.begin(), row.end());
+    for (const auto &[step, d] : row) {
+      pairs.push_back(s);
+      pairs.push_back(d);
+    }
+  }
+  return pairs;
+}
+
+} // namespace
+
+cudaqx::tensor<uint8_t> stabilizer_grid::get_cnot_schedule_x() const {
+  // A fault on the ancilla mid-round propagates onto the data qubits of the
+  // CNOTs that have not executed yet ("hook error"): an X⊗X pair on the last
+  // two data qubits in the schedule. If the X logical runs along the top row
+  // (horizontal), that pair must be vertical (column-major order) so hook
+  // chains run against the logical instead of along it; otherwise the pair
+  // must be horizontal (row-major order).
+  const bool x_logical_on_top_row =
+      orientation_ == sc_orientation::XV || orientation_ == sc_orientation::ZH;
+  return build_cnot_schedule(x_stab_coords, data_indices, distance * distance,
+                             /*use_col_major_order=*/x_logical_on_top_row);
+}
+
+cudaqx::tensor<uint8_t> stabilizer_grid::get_cnot_schedule_z() const {
+  // Same reasoning as get_cnot_schedule_x() for Z⊗Z hook pairs versus the Z
+  // logical, which runs perpendicular to the X logical: the Z plaquettes use
+  // the transposed order. The two orders also form a conflict-free pair if
+  // the X and Z timesteps are ever interleaved.
+  const bool x_logical_on_top_row =
+      orientation_ == sc_orientation::XV || orientation_ == sc_orientation::ZH;
+  return build_cnot_schedule(z_stab_coords, data_indices, distance * distance,
+                             /*use_col_major_order=*/!x_logical_on_top_row);
+}
+
+std::vector<std::size_t> stabilizer_grid::get_cnot_schedule_pairs_x() const {
+  return schedule_to_pairs(get_cnot_schedule_x());
+}
+
+std::vector<std::size_t> stabilizer_grid::get_cnot_schedule_pairs_z() const {
+  return schedule_to_pairs(get_cnot_schedule_z());
+}
+
 std::vector<cudaq::spin_op_term>
 stabilizer_grid::get_spin_op_observables() const {
   std::vector<cudaq::spin_op_term> spin_op_obs;
@@ -480,6 +600,14 @@ std::size_t surface_code::get_num_z_stabilizers() const {
   return (distance * distance - 1) / 2;
 }
 
+cudaqx::tensor<uint8_t> surface_code::get_stabilizer_schedule_x() const {
+  return grid.get_cnot_schedule_x();
+}
+
+cudaqx::tensor<uint8_t> surface_code::get_stabilizer_schedule_z() const {
+  return grid.get_cnot_schedule_z();
+}
+
 /// @brief Register the surace_code type
 CUDAQ_EXT_PT_REGISTER_TYPE(surface_code)
 

libs/qec/lib/codes/surface_code_device.cpp

@@ -49,26 +49,32 @@ __qpu__ void prepm(patch logicalQubit) {
 }
 
 __qpu__ std::vector<cudaq::measure_result>
-stabilizer(patch logicalQubit, const std::vector<std::size_t> &x_stabilizers,
-           const std::vector<std::size_t> &z_stabilizers) {
+stabilizer(patch logicalQubit,
+           const std::vector<std::size_t> &x_stabilizer_schedule,
+           const std::vector<std::size_t> &z_stabilizer_schedule) {
   for (std::size_t i = 0; i < logicalQubit.ancx.size(); i++)
     reset(logicalQubit.ancx[i]);
   for (std::size_t i = 0; i < logicalQubit.ancz.size(); i++)
     reset(logicalQubit.ancz[i]);
 
+  // The X and Z checks share the timesteps, and the surface-code schedule
+  // always fits in one depth-4 extraction round (see
+  // stabilizer_grid::get_cnot_schedule_x), so num_steps is a constant.
+  const std::size_t num_steps = 4;
+
   h(logicalQubit.ancx);
-  for (std::size_t xi = 0; xi < logicalQubit.ancx.size(); ++xi)
-    for (std::size_t di = 0; di < logicalQubit.data.size(); ++di)
-      if (x_stabilizers[xi * logicalQubit.data.size() + di] == 1)
-        cudaq::x<cudaq::ctrl>(logicalQubit.ancx[xi], logicalQubit.data[di]);
+  for (std::size_t step = 1; step <= num_steps; ++step) {
+    for (std::size_t xi = 0; xi < logicalQubit.ancx.size(); ++xi)
+      for (std::size_t di = 0; di < logicalQubit.data.size(); ++di)
+        if (x_stabilizer_schedule[xi * logicalQubit.data.size() + di] == step)
+          cudaq::x<cudaq::ctrl>(logicalQubit.ancx[xi], logicalQubit.data[di]);
+    for (std::size_t zi = 0; zi < logicalQubit.ancz.size(); ++zi)
+      for (std::size_t di = 0; di < logicalQubit.data.size(); ++di)
+        if (z_stabilizer_schedule[zi * logicalQubit.data.size() + di] == step)
+          cudaq::x<cudaq::ctrl>(logicalQubit.data[di], logicalQubit.ancz[zi]);
+  }
   h(logicalQubit.ancx);
 
-  // Now apply z_stabilizer circuit
-  for (size_t zi = 0; zi < logicalQubit.ancz.size(); ++zi)
-    for (size_t di = 0; di < logicalQubit.data.size(); ++di)
-      if (z_stabilizers[zi * logicalQubit.data.size() + di] == 1)
-        cudaq::x<cudaq::ctrl>(logicalQubit.data[di], logicalQubit.ancz[zi]);
-
   // S = (S_X, S_Z), (x flip syndromes, z flip syndromes).
   // x flips are triggered by z-stabilizers (ancz)
   // z flips are triggered by x-stabilizers (ancx)