[llvm] 738fcbe - [SROA] Preserve all GEP flags during speculation

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2024-06-14T11:48:35+02:00
New Revision: 738fcbee687a50bfa83ba30daf65bab41307211a

URL: https://github.com/llvm/llvm-project/commit/738fcbee687a50bfa83ba30daf65bab41307211a
DIFF: https://github.com/llvm/llvm-project/commit/738fcbee687a50bfa83ba30daf65bab41307211a.diff

LOG: [SROA] Preserve all GEP flags during speculation

Unlikely to matter in practice, as these GEPs are typically
promoted away.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/SROA.cpp
    llvm/test/Transforms/SROA/phi-gep.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index a3df89b355640..c2997913a0a25 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -3981,15 +3981,15 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     SmallVector<Value *> FalseOps = GetNewOps(False);
 
     IRB.SetInsertPoint(&GEPI);
-    bool IsInBounds = GEPI.isInBounds();
+    GEPNoWrapFlags NW = GEPI.getNoWrapFlags();
 
     Type *Ty = GEPI.getSourceElementType();
     Value *NTrue = IRB.CreateGEP(Ty, TrueOps[0], ArrayRef(TrueOps).drop_front(),
-                                 True->getName() + ".sroa.gep", IsInBounds);
+                                 True->getName() + ".sroa.gep", NW);
 
     Value *NFalse =
         IRB.CreateGEP(Ty, FalseOps[0], ArrayRef(FalseOps).drop_front(),
-                      False->getName() + ".sroa.gep", IsInBounds);
+                      False->getName() + ".sroa.gep", NW);
 
     Value *NSel = IRB.CreateSelect(Sel->getCondition(), NTrue, NFalse,
                                    Sel->getName() + ".sroa.sel");
@@ -4069,7 +4069,6 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     PHINode *NewPhi = IRB.CreatePHI(GEPI.getType(), Phi->getNumIncomingValues(),
                                     Phi->getName() + ".sroa.phi");
 
-    bool IsInBounds = GEPI.isInBounds();
     Type *SourceTy = GEPI.getSourceElementType();
     // We only handle arguments, constants, and static allocas here, so we can
     // insert GEPs at the end of the entry block.
@@ -4084,7 +4083,7 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
         SmallVector<Value *> NewOps = GetNewOps(Op);
         NewGEP =
             IRB.CreateGEP(SourceTy, NewOps[0], ArrayRef(NewOps).drop_front(),
-                          Phi->getName() + ".sroa.gep", IsInBounds);
+                          Phi->getName() + ".sroa.gep", GEPI.getNoWrapFlags());
       }
       NewPhi->addIncoming(NewGEP, BB);
     }

diff  --git a/llvm/test/Transforms/SROA/phi-gep.ll b/llvm/test/Transforms/SROA/phi-gep.ll
index c2dfa7578eb4e..776624c0798cf 100644
--- a/llvm/test/Transforms/SROA/phi-gep.ll
+++ b/llvm/test/Transforms/SROA/phi-gep.ll
@@ -363,6 +363,35 @@ exit:
   ret void
 }
 
+define void @test_sroa_gep_phi_select_same_block_nuw(i1 %c1, i1 %c2, ptr %ptr) {
+; CHECK-LABEL: @test_sroa_gep_phi_select_same_block_nuw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca [[PAIR:%.*]], align 8
+; CHECK-NEXT:    br label [[WHILE_BODY:%.*]]
+; CHECK:       while.body:
+; CHECK-NEXT:    [[PHI:%.*]] = phi ptr [ [[ALLOCA]], [[ENTRY:%.*]] ], [ [[SELECT:%.*]], [[WHILE_BODY]] ]
+; CHECK-NEXT:    [[SELECT]] = select i1 [[C1:%.*]], ptr [[PHI]], ptr [[PTR:%.*]]
+; CHECK-NEXT:    [[PHI_SROA_GEP:%.*]] = getelementptr nuw [[PAIR]], ptr [[PHI]], i64 1
+; CHECK-NEXT:    [[PTR_SROA_GEP:%.*]] = getelementptr nuw [[PAIR]], ptr [[PTR]], i64 1
+; CHECK-NEXT:    [[SELECT_SROA_SEL:%.*]] = select i1 [[C1]], ptr [[PHI_SROA_GEP]], ptr [[PTR_SROA_GEP]]
+; CHECK-NEXT:    br i1 [[C2:%.*]], label [[EXIT:%.*]], label [[WHILE_BODY]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %alloca = alloca %pair, align 8
+  br label %while.body
+
+while.body:
+  %phi = phi ptr [ %alloca, %entry ], [ %select, %while.body ]
+  %select = select i1 %c1, ptr %phi, ptr %ptr
+  %gep = getelementptr nuw %pair, ptr %select, i64 1
+  br i1 %c2, label %exit, label %while.body
+
+exit:
+  ret void
+}
+
 define i32 @test_sroa_gep_cast_phi_gep(i1 %cond) {
 ; CHECK-LABEL: @test_sroa_gep_cast_phi_gep(
 ; CHECK-NEXT:  entry: