Plan 9 SIMD — Forward and Backward

Loom's CPU SIMD path uses Plan 9 assembly (.s files) for hot GEMV-style kernels on amd64 (AVX2) and arm64 (NEON). Forward passes use DotTile; backward passes use SaxpyF32AccF64 (and strided variants). Both are wired through the same UseSimdForward flag and TrainingModeCPUSimd training mode.

Overview

Layers with SIMD forward + backward: Dense, SwiGLU, MHA, CNN1, CNN2, CNN3, RNN, LSTM, Embedding, Residual.

Native-exact SIMD (*_native_simd.go): same numerics as *_native.go, faster kernels — MAC dtypes via materializeF32Weights + DotTile; true integers via DotI8Tile / SaxpyI8*. Enabled when both UseExactDType and UseSimdForward are on (Lucy menus [14] and [15]).

Layers without heavy GEMV SIMD: RMSNorm, Softmax (attention softmax/RoPE in MHA remain scalar).

Enabling SIMD

// Per-network (training helper)
poly.ConfigureNetworkForMode(net, poly.TrainingModeCPUSimd)

// Or manually
net.SetSimdForwardRecursive(true) // enables forward DotTile + backward saxpy on all layers

TrainingModeCPUSimd is MC-tiled CPU training with SIMD: same tiling/parallelism as TrainingModeCPUMC, plus recursive UseSimdForward on every compute layer.

Package layout

poly/simd/
├── dot.go, saxpy.go           # Go wrappers, feature detection
├── avx2_amd64.s               # DotTile (AVX2)
├── dot_neon_arm64.s           # DotTile (NEON)
├── saxpy_avx2_amd64.s         # SaxpyF32AccF64 (AVX2)
├── saxpy_neon_arm64.s         # SaxpyF32AccF64 (NEON)
├── bitnet_ternary_amd64.s     # optional ternary forward
└── stub.go                    # scalar fallback when !amd64 && !arm64

poly/
├── simd_forward.go            # layerUseSimdForward, SetSimdForwardRecursive
├── {dense,swiglu,mha,cnn1,cnn2,cnn3,rnn,lstm}_simd.go          # default forward (GetActive FP32)
├── {dense,swiglu,mha,cnn1,cnn2,cnn3,rnn,lstm}_simd_backward.go # default backward
├── {dense,swiglu,mha,cnn1,cnn2,cnn3,rnn,lstm,embedding,residual}_native_simd.go  # native-exact SIMD
├── embedding_simd.go          # parallel lookup / scatter (MAC native bridge)
└── residual_simd.go           # polymorphic residual add (non-native-exact)

Forward path

When layerUseSimdForward(layer) && simd.SimdEnabled():

• Dense / SwiGLU / RNN / LSTM: simd.DotTile over weight rows (same tile structure as MC tiled path).
• MHA: DotTile on Q/K/V/O projections; attention softmax and RoPE stay scalar.
• CNN1/2/3: DotTile over filter patches where the receptive field is contiguous.

Accumulation uses float64 internally where the tiled path does, for dtype parity.

Backward path

When the same SIMD gate is true, each layer's *BackwardPolymorphic tries try*BackwardSimd first:

CNN backward uses an output-centric saxpy scatter for ∂L/∂X (matches tiled exactly; stride-1 full-kernel fast paths, scalar at edges).

RNN/LSTM: parallel over batch; ∂L/∂W_IH and ∂L/∂X via saxpy; hidden-hidden, bias, and cell state paths match tiled scalar code.

Native-exact SIMD (*_native_simd.go)

When UseExactDType and UseSimdForward are both enabled, native-exact layers try SIMD before scalar native fallback:

LayerForwardNativeExact
  └─ layerUseSimdForward → try*ForwardNativeSimd
       ├─ use*TrueNative  → DotI8Tile / SaxpyI8* (int8 MAC)
       └─ else             → materializeF32Weights → *SimdF32WithWeights
  └─ fallback → scalar *_native.go

Lucy [14] reports SIMD fwd/bwd speedup columns for layers that implement this bridge (Dense through Residual). Default menu [7] SIMD uses GetActive FP32 tiles — same fast kernels, QAT-like numerics. Full [14] pass matrix and native-exact speedup tables: native_layers.md.

Validation

Unit tests (poly/tests/*_backward_test.go): SIMD backward vs tiled forward/backward; SC vs MC where applicable.

Seven-layer suite (Lucy examples/seven_layer/):

• TestSimdParityAllLayers_Float32_1x1 — fwd/bwd SC ↔ MC ↔ SIMD, all seven compute layers
• Per-layer grid tests (e.g. TestCNN3SimdParityAllGrids_Float32, TestLSTMSimdParityAllGrids_Float32)
• Training table includes CPU-SIMD column alongside SC and MC

Run:

cd lucy_bloom_rivers && go test ./examples/seven_layer/ -run Simd -count=1
go test ./poly/tests/ -run Backward -count=1

Seven-layer benchmark results

Logs: seven_layer_amd.txt, seven_layer_arm.txt (Float32, 1×1×1 grid, avg 25 passes per cell). Speedup = SC time ÷ SIMD time (>1 means SIMD is faster).

Forward (SIMD vs SC)

Dense forward SIMD is the largest win on both platforms (wide GEMV). CNN3 forward can be slower than SC on small grids (patch layout / overhead); backward still wins (below).

Backward (SIMD vs SC)

CNN2/CNN3 backward saxpy paths show the strongest gains. SwiGLU backward on AMD can trail SC slightly (strided weights + gate structure); ARM sees a modest win.

Training (Dense, Float32, 50 epochs, 1×1×1)

On AMD, SIMD training is ~47% faster than SC for this Dense micro-benchmark. On ARM, SIMD sits between SC (fastest for this tiny grid) and MC.

Parity

On both platforms, seven-layer Float32 1×1×1 parity tables report 0 fwd/bwd diff across SC, MC, and SIMD for all seven layer types (21 dtype rows in full suite). Larger grids (2³, 3³) are covered by per-layer *SimdParityAllGrids* tests.

BitNet (optional)

Ternary / packed BitNet forward kernels live in the same poly/simd package (bitnet_ternary_*.s, TL1 variants on arm64). These are separate from the main float32 training SIMD path; see poly/simd/doc.go and BitNet tests.

See also

• Training — TrainingModeCPUSimd, UseExactDType, and training paradigms
• Quantization — three training/inference modes
• Testing and validation — seven-layer SC/MC/SIMD suite
• Bedrock validation — menu option [7] and artifact layout
• Cross-path layers — menu [15] QAT-SIMD vs Nat-SIMD duel (grid [5], archived amd64/arm64 logs)