zstd: Add arm64 decoder asm (#1160)

* zstd: add arm64 asm implementation of fseDecoder.buildDtable

arm64 previously fell through to the pure-Go fse_decoder_generic.go because
the decoder asm is avo-generated and avo has no arm64 backend. This adds a
hand-written scalar arm64 port of buildDtable, mirroring the avo-generated
fse_decoder_amd64.s and the pure-Go reference. The kernel is serial bitfield
work (no SIMD), so it needs no NEON/SVE and builds on the stock go1.x
assembler.

Correctness is covered by a differential test against an in-tree copy of the
generic algorithm over the three RFC predefined distributions, plus the full
existing decode suite (which runs on macos-latest/arm64 in CI).

darwin/arm64 (M4, go1.26), asm vs -tags noasm; the bench's table reset is
identical on both sides so the delta is the buildDtable work:

  litLength     ~166 vs ~201 ns/op   (+18%)
  offset        ~118 vs ~141 ns/op   (+16%)
  matchLength   ~175 vs ~210 ns/op   (+17%)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* zstd: add avo-generated arm64 asm sequence decoder

Generates the full non-BMI2 zstd sequence decoder for arm64 by lowering the
existing amd64 avo program (_generate/gen.go) through a new arm64 backend in
avo (a post-pass.Compile lowering printer; see mmcloughlin/avo#486). No asm is
hand-written: `go generate` now emits both seqdec_amd64.s and seqdec_arm64.s
from the same IR, so they cannot drift, and CI's generate+git-diff job enforces
it.

- seqdec_arm64.s: generated (6 non-BMI2 functions: decode/_56, executeSimple/
  _safe, decodeSync/_safe).
- seqdec_arm64.go: shim mirroring seqdec_amd64.go (context structs, noescape
  decls, dispatch) without the BMI2 paths.
- seqdec_generic.go: build tag now excludes arm64.

Scalar only — no NEON/SVE (one VLD1/VST1 pair for a 16-byte MOVUPS copy);
builds on go1.26. Full zstd decode suite passes on arm64. On Apple M4,
BenchmarkDecoder_DecodeAll is ~1.5x faster than the pure-Go fallback on
sequence-heavy input (kppkn: 469 vs 304 MB/s).

Regeneration requires the avo arm64 backend (mmcloughlin/avo#486) via a local
replace directive in _generate/go.mod (not committed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* zstd: generate arm64 fse decoder via avo + pin avo fork

Switches fse_decoder_arm64.s from the hand-written bootstrap to the
avo-generated lowering, so the entire arm64 zstd decoder (fse table build
+ sequence decoder) is produced from the existing, unmodified amd64 avo
generators via the arm64 lowering printer (mmcloughlin/avo#486). Supersedes
the hand-written fse PR #1159.

Pins the printer through a replace to the honeycombio/avo fork (which keeps
the github.com/mmcloughlin/avo module path, so the replace resolves) until
the printer merges upstream and is released, at which point the replace is
dropped for a normal require.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* zstd: asmfmt-style formatting for generated arm64 asm

Regenerate the arm64 decoder with the updated avo arm64 printer, which now
aligns operands to a common column and emits a blank line before labels —
matching avo's amd64 (goasm) output. asmfmt reports zero non-comment
differences on the arm64 files (the same standalone-comment indentation
residual goasm produces on amd64). Repins the avo fork to the printer commit.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* zstd: address review — dedup shims, -arch generation, asm-vs-nonasm fuzz

Responds to maintainer review on the arm64 decoder PR:

- Deduplicate the asm Go wrappers: shared seqdec_asm.go holds the context
  structs, error codes and the decode/decodeSync/executeSimple bodies; each
  arch keeps only its //go:noescape decls and a small dispatch helper (amd64
  selects the BMI2 / 56-bit / safe variant, arm64 the 56-bit / safe variant).
  The fse Go wrapper is identical across arches and collapses into one
  fse_decoder_asm.go (per-arch fse_decoder_{amd64,arm64}.go removed).

- Generate both architectures via avo's new -arch flag: one
  `go run gen.go -out ../seqdec.s -arch amd64,arm64` writes seqdec_amd64.s and
  seqdec_arm64.s from a single compile, so neither file's header references the
  other architecture.

- FuzzDecodeAll now cross-checks the assembly decoder against a non-assembly
  path (BMI2 disabled on amd64) for byte-identical output; run with -tags=noasm
  over the shared corpus to cover assembly-vs-generic on arm64.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* ci: run build matrix and zstd fuzzers on linux arm64

Adds ubuntu-24.04-arm to the build matrix (vet/test/noasm/race across Go
versions on Linux arm64) and gives the fuzz-zstd job an OS matrix so the zstd
decode fuzzers also run on arm64 — exercising the new arm64 decoder assembly
directly in CI, as requested in review.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* zstd: defer BMI2 restore in FuzzDecodeAll asm/non-asm comparison

Use the defer cpuinfo.DisableBMI2()() idiom so BMI2 is re-enabled on every exit
path (panic/return), matching FuzzDecAllNoBMI2. Addresses review feedback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Klaus Post <klauspost@gmail.com>

Author: 3 people

.github/workflows/go.yml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       matrix:
         go-version: [1.24.x, 1.25.x, 1.26.x]
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest, ubuntu-24.04-arm, macos-latest, windows-latest]
     env:
       CGO_ENABLED: 0
     runs-on: ${{ matrix.os }}
@@ -146,9 +146,10 @@ jobs:
   fuzz-zstd:
     env:
       CGO_ENABLED: 0
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
     strategy:
       matrix:
+        os: [ ubuntu-latest, ubuntu-24.04-arm ]
         tags: [ 'nounsafe', '"noasm,nounsafe"' ]
     steps:
       - name: Set up Go

zstd/_generate/gen.go

@@ -1,6 +1,6 @@
 package main
 
-//go:generate go run gen.go -out ../seqdec_amd64.s -pkg=zstd
+//go:generate go run gen.go -out ../seqdec.s -arch amd64,arm64 -pkg=zstd
 
 import (
 	"flag"

zstd/_generate/gen_fse.go

@@ -1,6 +1,6 @@
 package main
 
-//go:generate go run gen_fse.go -out ../fse_decoder_amd64.s -pkg=zstd
+//go:generate go run gen_fse.go -out ../fse_decoder.s -arch amd64,arm64 -pkg=zstd
 
 import (
 	"flag"

zstd/_generate/go.mod

@@ -10,9 +10,11 @@ require (
 )
 
 require (
-	golang.org/x/mod v0.21.0 // indirect
-	golang.org/x/sync v0.8.0 // indirect
-	golang.org/x/tools v0.25.1 // indirect
+	golang.org/x/mod v0.27.0 // indirect
+	golang.org/x/sync v0.16.0 // indirect
+	golang.org/x/tools v0.36.0 // indirect
 )
 
 replace github.com/klauspost/compress => ../..
+
+replace github.com/mmcloughlin/avo => github.com/honeycombio/avo v0.6.1-0.20260629192935-ff52dcb5f88d

zstd/_generate/go.sum

@@ -1,8 +1,10 @@
-github.com/mmcloughlin/avo v0.6.0 h1:QH6FU8SKoTLaVs80GA8TJuLNkUYl4VokHKlPhVDg4YY=
-github.com/mmcloughlin/avo v0.6.0/go.mod h1:8CoAGaCSYXtCPR+8y18Y9aB/kxb8JSS6FRI7mSkvD+8=
-golang.org/x/mod v0.21.0 h1:vvrHzRwRfVKSiLrG+d4FMl/Qi4ukBCE6kZlTUkDYRT0=
-golang.org/x/mod v0.21.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
-golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
-golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
-golang.org/x/tools v0.25.1 h1:YeIyhd0M7gStYR9jb2IFXVVT+QJhgXu1ZECOuRwofh4=
-golang.org/x/tools v0.25.1/go.mod h1:/vtpO8WL1N9cQC3FN5zPqb//fRXskFHbLKk4OW1Q7rg=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/honeycombio/avo v0.6.1-0.20260629192935-ff52dcb5f88d h1:946aSI2nbyDYa8iTjFmwoM4S4TKmE8heVLGlchLj7vU=
+github.com/honeycombio/avo v0.6.1-0.20260629192935-ff52dcb5f88d/go.mod h1:+Pk+j3KceMwAIGb32HZSCvkMCzJvxy6xCRk+urz7iPw=
+golang.org/x/mod v0.27.0 h1:kb+q2PyFnEADO2IEF935ehFUXlWiNjJWtRNgBLSfbxQ=
+golang.org/x/mod v0.27.0/go.mod h1:rWI627Fq0DEoudcK+MBkNkCe0EetEaDSwJJkCcjpazc=
+golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=
+golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
+golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg=
+golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s=

zstd/fse_decoder_amd64.s

@@ -1,4 +1,4 @@
-// Code generated by command: go run gen_fse.go -out ../fse_decoder_amd64.s -pkg=zstd. DO NOT EDIT.
+// Code generated by command: go run gen_fse.go -out ../fse_decoder.s -arch amd64,arm64 -pkg=zstd. DO NOT EDIT.
 
 //go:build !appengine && !noasm && gc && !noasm
 

zstd/fse_decoder_arm64.s

@@ -0,0 +1,146 @@
+// Code generated by command: go run gen_fse.go -out ../fse_decoder.s -arch amd64,arm64 -pkg=zstd. DO NOT EDIT.
+// EXPERIMENTAL arm64 output lowered from an amd64 avo program.
+
+//go:build arm64 && !appengine && !noasm && gc && !noasm
+
+// func buildDtable_asm(s *fseDecoder, ctx *buildDtableAsmContext) int
+TEXT ·buildDtable_asm(SB), $0-24
+	MOVD ctx+8(FP), R1
+	MOVD s+0(FP), R6
+
+	// Load values
+	MOVBU 4098(R6), R2
+	MOVD  $0, R0
+	MOVD  $1, R16
+	LSL   R2, R16, R16
+	ORR   R16, R0, R0
+	MOVD  (R1), R3
+	MOVD  16(R1), R5
+	SUB   $1, R0, R7
+	MOVD  8(R1), R1
+	MOVHU 4096(R6), R6
+
+	// End load values
+	// Init, lay down lowprob symbols
+	MOVD $0, R8
+	JMP  init_main_loop_condition
+
+init_main_loop:
+	ADD  R8<<1, R1, R15
+	MOVH (R15), R9
+	CMN  $1, R9
+	BNE  do_not_update_high_threshold
+	ADD  R7<<3, R5, R15
+	MOVB R8, 1(R15)
+	SUB  $1, R7, R7
+	MOVD $0x0000000000000001, R9
+
+do_not_update_high_threshold:
+	ADD  R8<<1, R3, R15
+	MOVH R9, (R15)
+	ADD  $1, R8, R8
+
+init_main_loop_condition:
+	CMP R6, R8
+	BLT init_main_loop
+
+	// Spread symbols
+	// Calculate table step
+	MOVD R0, R8
+	LSR  $0x01, R8, R8
+	MOVD R0, R9
+	LSR  $0x03, R9, R9
+	ADD  R9, R8, R8
+	ADD  $3, R8, R8
+
+	// Fill add bits values
+	SUB  $1, R0, R9
+	MOVD $0, R10
+	MOVD $0, R11
+	JMP  spread_main_loop_condition
+
+spread_main_loop:
+	MOVD $0, R12
+	ADD  R11<<1, R1, R15
+	MOVH (R15), R13
+	JMP  spread_inner_loop_condition
+
+spread_inner_loop:
+	ADD  R10<<3, R5, R15
+	MOVB R11, 1(R15)
+
+adjust_position:
+	ADD R8, R10, R10
+	AND R9, R10, R10
+	CMP R7, R10
+	BGT adjust_position
+	ADD $1, R12, R12
+
+spread_inner_loop_condition:
+	CMP R13, R12
+	BLT spread_inner_loop
+	ADD $1, R11, R11
+
+spread_main_loop_condition:
+	CMP  R6, R11
+	BLT  spread_main_loop
+	TST  R10, R10
+	BEQ  spread_check_ok
+	MOVD ctx+8(FP), R0
+	MOVD R10, 24(R0)
+	MOVD $+1, R16
+	MOVD R16, ret+16(FP)
+	RET
+
+spread_check_ok:
+	// Build Decoding table
+	MOVD $0, R6
+
+build_table_main_table:
+	ADD   R6<<3, R5, R15
+	MOVBU 1(R15), R1
+	ADD   R1<<1, R3, R15
+	MOVHU (R15), R7
+	ADD   $1, R7, R8
+	ADD   R1<<1, R3, R15
+	MOVH  R8, (R15)
+	MOVD  R7, R8
+	CLZ   R8, R16
+	MOVD  $63, R8
+	SUB   R16, R8, R8
+	MOVD  R2, R1
+	SUB   R8, R1, R1
+	LSL   R1, R7, R7
+	SUB   R0, R7, R7
+	ADD   R6<<3, R5, R15
+	MOVB  R1, (R15)
+	ADD   R6<<3, R5, R15
+	MOVH  R7, 2(R15)
+	CMP   R0, R7
+	BLE   build_table_check1_ok
+	MOVD  ctx+8(FP), R1
+	MOVD  R7, 24(R1)
+	MOVD  R0, 32(R1)
+	MOVD  $+2, R16
+	MOVD  R16, ret+16(FP)
+	RET
+
+build_table_check1_ok:
+	TST  R1, R1
+	BNE  build_table_check2_ok
+	CMP  R6, R7
+	BNE  build_table_check2_ok
+	MOVD ctx+8(FP), R0
+	MOVD R7, 24(R0)
+	MOVD R6, 32(R0)
+	MOVD $+3, R16
+	MOVD R16, ret+16(FP)
+	RET
+
+build_table_check2_ok:
+	ADD  $1, R6, R6
+	CMP  R0, R6
+	BLT  build_table_main_table
+	MOVD $+0, R16
+	MOVD R16, ret+16(FP)
+	RET

zstd/fse_decoder_arm64_test.go

@@ -0,0 +1,98 @@
+//go:build arm64 && !appengine && !noasm && gc
+
+package zstd
+
+import (
+	"errors"
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+// buildDtableRef is a copy of the pure-Go reference algorithm
+// (fse_decoder_generic.go), kept here so the arm64 asm implementation can be
+// differentially tested against it even though the generic file is not
+// compiled on arm64.
+func buildDtableRef(s *fseDecoder) error {
+	tableSize := uint32(1 << s.actualTableLog)
+	highThreshold := tableSize - 1
+	symbolNext := s.stateTable[:256]
+
+	{
+		for i, v := range s.norm[:s.symbolLen] {
+			if v == -1 {
+				s.dt[highThreshold].setAddBits(uint8(i))
+				highThreshold--
+				v = 1
+			}
+			symbolNext[i] = uint16(v)
+		}
+	}
+
+	{
+		tableMask := tableSize - 1
+		step := tableStep(tableSize)
+		position := uint32(0)
+		for ss, v := range s.norm[:s.symbolLen] {
+			for i := 0; i < int(v); i++ {
+				s.dt[position].setAddBits(uint8(ss))
+				for {
+					position = (position + step) & tableMask
+					if position <= highThreshold {
+						break
+					}
+				}
+			}
+		}
+		if position != 0 {
+			return errors.New("corrupted input (position != 0)")
+		}
+	}
+
+	{
+		tableSize := uint16(1 << s.actualTableLog)
+		for u, v := range s.dt[:tableSize] {
+			symbol := v.addBits()
+			nextState := symbolNext[symbol]
+			symbolNext[symbol] = nextState + 1
+			nBits := s.actualTableLog - byte(highBits(uint32(nextState)))
+			s.dt[u&maxTableMask].setNBits(nBits)
+			newState := (nextState << nBits) - tableSize
+			if newState > tableSize {
+				return fmt.Errorf("newState (%d) outside table size (%d)", newState, tableSize)
+			}
+			if newState == uint16(u) && nBits == 0 {
+				return fmt.Errorf("newState (%d) == oldState (%d) and no bits", newState, u)
+			}
+			s.dt[u&maxTableMask].setNewState(newState)
+		}
+	}
+	return nil
+}
+
+func TestBuildDtableARM64MatchesReference(t *testing.T) {
+	for name, mk := range predefinedFSEInputs() {
+		t.Run(name, func(t *testing.T) {
+			got := mk()
+			if err := got.buildDtable(); err != nil {
+				t.Fatalf("asm buildDtable: %v", err)
+			}
+			want := mk()
+			if err := buildDtableRef(want); err != nil {
+				t.Fatalf("reference buildDtable: %v", err)
+			}
+			tableSize := 1 << got.actualTableLog
+			if !reflect.DeepEqual(got.dt[:tableSize], want.dt[:tableSize]) {
+				for i := range tableSize {
+					if got.dt[i] != want.dt[i] {
+						t.Errorf("dt[%d]: asm %#016x, ref %#016x", i, uint64(got.dt[i]), uint64(want.dt[i]))
+					}
+				}
+				t.FailNow()
+			}
+			if got.stateTable != want.stateTable {
+				t.Errorf("stateTable mismatch:\nasm %v\nref %v", got.stateTable, want.stateTable)
+			}
+		})
+	}
+}

zstd/fse_decoder_amd64.go zstd/fse_decoder_asm.gozstd/fse_decoder_amd64.go renamed to zstd/fse_decoder_asm.go

@@ -1,11 +1,15 @@
-//go:build amd64 && !appengine && !noasm && gc
+//go:build (amd64 || arm64) && !appengine && !noasm && gc
 
 package zstd
 
 import (
 	"fmt"
 )
 
+// buildDtable_asm is generated by _generate/gen_fse.go and lowered to each
+// architecture (amd64 by goasm, arm64 by the avo arm64 lowering printer). The
+// Go side is identical across architectures, so it lives here.
+
 type buildDtableAsmContext struct {
 	// inputs
 	stateTable *uint16
@@ -18,7 +22,7 @@ type buildDtableAsmContext struct {
 	errParam2 uint64
 }
 
-// buildDtable_asm is an x86 assembly implementation of fseDecoder.buildDtable.
+// buildDtable_asm is an assembly implementation of fseDecoder.buildDtable.
 // Function returns non-zero exit code on error.
 //
 //go:noescape