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

[Simon Pilgrim via Phabricator via llvm-commits]

RKSimon created this revision.
RKSimon added reviewers: spatel, reames, lebedev.ri.
Herald added a subscriber: hiraditya.
RKSimon requested review of this revision.
Herald added a project: LLVM.

This extends the PR50183/D105901 <https://reviews.llvm.org/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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D106352

Files:
  llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
  llvm/test/Transforms/InstCombine/select-gep.ll


Index: llvm/test/Transforms/InstCombine/select-gep.ll
===================================================================
--- llvm/test/Transforms/InstCombine/select-gep.ll
+++ llvm/test/Transforms/InstCombine/select-gep.ll
@@ -102,10 +102,10 @@
 ; 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
@@ -118,10 +118,10 @@
 ; 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
Index: llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
===================================================================
--- llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -2937,14 +2937,33 @@
   // 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 =


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D106352.360072.patch
Type: text/x-patch
Size: 3560 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210720/86a518da/attachment.bin>