weaselab/conflict-set
1
Fork 1
Code Issues 5 Pull Requests Actions 1 Packages Releases 13 Activity

Move aarch64 Node16 SIMD index loads into assembly #76

Merged
andrew merged 2 commits from weaselbot/conflict-set:weaselbot/issue-68 into main 2026-08-03 13:42:51 +00:00
Conversation 0 Commits 2 Files Changed 4
+134 -49
4 changed files with 134 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+4
View File
@@ -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})
ConflictSet.cpp
+22 -48
View File
@@ -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)
simd.h
+24 -1
View File
@@ -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
#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
simd_aarch64.S
+84
View File
@@ -0,0 +1,84 @@
// SIMD operations on potentially-indeterminate Node16::index[16] bytes.
// Written in assembly because loading and operating on indeterminate values
andrew commented 2026-08-03 12:21:19 +00:00
Outdated
Review
Copy Link
Copy Source

It's not so msan doesn't track the loads, it's because it's not UB in asm

It's not so msan doesn't track the loads, it's because it's not UB in asm
// is undefined behavior in C++ ([basic.indet]) but well-defined in assembly.
// 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)
//
// These functions are only ever called directly (never indirectly), so they
// do not need BTI landing pads; the object is still marked BTI/PAC/GCS-aware
// below so a -z force-bti link keeps BTI enabled for the whole binary.
.text
// uint64_t find_eq_16(const uint8_t idx[16], uint8_t key)
// nibble i = 0xf iff idx[i] == key
andrew commented 2026-08-03 12:21:47 +00:00
Review
Copy Link
Copy Source

Do we need these? They shouldn't ever be called indirectly

Do we need these? They shouldn't ever be called indirectly
.globl find_eq_16
.type find_eq_16, %function
find_eq_16:
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:
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:
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 are only ever
// called directly, so they need no BTI landing pads; PAC/GCS compatibility
// holds trivially since they use no stack.
.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