[InstCombine] Avoid folding GEPs across loop boundaries

[clin1]

Summary

Folding a GEP from outside to inside a loop will materialize an add where there wasn't an equivalent operation before. Check the containing loops before making this fold.

[clin1]

Forgot the check lines...

[clin1]

Failing "task-dependency.c" test looks like unrelated flakiness. Observed it pass/failing randomly in local testing.

[RKSimon]

Adding some likely reviewers (I'm currently away with only limited web access).

[lebedev.ri]

Hm, i thought i commented here, but alas not. LGTM.

I suspect while we want to merge this, it's only half of the story - what if LICM didn't run yet?
We need some kind of an undo transform. (and no, SeparateConstOffsetFromGEP doesn't help, i tried)

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2137

[clin1]

Thank you for the review. Could I please ask someone to commit, with --author="Chang-Sun Lin, Jr. <[e-mail]>"?

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

if (Src->getNumOperands() == 2 && GEP.getNumOperands() == 2 &&
                   GetElementPtrInst::Create(GEPEltType, NewSrc, {SO1});
               NewGEP->setIsInBounds(GEP.isInBounds());
               return NewGEP;
             }
           }
         }
       }
 
+      // Guard the gep(gep) fold so we don't create an add inside a loop
+      // when there wasn't an equivalent instruction there before.
+      bool DifferentLoops = false;
+      if (LI)
+        if (auto *GEPLoop = LI->getLoopFor(GEP.getParent()))
+          if (auto *SrcOpI = dyn_cast<Instruction>(Src))
+            if (LI->getLoopFor(SrcOpI->getParent()) != GEPLoop)
+              DifferentLoops = true;
+
       // Fold (gep(gep(Ptr,Idx0),Idx1) -> gep(Ptr,add(Idx0,Idx1))
-      if (GO1->getType() == SO1->getType()) {
+      if (!DifferentLoops && GO1->getType() == SO1->getType()) {
         bool NewInBounds = GEP.isInBounds() && Src->isInBounds();
         auto *NewIdx =
             Builder.CreateAdd(GO1, SO1, GEP.getName() + ".idx",
                               /*HasNUW*/ false, /*HasNSW*/ NewInBounds);
         auto *NewGEP = GetElementPtrInst::Create(
             GEPEltType, Src->getPointerOperand(), {NewIdx});
         NewGEP->setIsInBounds(NewInBounds);
         return NewGEP;

llvm/test/Transforms/InstCombine/gep-combine-loop-invariant.ll

 }
 
 ; Avoid folding the GEP outside the loop to inside, and increasing loop
 ; instruction count.
 define float @gep_cross_loop(i64* %_arg_, float* %_arg_3, float %_arg_8)
 ; CHECK-LABEL: @gep_cross_loop(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i64, i64* [[_ARG_:%.*]], align 8
+; CHECK-NEXT:    [[ADD_PTR:%.*]] = getelementptr inbounds float, float* [[_ARG_3:%.*]], i64 [[TMP0]]
 ; CHECK-NEXT:    br label [[FOR_COND_I:%.*]]
 ; CHECK:       for.cond.i:
 ; CHECK-NEXT:    [[IDX:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[ADD11_I:%.*]], [[FOR_BODY_I:%.*]] ]
 ; CHECK-NEXT:    [[SUM:%.*]] = phi float [ 0.000000e+00, [[ENTRY]] ], [ [[ADD_I:%.*]], [[FOR_BODY_I]] ]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i64 [[IDX]], 17
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY_I]], label [[FOR_COND_I_I_I_PREHEADER:%.*]]
 ; CHECK:       for.cond.i.i.i.preheader:
 ; CHECK-NEXT:    ret float [[SUM]]
 ; CHECK:       for.body.i:
-; CHECK-NEXT:    [[ARRAYIDX_I84_I_IDX:%.*]] = add nsw i64 [[IDX]], [[TMP0]]
-; CHECK-NEXT:    [[ARRAYIDX_I84_I:%.*]] = getelementptr inbounds float, float* [[_ARG_3:%.*]], i64 [[ARRAYIDX_I84_I_IDX]]
+; CHECK-NEXT:    [[ARRAYIDX_I84_I:%.*]] = getelementptr inbounds float, float* [[ADD_PTR]], i64 [[IDX]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = load float, float* [[ARRAYIDX_I84_I]], align 4
 ; CHECK-NEXT:    [[ADD_I]] = fadd fast float [[SUM]], [[TMP1]]
 ; CHECK-NEXT:    [[ADD11_I]] = add nuw nsw i64 [[IDX]], 1
 ; CHECK-NEXT:    br label [[FOR_COND_I]]
 ;
 {
 entry:
   %0 = load i64, i64* %_arg_, align 8