From 18ddebe1a1a9bde349441631365f0472e9693520 Mon Sep 17 00:00:00 2001
From: DianQK <dianqk@dianqk.net>
Date: Sun, 16 Apr 2023 22:35:27 +0800
Subject: Remove UB-implying metadata when promoting speculative instruction.

After D138238 introduced the then/else blocks, we should remove UB-implying metadata for the promoted speculative instruction.

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D148456

(cherry picked from commit 2832d7941f4207f1fcf813b27cf08cecc3086959)
---
 llvm/lib/Transforms/Scalar/SROA.cpp      |  6 ++++--
 llvm/test/Transforms/SROA/select-load.ll | 24 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 8339981e1bdc..818d64725893 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -1581,15 +1581,17 @@ static void rewriteMemOpOfSelect(SelectInst &SI, T &I,
     bool IsThen = SuccBB == HeadBI->getSuccessor(0);
     int SuccIdx = IsThen ? 0 : 1;
     auto *NewMemOpBB = SuccBB == Tail ? Head : SuccBB;
+    auto &CondMemOp = cast<T>(*I.clone());
     if (NewMemOpBB != Head) {
       NewMemOpBB->setName(Head->getName() + (IsThen ? ".then" : ".else"));
       if (isa<LoadInst>(I))
         ++NumLoadsPredicated;
       else
         ++NumStoresPredicated;
-    } else
+    } else {
+      CondMemOp.dropUndefImplyingAttrsAndUnknownMetadata();
       ++NumLoadsSpeculated;
-    auto &CondMemOp = cast<T>(*I.clone());
+    }
     CondMemOp.insertBefore(NewMemOpBB->getTerminator());
     Value *Ptr = SI.getOperand(1 + SuccIdx);
     if (auto *PtrTy = Ptr->getType();
diff --git a/llvm/test/Transforms/SROA/select-load.ll b/llvm/test/Transforms/SROA/select-load.ll
index aac7af93d219..7df72419b0ce 100644
--- a/llvm/test/Transforms/SROA/select-load.ll
+++ b/llvm/test/Transforms/SROA/select-load.ll
@@ -409,7 +409,31 @@ entry:
   ret i32 %r
 }
 
+; When promoting speculative instruction, metadata that may trigger immediate UB should be dropped.
+define void @load_of_select_with_noundef_nonnull(ptr %buffer, i1 %b) {
+; CHECK-PRESERVE-CFG-LABEL: @load_of_select_with_noundef_nonnull(
+; CHECK-PRESERVE-CFG-NEXT:    [[UB_PTR:%.*]] = alloca ptr, align 8
+; CHECK-PRESERVE-CFG-NEXT:    [[SELECT_PTR:%.*]] = select i1 [[B:%.*]], ptr [[BUFFER:%.*]], ptr [[UB_PTR]]
+; CHECK-PRESERVE-CFG-NEXT:    [[LOAD_PTR:%.*]] = load ptr, ptr [[SELECT_PTR]], align 8, !nonnull !1, !noundef !1
+; CHECK-PRESERVE-CFG-NEXT:    ret void
+;
+; CHECK-MODIFY-CFG-LABEL: @load_of_select_with_noundef_nonnull(
+; CHECK-MODIFY-CFG-NEXT:    br i1 [[B:%.*]], label [[DOTTHEN:%.*]], label [[DOTCONT:%.*]]
+; CHECK-MODIFY-CFG:       .then:
+; CHECK-MODIFY-CFG-NEXT:    [[LOAD_PTR_THEN_VAL:%.*]] = load ptr, ptr [[BUFFER:%.*]], align 8, !nonnull !2, !noundef !2
+; CHECK-MODIFY-CFG-NEXT:    br label [[DOTCONT]]
+; CHECK-MODIFY-CFG:       .cont:
+; CHECK-MODIFY-CFG-NEXT:    [[LOAD_PTR:%.*]] = phi ptr [ [[LOAD_PTR_THEN_VAL]], [[DOTTHEN]] ], [ undef, [[TMP0:%.*]] ]
+; CHECK-MODIFY-CFG-NEXT:    ret void
+;
+  %ub_ptr = alloca ptr
+  %select_ptr = select i1 %b, ptr %buffer, ptr %ub_ptr
+  %load_ptr = load ptr, ptr %select_ptr, !nonnull !1, !noundef !1
+  ret void
+}
+
 !0  = !{!"branch_weights", i32 1,  i32 99}
+!1 = !{}
 
 ; Ensure that the branch metadata is reversed to match the reversals above.
 
-- 
cgit v1.2.1