SIMD Optimization Reference

Date: 2026-06-08 Latest benchmark host: landau Original AVX2 selection host: cantor

This document is the authoritative reference for SIMD work tried in the Rust BGEN and preprocessing paths. It merges the old planning notes, benchmark results, and follow-up opportunity audit so future work can quickly see what was tested, what was kept, and what should not be repeated without new profiling evidence.

Current Policy

• AVX2 is the only production x86 SIMD target.
• Scalar fallbacks remain for non-AVX2 platforms.
• AVX-512 is not a production target. It was tried before this reference was written and was slower than AVX2 for this workload. A fresh landau pass on 2026-06-08 found that AVX-512 is available on the target CPU, but wider AVX-512 decode/copy prototypes still regressed adjacent reader benchmark groups, so the AVX2-only policy remains.
• Runtime benchmark switches were removed after selecting the winning implementations.
• SIMD code should stay isolated in small Rust modules or helpers using std::arch and runtime is_x86_feature_detected!("avx2") checks.
• Keep only implementations that improve realistic reader or preprocessing benchmarks. Do not keep alternate production versions just because they are interesting.

Decision Matrix

2026-06-08 Landau Review

GLA-7 reopened SIMD review for instruction sets actually present on landau. The original implementation at fc0c49a54f82ad9463e29c04fb7507f3f7bd4657 was benchmarked before any source changes. The worktree used RUSTFLAGS="-C target-cpu=native" and CARGO_BUILD_JOBS=40; benchmarks used /mnt/beegfs/kirill/Projects/g/data/1kg_chr22_full.bgen with the matching .sample file.

landau CPU metadata from Slurm:

• CPU: Intel Xeon Gold 5220, 2 sockets, 18 cores per socket, 2 threads per core.
• Rust: rustc 1.96.0 (ac68faa20 2026-05-25), LLVM 22.1.2.
• SIMD flags include avx, avx2, avx512f, avx512dq, avx512cd, avx512bw, avx512vl, and avx512_vnni.

Original baseline medians in milliseconds:

AVX-512 was prototyped after the baseline using the same node, data, build flags, benchmark filter, sample size, warm-up, and measurement time. Negative percentages below mean faster than the original baseline; positive percentages mean slower.

After refreshing the documentation-only final branch to origin/main at 5a69b16c, the same landau Criterion command was rerun with --baseline gla7-original and no AVX-512 production code. That no-code comparison still showed mixed small changes and an isolated bgen_row_major_direct_write/8192 regression of about 23.8%, so that single benchmark point should be treated as noisy current-main drift rather than as the sole AVX-512 rejection signal. The AVX-512 decision rests on the full pattern: no reproducible clean reader-level win across the required groups, plus adjacent strided/subset regressions in the candidate runs.

Decision: no production code changed in GLA-7. The available AVX-512 support on landau does not justify replacing the current AVX2 production paths because the wider kernels do not produce clean reader-level wins without adjacent regressions. Keep AVX2 as the only x86 production SIMD target, with scalar fallbacks.

Core BGEN Math

For unphased 8-bit BGEN probability pairs:

p0 = homozygous reference probability byte
p1 = heterozygous probability byte

raw = 510 - 2*p0 - p1
dosage = raw / 255

The optimized paths compute output dosage values from raw and compute summary statistics in integer space:

dosage_sum = raw_sum / 255
dosage_square_sum = raw_square_sum / (255 * 255)

Counts are also derived from integer raw thresholds:

zero_count = raw == 0
nonzero_count = raw > 0
homozygous_reference_count = raw <= 127
heterozygous_count = 128 <= raw <= 382
homozygous_alternate_count = raw >= 383

This removes lookup-table reads, avoids per-sample floating-point accumulation in the hot identity path, and gives AVX2 a simple contiguous byte-decoding problem.

Implemented Paths

Trusted Identity Decode

Path:

crates/genotype/src/bgen/trusted.rs
crates/genotype/src/bgen/simd.rs

Shape:

trusted no-missing diploid
unphased
8-bit probabilities
variant-major output
identity/full-sample selection

The old trusted identity loop used lookup-table dosage decoding and per-sample selection logic. The retained path branches once on identity selection, decodes contiguous probability bytes with raw integer AVX2, writes a contiguous variant-major output row, and returns integer-derived summary stats. Scalar raw integer decode remains only as the fallback when AVX2 is unavailable.

Median reader times from the trusted benchmark:

Single 16,384-variant native profile:

Non-Trusted Full-Sample No-Missing Decode

Path:

crates/genotype/src/bgen/decode/mod.rs
crates/genotype/src/bgen/simd.rs

Shape:

trusted flag disabled
unphased
8-bit probabilities
variant-major output
identity/full-sample selection
all samples present and diploid

Once the generic decoder proves the row is all-present and identity-selected, it uses the same raw integer AVX2 helper as the trusted path.

bgen_preprocessed_variant_major_trusted_disabled, 40 CPUs, RUSTFLAGS="-C target-cpu=native":

All-Present Diploid Scan

Path:

crates/genotype/src/bgen/trusted.rs
crates/genotype/src/bgen/simd.rs

all_samples_present_diploid now uses AVX2 to compare 32 bytes at a time against byte value 2, falling back to the old scalar chunk scan on non-AVX2 platforms. This is not part of trusted hot decode after validation, but it is still used by trusted validation and non-trusted all-present detection.

Result:

• No statistically significant change for some smaller chunks.
• Criterion-reported small improvements at 4096 and 16384 in the non-trusted reader benchmark.
• Final full-sample benchmark after all changes stayed within noise for the main reader path.

Decision: kept as a small AVX2-only improvement with scalar fallback.

Dense Contiguous Selected Subsets

Path:

crates/genotype/src/bgen/sample_selection.rs
crates/genotype/src/bgen/decode/mod.rs
crates/genotype/src/bgen/trusted.rs

The sample-selection builder records contiguous_file_index_start when selected file indices form one monotonic contiguous run. For that shape, selected input probabilities and output dosages are both contiguous, so generic and trusted variant-major decode reuse the raw AVX2 identity decoder. Sparse or shuffled subsets continue to use the scalar lookup path.

bgen_preprocessed_variant_major_contiguous_subset_trusted_disabled:

bgen_preprocessed_variant_major_contiguous_subset_trusted_no_missing_diploid:

Variant-Major Preprocessing Summaries

Path:

crates/genotype/src/preprocess.rs
crates/genotype/benches/preprocess.rs

The variant-major summarizer now uses AVX2 masks over contiguous f32 rows to compute sums, square sums, observation counts, and dosage threshold counts. Scalar fallback remains for non-AVX2 platforms and row tails.

Focused preprocess_variant_major_summary benchmark over dense no-missing f32 rows:

Rejected Or Deferred Work

AVX-512

AVX-512 was tried before the final AVX2-only policy and was slower than AVX2 for this workload. Do not reintroduce an AVX-512 production path unless there is new hardware-specific evidence and an isolated benchmark branch.

Raw Scalar As A Replacement For AVX2

Raw scalar integer decode is useful because it provides the same semantics as the AVX2 path on non-AVX2 platforms. It was not faster than lookup at larger trusted reader chunk sizes, so it should not be treated as a separate optimized production mode.

Arbitrary Selected-Subset SIMD

Sparse or shuffled subsets require gather/scatter-style access through selected_file_indices or file_to_selected_index. That defeats the contiguous-load/contiguous-store shape that made the kept AVX2 paths fast. Only the dense single-run subset case is optimized.

Row-Major Decode SIMD

Row-major full-sample unphased 8-bit decode writes one dosage per sample with a stride of variant_count. SIMD decode could speed up probability math, but stores are not contiguous for a single variant. The current row-major tile copy already uses contiguous copy_nonoverlapping after decoding the tile. Do not prioritize this before new profile data.

Generic Bit-Packed Or Phased BGEN SIMD

The generic BGEN path supports phased variants and probability bit widths from 1 to 32. It is branchy and uses a bit reader abstraction. SIMD here would require specialized kernels per bit width and phased mode. Avoid this unless real input profiles show non-8-bit or phased BGEN is hot.

Missing-Value Imputation SIMD

Variant-major imputation scans rows to replace NaN with the mean only when missing values are present. The main trusted and no-missing paths avoid it. This is not a good SIMD target without missing-heavy profiles.

Non-Numeric Rust Paths

Chunk stats post-processing, Parquet output writing, sample alignment, LOCO prediction loading, and BGEN metadata or index parsing are not current hand-written SIMD targets. They are dominated by control flow, library code, allocation, strings, I/O, joins, or small per-variant scalar work.

Correctness Coverage

The SIMD work added or exercised Rust coverage for:

• sample counts around vector boundaries: 0, 1, 7, 8, 15, 16, 17, 31, 32, 33
• all dosage 0
• all dosage 2
• alternating raw dosage values
• rare-variant-like values
• deterministic valid probability pairs where p0 + p1 <= 255
• scalar versus AVX2 output equality or bounded deltas
• integer-derived sum and square-sum parity
• zero/nonzero and genotype-threshold counts
• dense contiguous selected-subset detection
• trusted and non-trusted variant-major BGEN decode behavior

Validation commands used in the final pass:

cargo test --lib
cargo test bgen --lib
cargo test trusted --lib
cargo clippy --workspace --all-targets -- -D warnings -W clippy::pedantic
cargo fmt --all --check

Benchmarking Notes

Use cantor or another CPU worker node for CPU benchmarks. Do not run heavy compilation or Criterion benchmarks on the head node.

Representative command shape:

srun --nodelist=cantor --cpus-per-task=40 --mem=16G --time=00:45:00 bash -lc 'cd /mnt/beegfs/kirill/Projects/g && . tooling/server/server_env.sh && CARGO_BUILD_JOBS=40 cargo bench --bench bgen_read <benchmark-filter> -- --sample-size 10 --measurement-time 1 --warm-up-time 1'

Keep benchmark-only mode switches temporary. Once a faster implementation is selected, remove alternate public runtime switches and leave a single production dispatch policy.