gobackend: SIMD FusedSDPA, Where implicit broadcast & AVX, and attention/matmul optimizations

[janpfeifer]

Summary

This PR introduces substantial performance optimizations and feature improvements across the pure Go backend (compute/internal/gobackend), with a major focus on attention workloads (Scaled Dot-Product Attention), operator SIMD vectorization and parallelization, implicit broadcasting for Where, and MatMul assembly enhancements:

1. SIMD Fused Scaled Dot-Product Attention (SDPA)

• Architecture Dispatch Mechanism:
  • Added SDPAFloat32ArchFn and registered architecture-specific dispatchers (SetSDPAArchDispatcher, GetSDPAArchDispatcher) in internal/gobackend/sdpa_arch.go.
  • Updated multi-head and Grouped Query Attention (GQA) generic routines to dispatch directly to architecture-specialized kernels when available with zero intermediate transpositions.
  • Multi-threaded head/batch processing across worker pools (batchSize * numKVHeads) with dedicated per-worker scratchpad buffer reuse.
• AVX2 & AVX-512 Kernels:
  • Implemented high-performance AVX2 and AVX-512 SIMD kernels in internal/gobackend/fusedops/avx2/sdpa.go and internal/gobackend/fusedops/avx512/sdpa.go.
  • Vectorized query-key dot products ($Q \cdot K^T$), scaling, causal masking, additive/boolean masking, and bias addition.
  • Vectorized online/stable softmax calculation using Cephes degree-7 Horner polynomial exponential approximations (exp256 / exp512), maximum reduction, and vector reciprocal normalization.
  • Vectorized attention-value weighted accumulation directly into output buffers.

2. Where Operator: Implicit Broadcasting & SIMD Acceleration

• Implicit Broadcasting Support:
  • Extended Where and shapeinference.Where to support standard implicit broadcasting across all three operands (condition, onTrue, onFalse), supporting scalars as well as size-1 dimension expansions across matching ranks.
  • Extended the ImplicitBroadcastFusion graph pass (internal/gobackend/passes/implicitbroadcast.go) to fuse BroadcastInDim directly into Where.
  • Updated public documentation and signatures in StandardOps.Where.
• Intra-op Parallelization & Chunking:
  • Added multi-threaded worker chunking for Where when tensor size exceeds single-threaded thresholds (totalElements > 16384).
• AVX2 & AVX-512 Assembly Kernels:
  • Added AMD64 assembly kernels (internal/gobackend/ops/where_amd64.s) for contiguous and trailing broadcast evaluation under AVX2 and AVX-512.
  • Added architecture dispatch hooks (where_amd64.go, where_other.go) and comprehensive unit tests (where_test.go, shapeinference_test.go).

3. MatMul Packing & Epilogue Optimizations

• Assembly PackLHS:
  • Implemented vectorized PackLHS assembly routines in internal/gobackend/dot/matmul/avx512/pack_amd64_float32.s and packing.go.
  • Fixed assembly packing to 8 rows for large matrix multiplication kernels.
• Chunked & Parallelized Epilogue:
  • Added intra-op parallelization and chunking for the MatMul epilogue (bias addition, scaling, residual addition, fused activations) across backend workers.
  • Added architecture-specific assembly epilogue kernels (epilogue_amd64.s, epilogue_amd64.go, epilogue_generic.go) and unit tests.

4. Fast-Paths and Intra-Op Parallelization for Transpose, SoftMax & Activations

• 4D & 2D Transposition Fast-Paths:
  • Added specialized, parallelized fast-path for 4D transpositions with permutation (0, 2, 1, 3) ([batch, seq, heads, dim] <-> [batch, heads, seq, dim]) in internal/gobackend/ops/transpose.go.
  • Added cache-tiled parallel 2D transposition and contiguous trailing-axes transposition paths.
• Chunked SoftMax & Activations:
  • Parallelized fusedops/softmax.go across worker pools for multi-head and batch dimensions.
  • Added intra-op chunking and parallelization for elementwise activations in internal/gobackend/activations/activations.go.

5. Runtime Hardware Capabilities

• Converted IsAVX2Allowed and IsAVX512Allowed in internal/gobackend/capabilities.go to mutable boolean variables, allowing runtime toggling and fine-grained SIMD test validation across environments.