Go Assembly & ARM64 Hardware Optimization

Overview

Go’s compiler generates efficient machine code for most workloads, but certain performance-critical operations—SIMD vector math, cryptographic primitives, hardware feature detection—require dropping to assembly. Go supports inline assembly through a unique mechanism: assembly files (*_arm64.s, *_amd64.s) that use Plan 9 assembly syntax and integrate directly with Go’s calling convention. This is not traditional inline assembly; it is a separate file that the Go toolchain assembles and links alongside Go source code.

ARM64 (AArch64) has become the dominant architecture for server workloads, driven by AWS Graviton, Apple Silicon, and Ampere Altra processors. ARM64 provides rich SIMD capabilities through NEON (128-bit vector operations) and, on newer cores, SVE/SME (Scalable Matrix Extensions) for matrix operations. As an example, an SME acceleration library can specifically target Apple Silicon’s SME instructions for accelerated matrix operations in vector search, using Go assembly to access instructions the compiler cannot generate.

This topic covers Go’s Plan 9 assembly for ARM64: the register model, function declarations with //go:noescape directives, WORD-encoded instructions for ISA extensions the assembler does not yet support, ILP-4 accumulator patterns for maximizing throughput, and runtime hardware feature detection via sysctl on macOS or /proc/cpuinfo on Linux.

Core Concepts

Plan 9 Assembly Syntax

Go uses Plan 9 assembly syntax, which differs from both Intel and AT&T syntax. Registers are named R0–R30 on ARM64 (instead of X0–X30). The syntax is INSTRUCTION src, dst (source first, like AT&T). Go-specific pseudo-registers include FP (frame pointer—references function arguments), SP (stack pointer—local variables), SB (static base—global symbols), and PC (program counter).

The //go:noescape Directive

Assembly functions in Go are declared in .go files as function signatures without bodies, paired with the implementation in .s files. The //go:noescape directive tells the compiler that pointer arguments do not escape to the heap, preventing unnecessary heap allocations when passing slices to assembly functions. This is critical for performance: without it, the compiler conservatively allocates slice backing arrays on the heap.

WORD-Encoded Instructions

Go’s assembler does not support every ARM64 instruction, particularly newer extensions like SME (Scalable Matrix Extension). When the assembler lacks support for an instruction, you encode the raw instruction bytes using WORD $0xNNNNNNNN. This embeds the literal 32-bit instruction into the output, bypassing the assembler’s instruction table. You derive the encoding from the ARM Architecture Reference Manual.

Key Concept: ILP-4 Accumulator Pattern

Instruction-Level Parallelism (ILP) maximizes CPU pipeline utilization by issuing independent instructions simultaneously. The ILP-4 pattern uses 4 independent accumulator registers (e.g., V0–V3) in a loop. Each iteration feeds results to a different accumulator, breaking data dependencies so the CPU can execute 4 multiply-accumulate operations per cycle instead of 1. The accumulators are summed after the loop.

Key Concept: FMOPA — Outer Product Accumulate

ARM’s SME extension provides FMOPA (floating-point outer product and accumulate), which computes the outer product of two vectors and adds it to an accumulator matrix tile in a single instruction. For vector search, this accelerates batched distance computations where a query vector is compared against multiple candidate vectors simultaneously, turning N dot products into a single matrix operation.

Key Concept: Hardware Feature Detection

Not all ARM64 CPUs support the same extensions (NEON is universal; SVE, SVE2, SME are not). At runtime, query the hardware: on macOS use sysctlbyname("hw.optional.arm.FEAT_SME"); on Linux read /proc/cpuinfo or use getauxval(AT_HWCAP2). Dispatch to the optimal implementation based on detected features, with a NEON fallback for maximum compatibility.

How It Works

Go Function Declaration with Assembly

The Go side declares the function signature. The assembly side provides the implementation. These must be in the same package:

// dot_arm64.go
package simd

//go:noescape
func dotProductNEON(a, b *float32, n int) float32

//go:noescape
func dotProductSME(a, b *float32, n int) float32

// DotProduct dispatches to the best available implementation
func DotProduct(a, b []float32) float32 {
    if len(a) != len(b) || len(a) == 0 {
        return 0
    }
    if hasSME {
        return dotProductSME(&a[0], &b[0], len(a))
    }
    return dotProductNEON(&a[0], &b[0], len(a))
}

ARM64 NEON Assembly with ILP-4

The NEON implementation processes 16 floats per iteration (4 NEON registers × 4 floats each) with 4 independent accumulators:

// dot_arm64.s
#include "textflag.h"

// func dotProductNEON(a, b *float32, n int) float32
TEXT ·dotProductNEON(SB), NOSPLIT, $0-28
    MOVD a+0(FP), R0       // R0 = pointer to a
    MOVD b+8(FP), R1       // R1 = pointer to b
    MOVD n+16(FP), R2      // R2 = length

    // Zero 4 accumulator registers (ILP-4 pattern)
    VEOR V0.B16, V0.B16, V0.B16
    VEOR V1.B16, V1.B16, V1.B16
    VEOR V2.B16, V2.B16, V2.B16
    VEOR V3.B16, V3.B16, V3.B16

    // Process 16 floats per iteration
    LSR  R3, R2, #4        // R3 = n / 16
    CBZ  R3, tail

loop:
    VLD1.P 64(R0), [V16.S4, V17.S4, V18.S4, V19.S4]
    VLD1.P 64(R1), [V20.S4, V21.S4, V22.S4, V23.S4]

    // 4 independent multiply-accumulates (ILP-4)
    VFMLA V16.S4, V20.S4, V0.S4
    VFMLA V17.S4, V21.S4, V1.S4
    VFMLA V18.S4, V22.S4, V2.S4
    VFMLA V19.S4, V23.S4, V3.S4

    SUB  R3, R3, #1
    CBNZ R3, loop

tail:
    // Combine 4 accumulators into V0
    VADD V0.S4, V1.S4, V0.S4
    VADD V2.S4, V3.S4, V2.S4
    VADD V0.S4, V2.S4, V0.S4

    // Horizontal sum of V0 lanes
    FADDP V0.S4, V0.S4, V0.S4
    FADDP V0.S2, V0.S2, V0.S2

    // Handle remaining elements (scalar)
    AND  R3, R2, #15
    CBZ  R3, done

scalar:
    FMOVS F1, (R0)
    FMOVS F2, (R1)
    FMADD F0, F1, F2, F0
    ADD  R0, R0, #4
    ADD  R1, R1, #4
    SUB  R3, R3, #1
    CBNZ R3, scalar

done:
    FMOVS ret+24(FP), F0
    RET

WORD-Encoded SME Instructions

For SME instructions not yet in Go’s assembler, encode them as raw WORD values. The FMOPA instruction computes an outer product into a ZA tile:

// sme_arm64.s — SME outer product for batched dot products
#include "textflag.h"

// SMSTART — enable streaming SVE mode + ZA storage
#define SMSTART WORD $0xd503437f

// SMSTOP — disable streaming mode
#define SMSTOP  WORD $0xd503427f

// FMOPA ZA0.S, P0/M, P0/M, Z0.S, Z1.S
// Outer product: ZA0 += Z0 * Z1^T
#define FMOPA_ZA0_Z0_Z1 WORD $0x80800000

// ZERO {ZA} — zero all ZA tiles
#define ZERO_ZA WORD $0xc00800ff

TEXT ·dotProductSME(SB), NOSPLIT, $0-28
    MOVD a+0(FP), R0
    MOVD b+8(FP), R1
    MOVD n+16(FP), R2

    SMSTART                 // Enter streaming SVE mode
    ZERO_ZA                 // Clear accumulator tiles

    // ... load vectors into Z0, Z1 and FMOPA ...
    // (simplified — full implementation handles tiling)

    // Extract result from ZA tile
    SMSTOP                  // Exit streaming mode

    FMOVS ret+24(FP), F0
    RET

Key Concept: NOSPLIT Stack Frames

The NOSPLIT flag in Go assembly indicates the function does not need stack split checking (it won’t grow the stack). This is required for leaf assembly functions that do not call other Go functions. The $0-28 declares a 0-byte local frame with 28 bytes of arguments (three 8-byte args + one 4-byte return). Getting frame sizes wrong causes subtle stack corruption.

Practical Applications

Vector Distance with SME Acceleration

An SME acceleration library implements SIMD-accelerated vector distance functions for the Go-based components of a vector search engine. On Apple Silicon (M1/M2/M3), it detects SME support and uses FMOPA-based matrix operations for batched distance computation. On Graviton instances without SME, it falls back to NEON. The dispatch happens once at initialization and is stored in a package-level function pointer.

Cryptographic Hashing in Distributed Databases

A distributed database may use ARM64 hardware crypto extensions (AES, SHA-256) for data integrity and consistent hashing. Go’s standard library already uses assembly for crypto/sha256 on ARM64, but custom hash functions for the hash ring use similar assembly patterns to access SHA256H and SHA256H2 hardware instructions directly.

Bitwise Operations for Bitmap Indices

Bitmap indices in the search engine use ARM64 NEON to perform bulk bitwise AND, OR, and ANDNOT operations on posting list bitmaps. Processing 128 bits per NEON instruction (vs 64 bits for scalar operations) doubles throughput for bitmap intersection queries. The VCNT instruction (population count per byte) accelerates cardinality estimation.

Runtime Feature Detection

Feature detection on macOS uses sysctlbyname through cgo or Go’s internal cpu package:

package simd

import (
    "os"
    "runtime"
    "strings"
)

var hasSME bool
var hasNEON bool = true // NEON is mandatory on ARM64

func init() {
    if runtime.GOARCH != "arm64" {
        return
    }
    switch runtime.GOOS {
    case "darwin":
        hasSME = checkDarwinFeature("hw.optional.arm.FEAT_SME")
    case "linux":
        hasSME = checkLinuxCPUInfo("sme")
    }
}

func checkLinuxCPUInfo(feature string) bool {
    data, err := os.ReadFile("/proc/cpuinfo")
    if err != nil {
        return false
    }
    for _, line := range strings.Split(string(data), "\n") {
        if strings.HasPrefix(line, "Features") {
            return strings.Contains(line, feature)
        }
    }
    return false
}

Implementation Patterns

The Dispatch Pattern

Store a function pointer selected at init time based on CPU features. All call sites use the function pointer, avoiding per-call feature checks:

type distanceFunc func(a, b *float32, n int) float32

var computeDistance distanceFunc

func init() {
    switch {
    case hasSME:
        computeDistance = dotProductSME
    default:
        computeDistance = dotProductNEON
    }
}

// Every call site uses the dispatched function
func (idx *Index) Score(query, candidate []float32) float32 {
    return computeDistance(&query[0], &candidate[0], len(query))
}

Benchmark-Driven Development

Assembly optimizations must be validated with benchmarks. Go’s testing framework makes this straightforward:

func BenchmarkDotProductNEON(b *testing.B) {
    a := make([]float32, 768)
    c := make([]float32, 768)
    for i := range a {
        a[i] = float32(i) * 0.001
        c[i] = float32(i) * 0.002
    }

    b.ResetTimer()
    b.ReportAllocs()
    for i := 0; i < b.N; i++ {
        _ = dotProductNEON(&a[0], &c[0], len(a))
    }
}

// Run: go test -bench=BenchmarkDotProduct -count=5 -benchmem

Testing Assembly Functions

Always test assembly implementations against a pure Go reference. Use property-based or golden-value tests to catch rounding differences and off-by-one bugs:

func TestDotProductNEON(t *testing.T) {
    a := []float32{1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}
    b := []float32{8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0}

    got := dotProductNEON(&a[0], &b[0], len(a))
    want := dotProductGo(a, b) // pure Go reference implementation

    if diff := math.Abs(float64(got - want)); diff > 1e-6 {
        t.Errorf("dotProductNEON = %f, want %f (diff %e)", got, want, diff)
    }
}

Key Concept: Assembly File Naming Convention

Go assembly files use build-tag-style naming: func_arm64.s is only assembled for ARM64, func_amd64.s for x86-64. The corresponding Go declaration file should be func_arm64.go with a matching build constraint. This ensures the pure-Go fallback is used on unsupported architectures without any build configuration.

Performance Considerations

ILP-4 vs ILP-2: Modern ARM64 cores (Firestorm, Neoverse V2) can sustain 4 floating-point multiply-accumulate operations per cycle. Using fewer than 4 accumulators leaves execution units idle. Using more than 4 provides no additional benefit and increases register pressure.
Memory alignment for NEON loads: ARM64 NEON loads from aligned addresses (VLD1) do not require explicit alignment in ARMv8, but aligned accesses still benefit from cache line utilization. Ensure vector data starts on 16-byte boundaries.
Branch prediction on ARM64: ARM64 cores have excellent branch predictors, but tight inner loops should still minimize branches. Use conditional instructions (CSEL, FCSEL) instead of branch-over patterns for simple conditionals inside loops.
SME streaming mode overhead: Entering and exiting SME streaming mode (SMSTART/SMSTOP) has significant overhead (hundreds of cycles). Batch all SME work between a single SMSTART/SMSTOP pair rather than toggling per operation.

Common Pitfalls

Wrong frame size in TEXT directive: The $N-M in TEXT ·func(SB), NOSPLIT, $N-M declares N bytes of local stack and M bytes of arguments+return. Getting M wrong causes the garbage collector to misidentify pointers on the stack, leading to memory corruption that manifests as random crashes far from the assembly code.
Missing //go:noescape: Without this directive, the compiler assumes pointer arguments escape to the heap. Every call to your assembly function triggers a heap allocation for the slice backing array, destroying the performance benefit of assembly in the first place.
Clobbering callee-saved registers: ARM64 requires preserving R19–R28 and V8–V15 across function calls. Go assembly that modifies these without saving/restoring them corrupts the calling Go code’s state. NEON temporary registers V16–V31 are caller-saved and free to clobber.
WORD encoding errors: A single wrong bit in a WORD-encoded instruction produces a different instruction that may execute without faulting but compute wrong results. Always verify WORD encodings against the ARM ARM (Architecture Reference Manual) and test with known-answer vectors.
Testing only on one architecture: Assembly functions are architecture-specific. The Go file must provide a pure-Go fallback for other architectures, and tests must run on both the assembly and fallback paths. CI pipelines should include both ARM64 and AMD64 runners.