[sanitizer] Implement TwoLevelByteMap and use it for the internal allocator on 64-bit.

Summary:
Implement TwoLevelByteMap and use it for the internal allocator on 64-bit.
This reduces bss on 64-bit by ~8Mb because we don't use FlatByteMap on 64-bits any more.

Dmitry, please check my understanding of atomics.

Reviewers: dvyukov

Reviewed By: dvyukov

CC: samsonov, llvm-commits

Differential Revision: http://llvm-reviews.chandlerc.com/D2259

llvm-svn: 195637

3 files changed, 133 insertions, 6 deletions

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_allocator.h b/compiler-rt/lib/sanitizer_common/sanitizer_allocator.h
index 5ea1edbc798..a8debd971be 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_allocator.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_allocator.h
@@ -587,7 +587,69 @@ class FlatByteMap {
   u8 map_[kSize];
 };
 
-// FIXME: Also implement TwoLevelByteMap.
+// TwoLevelByteMap maps integers in range [0, kSize1*kSize2) to u8 values.
+// It is implemented as a two-dimensional array: array of kSize1 pointers
+// to kSize2-byte arrays. The secondary arrays are mmaped on demand.
+// Each value is initially zero and can be set to something else only once.
+// Setting and getting values from multiple threads is safe w/o extra locking.
+template <u64 kSize1, u64 kSize2, class MapUnmapCallback = NoOpMapUnmapCallback>
+class TwoLevelByteMap {
+ public:
+  void TestOnlyInit() {
+    internal_memset(map1_, 0, sizeof(map1_));
+    mu_.Init();
+  }
+  void TestOnlyUnmap() {
+    for (uptr i = 0; i < kSize1; i++) {
+      u8 *p = Get(i);
+      if (!p) continue;
+      MapUnmapCallback().OnUnmap(reinterpret_cast<uptr>(p), kSize2);
+      UnmapOrDie(p, kSize2);
+    }
+  }
+
+  uptr size() const { return kSize1 * kSize2; }
+  uptr size1() const { return kSize1; }
+  uptr size2() const { return kSize2; }
+
+  void set(uptr idx, u8 val) {
+    CHECK_LT(idx, kSize1 * kSize2);
+    u8 *map2 = GetOrCreate(idx / kSize2);
+    CHECK_EQ(0U, map2[idx % kSize2]);
+    map2[idx % kSize2] = val;
+  }
+
+  u8 operator[] (uptr idx) const {
+    CHECK_LT(idx, kSize1 * kSize2);
+    u8 *map2 = Get(idx / kSize2);
+    if (!map2) return 0;
+    return map2[idx % kSize2];
+  }
+
+ private:
+  u8 *Get(uptr idx) const {
+    CHECK_LT(idx, kSize1);
+    return reinterpret_cast<u8 *>(
+        atomic_load(&map1_[idx], memory_order_acquire));
+  }
+
+  u8 *GetOrCreate(uptr idx) {
+    u8 *res = Get(idx);
+    if (!res) {
+      SpinMutexLock l(&mu_);
+      if (!(res = Get(idx))) {
+        res = (u8*)MmapOrDie(kSize2, "TwoLevelByteMap");
+        MapUnmapCallback().OnMap(reinterpret_cast<uptr>(res), kSize2);
+        atomic_store(&map1_[idx], reinterpret_cast<uptr>(res),
+                     memory_order_release);
+      }
+    }
+    return res;
+  }
+
+  atomic_uintptr_t map1_[kSize1];
+  StaticSpinMutex mu_;
+};
 
 // SizeClassAllocator32 -- allocator for 32-bit address space.
 // This allocator can theoretically be used on 64-bit arch, but there it is less
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_allocator_internal.h b/compiler-rt/lib/sanitizer_common/sanitizer_allocator_internal.h
index 5b24bfdafa3..53af10a6c4d 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_allocator_internal.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_allocator_internal.h
@@ -27,21 +27,25 @@ static const uptr kInternalAllocatorSpace = 0;
 #if SANITIZER_WORDSIZE == 32
 static const u64 kInternalAllocatorSize = (1ULL << 32);
 static const uptr kInternalAllocatorRegionSizeLog = 20;
+static const uptr kInternalAllocatorNumRegions =
+    kInternalAllocatorSize >> kInternalAllocatorRegionSizeLog;
+typedef FlatByteMap<kInternalAllocatorNumRegions> ByteMap;
 #else
 static const u64 kInternalAllocatorSize = (1ULL << 47);
 static const uptr kInternalAllocatorRegionSizeLog = 24;
-#endif
-static const uptr kInternalAllocatorFlatByteMapSize =
+static const uptr kInternalAllocatorNumRegions =
     kInternalAllocatorSize >> kInternalAllocatorRegionSizeLog;
+typedef TwoLevelByteMap<(kInternalAllocatorNumRegions >> 12), 1 << 12> ByteMap;
+#endif
 typedef SizeClassAllocator32<
     kInternalAllocatorSpace, kInternalAllocatorSize, 16, InternalSizeClassMap,
-    kInternalAllocatorRegionSizeLog,
-    FlatByteMap<kInternalAllocatorFlatByteMapSize> > PrimaryInternalAllocator;
+    kInternalAllocatorRegionSizeLog, ByteMap> PrimaryInternalAllocator;
 
 typedef SizeClassAllocatorLocalCache<PrimaryInternalAllocator>
     InternalAllocatorCache;
 
-// We don't want our internal allocator to do any map/unmap operations.
+// We don't want our internal allocator to do any map/unmap operations from
+// LargeMmapAllocator.
 struct CrashOnMapUnmap {
   void OnMap(uptr p, uptr size) const {
     RAW_CHECK_MSG(0, "Unexpected mmap in InternalAllocator!");
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator_test.cc b/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
index d92a07fe4c2..ee977682bf1 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_allocator_test.cc
@@ -794,4 +794,65 @@ TEST(SanitizerCommon, SizeClassAllocator64PopulateFreeListOOM) {
 }
 #endif
 
+TEST(SanitizerCommon, TwoLevelByteMap) {
+  const u64 kSize1 = 1 << 6, kSize2 = 1 << 12;
+  const u64 n = kSize1 * kSize2;
+  TwoLevelByteMap<kSize1, kSize2> m;
+  m.TestOnlyInit();
+  for (u64 i = 0; i < n; i += 7) {
+    m.set(i, (i % 100) + 1);
+  }
+  for (u64 j = 0; j < n; j++) {
+    if (j % 7)
+      EXPECT_EQ(m[j], 0);
+    else
+      EXPECT_EQ(m[j], (j % 100) + 1);
+  }
+
+  m.TestOnlyUnmap();
+}
+
+
+typedef TwoLevelByteMap<1 << 12, 1 << 13, TestMapUnmapCallback> TestByteMap;
+
+struct TestByteMapParam {
+  TestByteMap *m;
+  size_t shard;
+  size_t num_shards;
+};
+
+void *TwoLevelByteMapUserThread(void *param) {
+  TestByteMapParam *p = (TestByteMapParam*)param;
+  for (size_t i = p->shard; i < p->m->size(); i += p->num_shards) {
+    size_t val = (i % 100) + 1;
+    p->m->set(i, val);
+    EXPECT_EQ((*p->m)[i], val);
+  }
+  return 0;
+}
+
+TEST(SanitizerCommon, ThreadedTwoLevelByteMap) {
+  TestByteMap m;
+  m.TestOnlyInit();
+  TestMapUnmapCallback::map_count = 0;
+  TestMapUnmapCallback::unmap_count = 0;
+  static const int kNumThreads = 4;
+  pthread_t t[kNumThreads];
+  TestByteMapParam p[kNumThreads];
+  for (int i = 0; i < kNumThreads; i++) {
+    p[i].m = &m;
+    p[i].shard = i;
+    p[i].num_shards = kNumThreads;
+    EXPECT_EQ(0, pthread_create(&t[i], 0, TwoLevelByteMapUserThread, &p[i]));
+  }
+  for (int i = 0; i < kNumThreads; i++) {
+    EXPECT_EQ(0, pthread_join(t[i], 0));
+  }
+  EXPECT_EQ((uptr)TestMapUnmapCallback::map_count, m.size1());
+  EXPECT_EQ((uptr)TestMapUnmapCallback::unmap_count, 0UL);
+  m.TestOnlyUnmap();
+  EXPECT_EQ((uptr)TestMapUnmapCallback::map_count, m.size1());
+  EXPECT_EQ((uptr)TestMapUnmapCallback::unmap_count, m.size1());
+}
+
 #endif  // #if TSAN_DEBUG==0