[LICM] Reassociate GEPs to allow hoisting

[nikic]

Summary

Reassociate gep (gep ptr, idx1), idx2 to gep (gep ptr, idx2), idx1 if this would make the inner GEP loop invariant and thus hoistable.

This is intended to replace an InstCombine fold that does this (in https://github.com/llvm/llvm-project/blob/04f61fb73dc6a994ab267d431f2fdaedc67430ff/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp#L2006). The problem with the InstCombine fold is that LoopInfo is an optional dependency, so it is not performed reliably.

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2536

Can you please add a test where GEP and Src's indices have different dimensions? I'm a bit worried about type compatibility here.

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2540

Give them some name maybe?

[nikic]

Name instructions, add extra test.

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

2536

There is a test for different source element types in @different_elem_types and I've added @different_index_types with different index types. Not sure whether these are what you had in mind...

[nikic]

Add another test with different number of indices.

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2536

I meant something like:

Src = gep ptr %src_p, i32 %i1, i32 %i2

GEP = gep ptr %gep_p, i32 %j1

different_index_count will do, thanks!

llvm/lib/Transforms/Scalar/LICM.cpp

static bool pointerInvalidatedByLoop(MemorySSA *MSSA, MemoryUse *MU,
                                      bool InvariantGroup);
 static bool pointerInvalidatedByBlock(BasicBlock &BB, MemorySSA &MSSA,
                                       MemoryUse &MU);
 /// Try to simplify things like (A < INV_1 AND icmp A < INV_2) into (A <
 /// min(INV_1, INV_2)), if INV_1 and INV_2 are both loop invariants and their
 /// minimun can be computed outside of loop, and X is not a loop-invariant.
 static bool hoistMinMax(Instruction &I, Loop &L, ICFLoopSafetyInfo &SafetyInfo,
                         MemorySSAUpdater &MSSAU);
+static bool hoistGEP(Instruction &I, Loop &L, ICFLoopSafetyInfo &SafetyInfo,
+                     MemorySSAUpdater &MSSAU, AssumptionCache *AC,
+                     DominatorTree *DT);
 static Instruction *cloneInstructionInExitBlock(
     Instruction &I, BasicBlock &ExitBlock, PHINode &PN, const LoopInfo *LI,
     const LoopSafetyInfo *SafetyInfo, MemorySSAUpdater &MSSAU);
 
 static void eraseInstruction(Instruction &I, ICFLoopSafetyInfo &SafetyInfo,
                              MemorySSAUpdater &MSSAU);
 
 static void moveInstructionBefore(Instruction &I, Instruction &Dest,
for (Instruction &I : llvm::make_early_inc_range(*BB)) {
       // into (x < min(INV1, INV2)), and hoisting the invariant part of this
       // expression out of the loop.
       if (hoistMinMax(I, *CurLoop, *SafetyInfo, MSSAU)) {
         ++NumMinMaxHoisted;
         Changed = true;
         continue;
       }
 
+      // Try to hoist GEPs by reassociation.
+      if (hoistGEP(I, *CurLoop, *SafetyInfo, MSSAU, AC, DT)) {
+        ++NumHoisted;
+        Changed = true;
+        continue;
+      }
+
       // Remember possibly hoistable branches so we can actually hoist them
       // later if needed.
       if (BranchInst *BI = dyn_cast<BranchInst>(&I))
         CFH.registerPossiblyHoistableBranch(BI);
     }
   }
 
   // If we hoisted instructions to a conditional block they may not dominate
static bool hoistMinMax(Instruction &I, Loop &L, ICFLoopSafetyInfo &SafetyInfo,
   NewCond->takeName(&I);
   I.replaceAllUsesWith(NewCond);
   eraseInstruction(I, SafetyInfo, MSSAU);
   eraseInstruction(*cast<Instruction>(Cond1), SafetyInfo, MSSAU);
   eraseInstruction(*cast<Instruction>(Cond2), SafetyInfo, MSSAU);
   return true;
 }
 
+/// Reassociate gep (gep ptr, idx1), idx2 to gep (gep ptr, idx2), idx1 if
+/// this allows hoisting the inner GEP.
+static bool hoistGEP(Instruction &I, Loop &L, ICFLoopSafetyInfo &SafetyInfo,
+                     MemorySSAUpdater &MSSAU, AssumptionCache *AC,
+                     DominatorTree *DT) {
+  auto *GEP = dyn_cast<GetElementPtrInst>(&I);
+  if (!GEP)
+    return false;
+
+  auto *Src = dyn_cast<GetElementPtrInst>(GEP->getPointerOperand());
+  if (!Src || !Src->hasOneUse() || !L.contains(Src))
+    return false;
+
+  Value *SrcPtr = Src->getPointerOperand();
+  auto LoopInvariant = [&](Value *V) { return L.isLoopInvariant(V); };
+  if (!L.isLoopInvariant(SrcPtr) || !all_of(GEP->indices(), LoopInvariant))
+    return false;
+
+  assert(!all_of(Src->indices(), LoopInvariant) &&
+         "Would have been hoisted already");
+
+  // The swapped GEPs are inbounds if both original GEPs are inbounds
+  // and the sign of the offsets is the same. For simplicity, only
+  // handle both offsets being non-negative.
+  const DataLayout &DL = GEP->getModule()->getDataLayout();
+  auto NonNegative = [&](Value *V) {
+    return isKnownNonNegative(V, DL, 0, AC, GEP, DT);
+  };
+  bool IsInBounds = Src->isInBounds() && GEP->isInBounds() &&
+                    all_of(Src->indices(), NonNegative) &&
+                    all_of(GEP->indices(), NonNegative);
+
+  BasicBlock *Preheader = L.getLoopPreheader();
+  IRBuilder<> Builder(Preheader->getTerminator());
+  Value *NewSrc =
+      Builder.CreateGEP(GEP->getSourceElementType(), SrcPtr,
+                        SmallVector<Value *>(GEP->indices()), "", IsInBounds);
+  Builder.SetInsertPoint(GEP);
+  Value *NewGEP =
+      Builder.CreateGEP(Src->getSourceElementType(), NewSrc,
+                        SmallVector<Value *>(Src->indices()), "", IsInBounds);
+  GEP->replaceAllUsesWith(NewGEP);
+  eraseInstruction(*GEP, SafetyInfo, MSSAU);
+  eraseInstruction(*Src, SafetyInfo, MSSAU);
+  return true;
+}
+
 /// Little predicate that returns true if the specified basic block is in
 /// a subloop of the current one, not the current one itself.
 ///
 static bool inSubLoop(BasicBlock *BB, Loop *CurLoop, LoopInfo *LI) {
   assert(CurLoop->contains(BB) && "Only valid if BB is IN the loop");
   return LI->getLoopFor(BB) != CurLoop;
 }

llvm/test/Transforms/LICM/gep-reassociate.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -S -passes=licm < %s | FileCheck %s
 
 declare void @use(ptr)
 declare i32 @get.i32()
 declare i64 @get.i64()
 declare ptr @get.ptr()
 
 define void @only_one_inbounds(ptr %ptr, i1 %c, i32 %arg) {
 ; CHECK-LABEL: define void @only_one_inbounds
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i32 [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[ARG_EXT:%.*]] = zext i32 [[ARG]] to i64
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[PTR]], i64 [[ARG_EXT]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL:%.*]] = call i32 @get.i32()
 ; CHECK-NEXT:    [[VAL_EXT:%.*]] = zext i32 [[VAL]] to i64
-; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr inbounds i8, ptr [[PTR]], i64 [[VAL_EXT]]
-; CHECK-NEXT:    [[PTR3:%.*]] = getelementptr i8, ptr [[PTR2]], i64 [[ARG_EXT]]
-; CHECK-NEXT:    call void @use(ptr [[PTR3]])
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, ptr [[TMP0]], i64 [[VAL_EXT]]
+; CHECK-NEXT:    call void @use(ptr [[TMP1]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %arg.ext = zext i32 %arg to i64
   br label %loop
 
 loop:
   %val = call i32 @get.i32()
   %val.ext = zext i32 %val to i64
   %ptr2 = getelementptr inbounds i8, ptr %ptr, i64 %val.ext
   %ptr3 = getelementptr i8, ptr %ptr2, i64 %arg.ext
   call void @use(ptr %ptr3)
   br i1 %c, label %loop, label %exit
 
 exit:
   ret void
 }
 
 define void @both_inbounds_one_neg(ptr %ptr, i1 %c) {
 ; CHECK-LABEL: define void @both_inbounds_one_neg
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]]) {
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr [[PTR]], i64 -1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL:%.*]] = call i32 @get.i32()
 ; CHECK-NEXT:    [[VAL_EXT:%.*]] = zext i32 [[VAL]] to i64
-; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr inbounds i8, ptr [[PTR]], i64 [[VAL_EXT]]
-; CHECK-NEXT:    [[PTR3:%.*]] = getelementptr i8, ptr [[PTR2]], i64 -1
-; CHECK-NEXT:    call void @use(ptr [[PTR3]])
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, ptr [[TMP0]], i64 [[VAL_EXT]]
+; CHECK-NEXT:    call void @use(ptr [[TMP1]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:
   %val = call i32 @get.i32()
   %val.ext = zext i32 %val to i64
   %ptr2 = getelementptr inbounds i8, ptr %ptr, i64 %val.ext
   %ptr3 = getelementptr i8, ptr %ptr2, i64 -1
   call void @use(ptr %ptr3)
   br i1 %c, label %loop, label %exit
 
 exit:
   ret void
 }
 
 define void @both_inbounds_pos(ptr %ptr, i1 %c) {
 ; CHECK-LABEL: define void @both_inbounds_pos
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]]) {
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds i8, ptr [[PTR]], i64 1
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL:%.*]] = call i32 @get.i32()
 ; CHECK-NEXT:    [[VAL_EXT:%.*]] = zext i32 [[VAL]] to i64
-; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr inbounds i8, ptr [[PTR]], i64 [[VAL_EXT]]
-; CHECK-NEXT:    [[PTR3:%.*]] = getelementptr inbounds i8, ptr [[PTR2]], i64 1
-; CHECK-NEXT:    call void @use(ptr [[PTR3]])
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i8, ptr [[TMP0]], i64 [[VAL_EXT]]
+; CHECK-NEXT:    call void @use(ptr [[TMP1]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:
   %val = call i32 @get.i32()
   %val.ext = zext i32 %val to i64
   %ptr2 = getelementptr inbounds i8, ptr %ptr, i64 %val.ext
   %ptr3 = getelementptr inbounds i8, ptr %ptr2, i64 1
   call void @use(ptr %ptr3)
   br i1 %c, label %loop, label %exit
 
 exit:
   ret void
 }
 
 define void @different_elem_types(ptr %ptr, i1 %c, i64 %arg) {
 ; CHECK-LABEL: define void @different_elem_types
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i64 [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i64, ptr [[PTR]], i64 [[ARG]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL:%.*]] = call i64 @get.i64()
-; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr i32, ptr [[PTR]], i64 [[VAL]]
-; CHECK-NEXT:    [[PTR3:%.*]] = getelementptr i64, ptr [[PTR2]], i64 [[ARG]]
-; CHECK-NEXT:    call void @use(ptr [[PTR3]])
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[TMP0]], i64 [[VAL]]
+; CHECK-NEXT:    call void @use(ptr [[TMP1]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:

 exit:
   ret void
 }
 
 define void @multiple_indices(ptr %ptr, i1 %c, i64 %arg1, i64 %arg2) {
 ; CHECK-LABEL: define void @multiple_indices
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i64 [[ARG1:%.*]], i64 [[ARG2:%.*]]) {
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr [0 x i8], ptr [[PTR]], i64 [[ARG1]], i64 [[ARG2]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL1:%.*]] = call i64 @get.i64()
 ; CHECK-NEXT:    [[VAL2:%.*]] = call i64 @get.i64()
-; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr [0 x i8], ptr [[PTR]], i64 [[VAL1]], i64 [[VAL2]]
-; CHECK-NEXT:    [[PTR3:%.*]] = getelementptr [0 x i8], ptr [[PTR2]], i64 [[ARG1]], i64 [[ARG2]]
-; CHECK-NEXT:    call void @use(ptr [[PTR3]])
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [0 x i8], ptr [[TMP0]], i64 [[VAL1]], i64 [[VAL2]]
+; CHECK-NEXT:    call void @use(ptr [[TMP1]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:

 exit:
   ret void
 }
 
 define void @multiple_indices_very_invariant(ptr %ptr, i1 %c, i64 %arg1, i64 %arg2, i64 %arg3) {
 ; CHECK-LABEL: define void @multiple_indices_very_invariant
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i64 [[ARG1:%.*]], i64 [[ARG2:%.*]], i64 [[ARG3:%.*]]) {
 ; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr [0 x i8], ptr [[PTR]], i64 [[ARG1]], i64 [[ARG2]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL1:%.*]] = call i64 @get.i64()
-; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr [0 x i8], ptr [[PTR]], i64 [[ARG3]], i64 [[VAL1]]
-; CHECK-NEXT:    [[PTR3:%.*]] = getelementptr [0 x i8], ptr [[PTR2]], i64 [[ARG1]], i64 [[ARG2]]
-; CHECK-NEXT:    call void @use(ptr [[PTR3]])
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr [0 x i8], ptr [[TMP0]], i64 [[ARG3]], i64 [[VAL1]]
+; CHECK-NEXT:    call void @use(ptr [[TMP1]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop: