[InstCombine] Preserve all gep flags in gep of select fold

2 files changed, 23 insertions, 3 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index f9c1b72..d1b1ec8 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2332,10 +2332,10 @@ static Instruction *foldSelectGEP(GetElementPtrInst &GEP,
   // Propagate 'inbounds' and metadata from existing instructions.
   // Note: using IRBuilder to create the constants for efficiency.
   SmallVector<Value *, 4> IndexC(GEP.indices());
-  bool IsInBounds = GEP.isInBounds();
+  GEPNoWrapFlags NW = GEP.getNoWrapFlags();
   Type *Ty = GEP.getSourceElementType();
-  Value *NewTrueC = Builder.CreateGEP(Ty, TrueC, IndexC, "", IsInBounds);
-  Value *NewFalseC = Builder.CreateGEP(Ty, FalseC, IndexC, "", IsInBounds);
+  Value *NewTrueC = Builder.CreateGEP(Ty, TrueC, IndexC, "", NW);
+  Value *NewFalseC = Builder.CreateGEP(Ty, FalseC, IndexC, "", NW);
   return SelectInst::Create(Cond, NewTrueC, NewFalseC, "", nullptr, Sel);
 }
 
diff --git a/llvm/test/Transforms/InstCombine/getelementptr.ll b/llvm/test/Transforms/InstCombine/getelementptr.ll
index f968fa6..a842484 100644
--- a/llvm/test/Transforms/InstCombine/getelementptr.ll
+++ b/llvm/test/Transforms/InstCombine/getelementptr.ll
@@ -1760,4 +1760,24 @@ define ptr @gep_to_i8_nusw_nuw(ptr %p) {
   ret ptr %gep
 }
 
+define ptr @gep_sel_const(i1 %c) {
+; CHECK-LABEL: @gep_sel_const(
+; CHECK-NEXT:    [[GEP:%.*]] = select i1 [[C:%.*]], ptr getelementptr (i8, ptr @A, i64 5), ptr getelementptr (i8, ptr @B, i64 5)
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %sel = select i1 %c, ptr @A, ptr @B
+  %gep = getelementptr i8, ptr %sel, i64 5
+  ret ptr %gep
+}
+
+define ptr @gep_sel_const_nuw(i1 %c) {
+; CHECK-LABEL: @gep_sel_const_nuw(
+; CHECK-NEXT:    [[GEP:%.*]] = select i1 [[C:%.*]], ptr getelementptr nuw (i8, ptr @A, i64 5), ptr getelementptr nuw (i8, ptr @B, i64 5)
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+  %sel = select i1 %c, ptr @A, ptr @B
+  %gep = getelementptr nuw i8, ptr %sel, i64 5
+  ret ptr %gep
+}
+
 !0 = !{!"branch_weights", i32 2, i32 10}