SIMD Acceleration¶
tensai's fast kernels are AVX2 on amd64 and NEON (ARM's Advanced SIMD) on arm64, written with Go's experimental simd/archsimd package — still pure Go, no cgo, no assembly files.
Requirements: amd64 with Go 1.26 or 1.27 (both simd API generations are supported via build tags), or arm64 with Go 1.27, whose simd/archsimd is the first to carry an arm64 half. Every other build — other architectures, older Go, or GOEXPERIMENT unset — uses the portable fallbacks automatically, with identical results. Platforms has the per-kernel breakdown: the NEON build vectorizes the decode path but not yet the 4-bit and grouped-int8 matvecs, the batched prefill, or the dense float matmul.
What is vectorized¶
Where the AVX2 kernels apply today, and where they still could:
- Matmul (
Dot/DotInto) — used byDense,Conv2D(im2col product),knn.Classifierdistances, and autogradMatMul - ReLU / LeakyReLU forward & backward
- Sigmoid / Tanh forward & backward (vectorized polynomial
exp) - GELU forward & backward (vectorized
erf) - LayerNorm forward & backward (vector row reductions)
- Softmax / SoftmaxCrossEntropy exponentials and scaling
- Adam / AdamW parameter update
- SGD update (momentum form, same fused multiply-add loop as Adam)
- Slice add & scale primitives (bias add,
Embeddinggradient scatter-add) - Transpose-free gradient matmul (
DotTAInto) —Dense/Conv2Dweight gradients no longer materializeinput^T/im2col^T - Remaining transposes (
T/TInto) — cache-blocked 32x32 tiles - Softmax backward row dot products (autograd) — fused AVX2 dot and Jacobian-vector accumulation
- MSE / BinaryCrossEntropy losses (BCE needs a vectorized
log) - Autograd element-wise backward passes (gradients accumulate with
+=, so they need dedicated fused kernels) - BatchNorm statistics (column-strided access needs a restructure)
- MaxPool2D window scan
- im2col / col2im gather-scatter (contiguous runs could use bulk copies)
The unchecked items are ordered roughly by expected impact; none of them show up prominently in training profiles today.
The int8/int4 quantized matmuls have their own AVX2 paths built on the 256-bit u8 x s8 pairwise multiply-add — see Quantization.
A live benchmark¶
_example/plasma animates a demoscene-style plasma in the terminal where the plasma function is a randomly weighted network (a CPPN) evaluated for every pixel of every frame as one batch. The status line shows the per-frame network time: 120x90 pixels runs at ~32 fps on the portable build and ~100 fps with GOEXPERIMENT=simd on the same machine.