Skip to content

Commit 0875da4

Browse files
authored
Merge pull request #35 from ahrav/perf/32-neon-kernels
perf: NEON search kernels for pair-filter and two-value scans (arm64)
2 parents 0b533ea + 6e62ec2 commit 0875da4

10 files changed

Lines changed: 790 additions & 14 deletions

.github/workflows/ci.yml

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
name: CI
2+
3+
on:
4+
push:
5+
branches: [master]
6+
pull_request:
7+
8+
permissions:
9+
contents: read
10+
11+
jobs:
12+
test:
13+
strategy:
14+
fail-fast: false
15+
matrix:
16+
# ubuntu-24.04-arm assembles index_arm64.s and runs the NEON
17+
# kernel differential and guard-page tests natively; ubuntu-24.04
18+
# covers the portable paths on amd64.
19+
os: [ubuntu-24.04, ubuntu-24.04-arm]
20+
runs-on: ${{ matrix.os }}
21+
steps:
22+
- uses: actions/checkout@v4
23+
- uses: actions/setup-go@v5
24+
with:
25+
go-version-file: go.mod
26+
- name: gofmt
27+
run: test -z "$(gofmt -l .)"
28+
- name: vet
29+
run: go vet ./...
30+
- name: test
31+
run: go test ./...
32+
- name: vet and test with purego tag
33+
run: |
34+
go vet -tags purego ./...
35+
go test -tags purego ./...
36+
- name: cross-compile vet for other arm64 targets
37+
if: matrix.os == 'ubuntu-24.04-arm'
38+
run: |
39+
for goos in darwin windows freebsd netbsd openbsd; do
40+
echo "vet GOOS=$goos GOARCH=arm64"
41+
GOOS=$goos GOARCH=arm64 go vet ./...
42+
done

index_arm64.go

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
//go:build arm64 && !purego
2+
3+
package ahocorasick
4+
5+
// Per-kernel availability constants: call sites branch on these, so
6+
// unused strategies compile away. On arm64 both NEON kernels beat the
7+
// portable strategies (single-pass NEON vs SWAR or dual-pass windowed
8+
// IndexByte).
9+
const (
10+
hasPairKernel = true
11+
hasOr2Kernel = true
12+
)
13+
14+
// indexPair2Asm returns the smallest i in [0, m&^31) with p[i] == a and
15+
// p[i+d] == b, or -1 when the full 32-byte blocks contain no such
16+
// position (the caller scans the remaining tail with scalar code).
17+
// The caller must guarantee p[0 : m&^31 + d] is readable.
18+
//
19+
//go:noescape
20+
func indexPair2Asm(p *byte, m int, a, b byte, d int) int
21+
22+
// indexOr2Asm returns the smallest i in [0, m&^31) with p[i] == a or
23+
// p[i] == b, or -1 when the full blocks contain neither value. Reads
24+
// p[0 : m&^31] only.
25+
//
26+
//go:noescape
27+
func indexOr2Asm(p *byte, m int, a, b byte) int
28+
29+
func indexPair2(p []byte, m int, a, b byte, d int) int {
30+
return indexPair2Asm(&p[0], m, a, b, d)
31+
}
32+
33+
func indexOr2(p []byte, m int, a, b byte) int {
34+
return indexOr2Asm(&p[0], m, a, b)
35+
}

index_arm64.s

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
// NEON search kernels for the single-pattern fast path and the two-value
2+
// root skip. Both follow the internal/bytealg syndrome idiom: compare 32
3+
// input bytes per iteration, AND the match mask with the magic constant
4+
// 0x40100401 (bytes of each 4-byte group get distinct bits 1/4/16/64),
5+
// then two pairwise adds fold the 256-bit mask into a 64-bit syndrome
6+
// with two bits per input byte, in input order. RBIT+CLZ converts the
7+
// first set bit into the matching byte's offset.
8+
//
9+
// Neither kernel reads past its caller-proven region: both process only
10+
// full 32-byte blocks and report "no hit in the full blocks" with -1;
11+
// the Go callers finish the < 32-byte tail with scalar code.
12+
13+
//go:build !purego
14+
15+
#include "textflag.h"
16+
17+
// func indexPair2Asm(p *byte, m int, a byte, b byte, d int) int
18+
//
19+
// Returns the smallest i in [0, m&^31) with p[i] == a && p[i+d] == b,
20+
// or -1 if there is none. The caller guarantees p[0 : m&^31 + d] is
21+
// readable (stream B reads p[d : d+(m&^31)]).
22+
TEXT ·indexPair2Asm(SB), NOSPLIT, $0-40
23+
MOVD p+0(FP), R0
24+
MOVD m+8(FP), R1
25+
MOVBU a+16(FP), R2
26+
MOVBU b+17(FP), R3
27+
MOVD d+24(FP), R4
28+
ADD R0, R4, R4 // stream B cursor = p + d
29+
VMOV R2, V0.B16 // broadcast a
30+
VMOV R3, V7.B16 // broadcast b
31+
MOVD $0x40100401, R5
32+
VMOV R5, V5.S4 // syndrome magic
33+
MOVD ZR, R6 // positions consumed
34+
35+
loop:
36+
SUB R6, R1, R7 // remaining positions
37+
CMP $32, R7
38+
BLT notfound
39+
VLD1.P 32(R0), [V1.B16, V2.B16]
40+
VLD1.P 32(R4), [V3.B16, V4.B16]
41+
VCMEQ V0.B16, V1.B16, V1.B16
42+
VCMEQ V0.B16, V2.B16, V2.B16
43+
VCMEQ V7.B16, V3.B16, V3.B16
44+
VCMEQ V7.B16, V4.B16, V4.B16
45+
VAND V3.B16, V1.B16, V3.B16 // pair condition, bytes 0-15
46+
VAND V4.B16, V2.B16, V4.B16 // pair condition, bytes 16-31
47+
// Cheap existence check in the hot loop (bytealg idiom); the
48+
// positional syndrome is computed only after a hit.
49+
VORR V4.B16, V3.B16, V6.B16
50+
VADDP V6.D2, V6.D2, V6.D2
51+
VMOV V6.D[0], R8
52+
CBNZ R8, found
53+
ADD $32, R6
54+
B loop
55+
56+
found:
57+
VAND V5.B16, V3.B16, V3.B16
58+
VAND V5.B16, V4.B16, V4.B16
59+
VADDP V4.B16, V3.B16, V6.B16 // 256 -> 128
60+
VADDP V6.B16, V6.B16, V6.B16 // 128 -> 64
61+
VMOV V6.D[0], R8
62+
RBIT R8, R8
63+
CLZ R8, R8
64+
ADD R8>>1, R6, R0 // index = consumed + syndrome-bit/2
65+
MOVD R0, ret+32(FP)
66+
RET
67+
68+
notfound:
69+
MOVD $-1, R0
70+
MOVD R0, ret+32(FP)
71+
RET
72+
73+
// func indexOr2Asm(p *byte, m int, a byte, b byte) int
74+
//
75+
// Returns the smallest i in [0, m&^31) with p[i] == a || p[i] == b, or
76+
// -1 if there is none. Reads p[0 : m&^31] only.
77+
TEXT ·indexOr2Asm(SB), NOSPLIT, $0-32
78+
MOVD p+0(FP), R0
79+
MOVD m+8(FP), R1
80+
MOVBU a+16(FP), R2
81+
MOVBU b+17(FP), R3
82+
VMOV R2, V0.B16
83+
VMOV R3, V7.B16
84+
MOVD $0x40100401, R5
85+
VMOV R5, V5.S4
86+
MOVD ZR, R6
87+
88+
loop:
89+
SUB R6, R1, R7
90+
CMP $32, R7
91+
BLT notfound
92+
VLD1.P 32(R0), [V1.B16, V2.B16]
93+
VCMEQ V0.B16, V1.B16, V3.B16
94+
VCMEQ V0.B16, V2.B16, V4.B16
95+
VCMEQ V7.B16, V1.B16, V8.B16
96+
VCMEQ V7.B16, V2.B16, V9.B16
97+
VORR V8.B16, V3.B16, V3.B16
98+
VORR V9.B16, V4.B16, V4.B16
99+
// Cheap existence check; positional syndrome only after a hit.
100+
VORR V4.B16, V3.B16, V6.B16
101+
VADDP V6.D2, V6.D2, V6.D2
102+
VMOV V6.D[0], R8
103+
CBNZ R8, found
104+
ADD $32, R6
105+
B loop
106+
107+
found:
108+
VAND V5.B16, V3.B16, V3.B16
109+
VAND V5.B16, V4.B16, V4.B16
110+
VADDP V4.B16, V3.B16, V6.B16
111+
VADDP V6.B16, V6.B16, V6.B16
112+
VMOV V6.D[0], R8
113+
RBIT R8, R8
114+
CLZ R8, R8
115+
ADD R8>>1, R6, R0
116+
MOVD R0, ret+24(FP)
117+
RET
118+
119+
notfound:
120+
MOVD $-1, R0
121+
MOVD R0, ret+24(FP)
122+
RET

index_generic.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
//go:build !arm64 || purego
2+
3+
package ahocorasick
4+
5+
// No vector search kernels on this build: the scan paths use the
6+
// portable strategies (rare-byte IndexByte, SWAR pair scan, windowed
7+
// per-value IndexByte). The stubs below are never called; call sites
8+
// are gated on the constant and eliminated at compile time.
9+
const (
10+
hasPairKernel = false
11+
hasOr2Kernel = false
12+
)
13+
14+
func indexPair2(p []byte, m int, a, b byte, d int) int { panic("unreachable") }
15+
16+
func indexOr2(p []byte, m int, a, b byte) int { panic("unreachable") }

index_kernel_bench_test.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
//go:build arm64 && !purego
2+
3+
package ahocorasick
4+
5+
import (
6+
"bytes"
7+
"testing"
8+
)
9+
10+
func BenchmarkKernelPair2(b *testing.B) {
11+
p := bytes.Repeat([]byte{'x'}, 100032)
12+
b.SetBytes(100000)
13+
for n := 0; n < b.N; n++ {
14+
if indexPair2(p, 100000, 'A', 'B', 3) != -1 {
15+
b.Fatal("unexpected hit")
16+
}
17+
}
18+
}
19+
20+
func BenchmarkKernelOr2(b *testing.B) {
21+
p := bytes.Repeat([]byte{'x'}, 100032)
22+
b.SetBytes(100000)
23+
for n := 0; n < b.N; n++ {
24+
if indexOr2(p, 100000, 'A', 'B') != -1 {
25+
b.Fatal("unexpected hit")
26+
}
27+
}
28+
}

index_kernel_test.go

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
//go:build arm64 && !purego
2+
3+
package ahocorasick
4+
5+
// Kernel-level tests for the vector search kernels: exhaustive differential
6+
// against scalar oracles across sizes, positions, and pair distances. The
7+
// guard-page test proving the read contracts lives in
8+
// index_kernel_unix_test.go (mmap/mprotect are unix-only).
9+
10+
import (
11+
"math/rand"
12+
"testing"
13+
)
14+
15+
func oraclePair2(p []byte, m int, a, b byte, d int) int {
16+
m &^= 31
17+
for i := 0; i < m; i++ {
18+
if p[i] == a && p[i+d] == b {
19+
return i
20+
}
21+
}
22+
return -1
23+
}
24+
25+
func oracleOr2(p []byte, m int, a, b byte) int {
26+
m &^= 31
27+
for i := 0; i < m; i++ {
28+
if p[i] == a || p[i] == b {
29+
return i
30+
}
31+
}
32+
return -1
33+
}
34+
35+
func TestIndexPair2Exhaustive(t *testing.T) {
36+
// Sizes crossing every block boundary; hit planted at every position;
37+
// distances covering typical pattern spans plus the beyond-one-block
38+
// cases (d is bounded only by pattern length, so d > 32 must place
39+
// stream B more than a whole block past stream A).
40+
for _, m := range []int{0, 1, 31, 32, 33, 63, 64, 65, 96, 127, 130} {
41+
for _, d := range []int{1, 2, 5, 7, 15, 32, 33, 64, 100} {
42+
buf := make([]byte, m+d)
43+
for i := range buf {
44+
buf[i] = 'x'
45+
}
46+
// No hit anywhere.
47+
if got := indexPair2(append([]byte{}, buf...), m, 'A', 'B', d); m > 0 && got != -1 {
48+
t.Fatalf("m=%d d=%d empty: got %d want -1", m, d, got)
49+
}
50+
for pos := 0; pos < m; pos++ {
51+
p := append([]byte{}, buf...)
52+
p[pos] = 'A'
53+
p[pos+d] = 'B'
54+
want := oraclePair2(p, m, 'A', 'B', d)
55+
got := indexPair2(p, m, 'A', 'B', d)
56+
if got != want {
57+
t.Fatalf("m=%d d=%d pos=%d: got %d want %d", m, d, pos, got, want)
58+
}
59+
}
60+
}
61+
}
62+
}
63+
64+
func TestIndexPair2EqualBytes(t *testing.T) {
65+
// a == b and overlapping planted pairs.
66+
p := make([]byte, 128)
67+
for i := range p {
68+
p[i] = 'x'
69+
}
70+
p[40], p[43] = 'z', 'z'
71+
if got, want := indexPair2(p, 96, 'z', 'z', 3), 40; got != want {
72+
t.Fatalf("got %d want %d", got, want)
73+
}
74+
}
75+
76+
func TestIndexOr2Exhaustive(t *testing.T) {
77+
for _, m := range []int{0, 1, 31, 32, 33, 64, 65, 127, 130} {
78+
buf := make([]byte, m)
79+
for i := range buf {
80+
buf[i] = 'x'
81+
}
82+
for pos := 0; pos < m; pos++ {
83+
for _, c := range []byte{'A', 'B'} {
84+
p := append([]byte{}, buf...)
85+
p[pos] = c
86+
want := oracleOr2(p, m, 'A', 'B')
87+
got := indexOr2(p, m, 'A', 'B')
88+
if got != want {
89+
t.Fatalf("m=%d pos=%d c=%c: got %d want %d", m, pos, c, got, want)
90+
}
91+
}
92+
}
93+
}
94+
}
95+
96+
func TestIndexKernelsRandomDifferential(t *testing.T) {
97+
rng := rand.New(rand.NewSource(77))
98+
for iter := 0; iter < 5000; iter++ {
99+
m := rng.Intn(300)
100+
d := 1 + rng.Intn(16)
101+
p := make([]byte, m+d)
102+
for i := range p {
103+
p[i] = byte('a' + rng.Intn(4)) // small alphabet: many hits
104+
}
105+
a, b := byte('a'+rng.Intn(4)), byte('a'+rng.Intn(4))
106+
if m > 0 {
107+
if got, want := indexPair2(p, m, a, b, d), oraclePair2(p, m, a, b, d); got != want {
108+
t.Fatalf("pair m=%d d=%d a=%c b=%c: got %d want %d (%q)", m, d, a, b, got, want, p)
109+
}
110+
if got, want := indexOr2(p, m, a, b), oracleOr2(p, m, a, b); got != want {
111+
t.Fatalf("or m=%d a=%c b=%c: got %d want %d (%q)", m, a, b, got, want, p)
112+
}
113+
}
114+
}
115+
}

0 commit comments

Comments
 (0)