Add: Intel's Linear Address Masking

Author: ashvardanian

less_slow.cpp

@@ -2724,7 +2724,9 @@ struct arena_t {
     /// The total bytes "freed" so far
     std::size_t total_reclaimed = 0;
     /// The total number of unique allocations before a reset
-    std::size_t unique_allocations = 0;
+    std::size_t unique_allocs = 0;
+    // The maximum number of bytes allocated at once
+    std::size_t max_alloc_size = 0;
 };
 
 /**
@@ -2735,7 +2737,8 @@ inline std::byte *allocate_from_arena(arena_t &arena, std::size_t size) noexcept
     if (arena.total_allocated + size > arena_t::capacity_k) return nullptr; // Not enough space
     std::byte *ptr = arena.buffer + arena.total_allocated;
     arena.total_allocated += size;
-    arena.unique_allocations++;
+    arena.unique_allocs++;
+    arena.max_alloc_size = std::max(arena.max_alloc_size, size);
     return ptr;
 }
 
@@ -2751,7 +2754,7 @@ inline void deallocate_from_arena(arena_t &arena, std::byte *ptr, std::size_t si
     arena.total_reclaimed += size;
     // Reset completely if fully reclaimed
     if (arena.total_allocated == arena.total_reclaimed)
-        arena.total_allocated = 0, arena.total_reclaimed = 0, arena.unique_allocations = 0;
+        arena.total_allocated = 0, arena.total_reclaimed = 0, arena.unique_allocs = 0, arena.max_alloc_size = 0;
 }
 
 /**
@@ -2877,7 +2880,7 @@ yyjson_alc yyjson_wrap_arena_prepend(arena_t &arena) noexcept {
 }
 
 /**
- *  There is also an even cooler way to allocate memory! @b Pointer-tag! 🏷️
+ *  There is also an even cooler way to allocate memory! @b Pointer-tagging! 🏷️
  *  64-bit address space is a lie! Most systems only use 48 bits for addresses,
  *  some even less. So, we can use the remaining bits to store metadata about
  *  the allocated block, like its size, or the arena it came from.
@@ -2889,23 +2892,64 @@ yyjson_alc yyjson_wrap_arena_prepend(arena_t &arena) noexcept {
  *          Address sizes:          46 bits physical, 48 bits virtual
  *          Byte Order:             Little Endian
  *
- *  48-bit virtual addressing allows mapping up to @b 256-TiB of virtual space.
+ *  48-bit virtual addressing allows mapping up to @b 256-TiB of virtual space,
+ *  leaving 16 bits for metadata. On Armv8-A there is a Top Byte Ignore @b (TBI)
+ *  mode, that frees 8 bits for such metadata, but it may not be enough for our
+ *  current use-case.
+ *
+ *  There is a catch! On every OS and CPU vendor, the mechanic is different.
+ *  On Intel-based Linux systems, for example, the feature is called "Linear Address
+ *  Masking" or @b LAM for short. It can be configured in 2 modes
+ *
+ *  - LAM_U57: 57-bit linear addresses, 7 bits for metadata
+ *  - LAM_U48: 48-bit linear addresses, 16 bits for metadata
+ *
+ *  The Linux kernel itself has to be compiled with LAM support, and the feature must
+ *  also be enabled for the current running process.
+ *
+ *  @see "Support for Intel's Linear Address Masking" on Linux Weekly News:
+ *       https://lwn.net/Articles/902094/
  */
 
-constexpr std::uintptr_t pointer_tag_mask_k = 0xFFFF000000000000ull;
-
-inline void *pointer_tag(void *ptr, std::uint16_t size) noexcept {
-    std::uintptr_t addr = reinterpret_cast<std::uintptr_t>(ptr);
-    std::uintptr_t tagged = (addr & ~pointer_tag_mask_k) | (static_cast<std::uintptr_t>(size) << 48);
-    if (addr & (1ull << 47)) tagged |= pointer_tag_mask_k;
-    return reinterpret_cast<void *>(tagged);
+#if defined(__x86_64__) && defined(__linux__)
+#include <asm/prctl.h>   // `ARCH_ENABLE_TAGGED_ADDR`
+#include <sys/syscall.h> // `SYS_arch_prctl`
+static bool enable_pointer_tagging(unsigned long bits = 1) noexcept {
+    // The argument is required number of tag bits.
+    // It is rounded up to the nearest LAM mode that can provide it.
+    // For now only LAM_U57 is supported, with 6 tag bits.
+    return syscall(SYS_arch_prctl, ARCH_ENABLE_TAGGED_ADDR, bits) == 0;
 }
+#else
+static bool enable_pointer_tagging(unsigned long = 0) noexcept { return false; }
+#endif
 
+template <int start_bit_ = 48, int end_bit_ = 62>
+inline void *pointer_tag(void *ptr, std::uint16_t tag) noexcept {
+    static_assert(start_bit_ <= end_bit_);
+    // Number of bits available for the tag:
+    constexpr int bits_count = end_bit_ - start_bit_ + 1;
+    static_assert(bits_count <= 16, "We only store up to 16 bits in that range (std::uint16_t).");
+    // Convert pointer to a 64-bit integer:
+    std::uint64_t val = reinterpret_cast<std::uint64_t>(ptr);
+    // Create a mask that clears the bits in [start_bit_ .. end_bit_].
+    std::uint64_t const clear_mask = ~(((1ULL << bits_count) - 1ULL) << start_bit_);
+    val &= clear_mask;
+    // Insert our tag into those bits:
+    std::uint64_t const tag_val = (static_cast<std::uint64_t>(tag) & ((1ULL << bits_count) - 1ULL)) << start_bit_;
+    val |= tag_val;
+    return reinterpret_cast<void *>(val);
+}
+
+template <int start_bit_ = 48, int end_bit_ = 62>
 inline std::pair<void *, std::uint16_t> pointer_untag(void *ptr) noexcept {
-    std::uintptr_t tagged = reinterpret_cast<std::uintptr_t>(ptr);
-    std::uint16_t size = static_cast<std::uint16_t>(tagged >> 48);
-    std::uintptr_t addr = tagged & ~pointer_tag_mask_k;
-    return {reinterpret_cast<void *>(addr), size};
+    static_assert(start_bit_ <= end_bit_);
+    constexpr int bits_count = end_bit_ - start_bit_ + 1;
+    std::uint64_t val = reinterpret_cast<std::uint64_t>(ptr);
+    std::uint64_t extracted_tag = (val >> start_bit_) & ((1ULL << bits_count) - 1ULL);
+    std::uint64_t const clear_mask = ~(((1ULL << bits_count) - 1ULL) << start_bit_);
+    val &= clear_mask;
+    return {reinterpret_cast<void *>(val), static_cast<std::uint16_t>(extracted_tag)};
 }
 
 yyjson_alc yyjson_wrap_arena_tag(arena_t &arena) noexcept {
@@ -2925,9 +2969,11 @@ yyjson_alc yyjson_wrap_arena_tag(arena_t &arena) noexcept {
 
     alc.realloc = +[](void *ctx, void *ptr, size_t old_size_native, size_t size_native) noexcept -> void * {
         alc_size_t size = static_cast<alc_size_t>(size_native);
-        auto [real_ptr, _] = pointer_untag(ptr);
-        std::byte *new_ptr = reallocate_from_arena(*static_cast<arena_t *>(ctx), static_cast<std::byte *>(real_ptr),
-                                                   old_size_native, size_native);
+        auto [real_ptr, old_size_from_ptr] = pointer_untag(ptr);
+        assert(old_size_native == old_size_from_ptr);
+        std::byte *new_ptr = reallocate_from_arena(                           //
+            *static_cast<arena_t *>(ctx), static_cast<std::byte *>(real_ptr), //
+            old_size_from_ptr, size_native);
         if (!new_ptr) return nullptr;
         return pointer_tag(new_ptr, size);
     };
@@ -2939,7 +2985,7 @@ yyjson_alc yyjson_wrap_arena_tag(arena_t &arena) noexcept {
     return alc;
 }
 
-yyjson_alc yyjson_wrapp_malloc(arena_t &) noexcept {
+yyjson_alc yyjson_wrap_malloc(arena_t &) noexcept {
     yyjson_alc alc;
     alc.ctx = NULL;
     alc.malloc = +[](void *, size_t size) noexcept -> void * { return malloc(size); };
@@ -2950,7 +2996,10 @@ yyjson_alc yyjson_wrapp_malloc(arena_t &) noexcept {
 
 typedef yyjson_alc (*yyjson_alc_wrapper)(arena_t &);
 
-static void json_yyjson(bm::State &state, yyjson_alc_wrapper alc_wrapper = yyjson_wrapp_malloc) {
+static void json_yyjson(bm::State &state, yyjson_alc_wrapper alc_wrapper = yyjson_wrap_malloc) {
+
+    if (alc_wrapper == &yyjson_wrap_arena_tag)
+        if (!enable_pointer_tagging()) state.SkipWithError("Pointer tagging not supported");
 
     // Wrap our custom arena into a `yyjson_alc` structure, alternatively we could use:
     //
@@ -2962,8 +3011,9 @@ static void json_yyjson(bm::State &state, yyjson_alc_wrapper alc_wrapper = yyjso
 
     // Repeat the checks many times
     std::size_t bytes_processed = 0;
-    std::size_t peak_memory_usage = 0;
-    std::size_t peak_memory_calls = 0;
+    std::size_t peak_usage = 0;
+    std::size_t count_calls = 0;
+    std::size_t max_alloc = 0;
     std::size_t iteration = 0;
     for (auto _ : state) {
 
@@ -2978,18 +3028,22 @@ static void json_yyjson(bm::State &state, yyjson_alc_wrapper alc_wrapper = yyjso
             (char *)packet_json.data(), packet_json.size(), //
             YYJSON_READ_NOFLAG, &alc, &error);
         if (!error.code) bm::DoNotOptimize(contains_xss_in_yyjson(yyjson_doc_get_root(doc)));
-        peak_memory_usage = std::max(peak_memory_usage, arena.total_allocated);
-        peak_memory_calls = std::max(peak_memory_calls, arena.unique_allocations);
+        peak_usage = std::max(peak_usage, arena.total_allocated);
+        count_calls = std::max(count_calls, arena.unique_allocs);
+        max_alloc = std::max(max_alloc, arena.max_alloc_size);
         yyjson_doc_free(doc);
     }
     state.SetBytesProcessed(bytes_processed);
-    state.counters["peak_memory_usage"] = bm::Counter(peak_memory_usage, bm::Counter::kAvgThreads);
-    state.counters["mean_allocation_size"] =
-        bm::Counter(peak_memory_usage * 1.0 / peak_memory_calls, bm::Counter::kAvgThreads);
+
+    if (peak_usage) {
+        state.counters["peak_usage"] = bm::Counter(peak_usage, bm::Counter::kAvgThreads);
+        state.counters["mean_alloc"] = bm::Counter(peak_usage * 1.0 / count_calls, bm::Counter::kAvgThreads);
+        state.counters["max_alloc"] = bm::Counter(max_alloc, bm::Counter::kAvgThreads);
+    }
 }
 
-BENCHMARK_CAPTURE(json_yyjson, malloc, yyjson_wrapp_malloc)->MinTime(10)->Name("json_yyjson<malloc>");
-BENCHMARK_CAPTURE(json_yyjson, malloc, yyjson_wrapp_malloc)
+BENCHMARK_CAPTURE(json_yyjson, malloc, yyjson_wrap_malloc)->MinTime(10)->Name("json_yyjson<malloc>");
+BENCHMARK_CAPTURE(json_yyjson, malloc, yyjson_wrap_malloc)
     ->MinTime(10)
     ->Name("json_yyjson<malloc>")
     ->Threads(physical_cores());
@@ -3127,8 +3181,9 @@ enum class exception_handling_t { throw_k, noexcept_k };
 template <typename json_type_, exception_handling_t exception_handling_>
 static void json_nlohmann(bm::State &state) {
     std::size_t bytes_processed = 0;
-    std::size_t peak_memory_usage = 0;
-    std::size_t peak_memory_calls = 0;
+    std::size_t peak_usage = 0;
+    std::size_t count_calls = 0;
+    std::size_t max_alloc = 0;
     std::size_t iteration = 0;
     for (auto _ : state) {
 
@@ -3156,14 +3211,18 @@ static void json_nlohmann(bm::State &state) {
             if (!json.is_discarded()) bm::DoNotOptimize(contains_xss_nlohmann(json));
         }
         if constexpr (!std::is_same_v<json_type_, default_json>) {
-            peak_memory_usage = std::max(peak_memory_usage, thread_local_arena.total_allocated);
-            peak_memory_calls = std::max(peak_memory_calls, thread_local_arena.unique_allocations);
+            peak_usage = std::max(peak_usage, thread_local_arena.total_allocated);
+            count_calls = std::max(count_calls, thread_local_arena.unique_allocs);
+            max_alloc = std::max(max_alloc, thread_local_arena.max_alloc_size);
         }
     }
     state.SetBytesProcessed(bytes_processed);
-    state.counters["peak_memory_usage"] = bm::Counter(peak_memory_usage, bm::Counter::kAvgThreads);
-    state.counters["mean_allocation_size"] =
-        bm::Counter(peak_memory_usage * 1.0 / peak_memory_calls, bm::Counter::kAvgThreads);
+
+    if (peak_usage) {
+        state.counters["peak_usage"] = bm::Counter(peak_usage, bm::Counter::kAvgThreads);
+        state.counters["mean_alloc"] = bm::Counter(peak_usage * 1.0 / count_calls, bm::Counter::kAvgThreads);
+        state.counters["max_alloc"] = bm::Counter(max_alloc, bm::Counter::kAvgThreads);
+    }
 }
 
 BENCHMARK(json_nlohmann<default_json, exception_handling_t::throw_k>)