[InstCombine] Fold (select C, (gep (gep Ptr, Idx0), Idx1), (gep Ptr, Idx0)) -> (gep Ptr,(select C, Idx0+Idx1, Idx0)) (PR51069)

[RKSimon]

Summary

This extends the PR50183/D105901 "(select C, (gep Ptr, Idx), Ptr) -> (gep Ptr, (select C, Idx, 0))" fold to account if the inner Ptr was a base gep, allowing us to merge the geps and select between the base index and the offset'd index.

[lebedev.ri]

I feel like gep (gep Ptr, Idx0), Idx1 --> gep Ptr, Idx0+Idx1 is obviously a standalone fold that should be elsewhere in instcombine.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2950–2953

[lebedev.ri]

In D106352#2890021, @lebedev.ri wrote:

I feel like gep (gep Ptr, Idx0), Idx1 --> gep Ptr, Idx0+Idx1 is obviously a standalone fold that should be elsewhere in instcombine.

... and while we clearly don't have that fold (https://godbolt.org/z/56asoz3ch),
i realize it wouldn't help here, since %gep1 has extra use in select.

llvm/test/Transforms/InstCombine/select-gep.ll

132

Please add test where %gep1 has other uses (e.g. call void @use(i32* %gep1))

[reames]

In D106352#2890033, @lebedev.ri wrote:

In D106352#2890021, @lebedev.ri wrote:

I feel like gep (gep Ptr, Idx0), Idx1 --> gep Ptr, Idx0+Idx1 is obviously a standalone fold that should be elsewhere in instcombine.

... and while we clearly don't have that fold (https://godbolt.org/z/56asoz3ch),
i realize it wouldn't help here, since %gep1 has extra use in select.

I think Roman is onto something here. And maybe I'm missing something, but I don't see the extra use he's talking about in the *optimized* IR.

Here's the current output of test2c:
; CHECK-NEXT: [[GEP1:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 [[X:%.*]]
; CHECK-NEXT: [[ICMP:%.*]] = icmp ugt i64 [[X]], [[Y:%.*]]
; CHECK-NEXT: [[SEL_IDX:%.*]] = select i1 [[ICMP]], i64 0, i64 6
; CHECK-NEXT: [[SEL:%.*]] = getelementptr i32, i32* [[GEP1]], i64 [[SEL_IDX]]
; CHECK-NEXT: ret i32* [[SEL]]

Shouldn't a (gep (gep p, idx1), idx2) -> gep p, (idx1+idx2) rule catch this case just fine? The only real downside I see is that we loose the inbounds on the first gep.

[lebedev.ri]

llvm/test/Transforms/InstCombine/select-gep.ll

111–112

@reames if we look at this test, clearly we can't fold %gep2 into %p + (%x + 6),
because that results in two instructions, but %gep1 sticks around since it's used in select,
so we can't actually do this in instcombine.

[reames]

llvm/test/Transforms/InstCombine/select-gep.ll

111–112

Roman, I think you're looking at the wrong IR. The interesting IR is the result of running the current transforms, not the input to exercise the current transform.

[RKSimon]

I'm going to take a look at folding nested geps in another parallel patch - I agree there's a lot of overlap in the transforms but hopefully we'll end up with similar canonicalizations - I'm hoping the the multi-use cases might not be a major problem.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2958

Something that I've noticed while trying to get this tested with alive2 - I think we need to guarantee that the add won't overlap?

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2958

[EDIT] Something that I've noticed while trying to get this tested with alive2 - I think we need to guarantee that the add won't overflow?

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2958

No?

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2958

Cheers

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

2958

Two more bits:

• inbounds intersect
• add is never nuw
• add is nsw iff GEP ends up being inbounds

[lebedev.ri]

Rebase this?
Did that other patch handle everything here?
What about more complex patterns with more than to indexes?

[RKSimon]

In D106352#2896253, @lebedev.ri wrote:

Rebase this?
Did that other patch handle everything here?
What about more complex patterns with more than to indexes?

Yes the other patch handled these cases - I'm going to start investigating how to perform better select-of-geps with more than 2 operands (in foldSelectOpOp), but nested-geps need to be addressed as well.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
 
   if (Instruction *I = foldSelectExtConst(SI))
     return I;
 
   // Fold (select C, (gep Ptr, Idx), Ptr) -> (gep Ptr, (select C, Idx, 0))
   // Fold (select C, Ptr, (gep Ptr, Idx)) -> (gep Ptr, (select C, 0, Idx))
   auto SelectGepWithBase = [&](GetElementPtrInst *Gep, Value *Base,
                                bool Swap) -> GetElementPtrInst * {
-    Value *Ptr = Gep->getPointerOperand();
     if (Gep->getNumOperands() != 2 || Gep->getPointerOperand() != Base ||
         !Gep->hasOneUse())
       return nullptr;
+
+    auto *BaseGep = dyn_cast<GetElementPtrInst>(Base);
     Type *ElementType = Gep->getResultElementType();
     Value *Idx = Gep->getOperand(1);
-    Value *NewT = Idx;
-    Value *NewF = Constant::getNullValue(Idx->getType());
+    Value *Ptr, *NewT, *NewF;
+
+    // Handle nested geps special case.
+    // Fold (select C, (gep (gep Ptr, Idx0), Idx1), (gep Ptr, Idx0))
+    //  --> (gep Ptr,(select C, Idx0+Idx1, Idx0))
+    // Fold (select C, (gep Ptr, Idx0), (gep (gep Ptr, Idx0), Idx1))
+    //  --> (gep Ptr,(select C, Idx0, Idx0+Idx1))
+    if (BaseGep && BaseGep->getNumOperands() == 2 && 
+        ElementType == BaseGep->getResultElementType() &&
+        Idx->getType() == BaseGep->getOperand(1)->getType()) {
+      NewT =
+          Builder.CreateAdd(Idx, BaseGep->getOperand(1), SI.getName() + ".add");
+      NewF = BaseGep->getOperand(1);
+      Ptr = BaseGep->getPointerOperand();
+    } else {
+      NewT = Idx;
+      NewF = Constant::getNullValue(Idx->getType());
+      Ptr = Gep->getPointerOperand();
+    }
+
     if (Swap)
       std::swap(NewT, NewF);
     Value *NewSI =
         Builder.CreateSelect(CondVal, NewT, NewF, SI.getName() + ".idx", &SI);
     return GetElementPtrInst::Create(ElementType, Ptr, {NewSI});
   };
   if (auto *TrueGep = dyn_cast<GetElementPtrInst>(TrueVal))
     if (auto *NewGep = SelectGepWithBase(TrueGep, FalseVal, false))

llvm/test/Transforms/InstCombine/select-gep.ll

;
   %cmp = icmp ugt i64 %x, %y
   %select = select i1 %cmp, i32* %p, i32* %gep
   ret i32* %select
 }
 
 ; PR51069
 define i32* @test2c(i32* %p, i64 %x, i64 %y) {
 ; CHECK-LABEL: @test2c(
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 [[X:%.*]]
-; CHECK-NEXT:    [[ICMP:%.*]] = icmp ugt i64 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[SEL_IDX:%.*]] = select i1 [[ICMP]], i64 0, i64 6
-; CHECK-NEXT:    [[SEL:%.*]] = getelementptr i32, i32* [[GEP1]], i64 [[SEL_IDX]]
+; CHECK-NEXT:    [[ICMP:%.*]] = icmp ugt i64 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[SEL_ADD:%.*]] = add i64 [[X]], 6
+; CHECK-NEXT:    [[SEL_IDX:%.*]] = select i1 [[ICMP]], i64 [[X]], i64 [[SEL_ADD]]
+; CHECK-NEXT:    [[SEL:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[SEL_IDX]]
 ; CHECK-NEXT:    ret i32* [[SEL]]
 ;
   %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
   %gep2 = getelementptr inbounds i32, i32* %gep1, i64 6
   %icmp = icmp ugt i64 %x, %y
   %sel = select i1 %icmp, i32* %gep1, i32* %gep2
   ret i32* %sel
 }
 
 ; PR51069
 define i32* @test2d(i32* %p, i64 %x, i64 %y) {
 ; CHECK-LABEL: @test2d(
-; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i32, i32* [[P:%.*]], i64 [[X:%.*]]
-; CHECK-NEXT:    [[ICMP:%.*]] = icmp ugt i64 [[X]], [[Y:%.*]]
-; CHECK-NEXT:    [[SEL_IDX:%.*]] = select i1 [[ICMP]], i64 6, i64 0
-; CHECK-NEXT:    [[SEL:%.*]] = getelementptr i32, i32* [[GEP1]], i64 [[SEL_IDX]]
+; CHECK-NEXT:    [[ICMP:%.*]] = icmp ugt i64 [[X:%.*]], [[Y:%.*]]
+; CHECK-NEXT:    [[SEL_ADD:%.*]] = add i64 [[X]], 6
+; CHECK-NEXT:    [[SEL_IDX:%.*]] = select i1 [[ICMP]], i64 [[SEL_ADD]], i64 [[X]]
+; CHECK-NEXT:    [[SEL:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[SEL_IDX]]
 ; CHECK-NEXT:    ret i32* [[SEL]]
 ;
   %gep1 = getelementptr inbounds i32, i32* %p, i64 %x
   %gep2 = getelementptr inbounds i32, i32* %gep1, i64 6
   %icmp = icmp ugt i64 %x, %y
   %sel = select i1 %icmp, i32* %gep2, i32* %gep1
   ret i32* %sel
 }
 
 ; Three (or more) operand GEPs are currently expected to not be optimised,
 ; though they could be in principle.
 
 define i32* @test3a([4 x i32]* %p, i64 %x, i64 %y) {
 ; CHECK-LABEL: @test3a(
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds [4 x i32], [4 x i32]* [[P:%.*]], i64 2, i64 [[X:%.*]]
 ; CHECK-NEXT:    [[GEP2:%.*]] = getelementptr inbounds [4 x i32], [4 x i32]* [[P]], i64 2, i64 [[Y:%.*]]