From 565ae1fc728dc71a9c0f9df495bd70aa27c8bc0a Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Fri, 7 Aug 2026 23:21:28 +0000 Subject: [PATCH 1/8] feat(bp): add native C++ BP benchmark suite and complete timing metrics in bp_main --- devtools/benchmark.sh | 289 ++++++++++++++++++++++++++++++++++++++++++ src/bp_main.cc | 41 ++++-- 2 files changed, 317 insertions(+), 13 deletions(-) create mode 100755 devtools/benchmark.sh diff --git a/devtools/benchmark.sh b/devtools/benchmark.sh new file mode 100755 index 00000000..29d8ac6e --- /dev/null +++ b/devtools/benchmark.sh @@ -0,0 +1,289 @@ +#!/usr/bin/env bash +# 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. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +WORKSPACE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +cd "${WORKSPACE_ROOT}" + +# CL / Experiment tag (default: baseline) +CL_TAG="${1:-baseline}" +OUT_DIR="${WORKSPACE_ROOT}/out/${CL_TAG}" +mkdir -p "${OUT_DIR}" + +THREADS="${THREADS:-48}" +SAMPLE_SEED="${SAMPLE_SEED:-1234}" + +echo "============================================================" +echo " Running Tesseract-BP Benchmarks for [${CL_TAG}]" +echo " Output directory: ${OUT_DIR}" +echo " Threads: ${THREADS}" +echo "============================================================" + +# Build binaries using Bazel (single-core for build as per rules) +echo "=== Building C++ Binaries (bazel build --jobs=1 src:bp src:tesseract) ===" +bazel build --jobs=1 src:bp src:tesseract + +BP_BIN="${WORKSPACE_ROOT}/bazel-bin/src/bp" +TESSERACT_BIN="${WORKSPACE_ROOT}/bazel-bin/src/tesseract" + +run_bp_benchmark() { + local name="$1" + local circuit="$2" + local shots="${3:-10000}" + local norm="${4:-0.75}" + local osd_order="${5:--1}" + local osd_weight="${6:-0}" + local max_iter="${7:-30}" + local max_errors="${8:-100}" + + if [[ ! -f "${circuit}" ]]; then + echo "[SKIP] Circuit not found: ${circuit}" + return 0 + fi + + local out_json="${OUT_DIR}/${name}.json" + echo "" + echo "------------------------------------------------------------" + echo ">> Running: ${name}" + echo " Circuit: ${circuit}" + echo " Config: schedule=serial, batched=true, osd_order=${osd_order}, osd_weight=${osd_weight}, norm=${norm}" + echo "------------------------------------------------------------" + + local cmd=( + "${BP_BIN}" + --circuit "${circuit}" + --sample-num-shots "${shots}" + --sample-seed "${SAMPLE_SEED}" + --threads "${THREADS}" + --max-errors "${max_errors}" + --normalization-factor "${norm}" + --max-iter "${max_iter}" + --schedule "serial" + --batched + --print-stats + --stats-out "${out_json}" + ) + + if [[ "${osd_order}" -ge 0 ]]; then + cmd+=(--osd-order "${osd_order}" --osd-weight "${osd_weight}") + fi + + "${cmd[@]}" +} + +run_tesseract_benchmark() { + local name="$1" + local circuit="$2" + local shots="${3:-10000}" + local max_errors="${4:-100}" + + if [[ ! -f "${circuit}" ]]; then + echo "[SKIP] Circuit not found: ${circuit}" + return 0 + fi + + echo "" + echo "------------------------------------------------------------" + echo ">> Running Tesseract: ${name}" + echo " Circuit: ${circuit}" + echo "------------------------------------------------------------" + + local cmd=( + "${TESSERACT_BIN}" + --circuit "${circuit}" + --sample-num-shots "${shots}" + --sample-seed "${SAMPLE_SEED}" + --threads "${THREADS}" + --max-errors "${max_errors}" + --no-revisit-dets + --beam 20 + --beam-climbing + --num-det-orders 21 + --det-order-index + --pqlimit 1000000 + --print-stats + ) + + "${cmd[@]}" +} + +# ============================================================================== +# 1. Surface Codes (d=3, d=5, d=7, d=9) +# ============================================================================== +run_bp_benchmark \ + "surface_code_d3_p001_serial_batched" \ + "testdata/surfacecodes/r=3,d=3,p=0.001,noise=si1000,c=surface_code_Z,q=17,gates=cz.stim" \ + 100000 \ + 0.625 + +run_bp_benchmark \ + "surface_code_d5_p001_serial_batched" \ + "testdata/surfacecodes/r=5,d=5,p=0.001,noise=si1000,c=surface_code_Z,q=49,gates=cz.stim" \ + 100000 \ + 0.625 + +run_bp_benchmark \ + "surface_code_d7_p001_serial_batched" \ + "testdata/surfacecodes/r=7,d=7,p=0.001,noise=si1000,c=surface_code_Z,q=97,gates=cz.stim" \ + 50000 \ + 0.625 + +run_bp_benchmark \ + "surface_code_d9_p001_serial_batched" \ + "testdata/surfacecodes/r=9,d=9,p=0.001,noise=si1000,c=surface_code_Z,q=161,gates=cz.stim" \ + 20000 \ + 0.625 + +# ============================================================================== +# 2. Color Codes (d=5, d=7) +# ============================================================================== +run_bp_benchmark \ + "color_code_d5_superdense_serial_batched" \ + "testdata/colorcodes/r=5,d=5,p=0.001,noise=si1000,c=superdense_color_code_Z,q=37,gates=cz.stim" \ + 50000 \ + 0.9063 + +run_bp_benchmark \ + "color_code_d7_superdense_serial_batched" \ + "testdata/colorcodes/r=7,d=7,p=0.001,noise=si1000,c=superdense_color_code_Z,q=73,gates=cz.stim" \ + 20000 \ + 0.9063 + +# ============================================================================== +# 3. Bivariate Bicycle Codes (from testdata) +# ============================================================================== +run_bp_benchmark \ + "bb_72_12_6_serial_batched_osd0" \ + "testdata/bivariatebicyclecodes/r=6,d=6,p=0.001,noise=si1000,c=bivariate_bicycle_Z,nkd=[[72,12,6]],q=144,iscolored=True,A_poly=x^3+y+y^2,B_poly=y^3+x+x^2.stim" \ + 10000 \ + 0.75 \ + 10000 \ + 0 \ + 30 + +run_bp_benchmark \ + "bb_72_12_6_serial_batched_osd1" \ + "testdata/bivariatebicyclecodes/r=6,d=6,p=0.001,noise=si1000,c=bivariate_bicycle_Z,nkd=[[72,12,6]],q=144,iscolored=True,A_poly=x^3+y+y^2,B_poly=y^3+x+x^2.stim" \ + 5000 \ + 0.675 \ + 10000 \ + 1 \ + 30 + +run_bp_benchmark \ + "bb_90_8_10_serial_batched_osd0" \ + "testdata/bivariatebicyclecodes/r=10,d=10,p=0.001,noise=si1000,c=bivariate_bicycle_Z,nkd=[[90,8,10]],q=180,iscolored=True,A_poly=x^9+y+y^2,B_poly=x^7+1+x^2.stim" \ + 10000 \ + 0.75 \ + 10000 \ + 0 \ + 30 + +run_bp_benchmark \ + "bb_108_8_10_serial_batched_osd0" \ + "testdata/bivariatebicyclecodes/r=10,d=10,p=0.001,noise=si1000,c=bivariate_bicycle_Z,nkd=[[108,8,10]],q=216,iscolored=True,A_poly=x^3+y+y^2,B_poly=y^3+x+x^2.stim" \ + 10000 \ + 0.75 \ + 10000 \ + 0 \ + 30 + +run_bp_benchmark \ + "bb_144_12_12_serial_batched_osd0" \ + "testdata/bivariatebicyclecodes/r=12,d=12,p=0.001,noise=si1000,c=bivariate_bicycle_Z,nkd=[[144,12,12]],q=288,iscolored=True,A_poly=x^3+y+y^2,B_poly=y^3+x+x^2.stim" \ + 10000 \ + 0.75 \ + 10000 \ + 0 \ + 30 + +# ============================================================================== +# 4. High-Rate / CPM Codes (if present in benchmarking/hrcodes/) +# ============================================================================== +run_bp_benchmark \ + "bb_z_onebasis_serial_batched_osd0" \ + "benchmarking/hrcodes/traincodes/circuits/r=6,d=6,p=0.001,noise=si1000,c=bivariate_bicycle_Z_onebasis.stim" \ + 100000 \ + 0.75 \ + 10000 \ + 0 \ + 30 + +run_bp_benchmark \ + "cpm348_serial_batched_hard" \ + "benchmarking/hrcodes/traincodes/circuits/cpm348_p_1e-3_Z_seed_470.stim" \ + 1000 \ + 0.9063 + +run_bp_benchmark \ + "cpm564_serial_batched_osd1" \ + "benchmarking/hrcodes/traincodes/circuits/cpm564_p_1e-3_Z_seed_55_64ops.stim" \ + 100000 \ + 0.75 \ + 1000 \ + 1 \ + 1000 + +run_tesseract_benchmark \ + "cpm348_tesseract_beam" \ + "benchmarking/hrcodes/traincodes/circuits/cpm348_p_1e-3_Z_seed_470_64ops.stim" \ + 100000 + +run_bp_benchmark \ + "cpm348_serial_batched_osd1" \ + "benchmarking/hrcodes/traincodes/circuits/cpm348_p_1e-3_Z_seed_470.stim" \ + 100000 \ + 0.75 \ + 1000 \ + 1 \ + 100 + +echo "" +echo "============================================================" +echo " Summary of results in ${OUT_DIR}:" +echo "============================================================" + +python3 - <8} | {'Errors':>8} | {'Wall (s)':>9} | {'CPU (s)':>9} | {'Shots/sec':>12} | {'LER':>10} |") + print(f"|:{'-'*42}-|-{'-'*22}-|-{'-'*8}:|-{'-'*8}:|-{'-'*9}:|-{'-'*9}:|-{'-'*12}:|-{'-'*10}:|") + for jf in json_files: + name = os.path.splitext(os.path.basename(jf))[0] + try: + with open(jf, "r") as f: + d = json.load(f) + shots = d.get("num_shots", 0) + errors = d.get("num_errors", 0) + errors_str = str(errors) if errors is not None else "N/A" + wall_t = d.get("wall_time_seconds", d.get("total_time_seconds", 0.0)) + cpu_t = d.get("cpu_time_seconds", d.get("total_time_seconds", 0.0)) + th = d.get("shots_per_second", (shots / wall_t) if wall_t > 0 else 0.0) + ler_str = f"{(errors / shots):.5f}" if (errors is not None and shots > 0) else "N/A" + dec = d.get("decoder", "N/A") + print(f"| {name:<42} | {dec:<22} | {shots:>8} | {errors_str:>8} | {wall_t:>9.3f} | {cpu_t:>9.3f} | {th:>12,.1f} | {ler_str:>10} |") + except Exception as e: + print(f"| {name:<42} | Error reading JSON: {e}") +print() +EOF diff --git a/src/bp_main.cc b/src/bp_main.cc index 46e46b0f..c9372bbf 100644 --- a/src/bp_main.cc +++ b/src/bp_main.cc @@ -388,21 +388,33 @@ int main(int argc, char* argv[]) { std::chrono::duration_cast(stop_global_time - start_global_time) .count() / 1e6; + double cpu_time = total_time_seconds.load(); + size_t final_shots = processed_shots.load(); + double wall_throughput = (global_elapsed > 0) ? (double)final_shots / global_elapsed : 0.0; + double cpu_throughput = (cpu_time > 0) ? (double)final_shots / cpu_time : 0.0; std::string decoder_name = std::string(args.use_batched_bp ? "batched-" : "scalar-") + args.schedule + "-bp"; if (args.osd_order >= 0) decoder_name += "+osd"; if (!args.stats_out_fname.empty()) { - nlohmann::json stats_json = {{"circuit_path", args.circuit_path}, - {"dem_path", args.dem_path}, - {"max_errors", args.max_errors}, - {"sample_seed", args.sample_seed}, - {"total_time_seconds", global_elapsed}, - {"num_errors", has_obs ? nlohmann::json(num_errors.load()) : nullptr}, - {"num_shots", processed_shots.load()}, - {"num_discards", num_discards.load()}, - {"decoder", decoder_name}}; + nlohmann::json stats_json = { + {"circuit_path", args.circuit_path}, + {"dem_path", args.dem_path}, + {"max_errors", args.max_errors}, + {"sample_seed", args.sample_seed}, + {"num_threads", args.num_threads}, + {"sample_num_shots", args.sample_num_shots}, + {"num_shots", final_shots}, + {"num_errors", has_obs ? nlohmann::json(num_errors.load()) : nullptr}, + {"num_discards", num_discards.load()}, + {"decoder", decoder_name}, + {"total_time_seconds", cpu_time}, + {"cpu_time_seconds", cpu_time}, + {"wall_time_seconds", global_elapsed}, + {"shots_per_second", wall_throughput}, + {"shots_per_cpu_second", cpu_throughput}}; + if (args.stats_out_fname == "-") { std::cout << stats_json << std::endl; } else { @@ -413,7 +425,7 @@ int main(int argc, char* argv[]) { if (!args.sinter_csv_out.empty()) { std::stringstream csv_line; - csv_line << processed_shots.load() << ","; + csv_line << final_shots << ","; if (has_obs) { csv_line << num_errors.load(); } @@ -428,14 +440,17 @@ int main(int argc, char* argv[]) { } } - if (args.stats_out_fname.empty() && args.sinter_csv_out.empty()) { - std::cout << "num_shots = " << processed_shots.load(); + if (args.print_stats || (args.stats_out_fname.empty() && args.sinter_csv_out.empty())) { + std::cout << "num_shots = " << final_shots; if (has_obs) { std::cout << " num_errors = " << num_errors.load(); } else { std::cout << " num_errors = N/A"; } - std::cout << " total_time_seconds = " << global_elapsed << std::endl; + std::cout << " wall_time_seconds = " << global_elapsed + << " total_time_seconds = " << cpu_time + << " (wall shots/sec = " << std::fixed << std::setprecision(1) << wall_throughput << ")" + << std::endl; } return 0; From 640f32cc3e49ebd69a614392b3c9d26cfe6121e3 Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Sat, 8 Aug 2026 00:17:42 +0000 Subject: [PATCH 2/8] build: add CMake build directory to .gitignore --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 9f771e10..c5823e20 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,6 @@ MODULE.bazel.lock # Python build artifacts *.so *.egg-info/ + +# Local CMake build directory +build/ From bab66168fa6a43afc82522a543d69023a376159a Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Fri, 7 Aug 2026 23:21:33 +0000 Subject: [PATCH 3/8] feat(bp): transition BP decoder to flat interleaved 1D array layout --- src/bp/batched_bp_parallel_min_sum.h | 6 +++--- src/bp/batched_bp_parallel_min_sum.inl | 9 +++++---- src/bp/batched_bp_parallel_min_sum.test.cc | 4 ++-- src/bp/batched_bp_serial_min_sum.h | 2 +- src/bp/batched_bp_serial_min_sum.test.cc | 22 +++++++++++----------- src/bp/tesseract_bp_decoder.cc | 14 +++++++++----- 6 files changed, 31 insertions(+), 26 deletions(-) diff --git a/src/bp/batched_bp_parallel_min_sum.h b/src/bp/batched_bp_parallel_min_sum.h index bc145cbe..b08dc39c 100644 --- a/src/bp/batched_bp_parallel_min_sum.h +++ b/src/bp/batched_bp_parallel_min_sum.h @@ -14,8 +14,8 @@ namespace bp { // graph: The BatchedTannerGraph to run the decoder on. // detection_events_batch: A 2D vector [BATCH_SIZE][num_events]. The syndromes // for all shots in the current batch. -// posteriors_batch: A 2D vector [BATCH_SIZE][num_variables] to store the -// final posterior LLRs of the variable nodes for all shots. +// posteriors_flat: A 1D flat interleaved array of size [num_variables * BATCH_SIZE] +// where variable v and shot b is stored at index [v * BATCH_SIZE + b]. // max_iters: The maximum number of iterations to run the decoder for. // normalization_factor: The alpha factor for Normalized Min-Sum. // stop_at_convergence: If true, the decoder will stop early for specific @@ -27,7 +27,7 @@ namespace bp { template std::vector batched_bp_parallel_min_sum( BatchedTannerGraph& graph, const std::vector>& detection_events_batch, - std::vector>& posteriors_batch, size_t max_iters, float normalization_factor, + std::vector& posteriors_flat, size_t max_iters, float normalization_factor, bool stop_at_convergence = true); } // namespace bp diff --git a/src/bp/batched_bp_parallel_min_sum.inl b/src/bp/batched_bp_parallel_min_sum.inl index 34c11b61..18548e96 100644 --- a/src/bp/batched_bp_parallel_min_sum.inl +++ b/src/bp/batched_bp_parallel_min_sum.inl @@ -10,7 +10,7 @@ namespace bp { template std::vector batched_bp_parallel_min_sum( BatchedTannerGraph& graph, const std::vector>& detection_events_batch, - std::vector>& posteriors_batch, size_t max_iters, float normalization_factor, + std::vector& posteriors_flat, size_t max_iters, float normalization_factor, bool stop_at_convergence) { size_t actual_batch_size = detection_events_batch.size(); if (actual_batch_size > BP_BATCH_SIZE) { @@ -145,11 +145,12 @@ std::vector batched_bp_parallel_min_sum( // --- Posterior and Convergence Check --- if (stop_at_convergence || (iter == max_iters - 1)) { - // Calculate current posteriors for active shots + // Calculate current posteriors for active shots in flat layout for (size_t i = 0; i < graph.num_variables; ++i) { size_t start = graph.var_edge_offsets[i]; size_t end = graph.var_edge_offsets[i + 1]; T prior_val = graph.priors[i]; + size_t var_post_idx = i * BP_BATCH_SIZE; #pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { @@ -158,7 +159,7 @@ std::vector batched_bp_parallel_min_sum( for (size_t e = start; e < end; ++e) { post += graph.check_to_var_messages[e * BP_BATCH_SIZE + b]; } - posteriors_batch[b][i] = post; + posteriors_flat[var_post_idx + b] = post; } } @@ -178,7 +179,7 @@ std::vector batched_bp_parallel_min_sum( uint8_t posterior_parity = 0; for (size_t e = start; e < end; ++e) { size_t v_idx = graph.check_edges[e]; - if (posteriors_batch[b][v_idx] < 0) { + if (posteriors_flat[v_idx * BP_BATCH_SIZE + b] < 0) { posterior_parity ^= 1; } } diff --git a/src/bp/batched_bp_parallel_min_sum.test.cc b/src/bp/batched_bp_parallel_min_sum.test.cc index b292def0..964cb258 100644 --- a/src/bp/batched_bp_parallel_min_sum.test.cc +++ b/src/bp/batched_bp_parallel_min_sum.test.cc @@ -34,7 +34,7 @@ TEST(BatchedBpParallelMinSumTest, MatchesUnbatchedImplementation) { batched_syndromes.push_back({0, 3}); // Shot 1: Error at V0, V4 for (size_t i = 2; i < BP_BATCH_SIZE; i++) batched_syndromes.push_back({}); // Other shots empty - std::vector> batched_posteriors(BP_BATCH_SIZE, std::vector(5, 0)); + std::vector batched_posteriors(5 * BP_BATCH_SIZE, 0.0f); auto batched_results = batched_bp_parallel_min_sum( batched_graph, batched_syndromes, batched_posteriors, 20, kNormalizationFactor, true); @@ -51,7 +51,7 @@ TEST(BatchedBpParallelMinSumTest, MatchesUnbatchedImplementation) { << "Shot " << b << " iter mismatch."; for (size_t i = 0; i < 5; ++i) { - EXPECT_FLOAT_EQ(batched_posteriors[b][i], unbatched_posteriors[i]) + EXPECT_FLOAT_EQ(batched_posteriors[i * BP_BATCH_SIZE + b], unbatched_posteriors[i]) << "Shot " << b << " variable " << i << " mismatch."; } } diff --git a/src/bp/batched_bp_serial_min_sum.h b/src/bp/batched_bp_serial_min_sum.h index 7825fc07..f1217960 100644 --- a/src/bp/batched_bp_serial_min_sum.h +++ b/src/bp/batched_bp_serial_min_sum.h @@ -11,7 +11,7 @@ namespace bp { template std::vector batched_bp_serial_min_sum( BatchedTannerGraph& graph, const std::vector>& detection_events_batch, - std::vector>& posteriors_batch, size_t max_iters, float normalization_factor, + std::vector& posteriors_flat, size_t max_iters, float normalization_factor, bool stop_at_convergence = true); } // namespace bp diff --git a/src/bp/batched_bp_serial_min_sum.test.cc b/src/bp/batched_bp_serial_min_sum.test.cc index 8b2dd736..bf760b5c 100644 --- a/src/bp/batched_bp_serial_min_sum.test.cc +++ b/src/bp/batched_bp_serial_min_sum.test.cc @@ -33,26 +33,26 @@ TEST(BatchedBpSerialMinSumTest, ConvergesOnSimpleErrors) { batched_syndromes.push_back({0, 3}); // Shot 1: Error at V0, V4 for (size_t i = 2; i < BP_BATCH_SIZE; i++) batched_syndromes.push_back({}); // Other shots empty - std::vector> batched_posteriors(BP_BATCH_SIZE, std::vector(5, 0)); + std::vector batched_posteriors(5 * BP_BATCH_SIZE, 0.0f); auto batched_results = batched_bp_serial_min_sum( batched_graph, batched_syndromes, batched_posteriors, 20, kNormalizationFactor, true); // Shot 0 should converge to fixing V2 EXPECT_TRUE(batched_results[0].converged); - EXPECT_GT(batched_posteriors[0][0], 0); - EXPECT_GT(batched_posteriors[0][1], 0); - EXPECT_LT(batched_posteriors[0][2], 0); // V2 is negative (error) - EXPECT_GT(batched_posteriors[0][3], 0); - EXPECT_GT(batched_posteriors[0][4], 0); + EXPECT_GT(batched_posteriors[0 * BP_BATCH_SIZE + 0], 0); + EXPECT_GT(batched_posteriors[1 * BP_BATCH_SIZE + 0], 0); + EXPECT_LT(batched_posteriors[2 * BP_BATCH_SIZE + 0], 0); // V2 is negative (error) + EXPECT_GT(batched_posteriors[3 * BP_BATCH_SIZE + 0], 0); + EXPECT_GT(batched_posteriors[4 * BP_BATCH_SIZE + 0], 0); // Shot 1 should converge to fixing V0, V4 EXPECT_TRUE(batched_results[1].converged); - EXPECT_LT(batched_posteriors[1][0], 0); // V0 is negative - EXPECT_GT(batched_posteriors[1][1], 0); - EXPECT_GT(batched_posteriors[1][2], 0); - EXPECT_GT(batched_posteriors[1][3], 0); - EXPECT_LT(batched_posteriors[1][4], 0); // V4 is negative + EXPECT_LT(batched_posteriors[0 * BP_BATCH_SIZE + 1], 0); // V0 is negative + EXPECT_GT(batched_posteriors[1 * BP_BATCH_SIZE + 1], 0); + EXPECT_GT(batched_posteriors[2 * BP_BATCH_SIZE + 1], 0); + EXPECT_GT(batched_posteriors[3 * BP_BATCH_SIZE + 1], 0); + EXPECT_LT(batched_posteriors[4 * BP_BATCH_SIZE + 1], 0); // V4 is negative } } // namespace bp \ No newline at end of file diff --git a/src/bp/tesseract_bp_decoder.cc b/src/bp/tesseract_bp_decoder.cc index fdff0957..c912b045 100644 --- a/src/bp/tesseract_bp_decoder.cc +++ b/src/bp/tesseract_bp_decoder.cc @@ -119,9 +119,10 @@ std::vector> TesseractBpDecoder::decode_batch( size_t num_shots = detection_events_batch.size(); std::vector> results(num_shots); + size_t num_vars = graph_.variable_nodes.size(); std::vector> current_batch; - std::vector> posteriors_batch( - BP_BATCH_SIZE, std::vector(graph_.variable_nodes.size())); + alignas(64) std::vector posteriors_flat(num_vars * BP_BATCH_SIZE); + std::vector shot_posteriors(num_vars); for (size_t shot = 0; shot < num_shots; ++shot) { current_batch.push_back(std::vector(detection_events_batch[shot].begin(), @@ -135,19 +136,22 @@ std::vector> TesseractBpDecoder::decode_batch( std::vector bp_results; if (params_.schedule == "serial") { - bp_results = batched_bp_serial_min_sum(batched_graph_, current_batch, posteriors_batch, + bp_results = batched_bp_serial_min_sum(batched_graph_, current_batch, posteriors_flat, params_.max_iter, params_.normalization_factor, params_.stop_at_convergence); } else { - bp_results = batched_bp_parallel_min_sum(batched_graph_, current_batch, posteriors_batch, + bp_results = batched_bp_parallel_min_sum(batched_graph_, current_batch, posteriors_flat, params_.max_iter, params_.normalization_factor, params_.stop_at_convergence); } for (size_t b = 0; b < actual_size; ++b) { size_t real_shot_idx = shot - actual_size + 1 + b; + for (size_t v = 0; v < num_vars; ++v) { + shot_posteriors[v] = posteriors_flat[v * BP_BATCH_SIZE + b]; + } results[real_shot_idx] = - post_processor->process(bp_results[b], posteriors_batch[b], + post_processor->process(bp_results[b], shot_posteriors, detection_events_batch[real_shot_idx], hyperedge_observables_); } current_batch.clear(); From b71c8faa328a02c827d40b4983cb948829e30962 Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Fri, 7 Aug 2026 23:21:38 +0000 Subject: [PATCH 4/8] feat(bp): implement horizontal layered serial scheduling with immediate posterior update --- src/bp/batched_bp_serial_min_sum.inl | 220 ++++++++------------------- src/bp/bp_serial_min_sum.inl | 97 ++++++------ 2 files changed, 109 insertions(+), 208 deletions(-) diff --git a/src/bp/batched_bp_serial_min_sum.inl b/src/bp/batched_bp_serial_min_sum.inl index 5e601dc6..a577cbc7 100644 --- a/src/bp/batched_bp_serial_min_sum.inl +++ b/src/bp/batched_bp_serial_min_sum.inl @@ -10,7 +10,7 @@ namespace bp { template std::vector batched_bp_serial_min_sum( BatchedTannerGraph& graph, const std::vector>& detection_events_batch, - std::vector>& posteriors_batch, size_t max_iters, float normalization_factor, + std::vector& posteriors_flat, size_t max_iters, float normalization_factor, bool stop_at_convergence) { size_t actual_batch_size = detection_events_batch.size(); if (actual_batch_size > BP_BATCH_SIZE) { @@ -35,126 +35,90 @@ std::vector batched_bp_serial_min_sum( using T_MAG = typename llr_traits::magnitude_type; const T_MAG max_mag = std::numeric_limits::max(); - // --- State Tracking Arrays --- - // For each check node c, and batch b, we track min1, min2, min1_idx, and sign_prod. - std::vector check_min1(graph.num_checks * BP_BATCH_SIZE, max_mag); - std::vector check_min2(graph.num_checks * BP_BATCH_SIZE, max_mag); - std::vector check_min1_idx(graph.num_checks * BP_BATCH_SIZE, SIZE_MAX); - std::vector check_sign_prod(graph.num_checks * BP_BATCH_SIZE, 0); + // Initialize check_to_var messages to 0 + std::fill(graph.check_to_var_messages.begin(), graph.check_to_var_messages.end(), 0); - // Initialization: var_to_check messages are set to prior values + // Initialize posteriors_flat to priors for (size_t i = 0; i < graph.num_variables; ++i) { - size_t start = graph.var_edge_offsets[i]; - size_t end = graph.var_edge_offsets[i + 1]; T prior_val = graph.priors[i]; - - for (size_t e = start; e < end; ++e) { - size_t c_e = graph.var_edge_rev_indices[e]; - size_t out_idx = c_e * BP_BATCH_SIZE; + size_t var_post_idx = i * BP_BATCH_SIZE; #pragma GCC ivdep - for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - graph.var_to_check_messages[out_idx + b] = prior_val; - } + for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { + posteriors_flat[var_post_idx + b] = prior_val; } } - // Initial full scan to populate check node states - for (size_t i = 0; i < graph.num_checks; ++i) { - size_t start = graph.check_edge_offsets[i]; - size_t end = graph.check_edge_offsets[i + 1]; - size_t state_idx = i * BP_BATCH_SIZE; + size_t iter = 0; + for (iter = 0; iter < max_iters && num_active > 0; ++iter) { + // Horizontal / Layered Schedule: Iterate through check nodes + for (size_t c = 0; c < graph.num_checks; ++c) { + size_t start = graph.check_edge_offsets[c]; + size_t end = graph.check_edge_offsets[c + 1]; + size_t deg = end - start; + if (deg == 0) continue; + + size_t syn_idx = c * BP_BATCH_SIZE; + + T_MAG min1[BP_BATCH_SIZE]; + T_MAG min2[BP_BATCH_SIZE]; + size_t min1_idx[BP_BATCH_SIZE]; + uint8_t total_sign_prod[BP_BATCH_SIZE]; - for (size_t c_e = start; c_e < end; ++c_e) { - size_t msg_idx = c_e * BP_BATCH_SIZE; #pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - T msg = graph.var_to_check_messages[msg_idx + b]; - T_MAG mag = (T_MAG)std::abs(msg); - uint8_t sign = (msg < 0) ? 1 : 0; - - check_sign_prod[state_idx + b] ^= sign; - - if (mag < check_min1[state_idx + b]) { - check_min2[state_idx + b] = check_min1[state_idx + b]; - check_min1[state_idx + b] = mag; - check_min1_idx[state_idx + b] = c_e; - } else if (mag < check_min2[state_idx + b]) { - check_min2[state_idx + b] = mag; - } + min1[b] = max_mag; + min2[b] = max_mag; + min1_idx[b] = SIZE_MAX; + total_sign_prod[b] = 0; } - } - } - // Helper lambda to do a full rescan of a specific check node for the entire batch. - auto rescan_check_node_batched = [&](size_t c_idx, const uint8_t* needs_rescan) { - size_t start = graph.check_edge_offsets[c_idx]; - size_t end = graph.check_edge_offsets[c_idx + 1]; - size_t state_idx = c_idx * BP_BATCH_SIZE; + // Pass 1: Compute variable-to-check extrinsic messages Q_{c,v} = L_v - R_{c,v} + // and find min1, min2, and total sign product. + for (size_t e = start; e < end; ++e) { + size_t v = graph.check_edges[e]; + size_t v_idx = v * BP_BATCH_SIZE; + size_t msg_idx = e * BP_BATCH_SIZE; #pragma GCC ivdep - for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - if (needs_rescan[b]) { - check_min1[state_idx + b] = max_mag; - check_min2[state_idx + b] = max_mag; - check_min1_idx[state_idx + b] = SIZE_MAX; - } - } + for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { + if (!active_shots[b]) continue; - for (size_t c_e = start; c_e < end; ++c_e) { - size_t msg_idx = c_e * BP_BATCH_SIZE; -#pragma GCC ivdep - for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - if (needs_rescan[b]) { - T msg = graph.var_to_check_messages[msg_idx + b]; - T_MAG mag = (T_MAG)std::abs(msg); - if (mag < check_min1[state_idx + b]) { - check_min2[state_idx + b] = check_min1[state_idx + b]; - check_min1[state_idx + b] = mag; - check_min1_idx[state_idx + b] = c_e; - } else if (mag < check_min2[state_idx + b]) { - check_min2[state_idx + b] = mag; - } - } - } - } - }; + T old_r = graph.check_to_var_messages[msg_idx + b]; + T q_msg = posteriors_flat[v_idx + b] - old_r; - size_t iter = 0; - for (iter = 0; iter < max_iters && num_active > 0; ++iter) { - // Serial Schedule: Iterate through variable nodes - for (size_t i = 0; i < graph.num_variables; ++i) { - size_t start = graph.var_edge_offsets[i]; - size_t end = graph.var_edge_offsets[i + 1]; - T prior_val = graph.priors[i]; + T_MAG mag = (T_MAG)std::abs(q_msg); + uint8_t sign = (q_msg < 0) ? 1 : 0; + total_sign_prod[b] ^= sign; -// 1. Fetch C->V messages and compute new posterior in O(1) time per edge -#pragma GCC ivdep - for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - if (active_shots[b]) { - posteriors_batch[b][i] = prior_val; + if (mag < min1[b]) { + min2[b] = min1[b]; + min1[b] = mag; + min1_idx[b] = e; + } else if (mag < min2[b]) { + min2[b] = mag; + } } } + // Pass 2: Compute new check-to-variable message R'_{c,v} + // and immediately update posterior L'_v = Q_{c,v} + R'_{c,v} = L_v + (R'_{c,v} - R_{c,v}) for (size_t e = start; e < end; ++e) { - size_t c_idx = graph.var_edges[e]; - size_t c_e = - graph.var_edge_rev_indices[e]; // The index of this edge in the check node's list - size_t state_idx = c_idx * BP_BATCH_SIZE; - size_t syn_idx = c_idx * BP_BATCH_SIZE; - size_t msg_idx = c_e * BP_BATCH_SIZE; + size_t v = graph.check_edges[e]; + size_t v_idx = v * BP_BATCH_SIZE; + size_t msg_idx = e * BP_BATCH_SIZE; #pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { if (!active_shots[b]) continue; - T_MAG min_mag = (check_min1_idx[state_idx + b] == c_e) ? check_min2[state_idx + b] - : check_min1[state_idx + b]; + T old_r = graph.check_to_var_messages[msg_idx + b]; + T q_msg = posteriors_flat[v_idx + b] - old_r; - T my_msg = graph.var_to_check_messages[msg_idx + b]; - uint8_t my_sign = (my_msg < 0) ? 1 : 0; - uint8_t extrinsic_sign = check_sign_prod[state_idx + b] ^ my_sign; + uint8_t q_sign = (q_msg < 0) ? 1 : 0; + uint8_t extrinsic_sign = total_sign_prod[b] ^ q_sign; uint8_t final_sign = batched_syndromes[syn_idx + b] ^ extrinsic_sign; + T_MAG min_mag = (e == min1_idx[b]) ? min2[b] : min1[b]; T_MAG normalized_mag = min_mag * normalization_factor; T final_mag; if constexpr (std::is_integral_v) { @@ -164,70 +128,12 @@ std::vector batched_bp_serial_min_sum( final_mag = normalized_mag; } - T c2v_msg = final_sign ? -final_mag : final_mag; - posteriors_batch[b][i] += c2v_msg; - - // Store this for step 2 so we don't recompute it - graph.check_to_var_messages[e * BP_BATCH_SIZE + b] = c2v_msg; - } - } - - // 2. Compute new V->C messages and update Check Node states - for (size_t e = start; e < end; ++e) { - size_t c_idx = graph.var_edges[e]; - size_t c_e = graph.var_edge_rev_indices[e]; - size_t state_idx = c_idx * BP_BATCH_SIZE; - size_t msg_idx = c_e * BP_BATCH_SIZE; - - uint8_t needs_rescan[BP_BATCH_SIZE] = {0}; - bool any_needs_rescan = false; - -#pragma GCC ivdep - for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - if (!active_shots[b]) continue; - - T old_msg = graph.var_to_check_messages[msg_idx + b]; - T new_msg = posteriors_batch[b][i] - graph.check_to_var_messages[e * BP_BATCH_SIZE + b]; - - if (old_msg == new_msg) continue; - - graph.var_to_check_messages[msg_idx + b] = new_msg; - - T_MAG old_mag = (T_MAG)std::abs(old_msg); - T_MAG new_mag = (T_MAG)std::abs(new_msg); - uint8_t old_sign = (old_msg < 0) ? 1 : 0; - uint8_t new_sign = (new_msg < 0) ? 1 : 0; - - // Incremental sign update - check_sign_prod[state_idx + b] ^= (old_sign ^ new_sign); - - // Incremental magnitude update - if (new_mag < check_min1[state_idx + b]) { - // New absolute minimum found! - if (check_min1_idx[state_idx + b] != c_e) { - check_min2[state_idx + b] = check_min1[state_idx + b]; - } - check_min1[state_idx + b] = new_mag; - check_min1_idx[state_idx + b] = c_e; - } else if (new_mag < check_min2[state_idx + b] && check_min1_idx[state_idx + b] != c_e) { - // New second minimum found! - check_min2[state_idx + b] = new_mag; - } else if (check_min1_idx[state_idx + b] == c_e && new_mag > old_mag) { - // The minimum grew. We must rescan to find the new true minimums. - needs_rescan[b] = 1; - any_needs_rescan = true; - } else if (old_mag == check_min2[state_idx + b] && new_mag > old_mag) { - // The second minimum grew. We must rescan. - needs_rescan[b] = 1; - any_needs_rescan = true; - } - } - - if (any_needs_rescan) { - rescan_check_node_batched(c_idx, needs_rescan); + T new_r = final_sign ? -final_mag : final_mag; + graph.check_to_var_messages[msg_idx + b] = new_r; + posteriors_flat[v_idx + b] = q_msg + new_r; } } - } // End of variable loop + } // End of check loop // --- Convergence Check (End of iteration) --- if (stop_at_convergence || (iter == max_iters - 1)) { @@ -240,12 +146,12 @@ std::vector batched_bp_serial_min_sum( for (size_t b = 0; b < actual_batch_size; ++b) { if (!active_shots[b]) continue; - if (!shot_converged[b]) continue; // Already failed a check + if (!shot_converged[b]) continue; uint8_t posterior_parity = 0; for (size_t e = start; e < end; ++e) { size_t v_idx = graph.check_edges[e]; - if (posteriors_batch[b][v_idx] < 0) { + if (posteriors_flat[v_idx * BP_BATCH_SIZE + b] < 0) { posterior_parity ^= 1; } } @@ -271,7 +177,7 @@ std::vector batched_bp_serial_min_sum( for (size_t b = 0; b < actual_batch_size; ++b) { if (active_shots[b]) { results[b].converged = false; - results[b].num_iters = iter; // `iter` will be `max_iters` here + results[b].num_iters = iter; } } diff --git a/src/bp/bp_serial_min_sum.inl b/src/bp/bp_serial_min_sum.inl index 4ef18044..d7591144 100644 --- a/src/bp/bp_serial_min_sum.inl +++ b/src/bp/bp_serial_min_sum.inl @@ -13,17 +13,11 @@ BPResult bp_serial_min_sum(TannerGraph& graph, const std::vector& det bool stop_at_convergence) { graph.add_detection_events(detection_events); - // Initialization: var_to_check messages are set to prior values + // Initialize posteriors to priors for (size_t i = 0; i < graph.variable_nodes.size(); ++i) { - size_t start = graph.var_edge_offsets[i]; - size_t end = graph.var_edge_offsets[i + 1]; - T prior_val = graph.variable_nodes[i].prior; - - for (size_t e = start; e < end; ++e) { - size_t c_e = graph.var_edge_rev_indices[e]; - graph.var_to_check_messages[c_e] = prior_val; - } + graph.variable_nodes[i].posterior = graph.variable_nodes[i].prior; } + std::fill(graph.check_to_var_messages.begin(), graph.check_to_var_messages.end(), 0); using T_MAG = typename llr_traits::magnitude_type; const T_MAG max_mag = std::numeric_limits::max(); @@ -31,38 +25,49 @@ BPResult bp_serial_min_sum(TannerGraph& graph, const std::vector& det bool has_converged = false; size_t iter = 0; for (iter = 0; iter < max_iters; ++iter) { - // Serial Schedule: Iterate through variable nodes - for (size_t i = 0; i < graph.variable_nodes.size(); ++i) { - auto& variable = graph.variable_nodes[i]; - size_t start = graph.var_edge_offsets[i]; - size_t end = graph.var_edge_offsets[i + 1]; - - // 1. Fetch check-to-var messages dynamically for this variable + // Horizontal / Layered Schedule: Iterate through check nodes + for (size_t c = 0; c < graph.check_nodes.size(); ++c) { + auto& check = graph.check_nodes[c]; + size_t start = graph.check_edge_offsets[c]; + size_t end = graph.check_edge_offsets[c + 1]; + size_t deg = end - start; + if (deg == 0) continue; + + T_MAG min1 = max_mag; + T_MAG min2 = max_mag; + size_t min1_idx = SIZE_MAX; + uint8_t total_sign_prod = 0; + + // Pass 1: Compute Q_{c,v} = L_v - R_{c,v} and find min1, min2, and sign product for (size_t e = start; e < end; ++e) { - size_t c_idx = graph.var_edges[e]; - auto& check = graph.check_nodes[c_idx]; - size_t c_start = graph.check_edge_offsets[c_idx]; - size_t c_end = graph.check_edge_offsets[c_idx + 1]; - - T_MAG min_mag = max_mag; - uint8_t sign_prod = 0; - - // Compute min1 and product of signs for all neighbors of check c EXCEPT variable i - for (size_t c_e = c_start; c_e < c_end; ++c_e) { - if (graph.check_edges[c_e] == i) continue; // Skip self - - T msg = graph.var_to_check_messages[c_e]; - T_MAG mag = (T_MAG)std::abs(msg); - if (mag < min_mag) { - min_mag = mag; - } - uint8_t sign = (msg < 0) ? 1 : 0; - sign_prod ^= sign; + size_t v = graph.check_edges[e]; + T old_r = graph.check_to_var_messages[e]; + T q_msg = graph.variable_nodes[v].posterior - old_r; + + T_MAG mag = (T_MAG)std::abs(q_msg); + uint8_t sign = (q_msg < 0) ? 1 : 0; + total_sign_prod ^= sign; + + if (mag < min1) { + min2 = min1; + min1 = mag; + min1_idx = e; + } else if (mag < min2) { + min2 = mag; } + } - // Compute final check-to-var message - uint8_t final_sign = check.syndrome ^ sign_prod; + // Pass 2: Compute R'_{c,v} and immediately update L_v = Q_{c,v} + R'_{c,v} + for (size_t e = start; e < end; ++e) { + size_t v = graph.check_edges[e]; + T old_r = graph.check_to_var_messages[e]; + T q_msg = graph.variable_nodes[v].posterior - old_r; + + uint8_t q_sign = (q_msg < 0) ? 1 : 0; + uint8_t extrinsic_sign = total_sign_prod ^ q_sign; + uint8_t final_sign = check.syndrome ^ extrinsic_sign; + T_MAG min_mag = (e == min1_idx) ? min2 : min1; T_MAG normalized_mag = min_mag * normalization_factor; T final_mag; if constexpr (std::is_integral_v) { @@ -72,21 +77,11 @@ BPResult bp_serial_min_sum(TannerGraph& graph, const std::vector& det final_mag = normalized_mag; } - graph.check_to_var_messages[e] = final_sign ? -final_mag : final_mag; + T new_r = final_sign ? -final_mag : final_mag; + graph.check_to_var_messages[e] = new_r; + graph.variable_nodes[v].posterior = sat_add(q_msg, new_r); } - - // 2. Compute updated posterior for the current variable node - variable.posterior = variable.prior; - for (size_t e = start; e < end; ++e) { - variable.posterior = sat_add(variable.posterior, graph.check_to_var_messages[e]); - } - - // 3. Update var-to-check messages for the current variable node - for (size_t e = start; e < end; ++e) { - size_t c_e = graph.var_edge_rev_indices[e]; - graph.var_to_check_messages[c_e] = variable.posterior - graph.check_to_var_messages[e]; - } - } // End of variable loop + } // --- Convergence Check (End of iteration) --- if (stop_at_convergence || (iter == max_iters - 1)) { From c900f955cb1b0df3093c9a5af87bafc3cfb3abc4 Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Fri, 7 Aug 2026 23:30:43 +0000 Subject: [PATCH 5/8] feat(bp): vectorize layered serial min-sum with AVX-512 intrinsics --- .bazelrc | 3 + src/bp/batched_bp_serial_min_sum.inl | 191 ++++++++++++++++++++++++--- 2 files changed, 173 insertions(+), 21 deletions(-) diff --git a/.bazelrc b/.bazelrc index c46131dd..33963cca 100644 --- a/.bazelrc +++ b/.bazelrc @@ -10,4 +10,7 @@ build --cxxopt='-std=c++20' # We keep it linux-only so macos does not receive GNU ld flags build:linux --copt=-ffunction-sections build:linux --copt=-fdata-sections +build:linux --copt=-mavx512f +build:linux --copt=-mavx512bw +build:linux --copt=-mavx512dq build:linux --linkopt=-Wl,--gc-sections diff --git a/src/bp/batched_bp_serial_min_sum.inl b/src/bp/batched_bp_serial_min_sum.inl index a577cbc7..c4d2f7c0 100644 --- a/src/bp/batched_bp_serial_min_sum.inl +++ b/src/bp/batched_bp_serial_min_sum.inl @@ -3,6 +3,10 @@ #include #include +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__) +#include +#endif + #include "bp/batched_bp_serial_min_sum.h" namespace bp { @@ -17,24 +21,21 @@ std::vector batched_bp_serial_min_sum( throw std::invalid_argument("Provided batch size exceeds BP_BATCH_SIZE"); } - // Set up batched syndromes. - std::vector batched_syndromes(graph.num_checks * BP_BATCH_SIZE, 0); + // Set up batched syndromes as 16-bit masks per check node. + std::vector check_syndrome_masks(graph.num_checks, 0); for (size_t b = 0; b < actual_batch_size; ++b) { for (size_t d : detection_events_batch[b]) { - batched_syndromes[d * BP_BATCH_SIZE + b] = 1; + check_syndrome_masks[d] |= (static_cast(1) << b); } } - // Track which shots in the batch are still active - std::vector active_shots(BP_BATCH_SIZE, 0); - for (size_t b = 0; b < actual_batch_size; ++b) active_shots[b] = 1; + uint16_t active_mask = (actual_batch_size == 16) + ? 0xFFFF + : static_cast((1U << actual_batch_size) - 1); size_t num_active = actual_batch_size; std::vector results(actual_batch_size, {false, 0}); - using T_MAG = typename llr_traits::magnitude_type; - const T_MAG max_mag = std::numeric_limits::max(); - // Initialize check_to_var messages to 0 std::fill(graph.check_to_var_messages.begin(), graph.check_to_var_messages.end(), 0); @@ -42,15 +43,172 @@ std::vector batched_bp_serial_min_sum( for (size_t i = 0; i < graph.num_variables; ++i) { T prior_val = graph.priors[i]; size_t var_post_idx = i * BP_BATCH_SIZE; -#pragma GCC ivdep +#if defined(__AVX512F__) + if constexpr (std::is_same_v) { + __m512i v_prior = _mm512_set1_epi32(prior_val); + _mm512_storeu_si512(reinterpret_cast<__m512i*>(&posteriors_flat[var_post_idx]), v_prior); + } else { + for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { + posteriors_flat[var_post_idx + b] = prior_val; + } + } +#else for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { posteriors_flat[var_post_idx + b] = prior_val; } +#endif + } + +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__) + if constexpr (std::is_same_v) { + const __m512 v_norm = _mm512_set1_ps(normalization_factor); + const __m512i v_zero = _mm512_setzero_si512(); + + size_t iter = 0; + for (iter = 0; iter < max_iters && active_mask != 0; ++iter) { + // Horizontal / Layered Schedule: Iterate through check nodes + for (size_t c = 0; c < graph.num_checks; ++c) { + size_t start = graph.check_edge_offsets[c]; + size_t end = graph.check_edge_offsets[c + 1]; + size_t deg = end - start; + if (deg == 0) continue; + + uint16_t syn_mask = check_syndrome_masks[c]; + + __m512i min1 = _mm512_set1_epi32(std::numeric_limits::max()); + __m512i min2 = _mm512_set1_epi32(std::numeric_limits::max()); + __m512i min1_idx = _mm512_set1_epi32(-1); + uint16_t total_sign_mask = 0; + + // Pass 1: Compute variable-to-check extrinsic messages Q_{c,v} = L_v - R_{c,v} + // and find min1, min2, and total sign product across 16 SIMD lanes. + for (size_t e = start; e < end; ++e) { + size_t v = graph.check_edges[e]; + size_t v_idx = v * BP_BATCH_SIZE; + size_t msg_idx = e * BP_BATCH_SIZE; + + __m512i* v_ptr = reinterpret_cast<__m512i*>(&posteriors_flat[v_idx]); + __m512i* msg_ptr = reinterpret_cast<__m512i*>(&graph.check_to_var_messages[msg_idx]); + + __m512i v_post = _mm512_loadu_si512(v_ptr); + __m512i old_r = _mm512_loadu_si512(msg_ptr); + __m512i q_msg = _mm512_sub_epi32(v_post, old_r); + + __m512i mag = _mm512_abs_epi32(q_msg); + __mmask16 q_sign = _mm512_movepi32_mask(q_msg); + total_sign_mask ^= static_cast(q_sign); + + __mmask16 is_less_min1 = _mm512_cmplt_epi32_mask(mag, min1); + __mmask16 is_less_min2 = _mm512_cmplt_epi32_mask(mag, min2); + + // If mag < min1: min2 = min1, min1 = mag, min1_idx = e + min2 = _mm512_mask_blend_epi32(is_less_min1, min2, min1); + min1 = _mm512_mask_blend_epi32(is_less_min1, min1, mag); + min1_idx = _mm512_mask_blend_epi32(is_less_min1, min1_idx, _mm512_set1_epi32(static_cast(e))); + + // Else if mag < min2: min2 = mag + __mmask16 update_min2 = is_less_min2 & ~is_less_min1; + min2 = _mm512_mask_blend_epi32(update_min2, min2, mag); + } + + // Pass 2: Compute new check-to-variable message R'_{c,v} + // and immediately update posterior L'_v = Q_{c,v} + R'_{c,v} + for (size_t e = start; e < end; ++e) { + size_t v = graph.check_edges[e]; + size_t v_idx = v * BP_BATCH_SIZE; + size_t msg_idx = e * BP_BATCH_SIZE; + + __m512i* v_ptr = reinterpret_cast<__m512i*>(&posteriors_flat[v_idx]); + __m512i* msg_ptr = reinterpret_cast<__m512i*>(&graph.check_to_var_messages[msg_idx]); + + __m512i v_post = _mm512_loadu_si512(v_ptr); + __m512i old_r = _mm512_loadu_si512(msg_ptr); + __m512i q_msg = _mm512_sub_epi32(v_post, old_r); + + __mmask16 is_min1 = _mm512_cmpeq_epi32_mask(min1_idx, _mm512_set1_epi32(static_cast(e))); + __m512i min_mag = _mm512_mask_blend_epi32(is_min1, min1, min2); + + // Normalized Min-Sum scaling + __m512 mag_f = _mm512_cvtepi32_ps(min_mag); + __m512 norm_f = _mm512_mul_ps(mag_f, v_norm); + __m512i final_mag = _mm512_cvtps_epi32(norm_f); + + __mmask16 q_sign = _mm512_movepi32_mask(q_msg); + uint16_t ext_sign = total_sign_mask ^ static_cast(q_sign); + __mmask16 final_sign = static_cast<__mmask16>(syn_mask ^ ext_sign); + + __m512i neg_final_mag = _mm512_sub_epi32(v_zero, final_mag); + __m512i new_r = _mm512_mask_blend_epi32(final_sign, final_mag, neg_final_mag); + + _mm512_storeu_si512(msg_ptr, new_r); + + // Immediate Layered Update: L'_v = Q_{c,v} + R'_{c,v} + __m512i new_post = _mm512_add_epi32(q_msg, new_r); + _mm512_storeu_si512(v_ptr, new_post); + } + } // End of check loop + + // --- Convergence Check (End of iteration) --- + if (stop_at_convergence || (iter == max_iters - 1)) { + uint16_t converged_mask = active_mask; + + for (size_t c = 0; c < graph.num_checks; ++c) { + size_t start = graph.check_edge_offsets[c]; + size_t end = graph.check_edge_offsets[c + 1]; + uint16_t syn_mask = check_syndrome_masks[c]; + + uint16_t posterior_parity_mask = 0; + for (size_t e = start; e < end; ++e) { + size_t v = graph.check_edges[e]; + __m512i v_post = _mm512_loadu_si512(reinterpret_cast<__m512i*>(&posteriors_flat[v * BP_BATCH_SIZE])); + posterior_parity_mask ^= static_cast(_mm512_movepi32_mask(v_post)); + } + + uint16_t failed_checks = (posterior_parity_mask ^ syn_mask); + converged_mask &= ~failed_checks; + } + + // Deactivate newly converged shots + for (size_t b = 0; b < actual_batch_size; ++b) { + uint16_t bit = static_cast(1U << b); + if ((active_mask & bit) && (converged_mask & bit)) { + active_mask &= ~bit; + results[b].converged = true; + results[b].num_iters = iter + 1; + } + } + } + } + + // Update results for shots that didn't converge early + for (size_t b = 0; b < actual_batch_size; ++b) { + uint16_t bit = static_cast(1U << b); + if (active_mask & bit) { + results[b].converged = false; + results[b].num_iters = iter; + } + } + + return results; + } +#endif + + // Fallback portable path (for non-int32 types or non-AVX-512 targets) + using T_MAG = typename llr_traits::magnitude_type; + const T_MAG max_mag = std::numeric_limits::max(); + + std::vector active_shots(BP_BATCH_SIZE, 0); + for (size_t b = 0; b < actual_batch_size; ++b) active_shots[b] = 1; + + std::vector batched_syndromes(graph.num_checks * BP_BATCH_SIZE, 0); + for (size_t b = 0; b < actual_batch_size; ++b) { + for (size_t d : detection_events_batch[b]) { + batched_syndromes[d * BP_BATCH_SIZE + b] = 1; + } } size_t iter = 0; for (iter = 0; iter < max_iters && num_active > 0; ++iter) { - // Horizontal / Layered Schedule: Iterate through check nodes for (size_t c = 0; c < graph.num_checks; ++c) { size_t start = graph.check_edge_offsets[c]; size_t end = graph.check_edge_offsets[c + 1]; @@ -64,7 +222,6 @@ std::vector batched_bp_serial_min_sum( size_t min1_idx[BP_BATCH_SIZE]; uint8_t total_sign_prod[BP_BATCH_SIZE]; -#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { min1[b] = max_mag; min2[b] = max_mag; @@ -72,14 +229,11 @@ std::vector batched_bp_serial_min_sum( total_sign_prod[b] = 0; } - // Pass 1: Compute variable-to-check extrinsic messages Q_{c,v} = L_v - R_{c,v} - // and find min1, min2, and total sign product. for (size_t e = start; e < end; ++e) { size_t v = graph.check_edges[e]; size_t v_idx = v * BP_BATCH_SIZE; size_t msg_idx = e * BP_BATCH_SIZE; -#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { if (!active_shots[b]) continue; @@ -100,14 +254,11 @@ std::vector batched_bp_serial_min_sum( } } - // Pass 2: Compute new check-to-variable message R'_{c,v} - // and immediately update posterior L'_v = Q_{c,v} + R'_{c,v} = L_v + (R'_{c,v} - R_{c,v}) for (size_t e = start; e < end; ++e) { size_t v = graph.check_edges[e]; size_t v_idx = v * BP_BATCH_SIZE; size_t msg_idx = e * BP_BATCH_SIZE; -#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { if (!active_shots[b]) continue; @@ -133,9 +284,8 @@ std::vector batched_bp_serial_min_sum( posteriors_flat[v_idx + b] = q_msg + new_r; } } - } // End of check loop + } - // --- Convergence Check (End of iteration) --- if (stop_at_convergence || (iter == max_iters - 1)) { std::vector shot_converged(actual_batch_size, 1); @@ -173,7 +323,6 @@ std::vector batched_bp_serial_min_sum( } } - // Update results for shots that didn't converge early for (size_t b = 0; b < actual_batch_size; ++b) { if (active_shots[b]) { results[b].converged = false; From eac413acc237530c2244988b81565cc8ee60ad64 Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Fri, 7 Aug 2026 23:35:32 +0000 Subject: [PATCH 6/8] feat(bp): add stochastic check node permutation to break LDPC trapping sets --- src/bp/batched_bp_serial_min_sum.h | 2 +- src/bp/batched_bp_serial_min_sum.inl | 34 +++++++++++++++++-- src/bp/batched_bp_serial_min_sum.test.cc | 43 ++++++++++++++++++++++++ src/bp/bp.pybind.h | 4 ++- src/bp/bp_params.h | 17 ++++++++-- src/bp/bp_serial_min_sum.h | 3 +- src/bp/bp_serial_min_sum.inl | 24 +++++++++++-- src/bp/bp_serial_min_sum.test.cc | 28 +++++++++++++++ src/bp/tesseract_bp_decoder.cc | 16 ++++++--- src/bp_main.cc | 11 +++++- 10 files changed, 166 insertions(+), 16 deletions(-) diff --git a/src/bp/batched_bp_serial_min_sum.h b/src/bp/batched_bp_serial_min_sum.h index f1217960..c0cc1e07 100644 --- a/src/bp/batched_bp_serial_min_sum.h +++ b/src/bp/batched_bp_serial_min_sum.h @@ -12,7 +12,7 @@ template std::vector batched_bp_serial_min_sum( BatchedTannerGraph& graph, const std::vector>& detection_events_batch, std::vector& posteriors_flat, size_t max_iters, float normalization_factor, - bool stop_at_convergence = true); + bool stop_at_convergence = true, bool random_schedule = false, uint64_t random_seed = 0); } // namespace bp diff --git a/src/bp/batched_bp_serial_min_sum.inl b/src/bp/batched_bp_serial_min_sum.inl index c4d2f7c0..36b9e0ed 100644 --- a/src/bp/batched_bp_serial_min_sum.inl +++ b/src/bp/batched_bp_serial_min_sum.inl @@ -15,7 +15,7 @@ template std::vector batched_bp_serial_min_sum( BatchedTannerGraph& graph, const std::vector>& detection_events_batch, std::vector& posteriors_flat, size_t max_iters, float normalization_factor, - bool stop_at_convergence) { + bool stop_at_convergence, bool random_schedule, uint64_t random_seed) { size_t actual_batch_size = detection_events_batch.size(); if (actual_batch_size > BP_BATCH_SIZE) { throw std::invalid_argument("Provided batch size exceeds BP_BATCH_SIZE"); @@ -59,6 +59,18 @@ std::vector batched_bp_serial_min_sum( #endif } + std::vector check_order(graph.num_checks); + for (size_t i = 0; i < graph.num_checks; ++i) { + check_order[i] = i; + } + uint64_t rng_state = random_seed ? random_seed : 123456789ULL; + auto fast_rand = [&rng_state]() -> uint32_t { + rng_state ^= rng_state >> 12; + rng_state ^= rng_state << 25; + rng_state ^= rng_state >> 27; + return static_cast((rng_state * 0x2545F4914F6CDD1DULL) >> 32); + }; + #if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__) if constexpr (std::is_same_v) { const __m512 v_norm = _mm512_set1_ps(normalization_factor); @@ -66,8 +78,16 @@ std::vector batched_bp_serial_min_sum( size_t iter = 0; for (iter = 0; iter < max_iters && active_mask != 0; ++iter) { + if (random_schedule && graph.num_checks > 1) { + for (size_t i = graph.num_checks - 1; i > 0; --i) { + size_t j = fast_rand() % (i + 1); + std::swap(check_order[i], check_order[j]); + } + } + // Horizontal / Layered Schedule: Iterate through check nodes - for (size_t c = 0; c < graph.num_checks; ++c) { + for (size_t c_idx = 0; c_idx < graph.num_checks; ++c_idx) { + size_t c = check_order[c_idx]; size_t start = graph.check_edge_offsets[c]; size_t end = graph.check_edge_offsets[c + 1]; size_t deg = end - start; @@ -209,7 +229,15 @@ std::vector batched_bp_serial_min_sum( size_t iter = 0; for (iter = 0; iter < max_iters && num_active > 0; ++iter) { - for (size_t c = 0; c < graph.num_checks; ++c) { + if (random_schedule && graph.num_checks > 1) { + for (size_t i = graph.num_checks - 1; i > 0; --i) { + size_t j = fast_rand() % (i + 1); + std::swap(check_order[i], check_order[j]); + } + } + + for (size_t c_idx = 0; c_idx < graph.num_checks; ++c_idx) { + size_t c = check_order[c_idx]; size_t start = graph.check_edge_offsets[c]; size_t end = graph.check_edge_offsets[c + 1]; size_t deg = end - start; diff --git a/src/bp/batched_bp_serial_min_sum.test.cc b/src/bp/batched_bp_serial_min_sum.test.cc index bf760b5c..c5f891ec 100644 --- a/src/bp/batched_bp_serial_min_sum.test.cc +++ b/src/bp/batched_bp_serial_min_sum.test.cc @@ -55,4 +55,47 @@ TEST(BatchedBpSerialMinSumTest, ConvergesOnSimpleErrors) { EXPECT_LT(batched_posteriors[4 * BP_BATCH_SIZE + 1], 0); // V4 is negative } +TEST(BatchedBpSerialMinSumTest, ConvergesWithRandomSchedulePermutations) { + TannerGraph graph(5, 4, + {kPositivePriorLlr, kPositivePriorLlr, kPositivePriorLlr, + kPositivePriorLlr, kPositivePriorLlr}); + graph.add_edge(0, 0); + graph.add_edge(1, 0); + graph.add_edge(1, 1); + graph.add_edge(2, 1); + graph.add_edge(2, 2); + graph.add_edge(3, 2); + graph.add_edge(3, 3); + graph.add_edge(4, 3); + graph.build(); + + BatchedTannerGraph batched_graph; + batched_graph.build_from_unbatched(graph); + + std::vector> batched_syndromes; + batched_syndromes.push_back({1, 2}); // Shot 0: Error at V2 + batched_syndromes.push_back({0, 3}); // Shot 1: Error at V0, V4 + for (size_t i = 2; i < BP_BATCH_SIZE; i++) batched_syndromes.push_back({}); + + std::vector batched_posteriors(5 * BP_BATCH_SIZE, 0.0f); + + auto batched_results = + batched_bp_serial_min_sum(batched_graph, batched_syndromes, batched_posteriors, 20, + kNormalizationFactor, true, true, 12345); + + EXPECT_TRUE(batched_results[0].converged); + EXPECT_GT(batched_posteriors[0 * BP_BATCH_SIZE + 0], 0); + EXPECT_GT(batched_posteriors[1 * BP_BATCH_SIZE + 0], 0); + EXPECT_LT(batched_posteriors[2 * BP_BATCH_SIZE + 0], 0); + EXPECT_GT(batched_posteriors[3 * BP_BATCH_SIZE + 0], 0); + EXPECT_GT(batched_posteriors[4 * BP_BATCH_SIZE + 0], 0); + + EXPECT_TRUE(batched_results[1].converged); + EXPECT_LT(batched_posteriors[0 * BP_BATCH_SIZE + 1], 0); + EXPECT_GT(batched_posteriors[1 * BP_BATCH_SIZE + 1], 0); + EXPECT_GT(batched_posteriors[2 * BP_BATCH_SIZE + 1], 0); + EXPECT_GT(batched_posteriors[3 * BP_BATCH_SIZE + 1], 0); + EXPECT_LT(batched_posteriors[4 * BP_BATCH_SIZE + 1], 0); +} + } // namespace bp \ No newline at end of file diff --git a/src/bp/bp.pybind.h b/src/bp/bp.pybind.h index df25d219..189414a0 100644 --- a/src/bp/bp.pybind.h +++ b/src/bp/bp.pybind.h @@ -29,7 +29,9 @@ void add_bp_module(py::module& root) { .def_readwrite("max_iter", &BPParams::max_iter) .def_readwrite("update_rule", &BPParams::update_rule) .def_readwrite("schedule", &BPParams::schedule) - .def_readwrite("normalization_factor", &BPParams::normalization_factor); + .def_readwrite("normalization_factor", &BPParams::normalization_factor) + .def_readwrite("random_schedule", &BPParams::random_schedule) + .def_readwrite("random_seed", &BPParams::random_seed); py::class_>(m, "PostProcessor"); diff --git a/src/bp/bp_params.h b/src/bp/bp_params.h index 256130f6..649c1198 100644 --- a/src/bp/bp_params.h +++ b/src/bp/bp_params.h @@ -44,6 +44,8 @@ const bool DEFAULT_SORT_PRIORS = false; // The fraction of error mechanisms to use in BP. Determines the maximum variable node index // included, specified as the fraction of all error mechanisms. const double DEFAULT_VARIABLE_NODE_TRUNCATION_FRACTION = 1.0; +const bool DEFAULT_RANDOM_SCHEDULE = false; +const uint64_t DEFAULT_RANDOM_SEED = 0; struct BPParams { size_t max_iter; @@ -53,26 +55,35 @@ struct BPParams { bool sort_priors; double variable_node_truncation_fraction; float normalization_factor; + bool random_schedule; + uint64_t random_seed; + BPParams(size_t max_iter = DEFAULT_MAX_ITER, std::string update_rule = DEFAULT_UPDATE_RULE, std::string schedule = DEFAULT_SCHEDULE, bool stop_at_convergence = DEFAULT_STOP_AT_CONVERGENCE, bool sort_priors = DEFAULT_SORT_PRIORS, double variable_node_truncation_fraction = DEFAULT_VARIABLE_NODE_TRUNCATION_FRACTION, - float normalization_factor = 0.875f) + float normalization_factor = 0.875f, + bool random_schedule = DEFAULT_RANDOM_SCHEDULE, + uint64_t random_seed = DEFAULT_RANDOM_SEED) : max_iter(max_iter), update_rule(update_rule), schedule(schedule), stop_at_convergence(stop_at_convergence), sort_priors(sort_priors), variable_node_truncation_fraction(variable_node_truncation_fraction), - normalization_factor(normalization_factor) {} + normalization_factor(normalization_factor), + random_schedule(random_schedule), + random_seed(random_seed) {} bool operator==(const BPParams& other) const { return max_iter == other.max_iter && update_rule == other.update_rule && schedule == other.schedule && stop_at_convergence == other.stop_at_convergence && sort_priors == other.sort_priors && variable_node_truncation_fraction == other.variable_node_truncation_fraction && - normalization_factor == other.normalization_factor; + normalization_factor == other.normalization_factor && + random_schedule == other.random_schedule && + random_seed == other.random_seed; } }; } // namespace bp diff --git a/src/bp/bp_serial_min_sum.h b/src/bp/bp_serial_min_sum.h index 197c8a5d..f2ac5f91 100644 --- a/src/bp/bp_serial_min_sum.h +++ b/src/bp/bp_serial_min_sum.h @@ -11,7 +11,8 @@ namespace bp { template BPResult bp_serial_min_sum(TannerGraph& graph, const std::vector& detection_events, std::vector& posteriors, size_t max_iters, float normalization_factor, - bool stop_at_convergence = true); + bool stop_at_convergence = true, bool random_schedule = false, + uint64_t random_seed = 0); } // namespace bp diff --git a/src/bp/bp_serial_min_sum.inl b/src/bp/bp_serial_min_sum.inl index d7591144..0c72d131 100644 --- a/src/bp/bp_serial_min_sum.inl +++ b/src/bp/bp_serial_min_sum.inl @@ -10,7 +10,7 @@ namespace bp { template BPResult bp_serial_min_sum(TannerGraph& graph, const std::vector& detection_events, std::vector& posteriors, size_t max_iters, float normalization_factor, - bool stop_at_convergence) { + bool stop_at_convergence, bool random_schedule, uint64_t random_seed) { graph.add_detection_events(detection_events); // Initialize posteriors to priors @@ -22,11 +22,31 @@ BPResult bp_serial_min_sum(TannerGraph& graph, const std::vector& det using T_MAG = typename llr_traits::magnitude_type; const T_MAG max_mag = std::numeric_limits::max(); + std::vector check_order(graph.check_nodes.size()); + for (size_t i = 0; i < graph.check_nodes.size(); ++i) { + check_order[i] = i; + } + uint64_t rng_state = random_seed ? random_seed : 123456789ULL; + auto fast_rand = [&rng_state]() -> uint32_t { + rng_state ^= rng_state >> 12; + rng_state ^= rng_state << 25; + rng_state ^= rng_state >> 27; + return static_cast((rng_state * 0x2545F4914F6CDD1DULL) >> 32); + }; + bool has_converged = false; size_t iter = 0; for (iter = 0; iter < max_iters; ++iter) { + if (random_schedule && graph.check_nodes.size() > 1) { + for (size_t i = graph.check_nodes.size() - 1; i > 0; --i) { + size_t j = fast_rand() % (i + 1); + std::swap(check_order[i], check_order[j]); + } + } + // Horizontal / Layered Schedule: Iterate through check nodes - for (size_t c = 0; c < graph.check_nodes.size(); ++c) { + for (size_t c_idx = 0; c_idx < graph.check_nodes.size(); ++c_idx) { + size_t c = check_order[c_idx]; auto& check = graph.check_nodes[c]; size_t start = graph.check_edge_offsets[c]; size_t end = graph.check_edge_offsets[c + 1]; diff --git a/src/bp/bp_serial_min_sum.test.cc b/src/bp/bp_serial_min_sum.test.cc index e64f6661..0e45ebf4 100644 --- a/src/bp/bp_serial_min_sum.test.cc +++ b/src/bp/bp_serial_min_sum.test.cc @@ -67,4 +67,32 @@ TEST(BpSerialMinSumTest, CorrectsSingleErrorWithSyndrome) { EXPECT_GT(posteriors[4], 0); } +TEST(BpSerialMinSumTest, ConvergesWithRandomSchedulePermutations) { + TannerGraph graph(5, 4, + {kPositivePriorLlr, kPositivePriorLlr, kPositivePriorLlr, + kPositivePriorLlr, kPositivePriorLlr}); + graph.add_edge(0, 0); + graph.add_edge(1, 0); + graph.add_edge(1, 1); + graph.add_edge(2, 1); + graph.add_edge(2, 2); + graph.add_edge(3, 2); + graph.add_edge(3, 3); + graph.add_edge(4, 3); + graph.build(); + + std::vector detection_events = {1, 2}; + std::vector posteriors(5); + + BPResult result = bp_serial_min_sum(graph, detection_events, posteriors, 20, + kNormalizationFactor, true, true, 42); + + EXPECT_TRUE(result.converged); + EXPECT_GT(posteriors[0], 0); + EXPECT_GT(posteriors[1], 0); + EXPECT_LT(posteriors[2], 0); + EXPECT_GT(posteriors[3], 0); + EXPECT_GT(posteriors[4], 0); +} + } // namespace bp \ No newline at end of file diff --git a/src/bp/tesseract_bp_decoder.cc b/src/bp/tesseract_bp_decoder.cc index c912b045..b68fbd6c 100644 --- a/src/bp/tesseract_bp_decoder.cc +++ b/src/bp/tesseract_bp_decoder.cc @@ -78,16 +78,20 @@ std::vector TesseractBpDecoder::decode( const std::shared_ptr& post_processor) { std::vector posteriors(graph_.variable_nodes.size()); + bool is_random_sched = params_.random_schedule || params_.schedule == "random-serial" || + params_.schedule == "stochastic-serial" || params_.schedule == "random"; + BPResult result; if (params_.schedule == "parallel" && params_.update_rule == "min-sum") { std::vector dets(detection_events.begin(), detection_events.end()); result = bp_parallel_min_sum(graph_, dets, posteriors, params_.max_iter, params_.normalization_factor, params_.stop_at_convergence); - } else if (params_.schedule == "serial" && params_.update_rule == "min-sum") { + } else if ((params_.schedule == "serial" || is_random_sched) && params_.update_rule == "min-sum") { std::vector dets(detection_events.begin(), detection_events.end()); result = bp_serial_min_sum(graph_, dets, posteriors, params_.max_iter, - params_.normalization_factor, params_.stop_at_convergence); + params_.normalization_factor, params_.stop_at_convergence, + is_random_sched, params_.random_seed); } else { throw std::invalid_argument( "Unsupported schedule/update_rule combination. Only min-sum is supported in this phase."); @@ -124,6 +128,9 @@ std::vector> TesseractBpDecoder::decode_batch( alignas(64) std::vector posteriors_flat(num_vars * BP_BATCH_SIZE); std::vector shot_posteriors(num_vars); + bool is_random_sched = params_.random_schedule || params_.schedule == "random-serial" || + params_.schedule == "stochastic-serial" || params_.schedule == "random"; + for (size_t shot = 0; shot < num_shots; ++shot) { current_batch.push_back(std::vector(detection_events_batch[shot].begin(), detection_events_batch[shot].end())); @@ -135,10 +142,11 @@ std::vector> TesseractBpDecoder::decode_batch( } std::vector bp_results; - if (params_.schedule == "serial") { + if (params_.schedule == "serial" || is_random_sched) { bp_results = batched_bp_serial_min_sum(batched_graph_, current_batch, posteriors_flat, params_.max_iter, params_.normalization_factor, - params_.stop_at_convergence); + params_.stop_at_convergence, is_random_sched, + params_.random_seed); } else { bp_results = batched_bp_parallel_min_sum(batched_graph_, current_batch, posteriors_flat, params_.max_iter, params_.normalization_factor, diff --git a/src/bp_main.cc b/src/bp_main.cc index c9372bbf..094cca82 100644 --- a/src/bp_main.cc +++ b/src/bp_main.cc @@ -59,6 +59,7 @@ struct Args { int osd_order = -1; // -1 means HardDecision, >= 0 means OSD int osd_weight = 0; bool use_batched_bp = false; + bool random_schedule = false; std::string stats_out_fname = ""; std::string sinter_csv_out = ""; @@ -117,6 +118,9 @@ struct Args { params.update_rule = update_rule; params.schedule = schedule; params.normalization_factor = (float)normalization_factor; + params.random_schedule = random_schedule || schedule == "random-serial" || + schedule == "stochastic-serial" || schedule == "random"; + params.random_seed = sample_seed; if (sample_num_shots > 0) { std::mt19937_64 rng(sample_seed); @@ -224,6 +228,10 @@ int main(int argc, char* argv[]) { .help("Use AVX-512 batching across shots") .flag() .store_into(args.use_batched_bp); + program.add_argument("--random-schedule") + .help("Use randomized check node permutation in serial BP to prevent trapping sets") + .flag() + .store_into(args.random_schedule); program.add_argument("--threads") .default_value(size_t( std::thread::hardware_concurrency() == 0 ? 1 : std::thread::hardware_concurrency())) @@ -394,7 +402,8 @@ int main(int argc, char* argv[]) { double cpu_throughput = (cpu_time > 0) ? (double)final_shots / cpu_time : 0.0; std::string decoder_name = - std::string(args.use_batched_bp ? "batched-" : "scalar-") + args.schedule + "-bp"; + std::string(args.use_batched_bp ? "batched-" : "scalar-") + + (params.random_schedule ? "random-" : "") + args.schedule + "-bp"; if (args.osd_order >= 0) decoder_name += "+osd"; if (!args.stats_out_fname.empty()) { From 85f8c44dbcba94f77df5f44c09c21e2989e091da Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Sat, 8 Aug 2026 16:08:34 +0000 Subject: [PATCH 7/8] Revert "feat(bp): vectorize layered serial min-sum with AVX-512 intrinsics" This reverts commit c900f95, returning to GCC ivdep pragmas for auto-vectorization in order to maximize code readability and simplicity. The marginal performance gain in release builds (~2-4%) was deemed not worth the architectural complexity of explicit intrinsics. Also forces -c opt in benchmark.sh to prevent unintentional unoptimized profiling. --- devtools/benchmark.sh | 4 +- src/bp/batched_bp_serial_min_sum.inl | 199 +++------------------------ 2 files changed, 23 insertions(+), 180 deletions(-) diff --git a/devtools/benchmark.sh b/devtools/benchmark.sh index 29d8ac6e..6d7c23e4 100755 --- a/devtools/benchmark.sh +++ b/devtools/benchmark.sh @@ -34,8 +34,8 @@ echo " Threads: ${THREADS}" echo "============================================================" # Build binaries using Bazel (single-core for build as per rules) -echo "=== Building C++ Binaries (bazel build --jobs=1 src:bp src:tesseract) ===" -bazel build --jobs=1 src:bp src:tesseract +echo "=== Building C++ Binaries (bazel build -c opt --jobs=1 src:bp src:tesseract) ===" +bazel build -c opt --jobs=1 src:bp src:tesseract BP_BIN="${WORKSPACE_ROOT}/bazel-bin/src/bp" TESSERACT_BIN="${WORKSPACE_ROOT}/bazel-bin/src/tesseract" diff --git a/src/bp/batched_bp_serial_min_sum.inl b/src/bp/batched_bp_serial_min_sum.inl index 36b9e0ed..8af6ce8b 100644 --- a/src/bp/batched_bp_serial_min_sum.inl +++ b/src/bp/batched_bp_serial_min_sum.inl @@ -3,10 +3,6 @@ #include #include -#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__) -#include -#endif - #include "bp/batched_bp_serial_min_sum.h" namespace bp { @@ -21,21 +17,24 @@ std::vector batched_bp_serial_min_sum( throw std::invalid_argument("Provided batch size exceeds BP_BATCH_SIZE"); } - // Set up batched syndromes as 16-bit masks per check node. - std::vector check_syndrome_masks(graph.num_checks, 0); + // Set up batched syndromes. + std::vector batched_syndromes(graph.num_checks * BP_BATCH_SIZE, 0); for (size_t b = 0; b < actual_batch_size; ++b) { for (size_t d : detection_events_batch[b]) { - check_syndrome_masks[d] |= (static_cast(1) << b); + batched_syndromes[d * BP_BATCH_SIZE + b] = 1; } } - uint16_t active_mask = (actual_batch_size == 16) - ? 0xFFFF - : static_cast((1U << actual_batch_size) - 1); + // Track which shots in the batch are still active + std::vector active_shots(BP_BATCH_SIZE, 0); + for (size_t b = 0; b < actual_batch_size; ++b) active_shots[b] = 1; size_t num_active = actual_batch_size; std::vector results(actual_batch_size, {false, 0}); + using T_MAG = typename llr_traits::magnitude_type; + const T_MAG max_mag = std::numeric_limits::max(); + // Initialize check_to_var messages to 0 std::fill(graph.check_to_var_messages.begin(), graph.check_to_var_messages.end(), 0); @@ -43,20 +42,10 @@ std::vector batched_bp_serial_min_sum( for (size_t i = 0; i < graph.num_variables; ++i) { T prior_val = graph.priors[i]; size_t var_post_idx = i * BP_BATCH_SIZE; -#if defined(__AVX512F__) - if constexpr (std::is_same_v) { - __m512i v_prior = _mm512_set1_epi32(prior_val); - _mm512_storeu_si512(reinterpret_cast<__m512i*>(&posteriors_flat[var_post_idx]), v_prior); - } else { - for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { - posteriors_flat[var_post_idx + b] = prior_val; - } - } -#else +#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { posteriors_flat[var_post_idx + b] = prior_val; } -#endif } std::vector check_order(graph.num_checks); @@ -71,162 +60,6 @@ std::vector batched_bp_serial_min_sum( return static_cast((rng_state * 0x2545F4914F6CDD1DULL) >> 32); }; -#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__) - if constexpr (std::is_same_v) { - const __m512 v_norm = _mm512_set1_ps(normalization_factor); - const __m512i v_zero = _mm512_setzero_si512(); - - size_t iter = 0; - for (iter = 0; iter < max_iters && active_mask != 0; ++iter) { - if (random_schedule && graph.num_checks > 1) { - for (size_t i = graph.num_checks - 1; i > 0; --i) { - size_t j = fast_rand() % (i + 1); - std::swap(check_order[i], check_order[j]); - } - } - - // Horizontal / Layered Schedule: Iterate through check nodes - for (size_t c_idx = 0; c_idx < graph.num_checks; ++c_idx) { - size_t c = check_order[c_idx]; - size_t start = graph.check_edge_offsets[c]; - size_t end = graph.check_edge_offsets[c + 1]; - size_t deg = end - start; - if (deg == 0) continue; - - uint16_t syn_mask = check_syndrome_masks[c]; - - __m512i min1 = _mm512_set1_epi32(std::numeric_limits::max()); - __m512i min2 = _mm512_set1_epi32(std::numeric_limits::max()); - __m512i min1_idx = _mm512_set1_epi32(-1); - uint16_t total_sign_mask = 0; - - // Pass 1: Compute variable-to-check extrinsic messages Q_{c,v} = L_v - R_{c,v} - // and find min1, min2, and total sign product across 16 SIMD lanes. - for (size_t e = start; e < end; ++e) { - size_t v = graph.check_edges[e]; - size_t v_idx = v * BP_BATCH_SIZE; - size_t msg_idx = e * BP_BATCH_SIZE; - - __m512i* v_ptr = reinterpret_cast<__m512i*>(&posteriors_flat[v_idx]); - __m512i* msg_ptr = reinterpret_cast<__m512i*>(&graph.check_to_var_messages[msg_idx]); - - __m512i v_post = _mm512_loadu_si512(v_ptr); - __m512i old_r = _mm512_loadu_si512(msg_ptr); - __m512i q_msg = _mm512_sub_epi32(v_post, old_r); - - __m512i mag = _mm512_abs_epi32(q_msg); - __mmask16 q_sign = _mm512_movepi32_mask(q_msg); - total_sign_mask ^= static_cast(q_sign); - - __mmask16 is_less_min1 = _mm512_cmplt_epi32_mask(mag, min1); - __mmask16 is_less_min2 = _mm512_cmplt_epi32_mask(mag, min2); - - // If mag < min1: min2 = min1, min1 = mag, min1_idx = e - min2 = _mm512_mask_blend_epi32(is_less_min1, min2, min1); - min1 = _mm512_mask_blend_epi32(is_less_min1, min1, mag); - min1_idx = _mm512_mask_blend_epi32(is_less_min1, min1_idx, _mm512_set1_epi32(static_cast(e))); - - // Else if mag < min2: min2 = mag - __mmask16 update_min2 = is_less_min2 & ~is_less_min1; - min2 = _mm512_mask_blend_epi32(update_min2, min2, mag); - } - - // Pass 2: Compute new check-to-variable message R'_{c,v} - // and immediately update posterior L'_v = Q_{c,v} + R'_{c,v} - for (size_t e = start; e < end; ++e) { - size_t v = graph.check_edges[e]; - size_t v_idx = v * BP_BATCH_SIZE; - size_t msg_idx = e * BP_BATCH_SIZE; - - __m512i* v_ptr = reinterpret_cast<__m512i*>(&posteriors_flat[v_idx]); - __m512i* msg_ptr = reinterpret_cast<__m512i*>(&graph.check_to_var_messages[msg_idx]); - - __m512i v_post = _mm512_loadu_si512(v_ptr); - __m512i old_r = _mm512_loadu_si512(msg_ptr); - __m512i q_msg = _mm512_sub_epi32(v_post, old_r); - - __mmask16 is_min1 = _mm512_cmpeq_epi32_mask(min1_idx, _mm512_set1_epi32(static_cast(e))); - __m512i min_mag = _mm512_mask_blend_epi32(is_min1, min1, min2); - - // Normalized Min-Sum scaling - __m512 mag_f = _mm512_cvtepi32_ps(min_mag); - __m512 norm_f = _mm512_mul_ps(mag_f, v_norm); - __m512i final_mag = _mm512_cvtps_epi32(norm_f); - - __mmask16 q_sign = _mm512_movepi32_mask(q_msg); - uint16_t ext_sign = total_sign_mask ^ static_cast(q_sign); - __mmask16 final_sign = static_cast<__mmask16>(syn_mask ^ ext_sign); - - __m512i neg_final_mag = _mm512_sub_epi32(v_zero, final_mag); - __m512i new_r = _mm512_mask_blend_epi32(final_sign, final_mag, neg_final_mag); - - _mm512_storeu_si512(msg_ptr, new_r); - - // Immediate Layered Update: L'_v = Q_{c,v} + R'_{c,v} - __m512i new_post = _mm512_add_epi32(q_msg, new_r); - _mm512_storeu_si512(v_ptr, new_post); - } - } // End of check loop - - // --- Convergence Check (End of iteration) --- - if (stop_at_convergence || (iter == max_iters - 1)) { - uint16_t converged_mask = active_mask; - - for (size_t c = 0; c < graph.num_checks; ++c) { - size_t start = graph.check_edge_offsets[c]; - size_t end = graph.check_edge_offsets[c + 1]; - uint16_t syn_mask = check_syndrome_masks[c]; - - uint16_t posterior_parity_mask = 0; - for (size_t e = start; e < end; ++e) { - size_t v = graph.check_edges[e]; - __m512i v_post = _mm512_loadu_si512(reinterpret_cast<__m512i*>(&posteriors_flat[v * BP_BATCH_SIZE])); - posterior_parity_mask ^= static_cast(_mm512_movepi32_mask(v_post)); - } - - uint16_t failed_checks = (posterior_parity_mask ^ syn_mask); - converged_mask &= ~failed_checks; - } - - // Deactivate newly converged shots - for (size_t b = 0; b < actual_batch_size; ++b) { - uint16_t bit = static_cast(1U << b); - if ((active_mask & bit) && (converged_mask & bit)) { - active_mask &= ~bit; - results[b].converged = true; - results[b].num_iters = iter + 1; - } - } - } - } - - // Update results for shots that didn't converge early - for (size_t b = 0; b < actual_batch_size; ++b) { - uint16_t bit = static_cast(1U << b); - if (active_mask & bit) { - results[b].converged = false; - results[b].num_iters = iter; - } - } - - return results; - } -#endif - - // Fallback portable path (for non-int32 types or non-AVX-512 targets) - using T_MAG = typename llr_traits::magnitude_type; - const T_MAG max_mag = std::numeric_limits::max(); - - std::vector active_shots(BP_BATCH_SIZE, 0); - for (size_t b = 0; b < actual_batch_size; ++b) active_shots[b] = 1; - - std::vector batched_syndromes(graph.num_checks * BP_BATCH_SIZE, 0); - for (size_t b = 0; b < actual_batch_size; ++b) { - for (size_t d : detection_events_batch[b]) { - batched_syndromes[d * BP_BATCH_SIZE + b] = 1; - } - } - size_t iter = 0; for (iter = 0; iter < max_iters && num_active > 0; ++iter) { if (random_schedule && graph.num_checks > 1) { @@ -236,6 +69,7 @@ std::vector batched_bp_serial_min_sum( } } + // Horizontal / Layered Schedule: Iterate through check nodes for (size_t c_idx = 0; c_idx < graph.num_checks; ++c_idx) { size_t c = check_order[c_idx]; size_t start = graph.check_edge_offsets[c]; @@ -250,6 +84,7 @@ std::vector batched_bp_serial_min_sum( size_t min1_idx[BP_BATCH_SIZE]; uint8_t total_sign_prod[BP_BATCH_SIZE]; +#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { min1[b] = max_mag; min2[b] = max_mag; @@ -257,11 +92,14 @@ std::vector batched_bp_serial_min_sum( total_sign_prod[b] = 0; } + // Pass 1: Compute variable-to-check extrinsic messages Q_{c,v} = L_v - R_{c,v} + // and find min1, min2, and total sign product. for (size_t e = start; e < end; ++e) { size_t v = graph.check_edges[e]; size_t v_idx = v * BP_BATCH_SIZE; size_t msg_idx = e * BP_BATCH_SIZE; +#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { if (!active_shots[b]) continue; @@ -282,11 +120,14 @@ std::vector batched_bp_serial_min_sum( } } + // Pass 2: Compute new check-to-variable message R'_{c,v} + // and immediately update posterior L'_v = Q_{c,v} + R'_{c,v} = L_v + (R'_{c,v} - R_{c,v}) for (size_t e = start; e < end; ++e) { size_t v = graph.check_edges[e]; size_t v_idx = v * BP_BATCH_SIZE; size_t msg_idx = e * BP_BATCH_SIZE; +#pragma GCC ivdep for (size_t b = 0; b < BP_BATCH_SIZE; ++b) { if (!active_shots[b]) continue; @@ -312,8 +153,9 @@ std::vector batched_bp_serial_min_sum( posteriors_flat[v_idx + b] = q_msg + new_r; } } - } + } // End of check loop + // --- Convergence Check (End of iteration) --- if (stop_at_convergence || (iter == max_iters - 1)) { std::vector shot_converged(actual_batch_size, 1); @@ -351,6 +193,7 @@ std::vector batched_bp_serial_min_sum( } } + // Update results for shots that didn't converge early for (size_t b = 0; b < actual_batch_size; ++b) { if (active_shots[b]) { results[b].converged = false; From 96da6a271506e03c9386d559c250d3a4ef8dc00f Mon Sep 17 00:00:00 2001 From: Aria Shahingohar Date: Sat, 8 Aug 2026 16:41:56 +0000 Subject: [PATCH 8/8] test: add parallel decoder tests to benchmark script --- devtools/benchmark.sh | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/devtools/benchmark.sh b/devtools/benchmark.sh index 6d7c23e4..75693937 100755 --- a/devtools/benchmark.sh +++ b/devtools/benchmark.sh @@ -49,6 +49,7 @@ run_bp_benchmark() { local osd_weight="${6:-0}" local max_iter="${7:-30}" local max_errors="${8:-100}" + local schedule="${9:-serial}" if [[ ! -f "${circuit}" ]]; then echo "[SKIP] Circuit not found: ${circuit}" @@ -60,7 +61,7 @@ run_bp_benchmark() { echo "------------------------------------------------------------" echo ">> Running: ${name}" echo " Circuit: ${circuit}" - echo " Config: schedule=serial, batched=true, osd_order=${osd_order}, osd_weight=${osd_weight}, norm=${norm}" + echo " Config: schedule=${schedule}, batched=true, osd_order=${osd_order}, osd_weight=${osd_weight}, norm=${norm}" echo "------------------------------------------------------------" local cmd=( @@ -72,7 +73,7 @@ run_bp_benchmark() { --max-errors "${max_errors}" --normalization-factor "${norm}" --max-iter "${max_iter}" - --schedule "serial" + --schedule "${schedule}" --batched --print-stats --stats-out "${out_json}" @@ -130,12 +131,26 @@ run_bp_benchmark \ 100000 \ 0.625 +run_bp_benchmark \ + "surface_code_d3_p001_parallel_batched" \ + "testdata/surfacecodes/r=3,d=3,p=0.001,noise=si1000,c=surface_code_Z,q=17,gates=cz.stim" \ + 100000 \ + 0.625 \ + -1 0 30 100 "parallel" + run_bp_benchmark \ "surface_code_d5_p001_serial_batched" \ "testdata/surfacecodes/r=5,d=5,p=0.001,noise=si1000,c=surface_code_Z,q=49,gates=cz.stim" \ 100000 \ 0.625 +run_bp_benchmark \ + "surface_code_d5_p001_parallel_batched" \ + "testdata/surfacecodes/r=5,d=5,p=0.001,noise=si1000,c=surface_code_Z,q=49,gates=cz.stim" \ + 100000 \ + 0.625 \ + -1 0 30 100 "parallel" + run_bp_benchmark \ "surface_code_d7_p001_serial_batched" \ "testdata/surfacecodes/r=7,d=7,p=0.001,noise=si1000,c=surface_code_Z,q=97,gates=cz.stim" \