[NaryReassoc] reassociate GEP for CSE

[jingyue]

Summary

x = &a[i];
y = &a[i + j];

>

y = x + j;

along with some refactoring work such as extracting method
findClosestMatchingDominator.

Depends on D9786 which provides the ScalarEvolution::getGEPExpr interface.

[jingyue]

Following D9786, adjust the usage of getGEPExpr.

[broune]

lib/Transforms/Scalar/NaryReassociate.cpp

130

sss -> ss in reasssociates.

217

It's unclear to me what reassociable means here. Certainly other operations than add and GEPs can be reassociated in certain cases and I'd imagine that one can usually construct some example where that opens up further optimization opportunities. This seems more like a whitelist of instructions that nary reassociate knows how to deal with?

230–231

in the depth-first order -> in depth-first order
makes sure -> ensures that

You state more specifically that the traversal is in pre-order in another comment, and I think that that is necessary - it's not enough for the order to be depth-first. So this comment could also say pre-order.

236–237

Previously, the comment about SCEVable was helpful as it was not so clear that that was what isVectorTy was about. I think it's clear that isSCEVable checks if the type is SCEVable, so then the comment could be removed.

239

This could be moved into the previous if as

if (!SE->isSCEVable(I->getType()) || !isPotentiallyNaryReassociable(I)) {

243–244

Is this explaining why it is that we are not using the old I from above? If so, I'd say that it's not just because of tryReassociate, it's more because we called replaceAllUsesWith and then RecursivelyDeleteTriviallyDeadInstructions above. I'm not sure it requires a comment to explain that we then can't use the old I.

260

What's the motivation behind doing it this way instead of returning null here and then not needing isPotentiallyNaryReassociable?

337

I was wondering about the case where the index is unsigned so that hasNoUnsignedWrap would be relevant, but it looks like GEP always has signed indices. Is that how it works?

342

I think that one of these comments should say LHS + RHS.

453–470

"in the pre-order" -> "in pre-order" or perhaps "Because we do a pre-order traversal of the dominator tree"

[jingyue]

lib/Transforms/Scalar/NaryReassociate.cpp

130

Nice caaatch :)

217

Yes, you are right. It simply whitelists the instruction types we currently handle. Updated the comments accordingly.

230–231

Done.

239

Done.

243–244

It's tempting to put this statement before the "if tryReassociate" block. I was protecting code against that.

260

If !isPotentiallyNaryReassociable, we don't need to add I to SeenExprs. I could merge SeenExprs[SE->getSCEV(I)].push_back(I) to tryReassociate; then tryReassociate should probably be renamed to tryReassociateAndAddCandidate?

337

If the indexed is unsigned in the source code, then IR either zero-extends it (which we currently ignore) to pointer size or treats it as signed if it's already pointer size.

342

Done.

453–470

Done.

[jingyue]

Bjarke's comments

[broune]

lib/Transforms/Scalar/NaryReassociate.cpp

260

I'm fine with it as it is now. Merging it into tryReassociate seems fine too.

402

Suppose that RHS is 64 bits and pointers are 32 bits (or if RHS is >64 bits), do we then need CreateSExtOrTrunc here instead of CreateSExt?

[jingyue]

Yes, that's problematic. Fixed and added a test. Thank you!

lib/Transforms/Scalar/NaryReassociate.cpp

 // dominates
 //   Y = &arr[a + b]
 // we may rewrite Y into X + b.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/Local.h"
 using namespace llvm;
 using namespace PatternMatch;
 
 #define DEBUG_TYPE "nary-reassociate"
 
 namespace {
 class NaryReassociate : public FunctionPass {
 public:
   static char ID;
 
   NaryReassociate(): FunctionPass(ID) {
     initializeNaryReassociatePass(*PassRegistry::getPassRegistry());
   }
 
+  bool doInitialization(Module &M) override {
+    DL = &M.getDataLayout();
+    return false;
+  }
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addPreserved<DominatorTreeWrapperPass>();
     AU.addPreserved<ScalarEvolution>();
     AU.addPreserved<TargetLibraryInfoWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<ScalarEvolution>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
+    AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.setPreservesCFG();
   }
 
 private:
   // Runs only one iteration of the dominator-based algorithm. See the header
   // comments for why we need multiple iterations.
   bool doOneIteration(Function &F);
-  // Reasssociates I to a better form.
-  Instruction *tryReassociateAdd(Instruction *I);
+
+  // Reassociates I for better CSE.
+  Instruction *tryReassociate(Instruction *I);
+
+  // Reassociate GEP for better CSE.
+  Instruction *tryReassociateGEP(GetElementPtrInst *GEP);
+  // Try splitting GEP at the I-th index and see whether either part can be
+  // CSE'ed. This is a helper function for tryReassociateGEP.
+  //
+  // \p IndexedType The element type indexed by GEP's I-th index. This is
+  //                equivalent to
+  //                  GEP->getIndexedType(GEP->getPointerOperand(), 0-th index,
+  //                                      ..., i-th index).
+  GetElementPtrInst *tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
+                                              unsigned I, Type *IndexedType);
+  // Given GEP's I-th index = LHS + RHS, see whether &Base[..][LHS][..] or
+  // &Base[..][RHS][..] can be CSE'ed and rewrite GEP accordingly.
+  GetElementPtrInst *tryReassociateGEPAtIndex(GetElementPtrInst *GEP,
+                                              unsigned I, Value *LHS,
+                                              Value *RHS, Type *IndexedType);
+
+  // Reassociate Add for better CSE.
+  Instruction *tryReassociateAdd(BinaryOperator *I);
   // A helper function for tryReassociateAdd. LHS and RHS are explicitly passed.
   Instruction *tryReassociateAdd(Value *LHS, Value *RHS, Instruction *I);
   // Rewrites I to LHS + RHS if LHS is computed already.
   Instruction *tryReassociatedAdd(const SCEV *LHS, Value *RHS, Instruction *I);
 
+  // Returns the closest dominator of \c Dominatee that computes
+  // \c CandidateExpr. Returns null if not found.
+  Instruction *findClosestMatchingDominator(const SCEV *CandidateExpr,
+                                            Instruction *Dominatee);
+  // GetElementPtrInst implicitly sign-extends an index if the index is shorter
+  // than the pointer size. This function returns whether Index is shorter than
+  // GEP's pointer size, i.e., whether Index needs to be sign-extended in order
+  // to be an index of GEP.
+  bool requiresSignExtension(Value *Index, GetElementPtrInst *GEP);
+
   DominatorTree *DT;
   ScalarEvolution *SE;
   TargetLibraryInfo *TLI;
+  TargetTransformInfo *TTI;
+  const DataLayout *DL;
   // A lookup table quickly telling which instructions compute the given SCEV.
   // Note that there can be multiple instructions at different locations
   // computing to the same SCEV, so we map a SCEV to an instruction list.  For
   // example,
   //
   //   if (p1)
   //     foo(a + b);
   //   if (p2)
   //     bar(a + b);
   DenseMap<const SCEV *, SmallVector<Instruction *, 2>> SeenExprs;
 };
 } // anonymous namespace
 
 char NaryReassociate::ID = 0;
 INITIALIZE_PASS_BEGIN(NaryReassociate, "nary-reassociate", "Nary reassociation",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolution)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(NaryReassociate, "nary-reassociate", "Nary reassociation",
                     false, false)
 
 FunctionPass *llvm::createNaryReassociatePass() {
   return new NaryReassociate();
 }
 
 bool NaryReassociate::runOnFunction(Function &F) {
   if (skipOptnoneFunction(F))
     return false;
 
   DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   SE = &getAnalysis<ScalarEvolution>();
   TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
+  TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
 
   bool Changed = false, ChangedInThisIteration;
   do {
     ChangedInThisIteration = doOneIteration(F);
     Changed |= ChangedInThisIteration;
   } while (ChangedInThisIteration);
   return Changed;
 }
 
+// Whitelist the instruction types NaryReassociate handles for now.
+static bool isPotentiallyNaryReassociable(Instruction *I) {
+  switch (I->getOpcode()) {
+  case Instruction::Add:
+  case Instruction::GetElementPtr:
+    return true;
+  default:
+    return false;
+  }
+}
+
 bool NaryReassociate::doOneIteration(Function &F) {
   bool Changed = false;
   SeenExprs.clear();
-  // Traverse the dominator tree in the depth-first order. This order makes sure
-  // all bases of a candidate are in Candidates when we process it.
+  // Process the basic blocks in pre-order of the dominator tree. This order
+  // ensures that all bases of a candidate are in Candidates when we process it.
   for (auto Node = GraphTraits<DominatorTree *>::nodes_begin(DT);
        Node != GraphTraits<DominatorTree *>::nodes_end(DT); ++Node) {
     BasicBlock *BB = Node->getBlock();
     for (auto I = BB->begin(); I != BB->end(); ++I) {
-      // Skip vector types which are not SCEVable.
-      if (I->getOpcode() == Instruction::Add && !I->getType()->isVectorTy()) {
-        if (Instruction *NewI = tryReassociateAdd(I)) {
+      if (SE->isSCEVable(I->getType()) && isPotentiallyNaryReassociable(I)) {
+        if (Instruction *NewI = tryReassociate(I)) {
           Changed = true;
           SE->forgetValue(I);
           I->replaceAllUsesWith(NewI);
           RecursivelyDeleteTriviallyDeadInstructions(I, TLI);
           I = NewI;
         }
-        // We should add the rewritten instruction because tryReassociateAdd may
-        // have invalidated the original one.
+        // Add the rewritten instruction to SeenExprs; the original instruction
+        // is deleted.
         SeenExprs[SE->getSCEV(I)].push_back(I);
       }
     }
   }
   return Changed;
 }
 
-Instruction *NaryReassociate::tryReassociateAdd(Instruction *I) {
+Instruction *NaryReassociate::tryReassociate(Instruction *I) {
+  switch (I->getOpcode()) {
+  case Instruction::Add:
+    return tryReassociateAdd(cast<BinaryOperator>(I));
+  case Instruction::GetElementPtr:
+    return tryReassociateGEP(cast<GetElementPtrInst>(I));
+  default:
+    llvm_unreachable("should be filtered out by isPotentiallyNaryReassociable");
+  }
+}
+
+// FIXME: extract this method into TTI->getGEPCost.
+static bool isGEPFoldable(GetElementPtrInst *GEP,
+                          const TargetTransformInfo *TTI,
+                          const DataLayout *DL) {
+  GlobalVariable *BaseGV = nullptr;
+  int64_t BaseOffset = 0;
+  bool HasBaseReg = false;
+  int64_t Scale = 0;
+
+  if (GlobalVariable *GV = dyn_cast<GlobalVariable>(GEP->getPointerOperand()))
+    BaseGV = GV;
+  else
+    HasBaseReg = true;
+
+  gep_type_iterator GTI = gep_type_begin(GEP);
+  for (auto I = GEP->idx_begin(); I != GEP->idx_end(); ++I, ++GTI) {
+    if (isa<SequentialType>(*GTI)) {
+      int64_t ElementSize = DL->getTypeAllocSize(GTI.getIndexedType());
+      if (ConstantInt *ConstIdx = dyn_cast<ConstantInt>(*I)) {
+        BaseOffset += ConstIdx->getSExtValue() * ElementSize;
+      } else {
+        // Needs scale register.
+        if (Scale != 0) {
+          // No addressing mode takes two scale registers.
+          return false;
+        }
+        Scale = ElementSize;
+      }
+    } else {
+      StructType *STy = cast<StructType>(*GTI);
+      uint64_t Field = cast<ConstantInt>(*I)->getZExtValue();
+      BaseOffset += DL->getStructLayout(STy)->getElementOffset(Field);
+    }
+  }
+  return TTI->isLegalAddressingMode(GEP->getType()->getElementType(), BaseGV,
+                                    BaseOffset, HasBaseReg, Scale);
+}
+
+Instruction *NaryReassociate::tryReassociateGEP(GetElementPtrInst *GEP) {
+  // Not worth reassociating GEP if it is foldable.
+  if (isGEPFoldable(GEP, TTI, DL))
+    return nullptr;
+
+  gep_type_iterator GTI = gep_type_begin(*GEP);
+  for (unsigned I = 1, E = GEP->getNumOperands(); I != E; ++I) {
+    if (isa<SequentialType>(*GTI++)) {
+      if (auto *NewGEP = tryReassociateGEPAtIndex(GEP, I - 1, *GTI)) {
+        return NewGEP;
+      }
+    }
+  }
+  return nullptr;
+}
+
+bool NaryReassociate::requiresSignExtension(Value *Index,
+                                            GetElementPtrInst *GEP) {
+  unsigned PointerSizeInBits =
+      DL->getPointerSizeInBits(GEP->getType()->getPointerAddressSpace());
+  return cast<IntegerType>(Index->getType())->getBitWidth() < PointerSizeInBits;
+}
+
+GetElementPtrInst *
+NaryReassociate::tryReassociateGEPAtIndex(GetElementPtrInst *GEP, unsigned I,
+                                          Type *IndexedType) {
+  Value *IndexToSplit = GEP->getOperand(I + 1);
+  if (SExtInst *SExt = dyn_cast<SExtInst>(IndexToSplit))
+    IndexToSplit = SExt->getOperand(0);
+
+  if (AddOperator *AO = dyn_cast<AddOperator>(IndexToSplit)) {
+    // If the I-th index needs sext and the underlying add is not equipped with
+    // nsw, we cannot split the add because
+    //   sext(LHS + RHS) != sext(LHS) + sext(RHS).
+    if (requiresSignExtension(IndexToSplit, GEP) && !AO->hasNoSignedWrap())
+      return nullptr;
+    Value *LHS = AO->getOperand(0), *RHS = AO->getOperand(1);
+    // IndexToSplit = LHS + RHS.
+    if (auto *NewGEP = tryReassociateGEPAtIndex(GEP, I, LHS, RHS, IndexedType))
+      return NewGEP;
+    // Symmetrically, try IndexToSplit = RHS + LHS.
+    if (LHS != RHS) {
+      if (auto *NewGEP =
+              tryReassociateGEPAtIndex(GEP, I, RHS, LHS, IndexedType))
+        return NewGEP;
+    }
+  }
+  return nullptr;
+}
+
+GetElementPtrInst *
+NaryReassociate::tryReassociateGEPAtIndex(GetElementPtrInst *GEP, unsigned I,
+                                          Value *LHS, Value *RHS,
+                                          Type *IndexedType) {
+  // Look for GEP's closest dominator that has the same SCEV as GEP except that
+  // the I-th index is replaced with LHS.
+  SmallVector<const SCEV *, 4> IndexExprs;
+  for (auto Index = GEP->idx_begin(); Index != GEP->idx_end(); ++Index)
+    IndexExprs.push_back(SE->getSCEV(*Index));
+  // Replace the I-th index with LHS.
+  IndexExprs[I] = SE->getSCEV(LHS);
+  const SCEV *CandidateExpr = SE->getGEPExpr(
+      GEP->getSourceElementType(), SE->getSCEV(GEP->getPointerOperand()),
+      IndexExprs, GEP->isInBounds());
+
+  auto *Candidate = findClosestMatchingDominator(CandidateExpr, GEP);
+  if (Candidate == nullptr)
+    return nullptr;
+
+  PointerType *TypeOfCandidate = dyn_cast<PointerType>(Candidate->getType());
+  // Pretty rare but theoretically possible when a numeric value happens to
+  // share CandidateExpr.
+  if (TypeOfCandidate == nullptr)
+    return nullptr;
+
+  // NewGEP = (char *)Candidate + RHS * sizeof(IndexedType)
+  uint64_t IndexedSize = DL->getTypeAllocSize(IndexedType);
+  Type *ElementType = TypeOfCandidate->getElementType();
+  uint64_t ElementSize = DL->getTypeAllocSize(ElementType);
+  // Another less rare case: because I is not necessarily the last index of the
+  // GEP, the size of the type at the I-th index (IndexedSize) is not
+  // necessarily divisible by ElementSize. For example,
+  //
+  // #pragma pack(1)
+  // struct S {
+  //   int a[3];
+  //   int64 b[8];
+  // };
+  // #pragma pack()
+  //
+  // sizeof(S) = 100 is indivisible by sizeof(int64) = 8.
+  //
+  // TODO: bail out on this case for now. We could emit uglygep.
+  if (IndexedSize % ElementSize != 0)
+    return nullptr;
+
+  // NewGEP = &Candidate[RHS * (sizeof(IndexedType) / sizeof(Candidate[0])));
+  IRBuilder<> Builder(GEP);
+  Type *IntPtrTy = DL->getIntPtrType(TypeOfCandidate);
+  if (RHS->getType() != IntPtrTy)
+    RHS = Builder.CreateSExtOrTrunc(RHS, IntPtrTy);
+  if (IndexedSize != ElementSize) {
+    RHS = Builder.CreateMul(
+        RHS, ConstantInt::get(IntPtrTy, IndexedSize / ElementSize));
+  }
+  GetElementPtrInst *NewGEP =
+      cast<GetElementPtrInst>(Builder.CreateGEP(Candidate, RHS));
+  NewGEP->setIsInBounds(GEP->isInBounds());
+  NewGEP->takeName(GEP);
+  return NewGEP;
+}
+
+Instruction *NaryReassociate::tryReassociateAdd(BinaryOperator *I) {
   Value *LHS = I->getOperand(0), *RHS = I->getOperand(1);
   if (auto *NewI = tryReassociateAdd(LHS, RHS, I))
     return NewI;
   if (auto *NewI = tryReassociateAdd(RHS, LHS, I))
     return NewI;
   return nullptr;
 }
 

 
 Instruction *NaryReassociate::tryReassociatedAdd(const SCEV *LHSExpr,
                                                  Value *RHS, Instruction *I) {
   auto Pos = SeenExprs.find(LHSExpr);
   // Bail out if LHSExpr is not previously seen.
   if (Pos == SeenExprs.end())
     return nullptr;
 
-  auto &LHSCandidates = Pos->second;
   // Look for the closest dominator LHS of I that computes LHSExpr, and replace
   // I with LHS + RHS.
-  //
-  // Because we traverse the dominator tree in the pre-order, a
-  // candidate that doesn't dominate the current instruction won't dominate any
-  // future instruction either. Therefore, we pop it out of the stack. This
-  // optimization makes the algorithm O(n).
-  while (!LHSCandidates.empty()) {
-    Instruction *LHS = LHSCandidates.back();
-    if (DT->dominates(LHS, I)) {
-      Instruction *NewI = BinaryOperator::CreateAdd(LHS, RHS, "", I);
-      NewI->takeName(I);
-      return NewI;
-    }
-    LHSCandidates.pop_back();
+  auto *LHS = findClosestMatchingDominator(LHSExpr, I);
+  if (LHS == nullptr)
+    return nullptr;
+
+
+Instruction *
+NaryReassociate::findClosestMatchingDominator(const SCEV *CandidateExpr,
+                                              Instruction *Dominatee) {
+  auto Pos = SeenExprs.find(CandidateExpr);
+  if (Pos == SeenExprs.end())
+    return nullptr;
+
+  auto &Candidates = Pos->second;
+  // Because we process the basic blocks in pre-order of the dominator tree, a
+  // candidate that doesn't dominate the current instruction won't dominate any
+  // future instruction either. Therefore, we pop it out of the stack. This
+  // optimization makes the algorithm O(n).
+  while (!Candidates.empty()) {
+    Instruction *Candidate = Candidates.back();
+    if (DT->dominates(Candidate, Dominatee))
+      return Candidate;
+    Candidates.pop_back();
   }
   return nullptr;
 }

test/Transforms/NaryReassociate/NVPTX/nary-gep.ll

+; RUN: opt < %s -nary-reassociate -S | FileCheck %s
+
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-unknown-unknown"
+
+declare void @foo(float*)
+
+; foo(&a[i]);
+; foo(&a[i + j]);
+;   =>
+; t = &a[i];
+; foo(t);
+; foo(t + j);
+define void @reassociate_gep(float* %a, i64 %i, i64 %j) {
+; CHECK-LABEL: @reassociate_gep(
+  %1 = add i64 %i, %j
+  %2 = getelementptr float, float* %a, i64 %i
+; CHECK: [[t1:[^ ]+]] = getelementptr float, float* %a, i64 %i
+  call void @foo(float* %2)
+; CHECK: call void @foo(float* [[t1]])
+  %3 = getelementptr float, float* %a, i64 %1
+; CHECK: [[t2:[^ ]+]] = getelementptr float, float* [[t1]], i64 %j
+  call void @foo(float* %3)
+; CHECK: call void @foo(float* [[t2]])
+  ret void
+}
+
+; foo(&a[sext(j)]);
+; foo(&a[sext(i +nsw j)]);
+;   =>
+; t = &a[sext(j)];
+; foo(t);
+; foo(t + sext(i));
+define void @reassociate_gep_nsw(float* %a, i32 %i, i32 %j) {
+; CHECK-LABEL: @reassociate_gep_nsw(
+  %1 = add nsw i32 %i, %j
+  %idxprom.1 = sext i32 %1 to i64
+  %idxprom.j = sext i32 %j to i64
+  %2 = getelementptr float, float* %a, i64 %idxprom.j
+; CHECK: [[t1:[^ ]+]] = getelementptr float, float* %a, i64 %idxprom.j
+  call void @foo(float* %2)
+; CHECK: call void @foo(float* [[t1]])
+  %3 = getelementptr float, float* %a, i64 %idxprom.1
+; CHECK: [[sexti:[^ ]+]] = sext i32 %i to i64
+; CHECK: [[t2:[^ ]+]] = getelementptr float, float* [[t1]], i64 [[sexti]]
+  call void @foo(float* %3)
+; CHECK: call void @foo(float* [[t2]])
+  ret void
+}
+
+; Do not split the second GEP because sext(i + j) != sext(i) + sext(j).
+define void @reassociate_gep_no_nsw(float* %a, i32 %i, i32 %j) {
+; CHECK-LABEL: @reassociate_gep_no_nsw(
+  %1 = add i32 %i, %j
+  %2 = getelementptr float, float* %a, i32 %j
+; CHECK: getelementptr float, float* %a, i32 %j
+  call void @foo(float* %2)
+  %3 = getelementptr float, float* %a, i32 %1
+; CHECK: getelementptr float, float* %a, i32 %1
+  call void @foo(float* %3)
+  ret void
+}
+
+define void @reassociate_gep_128(float* %a, i128 %i, i128 %j) {
+; CHECK-LABEL: @reassociate_gep_128(
+  %1 = add i128 %i, %j
+  %2 = getelementptr float, float* %a, i128 %i
+; CHECK: [[t1:[^ ]+]] = getelementptr float, float* %a, i128 %i
+  call void @foo(float* %2)
+; CHECK: call void @foo(float* [[t1]])
+  %3 = getelementptr float, float* %a, i128 %1
+; CHECK: [[truncj:[^ ]+]] = trunc i128 %j to i64
+; CHECK: [[t2:[^ ]+]] = getelementptr float, float* [[t1]], i64 [[truncj]]
+  call void @foo(float* %3)
+; CHECK: call void @foo(float* [[t2]])
+  ret void
+}