Benchmarks: Flex vs NdArray

All benchmarks run on Apple M3 Max, comparing burn-flex against burn-ndarray. Default features enabled (std, simd, rayon); gemm is a required dependency.

Date: 2026-04-06

How to Read

• Median time reported (lower is better)
• Speedup = NdArray median / Flex median
• Mem = peak allocation (max alloc from divan)
• Bold speedup means Flex wins; plain means tie or NdArray wins

Binary Operations (f32)

Transposed

Binary Operations (i64)

Int Power

Transposed (i64)

Int Cast

Int Random

Matrix Multiplication

Square (f32)

Rectangular (f32)

Transposed (256x256)

Batched (f32)

Broadcast (f32)

Integer (i32)

Slice Operations

Basic Slicing

Narrow

Slice Assignment

Transposed Slicing

Slice with Step

Concatenation

Cat (dim 0, contiguous memcpy fast path)

Cat (dim 1, general path)

Dim-1 cat is much faster because NdArray's default uses N slice_assign calls while Flex copies contiguous chunks directly.

Reduce Operations

Full Tensor Sum

Full Tensor Max

Full Tensor Min

Int Max

Sum Along Dimension

3D Sum (Batched)

Sum Transposed

Sum Dim on Transposed

Mean Along Dimension

Argmax

Cumulative Operations

Cumsum

Cumprod

Cummin

Cummax

3D Cumsum (Batched)

Gather/Scatter Operations

Gather

Scatter Add

Select

Bool Select

Select Add

Unary Operations

Basic Math

Trigonometric

Transposed (Non-contiguous)

Comparison & Boolean Operations

Tensor-Tensor Comparisons

Scalar Comparisons

Transposed Comparisons

Broadcast Comparisons

Expand (Broadcasting)

Boolean Operations

Convolutions

Kernel Size Comparison (4x64x56x56, 64 to 128 channels)

ResNet Layers (batch=1, 3x3)

Small (batch=1, 3x3)

Large Batched (batch=16, 3x3)

Medium Batched (batch=8, 3x3)

Conv1d

Pooling

Max Pool 2D

Max Pool 2D (ResNet)

Avg Pool 2D

Adaptive Avg Pool 2D

Max Pool 1D

Kernel Size Comparison (4x64x56x56)

Transposed Convolutions

Conv Transpose 2D

DCGAN Generator

Conv Transpose 1D

Conv Transpose 3D

Interpolation

Nearest

Bilinear

Bicubic

Grid Sample 2D

Cross Product & Unfold

Cross Product

Unfold (1D)

Unfold (2D/3D)

Deformable Convolutions

Small/Tiny Inputs

Medium Inputs

Attention (Scaled Dot-Product)

Flex auto-selects between two gemm-backed strategies:

• Naive (score matrix <= 256K elements): Materializes full [seq_q, seq_kv] score matrix. Two large gemm calls per (batch, head) amortize dispatch overhead better than many small tiled calls.
• Flash (score matrix > 256K elements): Tiles over KV dimension with online softmax. O(seq_q * TILE_KV) memory per head instead of O(seq_q * seq_kv).

Both fuse scale + softcap + masking + bias + softmax into a single pass, reducing intermediate allocations from ~12 (NdArray fallback) to 3.

Self-Attention

Causal Attention

With Additive Bias (ALiBi-style)

Cross-Attention (seq_q != seq_k)

Quantized Tensor Operations

All quantized ops (except layout ops) go through a dequantize-op-quantize cycle. Flex stores scales separately and applies scale * x_q directly; NdArray reparses QuantizedBytes on every dequantize call, which dominates the cost.

Quantize (float to i8)

Dequantize (i8 to float)

q_add (dequant + add + requant)

q_matmul (dequant + matmul + requant)

q_sum (dequant + sum)

q_permute (zero-copy layout op)

q_argmax (operates on i8 directly)

q_argmin (operates on i8 directly)

q_gather (operates on i8 directly for tensor-level quant)

Default Ops (sort, repeat, creation, embedding, predicates)

These ops override burn's default trait implementations with direct storage operations.

Sort (f32, 1D)

Sort (f32, 2D along last dim)

Argsort (f32, 1D)

Repeat Dim (f32, 256x256)

Tensor Creation (f32, 1M elements)

Arange (i64)

Embedding (f32)

Predicates (f32, 1M elements)

FFT (Real FFT)

Compared against realfft v3 (backed by rustfft v6), the gold-standard pure-Rust FFT library. NdArray does not implement rfft. realfft requires std; Flex works in no_std.

1D rfft

Batched 2D rfft (along last dim)

1D irfft (inverse)

Flex implementation: Cooley-Tukey with mixed radix-4/radix-2, complex packing (forward) / inverse packing (inverse), compile-time twiddle tables via const fn, SIMD vectorization via macerator, unrolled small kernels (N=2,4,8), and rayon parallelism across fibers. The remaining gap to rustfft is due to their hand-tuned per-arch SIMD rewrites, split-radix algorithms, and strength-reduced modular arithmetic.

Running Benchmarks

cargo bench --bench attention
cargo bench --bench binary_ops
cargo bench --bench matmul
cargo bench --bench int_ops
cargo bench --bench slice_ops
cargo bench --bench reduce_ops
cargo bench --bench cumulative_ops
cargo bench --bench gather_scatter_ops
cargo bench --bench unary_ops
cargo bench --bench comparison_ops
cargo bench --bench conv_ops
cargo bench --bench pool_ops
cargo bench --bench conv_transpose_ops
cargo bench --bench interpolate_ops
cargo bench --bench cross_unfold_ops
cargo bench --bench deform_conv_ops
cargo bench --bench quantization_ops
cargo bench --bench cat_max_min_ops
cargo bench --bench default_ops
cargo bench --bench fft_ops