[llvm] [InstCombine] Fold chained GEP with constant base into single GEP (PR #170439)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Jianjian Guan (jacquesguan)

<details>
<summary>Changes</summary>

Fixes https://github.com/llvm/llvm-project/issues/167014.

---
Full diff: https://github.com/llvm/llvm-project/pull/170439.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+27) 
- (added) llvm/test/Transforms/InstCombine/gep-fold-chained-const-select.ll (+46) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 5bc9c28bed141..fe68b1866d535 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2803,6 +2803,33 @@ Instruction *InstCombinerImpl::visitGEPOfGEP(GetElementPtrInst &GEP,
   if (Src->getResultElementType() != GEP.getSourceElementType())
     return nullptr;
 
+  // Fold chained GEP with constant base into single GEP:
+  // gep i8, (gep i8, %base, C1), (select Cond, C2, C3)
+  // -> gep i8, %base, (select Cond, C1+C2, C1+C3)
+  if (GEP.getNumIndices() == 1 && Src->getNumIndices() == 1) {
+    Value *SrcIdx = *Src->idx_begin();
+    Value *GEPIdx = *GEP.idx_begin();
+    const APInt *ConstOffset, *TrueVal, *FalseVal;
+    Value *Cond;
+    if (match(SrcIdx, m_APInt(ConstOffset)) &&
+        match(GEPIdx,
+              m_Select(m_Value(Cond), m_APInt(TrueVal), m_APInt(FalseVal)))) {
+      if (TrueVal->getBitWidth() != ConstOffset->getBitWidth() ||
+          FalseVal->getBitWidth() != ConstOffset->getBitWidth())
+        return nullptr;
+      APInt NewTrueVal = *ConstOffset + *TrueVal;
+      APInt NewFalseVal = *ConstOffset + *FalseVal;
+      Value *NewSelect = Builder.CreateSelect(Cond, Builder.getInt(NewTrueVal),
+                                              Builder.getInt(NewFalseVal));
+      GEPNoWrapFlags Flags =
+          getMergedGEPNoWrapFlags(*Src, *cast<GEPOperator>(&GEP));
+      return replaceInstUsesWith(GEP,
+                                 Builder.CreateGEP(GEP.getSourceElementType(),
+                                                   Src->getPointerOperand(),
+                                                   NewSelect, "", Flags));
+    }
+  }
+
   // Find out whether the last index in the source GEP is a sequential idx.
   bool EndsWithSequential = false;
   for (gep_type_iterator I = gep_type_begin(*Src), E = gep_type_end(*Src);
diff --git a/llvm/test/Transforms/InstCombine/gep-fold-chained-const-select.ll b/llvm/test/Transforms/InstCombine/gep-fold-chained-const-select.ll
new file mode 100644
index 0000000000000..0591986ae4509
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/gep-fold-chained-const-select.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -S -passes=instcombine | FileCheck %s
+
+define ptr @fold(i32 %arg0, ptr %arg1) {
+; CHECK-LABEL: @fold(
+; CHECK-NEXT:    [[V1:%.*]] = icmp sgt i32 [[ARG0:%.*]], 3
+; CHECK-NEXT:    [[V2:%.*]] = select i1 [[V1]], i64 63252, i64 29452
+; CHECK-NEXT:    [[V3:%.*]] = getelementptr i8, ptr [[V0:%.*]], i64 [[V2]]
+; CHECK-NEXT:    ret ptr [[V3]]
+;
+  %v0 = getelementptr inbounds nuw i8, ptr %arg1, i64 8148
+  %v1 = icmp sgt i32 %arg0, 3
+  %v2 = select i1 %v1, i64 55104, i64 21304
+  %v3 = getelementptr i8, ptr %v0, i64 %v2
+  ret ptr %v3
+}
+
+define ptr @no_fold_0(i32 %arg0, ptr %arg1) {
+; CHECK-LABEL: @no_fold_0(
+; CHECK-NEXT:    [[V0:%.*]] = getelementptr inbounds nuw i8, ptr [[ARG1:%.*]], i64 8148
+; CHECK-NEXT:    [[V1:%.*]] = icmp sgt i32 [[ARG0:%.*]], 3
+; CHECK-NEXT:    [[V2:%.*]] = select i1 [[V1]], i64 55104, i64 21304
+; CHECK-NEXT:    [[V3:%.*]] = getelementptr i16, ptr [[V0]], i64 [[V2]]
+; CHECK-NEXT:    ret ptr [[V3]]
+;
+  %v0 = getelementptr inbounds nuw i8, ptr %arg1, i64 8148
+  %v1 = icmp sgt i32 %arg0, 3
+  %v2 = select i1 %v1, i64 55104, i64 21304
+  %v3 = getelementptr i16, ptr %v0, i64 %v2
+  ret ptr %v3
+}
+
+define ptr @no_fold_1(i32 %arg0, ptr %arg1, i64 %arg2) {
+; CHECK-LABEL: @no_fold_1(
+; CHECK-NEXT:    [[V0:%.*]] = getelementptr inbounds nuw i8, ptr [[ARG1:%.*]], i64 [[ARG2:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = icmp sgt i32 [[ARG0:%.*]], 3
+; CHECK-NEXT:    [[V2:%.*]] = select i1 [[V1]], i64 55104, i64 21304
+; CHECK-NEXT:    [[V3:%.*]] = getelementptr i8, ptr [[V0]], i64 [[V2]]
+; CHECK-NEXT:    ret ptr [[V3]]
+;
+  %v0 = getelementptr inbounds nuw i8, ptr %arg1, i64 %arg2
+  %v1 = icmp sgt i32 %arg0, 3
+  %v2 = select i1 %v1, i64 55104, i64 21304
+  %v3 = getelementptr i8, ptr %v0, i64 %v2
+  ret ptr %v3
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/170439