[llvm] [InstCombine] Convert load from LUT into a select (PR #98339)

[via llvm-commits]

================
@@ -998,6 +998,105 @@ static bool canSimplifyNullLoadOrGEP(LoadInst &LI, Value *Op) {
   return false;
 }
 
+static Value *foldLoadFromIndexedGlobal(LoadInst &LI, IRBuilderBase &Builder) {
+  if (LI.isVolatile())
+    return nullptr;
+
+  auto *GEP = dyn_cast<GetElementPtrInst>(LI.getPointerOperand());
+  if (!GEP)
+    return nullptr;
+
+  auto *GV = dyn_cast<GlobalVariable>(GEP->getPointerOperand());
+  if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer())
+    return nullptr;
+
+  Constant *Init = GV->getInitializer();
+  if (!isa<ConstantArray>(Init) && !isa<ConstantDataArray>(Init))
+    return nullptr;
+
+  Type *EltTy = Init->getType()->getArrayElementType();
+  if (EltTy != LI.getType())
+    return nullptr;
+
+  uint64_t ArrayElementCount = Init->getType()->getArrayNumElements();
+  // Don't blow up on huge arrays.
+  // This threshold is chosen based on statistics on a dataset
+  // which is collected from real-world applications.
+  constexpr uint64_t MaxArraySize = 16;
+  if (ArrayElementCount > MaxArraySize)
+    return nullptr;
+
+  auto &DL = LI.getDataLayout();
+  uint64_t IndexBW = DL.getIndexTypeSizeInBits(GEP->getType());
+  APInt ConstOffset(IndexBW, 0);
+  MapVector<Value *, APInt> VariableOffsets;
+  if (!GEP->collectOffset(DL, IndexBW, VariableOffsets, ConstOffset))
+    return nullptr;
+
+  if (!ConstOffset.isZero() || VariableOffsets.size() != 1 ||
+      VariableOffsets.front().second !=
+          DL.getTypeAllocSize(EltTy).getFixedValue())
+    return nullptr;
+
+  Value *Index = VariableOffsets.front().first;
+  if (Index->getType()->getScalarSizeInBits() != IndexBW)
+    return nullptr;
+
+  SmallMapVector<Constant *, uint64_t, 2> ValueMap;
+  // MultiMapIdx indicates that this value occurs more than once in the array.
+  constexpr uint64_t MultiMapIdx = static_cast<uint64_t>(-1);
+  uint32_t MultiMapElts = 0;
+  for (uint64_t I = 0; I < ArrayElementCount; ++I) {
+    Constant *Elt = Init->getAggregateElement(I);
+
+    // bail out if the array contains undef values
+    if (isa<UndefValue>(Elt))
+      return nullptr;
+
+    if (auto *It = ValueMap.find(Elt); It != ValueMap.end()) {
+      if (It->second == MultiMapIdx)
+        continue;
+      if (++MultiMapElts == 2)
----------------
goldsteinn wrote:

IIUC you are ensuring that at most 1 element value is reference by multiple indexes so your select condition can be `x == IndexOfEltWithOnlyOneIndex ? SingleReferencedElt : MultiReferencedOther`?

If so, IMO deserves a comment and a TODO as handling things like `x > Thresh ? EltsReferencedByXAboveThresh : EltsReferenceByXBelowThresh` shouldn't be too hard to handle as well.

https://github.com/llvm/llvm-project/pull/98339