From 891170e8636b312092486dee7c9117db7def8836 Mon Sep 17 00:00:00 2001
From: Nikita Popov <nikita.ppv@gmail.com>
Date: Sun, 29 Nov 2020 21:21:37 +0100
Subject: [PATCH] [DL] Optimize address space zero lookup (NFC)

Information for pointer size/alignment/etc is queried a lot, but
the binary search based implementation makes this fairly slow.

Add an explicit check for address space zero and skip the search
in that case -- we need to specially handle the zero address space
anyway, as it serves as the fallback for all address spaces that
were not explicitly defined.

I initially wanted to simply replace the binary search with a
linear search, which would handle both address space zero and the
general case efficiently, but I was not sure whether there are
any degenerate targets that use more than a handful of declared
address spaces (in-tree, even AMDGPU only declares six).
---
 llvm/include/llvm/IR/DataLayout.h |  7 +----
 llvm/lib/IR/DataLayout.cpp        | 52 +++++++++++++------------------
 2 files changed, 22 insertions(+), 37 deletions(-)

diff --git a/llvm/include/llvm/IR/DataLayout.h b/llvm/include/llvm/IR/DataLayout.h
index bf273038c51e..4dbca660d07b 100644
--- a/llvm/include/llvm/IR/DataLayout.h
+++ b/llvm/include/llvm/IR/DataLayout.h
@@ -161,12 +161,7 @@ private:
   using PointersTy = SmallVector<PointerAlignElem, 8>;
   PointersTy Pointers;
 
-  PointersTy::const_iterator
-  findPointerLowerBound(uint32_t AddressSpace) const {
-    return const_cast<DataLayout *>(this)->findPointerLowerBound(AddressSpace);
-  }
-
-  PointersTy::iterator findPointerLowerBound(uint32_t AddressSpace);
+  const PointerAlignElem &getPointerAlignElem(uint32_t AddressSpace) const;
 
   // The StructType -> StructLayout map.
   mutable void *LayoutMap = nullptr;
diff --git a/llvm/lib/IR/DataLayout.cpp b/llvm/lib/IR/DataLayout.cpp
index 66ae982e5120..3c9325a45395 100644
--- a/llvm/lib/IR/DataLayout.cpp
+++ b/llvm/lib/IR/DataLayout.cpp
@@ -582,12 +582,19 @@ Error DataLayout::setAlignment(AlignTypeEnum align_type, Align abi_align,
   return Error::success();
 }
 
-DataLayout::PointersTy::iterator
-DataLayout::findPointerLowerBound(uint32_t AddressSpace) {
-  return std::lower_bound(Pointers.begin(), Pointers.end(), AddressSpace,
-                          [](const PointerAlignElem &A, uint32_t AddressSpace) {
-    return A.AddressSpace < AddressSpace;
-  });
+const PointerAlignElem &
+DataLayout::getPointerAlignElem(uint32_t AddressSpace) const {
+  if (AddressSpace != 0) {
+    auto I = lower_bound(Pointers, AddressSpace,
+                         [](const PointerAlignElem &A, uint32_t AddressSpace) {
+      return A.AddressSpace < AddressSpace;
+    });
+    if (I != Pointers.end() && I->AddressSpace == AddressSpace)
+      return *I;
+  }
+
+  assert(Pointers[0].AddressSpace == 0);
+  return Pointers[0];
 }
 
 Error DataLayout::setPointerAlignment(uint32_t AddrSpace, Align ABIAlign,
@@ -597,7 +604,10 @@ Error DataLayout::setPointerAlignment(uint32_t AddrSpace, Align ABIAlign,
     return reportError(
         "Preferred alignment cannot be less than the ABI alignment");
 
-  PointersTy::iterator I = findPointerLowerBound(AddrSpace);
+  auto I = lower_bound(Pointers, AddrSpace,
+                       [](const PointerAlignElem &A, uint32_t AddressSpace) {
+    return A.AddressSpace < AddressSpace;
+  });
   if (I == Pointers.end() || I->AddressSpace != AddrSpace) {
     Pointers.insert(I, PointerAlignElem::get(AddrSpace, ABIAlign, PrefAlign,
                                              TypeByteWidth, IndexWidth));
@@ -712,30 +722,15 @@ const StructLayout *DataLayout::getStructLayout(StructType *Ty) const {
 }
 
 Align DataLayout::getPointerABIAlignment(unsigned AS) const {
-  PointersTy::const_iterator I = findPointerLowerBound(AS);
-  if (I == Pointers.end() || I->AddressSpace != AS) {
-    I = findPointerLowerBound(0);
-    assert(I->AddressSpace == 0);
-  }
-  return I->ABIAlign;
+  return getPointerAlignElem(AS).ABIAlign;
 }
 
 Align DataLayout::getPointerPrefAlignment(unsigned AS) const {
-  PointersTy::const_iterator I = findPointerLowerBound(AS);
-  if (I == Pointers.end() || I->AddressSpace != AS) {
-    I = findPointerLowerBound(0);
-    assert(I->AddressSpace == 0);
-  }
-  return I->PrefAlign;
+  return getPointerAlignElem(AS).PrefAlign;
 }
 
 unsigned DataLayout::getPointerSize(unsigned AS) const {
-  PointersTy::const_iterator I = findPointerLowerBound(AS);
-  if (I == Pointers.end() || I->AddressSpace != AS) {
-    I = findPointerLowerBound(0);
-    assert(I->AddressSpace == 0);
-  }
-  return I->TypeByteWidth;
+  return getPointerAlignElem(AS).TypeByteWidth;
 }
 
 unsigned DataLayout::getMaxPointerSize() const {
@@ -754,12 +749,7 @@ unsigned DataLayout::getPointerTypeSizeInBits(Type *Ty) const {
 }
 
 unsigned DataLayout::getIndexSize(unsigned AS) const {
-  PointersTy::const_iterator I = findPointerLowerBound(AS);
-  if (I == Pointers.end() || I->AddressSpace != AS) {
-    I = findPointerLowerBound(0);
-    assert(I->AddressSpace == 0);
-  }
-  return I->IndexWidth;
+  return getPointerAlignElem(AS).IndexWidth;
 }
 
 unsigned DataLayout::getIndexTypeSizeInBits(Type *Ty) const {