Move aarch64 Node16 SIMD index loads into assembly

The aarch64 NEON paths in getNodeIndex/getNodeIndexExists,
getChildGeq(Node16*), scan16, and checkMaxBetweenExclusiveImpl<Node16>
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

@@ -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

@@ -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
+  // The index load is done in assembly (find_eq_16) so that reading the
-  // https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon
+  // potentially-indeterminate unused index bytes is well-defined.
-
+  uint64_t bitfield = find_eq_16(self->index, index);
-  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));
   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
+  bitfield &= mask;
-  uint64_t bitfield =
-      vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(results, 4)), 0) & 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
+  // The index load is done in assembly (find_eq_16) so that reading the
-  // https://community.arm.com/arm-community-blogs/b/infrastructure-solutions-blog/posts/porting-x86-vector-bitmask-optimizations-to-arm-neon
+  // potentially-indeterminate unused index bytes is well-defined.
-
+  uint64_t bitfield = find_eq_16(self->index, index);
-  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));
   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
+  bitfield &= mask;
-  uint64_t bitfield =
-      vget_lane_u64(vreinterpret_u64_u8(vshrn_n_u16(results, 4)), 0) & 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;
+  // The index load is done in assembly (find_ge_16) so that reading the
-  memcpy(&indices, self->index, sizeof(self->index));
+  // potentially-indeterminate unused index bytes is well-defined.
-  // 0xff for each leq
+  uint64_t bitfield = find_ge_16(self->index, child);
-  auto results = vcleq_u8(vdupq_n_u8(child), indices);
   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)
+  bitfield &= mask;
-  uint64_t bitfield =
-      vget_lane_u64(
-          vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(results), 4)),
-          0) &
-      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;
+  // The index load is done in assembly (mask_in_range_16) so that reading
-  memcpy(&indices, is, 16);
+  // potentially-indeterminate unused index bytes is well-defined. `vs` slots
-  // 0xff for each in bounds
+  // beyond the in-use range are always initialized (to zero) by the allocator,
-  auto results =
+  // so the version compare below stays in C++.
-      vcltq_u8(vsubq_u8(indices, vdupq_n_u8(begin)), vdupq_n_u8(end - begin));
+  uint64_t mask = mask_in_range_16(is, begin, end);
-  // 0xf for each 0xff
-  uint64_t mask = vget_lane_u64(
-      vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(results), 4)), 0);
 
   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;
+  // The index load is done in assembly (mask_in_range_16) so that reading the
-  memcpy(&indices, self->index, 16);
+  // potentially-indeterminate unused index bytes is well-defined. The unused
-  // 0xff for each in bounds
+  // childMaxVersion slots are always initialized (to zero) by the allocator,
-  auto results =
+  // so the version compare below stays in C++.
-      vcltq_u8(vsubq_u8(indices, vdupq_n_u8(begin)), vdupq_n_u8(end - begin));
+  uint64_t mask = mask_in_range_16(self->index, begin, end);
-  // 0xf for each 0xff
-  uint64_t mask = vget_lane_u64(
-      vreinterpret_u64_u8(vshrn_n_u16(vreinterpretq_u16_u8(results), 4)), 0);
 
   mask &= self->numChildren == 16
               ? uint64_t(-1)

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);
 }
 
+#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

@@ -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