Merge pull request 'Move aarch64 Node16 SIMD index loads into assembly' (#76) from weaselbot/conflict-set:weaselbot/issue-68 into main

Reviewed-on: #76

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
-  // 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
@@ -2114,14 +2094,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);
 
   uint64_t compared = vget_lane_u64(
       vreinterpret_u64_u8(vshrn_n_u16(
@@ -2270,14 +2247,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

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

@@ -0,0 +1,84 @@
+// SIMD operations on potentially-indeterminate Node16::index[16] bytes.
+// Written in assembly because loading and operating on indeterminate values
+// 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
+    .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