From d040ed63a4c0c0576f58681a20f3a1f86729273b Mon Sep 17 00:00:00 2001 From: Weaselbot Date: Sun, 2 Aug 2026 23:48:28 -0400 Subject: [PATCH] Move aarch64 Node16 SIMD index loads into assembly The aarch64 NEON paths in getNodeIndex/getNodeIndexExists, getChildGeq(Node16*), scan16, and checkMaxBetweenExclusiveImpl loaded the full 16-element Node16::index array (and the Node48 reverseIndex array via scan16) with NEON intrinsics, then masked the result down to [0, numChildren). Only the in-use slots are initialized; the unused bytes are indeterminate, so the wide loads were undefined behavior in C++ (per [basic.indet]) even though the trailing lanes were discarded. MSan reports this on x86-64; on aarch64 it is the same UB but MSan's imprecise modeling doesn't flag it at -O0, so there is no red->green test. Mirror the existing x86-64 fix (commit 6fed133): implement the index operations in file-level assembly, where loading and operating on indeterminate values is well-defined. Add simd_aarch64.S with find_eq_16, find_ge_16, and mask_in_range_16. AArch64 lacks pmovmskb, so (like the prior NEON code) these return a 64-bit nibble mask rather than a 16-bit bitmask; the C++ call sites keep their existing nibble-mask arithmetic and only swap the inline NEON load/compare for the assembly call. The childMaxVersion compares stay in C++ NEON intrinsics, matching x86-64's compare16: those slots are always initialized to zero by the allocator, so the wide loads are defined. The assembly functions carry `bti c` landing pads and the same aeabi_feature_and_bits attributes the compiler emits for -mbranch-protection=standard, so the object stays BTI/PAC/GCS-compatible (and warning-free under -z force-bti). CMakeLists.txt builds simd_aarch64.S into the object library and the SIMD test/bench/fuzz targets on aarch64. Closes #68 --- CMakeLists.txt | 4 +++ ConflictSet.cpp | 70 +++++++++++++---------------------------- simd.h | 25 ++++++++++++++- simd_aarch64.S | 82 +++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 132 insertions(+), 49 deletions(-) create mode 100644 simd_aarch64.S diff --git a/CMakeLists.txt b/CMakeLists.txt index 19a2ada..a265e5d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -136,6 +136,10 @@ set(CMAKE_CXX_IMPLICIT_LINK_LIBRARIES "") set(SIMD_ASM_FILES) if(CMAKE_SYSTEM_PROCESSOR STREQUAL x86_64 AND NOT USE_SIMD_FALLBACK) set(SIMD_ASM_FILES ${CMAKE_CURRENT_SOURCE_DIR}/simd_x86_64.S) +elseif((CMAKE_SYSTEM_PROCESSOR STREQUAL aarch64 OR CMAKE_SYSTEM_PROCESSOR + STREQUAL arm64) + AND NOT USE_SIMD_FALLBACK) + set(SIMD_ASM_FILES ${CMAKE_CURRENT_SOURCE_DIR}/simd_aarch64.S) endif() add_library(${PROJECT_NAME}-object OBJECT ConflictSet.cpp ${SIMD_ASM_FILES}) diff --git a/ConflictSet.cpp b/ConflictSet.cpp index 8ebab6c..5d3a84d 100644 --- a/ConflictSet.cpp +++ b/ConflictSet.cpp @@ -918,21 +918,14 @@ int getNodeIndex(Node16 *self, uint8_t index) { return -1; return std::countr_zero(bitfield); #elif defined(HAS_ARM_NEON) - // Based on - // https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon - - uint8x16_t indices; - memcpy(&indices, self->index, Node16::kMaxNodes); - // 0xff for each match - uint16x8_t results = - vreinterpretq_u16_u8(vceqq_u8(vdupq_n_u8(index), indices)); + // The index load is done in assembly (find_eq_16) so that reading the + // potentially-indeterminate unused index bytes is well-defined. + uint64_t bitfield = find_eq_16(self->index, index); assume(self->numChildren <= Node16::kMaxNodes); uint64_t mask = self->numChildren == 16 ? uint64_t(-1) : (uint64_t(1) << (self->numChildren * 4)) - 1; - // 0xf for each match in valid range - uint64_t bitfield = - vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(results, 4)), 0) & mask; + bitfield &= mask; if (bitfield == 0) return -1; return std::countr_zero(bitfield) / 4; @@ -954,21 +947,14 @@ int getNodeIndexExists(Node16 *self, uint8_t index) { assume(bitfield != 0); return std::countr_zero(bitfield); #elif defined(HAS_ARM_NEON) - // Based on - // https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon - - uint8x16_t indices; - memcpy(&indices, self->index, Node16::kMaxNodes); - // 0xff for each match - uint16x8_t results = - vreinterpretq_u16_u8(vceqq_u8(vdupq_n_u8(index), indices)); + // The index load is done in assembly (find_eq_16) so that reading the + // potentially-indeterminate unused index bytes is well-defined. + uint64_t bitfield = find_eq_16(self->index, index); assume(self->numChildren <= Node16::kMaxNodes); uint64_t mask = self->numChildren == 16 ? uint64_t(-1) : (uint64_t(1) << (self->numChildren * 4)) - 1; - // 0xf for each match in valid range - uint64_t bitfield = - vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(results, 4)), 0) & mask; + bitfield &= mask; assume(bitfield != 0); return std::countr_zero(bitfield) / 4; #else @@ -1245,20 +1231,14 @@ TaggedNodePointer getChildGeq(Node16 *self, int child) { find_ge_16(self->index, child) & ((1 << self->numChildren) - 1); return bitfield == 0 ? nullptr : self->children[std::countr_zero(bitfield)]; #elif defined(HAS_ARM_NEON) - uint8x16_t indices; - memcpy(&indices, self->index, sizeof(self->index)); - // 0xff for each leq - auto results = vcleq_u8(vdupq_n_u8(child), indices); + // The index load is done in assembly (find_ge_16) so that reading the + // potentially-indeterminate unused index bytes is well-defined. + uint64_t bitfield = find_ge_16(self->index, child); assume(self->numChildren <= Node16::kMaxNodes); uint64_t mask = self->numChildren == 16 ? uint64_t(-1) : (uint64_t(1) << (self->numChildren * 4)) - 1; - // 0xf for each 0xff (within mask) - uint64_t bitfield = - vget_lane_u64( - vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(results), 4)), - 0) & - mask; + bitfield &= mask; return bitfield == 0 ? nullptr : self->children[std::countr_zero(bitfield) / 4]; #else @@ -2068,14 +2048,11 @@ bool scan16(const InternalVersionT *vs, const uint8_t *is, int begin, int end, #ifdef HAS_ARM_NEON - uint8x16_t indices; - memcpy(&indices, is, 16); - // 0xff for each in bounds - auto results = - vcltq_u8(vsubq_u8(indices, vdupq_n_u8(begin)), vdupq_n_u8(end - begin)); - // 0xf for each 0xff - uint64_t mask = vget_lane_u64( - vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(results), 4)), 0); + // The index load is done in assembly (mask_in_range_16) so that reading + // potentially-indeterminate unused index bytes is well-defined. `vs` slots + // beyond the in-use range are always initialized (to zero) by the allocator, + // so the version compare below stays in C++. + uint64_t mask = mask_in_range_16(is, begin, end); uint32x4_t w4[4]; memcpy(w4, vs, sizeof(w4)); @@ -2256,14 +2233,11 @@ bool checkMaxBetweenExclusiveImpl(Node16 *n, int begin, int end, #ifdef HAS_ARM_NEON - uint8x16_t indices; - memcpy(&indices, self->index, 16); - // 0xff for each in bounds - auto results = - vcltq_u8(vsubq_u8(indices, vdupq_n_u8(begin)), vdupq_n_u8(end - begin)); - // 0xf for each 0xff - uint64_t mask = vget_lane_u64( - vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(results), 4)), 0); + // The index load is done in assembly (mask_in_range_16) so that reading the + // potentially-indeterminate unused index bytes is well-defined. The unused + // childMaxVersion slots are always initialized (to zero) by the allocator, + // so the version compare below stays in C++. + uint64_t mask = mask_in_range_16(self->index, begin, end); mask &= self->numChildren == 16 ? uint64_t(-1) diff --git a/simd.h b/simd.h index 2af1289..ffa4f8b 100644 --- a/simd.h +++ b/simd.h @@ -22,4 +22,27 @@ uint32_t find_ge_16(const uint8_t idx[16], uint8_t child); uint32_t mask_in_range_16(const uint8_t idx[16], uint8_t begin, uint8_t end); } -#endif \ No newline at end of file +#elif defined(__aarch64__) && !defined(USE_SIMD_FALLBACK) + +// SIMD operations on potentially-indeterminate Node16::index[16] bytes. +// Implemented in file-level assembly (simd_aarch64.S) because loading and +// operating on indeterminate values is UB in C++ but well-defined in +// assembly. The caller must mask the returned bitfield to [0, numChildren) +// before using it. +// +// AArch64 has no pmovmskb-equivalent, so (unlike x86-64) each function returns +// a 64-bit "nibble mask": nibble i (bits [4i, 4i+4)) is 0xf iff the condition +// holds at index i. Bit (4i + 3) is the high bit of byte i's result. Callers +// locate a set lane with std::countr_zero(bitfield) / 4 and mask the valid +// lanes with (uint64_t(1) << (numChildren * 4)) - 1. + +extern "C" { +// Returns nibble i = 0xf iff idx[i] == key +uint64_t find_eq_16(const uint8_t idx[16], uint8_t key); +// Returns nibble i = 0xf iff idx[i] >= child +uint64_t find_ge_16(const uint8_t idx[16], uint8_t child); +// Returns nibble i = 0xf iff begin <= idx[i] < end +uint64_t mask_in_range_16(const uint8_t idx[16], uint8_t begin, uint8_t end); +} + +#endif diff --git a/simd_aarch64.S b/simd_aarch64.S new file mode 100644 index 0000000..609ac37 --- /dev/null +++ b/simd_aarch64.S @@ -0,0 +1,82 @@ +// SIMD operations on potentially-indeterminate Node16::index[16] bytes. +// Written in assembly so msan doesn't track the loads, and so that loading +// and operating on indeterminate values is well-defined (unlike C++). The +// caller is responsible for masking the returned bitfield to [0, numChildren) +// before using it. +// +// Unlike x86-64 (which has pmovmskb), AArch64 has no single instruction that +// produces a 1-bit-per-byte mask. Each function therefore returns a 64-bit +// "nibble mask" in x0: nibble i (bits [4i, 4i+4)) is 0xf iff the condition +// holds at index i. Bit (4i + 3) is the high bit of byte i's result. Callers +// locate a set lane with countr_zero(bitfield) / 4 and mask the valid lanes +// with (uint64_t(1) << (numChildren * 4)) - 1. +// +// AArch64 AAPCS: +// x0 = const uint8_t *idx (16 bytes, may contain indeterminate data) +// w1 = uint8_t key (find_eq_16, find_ge_16) +// w1 = uint8_t begin (mask_in_range_16) +// w2 = uint8_t end (mask_in_range_16) + + .text + +// uint64_t find_eq_16(const uint8_t idx[16], uint8_t key) +// nibble i = 0xf iff idx[i] == key + .globl find_eq_16 + .type find_eq_16, %function +find_eq_16: + hint 34 // bti c + dup v1.16b, w1 // broadcast key + ldr q0, [x0] // load 16 bytes (may be indeterminate) + cmeq v0.16b, v0.16b, v1.16b // 0xff for each match + shrn v0.8b, v0.8h, 4 // pack 16 byte-flags into 8 nibble-pairs + umov x0, v0.d[0] + ret + .size find_eq_16, .-find_eq_16 + +// uint64_t find_ge_16(const uint8_t idx[16], uint8_t child) +// nibble i = 0xf iff idx[i] >= child (unsigned) +// cmhs gives unsigned ">=" (higher-or-same): Vd = Vn >= Vm. + .globl find_ge_16 + .type find_ge_16, %function +find_ge_16: + hint 34 // bti c + dup v1.16b, w1 // broadcast child + ldr q0, [x0] // load 16 bytes + cmhs v0.16b, v0.16b, v1.16b // 0xff where idx[i] >= child (unsigned) + shrn v0.8b, v0.8h, 4 + umov x0, v0.d[0] + ret + .size find_ge_16, .-find_ge_16 + +// uint64_t mask_in_range_16(const uint8_t idx[16], uint8_t begin, uint8_t end) +// nibble i = 0xf iff begin <= idx[i] < end (unsigned, wrapping arithmetic) +// Logic: (idx[i] - begin) < (end - begin), valid when end - begin < 256. +// cmhi gives unsigned ">" (higher): Vd = Vn > Vm. We want +// (end - begin) > (idx - begin), so Vn = (end - begin). + .globl mask_in_range_16 + .type mask_in_range_16, %function +mask_in_range_16: + hint 34 // bti c + dup v1.16b, w1 // broadcast begin + dup v2.16b, w2 // broadcast end + ldr q0, [x0] // load 16 bytes + sub v0.16b, v0.16b, v1.16b // idx - begin (wrapping) + sub v2.16b, v2.16b, v1.16b // end - begin (range size) + cmhi v0.16b, v2.16b, v0.16b // 0xff where (end-begin) > (idx-begin) + shrn v0.8b, v0.8h, 4 + umov x0, v0.d[0] + ret + .size mask_in_range_16, .-mask_in_range_16 + + // Declare AArch64 branch-protection compatibility, matching what the + // compiler emits for -mbranch-protection=standard (BTI + PAC + GCS). This + // keeps the object indistinguishable from C/C++ translation units for + // linkers enforcing BTI (-z force-bti); the functions above only use + // `bti c` and no stack, so PAC/GCS compatibility holds trivially. + + .aeabi_subsection aeabi_feature_and_bits, optional, ULEB128 + .aeabi_attribute Tag_Feature_BTI, 1 + .aeabi_attribute Tag_Feature_PAC, 1 + .aeabi_attribute Tag_Feature_GCS, 1 + + .section .note.GNU-stack,"",@progbits