From c01b959723b73843f45e622424c09b8318c244c3 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sun, 11 Feb 2024 10:47:34 +0530 Subject: [PATCH] Fix Go unaligned index implementation --- tools/simdstring/generate.go | 57 ++++++++++++++++++++++++----- tools/simdstring/intrinsics_test.go | 34 ++++++++++------- tools/simdstring/test.sh | 2 +- 3 files changed, 69 insertions(+), 24 deletions(-) diff --git a/tools/simdstring/generate.go b/tools/simdstring/generate.go index 81b589b7b..8dea7f815 100644 --- a/tools/simdstring/generate.go +++ b/tools/simdstring/generate.go @@ -424,10 +424,18 @@ func (f *Function) MaskForCountDestructive(vec, ans Register) { } } -func (f *Function) ShiftSelfRight(self, amt any) { - op := "SHRQ" - if f.ISA.Goarch == ARM64 { - op = "LSR" +func (f *Function) shift_self(right bool, self, amt any) { + op := "" + if right { + op = "SHRQ" + if f.ISA.Goarch == ARM64 { + op = "LSR" + } + } else { + op = "SHLQ" + if f.ISA.Goarch == ARM64 { + op = "LSL" + } } switch v := amt.(type) { case Register: @@ -440,13 +448,31 @@ func (f *Function) ShiftSelfRight(self, amt any) { } } -func (f *Function) ShiftMaskRightDestructive(mask, amt Register) { +func (f *Function) ShiftSelfRight(self, amt any) { + f.shift_self(true, self, amt) +} + +func (f *Function) ShiftSelfLeft(self, amt any) { + f.shift_self(false, self, amt) +} + +func (f *Function) ShiftMaskRightDestructive(mask, amt any) { // The amt register is clobbered by this function - if f.ISA.Goarch == ARM64 { - f.Comment("The mask has 4 bits per byte, so multiply", amt, "by 4") - f.ShiftSelfRight(amt, 2) + switch n := amt.(type) { + case Register: + if f.ISA.Goarch == ARM64 { + f.Comment("The mask has 4 bits per byte, so multiply", n, "by 4") + f.ShiftSelfLeft(n, 2) + } + f.ShiftSelfRight(mask, n) + case int: + if f.ISA.Goarch == ARM64 { + n <<= 2 + } + f.ShiftSelfRight(mask, n) + default: + panic(fmt.Sprintf("Cannot shift by: %s", amt)) } - f.ShiftSelfRight(mask, amt) } func (f *Function) CountLeadingZeroBytesInMask(src, ans Register) { @@ -1101,6 +1127,15 @@ func (f *Function) AndSelf(self Register, val any) { f.AddTrailingComment(self, "&=", val) } +func (f *Function) NegateSelf(self Register) { + if f.ISA.Goarch == "ARM64" { + f.instr("NEG", self, self) + } else { + f.instr("NEGQ", self) + } + f.AddTrailingComment(self, "*= -1") +} + func (f *Function) AddToSelf(self Register, val any) { f.instr(f.ISA.NativeAdd(), val_repr_for_arithmetic(val), self) // pos += sizeof(vec) f.AddTrailingComment(self, "+=", val) @@ -1379,7 +1414,9 @@ func (s *State) index_func(f *Function, test_bytes func(bytes_to_test, test_ans f.MaskForCountDestructive(test_ans, mask) f.Comment("We need to shift out the possible extra bytes at the start of the string caused by the unaligned read") f.ShiftMaskRightDestructive(mask, unaligned_bytes) - f.JumpIfNonZero(mask, "byte_found_in_mask") + f.JumpIfZero(mask, "loop_start") + f.CopyRegister(data_start, pos) + f.JumpTo("byte_found_in_mask") }() f.Comment("Now loop over aligned blocks") diff --git a/tools/simdstring/intrinsics_test.go b/tools/simdstring/intrinsics_test.go index ac8397aeb..d99417a4d 100644 --- a/tools/simdstring/intrinsics_test.go +++ b/tools/simdstring/intrinsics_test.go @@ -5,6 +5,7 @@ package simdstring import ( "bytes" "fmt" + "kitty/tools/utils" "runtime" "strings" "testing" @@ -119,21 +120,25 @@ func addressof_data(b []byte) uintptr { return uintptr(unsafe.Pointer(&b[0])) } -func aligned_slice(sz, alignment int) []byte { - ans := make([]byte, sz+alignment) +func aligned_slice(sz, alignment int) ([]byte, []byte) { + ans := make([]byte, sz+alignment+512) a := addressof_data(ans) a &= uintptr(alignment - 1) extra := uintptr(alignment) - a - return ans[extra : extra+uintptr(sz)] + utils.Memset(ans, '<') + utils.Memset(ans[extra+uintptr(sz):], '>') + return ans[extra : extra+uintptr(sz)], ans } func TestSIMDStringOps(t *testing.T) { sizes := get_sizes(t) - test := func(haystack []byte, a, b byte) { + test := func(haystack []byte, a, b byte, dont_pad ...bool) { var actual int - safe_haystack := append(bytes.Repeat([]byte{'<'}, 64), haystack...) - safe_haystack = append(safe_haystack, bytes.Repeat([]byte{'>'}, 64)...) - haystack = safe_haystack[64 : 64+len(haystack)] + if len(dont_pad) == 0 || !dont_pad[0] { + safe_haystack := append(bytes.Repeat([]byte{'<'}, 64), haystack...) + safe_haystack = append(safe_haystack, bytes.Repeat([]byte{'>'}, 64)...) + haystack = safe_haystack[64 : 64+len(haystack)] + } expected := index_byte2_scalar(haystack, a, b) for _, sz := range sizes { @@ -144,21 +149,24 @@ func TestSIMDStringOps(t *testing.T) { actual = index_byte2_asm_256(haystack, a, b) } if actual != expected { - t.Fatalf("Failed to find '%c' or '%c' in: %#v (%d != %d) at size: %d", a, b, string(haystack), expected, actual, sz) + t.Fatalf("Failed to find '%c' or '%c' in: %#v at align: %d (expected: %d != actual: %d) at size: %d", + a, b, string(haystack), addressof_data(haystack)&uintptr(sz-1), expected, actual, sz) } } } // test alignment issues for sz := 0; sz < 32; sz++ { - q := aligned_slice(sz+3, 64)[sz:] + as, _ := aligned_slice(sz+3, 32) + q := as[sz:] q[0] = 'a' q[1] = 'b' q[2] = 'c' - test(q, '<', '>') - test(q, '<', 'a') - test(q, '<', 'b') - test(q, 'c', '>') + test(q, '<', '>', true) + test(q, ' ', 'b', true) + test(q, '<', 'a', true) + test(q, '<', 'b', true) + test(q, 'c', '>', true) } tests := func(h string, a, b byte) { diff --git a/tools/simdstring/test.sh b/tools/simdstring/test.sh index fc5dca898..b1066667e 100755 --- a/tools/simdstring/test.sh +++ b/tools/simdstring/test.sh @@ -7,6 +7,6 @@ # echo -e "\x1b[32mtesting amd64\x1b[m" && go test -v && - echo -e "\x1b[32mtesting arm64\x1b[m" && GOARCH=arm64 go test -c && qemu-aarch64-static simdstring.test -test.v + echo -e "\x1b[32mtesting arm64\x1b[m" && GOARCH=arm64 go test -v -exec qemu-aarch64-static