perf(math): accelerate FHT with SIMD

internal/ailego/math/Makefile

@@ -0,0 +1,29 @@
+GO ?= go
+GOAT ?= $(GO) tool goat
+GOAT_TARGET_AMD64 ?= amd64
+GOAT_TARGET_ARM64 ?= arm64
+GOAT_OUTPUT_DIR ?= .goat/mathutil
+MV ?= mv
+
+.PHONY: generate avx2 avx512 neon clean
+
+generate: avx2 avx512 neon clean
+
+avx2:
+	mkdir -p $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/fht_avx2.c --target $(GOAT_TARGET_AMD64) -O3 -mavx2 -o $(GOAT_OUTPUT_DIR)
+	$(MV) $(GOAT_OUTPUT_DIR)/fht_avx2.go $(GOAT_OUTPUT_DIR)/fht_avx2.s .
+
+avx512:
+	mkdir -p $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/fht_avx512.c --target $(GOAT_TARGET_AMD64) -O3 -mavx512f -mavx512dq -o $(GOAT_OUTPUT_DIR)
+	$(MV) $(GOAT_OUTPUT_DIR)/fht_avx512.go $(GOAT_OUTPUT_DIR)/fht_avx512.s .
+
+neon:
+	mkdir -p $(GOAT_OUTPUT_DIR)
+	$(GOAT) src/fht_neon.c --target $(GOAT_TARGET_ARM64) -O3 -o $(GOAT_OUTPUT_DIR)
+	$(MV) $(GOAT_OUTPUT_DIR)/fht_neon.go $(GOAT_OUTPUT_DIR)/fht_neon.s .
+
+clean:
+	$(RM) src/*.o src/*.s
+	$(RM) -r .goat

internal/ailego/math/fht.go

@@ -24,16 +24,29 @@ var (
 	ErrShortSignBits    = errors.New("ailego: sign-bit buffer is too short")
 )
 
+type fhtUnaryKernel func(data []float32)
+type fhtFlipSignsKernel func(signs []byte, data []float32)
+
+type fhtKernels struct {
+	flipSigns      fhtFlipSignsKernel
+	kacWalk        fhtUnaryKernel
+	inverseKacWalk fhtUnaryKernel
+	inPlace        fhtUnaryKernel
+}
+
+var activeFHTKernels = fhtKernels{
+	flipSigns:      fhtFlipSignsScalar,
+	kacWalk:        fhtKacWalkScalar,
+	inverseKacWalk: fhtInverseKacWalkScalar,
+	inPlace:        fhtInPlaceScalar,
+}
+
 // FHTFlipSigns negates elements selected by the little-endian bits in signs.
 func FHTFlipSigns(signs []byte, data []float32) error {
 	if len(signs) < (len(data)+7)/8 {
 		return ErrShortSignBits
 	}
-	for index := range data {
-		if signs[index/8]&(1<<uint(index%8)) != 0 {
-			data[index] = -data[index]
-		}
-	}
+	activeFHTKernels.flipSigns(signs, data)
 	return nil
 }
 
@@ -42,6 +55,45 @@ func FHTInPlace(data []float32) error {
 	if len(data) == 0 || len(data)&(len(data)-1) != 0 {
 		return ErrInvalidFHTLength
 	}
+	activeFHTKernels.inPlace(data)
+	return nil
+}
+
+// FHTKacWalk performs the baseline pairwise Kac reduction for an arbitrary
+// non-empty length.
+func FHTKacWalk(data []float32) error {
+	if len(data) == 0 {
+		return ErrInvalidFHTLength
+	}
+	activeFHTKernels.kacWalk(data)
+	return nil
+}
+
+// FHTInverseKacWalk reverses FHTKacWalk.
+func FHTInverseKacWalk(data []float32) error {
+	if len(data) == 0 {
+		return ErrInvalidFHTLength
+	}
+	activeFHTKernels.inverseKacWalk(data)
+	return nil
+}
+
+// ScaleFloat32 multiplies data in place.
+func ScaleFloat32(data []float32, factor float32) {
+	for index := range data {
+		data[index] *= factor
+	}
+}
+
+func fhtFlipSignsScalar(signs []byte, data []float32) {
+	for index := range data {
+		if signs[index/8]&(1<<uint(index%8)) != 0 {
+			data[index] = -data[index]
+		}
+	}
+}
+
+func fhtInPlaceScalar(data []float32) {
 	for width := 1; width < len(data); width <<= 1 {
 		for block := 0; block < len(data); block += width << 1 {
 			for index := block; index < block+width; index++ {
@@ -51,15 +103,9 @@ func FHTInPlace(data []float32) error {
 			}
 		}
 	}
-	return nil
 }
 
-// FHTKacWalk performs the baseline pairwise Kac reduction for an arbitrary
-// non-empty length.
-func FHTKacWalk(data []float32) error {
-	if len(data) == 0 {
-		return ErrInvalidFHTLength
-	}
+func fhtKacWalkScalar(data []float32) {
 	half := len(data) / 2
 	base := len(data) % 2
 	offset := base + half
@@ -71,14 +117,9 @@ func FHTKacWalk(data []float32) error {
 	if base != 0 {
 		data[half] *= float32(math.Sqrt(2))
 	}
-	return nil
 }
 
-// FHTInverseKacWalk reverses FHTKacWalk.
-func FHTInverseKacWalk(data []float32) error {
-	if len(data) == 0 {
-		return ErrInvalidFHTLength
-	}
+func fhtInverseKacWalkScalar(data []float32) {
 	half := len(data) / 2
 	base := len(data) % 2
 	offset := base + half
@@ -90,12 +131,4 @@ func FHTInverseKacWalk(data []float32) error {
 		data[index] = (left + right) * .5
 		data[index+offset] = (left - right) * .5
 	}
-	return nil
-}
-
-// ScaleFloat32 multiplies data in place.
-func ScaleFloat32(data []float32, factor float32) {
-	for index := range data {
-		data[index] *= factor
-	}
 }

internal/ailego/math/fht_amd64.go

@@ -0,0 +1,121 @@
+//go:build !noasm && amd64
+
+// Copyright 2026-present the xvec project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mathutil
+
+import (
+	"math"
+	"unsafe"
+
+	"golang.org/x/sys/cpu"
+)
+
+//go:generate make avx2 avx512
+
+func init() {
+	switch {
+	case cpu.X86.HasAVX512F && cpu.X86.HasAVX512DQ:
+		activeFHTKernels = fhtKernels{
+			flipSigns:      fhtFlipSignsAVX512,
+			kacWalk:        fhtKacWalkAVX512,
+			inverseKacWalk: fhtInverseKacWalkAVX512,
+			inPlace:        fhtInPlaceAVX512,
+		}
+	case cpu.X86.HasAVX2:
+		activeFHTKernels = fhtKernels{
+			flipSigns:      fhtFlipSignsAVX2,
+			kacWalk:        fhtKacWalkAVX2,
+			inverseKacWalk: fhtInverseKacWalkAVX2,
+			inPlace:        fhtInPlaceAVX2,
+		}
+	}
+}
+
+func fhtFlipSignsAVX2(signs []byte, data []float32) {
+	if len(data) < 8 {
+		fhtFlipSignsScalar(signs, data)
+		return
+	}
+	xvec_avx2_fht_flip_signs(unsafe.Pointer(&signs[0]), unsafe.Pointer(&data[0]), int64(len(data)))
+}
+
+func fhtKacWalkAVX2(data []float32) {
+	if len(data)/2 < 8 {
+		fhtKacWalkScalar(data)
+		return
+	}
+	xvec_avx2_fht_kac_walk(unsafe.Pointer(&data[0]), int64(len(data)))
+	if len(data)%2 != 0 {
+		data[len(data)/2] *= float32(math.Sqrt2)
+	}
+}
+
+func fhtInverseKacWalkAVX2(data []float32) {
+	if len(data)/2 < 8 {
+		fhtInverseKacWalkScalar(data)
+		return
+	}
+	if len(data)%2 != 0 {
+		data[len(data)/2] *= float32(math.Sqrt(.5))
+	}
+	xvec_avx2_fht_inverse_kac_walk(unsafe.Pointer(&data[0]), int64(len(data)))
+}
+
+func fhtInPlaceAVX2(data []float32) {
+	if len(data) < 16 {
+		fhtInPlaceScalar(data)
+		return
+	}
+	xvec_avx2_fht_in_place(unsafe.Pointer(&data[0]), int64(len(data)))
+}
+
+func fhtFlipSignsAVX512(signs []byte, data []float32) {
+	if len(data) < 16 {
+		fhtFlipSignsScalar(signs, data)
+		return
+	}
+	xvec_avx512_fht_flip_signs(unsafe.Pointer(&signs[0]), unsafe.Pointer(&data[0]), int64(len(data)))
+}
+
+func fhtKacWalkAVX512(data []float32) {
+	if len(data)/2 < 16 {
+		fhtKacWalkScalar(data)
+		return
+	}
+	xvec_avx512_fht_kac_walk(unsafe.Pointer(&data[0]), int64(len(data)))
+	if len(data)%2 != 0 {
+		data[len(data)/2] *= float32(math.Sqrt2)
+	}
+}
+
+func fhtInverseKacWalkAVX512(data []float32) {
+	if len(data)/2 < 16 {
+		fhtInverseKacWalkScalar(data)
+		return
+	}
+	if len(data)%2 != 0 {
+		data[len(data)/2] *= float32(math.Sqrt(.5))
+	}
+	xvec_avx512_fht_inverse_kac_walk(unsafe.Pointer(&data[0]), int64(len(data)))
+}
+
+func fhtInPlaceAVX512(data []float32) {
+	if len(data) < 32 {
+		fhtInPlaceScalar(data)
+		return
+	}
+	xvec_avx512_fht_in_place(unsafe.Pointer(&data[0]), int64(len(data)))
+}

internal/ailego/math/fht_amd64_test.go

@@ -0,0 +1,82 @@
+//go:build !noasm && amd64
+
+// Copyright 2026-present the xvec project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mathutil
+
+import (
+	"slices"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sys/cpu"
+)
+
+func TestFHTAVX2(t *testing.T) {
+	if !cpu.X86.HasAVX2 {
+		t.Skip("AVX2 is not supported")
+	}
+	testFHTKernels(t, fhtKernels{
+		flipSigns:      fhtFlipSignsAVX2,
+		kacWalk:        fhtKacWalkAVX2,
+		inverseKacWalk: fhtInverseKacWalkAVX2,
+		inPlace:        fhtInPlaceAVX2,
+	})
+}
+
+func TestFHTAVX512ScalarFallback(t *testing.T) {
+	signs := []byte{0b10100101, 0b01011010}
+	input := []float32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
+	wantSigns := slices.Clone(input)
+	fhtFlipSignsScalar(signs, wantSigns)
+	gotSigns := slices.Clone(input)
+	fhtFlipSignsAVX512(signs, gotSigns)
+	require.Equal(t, wantSigns, gotSigns)
+
+	walkInput := make([]float32, 31)
+	for index := range walkInput {
+		walkInput[index] = float32(index + 1)
+	}
+	wantWalk := slices.Clone(walkInput)
+	fhtKacWalkScalar(wantWalk)
+	gotWalk := slices.Clone(walkInput)
+	fhtKacWalkAVX512(gotWalk)
+	require.Equal(t, wantWalk, gotWalk)
+
+	wantInverse := slices.Clone(wantWalk)
+	fhtInverseKacWalkScalar(wantInverse)
+	gotInverse := slices.Clone(gotWalk)
+	fhtInverseKacWalkAVX512(gotInverse)
+	require.Equal(t, wantInverse, gotInverse)
+
+	inPlaceInput := walkInput[:16]
+	wantInPlace := slices.Clone(inPlaceInput)
+	fhtInPlaceScalar(wantInPlace)
+	gotInPlace := slices.Clone(inPlaceInput)
+	fhtInPlaceAVX512(gotInPlace)
+	require.Equal(t, wantInPlace, gotInPlace)
+}
+
+func TestFHTAVX512(t *testing.T) {
+	if !cpu.X86.HasAVX512F || !cpu.X86.HasAVX512DQ {
+		t.Skip("AVX-512F and AVX-512DQ are not supported")
+	}
+	testFHTKernels(t, fhtKernels{
+		flipSigns:      fhtFlipSignsAVX512,
+		kacWalk:        fhtKacWalkAVX512,
+		inverseKacWalk: fhtInverseKacWalkAVX512,
+		inPlace:        fhtInPlaceAVX512,
+	})
+}