Skip to content

Commit

Permalink
Add Surface Code (#7)
Browse files Browse the repository at this point in the history
This PR adds the `cudaq::qec::surface_code` namespace which includes the
`surface_code` type and the `stabilizer_grid` helper class. The
`stabilizer_grid` generates the 2d layout of the stabilizers, and which
data qubits each stabilizer has support on. Lastly it generates the
stabilizers and observables as vectors of `cudaq::spin_op` which the
`qec::code` expects.

---------

Co-authored-by: Ben Howe <141149032+bmhowe23@users.noreply.github.com>
justinlietz and bmhowe23 authored Jan 28, 2025
1 parent c303763 commit b8083bd
Showing 8 changed files with 947 additions and 13 deletions.
24 changes: 13 additions & 11 deletions docs/sphinx/api/core/cpp_api.rst
Original file line number Diff line number Diff line change
@@ -11,34 +11,36 @@ Namespaces
:desc-only:
.. doxygennamespace:: cudaq::qec::steane
:desc-only:
.. doxygennamespace:: cudaq::qec::surface_code
:desc-only:
.. doxygennamespace:: cudaq::qec::repetition
:desc-only:
.. doxygennamespace:: cudaq::solvers
:desc-only:
.. doxygennamespace:: cudaq::solvers::stateprep
.. doxygennamespace:: cudaq::solvers::stateprep
:desc-only:
.. doxygennamespace:: cudaq::solvers::adapt
:desc-only:
.. doxygennamespace:: cudaq::optim
:desc-only:

Core
Core
=============

.. doxygenclass:: cudaqx::extension_point
.. doxygenclass:: cudaqx::extension_point
:members:

.. doxygenclass:: cudaqx::heterogeneous_map
.. doxygenclass:: cudaqx::heterogeneous_map
:members:

.. doxygenclass:: cudaqx::tear_down
:members:

.. doxygenclass:: cudaqx::details::tensor_impl
:members:
.. doxygenclass:: cudaqx::details::tensor_impl
:members:

.. doxygenclass:: cudaqx::tensor
:members:
.. doxygenclass:: cudaqx::graph
:members:
.. doxygenclass:: cudaqx::tensor
:members:

.. doxygenclass:: cudaqx::graph
:members:
8 changes: 7 additions & 1 deletion docs/sphinx/api/qec/cpp_api.rst
Original file line number Diff line number Diff line change
@@ -15,14 +15,20 @@ CUDA-Q QEC C++ API
.. doxygenclass:: cudaq::qec::steane::steane
:members:

.. doxygenclass:: cudaq::qec::surface_code::stabilizer_grid
:members:

.. doxygenclass:: cudaq::qec::surface_code::surface_code
:members:

.. doxygenclass:: cudaq::qec::repetition::repetition
:members:

.. doxygenclass:: cudaq::qec::code
:members:

.. doxygenenum:: cudaq::qec::operation

.. doxygenfunction:: cudaq::qec::sample_code_capacity(const cudaqx::tensor<uint8_t> &, std::size_t, double)
.. doxygenfunction:: cudaq::qec::sample_code_capacity(const cudaqx::tensor<uint8_t> &, std::size_t, double, unsigned)
.. doxygenfunction:: cudaq::qec::sample_code_capacity(const code &, std::size_t, double)
269 changes: 269 additions & 0 deletions libs/qec/include/cudaq/qec/codes/surface_code.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,269 @@
/****************************************************************-*- C++ -*-****
* Copyright (c) 2024 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. *
******************************************************************************/
#pragma once

#include "cudaq/qec/code.h"
#include "cudaq/qec/patch.h"

using namespace cudaqx;

namespace cudaq::qec::surface_code {

/// @brief enumerates the role of a grid site in the surface codes stabilizer
/// grid
enum surface_role { amx, amz, empty };

/// @brief describes the 2d coordinate on the stabilizer grid
struct vec2d {
int row;
int col;

vec2d(int row_in, int col_in);
};

vec2d operator+(const vec2d &lhs, const vec2d &rhs);
vec2d operator-(const vec2d &lhs, const vec2d &rhs);
bool operator==(const vec2d &lhs, const vec2d &rhs);
bool operator<(const vec2d &lhs, const vec2d &rhs);

// clang-format off
/// @brief Generates and keeps track of the 2d grid of stabilizers in the
/// rotated surface code.
/// Following same layout convention as in: https://arxiv.org/abs/2311.10687
/// Grid layout is arranged from left to right, top to bottom (row major storage)
/// grid_length = 4 example:
/// ```
/// (0,0) (0,1) (0,2) (0,3)
/// (1,0) (1,1) (1,2) (1,3)
/// (2,0) (2,1) (2,2) (2,3)
/// (3,0) (3,1) (3,2) (3,3)
/// ```
// clang-format on

///
/// Each entry on the grid can be an X stabilizer, Z stabilizer,
/// or empty, as is needed on the edges.
/// The grid length of 4 corresponds to a distance 3 surface code, which results
/// in:
/// ```
/// e(0,0) e(0,1) Z(0,2) e(0,3)
/// X(1,0) Z(1,1) X(1,2) e(1,3)
/// e(2,0) X(2,1) Z(2,2) X(2,3)
/// e(3,0) Z(3,1) e(3,2) e(3,3)
/// ```
///
/// This is seen through the `print_stabilizer_grid()` member function.
/// To get rid of the empty sites, the `print_stabilizer_coords()` function is
/// used:
/// ```
/// Z(0,2)
/// X(1,0) Z(1,1) X(1,2)
/// X(2,1) Z(2,2) X(2,3)
/// Z(3,1)
/// ```
///
/// and to get the familiar visualization of the distance three surface code,
/// the `print_stabilizer_indices` results in:
/// ```
/// Z0
/// X0 Z1 X1
/// X2 Z2 X3
/// Z3
/// ```
///
/// The data qubits are located at the four corners of each of the weight-4
/// stabilizers. They are also organized with index increasing from left to
/// right, top to bottom:
/// ```
/// d0 d1 d2
/// d3 d4 d5
/// d6 d7 d8
/// ```
class stabilizer_grid {
private:
/// @brief Generates this->roles
void generate_grid_roles();
/// @brief Generates {x,z}_stab_coords and indices
void generate_grid_indices();
/// @brief Generates {x,z}_stabilizers
void generate_stabilizers();

public:
/// @brief The distance of the code
/// determines the number of data qubits per dimension
uint32_t distance = 0;

/// @brief length of the stabilizer grid
/// for distance = d data qubits,
/// the stabilizer grid has length d+1
uint32_t grid_length = 0;

/// @brief flattened vector of the stabilizer grid sites roles'
/// grid idx -> role
/// stored in row major order
std::vector<surface_role> roles;

/// @brief x stab index -> 2d coord
std::vector<vec2d> x_stab_coords;

/// @brief z stab index -> 2d coord
std::vector<vec2d> z_stab_coords;

/// @brief 2d coord -> z stab index
std::map<vec2d, size_t> x_stab_indices;

/// @brief 2d coord -> z stab index
std::map<vec2d, size_t> z_stab_indices;

/// @brief data index -> 2d coord
/// data qubits are in an offset 2D coord system from stabilizers
std::vector<vec2d> data_coords;

/// @brief 2d coord -> data index
std::map<vec2d, size_t> data_indices;

/// @brief Each element is an X stabilizer specified by the data qubits it has
/// support on
/// In surface code, can have weight 2 or weight 4 stabs
/// So {x,z}_stabilizer[i].size() == 2 || 4
std::vector<std::vector<size_t>> x_stabilizers;

/// @brief Each element is an Z stabilizer specified by the data qubits it has
/// support on
std::vector<std::vector<size_t>> z_stabilizers;

/// @brief Construct the grid from the code's distance
stabilizer_grid(uint32_t distance);
/// @brief Empty constructor
stabilizer_grid();

/// @brief Print a 2d grid of stabilizer roles
void print_stabilizer_grid() const;

/// @brief Print a 2d grid of stabilizer coords
void print_stabilizer_coords() const;

/// @brief Print a 2d grid of stabilizer indices
void print_stabilizer_indices() const;

/// @brief Print a 2d grid of data qubit indices
void print_data_grid() const;

/// @brief Print the coord <--> indices maps
void print_stabilizer_maps() const;

/// @brief Print the stabilizers in sparse pauli format
void print_stabilizers() const;

/// @brief Get the stabilizers as a vector of cudaq::spin_ops
std::vector<cudaq::spin_op> get_spin_op_stabilizers() const;

/// @brief Get the observables as a vector of cudaq::spin_ops
std::vector<cudaq::spin_op> get_spin_op_observables() const;
};

/// \pure_device_kernel
///
/// @brief Apply X gate to a surface_code patch
/// @param p The patch to apply the X gate to
__qpu__ void x(patch p);

/// \pure_device_kernel
///
/// @brief Apply Z gate to a surface_code patch
/// @param p The patch to apply the Z gate to
__qpu__ void z(patch p);

/// \pure_device_kernel
///
/// @brief Apply controlled-X gate between two surface_code patches
/// @param control The control patch
/// @param target The target patch
__qpu__ void cx(patch control, patch target);

/// \pure_device_kernel
///
/// @brief Apply controlled-Z gate between two surface_code patches
/// @param control The control patch
/// @param target The target patch
__qpu__ void cz(patch control, patch target);

/// \pure_device_kernel
///
/// @brief Prepare a surface_code patch in the |0⟩ state
/// @param p The patch to prepare
__qpu__ void prep0(patch p);

/// \pure_device_kernel
///
/// @brief Prepare a surface_code patch in the |1⟩ state
/// @param p The patch to prepare
__qpu__ void prep1(patch p);

/// \pure_device_kernel
///
/// @brief Prepare a surface_code patch in the |+⟩ state
/// @param p The patch to prepare
__qpu__ void prepp(patch p);

/// \pure_device_kernel
///
/// @brief Prepare a surface_code patch in the |-⟩ state
/// @param p The patch to prepare
__qpu__ void prepm(patch p);

/// \pure_device_kernel
///
/// @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
/// @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);

/// @brief surface_code implementation
class surface_code : public cudaq::qec::code {
protected:
/// @brief The code distance parameter
std::size_t distance;

/// @brief Get the number of data qubits in the surface_code
/// @return Number of data qubits (distance^2 for surface_code)
std::size_t get_num_data_qubits() const override;

/// @brief Get the number of total ancilla qubits in the surface_code
/// @return Number of data qubits (distance^2 - 1 for surface_code)
std::size_t get_num_ancilla_qubits() const override;

/// @brief Get the number of X ancilla qubits in the surface_code
/// @return Number of data qubits ((distance^2 - 1)/2 for surface_code)
std::size_t get_num_ancilla_x_qubits() const override;

/// @brief Get the number of Z ancilla qubits in the surface_code
/// @return Number of data qubits ((distance^2 - 1)/2 for surface_code)
std::size_t get_num_ancilla_z_qubits() const override;

public:
/// @brief Constructor for the surface_code
surface_code(const heterogeneous_map &);
// Grid constructor would be useful

/// @brief Extension creator function for the surface_code
CUDAQ_EXTENSION_CUSTOM_CREATOR_FUNCTION(
surface_code, static std::unique_ptr<cudaq::qec::code> create(
const cudaqx::heterogeneous_map &options) {
return std::make_unique<surface_code>(options);
})

/// @brief Grid to keep track of topological arrangement of qubits.
stabilizer_grid grid;
};

} // namespace cudaq::qec::surface_code
3 changes: 2 additions & 1 deletion libs/qec/lib/codes/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -9,7 +9,8 @@
cudaqx_add_device_code(cudaq-qec
SOURCES
steane_device.cpp
surface_code_device.cpp
repetition_device.cpp
)

target_sources(cudaq-qec PRIVATE steane.cpp repetition.cpp)
target_sources(cudaq-qec PRIVATE steane.cpp repetition.cpp surface_code.cpp)
392 changes: 392 additions & 0 deletions libs/qec/lib/codes/surface_code.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,392 @@
/*******************************************************************************
* Copyright (c) 2024 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 "cudaq/qec/codes/surface_code.h"
#include <iomanip>

using cudaq::qec::operation;

namespace cudaq::qec::surface_code {

vec2d::vec2d(int row_in, int col_in) : row(row_in), col(col_in) {}

vec2d operator+(const vec2d &lhs, const vec2d &rhs) {
return {lhs.row + rhs.row, lhs.col + rhs.col};
}

vec2d operator-(const vec2d &lhs, const vec2d &rhs) {
return {lhs.row - rhs.row, lhs.col - rhs.col};
}

bool operator==(const vec2d &lhs, const vec2d &rhs) {
return lhs.row == rhs.row && lhs.col == rhs.col;
}

// impose 2d ordering for reproducibility
bool operator<(const vec2d &lhs, const vec2d &rhs) {
// sort by row component.
// if row is tied, sort by col.
if (lhs.row != rhs.row) {
return lhs.row < rhs.row;
}
return lhs.col < rhs.col;
}

void stabilizer_grid::generate_grid_roles() {
// init grid to all empty
for (size_t row = 0; row < grid_length; ++row) {
for (size_t col = 0; col < grid_length; ++col) {
size_t idx = row * grid_length + col;
surface_role role = surface_role::empty;
roles[idx] = role;
}
}

// set alternating x/z interior
for (size_t row = 1; row < grid_length - 1; ++row) {
for (size_t col = 1; col < grid_length - 1; ++col) {
size_t idx = row * grid_length + col;
bool parity = (row + col) % 2;
if (!parity) {
roles[idx] = surface_role::amz;
} else {
roles[idx] = surface_role::amx;
}
}
}

// set top/bottom boundaries for weight 2 z-stabs
for (size_t row = 0; row < grid_length; row += grid_length - 1) {
for (size_t col = 1; col < grid_length - 1; ++col) {
size_t idx = row * grid_length + col;
bool parity = (row + col) % 2;
if (!parity)
roles[idx] = surface_role::amz;
}
}

// set left/right boundaries for weight 2 x-stabs
for (size_t row = 1; row < grid_length - 1; ++row) {
for (size_t col = 0; col < grid_length; col += grid_length - 1) {
size_t idx = row * grid_length + col;
bool parity = (row + col) % 2;
if (parity)
roles[idx] = surface_role::amx;
}
}
}

void stabilizer_grid::generate_grid_indices() {
size_t z_count = 0;
size_t x_count = 0;
for (size_t row = 0; row < grid_length; ++row) {
for (size_t col = 0; col < grid_length; ++col) {
size_t idx = row * grid_length + col;
switch (roles[idx]) {
case surface_role::amz:
z_stab_coords.push_back(vec2d(row, col));
z_stab_indices[vec2d(row, col)] = z_count;
z_count++;
break;
case surface_role::amx:
x_stab_coords.push_back(vec2d(row, col));
x_stab_indices[vec2d(row, col)] = x_count;
x_count++;
break;
case surface_role::empty:
// nothing
break;
default:
throw std::runtime_error(
"Grid index without role should be impossible\n");
}
}
}

// Data qubit grid
// This grid is a on a different coordinate system than the stabilizer grid
// Can think of this grid as offset from the stabilizer grid by a half unit
// right and down. data_row = stabilizer_row + 0.5 data_col = stabilizer_col +
// 0.5
for (size_t row = 0; row < distance; ++row) {
for (size_t col = 0; col < distance; ++col) {
size_t idx = row * distance + col;
data_coords.push_back(vec2d(row, col));
data_indices[vec2d(row, col)] = idx;
}
}
}

void stabilizer_grid::generate_stabilizers() {
for (size_t i = 0; i < x_stab_coords.size(); ++i) {
std::vector<size_t> current_stab;
for (int row_offset = -1; row_offset < 1; ++row_offset) {
for (int col_offset = -1; col_offset < 1; ++col_offset) {
int row = x_stab_coords[i].row + row_offset;
int col = x_stab_coords[i].col + col_offset;
vec2d trial_coord(row, col);
if (data_indices.find(trial_coord) != data_indices.end()) {
current_stab.push_back(data_indices[trial_coord]);
}
}
}
std::sort(current_stab.begin(), current_stab.end());
x_stabilizers.push_back(current_stab);
}

for (size_t i = 0; i < z_stab_coords.size(); ++i) {
std::vector<size_t> current_stab;
for (int row_offset = -1; row_offset < 1; ++row_offset) {
for (int col_offset = -1; col_offset < 1; ++col_offset) {
int row = z_stab_coords[i].row + row_offset;
int col = z_stab_coords[i].col + col_offset;
vec2d trial_coord(row, col);
if (data_indices.find(trial_coord) != data_indices.end()) {
current_stab.push_back(data_indices[trial_coord]);
}
}
}
std::sort(current_stab.begin(), current_stab.end());
z_stabilizers.push_back(current_stab);
}
}

stabilizer_grid::stabilizer_grid() {}

stabilizer_grid::stabilizer_grid(uint32_t distance)
: distance(distance), grid_length(distance + 1),
roles(grid_length * grid_length) {
// generate a 2d grid of roles
generate_grid_roles();
// now use grid to set coord
generate_grid_indices();
// now use coords to set the stabilizers
generate_stabilizers();
}

void stabilizer_grid::print_stabilizer_coords() const {
int width = std::to_string(grid_length).length();

for (size_t row = 0; row < grid_length; ++row) {
for (size_t col = 0; col < grid_length; ++col) {
size_t idx = row * grid_length + col;
switch (roles[idx]) {
case amz:
std::cout << "Z(" << std::setw(width) << row << "," << std::setw(width)
<< col << ") ";
break;
case amx:
std::cout << "X(" << std::setw(width) << row << "," << std::setw(width)
<< col << ") ";
break;
case empty:
std::cout << std::setw(2 * width + 6) << " ";
break;
default:
throw std::runtime_error(
"Grid index without role should be impossible\n");
}
}
std::cout << "\n";
}
std::cout << "\n";
}

void stabilizer_grid::print_stabilizer_indices() const {
int width = std::to_string(z_stab_indices.size()).length() + 2;
for (size_t row = 0; row < grid_length; ++row) {
for (size_t col = 0; col < grid_length; ++col) {
size_t idx = row * grid_length + col;
switch (roles[idx]) {
case amz:
std::cout << "Z" << std::left << std::setw(width)
<< z_stab_indices.at(vec2d(row, col));
break;
case amx:
std::cout << "X" << std::left << std::setw(width)
<< x_stab_indices.at(vec2d(row, col));
break;
case empty:
std::cout << std::setw(width + 1) << " ";
break;
default:
throw std::runtime_error(
"Grid index without role should be impossible\n");
}
}
std::cout << "\n";
}
std::cout << "\n";
}

void stabilizer_grid::print_stabilizer_maps() const {
std::cout << x_stab_coords.size() << " mx ancilla qubits:\n";
for (size_t i = 0; i < x_stab_coords.size(); ++i) {
std::cout << "amx[" << i << "] @ (" << x_stab_coords[i].row << ", "
<< x_stab_coords[i].col << ")\n";
}
std::cout << z_stab_coords.size() << " mz ancilla qubits:\n";
for (size_t i = 0; i < z_stab_coords.size(); ++i) {
std::cout << "amz[" << i << "] @ (" << z_stab_coords[i].row << ", "
<< z_stab_coords[i].col << ")\n";
}
std::cout << "\n";

std::cout << x_stab_coords.size() << " mx ancilla qubits:\n";
for (const auto &[k, v] : x_stab_indices) {
std::cout << "@(" << k.row << "," << k.col << "): amx[" << v << "]\n";
}
std::cout << z_stab_coords.size() << " mz ancilla qubits:\n";
for (const auto &[k, v] : z_stab_indices) {
std::cout << "@(" << k.row << "," << k.col << "): amz[" << v << "]\n";
}
std::cout << "\n";
}

void stabilizer_grid::print_data_grid() const {
int width = std::to_string(distance).length() + 2;

for (size_t row = 0; row < distance; ++row) {
for (size_t col = 0; col < distance; ++col) {
size_t idx = row * distance + col;
std::cout << "d" << std::left << std::setw(width) << idx;
}
std::cout << "\n";
}
std::cout << "\n";
}

void stabilizer_grid::print_stabilizer_grid() const {
int width = std::to_string(grid_length).length();

for (size_t row = 0; row < grid_length; ++row) {
for (size_t col = 0; col < grid_length; ++col) {
size_t idx = row * grid_length + col;
switch (roles[idx]) {
case amz:
std::cout << "Z(" << std::setw(width) << row << "," << std::setw(width)
<< col << ") ";
break;
case amx:
std::cout << "X(" << std::setw(width) << row << "," << std::setw(width)
<< col << ") ";
break;
case empty:
std::cout << "e(" << std::setw(width) << row << "," << std::setw(width)
<< col << ") ";
break;
default:
throw std::runtime_error(
"Grid index without role should be impossible\n");
}
}
std::cout << "\n";
}
std::cout << "\n";
}

void stabilizer_grid::print_stabilizers() const {
for (size_t s_i = 0; s_i < x_stabilizers.size(); ++s_i) {
std::cout << "s[" << s_i << "]: ";
for (size_t op_i = 0; op_i < x_stabilizers[s_i].size(); ++op_i) {
std::cout << "X" << x_stabilizers[s_i][op_i] << " ";
}
std::cout << "\n";
}

size_t offset = x_stabilizers.size();
for (size_t s_i = 0; s_i < z_stabilizers.size(); ++s_i) {
std::cout << "s[" << s_i + offset << "]: ";
for (size_t op_i = 0; op_i < z_stabilizers[s_i].size(); ++op_i) {
std::cout << "Z" << z_stabilizers[s_i][op_i] << " ";
}
std::cout << "\n";
}
std::cout << "\n";
}

std::vector<cudaq::spin_op> stabilizer_grid::get_spin_op_stabilizers() const {
std::vector<cudaq::spin_op> spin_op_stabs;
for (size_t s_i = 0; s_i < x_stabilizers.size(); ++s_i) {
std::string stab(data_coords.size(), 'I');
for (auto elem : x_stabilizers[s_i]) {
stab[elem] = 'X';
}
spin_op_stabs.emplace_back(cudaq::spin_op::from_word(stab));
}
for (size_t s_i = 0; s_i < z_stabilizers.size(); ++s_i) {
std::string stab(data_coords.size(), 'I');
for (auto elem : z_stabilizers[s_i]) {
stab[elem] = 'Z';
}
spin_op_stabs.emplace_back(cudaq::spin_op::from_word(stab));
}
return spin_op_stabs;
}

std::vector<cudaq::spin_op> stabilizer_grid::get_spin_op_observables() const {
std::vector<cudaq::spin_op> spin_op_obs;
// X obs runs along top row of data qubits
std::string xobs(data_coords.size(), 'I');
for (size_t i = 0; i < distance; ++i) {
xobs[i] = 'X';
}
spin_op_obs.emplace_back(cudaq::spin_op::from_word(xobs));

// Z ob runs along left col of data qubits
std::string zobs(data_coords.size(), 'I');
for (size_t i = 0; i < data_coords.size(); i += distance) {
zobs[i] = 'Z';
}
spin_op_obs.emplace_back(cudaq::spin_op::from_word(zobs));

return spin_op_obs;
}

surface_code::surface_code(const heterogeneous_map &options) : code() {
if (!options.contains("distance"))
throw std::runtime_error(
"[surface_code] distance not provided. distance must be provided via "
"qec::get_code(..., options) options map.");
distance = options.get<std::size_t>("distance");
grid = stabilizer_grid(distance);

m_stabilizers = grid.get_spin_op_stabilizers();
m_pauli_observables = grid.get_spin_op_observables();

operation_encodings.insert(std::make_pair(operation::x, x));
operation_encodings.insert(std::make_pair(operation::z, z));
operation_encodings.insert(std::make_pair(operation::cx, cx));
operation_encodings.insert(std::make_pair(operation::cz, cz));
operation_encodings.insert(
std::make_pair(operation::stabilizer_round, stabilizer));
operation_encodings.insert(std::make_pair(operation::prep0, prep0));
operation_encodings.insert(std::make_pair(operation::prep1, prep1));
operation_encodings.insert(std::make_pair(operation::prepp, prepp));
operation_encodings.insert(std::make_pair(operation::prepm, prepm));
}

std::size_t surface_code::get_num_data_qubits() const {
return distance * distance;
}

std::size_t surface_code::get_num_ancilla_qubits() const {
return distance * distance - 1;
}

std::size_t surface_code::get_num_ancilla_x_qubits() const {
return (distance * distance - 1) / 2;
}

std::size_t surface_code::get_num_ancilla_z_qubits() const {
return (distance * distance - 1) / 2;
}

/// @brief Register the surace_code type
CUDAQ_REGISTER_TYPE(surface_code)

} // namespace cudaq::qec::surface_code
80 changes: 80 additions & 0 deletions libs/qec/lib/codes/surface_code_device.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
/*******************************************************************************
* Copyright (c) 2024 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 "cudaq.h"
#include "cudaq/qec/patch.h"

namespace cudaq::qec::surface_code {

__qpu__ void x(patch logicalQubit) { x(logicalQubit.data); }
__qpu__ void z(patch logicalQubit) { z(logicalQubit.data); }

__qpu__ void cx(patch logicalQubitA, patch logicalQubitB) {
for (std::size_t i = 0; i < logicalQubitA.data.size(); i++) {
x<cudaq::ctrl>(logicalQubitA.data[i], logicalQubitB.data[i]);
}
}

__qpu__ void cz(patch logicalQubitA, patch logicalQubitB) {
for (std::size_t i = 0; i < logicalQubitA.data.size(); i++) {
z<cudaq::ctrl>(logicalQubitA.data[i], logicalQubitB.data[i]);
}
}

// Transversal state prep, turn on stabilizers after these ops
__qpu__ void prep0(patch logicalQubit) {
for (std::size_t i = 0; i < logicalQubit.data.size(); i++)
reset(logicalQubit.data[i]);
}

__qpu__ void prep1(patch logicalQubit) {
prep0(logicalQubit);
x(logicalQubit.data);
}

__qpu__ void prepp(patch logicalQubit) {
prep0(logicalQubit);
h(logicalQubit.data);
}

__qpu__ void prepm(patch logicalQubit) {
prep0(logicalQubit);
x(logicalQubit.data);
h(logicalQubit.data);
}

__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) {
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]);
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 syndrones).
// x flips are triggered by z-stabilizers (ancz)
// z flips are triggered by x-stabilizers (ancx)
auto results = mz(logicalQubit.ancz, logicalQubit.ancx);

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]);

return results;
}

} // namespace cudaq::qec::surface_code
1 change: 1 addition & 0 deletions libs/qec/unittests/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -29,6 +29,7 @@ endif()
include(GoogleTest)

# ==============================================================================
add_compile_options(-Wno-attributes)

add_executable(test_decoders test_decoders.cpp decoders/sample_decoder.cpp)
target_link_libraries(test_decoders PRIVATE GTest::gtest_main cudaq-qec)
183 changes: 183 additions & 0 deletions libs/qec/unittests/test_qec.cpp
Original file line number Diff line number Diff line change
@@ -11,6 +11,7 @@

#include "cudaq.h"

#include "cudaq/qec/codes/surface_code.h"
#include "cudaq/qec/experiments.h"

TEST(StabilizerTester, checkConstructFromSpinOps) {
@@ -502,6 +503,99 @@ TEST(QECCodeTester, checkRepetition) {
}
}

TEST(QECCodeTester, checkSurfaceCode) {
{
// must provide distance
EXPECT_THROW(cudaq::qec::get_code("surface_code"), std::runtime_error);
}
{
// with default stabilizers
auto surf_code = cudaq::qec::get_code("surface_code", {{"distance", 3}});
cudaqx::tensor<uint8_t> parity = surf_code->get_parity();
cudaqx::tensor<uint8_t> parity_x = surf_code->get_parity_x();
cudaqx::tensor<uint8_t> parity_z = surf_code->get_parity_z();
cudaqx::tensor<uint8_t> observables =
surf_code->get_pauli_observables_matrix();
cudaqx::tensor<uint8_t> Lx = surf_code->get_observables_x();
cudaqx::tensor<uint8_t> Lz = surf_code->get_observables_z();

{
// This is just a regression check, this has not been hand tested
std::vector<uint8_t> data = {
1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, /* row 0 */
0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, /* row 1 */
0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, /* row 2 */
0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, /* row 3 */
0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, /* row 4 */
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, /* row 5 */
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, /* row 6 */
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1}; /* row 7 */

std::vector<std::size_t> expected_shape{8, 18};
cudaqx::tensor<uint8_t> t(expected_shape);
t.borrow(data.data());
for (std::size_t i = 0; i < 8; i++)
for (std::size_t j = 0; j < 18; j++)
EXPECT_EQ(t.at({i, j}), parity.at({i, j}));
}
{
// This is just a regression check, this has not been hand tested
std::vector<uint8_t> data = {1, 0, 0, 1, 0, 0, 0, 0, 0, /* row 0 */
0, 1, 1, 0, 1, 1, 0, 0, 0, /* row 1 */
0, 0, 0, 1, 1, 0, 1, 1, 0, /* row 2 */
0, 0, 0, 0, 0, 1, 0, 0, 1}; /* row 3 */

std::vector<std::size_t> expected_shape{4, 9};
cudaqx::tensor<uint8_t> t(expected_shape);
t.borrow(data.data());
for (std::size_t i = 0; i < 4; i++)
for (std::size_t j = 0; j < 9; j++)
EXPECT_EQ(t.at({i, j}), parity_x.at({i, j}));
}
{
// This is just a regression check, this has not been hand tested
std::vector<uint8_t> data = {1, 1, 0, 1, 1, 0, 0, 0, 0, /* row 0 */
0, 1, 1, 0, 0, 0, 0, 0, 0, /* row 1 */
0, 0, 0, 0, 1, 1, 0, 1, 1, /* row 2 */
0, 0, 0, 0, 0, 0, 1, 1, 0}; /* row 3 */

std::vector<std::size_t> expected_shape{4, 9};
cudaqx::tensor<uint8_t> t(expected_shape);
t.borrow(data.data());
for (std::size_t i = 0; i < 4; i++)
for (std::size_t j = 0; j < 9; j++)
EXPECT_EQ(t.at({i, j}), parity_z.at({i, j}));
}
EXPECT_EQ(2, observables.rank());
EXPECT_EQ(2, observables.shape()[0]);
EXPECT_EQ(18, observables.shape()[1]);
EXPECT_EQ(2, Lx.rank());
EXPECT_EQ(1, Lx.shape()[0]);
EXPECT_EQ(9, Lx.shape()[1]);
EXPECT_EQ(2, Lz.rank());
EXPECT_EQ(1, Lz.shape()[0]);
EXPECT_EQ(9, Lz.shape()[1]);
{
std::vector<std::vector<uint8_t>> true_observables = {
{1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, // Z first
{0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0}}; // Then X
std::vector<std::vector<uint8_t>> true_Lx = {{1, 1, 1, 0, 0, 0, 0, 0, 0}};
std::vector<std::vector<uint8_t>> true_Lz = {{1, 0, 0, 1, 0, 0, 1, 0, 0}};
for (std::size_t i = 0; i < observables.shape()[0]; ++i)
for (std::size_t j = 0; j < observables.shape()[1]; ++j)
EXPECT_EQ(true_observables[i][j], observables.at({i, j}));

for (std::size_t i = 0; i < Lx.shape()[0]; ++i)
for (std::size_t j = 0; j < Lx.shape()[1]; ++j)
EXPECT_EQ(true_Lx[i][j], Lx.at({i, j}));

for (std::size_t i = 0; i < Lz.shape()[0]; ++i)
for (std::size_t j = 0; j < Lz.shape()[1]; ++j)
EXPECT_EQ(true_Lz[i][j], Lz.at({i, j}));
}
}
}

// expect |0>, |+> to measure out 0 in respective bases
// expect |1>, |-> to measure out 1 in respective bases
bool noiseless_logical_SPAM_test(const cudaq::qec::code &code,
@@ -582,3 +676,92 @@ TEST(QECCodeTester, checkRepetitionSPAM) {
EXPECT_FALSE(noiseless_logical_SPAM_test(*repetition,
cudaq::qec::operation::prep1, 0));
}

TEST(QECCodeTester, checkSurfaceCodeSPAM) {
// Must compile with stim for larger distances
auto surf_code = cudaq::qec::get_code("surface_code", {{"distance", 3}});
EXPECT_TRUE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prep0, 0));
EXPECT_TRUE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prep1, 1));
EXPECT_TRUE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prepp, 0));
EXPECT_TRUE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prepm, 1));
EXPECT_FALSE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prep0, 1));
EXPECT_FALSE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prep1, 0));
EXPECT_FALSE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prepp, 1));
EXPECT_FALSE(
noiseless_logical_SPAM_test(*surf_code, cudaq::qec::operation::prepm, 0));
}

TEST(QECCodeTester, checkStabilizerGrid) {
{
int distance = 3;

cudaq::qec::surface_code::stabilizer_grid grid(distance);

grid.print_stabilizer_grid();
grid.print_stabilizer_coords();
grid.print_stabilizer_indices();
grid.print_stabilizer_maps();
grid.print_data_grid();
grid.print_stabilizers();

EXPECT_EQ(3, grid.distance);
EXPECT_EQ(4, grid.grid_length);
EXPECT_EQ(16, grid.roles.size());
EXPECT_EQ(4, grid.x_stab_coords.size());
EXPECT_EQ(4, grid.z_stab_coords.size());
EXPECT_EQ(4, grid.x_stab_indices.size());
EXPECT_EQ(4, grid.z_stab_indices.size());
EXPECT_EQ(9, grid.data_coords.size());
EXPECT_EQ(9, grid.data_indices.size());
EXPECT_EQ(4, grid.x_stabilizers.size());
EXPECT_EQ(4, grid.z_stabilizers.size());
}
{
int distance = 5;

cudaq::qec::surface_code::stabilizer_grid grid(distance);

grid.print_stabilizer_grid();
grid.print_stabilizer_coords();
grid.print_stabilizer_indices();
grid.print_stabilizer_maps();
grid.print_data_grid();
grid.print_stabilizers();

EXPECT_EQ(5, grid.distance);
EXPECT_EQ(6, grid.grid_length);
EXPECT_EQ(36, grid.roles.size());
EXPECT_EQ(12, grid.x_stab_coords.size());
EXPECT_EQ(12, grid.z_stab_coords.size());
EXPECT_EQ(12, grid.x_stab_indices.size());
EXPECT_EQ(12, grid.z_stab_indices.size());
EXPECT_EQ(25, grid.data_coords.size());
EXPECT_EQ(25, grid.data_indices.size());
EXPECT_EQ(12, grid.x_stabilizers.size());
EXPECT_EQ(12, grid.z_stabilizers.size());
}
{
int distance = 17;

cudaq::qec::surface_code::stabilizer_grid grid(distance);

EXPECT_EQ(17, grid.distance);
EXPECT_EQ(18, grid.grid_length);
EXPECT_EQ(324, grid.roles.size());
EXPECT_EQ(144, grid.x_stab_coords.size());
EXPECT_EQ(144, grid.z_stab_coords.size());
EXPECT_EQ(144, grid.x_stab_indices.size());
EXPECT_EQ(144, grid.z_stab_indices.size());
EXPECT_EQ(289, grid.data_coords.size());
EXPECT_EQ(289, grid.data_indices.size());
EXPECT_EQ(144, grid.x_stabilizers.size());
EXPECT_EQ(144, grid.z_stabilizers.size());
}
}

0 comments on commit b8083bd

Please sign in to comment.