[NFC][sanitizer] Atomix relaxed in TwoLevelMap

[vitalybuka]

Summary

This is NOOP in x86_64.
On arch64 it avoids Data Memory Barrier with visible improvements on micro benchmarks.

[dvyukov]

compiler-rt/lib/sanitizer_common/sanitizer_flat_map.h

128

The new code is incorrect. So I would add an explanatory comment that this is intentional and why we assume this will work in practice. Otherwise the next maintainer may revert this.

154

I assume this part does not affect performance, right? And the store side has much higher chances of causing real problems. So I would keep the release here.

[vitalybuka]

fixes

[vitalybuka]

update

[dvyukov]

LGTM with a nit

compiler-rt/lib/sanitizer_common/sanitizer_flat_map.h

133

The problem is with the other case: the load returns non-nullptr, but we don't have proper visibility over the returned memory.
I think we also need to say that this code is incorrect and the reason for keeping incorrect code is performance.
Something along the following lines:

This code needs to use memory_order_acquire/consume, but we use memory_order_relaxed for performance reasons (matters for arm64). We expect memory_order_relaxed to be effectively equivalent to memory_order_consume in this case for all relevant architectures: all dependent data is reachable only by dereferencing the resulting pointer.

compiler-rt/lib/sanitizer_common/sanitizer_flat_map.h

  private:
   constexpr uptr MmapSize() const {
     return RoundUpTo(kSize2 * sizeof(T), GetPageSizeCached());
   }
 
   T *Get(uptr idx) const {
     DCHECK_LT(idx, kSize1);
     return reinterpret_cast<T *>(
         atomic_load(&map1_[idx], memory_order_acquire));
   }
 
   T *GetOrCreate(uptr idx) const {
-    T *res = Get(idx);
+    DCHECK_LT(idx, kSize1);
+    // This code needs to use memory_order_acquire/consume, but we use
+    // memory_order_relaxed for performance reasons (matters for arm64). We
+    // expect memory_order_relaxed to be effectively equivalent to
+    // memory_order_consume in this case for all relevant architectures: all
+    // dependent data is reachable only by dereferencing the resulting pointer.
+    // If relaxed load fails to see stored ptr, the code will fall back to
+    // Create() and reload the value again with locked mutex as a memory
+    // barrier.
+    T *res = reinterpret_cast<T *>(atomic_load_relaxed(&map1_[idx]));
     if (LIKELY(res))
       return res;
     return Create(idx);
   }
 
   NOINLINE T *Create(uptr idx) const {
     SpinMutexLock l(&mu_);
     T *res = Get(idx);
     if (!res) {
       res = reinterpret_cast<T *>(MmapOrDie(MmapSize(), "TwoLevelMap"));
       MapUnmapCallback().OnMap(reinterpret_cast<uptr>(res), kSize2);
       atomic_store(&map1_[idx], reinterpret_cast<uptr>(res),
                    memory_order_release);
     }
     return res;
   }
 
   mutable StaticSpinMutex mu_;
   mutable atomic_uintptr_t map1_[kSize1];
 };