From d6695fcaacdef249407cbd85ffa283ece9a9a98c Mon Sep 17 00:00:00 2001 From: exaclior Date: Wed, 15 Jul 2026 11:06:14 +0800 Subject: [PATCH 1/5] Add complex TDVP GEMM benchmark plan --- Cargo.toml | 4 + Makefile | 6 +- benches/complex_tdvp.rs | 181 +++++++++++++++++++++++++++++++ benchmarks/complex_tdvp.md | 211 +++++++++++++++++++++++++++++++++++++ 4 files changed, 401 insertions(+), 1 deletion(-) create mode 100644 benches/complex_tdvp.rs create mode 100644 benchmarks/complex_tdvp.md diff --git a/Cargo.toml b/Cargo.toml index b506867..017d068 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -99,6 +99,10 @@ members = [".", "omeinsum-cli"] name = "binary" harness = false +[[bench]] +name = "complex_tdvp" +harness = false + [[bench]] name = "network" harness = false diff --git a/Makefile b/Makefile index d84000e..750bded 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ BENCH_WARMUP ?= 5 BENCH_DIM ?= 128 BENCH_BATCH ?= 24 -.PHONY: all build build-debug cargo-check check test test-gpu test-gpu-tropical test-release bench bench-binary bench-network bench-cpu-contract bench-julia bench-compare docs clean help +.PHONY: all build build-debug cargo-check check test test-gpu test-gpu-tropical test-release bench bench-binary bench-complex-tdvp bench-network bench-cpu-contract bench-julia bench-compare docs clean help .PHONY: setup setup-rust .PHONY: docs-build docs-serve docs-book docs-book-serve .PHONY: fmt fmt-check clippy lint coverage @@ -48,6 +48,7 @@ help: @echo "Benchmark targets:" @echo " bench - Run all Rust benchmarks" @echo " bench-binary - Run binary contraction benchmarks" + @echo " bench-complex-tdvp - Run Complex64 TDVP-shaped binary contractions" @echo " bench-network - Run tensor-network benchmarks" @echo " bench-cpu-contract - Run the CPU contraction benchmark example" @echo " Override BENCH_SCENARIO, BENCH_ITERATIONS, BENCH_WARMUP," @@ -154,6 +155,9 @@ bench: bench-binary: cargo bench --bench binary +bench-complex-tdvp: + cargo bench --bench complex_tdvp + bench-network: cargo bench --bench network cargo run --release --example profile_network -- --scenario 3reg_150 --iterations 1 --output benchmarks/data/rust_network_timings.json diff --git a/benches/complex_tdvp.rs b/benches/complex_tdvp.rs new file mode 100644 index 0000000..df2ebae --- /dev/null +++ b/benches/complex_tdvp.rs @@ -0,0 +1,181 @@ +use std::time::Duration; + +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use num_complex::Complex64; +use omeinsum::{Cpu, Standard, Tensor}; + +const MPO_BOND_DIM: usize = 9; +const CHI_VALUES: [usize; 3] = [32, 64, 128]; + +#[derive(Clone, Copy)] +struct TdvpContraction { + name: &'static str, + left_shape: fn(usize) -> Vec, + right_shape: fn(usize) -> Vec, + left_labels: &'static [usize], + right_labels: &'static [usize], + output_labels: &'static [usize], + output_shape: fn(usize) -> Vec, + scalar_fmas: fn(usize) -> u64, +} + +struct PreparedContraction { + left: Tensor, + right: Tensor, + left_labels: &'static [usize], + right_labels: &'static [usize], + output_labels: &'static [usize], +} + +fn patterned_complex(len: usize, seed: usize) -> Vec { + (0..len) + .map(|index| { + let real = ((index.wrapping_mul(17) + seed.wrapping_mul(31)) % 257) as f64; + let imag = ((index.wrapping_mul(29) + seed.wrapping_mul(13)) % 251) as f64; + Complex64::new((real - 128.0) / 37.0, (imag - 125.0) / 41.0) + }) + .collect() +} + +fn prepare(case: TdvpContraction, chi: usize) -> PreparedContraction { + let left_shape = (case.left_shape)(chi); + let right_shape = (case.right_shape)(chi); + let left_len = left_shape.iter().product(); + let right_len = right_shape.iter().product(); + + PreparedContraction { + left: Tensor::from_data(&patterned_complex(left_len, 1), &left_shape), + right: Tensor::from_data(&patterned_complex(right_len, 2), &right_shape), + left_labels: case.left_labels, + right_labels: case.right_labels, + output_labels: case.output_labels, + } +} + +fn run(prepared: &PreparedContraction) -> Tensor { + prepared.left.contract_binary::>( + &prepared.right, + prepared.left_labels, + prepared.right_labels, + prepared.output_labels, + ) +} + +fn h1_left_shape(chi: usize) -> Vec { + vec![chi, MPO_BOND_DIM, chi] +} + +fn h1_wavefunction_shape(chi: usize) -> Vec { + vec![chi, 2, chi] +} + +fn h1_left_output_shape(chi: usize) -> Vec { + vec![MPO_BOND_DIM, chi, 2, chi] +} + +fn h1_right_intermediate_shape(chi: usize) -> Vec { + vec![chi, chi, 2, MPO_BOND_DIM] +} + +fn h1_right_output_shape(chi: usize) -> Vec { + vec![chi, 2, chi] +} + +fn right_environment_shape(chi: usize) -> Vec { + vec![chi, MPO_BOND_DIM, chi] +} + +fn h2_wavefunction_shape(chi: usize) -> Vec { + vec![chi, 2, 2, chi] +} + +fn h2_left_output_shape(chi: usize) -> Vec { + vec![MPO_BOND_DIM, chi, 2, 2, chi] +} + +fn h2_right_intermediate_shape(chi: usize) -> Vec { + vec![chi, 2, chi, 2, MPO_BOND_DIM] +} + +fn h2_right_output_shape(chi: usize) -> Vec { + vec![chi, 2, 2, chi] +} + +const TDVP_CONTRACTIONS: [TdvpContraction; 4] = [ + TdvpContraction { + name: "h1-left-environment", + left_shape: h1_left_shape, + right_shape: h1_wavefunction_shape, + left_labels: &[0, 1, 2], + right_labels: &[0, 3, 4], + output_labels: &[1, 2, 3, 4], + output_shape: h1_left_output_shape, + scalar_fmas: |chi| 2 * MPO_BOND_DIM as u64 * (chi as u64).pow(3), + }, + TdvpContraction { + name: "h1-right-environment", + left_shape: h1_right_intermediate_shape, + right_shape: right_environment_shape, + left_labels: &[2, 4, 5, 6], + right_labels: &[4, 6, 7], + output_labels: &[2, 5, 7], + output_shape: h1_right_output_shape, + scalar_fmas: |chi| 2 * MPO_BOND_DIM as u64 * (chi as u64).pow(3), + }, + TdvpContraction { + name: "h2-left-environment", + left_shape: h1_left_shape, + right_shape: h2_wavefunction_shape, + left_labels: &[0, 1, 2], + right_labels: &[0, 3, 4, 5], + output_labels: &[1, 2, 3, 4, 5], + output_shape: h2_left_output_shape, + scalar_fmas: |chi| 4 * MPO_BOND_DIM as u64 * (chi as u64).pow(3), + }, + TdvpContraction { + name: "h2-right-environment", + left_shape: h2_right_intermediate_shape, + right_shape: right_environment_shape, + left_labels: &[2, 6, 5, 8, 9], + right_labels: &[5, 9, 10], + output_labels: &[2, 6, 8, 10], + output_shape: h2_right_output_shape, + scalar_fmas: |chi| 4 * MPO_BOND_DIM as u64 * (chi as u64).pow(3), + }, +]; + +fn bench_tdvp_complex_contractions(criterion: &mut Criterion) { + let mut group = criterion.benchmark_group("tdvp-complex-binary"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(2)); + group.measurement_time(Duration::from_secs(10)); + + for case in TDVP_CONTRACTIONS { + for chi in CHI_VALUES { + let prepared = prepare(case, chi); + let actual_shape = run(&prepared).shape().to_vec(); + assert_eq!( + actual_shape, + (case.output_shape)(chi), + "benchmark contraction produced the wrong shape" + ); + + group.throughput(Throughput::Elements((case.scalar_fmas)(chi))); + group.bench_with_input( + BenchmarkId::new(case.name, format!("chi{chi}-d{MPO_BOND_DIM}")), + &prepared, + |bencher, prepared| { + bencher.iter(|| { + let output = run(black_box(prepared)); + black_box(output); + }); + }, + ); + } + } + + group.finish(); +} + +criterion_group!(benches, bench_tdvp_complex_contractions); +criterion_main!(benches); diff --git a/benchmarks/complex_tdvp.md b/benchmarks/complex_tdvp.md new file mode 100644 index 0000000..d407b25 --- /dev/null +++ b/benchmarks/complex_tdvp.md @@ -0,0 +1,211 @@ +# Optimized CPU GEMM for standard complex contractions + +## Status + +- Branch: `perf/complex-gemm` +- Owner workload: `rydbergsim-rs` issue #272 (Keesling 2019 TDVP reproduction) +- Target repository: `tensor4all/omeinsum-rs` +- Current state: benchmark workload added; baseline measurement and implementation pending + +## Problem + +OMEinsum supports `Standard` and `Standard` semantically, but its CPU GEMM dispatch only sends `Standard` and `Standard` to faer. Complex contractions therefore fall through to `generic_gemm`, a scalar triple loop. + +The contraction planner is not the problem. It already: + +1. classifies free, contracted, and batch labels; +2. recognizes directly usable matrix layouts; +3. materializes permutations when required; +4. lowers binary contractions to GEMM. + +The missing piece is optimized complex dispatch at the final GEMM boundary. + +This gap dominates local Hamiltonian applications in matrix-product-state TDVP. The motivating finite-range Rydberg Hamiltonian has MPO bond dimension `D=9`; the expensive contractions repeatedly multiply complex matrices whose dimensions scale with MPS bond dimension `chi`. + +## Scope + +### In scope + +- Optimized faer GEMM for `Standard` and `Standard` on CPU. +- Contiguous, strided/transposed, and batched execution paths. +- Concrete value tests against the generic implementation. +- A deterministic benchmark reproducing the TDVP contraction shapes. +- Before/after measurements on one pinned machine and toolchain. +- Downstream validation in `rydbergsim-rs` without its local mini-einsum workaround. + +### Out of scope + +- Contraction-order changes. +- Tropical algebra. +- CUDA/cuTENSOR. +- A new public API. +- TDVP-specific logic inside OMEinsum. +- Thread-pool policy changes; initial implementation should match the existing real-valued faer paths and use `Par::Seq`. + +## Target computations + +The benchmark at `benches/complex_tdvp.rs` calls the public `Tensor::contract_binary` path. It fixes the MPO bond dimension at `D=9` and exercises `chi in {32, 64, 128}`. + +| Case | Left shape | Right shape | Contracted modes | Equivalent GEMM work | +|---|---|---|---|---| +| `h1-left-environment` | `[chi,D,chi]` | `[chi,2,chi]` | one `chi` mode | `(D chi x chi) @ (chi x 2 chi)` | +| `h1-right-environment` | `[chi,chi,2,D]` | `[chi,D,chi]` | `chi*D` | `(2 chi x chi D) @ (chi D x chi)` | +| `h2-left-environment` | `[chi,D,chi]` | `[chi,2,2,chi]` | one `chi` mode | `(D chi x chi) @ (chi x 4 chi)` | +| `h2-right-environment` | `[chi,2,chi,2,D]` | `[chi,D,chi]` | `chi*D` | `(4 chi x chi D) @ (chi D x chi)` | + +These are library-level tensor contractions, not a paper simulator embedded in OMEinsum. The labels and shapes come from the downstream workload; the benchmark data are deterministic synthetic complex values. + +## Existing dispatch gap + +Four CPU entry points need complex support: + +1. `Cpu::gemm_internal` — contiguous GEMM. +2. `Cpu::gemm_standard_layout_internal` — directly usable strided/transposed layouts. +3. `Cpu::gemm_batched_internal` — materialized batched GEMM. +4. `Cpu::gemm_batched_standard_layout_internal` — direct batched layouts. + +Today each recognizes only `Standard` and `Standard`. Complex types return `None` from the layout fast path and eventually reach `generic_gemm`. + +## Implementation route + +### Stage 0 — establish the baseline + +Run each benchmark filter separately in release mode. The scalar baseline can be slow at `chi=128`, so do not start with the whole matrix. + +Required order: + +1. `chi=32` sanity run. +2. `chi=64` primary baseline. +3. `chi=128` paper-scale kernel baseline after the first two complete. + +All performance runs must follow this repository's runscribe protocol. Declare a goal and hypothesis, fill its `Why this path`, obtain owner approval, then wrap every command with `runscribe run`. + +Example command after a hypothesis code exists: + +```bash +runscribe run --hyp --tag baseline-chi64 -- \ + cargo bench --bench complex_tdvp -- 'chi64-d9' +``` + +Record the run directory, median time, dispersion, and throughput for every case. + +### Stage 1 — contiguous complex GEMM + +Add faer helpers for `Complex32` and `Complex64`, then dispatch them from `gemm_internal` using the same `TypeId` pattern as real scalars. + +Prefer borrowed column-major `MatRef` inputs and one owned output allocation. Do not copy inputs into `Mat::from_fn` unless faer's type constraints make borrowing impossible. + +Correctness tests: + +- hand-checkable 2x2 complex products; +- rectangular products; +- both complex widths; +- comparison with `generic_gemm` on deterministic inputs; +- zero/degenerate dimensions if currently supported by the real path. + +### Stage 2 — direct layout fast path + +Add complex equivalents of the existing layout helpers and dispatch them from `gemm_standard_layout_internal`. + +Required layouts: + +- ordinary column-major; +- transposed right operand; +- non-unit positive strides already accepted by `faer_mat_ref`; +- negative strides if the existing layout contract permits them. + +The output should require one allocation. Inputs must remain borrowed. + +### Stage 3 — batched paths + +Add complex dispatch to both batched entry points. + +Allocate the complete output once and write each batch into its destination slice. Avoid the current pattern of allocating a temporary result vector per batch and copying it into the final output; if this cleanup is generalized to real scalars, benchmark real paths to prove no regression. + +Test both layouts used by OMEinsum: + +- column-major batch views; +- batch-major views; +- transposed operand within a batch; +- batch size one and multiple batches. + +### Stage 4 — remove accidental duplication + +Only after all four paths work, inspect whether the f32/f64/c32/c64 helpers can share a small generic implementation. Do not build a trait hierarchy merely to avoid four short dispatch arms. The simple version wins unless genericization clearly reduces unsafe casts and layout code. + +### Stage 5 — downstream validation + +In `rydbergsim-rs`: + +1. temporarily pin OMEinsum to the candidate commit; +2. remove the local `contract_binary_faer` implementation and restore ordinary OMEinsum calls; +3. run the TDVP kernel reference tests; +4. run the N=21, chi=64 saturated-step benchmark; +5. rerun the partial N=51, chi=128 probe only after the kernel result justifies its cost. + +This confirms that the upstream optimization survives contraction planning, environment updates, and Krylov repetition. + +## Benchmark protocol + +### Controlled variables + +- Same host, CPU governor, Rust toolchain, commit mode, and feature set. +- Release profile. +- `D=9` and identical deterministic input values. +- Criterion sample size, warm-up, and measurement duration fixed by the benchmark. +- Run one benchmark process at a time. + +### Measurements + +For every case and `chi`: + +- Criterion median latency and confidence interval. +- Scalar fused multiply-add count reported as throughput. +- Candidate/baseline speedup. +- Peak resident memory for the complete benchmark process when practical. +- Allocation count in focused unit tests for direct-layout helpers. + +### Acceptance criteria + +Correctness is mandatory: + +- `make check` passes. +- Complex fast paths agree with generic GEMM within dtype-appropriate relative and absolute tolerances. +- Existing f32/f64, tropical, repeated-label, scalar-output, and backward tests remain green. + +Performance acceptance for the motivating workload: + +- At least 5x lower median latency than the scalar baseline for all four `chi=64, D=9` contractions. +- No target case regresses relative to baseline. +- `chi=128` shows the same direction of improvement and completes without pathological memory growth. +- Existing real-valued binary benchmarks do not regress by more than 5% beyond measurement noise. + +The 5x threshold is a minimum useful result, not a promise that it makes the full Keesling grid feasible. End-to-end TDVP timing decides that separately. + +## Risks and controls + +| Risk | Control | +|---|---| +| `TypeId` dispatch and `transmute` mismatch scalar storage | Mirror existing real dispatch, keep casts adjacent to checked type IDs, and test both complex widths. | +| Strided `MatRef` points outside backing storage | Reuse `layout_offset_bounds`/`faer_mat_ref`; add transposed and negative-stride regressions. | +| Complex convention accidentally conjugates an operand | Test against plain sum-product GEMM with nontrivial imaginary values; use ordinary views, never adjoints. | +| Batched offsets are wrong | Compare every batch element against independent unbatched contractions. | +| Faer parallelism oversubscribes downstream workloads | Keep `Par::Seq` in this change; treat parallel policy as a separate measured hypothesis. | +| Benchmark only measures synthetic GEMM | Validate the candidate commit in the downstream TDVP step before claiming user-visible speedup. | +| Local workaround and upstream path drift | Delete the downstream workaround once the candidate passes; maintain one implementation in OMEinsum. | + +## Deliverables + +- `benches/complex_tdvp.rs` with the four target contractions. +- Complex faer dispatch in all four CPU GEMM routes. +- Unit and integration regressions with concrete complex values. +- runscribe baseline and candidate runs, including the surviving alternative analysis. +- Before/after summary in the pull request. +- Downstream `rydbergsim-rs` validation result and candidate revision. + +## Decision points + +1. After baseline: confirm the scalar fallback is reproduced at target shapes. +2. After contiguous/layout implementation: decide whether batched cleanup belongs in the same PR. +3. After OMEinsum benchmarks: decide whether downstream N=51 probing is justified. +4. After downstream validation: merge, revise, or abandon based on measured end-to-end gain. From eaf29fe4d0cf5eaba5f582d040e4ce3b0c3d2344 Mon Sep 17 00:00:00 2001 From: exaclior Date: Wed, 15 Jul 2026 11:31:16 +0800 Subject: [PATCH 2/5] Link complex GEMM tracking issue --- benchmarks/complex_tdvp.md | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/complex_tdvp.md b/benchmarks/complex_tdvp.md index d407b25..0ff53ad 100644 --- a/benchmarks/complex_tdvp.md +++ b/benchmarks/complex_tdvp.md @@ -5,6 +5,7 @@ - Branch: `perf/complex-gemm` - Owner workload: `rydbergsim-rs` issue #272 (Keesling 2019 TDVP reproduction) - Target repository: `tensor4all/omeinsum-rs` +- Tracking issue: [`tensor4all/omeinsum-rs#53`](https://github.com/tensor4all/omeinsum-rs/issues/53) - Current state: benchmark workload added; baseline measurement and implementation pending ## Problem From 6c273ae2f66c7997b2b6052c9a02b8abb2889b36 Mon Sep 17 00:00:00 2001 From: exaclior Date: Wed, 15 Jul 2026 13:03:54 +0800 Subject: [PATCH 3/5] Optimize complex CPU GEMM with faer --- benchmarks/complex_tdvp.md | 45 +- src/backend/cpu/contract.rs | 16 +- src/backend/cpu/mod.rs | 846 ++++++++++++++++++++++++++--------- tests/suites/binary_rules.rs | 19 +- 4 files changed, 722 insertions(+), 204 deletions(-) diff --git a/benchmarks/complex_tdvp.md b/benchmarks/complex_tdvp.md index 0ff53ad..0851c79 100644 --- a/benchmarks/complex_tdvp.md +++ b/benchmarks/complex_tdvp.md @@ -6,7 +6,7 @@ - Owner workload: `rydbergsim-rs` issue #272 (Keesling 2019 TDVP reproduction) - Target repository: `tensor4all/omeinsum-rs` - Tracking issue: [`tensor4all/omeinsum-rs#53`](https://github.com/tensor4all/omeinsum-rs/issues/53) -- Current state: benchmark workload added; baseline measurement and implementation pending +- Current state: CPU implementation and upstream validation complete; downstream validation pending ## Problem @@ -166,6 +166,49 @@ For every case and `chi`: - Peak resident memory for the complete benchmark process when practical. - Allocation count in focused unit tests for direct-layout helpers. +## Validation record (2026-07-15) + +Measurements ran on remote host `6xa800` (2-socket Intel Xeon Platinum 8378A, +128 logical CPUs) with Rust 1.88.0. Baseline and candidate used the same release +profile, benchmark inputs, and Criterion settings. Criterion artifacts are under +`~/projects/omeinsum-rs/target/criterion/tdvp-complex-binary` on that host. + +The repository-mandated `runscribe` executable was unavailable locally and on +the remote host; the similarly named PyPI package is an unrelated terminal +recorder. With the owner-requested remote workflow, runs were submitted through +`easy-ssh submit` and their full job output was captured instead. There is +therefore no runscribe run directory for these measurements. + +Times below are Criterion point estimates in milliseconds; parenthesized ranges +are the default 95% confidence intervals. Speedup is baseline divided by +candidate. + +| Case | chi | Scalar baseline | faer candidate | Speedup | Candidate throughput | +|---|---:|---:|---:|---:|---:| +| `h1-left-environment` | 32 | 1.0503 (1.0419-1.0651) | 0.084151 (0.084054-0.084254) | 12.48x | 7.0091 Gelem/s | +| `h1-right-environment` | 32 | 1.1197 (1.1158-1.1286) | 0.13741 (0.13709-0.13759) | 8.15x | 4.2924 Gelem/s | +| `h2-left-environment` | 32 | 1.8304 (1.8032-1.9023) | 0.15972 (0.15964-0.15988) | 11.46x | 7.3856 Gelem/s | +| `h2-right-environment` | 32 | 2.4375 (2.3791-2.5324) | 0.30246 (0.30238-0.30255) | 8.06x | 3.9001 Gelem/s | +| `h1-left-environment` | 64 | 8.0044 (7.9638-8.1163) | 0.61142 (0.61091-0.61199) | 13.09x | 7.7175 Gelem/s | +| `h1-right-environment` | 64 | 12.281 (11.996-12.757) | 0.96001 (0.95767-0.96541) | 12.79x | 4.9151 Gelem/s | +| `h2-left-environment` | 64 | 16.475 (15.980-17.505) | 1.1750 (1.1745-1.1756) | 14.02x | 8.0319 Gelem/s | +| `h2-right-environment` | 64 | 33.899 (32.550-35.217) | 1.9024 (1.8967-1.9183) | 17.82x | 4.9606 Gelem/s | +| `h1-left-environment` | 128 | 92.408 (88.993-95.673) | 4.5194 (4.4687-4.6194) | 20.45x | 8.3525 Gelem/s | +| `h1-right-environment` | 128 | 130.87 (129.61-132.46) | 5.9194 (5.9064-5.9326) | 22.11x | 6.3772 Gelem/s | +| `h2-left-environment` | 128 | 164.62 (163.87-166.52) | 8.6727 (8.6685-8.6766) | 18.98x | 8.7051 Gelem/s | +| `h2-right-environment` | 128 | 248.54 (247.66-249.35) | 11.647 (11.623-11.670) | 21.34x | 6.4822 Gelem/s | + +All four chi=64 cases exceed the required 5x threshold. The existing real f32 +`binary` benchmark suite was also compared against a clean `eaf29fe` worktree. +Two noisy initial outliers were rerun: `high_d_12x12_contract_6` measured 40.35 +microseconds versus 40.21 microseconds (+0.34%), and +`high_d_20x20_contract_9` measured 28.51 milliseconds versus 28.83 milliseconds +(-1.1%). No repeatable real-valued regression exceeded 5%. + +`make check` passed after the final change: clippy with `tropical parallel`, 158 +library tests passed (11 ignored), 333 integration tests passed, and 16 doctests +passed (4 ignored). + ### Acceptance criteria Correctness is mandatory: diff --git a/src/backend/cpu/contract.rs b/src/backend/cpu/contract.rs index d8ebcc7..3b1ff54 100644 --- a/src/backend/cpu/contract.rs +++ b/src/backend/cpu/contract.rs @@ -3,9 +3,22 @@ use std::collections::HashSet; use crate::backend::contract_plan::{ - classify_modes, compute_permutation, mode_position, product_of_dims, reduce_trace, + classify_modes, compute_permutation, mode_position, product_of_dims as nonzero_product_of_dims, + reduce_trace, }; +// Empty mode groups are scalar dimensions; actual zero-sized modes stay zero. +fn product_of_dims(modes: &[i32], all_modes: &[i32], shape: &[usize]) -> usize { + if modes + .iter() + .any(|&mode| shape[mode_position(all_modes, mode)] == 0) + { + 0 + } else { + nonzero_product_of_dims(modes, all_modes, shape) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] enum MaterializationPlan { NoCopy, @@ -764,6 +777,7 @@ mod tests { assert_eq!(product_of_dims(&[0], modes, shape), 2); assert_eq!(product_of_dims(&[1, 2], modes, shape), 12); assert_eq!(product_of_dims(&[], modes, shape), 1); + assert_eq!(product_of_dims(&[1], modes, &[2, 0, 4]), 0); } #[test] diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index 73eefd9..63886e8 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -5,6 +5,7 @@ mod contract; use super::traits::{Backend, BackendScalar, Storage}; use crate::algebra::{Algebra, Scalar, Standard}; +use num_complex::{Complex32, Complex64}; use std::any::TypeId; /// CPU backend using Vec storage. @@ -21,7 +22,6 @@ pub(crate) struct MatrixLayout<'a, T> { } impl<'a, T> MatrixLayout<'a, T> { - #[cfg(test)] pub(crate) fn column_major(data: &'a [T], rows: usize, cols: usize) -> Self { Self { data, @@ -193,7 +193,7 @@ impl Cpu { row_stride: b.row_stride, col_stride: b.col_stride, }; - let result = faer_gemm_f32_layout(a_f32, b_f32); + let result = faer_gemm_layout(a_f32, b_f32); return Some(unsafe { std::mem::transmute::, Vec>(result) }); } if TypeId::of::() == TypeId::of::>() { @@ -211,9 +211,47 @@ impl Cpu { row_stride: b.row_stride, col_stride: b.col_stride, }; - let result = faer_gemm_f64_layout(a_f64, b_f64); + let result = faer_gemm_layout(a_f64, b_f64); return Some(unsafe { std::mem::transmute::, Vec>(result) }); } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: TypeId proves A::Scalar is Complex32 for this branch. + let a_c32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex32]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_c32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex32]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let result = faer_gemm_layout(a_c32, b_c32); + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: TypeId proves A::Scalar is Complex64 for this branch. + let a_c64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex64]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_c64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex64]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let result = faer_gemm_layout(a_c64, b_c64); + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } None } @@ -224,8 +262,6 @@ impl Cpu { a: MatrixLayout<'_, A::Scalar>, b: MatrixLayout<'_, A::Scalar>, ) -> Option> { - let c_batch_stride = a.rows * b.cols; - if TypeId::of::() == TypeId::of::>() { let a_f32 = MatrixLayout { data: unsafe { std::mem::transmute::<&[A::Scalar], &[f32]>(a.data) }, @@ -241,17 +277,7 @@ impl Cpu { row_stride: b.row_stride, col_stride: b.col_stride, }; - let mut result = vec![0.0f32; batch_size * c_batch_stride]; - - for batch in 0..batch_size { - let c_offset = batch * c_batch_stride; - let c_batch = faer_gemm_f32_layout( - matrix_layout_batch_view(a_f32, batch), - matrix_layout_batch_view(b_f32, batch), - ); - result[c_offset..c_offset + c_batch_stride].copy_from_slice(&c_batch); - } - + let result = faer_batched_gemm_layout(batch_size, a_f32, b_f32); return Some(unsafe { std::mem::transmute::, Vec>(result) }); } if TypeId::of::() == TypeId::of::>() { @@ -269,19 +295,47 @@ impl Cpu { row_stride: b.row_stride, col_stride: b.col_stride, }; - let mut result = vec![0.0f64; batch_size * c_batch_stride]; - - for batch in 0..batch_size { - let c_offset = batch * c_batch_stride; - let c_batch = faer_gemm_f64_layout( - matrix_layout_batch_view(a_f64, batch), - matrix_layout_batch_view(b_f64, batch), - ); - result[c_offset..c_offset + c_batch_stride].copy_from_slice(&c_batch); - } - + let result = faer_batched_gemm_layout(batch_size, a_f64, b_f64); return Some(unsafe { std::mem::transmute::, Vec>(result) }); } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: TypeId proves A::Scalar is Complex32 for this branch. + let a_c32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex32]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_c32 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex32]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let result = faer_batched_gemm_layout(batch_size, a_c32, b_c32); + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: TypeId proves A::Scalar is Complex64 for this branch. + let a_c64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex64]>(a.data) }, + rows: a.rows, + cols: a.cols, + row_stride: a.row_stride, + col_stride: a.col_stride, + }; + let b_c64 = MatrixLayout { + data: unsafe { std::mem::transmute::<&[A::Scalar], &[Complex64]>(b.data) }, + rows: b.rows, + cols: b.cols, + row_stride: b.row_stride, + col_stride: b.col_stride, + }; + let result = faer_batched_gemm_layout(batch_size, a_c64, b_c64); + return Some(unsafe { std::mem::transmute::, Vec>(result) }); + } None } @@ -301,20 +355,34 @@ impl Cpu { b: &[A::Scalar], n: usize, ) -> Vec { - // Fast path: faer for Standard f32/f64 + // Fast path: faer for native real and complex Standard scalars. if TypeId::of::() == TypeId::of::>() { // SAFETY: A::Scalar is f32 when A is Standard let a_f32: &[f32] = unsafe { std::mem::transmute(a) }; let b_f32: &[f32] = unsafe { std::mem::transmute(b) }; - let result = faer_gemm_f32(a_f32, m, k, b_f32, n); + let result = faer_gemm(a_f32, m, k, b_f32, n); return unsafe { std::mem::transmute::, Vec>(result) }; } if TypeId::of::() == TypeId::of::>() { let a_f64: &[f64] = unsafe { std::mem::transmute(a) }; let b_f64: &[f64] = unsafe { std::mem::transmute(b) }; - let result = faer_gemm_f64(a_f64, m, k, b_f64, n); + let result = faer_gemm(a_f64, m, k, b_f64, n); return unsafe { std::mem::transmute::, Vec>(result) }; } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: A::Scalar is Complex32 when A is Standard. + let a_c32: &[Complex32] = unsafe { std::mem::transmute(a) }; + let b_c32: &[Complex32] = unsafe { std::mem::transmute(b) }; + let result = faer_gemm(a_c32, m, k, b_c32, n); + return unsafe { std::mem::transmute::, Vec>(result) }; + } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: A::Scalar is Complex64 when A is Standard. + let a_c64: &[Complex64] = unsafe { std::mem::transmute(a) }; + let b_c64: &[Complex64] = unsafe { std::mem::transmute(b) }; + let result = faer_gemm(a_c64, m, k, b_c64, n); + return unsafe { std::mem::transmute::, Vec>(result) }; + } // Try to use optimized tropical-gemm if available #[cfg(feature = "tropical-kernels")] @@ -423,15 +491,29 @@ impl Cpu { if TypeId::of::() == TypeId::of::>() { let a_f32: &[f32] = unsafe { std::mem::transmute(a) }; let b_f32: &[f32] = unsafe { std::mem::transmute(b) }; - let result = standard_batched_gemm_f32(a_f32, batch_size, m, k, b_f32, n); + let result = standard_batched_gemm(a_f32, batch_size, m, k, b_f32, n); return unsafe { std::mem::transmute::, Vec>(result) }; } if TypeId::of::() == TypeId::of::>() { let a_f64: &[f64] = unsafe { std::mem::transmute(a) }; let b_f64: &[f64] = unsafe { std::mem::transmute(b) }; - let result = standard_batched_gemm_f64(a_f64, batch_size, m, k, b_f64, n); + let result = standard_batched_gemm(a_f64, batch_size, m, k, b_f64, n); return unsafe { std::mem::transmute::, Vec>(result) }; } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: A::Scalar is Complex32 when A is Standard. + let a_c32: &[Complex32] = unsafe { std::mem::transmute(a) }; + let b_c32: &[Complex32] = unsafe { std::mem::transmute(b) }; + let result = standard_batched_gemm(a_c32, batch_size, m, k, b_c32, n); + return unsafe { std::mem::transmute::, Vec>(result) }; + } + if TypeId::of::() == TypeId::of::>() { + // SAFETY: A::Scalar is Complex64 when A is Standard. + let a_c64: &[Complex64] = unsafe { std::mem::transmute(a) }; + let b_c64: &[Complex64] = unsafe { std::mem::transmute(b) }; + let result = standard_batched_gemm(a_c64, batch_size, m, k, b_c64, n); + return unsafe { std::mem::transmute::, Vec>(result) }; + } let a_batch_stride = m * k; let b_batch_stride = k * n; @@ -630,79 +712,45 @@ impl Backend for Cpu { } } -/// GEMM using faer for f32 (column-major layout). +/// GEMM using faer's native real and complex kernels. /// -/// Computes C = A @ B where A is m×k, B is k×n, C is m×n. -fn faer_gemm_f32(a: &[f32], m: usize, k: usize, b: &[f32], n: usize) -> Vec { - use faer::Mat; - - // Create matrices from column-major data - // Column-major: element (i, j) is at index j * nrows + i - let a_mat = Mat::from_fn(m, k, |i, j| a[j * m + i]); - let b_mat = Mat::from_fn(k, n, |i, j| b[j * k + i]); - - // Multiply - let c_mat = &a_mat * &b_mat; - - // Convert back to column-major Vec - let mut c = vec![0.0f32; m * n]; - for j in 0..n { - for i in 0..m { - c[j * m + i] = c_mat[(i, j)]; - } - } - c -} - -fn faer_gemm_f32_layout(a: MatrixLayout<'_, f32>, b: MatrixLayout<'_, f32>) -> Vec { - let mut c = vec![0.0f32; a.rows * b.cols]; - faer_gemm_f32_layout_into(a, b, &mut c); - c -} - -fn faer_gemm_f32_layout_into(a: MatrixLayout<'_, f32>, b: MatrixLayout<'_, f32>, c: &mut [f32]) { - use faer::{linalg::matmul::matmul, Accum, MatMut, Par}; - - assert_eq!(c.len(), a.rows * b.cols); - let a_mat = faer_mat_ref(a); - let b_mat = faer_mat_ref(b); - let mut c_mat = - unsafe { MatMut::from_raw_parts_mut(c.as_mut_ptr(), a.rows, b.cols, 1, a.rows as isize) }; - matmul( - c_mat.as_mut(), - Accum::Replace, - a_mat, - b_mat, - 1.0f32, - Par::Seq, - ); -} - -/// GEMM using faer for f64 (column-major layout). -fn faer_gemm_f64(a: &[f64], m: usize, k: usize, b: &[f64], n: usize) -> Vec { - use faer::Mat; - - let a_mat = Mat::from_fn(m, k, |i, j| a[j * m + i]); - let b_mat = Mat::from_fn(k, n, |i, j| b[j * k + i]); - - let c_mat = &a_mat * &b_mat; - - let mut c = vec![0.0f64; m * n]; - for j in 0..n { - for i in 0..m { - c[j * m + i] = c_mat[(i, j)]; - } - } - c +/// Inputs and output are column-major. Inputs are borrowed; only the output is +/// allocated. +fn faer_gemm(a: &[T], m: usize, k: usize, b: &[T], n: usize) -> Vec +where + T: faer::traits::ComplexField + Copy, +{ + faer_gemm_layout( + MatrixLayout { + data: a, + rows: m, + cols: k, + row_stride: 1, + col_stride: m as isize, + }, + MatrixLayout { + data: b, + rows: k, + cols: n, + row_stride: 1, + col_stride: k as isize, + }, + ) } -fn faer_gemm_f64_layout(a: MatrixLayout<'_, f64>, b: MatrixLayout<'_, f64>) -> Vec { - let mut c = vec![0.0f64; a.rows * b.cols]; - faer_gemm_f64_layout_into(a, b, &mut c); +fn faer_gemm_layout(a: MatrixLayout<'_, T>, b: MatrixLayout<'_, T>) -> Vec +where + T: faer::traits::ComplexField + Copy, +{ + let mut c = vec![faer::traits::math_utils::zero::(); a.rows * b.cols]; + faer_gemm_layout_into(a, b, &mut c); c } -fn faer_gemm_f64_layout_into(a: MatrixLayout<'_, f64>, b: MatrixLayout<'_, f64>, c: &mut [f64]) { +fn faer_gemm_layout_into(a: MatrixLayout<'_, T>, b: MatrixLayout<'_, T>, c: &mut [T]) +where + T: faer::traits::ComplexField + Copy, +{ use faer::{linalg::matmul::matmul, Accum, MatMut, Par}; assert_eq!(c.len(), a.rows * b.cols); @@ -715,63 +763,30 @@ fn faer_gemm_f64_layout_into(a: MatrixLayout<'_, f64>, b: MatrixLayout<'_, f64>, Accum::Replace, a_mat, b_mat, - 1.0f64, + faer::traits::math_utils::one::(), Par::Seq, ); } -fn standard_batched_gemm_f32( - a: &[f32], - batch_size: usize, - m: usize, - k: usize, - b: &[f32], - n: usize, -) -> Vec { - if should_use_standard_batched_gemm(batch_size, m, k, n) { - return faer_batched_gemm_f32(a, batch_size, m, k, b, n); - } - - let a_batch_stride = m * k; - let b_batch_stride = k * n; - let c_batch_stride = m * n; - let mut c = vec![0.0f32; batch_size * c_batch_stride]; - - for batch in 0..batch_size { - let a_offset = batch * a_batch_stride; - let b_offset = batch * b_batch_stride; - let c_offset = batch * c_batch_stride; - - for j in 0..n { - for i in 0..m { - let mut acc = 0.0f32; - for kk in 0..k { - acc += a[a_offset + kk * m + i] * b[b_offset + j * k + kk]; - } - c[c_offset + j * m + i] = acc; - } - } - } - - c -} - -fn standard_batched_gemm_f64( - a: &[f64], +fn standard_batched_gemm( + a: &[T], batch_size: usize, m: usize, k: usize, - b: &[f64], + b: &[T], n: usize, -) -> Vec { +) -> Vec +where + T: faer::traits::ComplexField + Copy + std::ops::AddAssign + std::ops::Mul, +{ if should_use_standard_batched_gemm(batch_size, m, k, n) { - return faer_batched_gemm_f64(a, batch_size, m, k, b, n); + return faer_batched_gemm(a, batch_size, m, k, b, n); } let a_batch_stride = m * k; let b_batch_stride = k * n; let c_batch_stride = m * n; - let mut c = vec![0.0f64; batch_size * c_batch_stride]; + let mut c = vec![faer::traits::math_utils::zero::(); batch_size * c_batch_stride]; for batch in 0..batch_size { let a_offset = batch * a_batch_stride; @@ -780,7 +795,7 @@ fn standard_batched_gemm_f64( for j in 0..n { for i in 0..m { - let mut acc = 0.0f64; + let mut acc = faer::traits::math_utils::zero::(); for kk in 0..k { acc += a[a_offset + kk * m + i] * b[b_offset + j * k + kk]; } @@ -792,38 +807,22 @@ fn standard_batched_gemm_f64( c } -fn faer_batched_gemm_f32( - a: &[f32], - batch_size: usize, - m: usize, - k: usize, - b: &[f32], - n: usize, -) -> Vec { +fn faer_batched_gemm(a: &[T], batch_size: usize, m: usize, k: usize, b: &[T], n: usize) -> Vec +where + T: faer::traits::ComplexField + Copy, +{ let a_batch_stride = m * k; let b_batch_stride = k * n; let c_batch_stride = m * n; - let mut c = vec![0.0f32; batch_size * c_batch_stride]; + let mut c = vec![faer::traits::math_utils::zero::(); batch_size * c_batch_stride]; for batch in 0..batch_size { let a_offset = batch * a_batch_stride; let b_offset = batch * b_batch_stride; let c_offset = batch * c_batch_stride; - faer_gemm_f32_layout_into( - MatrixLayout { - data: &a[a_offset..a_offset + a_batch_stride], - rows: m, - cols: k, - row_stride: 1, - col_stride: m as isize, - }, - MatrixLayout { - data: &b[b_offset..b_offset + b_batch_stride], - rows: k, - cols: n, - row_stride: 1, - col_stride: k as isize, - }, + faer_gemm_layout_into( + MatrixLayout::column_major(&a[a_offset..a_offset + a_batch_stride], m, k), + MatrixLayout::column_major(&b[b_offset..b_offset + b_batch_stride], k, n), &mut c[c_offset..c_offset + c_batch_stride], ); } @@ -831,38 +830,26 @@ fn faer_batched_gemm_f32( c } -fn faer_batched_gemm_f64( - a: &[f64], +fn faer_batched_gemm_layout( batch_size: usize, - m: usize, - k: usize, - b: &[f64], - n: usize, -) -> Vec { - let a_batch_stride = m * k; - let b_batch_stride = k * n; - let c_batch_stride = m * n; - let mut c = vec![0.0f64; batch_size * c_batch_stride]; + a: MatrixLayout<'_, T>, + b: MatrixLayout<'_, T>, +) -> Vec +where + T: faer::traits::ComplexField + Copy, +{ + assert_eq!(a.cols, b.rows, "GEMM inner dimensions must match"); + let c_batch_stride = a.rows * b.cols; + let mut c = vec![faer::traits::math_utils::zero::(); batch_size * c_batch_stride]; + if batch_size == 0 || a.rows == 0 || a.cols == 0 || b.cols == 0 { + return c; + } for batch in 0..batch_size { - let a_offset = batch * a_batch_stride; - let b_offset = batch * b_batch_stride; let c_offset = batch * c_batch_stride; - faer_gemm_f64_layout_into( - MatrixLayout { - data: &a[a_offset..a_offset + a_batch_stride], - rows: m, - cols: k, - row_stride: 1, - col_stride: m as isize, - }, - MatrixLayout { - data: &b[b_offset..b_offset + b_batch_stride], - rows: k, - cols: n, - row_stride: 1, - col_stride: k as isize, - }, + faer_gemm_layout_into( + matrix_layout_batch_view(a, batch), + matrix_layout_batch_view(b, batch), &mut c[c_offset..c_offset + c_batch_stride], ); } @@ -1158,6 +1145,27 @@ mod tests { use super::*; use crate::algebra::Standard; + fn generic_batched_gemm_for_test( + a: &[A::Scalar], + batch_size: usize, + m: usize, + k: usize, + b: &[A::Scalar], + n: usize, + ) -> Vec { + let mut result = Vec::with_capacity(batch_size * m * n); + for batch in 0..batch_size { + result.extend(generic_gemm::( + &a[batch * m * k..(batch + 1) * m * k], + m, + k, + &b[batch * k * n..(batch + 1) * k * n], + n, + )); + } + result + } + #[cfg(feature = "tropical")] use crate::algebra::MaxPlus; @@ -1174,17 +1182,266 @@ mod tests { assert_eq!(c, vec![7.0, 10.0, 15.0, 22.0]); } + #[test] + fn test_cpu_gemm_standard_complex_hand_checked() { + let cpu = Cpu; + let a32 = vec![ + Complex32::new(1.0, 1.0), + Complex32::new(3.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(4.0, -1.0), + ]; + let b32 = vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, -1.0), + Complex32::new(0.0, 1.0), + Complex32::new(-1.0, 0.0), + ]; + let expected32 = vec![ + Complex32::new(5.0, -1.0), + Complex32::new(10.0, -6.0), + Complex32::new(-3.0, 1.0), + Complex32::new(-4.0, 4.0), + ]; + assert_eq!( + cpu.gemm_internal::>(&a32, 2, 2, &b32, 2), + expected32 + ); + + let a64: Vec = a32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let b64: Vec = b32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let expected64: Vec = expected32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + assert_eq!( + cpu.gemm_internal::>(&a64, 2, 2, &b64, 2), + expected64 + ); + } + + #[test] + fn test_cpu_gemm_standard_complex_rectangular_matches_generic() { + let cpu = Cpu; + let (m, k, n) = (3usize, 2usize, 4usize); + let a32: Vec = (0..m * k) + .map(|index| Complex32::new(index as f32 * 0.25 - 0.5, index as f32 * -0.125 + 0.25)) + .collect(); + let b32: Vec = (0..k * n) + .map(|index| Complex32::new(index as f32 * -0.2 + 0.75, index as f32 * 0.15 - 0.3)) + .collect(); + let actual32 = cpu.gemm_internal::>(&a32, m, k, &b32, n); + let expected32 = generic_gemm::>(&a32, m, k, &b32, n); + for (actual, expected) in actual32.iter().zip(&expected32) { + assert!((*actual - *expected).norm() <= 1e-5); + } + + let a64: Vec = a32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let b64: Vec = b32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let actual64 = cpu.gemm_internal::>(&a64, m, k, &b64, n); + let expected64 = generic_gemm::>(&a64, m, k, &b64, n); + for (actual, expected) in actual64.iter().zip(&expected64) { + assert!((*actual - *expected).norm() <= 1e-12); + } + } + + #[test] + fn test_cpu_gemm_standard_complex_degenerate_dimensions() { + let cpu = Cpu; + let empty32 = cpu.gemm_internal::>(&[], 0, 3, &[], 0); + assert!(empty32.is_empty()); + + let zeros64 = cpu.gemm_internal::>(&[], 2, 0, &[], 3); + assert_eq!(zeros64, vec![Complex64::new(0.0, 0.0); 6]); + } + + #[test] + fn test_complex_layout_gemm_accepts_contiguous_and_transposed_views() { + let cpu = Cpu; + let a32 = vec![ + Complex32::new(1.0, 1.0), + Complex32::new(3.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(4.0, -1.0), + ]; + let b32 = vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, -1.0), + Complex32::new(0.0, 1.0), + Complex32::new(-1.0, 0.0), + ]; + let actual32 = cpu + .gemm_standard_layout_internal::>( + MatrixLayout::column_major(&a32, 2, 2), + MatrixLayout::column_major(&b32, 2, 2), + ) + .expect("Complex32 layouts should use faer"); + let expected32 = generic_gemm::>(&a32, 2, 2, &b32, 2); + assert_eq!(actual32, expected32); + + let a64: Vec = a32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let b64: Vec = b32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let actual64 = cpu + .gemm_standard_layout_internal::>( + MatrixLayout::column_major(&a64, 2, 2), + MatrixLayout::column_major_transposed(&b64, 2, 2), + ) + .expect("Complex64 transpose layouts should use faer"); + let b64_transposed = vec![b64[0], b64[2], b64[1], b64[3]]; + let expected64 = generic_gemm::>(&a64, 2, 2, &b64_transposed, 2); + assert_eq!(actual64, expected64); + } + + #[test] + fn test_complex_layout_gemm_accepts_nonunit_and_negative_strides() { + let cpu = Cpu; + let a32 = vec![ + Complex32::new(1.0, 1.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, -1.0), + Complex32::new(-1.0, 2.0), + Complex32::new(0.5, 0.0), + Complex32::new(4.0, 1.0), + ]; + let b32 = vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 1.0), + Complex32::new(-1.0, 0.0), + Complex32::new(0.0, 0.5), + Complex32::new(3.0, 0.0), + Complex32::new(-2.0, 1.0), + ]; + let expected32 = generic_gemm::>(&a32, 2, 3, &b32, 2); + + let a_negative = vec![a32[1], a32[0], a32[3], a32[2], a32[5], a32[4]]; + let b_negative = vec![b32[3], b32[4], b32[5], b32[0], b32[1], b32[2]]; + let actual32 = cpu + .gemm_standard_layout_internal::>( + MatrixLayout { + data: &a_negative, + rows: 2, + cols: 3, + row_stride: -1, + col_stride: 2, + }, + MatrixLayout { + data: &b_negative, + rows: 3, + cols: 2, + row_stride: 1, + col_stride: -3, + }, + ) + .expect("negative Complex32 strides should use faer"); + for (actual, expected) in actual32.iter().zip(&expected32) { + assert!((*actual - *expected).norm() <= 1e-5); + } + + let a64: Vec = a32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let b64: Vec = b32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let expected64 = generic_gemm::>(&a64, 2, 3, &b64, 2); + let mut a_positive = vec![Complex64::new(99.0, 99.0); 13]; + for (index, offset) in [0usize, 2, 5, 7, 10, 12].into_iter().enumerate() { + a_positive[offset] = a64[index]; + } + let mut b_positive = vec![Complex64::new(99.0, 99.0); 12]; + for (index, offset) in [0usize, 2, 4, 7, 9, 11].into_iter().enumerate() { + b_positive[offset] = b64[index]; + } + let actual64 = cpu + .gemm_standard_layout_internal::>( + MatrixLayout { + data: &a_positive, + rows: 2, + cols: 3, + row_stride: 2, + col_stride: 5, + }, + MatrixLayout { + data: &b_positive, + rows: 3, + cols: 2, + row_stride: 2, + col_stride: 7, + }, + ) + .expect("non-unit Complex64 strides should use faer"); + for (actual, expected) in actual64.iter().zip(&expected64) { + assert!((*actual - *expected).norm() <= 1e-12); + } + } + + #[test] + fn test_complex_layout_gemm_into_does_not_copy_inputs() { + let a32 = vec![Complex32::new(1.0, 1.0); 4]; + let b32 = vec![Complex32::new(2.0, -1.0); 4]; + let mut c32 = vec![Complex32::new(0.0, 0.0); 4]; + let a64 = vec![Complex64::new(1.0, 1.0); 4]; + let b64 = vec![Complex64::new(2.0, -1.0); 4]; + let mut c64 = vec![Complex64::new(0.0, 0.0); 4]; + + faer_gemm_layout_into( + MatrixLayout::column_major(&a32, 2, 2), + MatrixLayout::column_major(&b32, 2, 2), + &mut c32, + ); + faer_gemm_layout_into( + MatrixLayout::column_major(&a64, 2, 2), + MatrixLayout::column_major(&b64, 2, 2), + &mut c64, + ); + + let ((), allocations) = allocation_counting::with_allocation_counting(|| { + faer_gemm_layout_into( + MatrixLayout::column_major(&a32, 2, 2), + MatrixLayout::column_major(&b32, 2, 2), + &mut c32, + ); + faer_gemm_layout_into( + MatrixLayout::column_major(&a64, 2, 2), + MatrixLayout::column_major(&b64, 2, 2), + &mut c64, + ); + }); + + assert_eq!(allocations, 0, "borrowed complex GEMM must not copy inputs"); + } + #[test] fn test_faer_layout_gemm_accepts_rhs_transpose_view() { let a = vec![1.0f32, 2.0, 3.0, 4.0]; let b = vec![1.0f32, 2.0, 3.0, 4.0]; - let c = faer_gemm_f32_layout( + let c = faer_gemm_layout( MatrixLayout::column_major(&a, 2, 2), MatrixLayout::column_major_transposed(&b, 2, 2), ); - let expected = faer_gemm_f32(&a, 2, 2, &[1.0, 3.0, 2.0, 4.0], 2); + let expected = faer_gemm(&a, 2, 2, &[1.0, 3.0, 2.0, 4.0], 2); assert_eq!(c, expected); } @@ -1194,7 +1451,7 @@ mod tests { let b = vec![1.0f32, 2.0, 3.0, 4.0]; let mut c = vec![0.0f32; 4]; - faer_gemm_f32_layout_into( + faer_gemm_layout_into( MatrixLayout::column_major(&a, 2, 2), MatrixLayout::column_major(&b, 2, 2), &mut c, @@ -1202,7 +1459,7 @@ mod tests { c.fill(-1.0); let ((), allocations) = allocation_counting::with_allocation_counting(|| { - faer_gemm_f32_layout_into( + faer_gemm_layout_into( MatrixLayout::column_major(&a, 2, 2), MatrixLayout::column_major(&b, 2, 2), &mut c, @@ -1222,7 +1479,7 @@ mod tests { let b = vec![1.0f64, 2.0, 3.0, 4.0]; let mut c = vec![0.0f64; 4]; - faer_gemm_f64_layout_into( + faer_gemm_layout_into( MatrixLayout::column_major(&a, 2, 2), MatrixLayout::column_major(&b, 2, 2), &mut c, @@ -1230,7 +1487,7 @@ mod tests { c.fill(-1.0); let ((), allocations) = allocation_counting::with_allocation_counting(|| { - faer_gemm_f64_layout_into( + faer_gemm_layout_into( MatrixLayout::column_major(&a, 2, 2), MatrixLayout::column_major(&b, 2, 2), &mut c, @@ -1244,6 +1501,193 @@ mod tests { ); } + #[test] + fn test_complex_batched_gemm_contiguous_matches_generic() { + let cpu = Cpu; + let (batch_size, m, k, n) = (2usize, 11usize, 10usize, 10usize); + let a32: Vec = (0..batch_size * m * k) + .map(|index| { + Complex32::new( + (index % 17) as f32 * 0.0625 - 0.5, + (index % 11) as f32 * 0.03125 - 0.125, + ) + }) + .collect(); + let b32: Vec = (0..batch_size * k * n) + .map(|index| { + Complex32::new( + (index % 13) as f32 * -0.05 + 0.3, + (index % 7) as f32 * 0.04 - 0.1, + ) + }) + .collect(); + let actual32 = + cpu.gemm_batched_internal::>(&a32, batch_size, m, k, &b32, n); + let expected32 = + generic_batched_gemm_for_test::>(&a32, batch_size, m, k, &b32, n); + for (actual, expected) in actual32.iter().zip(&expected32) { + assert!((*actual - *expected).norm() <= 1e-4); + } + + let a64: Vec = a32[..m * k] + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let b64: Vec = b32[..k * n] + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let actual64 = cpu.gemm_batched_internal::>(&a64, 1, m, k, &b64, n); + let expected64 = + generic_batched_gemm_for_test::>(&a64, 1, m, k, &b64, n); + for (actual, expected) in actual64.iter().zip(&expected64) { + assert!((*actual - *expected).norm() <= 1e-12); + } + + assert!(cpu + .gemm_batched_internal::>(&[], 0, m, k, &[], n) + .is_empty()); + } + + #[test] + fn test_complex_batched_layout_gemm_accepts_interleaved_and_transposed_views() { + let cpu = Cpu; + let (batch_size, m, k, n) = (3usize, 2usize, 3usize, 2usize); + let a32: Vec = (0..batch_size * m * k) + .map(|index| Complex32::new(index as f32 * 0.1 - 0.4, index as f32 * 0.03 - 0.2)) + .collect(); + let b32: Vec = (0..batch_size * k * n) + .map(|index| Complex32::new(index as f32 * -0.07 + 0.5, index as f32 * 0.02)) + .collect(); + let expected32 = + generic_batched_gemm_for_test::>(&a32, batch_size, m, k, &b32, n); + + let mut a_interleaved = vec![Complex32::new(0.0, 0.0); a32.len()]; + let mut b_interleaved = vec![Complex32::new(0.0, 0.0); b32.len()]; + for batch in 0..batch_size { + for col in 0..k { + for row in 0..m { + a_interleaved[batch + row * batch_size + col * batch_size * m] = + a32[batch * m * k + col * m + row]; + } + } + for col in 0..n { + for row in 0..k { + b_interleaved[batch + row * batch_size + col * batch_size * k] = + b32[batch * k * n + col * k + row]; + } + } + } + let actual32 = cpu + .gemm_batched_standard_layout_internal::>( + batch_size, + MatrixLayout { + data: &a_interleaved, + rows: m, + cols: k, + row_stride: batch_size as isize, + col_stride: (batch_size * m) as isize, + }, + MatrixLayout { + data: &b_interleaved, + rows: k, + cols: n, + row_stride: batch_size as isize, + col_stride: (batch_size * k) as isize, + }, + ) + .expect("interleaved Complex32 batches should use faer"); + for (actual, expected) in actual32.iter().zip(&expected32) { + assert!((*actual - *expected).norm() <= 1e-5); + } + + let a64: Vec = a32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let b64: Vec = b32 + .iter() + .map(|value| Complex64::new(value.re as f64, value.im as f64)) + .collect(); + let expected64 = + generic_batched_gemm_for_test::>(&a64, batch_size, m, k, &b64, n); + let mut a_transposed = vec![Complex64::new(0.0, 0.0); a64.len()]; + let mut b_transposed = vec![Complex64::new(0.0, 0.0); b64.len()]; + for batch in 0..batch_size { + for col in 0..k { + for row in 0..m { + a_transposed[batch + row * batch_size * k + col * batch_size] = + a64[batch * m * k + col * m + row]; + } + } + for col in 0..n { + for row in 0..k { + b_transposed[batch + row * batch_size * n + col * batch_size] = + b64[batch * k * n + col * k + row]; + } + } + } + let actual64 = cpu + .gemm_batched_standard_layout_internal::>( + batch_size, + MatrixLayout { + data: &a_transposed, + rows: m, + cols: k, + row_stride: (batch_size * k) as isize, + col_stride: batch_size as isize, + }, + MatrixLayout { + data: &b_transposed, + rows: k, + cols: n, + row_stride: (batch_size * n) as isize, + col_stride: batch_size as isize, + }, + ) + .expect("transposed Complex64 batches should use faer"); + for (actual, expected) in actual64.iter().zip(&expected64) { + assert!((*actual - *expected).norm() <= 1e-12); + } + } + + #[test] + fn test_complex_batched_layout_gemm_allocates_only_output() { + let cpu = Cpu; + let batch_size = 3usize; + let a = vec![Complex32::new(1.0, 0.5); batch_size * 4]; + let b = vec![Complex32::new(0.5, -1.0); batch_size * 4]; + let a_layout = MatrixLayout { + data: &a, + rows: 2, + cols: 2, + row_stride: batch_size as isize, + col_stride: (batch_size * 2) as isize, + }; + let b_layout = MatrixLayout { + data: &b, + rows: 2, + cols: 2, + row_stride: batch_size as isize, + col_stride: (batch_size * 2) as isize, + }; + + drop( + cpu.gemm_batched_standard_layout_internal::>( + batch_size, a_layout, b_layout, + ), + ); + let (result, allocations) = allocation_counting::with_allocation_counting(|| { + cpu.gemm_batched_standard_layout_internal::>( + batch_size, a_layout, b_layout, + ) + .expect("Complex32 batches should use faer") + }); + + assert_eq!(result.len(), batch_size * 4); + assert_eq!(allocations, 1, "batched GEMM should allocate only output"); + } + #[test] fn test_gemm_batched_standard_layout_internal_accepts_batch_major_views() { let cpu = Cpu; diff --git a/tests/suites/binary_rules.rs b/tests/suites/binary_rules.rs index f0a3727..943d660 100644 --- a/tests/suites/binary_rules.rs +++ b/tests/suites/binary_rules.rs @@ -6,12 +6,29 @@ use std::collections::HashMap; use omeinsum::backend::Cpu; use omeinsum::einsum::Einsum; -use omeinsum::{einsum, Standard, Tensor}; +use omeinsum::{einsum, Complex32, Standard, Tensor}; // ============================================================================ // Basic Binary Contractions // ============================================================================ +#[test] +fn test_binary_batched_complex_zero_dimensions() { + let a = Tensor::::from_data(&[], &[2, 0, 3]); + let b = Tensor::::from_data(&[Complex32::new(1.0, -0.5); 24], &[2, 3, 4]); + let zero_free = + einsum::, _, _>(&[&a, &b], &[&[0, 1, 2], &[0, 2, 3]], &[0, 1, 3]); + assert_eq!(zero_free.shape(), &[2, 0, 4]); + assert!(zero_free.to_vec().is_empty()); + + let a = Tensor::::from_data(&[], &[2, 3, 0]); + let b = Tensor::::from_data(&[], &[2, 0, 4]); + let zero_contract = + einsum::, _, _>(&[&a, &b], &[&[0, 1, 2], &[0, 2, 3]], &[0, 1, 3]); + assert_eq!(zero_contract.shape(), &[2, 3, 4]); + assert_eq!(zero_contract.to_vec(), vec![Complex32::new(0.0, 0.0); 24]); +} + #[test] fn test_binary_matmul_ij_jk_ik() { // Standard matrix multiplication: ij,jk->ik From c2866c9caa4b918367002c0478d47a41050e72e4 Mon Sep 17 00:00:00 2001 From: exaclior Date: Wed, 15 Jul 2026 13:43:53 +0800 Subject: [PATCH 4/5] Strengthen complex GEMM validation --- benchmarks/complex_tdvp.md | 38 ++++++++++++++++++++------------------ src/backend/cpu/mod.rs | 38 +++++++++++++++++++++++++------------- 2 files changed, 45 insertions(+), 31 deletions(-) diff --git a/benchmarks/complex_tdvp.md b/benchmarks/complex_tdvp.md index 0851c79..d9314ae 100644 --- a/benchmarks/complex_tdvp.md +++ b/benchmarks/complex_tdvp.md @@ -169,9 +169,11 @@ For every case and `chi`: ## Validation record (2026-07-15) Measurements ran on remote host `6xa800` (2-socket Intel Xeon Platinum 8378A, -128 logical CPUs) with Rust 1.88.0. Baseline and candidate used the same release -profile, benchmark inputs, and Criterion settings. Criterion artifacts are under -`~/projects/omeinsum-rs/target/criterion/tdvp-complex-binary` on that host. +128 logical CPUs) with Rust 1.88.0. The scalar baseline was commit `eaf29fe`; the +faer implementation was commit `6c273ae`. Both used the same release profile, +benchmark inputs, and Criterion settings. Criterion artifacts are under +`~/projects/omeinsum-rs-complex-median/target/criterion/tdvp-complex-binary` on +that host. The repository-mandated `runscribe` executable was unavailable locally and on the remote host; the similarly named PyPI package is an unrelated terminal @@ -179,24 +181,24 @@ recorder. With the owner-requested remote workflow, runs were submitted through `easy-ssh submit` and their full job output was captured instead. There is therefore no runscribe run directory for these measurements. -Times below are Criterion point estimates in milliseconds; parenthesized ranges -are the default 95% confidence intervals. Speedup is baseline divided by -candidate. +Times below are Criterion median estimates in milliseconds; parenthesized ranges +are the median's bootstrap 95% confidence intervals from `estimates.json`. +Speedup is baseline median divided by candidate median. | Case | chi | Scalar baseline | faer candidate | Speedup | Candidate throughput | |---|---:|---:|---:|---:|---:| -| `h1-left-environment` | 32 | 1.0503 (1.0419-1.0651) | 0.084151 (0.084054-0.084254) | 12.48x | 7.0091 Gelem/s | -| `h1-right-environment` | 32 | 1.1197 (1.1158-1.1286) | 0.13741 (0.13709-0.13759) | 8.15x | 4.2924 Gelem/s | -| `h2-left-environment` | 32 | 1.8304 (1.8032-1.9023) | 0.15972 (0.15964-0.15988) | 11.46x | 7.3856 Gelem/s | -| `h2-right-environment` | 32 | 2.4375 (2.3791-2.5324) | 0.30246 (0.30238-0.30255) | 8.06x | 3.9001 Gelem/s | -| `h1-left-environment` | 64 | 8.0044 (7.9638-8.1163) | 0.61142 (0.61091-0.61199) | 13.09x | 7.7175 Gelem/s | -| `h1-right-environment` | 64 | 12.281 (11.996-12.757) | 0.96001 (0.95767-0.96541) | 12.79x | 4.9151 Gelem/s | -| `h2-left-environment` | 64 | 16.475 (15.980-17.505) | 1.1750 (1.1745-1.1756) | 14.02x | 8.0319 Gelem/s | -| `h2-right-environment` | 64 | 33.899 (32.550-35.217) | 1.9024 (1.8967-1.9183) | 17.82x | 4.9606 Gelem/s | -| `h1-left-environment` | 128 | 92.408 (88.993-95.673) | 4.5194 (4.4687-4.6194) | 20.45x | 8.3525 Gelem/s | -| `h1-right-environment` | 128 | 130.87 (129.61-132.46) | 5.9194 (5.9064-5.9326) | 22.11x | 6.3772 Gelem/s | -| `h2-left-environment` | 128 | 164.62 (163.87-166.52) | 8.6727 (8.6685-8.6766) | 18.98x | 8.7051 Gelem/s | -| `h2-right-environment` | 128 | 248.54 (247.66-249.35) | 11.647 (11.623-11.670) | 21.34x | 6.4822 Gelem/s | +| `h1-left-environment` | 32 | 1.02905 (1.02593-1.10051) | 0.0850231 (0.0846379-0.0851403) | 12.10x | 6.9372 Gelem/s | +| `h1-right-environment` | 32 | 1.11454 (1.10367-1.19000) | 0.137350 (0.137313-0.137598) | 8.11x | 4.2943 Gelem/s | +| `h2-left-environment` | 32 | 1.92220 (1.82214-2.11122) | 0.166644 (0.162343-0.178066) | 11.53x | 7.0788 Gelem/s | +| `h2-right-environment` | 32 | 2.40304 (2.37614-2.73987) | 0.309394 (0.308372-0.309727) | 7.77x | 3.8128 Gelem/s | +| `h1-left-environment` | 64 | 7.94239 (7.88650-9.13513) | 0.613996 (0.613023-0.615307) | 12.94x | 7.6851 Gelem/s | +| `h1-right-environment` | 64 | 12.0928 (11.9707-12.9674) | 0.944706 (0.944430-0.945360) | 12.80x | 4.9948 Gelem/s | +| `h2-left-environment` | 64 | 17.5628 (15.8957-18.3190) | 1.18860 (1.18437-1.19116) | 14.78x | 7.9397 Gelem/s | +| `h2-right-environment` | 64 | 31.9260 (29.1754-32.3224) | 1.88922 (1.88402-1.90069) | 16.90x | 4.9953 Gelem/s | +| `h1-left-environment` | 128 | 81.5555 (81.3140-92.3979) | 4.47290 (4.46816-4.48079) | 18.23x | 8.4394 Gelem/s | +| `h1-right-environment` | 128 | 134.061 (128.173-137.447) | 5.78016 (5.77582-5.78916) | 23.19x | 6.5307 Gelem/s | +| `h2-left-environment` | 128 | 164.353 (163.578-182.110) | 8.77822 (8.75957-8.87026) | 18.72x | 8.6005 Gelem/s | +| `h2-right-environment` | 128 | 252.952 (250.701-257.196) | 11.7994 (11.7127-11.8462) | 21.44x | 6.3984 Gelem/s | All four chi=64 cases exceed the required 5x threshold. The existing real f32 `binary` benchmark suite was also compared against a clean `eaf29fe` worktree. diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index 63886e8..01cd904 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -1652,40 +1652,52 @@ mod tests { } #[test] - fn test_complex_batched_layout_gemm_allocates_only_output() { + fn test_complex_batched_layout_gemm_does_not_allocate_temporary_outputs() { let cpu = Cpu; - let batch_size = 3usize; - let a = vec![Complex32::new(1.0, 0.5); batch_size * 4]; - let b = vec![Complex32::new(0.5, -1.0); batch_size * 4]; + let (batch_size, dim) = (3usize, 64usize); + let matrix_len = dim * dim; + let a = vec![Complex32::new(1.0, 0.5); batch_size * matrix_len]; + let b = vec![Complex32::new(0.5, -1.0); batch_size * matrix_len]; let a_layout = MatrixLayout { data: &a, - rows: 2, - cols: 2, + rows: dim, + cols: dim, row_stride: batch_size as isize, - col_stride: (batch_size * 2) as isize, + col_stride: (batch_size * dim) as isize, }; let b_layout = MatrixLayout { data: &b, - rows: 2, - cols: 2, + rows: dim, + cols: dim, row_stride: batch_size as isize, - col_stride: (batch_size * 2) as isize, + col_stride: (batch_size * dim) as isize, }; + let mut one_output = vec![Complex32::new(0.0, 0.0); matrix_len]; + // Warm faer's dispatch before measuring its per-call workspace behavior. + faer_gemm_layout_into(a_layout, b_layout, &mut one_output); drop( cpu.gemm_batched_standard_layout_internal::>( batch_size, a_layout, b_layout, ), ); - let (result, allocations) = allocation_counting::with_allocation_counting(|| { + + let ((), workspace_allocations) = allocation_counting::with_allocation_counting(|| { + faer_gemm_layout_into(a_layout, b_layout, &mut one_output); + }); + let (result, batched_allocations) = allocation_counting::with_allocation_counting(|| { cpu.gemm_batched_standard_layout_internal::>( batch_size, a_layout, b_layout, ) .expect("Complex32 batches should use faer") }); - assert_eq!(result.len(), batch_size * 4); - assert_eq!(allocations, 1, "batched GEMM should allocate only output"); + assert_eq!(result.len(), batch_size * matrix_len); + assert_eq!( + batched_allocations, + 1 + batch_size * workspace_allocations, + "batched GEMM should allocate one result plus only faer's per-batch workspace" + ); } #[test] From ef80db8d44db492854ba52122d4980d80334f50d Mon Sep 17 00:00:00 2001 From: exaclior Date: Wed, 15 Jul 2026 16:36:49 +0800 Subject: [PATCH 5/5] Fix allocation counter race in tests --- src/backend/cpu/mod.rs | 71 +++++++++++++++++++++++++++++++----------- 1 file changed, 52 insertions(+), 19 deletions(-) diff --git a/src/backend/cpu/mod.rs b/src/backend/cpu/mod.rs index 01cd904..c9f8196 100644 --- a/src/backend/cpu/mod.rs +++ b/src/backend/cpu/mod.rs @@ -116,31 +116,30 @@ fn should_use_standard_batched_gemm(batch_size: usize, m: usize, k: usize, n: us mod allocation_counting { use std::alloc::{GlobalAlloc, Layout, System}; use std::cell::Cell; - use std::sync::atomic::{AtomicUsize, Ordering}; pub(crate) struct CountingAllocator; - static ALLOCATION_COUNT: AtomicUsize = AtomicUsize::new(0); thread_local! { + static ALLOCATION_COUNT: Cell = const { Cell::new(0) }; static COUNT_ALLOCATIONS: Cell = const { Cell::new(false) }; } + fn record_allocation() { + COUNT_ALLOCATIONS.with(|active| { + if active.get() { + ALLOCATION_COUNT.with(|count| count.set(count.get() + 1)); + } + }); + } + unsafe impl GlobalAlloc for CountingAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - COUNT_ALLOCATIONS.with(|active| { - if active.get() { - ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed); - } - }); + record_allocation(); unsafe { System.alloc(layout) } } unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - COUNT_ALLOCATIONS.with(|active| { - if active.get() { - ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed); - } - }); + record_allocation(); unsafe { System.alloc_zeroed(layout) } } @@ -149,21 +148,55 @@ mod allocation_counting { } unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - COUNT_ALLOCATIONS.with(|active| { - if active.get() { - ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed); - } - }); + record_allocation(); unsafe { System.realloc(ptr, layout, new_size) } } } pub(crate) fn with_allocation_counting(f: impl FnOnce() -> T) -> (T, usize) { - ALLOCATION_COUNT.store(0, Ordering::Relaxed); + ALLOCATION_COUNT.with(|count| count.set(0)); COUNT_ALLOCATIONS.with(|active| active.set(true)); let result = f(); COUNT_ALLOCATIONS.with(|active| active.set(false)); - (result, ALLOCATION_COUNT.load(Ordering::Relaxed)) + let allocations = ALLOCATION_COUNT.with(Cell::get); + (result, allocations) + } + + #[cfg(test)] + mod tests { + use super::with_allocation_counting; + use std::hint::spin_loop; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::thread; + + #[test] + fn allocation_counts_are_isolated_between_threads() { + let phase = Arc::new(AtomicUsize::new(0)); + let worker_phase = Arc::clone(&phase); + let worker = thread::spawn(move || { + let (buffer, allocations) = with_allocation_counting(|| { + let buffer = Vec::::with_capacity(64); + worker_phase.store(1, Ordering::Release); + while worker_phase.load(Ordering::Acquire) != 2 { + spin_loop(); + } + buffer + }); + drop(buffer); + allocations + }); + + while phase.load(Ordering::Acquire) != 1 { + spin_loop(); + } + let ((), allocations) = with_allocation_counting(|| { + phase.store(2, Ordering::Release); + }); + + assert_eq!(allocations, 0); + assert_eq!(worker.join().unwrap(), 1); + } } }