[JumpThreading] Copy metadata when inserting preload into preds

llvm/lib/Transforms/Scalar/JumpThreading.cpp

@@ -1362,11 +1362,14 @@ bool JumpThreadingPass::simplifyPartiallyRedundantLoad(LoadInst *LoadI) {
   // farther than to a predecessor, we need to reuse the code from GVN's PRE.
   // It requires domination tree analysis, so for this simple case it is an
   // overkill.
+  std::optional<bool> TransfersExecution = std::nullopt;
   if (PredsScanned.size() != AvailablePreds.size() &&
-      !isSafeToSpeculativelyExecute(LoadI))
-    for (auto I = LoadBB->begin(); &*I != LoadI; ++I)
-      if (!isGuaranteedToTransferExecutionToSuccessor(&*I))
-        return false;
+      !isSafeToSpeculativelyExecute(LoadI)) {
+    if (!isGuaranteedToTransferExecutionToSuccessor(LoadBB->begin(),
+                                                    LoadI->getIterator()))
+      return false;
+    TransfersExecution = true;
+  }
 
   // If there is exactly one predecessor where the value is unavailable, the
   // already computed 'OneUnavailablePred' block is it.  If it ends in an
@@ -1407,8 +1410,19 @@ bool JumpThreadingPass::simplifyPartiallyRedundantLoad(LoadInst *LoadI) {
         LoadI->getOrdering(), LoadI->getSyncScopeID(),
         UnavailablePred->getTerminator()->getIterator());
     NewVal->setDebugLoc(LoadI->getDebugLoc());
-    if (AATags)
-      NewVal->setAAMetadata(AATags);
+    NewVal->copyMetadata(*LoadI);
+    NewVal->eraseMetadataIf([&](unsigned Kind, const MDNode *MD) {
+      if (Kind == LLVMContext::MD_dbg || Kind == LLVMContext::MD_annotation)
+        return false;
+      if (is_contained(Metadata::PoisonGeneratingIDs, Kind))
+        return false;
+      // Try to salvage UB-implying metadata if we know it is guaranteed to
+      // transfer the execution to the original load.
+      if (!TransfersExecution.has_value())
+        TransfersExecution = isGuaranteedToTransferExecutionToSuccessor(
+            LoadBB->begin(), LoadI->getIterator());
+      return !*TransfersExecution;
+    });
 
     AvailablePreds.emplace_back(UnavailablePred, NewVal);
   }

llvm/test/Transforms/JumpThreading/pre-load.ll

@@ -1,4 +1,4 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
 ; RUN: opt -passes=jump-threading -S < %s | FileCheck %s
 
 @x = global i32 0
@@ -7,6 +7,10 @@
 declare void @f()
 declare void @g()
 
+;.
+; CHECK: @x = global i32 0
+; CHECK: @y = global i32 0
+;.
 define i32 @pre(i1 %cond, i32 %n) {
 ; CHECK-LABEL: @pre(
 ; CHECK-NEXT:    br i1 [[COND:%.*]], label [[C_THREAD:%.*]], label [[C:%.*]]
@@ -82,3 +86,97 @@ NO:
   call void @g()
   ret i32 1
 }
+
+define i32 @pre_metadata(i1 %cond) {
+; CHECK-LABEL: @pre_metadata(
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[C_THREAD:%.*]], label [[C:%.*]]
+; CHECK:       C.thread:
+; CHECK-NEXT:    store i32 0, ptr @x, align 4
+; CHECK-NEXT:    br label [[YES:%.*]]
+; CHECK:       C:
+; CHECK-NEXT:    [[A_PR:%.*]] = load i32, ptr @y, align 4, !range [[RNG0:![0-9]+]], !noundef [[META1:![0-9]+]]
+; CHECK-NEXT:    [[COND2:%.*]] = icmp eq i32 [[A_PR]], 0
+; CHECK-NEXT:    br i1 [[COND2]], label [[YES]], label [[NO:%.*]]
+; CHECK:       YES:
+; CHECK-NEXT:    [[A4:%.*]] = phi i32 [ 0, [[C_THREAD]] ], [ [[A_PR]], [[C]] ]
+; CHECK-NEXT:    call void @f()
+; CHECK-NEXT:    ret i32 [[A4]]
+; CHECK:       NO:
+; CHECK-NEXT:    call void @g()
+; CHECK-NEXT:    ret i32 1
+;
+  br i1 %cond, label %A, label %B
+
+A:
+  store i32 0, ptr @x, align 4
+  br label %C
+
+B:
+  br label %C
+
+C:
+  %ptr = phi ptr [@x, %A], [@y, %B]
+  %a = load i32, ptr %ptr, align 4, !range !{i32 0, i32 2}, !noundef !{}
+  %cond2 = icmp eq i32 %a, 0
+  br i1 %cond2, label %YES, label %NO
+
+YES:
+  call void @f()
+  ret i32 %a
+
+NO:
+  call void @g()
+  ret i32 1
+}
+
+declare void @callee() memory(none)
+
+define i32 @pre_metadata_may_throw_speculative(i1 %cond) {
+; CHECK-LABEL: @pre_metadata_may_throw_speculative(
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[A:%.*]], label [[C:%.*]]
+; CHECK:       C.thread:
+; CHECK-NEXT:    store i32 0, ptr @x, align 4
+; CHECK-NEXT:    call void @callee()
+; CHECK-NEXT:    br label [[YES:%.*]]
+; CHECK:       C:
+; CHECK-NEXT:    [[A_PR:%.*]] = load i32, ptr @x, align 4, !range [[RNG0]]
+; CHECK-NEXT:    call void @callee()
+; CHECK-NEXT:    [[COND2:%.*]] = icmp eq i32 [[A_PR]], 0
+; CHECK-NEXT:    br i1 [[COND2]], label [[YES]], label [[NO:%.*]]
+; CHECK:       YES:
+; CHECK-NEXT:    [[A3:%.*]] = phi i32 [ 0, [[A]] ], [ [[A_PR]], [[C]] ]
+; CHECK-NEXT:    call void @f()
+; CHECK-NEXT:    ret i32 [[A3]]
+; CHECK:       NO:
+; CHECK-NEXT:    call void @g()
+; CHECK-NEXT:    ret i32 1
+;
+  br i1 %cond, label %A, label %B
+
+A:
+  store i32 0, ptr @x, align 4
+  br label %C
+
+B:
+  br label %C
+
+C:
+  call void @callee()
+  %a = load i32, ptr @x, align 4, !range !{i32 0, i32 2}, !noundef !{}
+  %cond2 = icmp eq i32 %a, 0
+  br i1 %cond2, label %YES, label %NO
+
+YES:
+  call void @f()
+  ret i32 %a
+
+NO:
+  call void @g()
+  ret i32 1
+}
+;.
+; CHECK: attributes #[[ATTR0:[0-9]+]] = { memory(none) }
+;.
+; CHECK: [[RNG0]] = !{i32 0, i32 2}
+; CHECK: [[META1]] = !{}
+;.