[GVN][JumpThreading] Try to salvage UB-implying metadata

Author: dtcxzyw

llvm/lib/Transforms/Scalar/GVN.cpp

@@ -1517,9 +1517,23 @@ void GVNPass::eliminatePartiallyRedundantLoad(
     }
 
     NewLoad->copyMetadata(*Load);
-    // Drop UB-implying metadata as we do not know if it is guaranteed to
-    // transfer the execution to the original load.
-    NewLoad->dropUBImplyingAttrsAndMetadata();
+    std::optional<bool> TransfersExecution = std::nullopt;
+    NewLoad->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()) {
+        // TEST ONLY
+        assert(
+            is_contained(successors(NewLoad->getParent()), Load->getParent()));
+        TransfersExecution = isGuaranteedToTransferExecutionToSuccessor(
+            Load->getParent()->begin(), Load->getIterator());
+      }
+      return !*TransfersExecution;
+    });
 
     // Add the newly created load.
     ValuesPerBlock.push_back(

llvm/lib/Transforms/Scalar/JumpThreading.cpp

@@ -1362,12 +1362,12 @@ 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.
-  bool TransfersExecution = false;
+  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;
+    if (!isGuaranteedToTransferExecutionToSuccessor(LoadBB->begin(),
+                                                    LoadI->getIterator()))
+      return false;
     TransfersExecution = true;
   }
 
@@ -1411,10 +1411,18 @@ bool JumpThreadingPass::simplifyPartiallyRedundantLoad(LoadInst *LoadI) {
         UnavailablePred->getTerminator()->getIterator());
     NewVal->setDebugLoc(LoadI->getDebugLoc());
     NewVal->copyMetadata(*LoadI);
-    // Drop UB-implying metadata if we do not know it is guaranteed to transfer
-    // the execution to the original load.
-    if (!TransfersExecution)
-      NewVal->dropUBImplyingAttrsAndMetadata();
+    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/GVN/pr64598.ll

@@ -6,20 +6,23 @@ define i32 @main(i64 %x, ptr %d, ptr noalias %p) {
 ; CHECK-SAME: (i64 [[X:%.*]], ptr [[D:%.*]], ptr noalias [[P:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[T1_PRE_PRE_PRE:%.*]] = load ptr, ptr [[P]], align 8
-; CHECK-NEXT:    [[T2_PRE_PRE_PRE:%.*]] = load ptr, ptr [[T1_PRE_PRE_PRE]], align 8
+; CHECK-NEXT:    [[T2_PRE_PRE_PRE:%.*]] = load ptr, ptr [[T1_PRE_PRE_PRE]], align 8, !tbaa [[TBAA0:![0-9]+]]
 ; CHECK-NEXT:    [[T3_PRE_PRE_PRE:%.*]] = load ptr, ptr [[T2_PRE_PRE_PRE]], align 8
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[T1_PRE_PRE:%.*]] = phi ptr [ [[T1_PRE_PRE19:%.*]], [[LOOP_LATCH:%.*]] ], [ [[T1_PRE_PRE_PRE]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[T2_PRE_PRE:%.*]] = phi ptr [ [[T2_PRE_PRE23:%.*]], [[LOOP_LATCH:%.*]] ], [ [[T2_PRE_PRE_PRE]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[T1_PRE_PRE:%.*]] = phi ptr [ [[T1_PRE_PRE19:%.*]], [[LOOP_LATCH]] ], [ [[T1_PRE_PRE_PRE]], [[ENTRY]] ]
 ; CHECK-NEXT:    br label [[LOOP2:%.*]]
 ; CHECK:       loop2:
-; CHECK-NEXT:    [[T1_PRE_PRE21:%.*]] = phi ptr [ [[T1_PRE_PRE19]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE:%.*]] ], [ [[T1_PRE_PRE]], [[LOOP]] ]
+; CHECK-NEXT:    [[T2_PRE_PRE25:%.*]] = phi ptr [ [[T2_PRE_PRE23]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE:%.*]] ], [ [[T2_PRE_PRE]], [[LOOP]] ]
+; CHECK-NEXT:    [[T1_PRE_PRE21:%.*]] = phi ptr [ [[T1_PRE_PRE19]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE]] ], [ [[T1_PRE_PRE]], [[LOOP]] ]
 ; CHECK-NEXT:    [[T3_PRE:%.*]] = phi ptr [ [[T3_PRE16:%.*]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE]] ], [ [[T3_PRE_PRE_PRE]], [[LOOP]] ]
-; CHECK-NEXT:    [[T2_PRE:%.*]] = phi ptr [ [[T2_PRE13:%.*]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE]] ], [ [[T2_PRE_PRE_PRE]], [[LOOP]] ]
+; CHECK-NEXT:    [[T2_PRE:%.*]] = phi ptr [ [[T2_PRE13:%.*]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE]] ], [ [[T2_PRE_PRE]], [[LOOP]] ]
 ; CHECK-NEXT:    [[T1_PRE:%.*]] = phi ptr [ [[T1_PRE10:%.*]], [[LOOP2_LATCH_LOOP2_CRIT_EDGE]] ], [ [[T1_PRE_PRE]], [[LOOP]] ]
 ; CHECK-NEXT:    br label [[LOOP3:%.*]]
 ; CHECK:       loop3:
-; CHECK-NEXT:    [[T1_PRE_PRE20:%.*]] = phi ptr [ [[T1_PRE_PRE19]], [[LOOP3_LATCH:%.*]] ], [ [[T1_PRE_PRE21]], [[LOOP2]] ]
+; CHECK-NEXT:    [[T2_PRE_PRE24:%.*]] = phi ptr [ [[T2_PRE_PRE23]], [[LOOP3_LATCH:%.*]] ], [ [[T2_PRE_PRE25]], [[LOOP2]] ]
+; CHECK-NEXT:    [[T1_PRE_PRE20:%.*]] = phi ptr [ [[T1_PRE_PRE19]], [[LOOP3_LATCH]] ], [ [[T1_PRE_PRE21]], [[LOOP2]] ]
 ; CHECK-NEXT:    [[T3_PRE17:%.*]] = phi ptr [ [[T3_PRE16]], [[LOOP3_LATCH]] ], [ [[T3_PRE]], [[LOOP2]] ]
 ; CHECK-NEXT:    [[T2_PRE14:%.*]] = phi ptr [ [[T2_PRE13]], [[LOOP3_LATCH]] ], [ [[T2_PRE]], [[LOOP2]] ]
 ; CHECK-NEXT:    [[T1_PRE11:%.*]] = phi ptr [ [[T1_PRE10]], [[LOOP3_LATCH]] ], [ [[T1_PRE]], [[LOOP2]] ]
@@ -33,10 +36,11 @@ define i32 @main(i64 %x, ptr %d, ptr noalias %p) {
 ; CHECK:       for.body.lr.ph.i:
 ; CHECK-NEXT:    store i32 0, ptr [[P]], align 4
 ; CHECK-NEXT:    [[T5_PRE:%.*]] = load ptr, ptr [[P]], align 8
-; CHECK-NEXT:    [[T6_PRE:%.*]] = load ptr, ptr [[T5_PRE]], align 8
+; CHECK-NEXT:    [[T6_PRE:%.*]] = load ptr, ptr [[T5_PRE]], align 8, !tbaa [[TBAA0]]
 ; CHECK-NEXT:    [[T7_PRE:%.*]] = load ptr, ptr [[T6_PRE]], align 8
 ; CHECK-NEXT:    br label [[LOOP3_LATCH]]
 ; CHECK:       loop3.latch:
+; CHECK-NEXT:    [[T2_PRE_PRE23]] = phi ptr [ [[T2_PRE_PRE24]], [[LOOP3_LOOP3_LATCH_CRIT_EDGE]] ], [ [[T6_PRE]], [[FOR_BODY_LR_PH_I]] ]
 ; CHECK-NEXT:    [[T1_PRE_PRE19]] = phi ptr [ [[T1_PRE_PRE20]], [[LOOP3_LOOP3_LATCH_CRIT_EDGE]] ], [ [[T5_PRE]], [[FOR_BODY_LR_PH_I]] ]
 ; CHECK-NEXT:    [[T3_PRE16]] = phi ptr [ [[T3_PRE17]], [[LOOP3_LOOP3_LATCH_CRIT_EDGE]] ], [ [[T7_PRE]], [[FOR_BODY_LR_PH_I]] ]
 ; CHECK-NEXT:    [[T2_PRE13]] = phi ptr [ [[T2_PRE14]], [[LOOP3_LOOP3_LATCH_CRIT_EDGE]] ], [ [[T6_PRE]], [[FOR_BODY_LR_PH_I]] ]
@@ -50,7 +54,7 @@ define i32 @main(i64 %x, ptr %d, ptr noalias %p) {
 ; CHECK:       loop2.latch.loop2_crit_edge:
 ; CHECK-NEXT:    br label [[LOOP2]]
 ; CHECK:       loop.latch:
-; CHECK-NEXT:    store i32 0, ptr [[D]], align 4, !tbaa [[TBAA0:![0-9]+]]
+; CHECK-NEXT:    store i32 0, ptr [[D]], align 4, !tbaa [[TBAA4:![0-9]+]]
 ; CHECK-NEXT:    br label [[LOOP]]
 ;
 entry: