diff --git a/AGENTS.md b/AGENTS.md index 6990ed4..cfbe9bc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,30 +3,33 @@ - Use the Bazel build system when interacting with this repository. - The CMake build system is available to help users who need it. - Keep both the CMake and Bazel builds working at all times. +- Use at most one CPU core for all builds, tests, and other expensive commands. + For Bazel, pass `--jobs=1`. For CMake builds, pass `--parallel 1`. For + `make`, pass `-j1`. ## Building with Bazel To build all code with bazel: ```bash -bazel build src:all +bazel build --jobs=1 src:all ``` To build the Tesseract and Simplex main binaries: ```bash -bazel build src:tesseract src:simplex +bazel build --jobs=1 src:tesseract src:simplex ``` ## Running Tests with Bazel ```bash -bazel test src/... +bazel test --jobs=1 src/... ``` ## Building with CMake -In case you need to, when building with CMake, use parallel flags to speed up the process. +In case you need to build with CMake, keep the build single-core. -- When using `cmake --build`, add the `--parallel` flag. -- When using `make`, add the `-j` flag (e.g., `make -j$(nproc)`). +- When using `cmake --build`, add `--parallel 1`. +- When using `make`, add `-j1`. ## Running Tests with CMake @@ -36,7 +39,7 @@ To run the tests, execute the following commands from the root of the repository mkdir -p build cd build cmake .. -cmake --build . --parallel +cmake --build . --parallel 1 ctest ``` @@ -49,7 +52,7 @@ To build the Python wheel for `tesseract_decoder` locally, you will need to use Use the following command: ```bash -bazel build //:tesseract_decoder_wheel --define=VERSION=0.1.1 --define=TARGET_VERSION=py313 +bazel build --jobs=1 //:tesseract_decoder_wheel --define=VERSION=0.1.1 --define=TARGET_VERSION=py313 ``` - `--define=VERSION=0.1.1`: Sets the version of the wheel. You should replace `0.1.1` with the current version from the `_version.py` file. @@ -62,7 +65,7 @@ The resulting wheel file will be located in the `bazel-bin/` directory. If you change the Python version in `MODULE.bazel`, you will need to regenerate the `requirements_lock.txt` file to ensure all dependencies are compatible. To do this, run the following command: ```bash -bazel run //src/py:requirements.update +bazel run --jobs=1 //src/py:requirements.update ``` This will update the `src/py/requirements_lock.txt` file with the correct dependency versions for the new Python environment. diff --git a/CMakeLists.txt b/CMakeLists.txt index b7a4270..23e560f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -98,6 +98,14 @@ target_include_directories(tesseract_lib PUBLIC ${TESSERACT_SRC_DIR}) target_compile_options(tesseract_lib PRIVATE ${OPT_COPTS}) target_link_libraries(tesseract_lib PUBLIC utils boost_headers visualization) +add_library(tesseract_trellis_lib + ${TESSERACT_SRC_DIR}/tesseract_trellis.cc + ${TESSERACT_SRC_DIR}/tesseract_trellis.h +) +target_include_directories(tesseract_trellis_lib PUBLIC ${TESSERACT_SRC_DIR}) +target_compile_options(tesseract_trellis_lib PRIVATE ${OPT_COPTS}) +target_link_libraries(tesseract_trellis_lib PUBLIC common utils libstim Threads::Threads) + add_library(simplex ${TESSERACT_SRC_DIR}/simplex.cc ${TESSERACT_SRC_DIR}/simplex.h) target_include_directories(simplex PUBLIC ${TESSERACT_SRC_DIR}) target_compile_options(simplex PRIVATE ${OPT_COPTS}) @@ -108,6 +116,12 @@ add_executable(tesseract ${TESSERACT_SRC_DIR}/tesseract_main.cc) target_compile_options(tesseract PRIVATE ${OPT_COPTS}) target_link_libraries(tesseract PRIVATE tesseract_lib argparse::argparse nlohmann_json::nlohmann_json) +add_executable(tesseract_trellis ${TESSERACT_SRC_DIR}/tesseract_trellis_main.cc) +target_compile_options(tesseract_trellis PRIVATE ${OPT_COPTS}) +target_link_libraries(tesseract_trellis + PRIVATE tesseract_trellis_lib argparse::argparse nlohmann_json::nlohmann_json +) + add_executable(simplex_bin ${TESSERACT_SRC_DIR}/simplex_main.cc) set_target_properties(simplex_bin PROPERTIES OUTPUT_NAME simplex) target_compile_options(simplex_bin PRIVATE ${OPT_COPTS}) @@ -137,3 +151,6 @@ add_executable(tesseract_test ${TESSERACT_SRC_DIR}/tesseract.test.cc) target_link_libraries(tesseract_test PRIVATE tesseract_lib simplex GTest::gtest_main) add_test(NAME tesseract_test COMMAND tesseract_test) +add_executable(tesseract_trellis_test ${TESSERACT_SRC_DIR}/tesseract_trellis.test.cc) +target_link_libraries(tesseract_trellis_test PRIVATE tesseract_trellis_lib GTest::gtest_main) +add_test(NAME tesseract_trellis_test COMMAND tesseract_trellis_test) diff --git a/src/BUILD b/src/BUILD index ebac6a5..9aa162b 100644 --- a/src/BUILD +++ b/src/BUILD @@ -140,6 +140,19 @@ cc_library( ], ) +cc_library( + name = "libtesseract_trellis", + srcs = ["tesseract_trellis.cc"], + hdrs = ["tesseract_trellis.h"], + copts = OPT_COPTS, + linkopts = OPT_LINKOPTS, + deps = [ + ":libcommon", + ":libutils", + "@stim//:stim_lib", + ], +) + cc_binary( name = "tesseract", srcs = ["tesseract_main.cc"], @@ -153,6 +166,19 @@ cc_binary( ], ) +cc_binary( + name = "tesseract_trellis", + srcs = ["tesseract_trellis_main.cc"], + copts = OPT_COPTS, + linkopts = OPT_LINKOPTS, + deps = [ + ":libtesseract_trellis", + "@argparse", + "@nlohmann_json//:json", + "@stim//:stim_lib", + ], +) + cc_test( name = "tesseract_tests", timeout = "eternal", @@ -169,6 +195,19 @@ cc_test( ], ) +cc_test( + name = "tesseract_trellis_tests", + srcs = ["tesseract_trellis.test.cc"], + copts = OPT_COPTS, + linkopts = OPT_LINKOPTS, + deps = [ + ":libtesseract_trellis", + "@gtest", + "@gtest//:gtest_main", + "@stim//:stim_lib", + ], +) + cc_test( name = "common_tests", srcs = ["common.test.cc"], diff --git a/src/simplex_main.cc b/src/simplex_main.cc index 7939a91..95437ae 100644 --- a/src/simplex_main.cc +++ b/src/simplex_main.cc @@ -251,6 +251,7 @@ struct Args { // Create a writer instance to write the predicted obs to a file stim::FileFormatData predictions_out_format = stim::format_name_to_enum_map().at(out_format); FILE* predictions_file = stdout; + // An output path of "-" means stdout. if (out_fname != "-") { predictions_file = fopen(out_fname.c_str(), "w"); } diff --git a/src/tesseract_main.cc b/src/tesseract_main.cc index 2afbce4..c54f82b 100644 --- a/src/tesseract_main.cc +++ b/src/tesseract_main.cc @@ -317,6 +317,7 @@ struct Args { // Create a writer instance to write the predicted obs to a file stim::FileFormatData predictions_out_format = stim::format_name_to_enum_map().at(out_format); FILE* predictions_file = stdout; + // An output path of "-" means stdout. if (out_fname != "-") { predictions_file = fopen(out_fname.c_str(), "w"); } diff --git a/src/tesseract_trellis.cc b/src/tesseract_trellis.cc new file mode 100644 index 0000000..b48a3fa --- /dev/null +++ b/src/tesseract_trellis.cc @@ -0,0 +1,1288 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tesseract_trellis.h" + +#include +#include +#include +#include +#include +#include +#if defined(__BMI2__) && \ + (defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || defined(_M_IX86)) +#include +#endif +#include +#include +#include +#include +#include +#include + +#include "utils.h" + +struct TesseractTrellisWideKernelBase { + virtual ~TesseractTrellisWideKernelBase() = default; + virtual void decode_shot(TesseractTrellisDecoder* decoder, + const std::vector& detections) const = 0; +}; + +namespace { + +constexpr size_t kMaxCompiledWideStateWords = 4; + +#if defined(__GNUC__) || defined(__clang__) +#define TESSERACT_ALWAYS_INLINE inline __attribute__((always_inline)) +#define TESSERACT_HOT __attribute__((hot)) +#else +#define TESSERACT_ALWAYS_INLINE inline +#define TESSERACT_HOT +#endif + +struct Fault { + size_t error_index; + double likelihood_cost; + double log_q; + double log_p; + uint64_t obs_mask; + std::vector detectors; +}; + +template +using FixedWideStateWords = std::array; + +template +struct FixedWideStateEntry { + FixedWideStateWords state_words{}; + double mass0 = 0.0; + double mass1 = 0.0; + double penalty = 0.0; + double score = -INF; +}; + +template +struct FixedWidePairBucket { + FixedWideStateWords key{}; + double mass0[2]{}; + double mass1[2]{}; + double penalty[2]{}; + uint8_t used_mask = 0; + bool occupied = false; +}; + +template +struct CompiledWideLayerTemplate { + double q = 0.0; + double p = 0.0; + bool toggles_observable = false; + bool has_retiring_terms = false; + size_t surviving_term_count = 0; + std::array surviving_masks{}; + std::array projection_dst_words{}; + std::array projection_dst_offsets{}; + std::array projected_fault_mask_words{}; + std::vector fault_target_word_indices; + std::vector fault_target_bit_masks; + std::vector fault_word_indices; + std::vector fault_bit_masks; + std::vector fault_was_active_before; + std::vector current_costs; + std::vector next_costs; +}; + +struct BranchPenaltyUpdate { + bool absent_valid = true; + bool present_valid = true; + double absent_penalty = 0.0; + double present_penalty = 0.0; +}; + +struct FinalizeKeptStateStatsOnExit { + TesseractTrellisDecoder* decoder; + + ~FinalizeKeptStateStatsOnExit(); +}; + +std::vector parse_faults(const std::vector& errors, size_t num_observables) { + std::vector faults; + faults.reserve(errors.size()); + for (size_t error_index = 0; error_index < errors.size(); ++error_index) { + const auto& error = errors[error_index]; + const double p = error.get_probability(); + if (p <= 0) { + continue; + } + Fault fault; + fault.error_index = error_index; + fault.likelihood_cost = error.likelihood_cost; + fault.log_q = std::log1p(-p); + fault.log_p = std::log(p); + fault.obs_mask = 0; + for (int obs : error.symptom.observables) { + if (obs >= 64) { + throw std::invalid_argument("tesseract_trellis currently supports at most 64 observables"); + } + if (size_t(obs) >= num_observables) { + throw std::invalid_argument("Observable index out of range in DEM"); + } + fault.obs_mask ^= uint64_t{1} << obs; + } + fault.detectors = error.symptom.detectors; + faults.push_back(std::move(fault)); + } + return faults; +} + +void build_wide_layer_templates(const std::vector& faults, size_t num_detectors, + std::vector* layers, + size_t* max_frontier_width_seen) { + std::vector last_seen(num_detectors, std::numeric_limits::max()); + for (size_t i = 0; i < faults.size(); ++i) { + for (int d : faults[i].detectors) { + last_seen[d] = i; + } + } + + std::vector active_detectors; + active_detectors.reserve(num_detectors); + std::vector global_to_local(num_detectors, -1); + layers->clear(); + layers->reserve(faults.size()); + *max_frontier_width_seen = 0; + + for (size_t i = 0; i < faults.size(); ++i) { + const size_t previous_width = active_detectors.size(); + for (int d : faults[i].detectors) { + if (global_to_local[d] == -1) { + global_to_local[d] = active_detectors.size(); + active_detectors.push_back(d); + } + } + + *max_frontier_width_seen = std::max(*max_frontier_width_seen, active_detectors.size()); + TesseractTrellisWideLayerTemplate layer{ + .q = std::exp(faults[i].log_q), + .p = std::exp(faults[i].log_p), + .obs_mask = faults[i].obs_mask, + .previous_width = previous_width, + .surviving_local_indices = {}, + .current_active_detectors = active_detectors, + .projected_fault_mask_words = {}, + .next_frontier_costs = {}, + .detcost_transition = {}, + }; + + for (size_t local = 0; local < active_detectors.size(); ++local) { + const int d = active_detectors[local]; + if (last_seen[d] != i) { + layer.surviving_local_indices.push_back((uint32_t)local); + } + } + + std::vector next_active; + next_active.reserve(layer.surviving_local_indices.size()); + std::fill(global_to_local.begin(), global_to_local.end(), -1); + for (size_t next_local = 0; next_local < layer.surviving_local_indices.size(); ++next_local) { + int d = active_detectors[layer.surviving_local_indices[next_local]]; + global_to_local[d] = next_local; + next_active.push_back(d); + } + active_detectors = std::move(next_active); + layers->push_back(std::move(layer)); + } +} + +template +void build_future_detcost_transitions(const std::vector& faults, size_t num_detectors, + std::vector* layers, + std::vector* initial_future_detcost) { + std::vector current_row(num_detectors, INF); + for (size_t fault_index = faults.size(); fault_index-- > 0;) { + auto& layer = (*layers)[fault_index]; + const auto& fault = faults[fault_index]; + + layer.next_frontier_costs.resize(layer.surviving_local_indices.size(), INF); + for (size_t next_local = 0; next_local < layer.surviving_local_indices.size(); ++next_local) { + size_t current_local = (size_t)layer.surviving_local_indices[next_local]; + int global_detector = layer.current_active_detectors[current_local]; + layer.next_frontier_costs[next_local] = current_row[(size_t)global_detector]; + } + + std::vector current_to_next(layer.current_active_detectors.size(), -1); + for (size_t next_local = 0; next_local < layer.surviving_local_indices.size(); ++next_local) { + current_to_next[(size_t)layer.surviving_local_indices[next_local]] = (int32_t)next_local; + } + + auto& transition = layer.detcost_transition; + transition.fault_local_indices.clear(); + transition.next_local_indices.clear(); + transition.current_costs.clear(); + transition.next_costs.clear(); + transition.fault_local_indices.reserve(fault.detectors.size()); + transition.next_local_indices.reserve(fault.detectors.size()); + transition.current_costs.reserve(fault.detectors.size()); + transition.next_costs.reserve(fault.detectors.size()); + + if (!fault.detectors.empty()) { + double ecost = fault.likelihood_cost / fault.detectors.size(); + for (int detector : fault.detectors) { + auto it = std::find(layer.current_active_detectors.begin(), + layer.current_active_detectors.end(), detector); + if (it == layer.current_active_detectors.end()) { + throw std::runtime_error("Missing detector in active frontier while preparing detcost."); + } + uint32_t local = (uint32_t)std::distance(layer.current_active_detectors.begin(), it); + double next_cost = current_row[(size_t)detector]; + double current_cost = std::min(ecost, next_cost); + transition.fault_local_indices.push_back(local); + transition.next_local_indices.push_back(current_to_next[local]); + transition.current_costs.push_back(current_cost); + transition.next_costs.push_back(next_cost); + current_row[(size_t)detector] = current_cost; + } + } + } + + if (initial_future_detcost != nullptr) { + *initial_future_detcost = std::move(current_row); + } +} + +struct ActiveDetcostEvent { + int detector = -1; + size_t valid_low = 0; + double cost = INF; +}; + +struct ActiveDetcostHeapEntry { + size_t valid_low = 0; + double cost = INF; + + bool operator>(const ActiveDetcostHeapEntry& other) const { + return cost > other.cost || (cost == other.cost && valid_low > other.valid_low); + } +}; + +template +void build_future_active_detcost_transitions(const std::vector& faults, size_t num_detectors, + std::vector* layers, + std::vector* initial_future_detcost) { + std::vector first_seen(num_detectors, std::numeric_limits::max()); + for (size_t fault_index = 0; fault_index < faults.size(); ++fault_index) { + for (int detector : faults[fault_index].detectors) { + size_t& seen = first_seen[(size_t)detector]; + if (seen == std::numeric_limits::max()) { + seen = fault_index; + } + } + } + + std::vector> events_by_high(faults.size()); + std::vector> first_seen_and_detector; + std::vector active_detectors; + for (size_t fault_index = 0; fault_index < faults.size(); ++fault_index) { + const auto& fault = faults[fault_index]; + if (fault.detectors.empty()) { + continue; + } + + first_seen_and_detector.clear(); + first_seen_and_detector.reserve(fault.detectors.size()); + for (int detector : fault.detectors) { + size_t seen = first_seen[(size_t)detector]; + if (seen == std::numeric_limits::max() || seen > fault_index) { + throw std::runtime_error("Invalid first-seen detector state while preparing detcost."); + } + first_seen_and_detector.push_back({seen, detector}); + } + std::sort(first_seen_and_detector.begin(), first_seen_and_detector.end()); + first_seen_and_detector.erase( + std::unique(first_seen_and_detector.begin(), first_seen_and_detector.end(), + [](const auto& a, const auto& b) { return a.second == b.second; }), + first_seen_and_detector.end()); + + active_detectors.clear(); + for (size_t pos = 0; pos < first_seen_and_detector.size();) { + const size_t low = first_seen_and_detector[pos].first; + while (pos < first_seen_and_detector.size() && first_seen_and_detector[pos].first == low) { + active_detectors.push_back(first_seen_and_detector[pos].second); + ++pos; + } + size_t high = fault_index; + if (pos < first_seen_and_detector.size()) { + high = std::min(high, first_seen_and_detector[pos].first - 1); + } + if (low > high || active_detectors.empty()) { + continue; + } + const double cost = fault.likelihood_cost / active_detectors.size(); + auto& events = events_by_high[high]; + events.reserve(events.size() + active_detectors.size()); + for (int detector : active_detectors) { + events.push_back({detector, low, cost}); + } + } + } + + std::vector, + std::greater>> + heaps(num_detectors); + std::vector> future_costs_by_layer(faults.size()); + + for (size_t layer_index = faults.size(); layer_index-- > 0;) { + for (const auto& event : events_by_high[layer_index]) { + heaps[(size_t)event.detector].push({event.valid_low, event.cost}); + } + + const auto& active = (*layers)[layer_index].current_active_detectors; + auto& costs = future_costs_by_layer[layer_index]; + costs.resize(active.size(), INF); + for (size_t local = 0; local < active.size(); ++local) { + auto& heap = heaps[(size_t)active[local]]; + while (!heap.empty() && heap.top().valid_low > layer_index) { + heap.pop(); + } + if (!heap.empty()) { + costs[local] = heap.top().cost; + } + } + } + + std::vector initial(num_detectors, INF); + if (!layers->empty()) { + const auto& active = layers->front().current_active_detectors; + const auto& costs = future_costs_by_layer.front(); + for (size_t local = 0; local < active.size(); ++local) { + initial[(size_t)active[local]] = costs[local]; + } + } + + for (size_t fault_index = 0; fault_index < faults.size(); ++fault_index) { + auto& layer = (*layers)[fault_index]; + const auto& fault = faults[fault_index]; + + std::vector current_to_next(layer.current_active_detectors.size(), -1); + for (size_t next_local = 0; next_local < layer.surviving_local_indices.size(); ++next_local) { + current_to_next[(size_t)layer.surviving_local_indices[next_local]] = (int32_t)next_local; + } + + layer.next_frontier_costs.resize(layer.surviving_local_indices.size(), INF); + if (fault_index + 1 < faults.size()) { + const auto& next_costs = future_costs_by_layer[fault_index + 1]; + for (size_t next_local = 0; next_local < layer.surviving_local_indices.size(); ++next_local) { + if (next_local < next_costs.size()) { + layer.next_frontier_costs[next_local] = next_costs[next_local]; + } + } + } + + auto& transition = layer.detcost_transition; + transition.fault_local_indices.clear(); + transition.next_local_indices.clear(); + transition.current_costs.clear(); + transition.next_costs.clear(); + transition.fault_local_indices.reserve(fault.detectors.size()); + transition.next_local_indices.reserve(fault.detectors.size()); + transition.current_costs.reserve(fault.detectors.size()); + transition.next_costs.reserve(fault.detectors.size()); + + for (int detector : fault.detectors) { + auto it = std::find(layer.current_active_detectors.begin(), + layer.current_active_detectors.end(), detector); + if (it == layer.current_active_detectors.end()) { + throw std::runtime_error("Missing detector in active frontier while preparing detcost."); + } + uint32_t local = (uint32_t)std::distance(layer.current_active_detectors.begin(), it); + int32_t next_local = current_to_next[local]; + double next_cost = INF; + if (next_local >= 0 && fault_index + 1 < faults.size() && + (size_t)next_local < future_costs_by_layer[fault_index + 1].size()) { + next_cost = future_costs_by_layer[fault_index + 1][(size_t)next_local]; + } + transition.fault_local_indices.push_back(local); + transition.next_local_indices.push_back(next_local); + transition.current_costs.push_back(future_costs_by_layer[fault_index][local]); + transition.next_costs.push_back(next_cost); + } + } + + if (initial_future_detcost != nullptr) { + *initial_future_detcost = std::move(initial); + } +} + +template +void scale_future_detcost_transitions(std::vector* layers, + std::vector* initial_future_detcost, double scale) { + if (scale == 1.0) { + return; + } + auto scale_value = [scale](double value) { return std::isfinite(value) ? value * scale : value; }; + for (auto& value : *initial_future_detcost) { + value = scale_value(value); + } + for (auto& layer : *layers) { + for (auto& value : layer.next_frontier_costs) { + value = scale_value(value); + } + for (auto& value : layer.detcost_transition.current_costs) { + value = scale_value(value); + } + for (auto& value : layer.detcost_transition.next_costs) { + value = scale_value(value); + } + } +} + +size_t num_state_words(size_t num_bits) { + return (num_bits + 63) >> 6; +} + +TESSERACT_ALWAYS_INLINE size_t detector_word_index(size_t detector) { + return detector >> 6; +} + +TESSERACT_ALWAYS_INLINE uint64_t detector_word_mask(size_t detector) { + return uint64_t{1} << (detector & 63); +} + +TESSERACT_ALWAYS_INLINE uint64_t compact_bits_u64(uint64_t value, uint64_t mask) { +#if defined(__BMI2__) && \ + (defined(__x86_64__) || defined(__i386__) || defined(_M_X64) || defined(_M_IX86)) + return _pext_u64(value, mask); +#else + uint64_t out = 0; + uint64_t out_bit = 1; + while (mask) { + uint64_t keep = mask & -mask; + if (value & keep) { + out |= out_bit; + } + mask ^= keep; + out_bit <<= 1; + } + return out; +#endif +} + +double compute_initial_penalty_for_active_detectors( + const std::vector& active_detector_word_indices, + const std::vector& active_detector_bit_masks, + const std::vector& active_detector_costs, + const std::vector& actual_detector_words) { + double total = 0.0; + for (size_t k = 0; k < active_detector_costs.size(); ++k) { + if ((actual_detector_words[active_detector_word_indices[k]] & active_detector_bit_masks[k]) == + 0) { + continue; + } + double best = active_detector_costs[k]; + if (best == INF) { + return INF; + } + total += best; + } + return total; +} + +double score_mass_and_penalty(double mass, double penalty, + TesseractTrellisRankingMode ranking_mode) { + if (ranking_mode == TesseractTrellisRankingMode::MassOnly) { + return mass; + } + if (penalty == INF || mass == 0.0) { + return -INF; + } + return std::log(mass) - penalty; +} + +template +TESSERACT_ALWAYS_INLINE double total_entry_mass(const FixedWideStateEntry& entry) { + return entry.mass0 + entry.mass1; +} + +TESSERACT_ALWAYS_INLINE uint64_t mix_splitmix64(uint64_t value) { + value += 0x9e3779b97f4a7c15ULL; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +void reset_kept_state_stats(TesseractTrellisDecoder* decoder) { + decoder->kept_state_sample_count = 0; + decoder->kept_state_min = 0; + decoder->kept_state_median = 0; + decoder->kept_state_mean = 0; + decoder->kept_state_max = 0; + if (!decoder->config.track_kept_state_stats) { + return; + } + + const size_t histogram_size = decoder->config.beam_width + 1; + if (decoder->kept_state_histogram_scratch.size() != histogram_size) { + decoder->kept_state_histogram_scratch.assign(histogram_size, 0); + } else { + std::fill(decoder->kept_state_histogram_scratch.begin(), + decoder->kept_state_histogram_scratch.end(), 0); + } +} + +void record_kept_state_count(TesseractTrellisDecoder* decoder, size_t kept_states) { + if (!decoder->config.track_kept_state_stats) { + return; + } + + kept_states = std::min(kept_states, decoder->config.beam_width); + if (decoder->kept_state_sample_count == 0) { + decoder->kept_state_min = kept_states; + decoder->kept_state_max = kept_states; + } else { + decoder->kept_state_min = std::min(decoder->kept_state_min, kept_states); + decoder->kept_state_max = std::max(decoder->kept_state_max, kept_states); + } + ++decoder->kept_state_sample_count; + decoder->kept_state_mean += kept_states; + ++decoder->kept_state_histogram_scratch[kept_states]; +} + +void finalize_kept_state_stats(TesseractTrellisDecoder* decoder) { + if (!decoder->config.track_kept_state_stats || decoder->kept_state_sample_count == 0) { + return; + } + + decoder->kept_state_mean /= decoder->kept_state_sample_count; + const size_t lower_target = (decoder->kept_state_sample_count - 1) >> 1; + const size_t upper_target = decoder->kept_state_sample_count >> 1; + size_t seen = 0; + size_t lower = 0; + size_t upper = 0; + bool lower_found = false; + for (size_t kept_states = 0; kept_states < decoder->kept_state_histogram_scratch.size(); + ++kept_states) { + seen += decoder->kept_state_histogram_scratch[kept_states]; + if (!lower_found && seen > lower_target) { + lower = kept_states; + lower_found = true; + } + if (seen > upper_target) { + upper = kept_states; + break; + } + } + decoder->kept_state_median = 0.5 * (lower + upper); +} + +FinalizeKeptStateStatsOnExit::~FinalizeKeptStateStatsOnExit() { + finalize_kept_state_stats(decoder); +} + +void prepare_projected_fault_masks(std::vector* layers) { + for (auto& layer : *layers) { + layer.projected_fault_mask_words.assign(num_state_words(layer.surviving_local_indices.size()), + 0); + for (int32_t next_local : layer.detcost_transition.next_local_indices) { + if (next_local >= 0) { + size_t local = (size_t)next_local; + layer.projected_fault_mask_words[local >> 6] ^= uint64_t{1} << (local & 63); + } + } + } +} + +template +bool fixed_wide_state_less(const FixedWideStateWords& a, + const FixedWideStateWords& b) { + for (size_t k = Words; k-- > 0;) { + if (a[k] != b[k]) { + return a[k] < b[k]; + } + } + return false; +} + +template +bool fixed_wide_state_zero(const FixedWideStateWords& state_words) { + for (size_t k = 0; k < Words; ++k) { + if (state_words[k] != 0) { + return false; + } + } + return true; +} + +template +void xor_compiled_wide_state(FixedWideStateWords* state_words, + const std::array& mask_words) { + for (size_t k = 0; k < Words; ++k) { + (*state_words)[k] ^= mask_words[k]; + } +} + +template +TESSERACT_ALWAYS_INLINE uint64_t +hash_fixed_wide_state(const FixedWideStateWords& state_words) { + uint64_t hash = 0x123456789abcdef0ULL; + for (size_t k = 0; k < Words; ++k) { + hash ^= mix_splitmix64(state_words[k] + 0x9e3779b97f4a7c15ULL * (k + 1)); + hash = std::rotl(hash, 21); + } + return hash; +} + +template +void ensure_pair_bucket_capacity(std::vector>* buckets, + size_t num_parents) { + const size_t required = std::bit_ceil(std::max(16, num_parents * 2)); + if (buckets->size() < required) { + buckets->resize(required); + } +} + +template +void clear_pair_buckets(std::vector>* buckets, + std::vector* used_bucket_indices) { + for (size_t index : *used_bucket_indices) { + (*buckets)[index].occupied = false; + (*buckets)[index].used_mask = 0; + } + used_bucket_indices->clear(); +} + +template +TESSERACT_ALWAYS_INLINE size_t find_or_insert_pair_bucket( + std::vector>* buckets, std::vector* used_bucket_indices, + const FixedWideStateWords& key) { + const size_t mask = buckets->size() - 1; + size_t index = hash_fixed_wide_state(key) & mask; + while ((*buckets)[index].occupied) { + if ((*buckets)[index].key == key) { + return index; + } + index = (index + 1) & mask; + } + + auto& bucket = (*buckets)[index]; + bucket.occupied = true; + bucket.key = key; + bucket.used_mask = 0; + used_bucket_indices->push_back(index); + return index; +} + +template +TESSERACT_ALWAYS_INLINE void accumulate_pair_bucket_slot(FixedWidePairBucket* bucket, + uint8_t slot, double mass0, double mass1, + double penalty) { + const uint8_t bit = (uint8_t)(1u << slot); + if ((bucket->used_mask & bit) == 0) { + bucket->mass0[slot] = mass0; + bucket->mass1[slot] = mass1; + bucket->penalty[slot] = penalty; + bucket->used_mask |= bit; + } else { + bucket->mass0[slot] += mass0; + bucket->mass1[slot] += mass1; + } +} + +template +FixedWideStateWords project_compiled_wide_state( + const FixedWideStateWords& state_words, const CompiledWideLayerTemplate& layer) { + FixedWideStateWords out{}; + for (size_t src_word = 0; src_word < Words; ++src_word) { + const uint64_t mask = layer.surviving_masks[src_word]; + if (mask == 0) { + continue; + } + const uint64_t packed = compact_bits_u64(state_words[src_word], mask); + const size_t dst_word = layer.projection_dst_words[src_word]; + const uint8_t shift = layer.projection_dst_offsets[src_word]; + out[dst_word] |= packed << shift; + if constexpr (Words > 1) { + if (shift != 0 && dst_word + 1 < Words) { + out[dst_word + 1] |= packed >> (64 - shift); + } + } + } + return out; +} + +template +TESSERACT_ALWAYS_INLINE BranchPenaltyUpdate compute_compiled_wide_branch_update( + const FixedWideStateWords& base_state_words, double current_penalty, + const std::vector& actual_detector_words, + const CompiledWideLayerTemplate& layer) { + BranchPenaltyUpdate update; + update.absent_penalty = ComputePenalties ? current_penalty : 0.0; + update.present_penalty = ComputePenalties ? current_penalty : 0.0; + + for (size_t k = 0; k < layer.surviving_term_count; ++k) { + const bool state_bit = + layer.fault_was_active_before[k] && + ((base_state_words[layer.fault_word_indices[k]] & layer.fault_bit_masks[k]) != 0); + const bool target_bit = (actual_detector_words[layer.fault_target_word_indices[k]] & + layer.fault_target_bit_masks[k]) != 0; + const bool mismatch = state_bit ^ target_bit; + + if constexpr (ComputePenalties) { + const double prev_contrib = + (layer.fault_was_active_before[k] && mismatch) ? layer.current_costs[k] : 0.0; + const double next_contrib = mismatch ? layer.next_costs[k] : 0.0; + update.absent_penalty += next_contrib - prev_contrib; + update.present_penalty += (layer.next_costs[k] - next_contrib) - prev_contrib; + } + } + + if constexpr (CheckRetiringTerms) { + for (size_t k = layer.surviving_term_count; k < layer.fault_target_word_indices.size(); ++k) { + const bool state_bit = + layer.fault_was_active_before[k] && + ((base_state_words[layer.fault_word_indices[k]] & layer.fault_bit_masks[k]) != 0); + const bool target_bit = (actual_detector_words[layer.fault_target_word_indices[k]] & + layer.fault_target_bit_masks[k]) != 0; + const bool mismatch = state_bit ^ target_bit; + + if (mismatch) { + update.absent_valid = false; + } else { + update.present_valid = false; + } + + if constexpr (ComputePenalties) { + const double prev_contrib = + (layer.fault_was_active_before[k] && mismatch) ? layer.current_costs[k] : 0.0; + update.absent_penalty -= prev_contrib; + update.present_penalty -= prev_contrib; + } + } + } + + return update; +} + +template +void expand_compiled_layer_into_pair_buckets( + const std::vector>& beam_entries, + std::vector>* pair_buckets, std::vector* used_bucket_indices, + const std::vector& actual_detector_words, + const CompiledWideLayerTemplate& layer, TesseractTrellisDecoder* decoder) { + for (const auto& item : beam_entries) { + ++decoder->num_states_expanded; + BranchPenaltyUpdate update = + compute_compiled_wide_branch_update( + item.state_words, item.penalty, actual_detector_words, layer); + + if (!update.absent_valid && !update.present_valid) { + continue; + } + + FixedWideStateWords projected_state = + project_compiled_wide_state(item.state_words, layer); + FixedWideStateWords projected_toggled = projected_state; + xor_compiled_wide_state(&projected_toggled, layer.projected_fault_mask_words); + const bool projected_is_key = !fixed_wide_state_less(projected_toggled, projected_state); + const auto& bucket_key = projected_is_key ? projected_state : projected_toggled; + const uint8_t absent_slot = projected_is_key ? 0 : 1; + const uint8_t present_slot = projected_toggled == bucket_key ? 0 : 1; + const size_t bucket_index = + find_or_insert_pair_bucket(pair_buckets, used_bucket_indices, bucket_key); + auto& bucket = (*pair_buckets)[bucket_index]; + const bool keep_absent = update.absent_valid && layer.q != 0.0; + const bool keep_present = update.present_valid && layer.p != 0.0; + + if (keep_absent) { + accumulate_pair_bucket_slot(&bucket, absent_slot, item.mass0 * layer.q, item.mass1 * layer.q, + update.absent_penalty); + } + if (keep_present) { + if (layer.toggles_observable) { + accumulate_pair_bucket_slot(&bucket, present_slot, item.mass1 * layer.p, + item.mass0 * layer.p, update.present_penalty); + } else { + accumulate_pair_bucket_slot(&bucket, present_slot, item.mass0 * layer.p, + item.mass1 * layer.p, update.present_penalty); + } + } + } +} + +template +void normalize_compiled_items(std::vector>* items) { + double total = 0.0; + for (const auto& item : *items) { + total += total_entry_mass(item); + } + if (total == 0.0) { + return; + } + for (auto& item : *items) { + item.mass0 /= total; + item.mass1 /= total; + } +} + +template +void merge_equal_compiled_keys_inplace(std::vector>* items) { + if (items->empty()) { + return; + } + std::sort(items->begin(), items->end(), + [](const FixedWideStateEntry& a, const FixedWideStateEntry& b) { + return fixed_wide_state_less(a.state_words, b.state_words); + }); + + size_t out = 0; + for (size_t i = 1; i < items->size(); ++i) { + if ((*items)[i].state_words == (*items)[out].state_words) { + (*items)[out].mass0 += (*items)[i].mass0; + (*items)[out].mass1 += (*items)[i].mass1; + } else { + ++out; + if (out != i) { + (*items)[out] = std::move((*items)[i]); + } + } + } + items->resize(out + 1); +} + +template +bool compiled_state_score_greater(const FixedWideStateEntry& a, + const FixedWideStateEntry& b) { + if (a.score != b.score) { + return a.score > b.score; + } + return fixed_wide_state_less(a.state_words, b.state_words); +} + +template +size_t keep_top_compiled_states(std::vector>* entries, size_t beam_width, + double beam_eps, TesseractTrellisRankingMode ranking_mode) { + if (entries->empty()) { + return 0; + } + + double total_mass = 0.0; + for (auto& entry : *entries) { + const double mass = total_entry_mass(entry); + entry.score = score_mass_and_penalty(mass, entry.penalty, ranking_mode); + if (beam_eps > 0.0) { + total_mass += mass; + } + } + + if (entries->size() > beam_width) { + std::nth_element(entries->begin(), entries->begin() + beam_width, entries->end(), + [](const FixedWideStateEntry& a, const FixedWideStateEntry& b) { + return compiled_state_score_greater(a, b); + }); + entries->resize(beam_width); + } else if (beam_eps <= 0.0) { + return entries->size(); + } + + if (beam_eps <= 0.0 || total_mass <= 0.0) { + return entries->size(); + } + + std::sort(entries->begin(), entries->end(), + [](const FixedWideStateEntry& a, const FixedWideStateEntry& b) { + return compiled_state_score_greater(a, b); + }); + const double retained_target_mass = total_mass * (1.0 - beam_eps); + double retained_mass = 0.0; + size_t keep_count = 0; + while (keep_count < entries->size()) { + retained_mass += total_entry_mass((*entries)[keep_count]); + ++keep_count; + if (retained_mass >= retained_target_mass) { + break; + } + } + entries->resize(keep_count); + return keep_count; +} + +template +std::vector> compile_wide_layers( + const std::vector& layers) { + std::vector> compiled_layers; + compiled_layers.reserve(layers.size()); + for (const auto& layer : layers) { + if (num_state_words(layer.current_active_detectors.size()) > Words || + layer.projected_fault_mask_words.size() > Words) { + throw std::invalid_argument("Compiled wide kernel word count is smaller than the frontier."); + } + + CompiledWideLayerTemplate compiled; + compiled.q = layer.q; + compiled.p = layer.p; + if (layer.obs_mask > 1) { + throw std::invalid_argument("tesseract_trellis currently supports at most one observable"); + } + compiled.toggles_observable = layer.obs_mask != 0; + + std::array surviving_masks{}; + for (uint32_t current_local : layer.surviving_local_indices) { + surviving_masks[current_local >> 6] |= uint64_t{1} << (current_local & 63); + } + size_t next_offset = 0; + for (size_t src_word = 0; src_word < Words; ++src_word) { + compiled.surviving_masks[src_word] = surviving_masks[src_word]; + compiled.projection_dst_words[src_word] = static_cast(next_offset >> 6); + compiled.projection_dst_offsets[src_word] = static_cast(next_offset & 63); + next_offset += std::popcount(surviving_masks[src_word]); + } + + for (size_t k = 0; k < layer.projected_fault_mask_words.size(); ++k) { + compiled.projected_fault_mask_words[k] = layer.projected_fault_mask_words[k]; + } + + const auto& transition = layer.detcost_transition; + const size_t term_count = transition.fault_local_indices.size(); + compiled.fault_target_word_indices.reserve(term_count); + compiled.fault_target_bit_masks.reserve(term_count); + compiled.fault_word_indices.reserve(term_count); + compiled.fault_bit_masks.reserve(term_count); + compiled.fault_was_active_before.reserve(term_count); + compiled.current_costs.reserve(term_count); + compiled.next_costs.reserve(term_count); + auto append_term = [&](size_t idx) { + const uint32_t local = transition.fault_local_indices[idx]; + const uint32_t detector = (uint32_t)layer.current_active_detectors[local]; + compiled.fault_target_word_indices.push_back((uint32_t)detector_word_index(detector)); + compiled.fault_target_bit_masks.push_back(detector_word_mask(detector)); + compiled.fault_word_indices.push_back(static_cast(local >> 6)); + compiled.fault_bit_masks.push_back(uint64_t{1} << (local & 63)); + compiled.fault_was_active_before.push_back(local < layer.previous_width); + compiled.current_costs.push_back(transition.current_costs[idx]); + compiled.next_costs.push_back(transition.next_costs[idx]); + }; + for (size_t idx = 0; idx < term_count; ++idx) { + if (transition.next_local_indices[idx] >= 0) { + append_term(idx); + } + } + compiled.surviving_term_count = compiled.fault_target_word_indices.size(); + compiled.has_retiring_terms = compiled.surviving_term_count != term_count; + for (size_t idx = 0; idx < term_count; ++idx) { + if (transition.next_local_indices[idx] < 0) { + append_term(idx); + } + } + + compiled_layers.push_back(std::move(compiled)); + } + return compiled_layers; +} + +template +struct CompiledWideKernel final : TesseractTrellisWideKernelBase { + explicit CompiledWideKernel(std::vector> layers_, + std::vector initial_detector_word_indices_, + std::vector initial_detector_bit_masks_, + std::vector initial_detector_costs_, + size_t max_frontier_width_) + : layers(std::move(layers_)), + initial_detector_word_indices(std::move(initial_detector_word_indices_)), + initial_detector_bit_masks(std::move(initial_detector_bit_masks_)), + initial_detector_costs(std::move(initial_detector_costs_)), + max_frontier_width(max_frontier_width_) {} + + void decode_shot(TesseractTrellisDecoder* decoder, + const std::vector& detections) const override { + auto& actual_detector_words = decoder->actual_detector_words_scratch; + std::fill(actual_detector_words.begin(), actual_detector_words.end(), 0); + for (uint64_t d : detections) { + if (d >= decoder->num_detectors) { + decoder->low_confidence_flag = true; + return; + } + const size_t word = detector_word_index((size_t)d); + const uint64_t mask = detector_word_mask((size_t)d); + if ((decoder->all_possible_detector_words[word] & mask) == 0) { + decoder->low_confidence_flag = true; + return; + } + actual_detector_words[word] ^= mask; + } + + decoder->max_frontier_width_seen = max_frontier_width; + + double initial_penalty = 0.0; + if (decoder->config.ranking_mode != TesseractTrellisRankingMode::MassOnly && !layers.empty()) { + initial_penalty = compute_initial_penalty_for_active_detectors( + initial_detector_word_indices, initial_detector_bit_masks, initial_detector_costs, + actual_detector_words); + } + + std::vector> beam_entries; + std::vector> next_entries; + std::vector> pair_buckets; + std::vector used_bucket_indices; + beam_entries.reserve(decoder->config.beam_width * 2 + 2); + next_entries.reserve(decoder->config.beam_width * 4 + 4); + beam_entries.push_back({{}, 1.0, 0.0, initial_penalty}); + decoder->max_beam_size_seen = 1; + + const bool compute_penalties = + decoder->config.ranking_mode != TesseractTrellisRankingMode::MassOnly; + for (size_t layer_index = 0; layer_index < layers.size(); ++layer_index) { + const auto& layer = layers[layer_index]; + + ensure_pair_bucket_capacity(&pair_buckets, beam_entries.size()); + clear_pair_buckets(&pair_buckets, &used_bucket_indices); + + auto t0 = std::chrono::high_resolution_clock::now(); + + if (decoder->config.verbose) { + std::cout << "expanding layer " << layer_index << " / " << (layers.size() - 1) << std::endl; + std::cout << "states to expand = " << beam_entries.size() << std::endl; + } + if (compute_penalties) { + if (layer.has_retiring_terms) { + expand_compiled_layer_into_pair_buckets( + beam_entries, &pair_buckets, &used_bucket_indices, actual_detector_words, layer, + decoder); + } else { + expand_compiled_layer_into_pair_buckets( + beam_entries, &pair_buckets, &used_bucket_indices, actual_detector_words, layer, + decoder); + } + } else if (layer.has_retiring_terms) { + expand_compiled_layer_into_pair_buckets( + beam_entries, &pair_buckets, &used_bucket_indices, actual_detector_words, layer, + decoder); + } else { + expand_compiled_layer_into_pair_buckets( + beam_entries, &pair_buckets, &used_bucket_indices, actual_detector_words, layer, + decoder); + } + auto t1 = std::chrono::high_resolution_clock::now(); + decoder->time_expand_seconds += + std::chrono::duration_cast(t1 - t0).count() / 1e6; + + auto t2a = std::chrono::high_resolution_clock::now(); + next_entries.clear(); + next_entries.reserve(used_bucket_indices.size() * 2); + for (size_t index : used_bucket_indices) { + auto& bucket = pair_buckets[index]; + if ((bucket.used_mask & 1u) != 0) { + next_entries.push_back({bucket.key, bucket.mass0[0], bucket.mass1[0], bucket.penalty[0]}); + } + if ((bucket.used_mask & 2u) != 0) { + auto other_state = bucket.key; + xor_compiled_wide_state(&other_state, layer.projected_fault_mask_words); + next_entries.push_back( + {std::move(other_state), bucket.mass0[1], bucket.mass1[1], bucket.penalty[1]}); + } + } + beam_entries.swap(next_entries); + auto t2 = std::chrono::high_resolution_clock::now(); + decoder->time_collapse_seconds += + std::chrono::duration_cast(t2 - t2a).count() / 1e6; + + const size_t kept_states = + keep_top_compiled_states(&beam_entries, decoder->config.beam_width, + decoder->config.beam_eps, decoder->config.ranking_mode); + normalize_compiled_items(&beam_entries); + record_kept_state_count(decoder, beam_entries.empty() ? 0 : kept_states); + if (beam_entries.empty()) { + decoder->low_confidence_flag = true; + return; + } + decoder->num_states_merged += kept_states; + decoder->max_beam_size_seen = std::max(decoder->max_beam_size_seen, kept_states); + auto t3 = std::chrono::high_resolution_clock::now(); + decoder->time_truncate_seconds += + std::chrono::duration_cast(t3 - t2).count() / 1e6; + } + + auto tr0 = std::chrono::high_resolution_clock::now(); + for (const auto& item : beam_entries) { + if (!fixed_wide_state_zero(item.state_words)) { + continue; + } + decoder->total_mass_obs0 += item.mass0; + decoder->total_mass_obs1 += item.mass1; + } + if (decoder->total_mass_obs0 == 0.0 && decoder->total_mass_obs1 == 0.0) { + decoder->low_confidence_flag = true; + return; + } + decoder->predicted_obs_mask = decoder->total_mass_obs1 > decoder->total_mass_obs0 ? 1 : 0; + auto tr1 = std::chrono::high_resolution_clock::now(); + decoder->time_reconstruct_seconds += + std::chrono::duration_cast(tr1 - tr0).count() / 1e6; + } + + std::vector> layers; + std::vector initial_detector_word_indices; + std::vector initial_detector_bit_masks; + std::vector initial_detector_costs; + size_t max_frontier_width; +}; + +std::unique_ptr build_compiled_wide_kernel( + const std::vector& layers, size_t max_frontier_width, + const std::vector& initial_future_detcost) { + const size_t required_words = std::max(1, num_state_words(max_frontier_width)); + if (required_words > kMaxCompiledWideStateWords) { + throw std::invalid_argument("Wide trellis frontier requires " + std::to_string(required_words) + + " words, but only " + std::to_string(kMaxCompiledWideStateWords) + + " compiled words are enabled."); + } + + std::vector initial_detector_word_indices; + std::vector initial_detector_bit_masks; + std::vector initial_detector_costs; + if (!layers.empty()) { + const auto& initial_active_detectors = layers.front().current_active_detectors; + initial_detector_word_indices.reserve(initial_active_detectors.size()); + initial_detector_bit_masks.reserve(initial_active_detectors.size()); + initial_detector_costs.reserve(initial_active_detectors.size()); + for (int detector : initial_active_detectors) { + initial_detector_word_indices.push_back((uint32_t)detector_word_index((size_t)detector)); + initial_detector_bit_masks.push_back(detector_word_mask((size_t)detector)); + initial_detector_costs.push_back(initial_future_detcost[(size_t)detector]); + } + } + switch (required_words) { + case 1: + return std::make_unique>( + compile_wide_layers<1>(layers), initial_detector_word_indices, initial_detector_bit_masks, + initial_detector_costs, max_frontier_width); + case 2: + return std::make_unique>( + compile_wide_layers<2>(layers), initial_detector_word_indices, initial_detector_bit_masks, + initial_detector_costs, max_frontier_width); + case 3: + return std::make_unique>( + compile_wide_layers<3>(layers), initial_detector_word_indices, initial_detector_bit_masks, + initial_detector_costs, max_frontier_width); + case 4: + return std::make_unique>( + compile_wide_layers<4>(layers), initial_detector_word_indices, initial_detector_bit_masks, + initial_detector_costs, max_frontier_width); + default: + throw std::invalid_argument("Unsupported compiled wide trellis word count."); + } +} + +} // namespace + +TesseractTrellisDecoder::~TesseractTrellisDecoder() = default; + +TesseractTrellisDecoder::TesseractTrellisDecoder(TesseractTrellisConfig config_) + : config(std::move(config_)) { + // Maps original flattened DEM error indices to currently preprocessed indices. + std::vector dem_error_map(config.dem.flattened().count_errors()); + std::iota(dem_error_map.begin(), dem_error_map.end(), 0); + + if (config.merge_errors) { + std::vector merge_map; + config.dem = common::merge_indistinguishable_errors(config.dem, merge_map); + common::chain_error_maps(dem_error_map, merge_map); + } + + std::vector nonzero_map; + config.dem = common::remove_zero_probability_errors(config.dem, nonzero_map); + common::chain_error_maps(dem_error_map, nonzero_map); + + dem_error_to_error = std::move(dem_error_map); + error_to_dem_error = common::invert_error_map(dem_error_to_error, config.dem.count_errors()); + + errors = get_errors_from_dem(config.dem.flattened()); + num_detectors = config.dem.count_detectors(); + num_observables = config.dem.count_observables(); + if (num_observables > 1) { + throw std::invalid_argument("tesseract_trellis currently supports at most one observable"); + } + + all_possible_detector_words.assign(num_state_words(num_detectors), 0); + actual_detector_words_scratch.assign(all_possible_detector_words.size(), 0); + for (const auto& error : errors) { + for (int d : error.symptom.detectors) { + all_possible_detector_words[detector_word_index((size_t)d)] |= detector_word_mask((size_t)d); + } + } + + auto faults = parse_faults(errors, num_observables); + + size_t wide_frontier_width = 0; + build_wide_layer_templates(faults, num_detectors, &wide_layer_templates, &wide_frontier_width); + if (config.ranking_mode == TesseractTrellisRankingMode::FutureActiveDetcostRanked) { + build_future_active_detcost_transitions(faults, num_detectors, &wide_layer_templates, + &initial_future_detcost); + } else { + build_future_detcost_transitions(faults, num_detectors, &wide_layer_templates, + &initial_future_detcost); + } + if (config.ranking_mode != TesseractTrellisRankingMode::MassOnly) { + scale_future_detcost_transitions(&wide_layer_templates, &initial_future_detcost, + config.future_detcost_scale); + } + prepare_projected_fault_masks(&wide_layer_templates); + wide_kernel = + build_compiled_wide_kernel(wide_layer_templates, wide_frontier_width, initial_future_detcost); + wide_layer_templates.clear(); +} + +TESSERACT_HOT void TesseractTrellisDecoder::decode_shot(const std::vector& detections) { + low_confidence_flag = false; + num_states_expanded = 0; + num_states_merged = 0; + max_beam_size_seen = 0; + max_frontier_width_seen = 0; + reset_kept_state_stats(this); + time_expand_seconds = 0; + time_collapse_seconds = 0; + time_truncate_seconds = 0; + time_reconstruct_seconds = 0; + predicted_obs_mask = 0; + total_mass_obs0 = 0; + total_mass_obs1 = 0; + FinalizeKeptStateStatsOnExit kept_state_stats_guard{this}; + wide_kernel->decode_shot(this, detections); + + if (config.verbose) { + std::cout << "trellis beam_width=" << config.beam_width + << " frontier_width=" << max_frontier_width_seen + << " states_expanded=" << num_states_expanded + << " states_merged=" << num_states_merged << " max_beam=" << max_beam_size_seen + << std::endl; + } +} + +double TesseractTrellisDecoder::observable_probability() const { + const double total_mass = total_mass_obs0 + total_mass_obs1; + if (!std::isfinite(total_mass) || total_mass <= 0) { + return std::numeric_limits::quiet_NaN(); + } + return total_mass_obs1 / total_mass; +} + +std::vector TesseractTrellisDecoder::decode(const std::vector& detections) { + decode_shot(detections); + return predicted_obs_mask ? std::vector{0} : std::vector{}; +} + +void TesseractTrellisDecoder::decode_shots(std::vector& shots, + std::vector>& obs_predicted) { + obs_predicted.resize(shots.size()); + for (size_t i = 0; i < shots.size(); ++i) { + obs_predicted[i] = decode(shots[i].hits); + } +} diff --git a/src/tesseract_trellis.h b/src/tesseract_trellis.h new file mode 100644 index 0000000..fec0a5e --- /dev/null +++ b/src/tesseract_trellis.h @@ -0,0 +1,105 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TESSERACT_TRELLIS_DECODER_H +#define TESSERACT_TRELLIS_DECODER_H + +#include +#include +#include + +#include "common.h" +#include "stim.h" + +struct TesseractTrellisWideKernelBase; + +enum class TesseractTrellisRankingMode { + MassOnly, + FutureDetcostRanked, + FutureActiveDetcostRanked, +}; + +struct TesseractTrellisDetcostTransition { + std::vector fault_local_indices; + std::vector next_local_indices; + std::vector current_costs; + std::vector next_costs; +}; + +struct TesseractTrellisWideLayerTemplate { + double q = 0; + double p = 0; + uint64_t obs_mask = 0; + size_t previous_width = 0; + std::vector surviving_local_indices; + std::vector current_active_detectors; + std::vector projected_fault_mask_words; + std::vector next_frontier_costs; + TesseractTrellisDetcostTransition detcost_transition; +}; + +struct TesseractTrellisConfig { + stim::DetectorErrorModel dem; + size_t beam_width = 1024; + double beam_eps = 0.0; + double future_detcost_scale = 2.0; + bool verbose = false; + bool merge_errors = true; + bool track_kept_state_stats = false; + TesseractTrellisRankingMode ranking_mode = TesseractTrellisRankingMode::MassOnly; +}; + +struct TesseractTrellisDecoder { + explicit TesseractTrellisDecoder(TesseractTrellisConfig config); + ~TesseractTrellisDecoder(); + + void decode_shot(const std::vector& detections); + double observable_probability() const; + std::vector decode(const std::vector& detections); + void decode_shots(std::vector& shots, + std::vector>& obs_predicted); + + TesseractTrellisConfig config; + bool low_confidence_flag = false; + size_t num_states_expanded = 0; + size_t num_states_merged = 0; + size_t max_beam_size_seen = 0; + size_t max_frontier_width_seen = 0; + size_t kept_state_sample_count = 0; + size_t kept_state_min = 0; + double kept_state_median = 0; + double kept_state_mean = 0; + size_t kept_state_max = 0; + double time_expand_seconds = 0; + double time_collapse_seconds = 0; + double time_truncate_seconds = 0; + double time_reconstruct_seconds = 0; + uint64_t predicted_obs_mask = 0; + double total_mass_obs0 = 0; + double total_mass_obs1 = 0; + + std::vector dem_error_to_error; + std::vector error_to_dem_error; + std::vector errors; + size_t num_observables = 0; + size_t num_detectors = 0; + std::vector all_possible_detector_words; + std::vector actual_detector_words_scratch; + std::vector wide_layer_templates; + std::vector initial_future_detcost; + std::unique_ptr wide_kernel; + std::vector kept_state_histogram_scratch; +}; + +#endif // TESSERACT_TRELLIS_DECODER_H diff --git a/src/tesseract_trellis.test.cc b/src/tesseract_trellis.test.cc new file mode 100644 index 0000000..e375212 --- /dev/null +++ b/src/tesseract_trellis.test.cc @@ -0,0 +1,223 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tesseract_trellis.h" + +#include + +#include +#include +#include + +#include "stim.h" + +TEST(TesseractTrellisDecoderTest, ComputesObservableProbabilityForAmbiguousSyndrome) { + stim::DetectorErrorModel dem(R"DEM( + error(0.1) D0 + error(0.2) D0 L0 + detector(0, 0, 0) D0 + )DEM"); + + TesseractTrellisConfig config; + config.dem = dem; + config.beam_width = 16; + TesseractTrellisDecoder decoder(config); + + decoder.decode_shot({0}); + EXPECT_FALSE(decoder.low_confidence_flag); + EXPECT_EQ(decoder.predicted_obs_mask, 1); + EXPECT_NEAR(decoder.observable_probability(), 0.18 / 0.26, 1e-12); + + decoder.decode_shot({}); + EXPECT_FALSE(decoder.low_confidence_flag); + EXPECT_EQ(decoder.predicted_obs_mask, 0); + EXPECT_NEAR(decoder.observable_probability(), 0.02 / 0.74, 1e-12); +} + +TEST(TesseractTrellisDecoderTest, SumsProbabilityMassAcrossMultipleExplanations) { + // A min-cost decoder would prefer the single right branch over any one left + // chain. The trellis should sum the three distinct left chains' probability + // mass and predict L0. These chains are not mergeable because they touch + // different hidden detectors. + // + // L0 chains no-L0 chain + // D0 --a1(L0)-- D1 --b1-- * D0 --r-- * + // D0 --a2(L0)-- D2 --b2-- * + // D0 --a3(L0)-- D3 --b3-- * + stim::DetectorErrorModel dem(R"DEM( + error(0.2) D0 D1 L0 + error(0.2) D1 + error(0.2) D0 D2 L0 + error(0.2) D2 + error(0.2) D0 D3 L0 + error(0.2) D3 + error(0.1) D0 + detector(0, 0, 0) D0 + detector(1, 0, 0) D1 + detector(2, 0, 0) D2 + detector(3, 0, 0) D3 + )DEM"); + + TesseractTrellisConfig config; + config.dem = dem; + config.beam_width = 64; + TesseractTrellisDecoder decoder(config); + + decoder.decode_shot({0}); + + const double p_chain_edge = 0.2; + const double p_right = 0.1; + const double q_chain_edge = 1 - p_chain_edge; + const double chain_present = p_chain_edge * p_chain_edge; + const double chain_absent = q_chain_edge * q_chain_edge; + const double odd_chain_mass = 3 * chain_present * chain_absent * chain_absent + + chain_present * chain_present * chain_present; + const double even_chain_mass = + chain_absent * chain_absent * chain_absent + 3 * chain_present * chain_present * chain_absent; + const double mass_l0 = odd_chain_mass * (1 - p_right); + const double mass_no_l0 = even_chain_mass * p_right; + const double expected_probability = mass_l0 / (mass_l0 + mass_no_l0); + const double right_cost = -std::log(p_right / (1 - p_right)); + const double left_chain_cost = 2 * -std::log(p_chain_edge / (1 - p_chain_edge)); + + EXPECT_FALSE(decoder.low_confidence_flag); + EXPECT_EQ(decoder.predicted_obs_mask, 1); + EXPECT_TRUE(decoder.config.merge_errors); + EXPECT_LT(right_cost, left_chain_cost); + EXPECT_GT(expected_probability, 0.5); + EXPECT_NEAR(decoder.observable_probability(), expected_probability, 1e-12); +} + +TEST(TesseractTrellisDecoderTest, ReportsNanProbabilityForInvalidDetector) { + stim::DetectorErrorModel dem(R"DEM( + error(0.1) D0 L0 + detector(0, 0, 0) D0 + )DEM"); + + TesseractTrellisConfig config; + config.dem = dem; + TesseractTrellisDecoder decoder(config); + + decoder.decode_shot({1}); + EXPECT_TRUE(decoder.low_confidence_flag); + EXPECT_TRUE(std::isnan(decoder.observable_probability())); +} + +TEST(TesseractTrellisDecoderTest, RankingModesDecodeAmbiguousSyndrome) { + stim::DetectorErrorModel dem(R"DEM( + error(0.1) D0 + error(0.2) D0 L0 + detector(0, 0, 0) D0 + )DEM"); + + for (auto ranking_mode : + {TesseractTrellisRankingMode::MassOnly, TesseractTrellisRankingMode::FutureDetcostRanked, + TesseractTrellisRankingMode::FutureActiveDetcostRanked}) { + TesseractTrellisConfig config; + config.dem = dem; + config.beam_width = 16; + config.ranking_mode = ranking_mode; + TesseractTrellisDecoder decoder(config); + + decoder.decode_shot({0}); + EXPECT_FALSE(decoder.low_confidence_flag); + EXPECT_EQ(decoder.predicted_obs_mask, 1); + EXPECT_NEAR(decoder.observable_probability(), 0.18 / 0.26, 1e-12); + } +} + +TEST(TesseractTrellisDecoderTest, BeamEpsSmokeTest) { + stim::DetectorErrorModel dem(R"DEM( + error(0.1) D0 + error(0.2) D0 L0 + detector(0, 0, 0) D0 + )DEM"); + + TesseractTrellisConfig config; + config.dem = dem; + config.beam_width = 16; + config.beam_eps = 0.1; + TesseractTrellisDecoder decoder(config); + + decoder.decode_shot({0}); + EXPECT_FALSE(decoder.low_confidence_flag); + EXPECT_TRUE(std::isfinite(decoder.observable_probability())); +} + +TEST(TesseractTrellisDecoderTest, MergeErrorsMatchesOtherDecoders) { + stim::DetectorErrorModel dem(R"DEM( + error(0.1) D0 L0 + error(0.2) D0 L0 + error(0) D0 + detector(0, 0, 0) D0 + )DEM"); + + TesseractTrellisConfig merged_config; + merged_config.dem = dem; + merged_config.beam_width = 16; + TesseractTrellisDecoder merged_decoder(merged_config); + + EXPECT_TRUE(merged_decoder.config.merge_errors); + EXPECT_EQ(merged_decoder.errors.size(), 1); + EXPECT_EQ(merged_decoder.config.dem.count_errors(), 1); + EXPECT_EQ(merged_decoder.dem_error_to_error.size(), 3); + EXPECT_EQ(merged_decoder.dem_error_to_error[0], 0); + EXPECT_EQ(merged_decoder.dem_error_to_error[1], 0); + EXPECT_EQ(merged_decoder.dem_error_to_error[2], std::numeric_limits::max()); + EXPECT_EQ(merged_decoder.error_to_dem_error.size(), 1); + EXPECT_EQ(merged_decoder.error_to_dem_error[0], 0); + + merged_decoder.decode_shot({0}); + EXPECT_FALSE(merged_decoder.low_confidence_flag); + EXPECT_EQ(merged_decoder.predicted_obs_mask, 1); + EXPECT_NEAR(merged_decoder.observable_probability(), 1.0, 1e-12); + + TesseractTrellisConfig unmerged_config; + unmerged_config.dem = dem; + unmerged_config.beam_width = 16; + unmerged_config.merge_errors = false; + TesseractTrellisDecoder unmerged_decoder(unmerged_config); + + EXPECT_FALSE(unmerged_decoder.config.merge_errors); + EXPECT_EQ(unmerged_decoder.errors.size(), 2); + EXPECT_EQ(unmerged_decoder.config.dem.count_errors(), 2); + EXPECT_EQ(unmerged_decoder.dem_error_to_error.size(), 3); + EXPECT_EQ(unmerged_decoder.dem_error_to_error[0], 0); + EXPECT_EQ(unmerged_decoder.dem_error_to_error[1], 1); + EXPECT_EQ(unmerged_decoder.dem_error_to_error[2], std::numeric_limits::max()); + + unmerged_decoder.decode_shot({0}); + EXPECT_FALSE(unmerged_decoder.low_confidence_flag); + EXPECT_EQ(unmerged_decoder.predicted_obs_mask, 1); + EXPECT_NEAR(unmerged_decoder.observable_probability(), merged_decoder.observable_probability(), + 1e-12); +} + +TEST(TesseractTrellisDecoderTest, RejectsMoreThanOneObservable) { + stim::DetectorErrorModel dem(R"DEM( + error(0.1) D0 L0 + error(0.1) D0 L1 + detector(0, 0, 0) D0 + )DEM"); + + TesseractTrellisConfig config; + config.dem = dem; + + try { + TesseractTrellisDecoder decoder(config); + FAIL() << "Expected TesseractTrellisDecoder construction to fail."; + } catch (const std::invalid_argument& err) { + EXPECT_NE(std::string(err.what()).find("supports at most one observable"), std::string::npos); + } +} diff --git a/src/tesseract_trellis_main.cc b/src/tesseract_trellis_main.cc new file mode 100644 index 0000000..eec3bb6 --- /dev/null +++ b/src/tesseract_trellis_main.cc @@ -0,0 +1,490 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "stim.h" +#include "tesseract_trellis.h" +#include "utils.h" + +namespace { + +TesseractTrellisRankingMode parse_ranking_mode(const std::string& value) { + if (value == "mass") return TesseractTrellisRankingMode::MassOnly; + if (value == "future-detcost") return TesseractTrellisRankingMode::FutureDetcostRanked; + if (value == "future-active-detcost") { + return TesseractTrellisRankingMode::FutureActiveDetcostRanked; + } + throw std::invalid_argument("Unknown trellis ranking mode: " + value); +} + +} // namespace + +struct Args { + std::string circuit_path; + std::string dem_path; + bool no_merge_errors = false; + + size_t sample_num_shots = 0; + size_t max_errors = SIZE_MAX; + uint64_t sample_seed; + + size_t shot_range_begin = 0; + size_t shot_range_end = 0; + + std::string in_fname = ""; + std::string in_format = ""; + std::string obs_in_fname = ""; + std::string obs_in_format = ""; + bool append_observables = false; + std::string out_fname = ""; + std::string out_format = ""; + std::string obs_probs_out_fname = ""; + + std::string dem_out_fname = ""; + std::string stats_out_fname = ""; + + size_t num_threads = 1; + size_t beam_width = 1024; + double beam_eps = 0.0; + double future_detcost_scale = 2.0; + std::string ranking_mode = "mass"; + + bool verbose = false; + bool print_stats = false; + + bool has_observables() { + return append_observables || !obs_in_fname.empty() || (sample_num_shots > 0); + } + + void validate() { + if (circuit_path.empty() && dem_path.empty()) { + throw std::invalid_argument("Must provide at least one of --circuit or --dem"); + } + int num_data_sources = int(sample_num_shots > 0) + int(!in_fname.empty()); + if (num_data_sources != 1) { + throw std::invalid_argument("Requires exactly 1 source of shots."); + } + if (!in_fname.empty() && in_format.empty()) { + throw std::invalid_argument("If --in is provided, must also specify --in-format."); + } + if (!out_fname.empty() && out_format.empty()) { + throw std::invalid_argument("If --out is provided, must also specify --out-format."); + } + if (!dem_out_fname.empty()) { + throw std::invalid_argument("--dem-out is not supported by tesseract_trellis."); + } + if (obs_probs_out_fname == "-") { + throw std::invalid_argument("--obs-probs-out must be a file path, not stdout."); + } + if (!in_format.empty() && !stim::format_name_to_enum_map().contains(in_format)) { + throw std::invalid_argument("Invalid format: " + in_format); + } + if (!obs_in_format.empty() && !stim::format_name_to_enum_map().contains(obs_in_format)) { + throw std::invalid_argument("Invalid format: " + obs_in_format); + } + if (!out_format.empty() && !stim::format_name_to_enum_map().contains(out_format)) { + throw std::invalid_argument("Invalid format: " + out_format); + } + if (!obs_in_fname.empty() && in_fname.empty()) { + throw std::invalid_argument( + "Cannot load observable flips without a corresponding detection event data file."); + } + if (num_threads == 0) { + throw std::invalid_argument("--threads must be at least 1."); + } + if (num_threads > 1000) { + throw std::invalid_argument("There is a maximum limit of 1000 threads."); + } + if ((shot_range_begin || shot_range_end) && shot_range_end < shot_range_begin) { + throw std::invalid_argument("Provided shot range must have end >= begin."); + } + if (sample_num_shots > 0 && circuit_path.empty()) { + throw std::invalid_argument("Cannot sample shots without a circuit."); + } + if (beam_width == 0) { + throw std::invalid_argument("--beam must be at least 1."); + } + if (!std::isfinite(beam_eps) || beam_eps < 0.0 || beam_eps >= 1.0) { + throw std::invalid_argument("--beam-eps must satisfy 0 <= beam-eps < 1."); + } + if (!std::isfinite(future_detcost_scale) || future_detcost_scale < 0.0) { + throw std::invalid_argument("--future-detcost-scale must be finite and nonnegative."); + } + parse_ranking_mode(ranking_mode); + } + + void extract(TesseractTrellisConfig& config, std::vector& shots, + std::unique_ptr& writer) { + stim::Circuit circuit; + if (!circuit_path.empty()) { + FILE* file = fopen(circuit_path.c_str(), "r"); + if (!file) { + throw std::invalid_argument("Could not open the file: " + circuit_path); + } + circuit = stim::Circuit::from_file(file); + fclose(file); + } + + if (!dem_path.empty()) { + FILE* file = fopen(dem_path.c_str(), "r"); + if (!file) { + throw std::invalid_argument("Could not open the file: " + dem_path); + } + config.dem = stim::DetectorErrorModel::from_file(file); + fclose(file); + } else { + assert(!circuit_path.empty()); + config.dem = stim::ErrorAnalyzer::circuit_to_detector_error_model( + circuit, /*decompose_errors=*/false, /*fold_loops=*/true, + /*allow_gauge_detectors=*/true, + /*approximate_disjoint_errors_threshold=*/1, + /*ignore_decomposition_failures=*/false, + /*block_decomposition_from_introducing_remnant_edges=*/false); + } + + config.merge_errors = !no_merge_errors; + config.beam_width = beam_width; + config.beam_eps = beam_eps; + config.future_detcost_scale = future_detcost_scale; + config.verbose = verbose; + config.track_kept_state_stats = print_stats; + config.ranking_mode = parse_ranking_mode(ranking_mode); + + if (sample_num_shots > 0) { + assert(!circuit_path.empty()); + std::mt19937_64 rng(sample_seed); + size_t num_detectors = circuit.count_detectors(); + const auto [dets, obs] = + stim::sample_batch_detection_events<64>(circuit, sample_num_shots, rng); + stim::simd_bit_table<64> obs_T = obs.transposed(); + shots.resize(sample_num_shots); + for (size_t k = 0; k < sample_num_shots; k++) { + shots[k].obs_mask = obs_T[k]; + for (size_t d = 0; d < num_detectors; d++) { + if (dets[d][k]) { + shots[k].hits.push_back(d); + } + } + } + } + + if (!in_fname.empty()) { + FILE* shots_file = fopen(in_fname.c_str(), "r"); + if (!shots_file) { + throw std::invalid_argument("Could not open the file: " + in_fname); + } + stim::FileFormatData shots_in_format = stim::format_name_to_enum_map().at(in_format); + auto reader = stim::MeasureRecordReader::make( + shots_file, shots_in_format.id, 0, config.dem.count_detectors(), + append_observables * config.dem.count_observables()); + stim::SparseShot sparse_shot; + sparse_shot.clear(); + while (reader->start_and_read_entire_record(sparse_shot)) { + shots.push_back(sparse_shot); + sparse_shot.clear(); + } + fclose(shots_file); + } + + if (!obs_in_fname.empty()) { + FILE* obs_file = fopen(obs_in_fname.c_str(), "r"); + if (!obs_file) { + throw std::invalid_argument("Could not open the file: " + obs_in_fname); + } + stim::FileFormatData obs_format = stim::format_name_to_enum_map().at(obs_in_format); + auto obs_reader = stim::MeasureRecordReader::make( + obs_file, obs_format.id, 0, 0, config.dem.count_observables()); + stim::SparseShot sparse_shot; + sparse_shot.clear(); + size_t num_obs_shots = 0; + while (obs_reader->start_and_read_entire_record(sparse_shot)) { + if (num_obs_shots >= shots.size()) { + throw std::invalid_argument("Shot data ended before obs data."); + } + shots[num_obs_shots].obs_mask = sparse_shot.obs_mask; + sparse_shot.clear(); + ++num_obs_shots; + } + if (num_obs_shots != shots.size()) { + throw std::invalid_argument("Obs data ended before shot data ended."); + } + fclose(obs_file); + } + + if (shot_range_begin || shot_range_end) { + if (shot_range_end > shots.size()) { + throw std::invalid_argument("Shot range end is past end of shots array."); + } + std::vector shots_in_range(shots.begin() + shot_range_begin, + shots.begin() + shot_range_end); + std::swap(shots_in_range, shots); + } + + if (!out_fname.empty()) { + stim::FileFormatData predictions_out_format = stim::format_name_to_enum_map().at(out_format); + FILE* predictions_file = stdout; + // An output path of "-" means stdout. + if (out_fname != "-") { + predictions_file = fopen(out_fname.c_str(), "w"); + if (!predictions_file) { + throw std::invalid_argument("Could not open the file: " + out_fname); + } + } + writer = stim::MeasureRecordWriter::make(predictions_file, predictions_out_format.id); + writer->begin_result_type('L'); + } + } +}; + +int main(int argc, char* argv[]) { + std::cout.precision(16); + argparse::ArgumentParser program("tesseract_trellis"); + Args args; + program.add_argument("--circuit").help("Stim circuit file path").store_into(args.circuit_path); + program.add_argument("--dem").help("Stim dem file path").store_into(args.dem_path); + program.add_argument("--no-merge-errors") + .help("If provided, will not merge identical error mechanisms.") + .store_into(args.no_merge_errors); + program.add_argument("--sample-num-shots").store_into(args.sample_num_shots); + program.add_argument("--max-errors").store_into(args.max_errors); + program.add_argument("--sample-seed") + .default_value(static_cast(std::random_device()())) + .store_into(args.sample_seed); + program.add_argument("--shot-range-begin") + .default_value(size_t(0)) + .store_into(args.shot_range_begin); + program.add_argument("--shot-range-end").default_value(size_t(0)).store_into(args.shot_range_end); + program.add_argument("--in").default_value(std::string("")).store_into(args.in_fname); + program.add_argument("--in-format", "--in_format") + .default_value(std::string("")) + .store_into(args.in_format); + program.add_argument("--in-includes-appended-observables", "--in_includes_appended_observables") + .default_value(false) + .store_into(args.append_observables) + .flag(); + program.add_argument("--obs_in", "--obs-in") + .default_value(std::string("")) + .store_into(args.obs_in_fname); + program.add_argument("--obs-in-format", "--obs_in_format") + .default_value(std::string("")) + .store_into(args.obs_in_format); + program.add_argument("--out").default_value(std::string("")).store_into(args.out_fname); + program.add_argument("--out-format").default_value(std::string("")).store_into(args.out_format); + program.add_argument("--obs-probs-out") + .help( + "File to write headerless binary doubles containing P(L0=1) for each decoded shot. " + "Requires exactly one observable.") + .default_value(std::string("")) + .store_into(args.obs_probs_out_fname); + program.add_argument("--dem-out").default_value(std::string("")).store_into(args.dem_out_fname); + program.add_argument("--stats-out") + .default_value(std::string("")) + .store_into(args.stats_out_fname); + program.add_argument("--threads") + .default_value(size_t( + std::thread::hardware_concurrency() == 0 ? 1 : std::thread::hardware_concurrency())) + .store_into(args.num_threads); + program.add_argument("--beam").default_value(size_t(1024)).store_into(args.beam_width); + program.add_argument("--beam-eps") + .help( + "Keep at most --beam merged states and also drop the suffix once the kept prefix has " + "accumulated at least (1 - beam-eps) of the total merged-state mass. Use 0 to disable " + "the mass-threshold cutoff.") + .default_value(0.0) + .store_into(args.beam_eps); + program.add_argument("--ranking-mode") + .help("Trellis ranking mode: mass, future-detcost, or future-active-detcost") + .default_value(std::string("mass")) + .store_into(args.ranking_mode); + program.add_argument("--future-detcost-scale") + .help("Multiplier applied to future detector-cost ranking penalties.") + .default_value(2.0) + .store_into(args.future_detcost_scale); + program.add_argument("--verbose").flag().store_into(args.verbose); + program.add_argument("--print-stats").flag().store_into(args.print_stats); + + try { + program.parse_args(argc, argv); + } catch (const std::exception& err) { + std::cerr << err.what() << std::endl; + std::cerr << program; + return EXIT_FAILURE; + } + + args.validate(); + TesseractTrellisConfig config; + std::vector shots; + std::unique_ptr writer; + args.extract(config, shots, writer); + + std::vector obs_predicted(shots.size()); + std::vector mass0_predicted(shots.size()); + std::vector mass1_predicted(shots.size()); + std::vector obs_probability_predicted(shots.size()); + std::vector decoding_time_seconds(shots.size()); + std::vector num_states_expanded_per_shot(shots.size()); + std::vector num_states_merged_per_shot(shots.size()); + std::vector max_beam_size_per_shot(shots.size()); + std::vector max_frontier_width_per_shot(shots.size()); + std::vector kept_state_min_per_shot(shots.size()); + std::vector kept_state_median_per_shot(shots.size()); + std::vector kept_state_mean_per_shot(shots.size()); + std::vector kept_state_max_per_shot(shots.size()); + std::vector time_expand_per_shot(shots.size()); + std::vector time_collapse_per_shot(shots.size()); + std::vector time_truncate_per_shot(shots.size()); + std::vector time_reconstruct_per_shot(shots.size()); + std::vector> low_confidence(shots.size()); + const stim::DetectorErrorModel original_dem = config.dem.flattened(); + std::vector> decoders(args.num_threads); + + bool has_obs = args.has_observables(); + size_t num_errors = 0; + size_t num_low_confidence = 0; + double total_time_seconds = 0; + size_t num_observables = config.dem.count_observables(); + if (num_observables > 1) { + std::cerr << "tesseract_trellis currently supports at most one observable; DEM has " + << num_observables << "." << std::endl; + return EXIT_FAILURE; + } + if (!args.obs_probs_out_fname.empty() && num_observables != 1) { + std::cerr << "--obs-probs-out requires a DEM with exactly one observable." << std::endl; + return EXIT_FAILURE; + } + + size_t shot = parallel_for_shots_in_order( + shots.size(), args.num_threads, + [&](size_t thread_index, size_t shot_index) { + if (!decoders[thread_index]) { + decoders[thread_index] = std::make_unique(config); + } + auto& decoder = *decoders[thread_index]; + auto start_time = std::chrono::high_resolution_clock::now(); + decoder.decode_shot(shots[shot_index].hits); + auto stop_time = std::chrono::high_resolution_clock::now(); + decoding_time_seconds[shot_index] = + std::chrono::duration_cast(stop_time - start_time).count() / + 1e6; + obs_predicted[shot_index] = decoder.predicted_obs_mask; + low_confidence[shot_index] = decoder.low_confidence_flag; + mass0_predicted[shot_index] = decoder.total_mass_obs0; + mass1_predicted[shot_index] = decoder.total_mass_obs1; + obs_probability_predicted[shot_index] = decoder.observable_probability(); + num_states_expanded_per_shot[shot_index] = decoder.num_states_expanded; + num_states_merged_per_shot[shot_index] = decoder.num_states_merged; + max_beam_size_per_shot[shot_index] = decoder.max_beam_size_seen; + max_frontier_width_per_shot[shot_index] = decoder.max_frontier_width_seen; + kept_state_min_per_shot[shot_index] = decoder.kept_state_min; + kept_state_median_per_shot[shot_index] = decoder.kept_state_median; + kept_state_mean_per_shot[shot_index] = decoder.kept_state_mean; + kept_state_max_per_shot[shot_index] = decoder.kept_state_max; + time_expand_per_shot[shot_index] = decoder.time_expand_seconds; + time_collapse_per_shot[shot_index] = decoder.time_collapse_seconds; + time_truncate_per_shot[shot_index] = decoder.time_truncate_seconds; + time_reconstruct_per_shot[shot_index] = decoder.time_reconstruct_seconds; + }, + [&](size_t shot_index) { + if (writer) { + writer->write_bits((uint8_t*)&obs_predicted[shot_index], num_observables); + writer->write_end(); + } + if (low_confidence[shot_index]) { + ++num_low_confidence; + } else if (has_obs && obs_predicted[shot_index] != shots[shot_index].obs_mask_as_u64()) { + ++num_errors; + } + total_time_seconds += decoding_time_seconds[shot_index]; + if (args.print_stats) { + std::cout << "num_shots = " << (shot_index + 1) + << " num_low_confidence = " << num_low_confidence + << " num_errors = " << num_errors + << " states_expanded = " << num_states_expanded_per_shot[shot_index] + << " states_merged = " << num_states_merged_per_shot[shot_index] + << " max_beam = " << max_beam_size_per_shot[shot_index] + << " frontier_width = " << max_frontier_width_per_shot[shot_index] + << " total_time_seconds = " << total_time_seconds << '\n'; + std::cout << "kept_states" << " min=" << kept_state_min_per_shot[shot_index] + << " median=" << kept_state_median_per_shot[shot_index] + << " mean=" << kept_state_mean_per_shot[shot_index] + << " max=" << kept_state_max_per_shot[shot_index] << '\n'; + std::cout << "branch_masses" << " obs0=" << mass0_predicted[shot_index] + << " obs1=" << mass1_predicted[shot_index] << '\n'; + std::cout << "phase_times_seconds" << " expand=" << time_expand_per_shot[shot_index] + << " collapse=" << time_collapse_per_shot[shot_index] + << " truncate=" << time_truncate_per_shot[shot_index] + << " reconstruct=" << time_reconstruct_per_shot[shot_index] << '\n'; + } + return !has_obs || num_errors < args.max_errors; + }); + + if (!args.obs_probs_out_fname.empty()) { + std::ofstream out(args.obs_probs_out_fname, std::ios::binary); + if (!out.is_open()) { + throw std::invalid_argument("Failed to open " + args.obs_probs_out_fname); + } + out.write(reinterpret_cast(obs_probability_predicted.data()), + static_cast(shot * sizeof(double))); + if (!out) { + throw std::runtime_error("Failed to write observable probabilities."); + } + } + + bool print_final_stats = true; + if (!args.stats_out_fname.empty()) { + nlohmann::json stats_json = {{"circuit_path", args.circuit_path}, + {"dem_path", args.dem_path}, + {"beam_width", args.beam_width}, + {"beam_eps", args.beam_eps}, + {"future_detcost_scale", args.future_detcost_scale}, + {"ranking_mode", args.ranking_mode}, + {"merge_errors", config.merge_errors}, + {"obs_probs_out", args.obs_probs_out_fname}, + {"sample_seed", args.sample_seed}, + {"sample_num_shots", args.sample_num_shots}, + {"num_threads", args.num_threads}, + {"num_low_confidence", num_low_confidence}, + {"num_shots", shot}, + {"total_time_seconds", total_time_seconds}}; + if (has_obs) { + stats_json["num_errors"] = num_errors; + } else { + stats_json["num_errors"] = nullptr; + } + if (args.stats_out_fname == "-") { + std::cout << stats_json << std::endl; + print_final_stats = false; + } else { + std::ofstream out(args.stats_out_fname, std::ofstream::out); + out << stats_json << std::endl; + } + } + + if (print_final_stats) { + std::cout << "num_shots = " << shot << " num_low_confidence = " << num_low_confidence; + if (has_obs) { + std::cout << " num_errors = " << num_errors; + } + std::cout << " total_time_seconds = " << total_time_seconds << std::endl; + } +}