Feat/bp decoder optimizations

.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

.gitignore

@@ -43,3 +43,6 @@ MODULE.bazel.lock
 # Python build artifacts
 *.so
 *.egg-info/
+
+# Local CMake build directory
+build/

devtools/benchmark.sh

@@ -0,0 +1,304 @@
+#!/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 -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"
+
+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}"
+    local schedule="${9:-serial}"
+
+    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=${schedule}, 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 "${schedule}"
+        --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_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" \
+    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 - <<EOF
+import os, glob, json
+
+out_dir = "${OUT_DIR}"
+cl_tag = "${CL_TAG}"
+json_files = sorted(glob.glob(os.path.join(out_dir, "*.json")))
+
+if not json_files:
+    print("No JSON files found in " + out_dir)
+else:
+    print(f"| {'Benchmark':<42} | {'Decoder':<22} | {'Shots':>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

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 <typename T>
 std::vector<BPResult> batched_bp_parallel_min_sum(
     BatchedTannerGraph<T>& graph, const std::vector<std::vector<size_t>>& detection_events_batch,
-    std::vector<std::vector<T>>& posteriors_batch, size_t max_iters, float normalization_factor,
+    std::vector<T>& posteriors_flat, size_t max_iters, float normalization_factor,
     bool stop_at_convergence = true);
 
 }  // namespace bp

src/bp/batched_bp_parallel_min_sum.inl

@@ -10,7 +10,7 @@ namespace bp {
 template <typename T>
 std::vector<BPResult> batched_bp_parallel_min_sum(
     BatchedTannerGraph<T>& graph, const std::vector<std::vector<size_t>>& detection_events_batch,
-    std::vector<std::vector<T>>& posteriors_batch, size_t max_iters, float normalization_factor,
+    std::vector<T>& 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<BPResult> 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<BPResult> 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<BPResult> 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;
             }
           }

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<std::vector<float>> batched_posteriors(BP_BATCH_SIZE, std::vector<float>(5, 0));
+  std::vector<float> 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.";
     }
   }