From f9fd6ffd465a9bc86e3ebc146a21d33cb9c54f72 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sun, 11 Feb 2024 06:10:33 +0530 Subject: [PATCH] Use only aligned loads for index funcs Also obviates the necessity for safe slice wrappers --- tools/simdstring/generate.go | 143 ++++++++++++++++++---------- tools/simdstring/intrinsics.go | 89 ++++------------- tools/simdstring/intrinsics_test.go | 36 +------ 3 files changed, 110 insertions(+), 158 deletions(-) diff --git a/tools/simdstring/generate.go b/tools/simdstring/generate.go index 6a34d5309..d820181e3 100644 --- a/tools/simdstring/generate.go +++ b/tools/simdstring/generate.go @@ -81,7 +81,7 @@ func (ans *ISA) add_x86_regs() { if ans.GeneralPurposeRegisterSize == 64 { ans.add_regs(ans.GeneralPurposeRegisterSize, `R8`, `R9`, `R10`, `R11`, `R12`, `R13`, `R14`, `R15`) } - // CX is used by the assembler in some modes + // CX is used for shift and rotate instructions ans.Registers = append(ans.Registers, Register{`CX`, ans.GeneralPurposeRegisterSize, true}) // SP is the stack pointer used by the Go runtime ans.Registers = append(ans.Registers, Register{`SP`, ans.GeneralPurposeRegisterSize, true}) @@ -179,6 +179,21 @@ func (f *Function) Reg() Register { panic(fmt.Sprint("No available general purpose registers, used registers: ", strings.Join(b, ", "))) } +func (f *Function) RegForShifts() Register { + if f.ISA.Goarch == ARM64 { + return f.Reg() + } + for _, r := range f.ISA.Registers { + if r.Name == "CX" { + if f.UsedRegisters[r] { + panic("The register for shifts is already used") + } + return r + } + } + panic("No register for shifts found") +} + func (f *Function) Vec(size ...int) Register { szq := f.ISA.Bits if len(size) > 0 { @@ -389,9 +404,9 @@ func encode_cmgt16b(a, b, dest Register) (ans uint32) { return 0x271<<21 | b.ARMId()<<16 | 0xd<<10 | a.ARMId()<<5 | dest.ARMId() } -func (f *Function) CountBytesToFirstMatchDestructive(vec, ans Register) { +func (f *Function) MaskForCountDestructive(vec, ans Register) { + // vec is clobbered by this function f.Comment("Count the number of bytes to the first 0xff byte and put the result in", ans) - defer f.ReleaseReg(vec) if f.ISA.Goarch == ARM64 { // See https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon f.Comment("Go assembler doesn't support the shrn instruction, below we have: shrn.8b", vec, vec, "#4") @@ -399,12 +414,6 @@ func (f *Function) CountBytesToFirstMatchDestructive(vec, ans Register) { f.instr("WORD", fmt.Sprintf("$0x%x", shrn8b_immediate4(vec, vec))) f.instr("FMOVD", "F"+vec.Name[1:], ans) f.AddTrailingComment("Extract the lower 64 bits from", vec, "and put them into", ans) - f.instr("RBIT", ans, ans) - f.AddTrailingComment("reverse bits in", ans) - f.instr("CLZ", ans, ans) - f.AddTrailingComment(ans, "= count of leading zeros in", ans) - f.instr("UBFX", "$2", ans, "$30", ans) - f.AddTrailingComment(ans, ">>= 2 (divide by 4)") } else { if f.ISA.Bits == 128 { f.instr("PMOVMSKB", vec, ans) @@ -412,8 +421,49 @@ func (f *Function) CountBytesToFirstMatchDestructive(vec, ans Register) { f.instr("VPMOVMSKB", vec, ans) } f.AddTrailingComment(ans, "= mask of the highest bit in every byte in", vec) - f.CountTrailingZeros(ans, ans) } +} + +func (f *Function) ShiftSelfRight(self, amt any) { + op := "SHRQ" + if f.ISA.Goarch == ARM64 { + op = "LSR" + } + switch v := amt.(type) { + case Register: + if f.ISA.Goarch != ARM64 && v.Name != "CX" { + panic("On Intel only the CX register can be used for shifts") + } + f.instr(op, v, self) + default: + f.instr(op, val_repr_for_arithmetic(v), self) + } +} + +func (f *Function) ShiftMaskRightDestructive(mask, amt Register) { + // 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) + } + f.ShiftSelfRight(mask, amt) +} + +func (f *Function) CountLeadingZeroBytesInMask(src, ans Register) { + if f.ISA.Goarch == ARM64 { + f.CountTrailingZeros(src, ans) + f.instr("UBFX", "$2", ans, "$30", ans) + f.AddTrailingComment(ans, ">>= 2 (divide by 4)") + } else { + f.CountTrailingZeros(src, ans) + } +} + +func (f *Function) CountBytesToFirstMatchDestructive(vec, ans Register) { + // vec is clobbered by this function + f.Comment("Count the number of bytes to the first 0xff byte and put the result in", ans) + f.MaskForCountDestructive(vec, ans) + f.CountLeadingZeroBytesInMask(ans, ans) f.BlankLine() } @@ -1294,70 +1344,59 @@ func (isa *ISA) LEA() string { return "LEAQ" } -func (s *State) index_func(f *Function, test_bytes_impl func(bytes_to_test, test_ans Register)) { +func (s *State) index_func(f *Function, test_bytes func(bytes_to_test, test_ans Register)) { pos := f.Reg() test_ans := f.Vec() + bytes_to_test := f.Vec() data_start := f.LoadParam(`data`) limit := f.LoadParamLen(`data`) f.AddToSelf(limit, data_start) + mask := f.Reg() vecsz := f.ISA.Bits / 8 f.CopyRegister(data_start, pos) - test_bytes := func(aligned bool) { - vec := f.Vec() - defer f.ReleaseReg(vec) - if aligned { - f.LoadPointerAligned(pos, vec) - } else { - f.LoadPointerUnaligned(pos, vec) - } - test_bytes_impl(vec, test_ans) - f.JumpIfNonZero(test_ans, "byte_found") - } - - f.Comment("Load next vector of possibly unaligned bytes and check if either of the bytes are present in it") - test_bytes(false) - func() { - f.Comment("Increment address to aligned position") - unaligned_bytes := f.Reg() - defer f.ReleaseReg(unaligned_bytes) - f.CopyRegister(data_start, pos) - f.AndSelf(unaligned_bytes, vecsz-1) - f.Comment(unaligned_bytes, "now has the number of unaligned bytes") - f.AddToSelf(pos, vecsz) - f.SubtractFromSelf(pos, unaligned_bytes) - f.Comment("") + extra_bytes := f.RegForShifts() + defer f.ReleaseReg(extra_bytes) + f.CopyRegister(data_start, extra_bytes) + f.AndSelf(extra_bytes, vecsz-1) + f.SubtractFromSelf(pos, extra_bytes) + f.Comment(fmt.Sprintf("%s is now aligned to a %d byte boundary so loading from it is safe", pos, vecsz)) + f.LoadPointerAligned(pos, bytes_to_test) + test_bytes(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, extra_bytes) + f.JumpIfNonZero(mask, "byte_found") }() - f.Comment("Now loop over aligned blocks (this will repeat checking of", vecsz, "- unaligned_bytes initial bytes, but better than a branch") + f.Comment("Now loop over aligned blocks") fail := func() { f.SetReturnValue("ans", -1) f.Return() } f.Label("loop_start") + f.AddToSelf(pos, vecsz) f.JumpIfLessThan(pos, limit, "loop_body") fail() f.Label("loop_body") - test_bytes(true) - f.AddToSelf(pos, vecsz) + f.LoadPointerAligned(pos, bytes_to_test) + test_bytes(bytes_to_test, test_ans) + f.MaskForCountDestructive(test_ans, mask) + f.JumpIfNonZero(mask, "byte_found") f.JumpTo("loop_start") - func() { - f.Comment("Get the result from", test_ans, "and return it") - f.Label("byte_found") - res := f.Reg() - defer f.ReleaseReg(res) - f.CountBytesToFirstMatchDestructive(test_ans, res) - f.AddToSelf(res, pos) - f.JumpIfLessThan(res, limit, "result_in_bounds") - fail() - f.Label("result_in_bounds") - f.SubtractFromSelf(res, data_start) - f.SetReturnValue("ans", res) - f.Return() - }() + f.Comment("Get the result from", mask, "and return it") + f.Label("byte_found") + f.CountLeadingZeroBytesInMask(mask, mask) + f.AddToSelf(mask, pos) + f.JumpIfLessThan(mask, limit, "result_in_bounds") + fail() + f.Label("result_in_bounds") + f.SubtractFromSelf(mask, data_start) + f.SetReturnValue("ans", mask) + f.Return() } diff --git a/tools/simdstring/intrinsics.go b/tools/simdstring/intrinsics.go index e0ede6b0b..acf33e034 100644 --- a/tools/simdstring/intrinsics.go +++ b/tools/simdstring/intrinsics.go @@ -4,7 +4,6 @@ package simdstring import ( "runtime" - "unsafe" "golang.org/x/sys/cpu" ) @@ -14,78 +13,18 @@ var Have256bit = false var VectorSize = 1 // Return the index at which either a or b first occurs in data. If neither is -// found -1 is returned. Can read upto VectorSize-1 bytes beyond the end of data. -var UnsafeIndexByte2 func(data []byte, a, b byte) int = index_byte2_scalar +// found -1 is returned. +var IndexByte2 func(data []byte, a, b byte) int = index_byte2_scalar // Return the index at which either a or b first occurs in text. If neither is -// found -1 is returned. Can read upto VectorSize-1 bytes beyond the end of text. -var UnsafeIndexByte2String func(text string, a, b byte) int = index_byte2_string_scalar - -func get_safe_slice(addr uintptr, mask, datalen int) (align_len, vecsafe_len int) { - vecsize := mask + 1 - if datalen < vecsize { - return datalen, 0 - } - align_len = vecsize - int(addr&uintptr(mask)) - datalen -= align_len - vecsafe_len = datalen - datalen&mask - return -} - -func get_safe_slice_with_cap(addr uintptr, mask, datalen, datacap int) (align_len, vecsafe_len int) { - if datacap > datalen+mask { - return 0, datalen - } - return get_safe_slice(addr, mask, datalen) -} - -// Return the index at which either a or b first occurs in data. If neither is // found -1 is returned. -func IndexByte2(data []byte, a, b byte) int { - align_len, vecsafe_len := get_safe_slice_with_cap(uintptr(unsafe.Pointer(&data)), VectorSize-1, len(data), cap(data)) - if align_len > 0 { - if ans := index_byte2_scalar(data[:align_len], a, b); ans > -1 { - return ans - } - data = data[align_len:] - } - if vecsafe_len > 0 { - if ans := UnsafeIndexByte2(data[:vecsafe_len], a, b); ans > -1 { - return align_len + ans - } - data = data[vecsafe_len:] - } - if len(data) > 0 { - if ans := index_byte2_scalar(data, a, b); ans > -1 { - return ans + align_len + vecsafe_len - } - } - return -1 -} +var IndexByte2String func(text string, a, b byte) int = index_byte2_string_scalar -// Return the index at which either a or b first occurs in data. If neither is -// found -1 is returned. -func IndexByte2String(data string, a, b byte) int { - align_len, vecsafe_len := get_safe_slice(uintptr(unsafe.Pointer(&data)), VectorSize-1, len(data)) - if align_len > 0 { - if ans := index_byte2_string_scalar(data[:align_len], a, b); ans > -1 { - return ans - } - data = data[align_len:] - } - if vecsafe_len > 0 { - if ans := UnsafeIndexByte2String(data[:vecsafe_len], a, b); ans > -1 { - return align_len + ans - } - data = data[vecsafe_len:] - } - if len(data) > 0 { - if ans := index_byte2_string_scalar(data, a, b); ans > -1 { - return ans + align_len + vecsafe_len - } - } - return -1 -} +// Return the index at which the first C0 byte is found or -1 when no such bytes are present. +var IndexC0 func(data []byte) int = index_c0_scalar + +// Return the index at which the first C0 byte is found or -1 when no such bytes are present. +var IndexC0String func(data string) int = index_c0_string_scalar func init() { switch runtime.GOARCH { @@ -99,12 +38,16 @@ func init() { Have256bit = HasSIMD256Code } if Have256bit { - UnsafeIndexByte2 = index_byte2_asm_256 - UnsafeIndexByte2String = index_byte2_string_asm_256 + IndexByte2 = index_byte2_asm_256 + IndexByte2String = index_byte2_string_asm_256 + IndexC0 = index_c0_asm_256 + IndexC0String = index_c0_string_asm_256 VectorSize = 32 } else if Have128bit { - UnsafeIndexByte2 = index_byte2_asm_128 - UnsafeIndexByte2String = index_byte2_string_asm_128 + IndexByte2 = index_byte2_asm_128 + IndexByte2String = index_byte2_string_asm_128 + IndexC0 = index_c0_asm_128 + IndexC0String = index_c0_string_asm_128 VectorSize = 16 } } diff --git a/tools/simdstring/intrinsics_test.go b/tools/simdstring/intrinsics_test.go index 5e0eab314..b567ff9ea 100644 --- a/tools/simdstring/intrinsics_test.go +++ b/tools/simdstring/intrinsics_test.go @@ -137,14 +137,13 @@ func TestSIMDStringOps(t *testing.T) { } tests := func(h string, a, b byte) { - test([]byte(h), a, b) - for _, sz := range []int{16, 32, 64, 79} { + for _, sz := range []int{0, 16, 32, 64, 79} { q := strings.Repeat(" ", sz) + h test([]byte(q), a, b) } } - test(nil, '1', '2') - test([]byte{}, '1', '2') + test(nil, '<', '>') + test([]byte{}, '<', '>') tests("", '<', '>') tests("a", 0, 0) @@ -155,35 +154,6 @@ func TestSIMDStringOps(t *testing.T) { tests("bba", 'a', '<') tests("baa", '>', 'a') - tbs := func(addr, datalen int) { - align_len, vecsafelen := get_safe_slice(uintptr(addr), 15, datalen) - if vecsafelen < 0 || align_len+vecsafelen > datalen || datalen-vecsafelen-align_len > 15 || align_len < 0 { - t.Fatalf("Invalid bounds for addr=%d datalen=%d (align_len=%d vecsafelen=%d)", addr, datalen, align_len, vecsafelen) - } - if vecsafelen > 0 { - pos := addr + align_len - if pos&15 != 0 { - t.Fatalf("Non-aligned vector read for addr=%d datalen=%d (align_len=%d vecsafelen=%d)", addr, datalen, align_len, vecsafelen) - } - limit := pos + vecsafelen - read := func() { - if pos+16 > addr+datalen { - t.Fatalf("Read past limit for addr=%d datalen=%d (align_len=%d vecsafelen=%d pos=%d)", addr, datalen, align_len, vecsafelen, pos) - } - } - read() - for ; pos < limit; pos += 16 { - read() - } - } - } - - for datalen := 0; datalen < 33; datalen++ { - for addr := 0; addr < 32; addr++ { - tbs(addr, datalen) - } - } - c0test := func(haystack []byte) { var actual int safe_haystack := append(bytes.Repeat([]byte{'<'}, 64), haystack...)