Skip to content

Commit 7ec979a

Browse files
fpetkovskifpetkovski
authored
Merge pull request #10 from MichaHoffmann/mhoffmann/improve-bitpack-on-amd64
pack, unpack: improve amd64 paths
2 parents fa1aca9 + 8cc3230 commit 7ec979a

27 files changed

Lines changed: 12015 additions & 1421 deletions

‎.github/workflows/test.yml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,3 +39,15 @@ jobs:
3939
run: |
4040
go test -v -tags=${{ matrix.tags }} ./...
4141
cd examples && go test -v
42+
43+
- name: Run tests without AVX2
44+
if: matrix.arch == 'amd64' && matrix.tags == ''
45+
env:
46+
GODEBUG: cpu.avx2=off
47+
run: go test -v ./...
48+
49+
- name: Check generated files
50+
if: matrix.arch == 'amd64' && matrix.tags == ''
51+
run: |
52+
go generate ./...
53+
git diff --exit-code

‎gen_pack_amd64.go‎

Lines changed: 318 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,318 @@
1+
//go:build ignore
2+
3+
package main
4+
5+
import (
6+
"bytes"
7+
"fmt"
8+
"io"
9+
"os"
10+
)
11+
12+
var widths = []int{
13+
2, 3, 4, 5, 6, 7,
14+
9, 10, 11, 12, 13, 14, 15,
15+
17, 18, 19, 20, 21, 22, 23,
16+
25, 26, 27, 28, 29, 30, 31,
17+
}
18+
19+
type integerType struct {
20+
name string
21+
fileName string
22+
goType string
23+
load string
24+
bitwiseAnd string
25+
valueStride int
26+
valueShift int
27+
}
28+
29+
var integerTypes = []integerType{
30+
{name: "Int32", fileName: "int32", goType: "int32", load: "MOVL", bitwiseAnd: "ANDL", valueStride: 4, valueShift: 2},
31+
{name: "Int64", fileName: "int64", goType: "int64", load: "MOVQ", bitwiseAnd: "ANDQ", valueStride: 8, valueShift: 3},
32+
}
33+
34+
func main() {
35+
for _, integerType := range integerTypes {
36+
generateGo(integerType)
37+
generateAssembly(integerType)
38+
}
39+
}
40+
41+
func generatedWidths(integerType integerType) []int {
42+
result := append([]int(nil), widths...)
43+
if integerType.name == "Int32" {
44+
return append(result, 32)
45+
}
46+
for width := 33; width <= 64; width++ {
47+
result = append(result, width)
48+
}
49+
return result
50+
}
51+
52+
func generateGo(integerType integerType) {
53+
var out bytes.Buffer
54+
f := &out
55+
56+
fmt.Fprintln(f, "//go:build !purego")
57+
fmt.Fprintln(f)
58+
fmt.Fprintln(f, "// Code generated by gen_pack_amd64.go; DO NOT EDIT.")
59+
fmt.Fprintln(f)
60+
fmt.Fprintln(f, "package bitpack")
61+
fmt.Fprintln(f)
62+
for _, width := range generatedWidths(integerType) {
63+
fmt.Fprintln(f, "//go:noescape")
64+
fmt.Fprintf(f, "func pack%sx%dbitsAMD64(dst []byte, src []%s)\n\n", integerType.name, width, integerType.goType)
65+
}
66+
fmt.Fprintf(f, "func pack%sGeneratedAMD64(dst []byte, src []%s, bitWidth uint) {\n", integerType.name, integerType.goType)
67+
fmt.Fprintln(f, "\tswitch bitWidth {")
68+
for _, width := range generatedWidths(integerType) {
69+
fmt.Fprintf(f, "\tcase %d:\n\t\tpack%sx%dbitsAMD64(dst, src)\n", width, integerType.name, width)
70+
}
71+
fmt.Fprintln(f, "\t}")
72+
fmt.Fprintln(f, "}")
73+
74+
fileName := fmt.Sprintf("pack_%s_generated_amd64.go", integerType.fileName)
75+
if err := os.WriteFile(fileName, out.Bytes(), 0o644); err != nil {
76+
panic(err)
77+
}
78+
}
79+
80+
func generateAssembly(integerType integerType) {
81+
var out bytes.Buffer
82+
f := &out
83+
84+
fmt.Fprintln(f, "//go:build !purego")
85+
fmt.Fprintln(f)
86+
fmt.Fprintln(f, "// Code generated by gen_pack_amd64.go; DO NOT EDIT.")
87+
fmt.Fprintln(f)
88+
fmt.Fprintln(f, "#include \"textflag.h\"")
89+
for _, width := range generatedWidths(integerType) {
90+
fmt.Fprintln(f)
91+
emitKernel(f, integerType, width)
92+
}
93+
94+
fileName := fmt.Sprintf("pack_%s_generated_amd64.s", integerType.fileName)
95+
if err := os.WriteFile(fileName, out.Bytes(), 0o644); err != nil {
96+
panic(err)
97+
}
98+
}
99+
100+
func emitKernel(f io.Writer, integerType integerType, width int) {
101+
fmt.Fprintf(f, "// func pack%sx%dbitsAMD64(dst []byte, src []%s)\n", integerType.name, width, integerType.goType)
102+
fmt.Fprintf(f, "TEXT ·pack%sx%dbitsAMD64(SB), NOSPLIT, $0-48\n", integerType.name, width)
103+
fmt.Fprintln(f, "\tMOVQ dst_base+0(FP), AX")
104+
fmt.Fprintln(f, "\tMOVQ src_base+24(FP), BX")
105+
fmt.Fprintln(f, "\tMOVQ src_len+32(FP), DX")
106+
fmt.Fprintf(f, "\tSHLQ $%d, DX\n", integerType.valueShift)
107+
fmt.Fprintln(f, "\tADDQ BX, DX")
108+
if width > 32 && width < 64 {
109+
fmt.Fprintf(f, "\tMOVQ $0x%x, R10\n", uint64(1)<<width-1)
110+
}
111+
fmt.Fprintf(f, "pack_%s_%dbit_amd64_loop:\n", integerType.fileName, width)
112+
113+
switch {
114+
case width == integerType.valueStride*8:
115+
emitCopyKernel(f, integerType)
116+
case width <= 7:
117+
emitSmallKernel(f, integerType, width)
118+
case width <= 15:
119+
emitMediumKernel(f, integerType, width)
120+
case width <= 21:
121+
emitLargeKernel(f, integerType, width)
122+
case width <= 31:
123+
emitPairKernel(f, integerType, width)
124+
default:
125+
emitWideKernel(f, integerType, width)
126+
}
127+
128+
fmt.Fprintf(f, "\tADDQ $%d, BX\n", 8*integerType.valueStride)
129+
fmt.Fprintf(f, "\tADDQ $%d, AX\n", width)
130+
fmt.Fprintln(f, "\tCMPQ BX, DX")
131+
fmt.Fprintf(f, "\tJB pack_%s_%dbit_amd64_loop\n", integerType.fileName, width)
132+
fmt.Fprintln(f, "\tRET")
133+
}
134+
135+
func emitCopyKernel(f io.Writer, integerType integerType) {
136+
for offset := 0; offset < 8*integerType.valueStride; offset += 16 {
137+
reg := fmt.Sprintf("X%d", offset/16)
138+
fmt.Fprintf(f, "\tMOVOU %d(BX), %s\n", offset, reg)
139+
}
140+
for offset := 0; offset < 8*integerType.valueStride; offset += 16 {
141+
reg := fmt.Sprintf("X%d", offset/16)
142+
fmt.Fprintf(f, "\tMOVOU %s, %d(AX)\n", reg, offset)
143+
}
144+
}
145+
146+
func emitWideKernel(f io.Writer, integerType integerType, width int) {
147+
wordCount := (width + 7) / 8
148+
for word := 0; word < wordCount; word++ {
149+
wordOffset := word * 64
150+
firstValue := wordOffset / width
151+
lastValue := min(7, (wordOffset+63)/width)
152+
for value := firstValue; value <= lastValue; value++ {
153+
reg := "R9"
154+
if value == firstValue {
155+
reg = "R8"
156+
}
157+
fmt.Fprintf(f, "\t%s %d(BX), %s\n", integerType.load, value*integerType.valueStride, reg)
158+
fmt.Fprintln(f, "\tANDQ R10, "+reg)
159+
shift := value*width - wordOffset
160+
switch {
161+
case shift > 0:
162+
fmt.Fprintf(f, "\tSHLQ $%d, %s\n", shift, reg)
163+
case shift < 0:
164+
fmt.Fprintf(f, "\tSHRQ $%d, %s\n", -shift, reg)
165+
}
166+
if value != firstValue {
167+
fmt.Fprintln(f, "\tORQ R9, R8")
168+
}
169+
}
170+
171+
offset := word * 8
172+
if word < width/8 {
173+
fmt.Fprintf(f, "\tMOVQ R8, %d(AX)\n", offset)
174+
} else {
175+
emitTailStore(f, "R8", offset, width%8)
176+
}
177+
}
178+
}
179+
180+
func emitSmallKernel(f io.Writer, integerType integerType, width int) {
181+
regs := []string{"R8", "R9", "R10", "R11", "R12", "R13", "R14", "R15"}
182+
for i, reg := range regs {
183+
emitValue(f, integerType, reg, i, i*width, width)
184+
}
185+
for stride := 1; stride < len(regs); stride *= 2 {
186+
for i := 0; i < len(regs); i += 2 * stride {
187+
fmt.Fprintf(f, "\tORQ %s, %s\n", regs[i+stride], regs[i])
188+
}
189+
}
190+
emitTailStore(f, regs[0], 0, width)
191+
}
192+
193+
func emitMediumKernel(f io.Writer, integerType integerType, width int) {
194+
low := []string{"R8", "R9", "R10", "R11"}
195+
high := []string{"R12", "R13", "R14", "R15"}
196+
for i, reg := range low {
197+
emitValue(f, integerType, reg, i, i*width, width)
198+
}
199+
for i, reg := range high {
200+
emitValue(f, integerType, reg, i+4, i*width, width)
201+
}
202+
emitFourValueReduction(f, low)
203+
emitFourValueReduction(f, high)
204+
205+
shift := 4 * width
206+
fmt.Fprintln(f, "\tMOVQ R12, R9")
207+
fmt.Fprintf(f, "\tSHLQ $%d, R9\n", shift)
208+
fmt.Fprintln(f, "\tORQ R9, R8")
209+
fmt.Fprintf(f, "\tSHRQ $%d, R12\n", 64-shift)
210+
fmt.Fprintln(f, "\tMOVQ R8, (AX)")
211+
emitTailStore(f, "R12", 8, width-8)
212+
}
213+
214+
func emitLargeKernel(f io.Writer, integerType integerType, width int) {
215+
groups := [][]string{
216+
{"R8", "R9", "R10"},
217+
{"R11", "R12", "R13"},
218+
{"R14", "R15"},
219+
}
220+
value := 0
221+
for _, group := range groups {
222+
for i, reg := range group {
223+
emitValue(f, integerType, reg, value, i*width, width)
224+
value++
225+
}
226+
for i := 1; i < len(group); i++ {
227+
fmt.Fprintf(f, "\tORQ %s, %s\n", group[i], group[0])
228+
}
229+
}
230+
231+
groupBits := 3 * width
232+
thirdOffset := 6*width - 64
233+
fmt.Fprintln(f, "\tMOVQ R11, R9")
234+
fmt.Fprintf(f, "\tSHLQ $%d, R9\n", groupBits)
235+
fmt.Fprintln(f, "\tORQ R9, R8")
236+
fmt.Fprintf(f, "\tSHRQ $%d, R11\n", 64-groupBits)
237+
fmt.Fprintln(f, "\tMOVQ R14, R15")
238+
fmt.Fprintf(f, "\tSHLQ $%d, R15\n", thirdOffset)
239+
fmt.Fprintln(f, "\tORQ R15, R11")
240+
fmt.Fprintf(f, "\tSHRQ $%d, R14\n", 64-thirdOffset)
241+
fmt.Fprintln(f, "\tMOVQ R8, (AX)")
242+
fmt.Fprintln(f, "\tMOVQ R11, 8(AX)")
243+
emitTailStore(f, "R14", 16, width-16)
244+
}
245+
246+
func emitPairKernel(f io.Writer, integerType integerType, width int) {
247+
pairs := [][2]string{
248+
{"R8", "R9"},
249+
{"R10", "R11"},
250+
{"R12", "R13"},
251+
{"R14", "R15"},
252+
}
253+
for i, pair := range pairs {
254+
emitValue(f, integerType, pair[0], 2*i, 0, width)
255+
emitValue(f, integerType, pair[1], 2*i+1, width, width)
256+
fmt.Fprintf(f, "\tORQ %s, %s\n", pair[1], pair[0])
257+
}
258+
259+
pairBits := 2 * width
260+
fmt.Fprintln(f, "\tMOVQ R10, R9")
261+
fmt.Fprintf(f, "\tSHLQ $%d, R9\n", pairBits)
262+
fmt.Fprintln(f, "\tORQ R9, R8")
263+
fmt.Fprintln(f, "\tMOVQ R10, R11")
264+
fmt.Fprintf(f, "\tSHRQ $%d, R11\n", 64-pairBits)
265+
fmt.Fprintln(f, "\tMOVQ R12, R13")
266+
fmt.Fprintf(f, "\tSHLQ $%d, R13\n", 2*pairBits-64)
267+
fmt.Fprintln(f, "\tORQ R13, R11")
268+
fmt.Fprintf(f, "\tSHRQ $%d, R12\n", 128-2*pairBits)
269+
fmt.Fprintln(f, "\tMOVQ R14, R15")
270+
fmt.Fprintf(f, "\tSHLQ $%d, R15\n", 3*pairBits-128)
271+
fmt.Fprintln(f, "\tORQ R15, R12")
272+
273+
fmt.Fprintln(f, "\tMOVQ R8, (AX)")
274+
fmt.Fprintln(f, "\tMOVQ R11, 8(AX)")
275+
if width <= 23 {
276+
emitTailStore(f, "R12", 16, width-16)
277+
return
278+
}
279+
fmt.Fprintf(f, "\tSHRQ $%d, R14\n", 192-3*pairBits)
280+
fmt.Fprintln(f, "\tMOVQ R12, 16(AX)")
281+
emitTailStore(f, "R14", 24, width-24)
282+
}
283+
284+
func emitValue(f io.Writer, integerType integerType, reg string, index, shift, width int) {
285+
fmt.Fprintf(f, "\t%s %d(BX), %s\n", integerType.load, index*integerType.valueStride, reg)
286+
fmt.Fprintf(f, "\t%s $0x%x, %s\n", integerType.bitwiseAnd, (1<<width)-1, reg)
287+
if shift != 0 {
288+
fmt.Fprintf(f, "\tSHLQ $%d, %s\n", shift, reg)
289+
}
290+
}
291+
292+
func emitFourValueReduction(f io.Writer, regs []string) {
293+
fmt.Fprintf(f, "\tORQ %s, %s\n", regs[1], regs[0])
294+
fmt.Fprintf(f, "\tORQ %s, %s\n", regs[3], regs[2])
295+
fmt.Fprintf(f, "\tORQ %s, %s\n", regs[2], regs[0])
296+
}
297+
298+
func emitTailStore(f io.Writer, reg string, offset, byteCount int) {
299+
if byteCount >= 4 {
300+
fmt.Fprintf(f, "\tMOVL %s, %d(AX)\n", reg, offset)
301+
offset += 4
302+
byteCount -= 4
303+
if byteCount > 0 {
304+
fmt.Fprintf(f, "\tSHRQ $32, %s\n", reg)
305+
}
306+
}
307+
if byteCount >= 2 {
308+
fmt.Fprintf(f, "\tMOVW %s, %d(AX)\n", reg, offset)
309+
offset += 2
310+
byteCount -= 2
311+
if byteCount > 0 {
312+
fmt.Fprintf(f, "\tSHRQ $16, %s\n", reg)
313+
}
314+
}
315+
if byteCount == 1 {
316+
fmt.Fprintf(f, "\tMOVB %s, %d(AX)\n", reg, offset)
317+
}
318+
}

0 commit comments

Comments
 (0)