first attempt at NEON intrinsics

Author: struct

Makefile

@@ -224,6 +224,10 @@ PROTECT_FREE_BIG_ZONES = -DPROTECT_FREE_BIG_ZONES=0
 ## incurs a small performance cost
 MASK_PTRS = -DMASK_PTRS=1
 
+## IsoAlloc uses ARM64 Neon instructions where possible. You can
+## explicitly disable that here
+DONT_USE_NEON = -DDONT_USE_NEON=1
+
 ## We start with the standard C++ specifics but giving
 ## the liberty to choose the gnu++* variants and/or
 ## higher than C++17
@@ -315,7 +319,7 @@ CFLAGS += $(COMMON_CFLAGS) $(DISABLE_CANARY) $(BUILD_ERROR_FLAGS) $(HOOKS) $(HEA
 	$(ABORT_NO_ENTROPY) $(ISO_DTOR_CLEANUP) $(RANDOMIZE_FREELIST) $(USE_SPINLOCK) $(HUGE_PAGES) $(USE_MLOCK) \
 	$(MEMORY_TAGGING) $(STRONG_SIZE_ISOLATION) $(MEMSET_SANITY) $(AUTO_CTOR_DTOR) $(SIGNAL_HANDLER) \
 	$(BIG_ZONE_META_DATA_GUARD) $(BIG_ZONE_GUARD) $(PROTECT_UNUSED_BIG_ZONE) $(MASK_PTRS) $(SANITIZE_CHUNKS) $(FUZZ_MODE) \
-	$(PERM_FREE_REALLOC) $(ARM_MTE)
+	$(PERM_FREE_REALLOC) $(ARM_MTE) $(DONT_USE_NEON)
 CXXFLAGS = $(COMMON_CFLAGS) -DCPP_SUPPORT=1 -std=$(STDCXX) $(SANITIZER_SUPPORT) $(HOOKS)
 
 EXE_CFLAGS = -fPIE

PERFORMANCE.md

@@ -67,6 +67,10 @@ This thread local cache speeds up the free hot path by quarantining chunks until
 
 Zones are linked by their `next_sz_index` member which tells the allocator where in the `_root->zones` array it can find the next zone that holds the same size chunks. This lookup table helps us find the first zone that holds a specific size in O(1) time. This is achieved by placing a zone's index value at that zones size index in the table, e.g. `zone_lookup_table[zone->size] = zone->index`, from there we just need to use the next zone's index member and walk it like a singly linked list to find other zones of that size. Zones are added to the front of the list as they are created.
 
+## ARM NEON
+
+IsoAlloc tries to use ARM Neon instructions to make loops in the hot path faster. If you want to disable any of this code at compile time and rely instead on more portable, but potentially slower code just set `DONT_USE_NEON` to 1 in the Makefile.
+
 ## Tests
 
 I've spent a good amount of time testing IsoAlloc to ensure its reasonably fast compared to glibc/ptmalloc. But it is impossible for me to model or simulate all the different states a program that uses IsoAlloc may be in. This section briefly covers the existing performance related tests for IsoAlloc and the data I have collected so far.

include/compiler.h

@@ -38,6 +38,13 @@
 #endif
 #endif
 
+#if DONT_USE_NEON == 0 && __ARM_NEON
+#include <arm_neon.h>
+#define USE_NEON 1
+#else
+#define USE_NEON 0
+#endif
+
 #if defined(__SANITIZE_ADDRESS__)
 static_assert(ENABLE_ASAN == 1, "ENABLE_ASAN should be 1 to enable asan instead");
 #endif

include/iso_alloc_internal.h

@@ -164,8 +164,13 @@ assert(sizeof(size_t) >= 64)
 #define BITS_PER_QWORD 64
 #define BITS_PER_QWORD_SHIFT 6
 
+#define BITS_PER_ODWORD 128
+#define BITS_PER_ODWORD_SHIFT 7
+
 #define CANARY_SIZE 8
 
+#define USED_BIT_VECTOR 0x5555555555555555
+
 /* All chunks are 8 byte aligned */
 #define CHUNK_ALIGNMENT 8
 

src/iso_alloc.c

@@ -678,31 +678,105 @@ INTERNAL_HIDDEN bit_slot_t get_next_free_bit_slot(iso_alloc_zone_t *zone) {
 /* Iterate through a zone bitmap a qword at
  * a time looking for empty slots */
 INTERNAL_HIDDEN bit_slot_t iso_scan_zone_free_slot(iso_alloc_zone_t *zone) {
+
+#if USE_NEON
+    const bitmap_index_t *bm = (bitmap_index_t *) zone->bitmap_start;
+    const bitmap_index_t max = (zone->max_bitmap_idx & ~(2 - 1));
+
+    for(int i = 0; i < max; i += 2) {
+        int64x2_t im = vld1q_s64(&bm[i]);
+
+        if(vgetq_lane_s64(im, 0) == 0x0) {
+            return (i << BITS_PER_QWORD_SHIFT);
+        }
+        if(vgetq_lane_s64(im, 1) == 0x0) {
+            return ((i + 1) << BITS_PER_QWORD_SHIFT);
+        }
+    }
+#elif __SIZEOF_INT128__
+    const __int128 *bm = (__int128 *) zone->bitmap_start;
+    const size_t max = (zone->max_bitmap_idx >> 1);
+    for(size_t i = 0; i < max; i++) {
+        if(bm[i] == 0x0) {
+            return (i << BITS_PER_ODWORD_SHIFT);
+        }
+    }
+#else
     const bitmap_index_t *bm = (bitmap_index_t *) zone->bitmap_start;
+    const bitmap_index_t max = zone->max_bitmap_idx;
 
     /* Iterate the entire bitmap a qword at a time */
-    for(bitmap_index_t i = 0; i < zone->max_bitmap_idx; i++) {
-        /* If the byte is 0 then there are some free
-         * slots we can use at this location */
+    for(bitmap_index_t i = 0; i < max; i++) {
         if(bm[i] == 0x0) {
             return (i << BITS_PER_QWORD_SHIFT);
         }
     }
+#endif
 
     return BAD_BIT_SLOT;
 }
 
 /* This function scans an entire bitmap bit-by-bit
  * and returns the first free bit position. In a heavily
- * used zone this function will be slow to search. */
+ * used zone this function may be slow to search. */
 INTERNAL_HIDDEN bit_slot_t iso_scan_zone_free_slot_slow(iso_alloc_zone_t *zone) {
-    const bitmap_index_t *bm = (bitmap_index_t *) zone->bitmap_start;
+    size_t max = 0;
+    bitmap_index_t *bm;
+
+#if USE_NEON
+    bm = (bitmap_index_t *) zone->bitmap_start;
+    max = (zone->max_bitmap_idx & ~(2 - 1));
+
+    for(bitmap_index_t i = 0; i < max; i += 2) {
+        int64x2_t im = vld1q_s64(&bm[i]);
+        int64_t i0 = vgetq_lane_s64(im, 0);
+
+        if(i0 != USED_BIT_VECTOR) {
+            for(int64_t j = 0; j < BITS_PER_QWORD; j += BITS_PER_CHUNK) {
+                if((GET_BIT(i0, j)) == 0) {
+                    return ((i << BITS_PER_QWORD_SHIFT) + j);
+                }
+            }
+        }
 
-    for(bitmap_index_t i = 0; i < zone->max_bitmap_idx; i++) {
+        int64_t i1 = vgetq_lane_s64(im, 1);
+
+        if(i1 != USED_BIT_VECTOR) {
+            for(int64_t j = 0; j < BITS_PER_QWORD; j += BITS_PER_CHUNK) {
+                if((GET_BIT(i1, j)) == 0) {
+                    return (((i + 1) << BITS_PER_QWORD_SHIFT) + j);
+                }
+            }
+        }
+    }
+
+    if(max == zone->max_bitmap_idx) {
+        return BAD_BIT_SLOT;
+    }
+#elif __SIZEOF_INT128__
+    const __int128 *ebm = (__int128 *) zone->bitmap_start;
+    max = (zone->max_bitmap_idx >> 1);
+
+    for(size_t i = 0; i < max; i++) {
+        __int128_t bts = ebm[i];
+
+        for(int64_t j = 0; j < BITS_PER_ODWORD; j += BITS_PER_CHUNK) {
+            if((GET_BIT(bts, j)) == 0) {
+                return ((i << BITS_PER_ODWORD_SHIFT) + j);
+            }
+        }
+    }
+#endif
+    bm = (bitmap_index_t *) zone->bitmap_start;
+
+    for(bitmap_index_t i = max; i < zone->max_bitmap_idx; i++) {
         bit_slot_t bts = bm[i];
 
+        if(bts != USED_BIT_VECTOR) {
+            continue;
+        }
+
         for(int64_t j = 0; j < BITS_PER_QWORD; j += BITS_PER_CHUNK) {
-            /* Check each bit to see if its available */
             if((GET_BIT(bts, j)) == 0) {
                 return ((i << BITS_PER_QWORD_SHIFT) + j);
             }

tests/pool_test.c

@@ -13,7 +13,6 @@ static const uint32_t array_sizes[] = {16, 32, 64, 128, 256, 512, 1024, 2048, 40
 int allocate(size_t array_size, size_t allocation_size) {
     iso_alloc_zone_handle *zone = iso_alloc_new_zone(allocation_size);
     size_t total_chunks = iso_zone_chunk_count(zone);
-    int32_t alloc_count = 0;
 
     /* We can treat private zones like pools of chunks
      * that don't need to be freed. Instead we can just
@@ -25,10 +24,8 @@ int allocate(size_t array_size, size_t allocation_size) {
 
         if(p == NULL) {
             LOG_AND_ABORT("Failed to allocate %ld bytes after %d total allocations from zone with %d total chunks",
-                          allocation_size, alloc_count, total_chunks);
+                          allocation_size, i, total_chunks);
         }
-
-        alloc_count++;
     }
 
     iso_alloc_destroy_zone(zone);