[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

2535

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

2539

Give them some name maybe?

[nikic]

Name instructions, add extra test.

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

2535

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

2535

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,
                                      Loop *CurLoop, Instruction &I,
                                      SinkAndHoistLICMFlags &Flags,
                                      bool InvariantGroup);
 static bool pointerInvalidatedByBlock(BasicBlock &BB, MemorySSA &MSSA,
                                       MemoryUse &MU);
 /// Aggregates various functions for hoisting computations out of loop.
 static bool hoistArithmetics(Instruction &I, Loop &L,
                              ICFLoopSafetyInfo &SafetyInfo,
-                             MemorySSAUpdater &MSSAU);
+                             MemorySSAUpdater &MSSAU, AssumptionCache *AC,
+                             DominatorTree *DT);
 /// 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 Instruction *cloneInstructionInExitBlock(
     Instruction &I, BasicBlock &ExitBlock, PHINode &PN, const LoopInfo *LI,
     const LoopSafetyInfo *SafetyInfo, MemorySSAUpdater &MSSAU);
for (Instruction &I : llvm::make_early_inc_range(*BB)) {
           assert(DT->dominates(PN, BB) && "Conditional PHIs not expected");
           Changed = true;
           continue;
         }
       }
 
       // Try to reassociate instructions so that part of computations can be
       // done out of loop.
-      if (hoistArithmetics(I, *CurLoop, *SafetyInfo, MSSAU)) {
+      if (hoistArithmetics(I, *CurLoop, *SafetyInfo, MSSAU, AC, DT)) {
         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);
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()),
+                                    "invariant.gep", IsInBounds);
+  Builder.SetInsertPoint(GEP);
+  Value *NewGEP = Builder.CreateGEP(Src->getSourceElementType(), NewSrc,
+                                    SmallVector<Value *>(Src->indices()), "gep",
+                                    IsInBounds);
+  GEP->replaceAllUsesWith(NewGEP);
+  eraseInstruction(*GEP, SafetyInfo, MSSAU);
+  eraseInstruction(*Src, SafetyInfo, MSSAU);
+  return true;
+}
+
 static bool hoistArithmetics(Instruction &I, Loop &L,
                              ICFLoopSafetyInfo &SafetyInfo,
-                             MemorySSAUpdater &MSSAU) {
+                             MemorySSAUpdater &MSSAU,
+                             AssumptionCache *AC, DominatorTree *DT) {
   // Optimize complex patterns, such as (x < INV1 && x < INV2), turning them
   // into (x < min(INV1, INV2)), and hoisting the invariant part of this
   // expression out of the loop.
   if (hoistMinMax(I, L, SafetyInfo, MSSAU)) {
     ++NumMinMaxHoisted;
     return true;
   }
+
+  // Try to hoist GEPs by reassociation.
+  if (hoistGEP(I, L, SafetyInfo, MSSAU, AC, DT)) {
+    ++NumHoisted;
+    return true;
+  }
+
   return false;
 }
 
 /// 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/CodeGen/PowerPC/no-ctr-loop-if-exit-in-nested-loop.ll

 ; CHECK-NEXT:    addi 8, 8, 1
 ; CHECK-NEXT:    stw 12, 8(10)
 ; CHECK-NEXT:    mr 10, 11
 ; CHECK-NEXT:    bc 4, 20, .LBB0_2
 ; CHECK-NEXT:    b .LBB0_5
 ; CHECK-NEXT:    .p2align 4
 ; CHECK-NEXT:  .LBB0_4: # %if.end9
 ; CHECK-NEXT:    #
-; CHECK-NEXT:    add 9, 3, 9
-; CHECK-NEXT:    lwz 10, 4(9)
+; CHECK-NEXT:    lwzx 10, 7, 9
 ; CHECK-NEXT:    addi 10, 10, 1
-; CHECK-NEXT:    stw 10, 4(9)
+; CHECK-NEXT:    stwx 10, 7, 9
 ; CHECK-NEXT:    b .LBB0_1
 ; CHECK-NEXT:  .LBB0_5: # %if.then
 ; CHECK-NEXT:    lwax 3, 9, 3
 ; CHECK-NEXT:    blr
 entry:
   br label %block2
 
 block2:                                           ; preds = %entry

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:    [[INVARIANT_GEP:%.*]] = 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:    [[GEP:%.*]] = getelementptr i8, ptr [[INVARIANT_GEP]], i64 [[VAL_EXT]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
 ; 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:    [[INVARIANT_GEP:%.*]] = 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:    [[GEP:%.*]] = getelementptr i8, ptr [[INVARIANT_GEP]], i64 [[VAL_EXT]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
 ; 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:    [[INVARIANT_GEP:%.*]] = 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:    [[GEP:%.*]] = getelementptr inbounds i8, ptr [[INVARIANT_GEP]], i64 [[VAL_EXT]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
 ; 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:    [[INVARIANT_GEP:%.*]] = 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:    [[GEP:%.*]] = getelementptr i32, ptr [[INVARIANT_GEP]], i64 [[VAL]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop:
   %val = call i64 @get.i64()
   %ptr2 = getelementptr i32, ptr %ptr, i64 %val
   %ptr3 = getelementptr i64, ptr %ptr2, i64 %arg
   call void @use(ptr %ptr3)
   br i1 %c, label %loop, label %exit
 
 exit:
   ret void
 }
 
+define void @different_index_types(ptr %ptr, i1 %c, i32 %arg) {
+; CHECK-LABEL: define void @different_index_types
+; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i32 [[ARG:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INVARIANT_GEP:%.*]] = getelementptr i8, ptr [[PTR]], i32 [[ARG]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[VAL:%.*]] = call i64 @get.i64()
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[INVARIANT_GEP]], i64 [[VAL]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
+; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %val = call i64 @get.i64()
+  %ptr2 = getelementptr i8, ptr %ptr, i64 %val
+  %ptr3 = getelementptr i8, ptr %ptr2, i32 %arg
+  call void @use(ptr %ptr3)
+  br i1 %c, label %loop, label %exit
+
+exit:
+  ret void
+}
+
+define void @different_index_count(ptr %ptr, i1 %c, i64 %arg1, i64 %arg2) {
+; CHECK-LABEL: define void @different_index_count
+; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i64 [[ARG1:%.*]], i64 [[ARG2:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[INVARIANT_GEP:%.*]] = getelementptr [0 x i8], ptr [[PTR]], i64 [[ARG1]], i64 [[ARG2]]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[VAL:%.*]] = call i64 @get.i64()
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[INVARIANT_GEP]], i64 [[VAL]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
+; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %val = call i64 @get.i64()
+  %ptr2 = getelementptr i8, ptr %ptr, i64 %val
+  %ptr3 = getelementptr [0 x i8], ptr %ptr2, i64 %arg1, i64 %arg2
+  call void @use(ptr %ptr3)
+  br i1 %c, label %loop, label %exit
+
+exit:
+  ret void
+}
+
 define void @src_has_extra_use(ptr %ptr, i1 %c, i64 %arg) {
 ; CHECK-LABEL: define void @src_has_extra_use
 ; CHECK-SAME: (ptr [[PTR:%.*]], i1 [[C:%.*]], i64 [[ARG:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[VAL:%.*]] = call i64 @get.i64()
 ; CHECK-NEXT:    [[PTR2:%.*]] = getelementptr i8, ptr [[PTR]], i64 [[VAL]]

 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:    [[INVARIANT_GEP:%.*]] = 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:    [[GEP:%.*]] = getelementptr [0 x i8], ptr [[INVARIANT_GEP]], i64 [[VAL1]], i64 [[VAL2]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
 ; 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:    [[INVARIANT_GEP:%.*]] = 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:    [[GEP:%.*]] = getelementptr [0 x i8], ptr [[INVARIANT_GEP]], i64 [[ARG3]], i64 [[VAL1]]
+; CHECK-NEXT:    call void @use(ptr [[GEP]])
 ; CHECK-NEXT:    br i1 [[C]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %loop
 
 loop: