[IR] Add MD_prof to the Keep list of dropUBImplyingAttrsAndMetadata

[mtrofin]

MD_prof is safe to keep when e.g. hoisting instructions.

Issue #147390

[mtrofin]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• [NFC] Leave a comment in Local.cpp about debug info & sample profiling #155296
• [IR] Add MD_prof to the Keep list of dropUBImplyingAttrsAndMetadata #154635 👈 (View in Graphite)
• [SimplfyCFG] Set MD_prof for select used for certain conditional simplifications #154426
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[mtrofin]

I believe this is different from #152420 (comment) because this is less generic - there, it was about whether a caller asking to dropUnknown...Metadata would be surprised to get MD_prof still there. Here, I'd argue the API is specifically about hoisting, and it's more reasonable to keep around some metadata. @nikic

[nikic]

I think it's actually fine to always handle MD_prof in dropUBImplyingAttrsAndMetadata() (that is push this one level down so you don't need to deal with this in each caller). My concern was only about doing it inside dropUnknownNonDebugMetadata().

[mtrofin]

Oh - that will simplify a few more things then.

[mtrofin]

Done, ptal.

[llvmbot]

@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/154635.diff

3 Files Affected:

• (modified) llvm/include/llvm/Transforms/Utils/Local.h (+1)
• (modified) llvm/lib/Transforms/Utils/Local.cpp (+11-3)
• (modified) llvm/test/Transforms/SimplifyCFG/PhiBlockMerge.ll (+14-7)

diff --git a/llvm/include/llvm/Transforms/Utils/Local.h b/llvm/include/llvm/Transforms/Utils/Local.h
index 3f5f4278a2766..605f8d6a19ebe 100644
--- a/llvm/include/llvm/Transforms/Utils/Local.h
+++ b/llvm/include/llvm/Transforms/Utils/Local.h
@@ -498,6 +498,7 @@ LLVM_ABI void dropDebugUsers(Instruction &I);
 ///
 /// The moved instructions receive the insertion point debug location values
 /// (DILocations) and their debug intrinsic instructions are removed.
+/// Selects and indirect calls keep their MD_prof metadata.
 LLVM_ABI void hoistAllInstructionsInto(BasicBlock *DomBlock,
                                        Instruction *InsertPt, BasicBlock *BB);
 
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index ac344904f90f0..ad796f8605de0 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -116,6 +116,8 @@ static cl::opt<unsigned> MaxPhiEntriesIncreaseAfterRemovingEmptyBlock(
     cl::desc("Stop removing an empty block if removing it will introduce more "
              "than this number of phi entries in its successor"));
 
+extern cl::opt<bool> ProfcheckDisableMetadataFixes;
+
 // Max recursion depth for collectBitParts used when detecting bswap and
 // bitreverse idioms.
 static const unsigned BitPartRecursionMaxDepth = 48;
@@ -3342,8 +3344,11 @@ void llvm::hoistAllInstructionsInto(BasicBlock *DomBlock, Instruction *InsertPt,
   // retain their original debug locations (DILocations) and debug intrinsic
   // instructions.
   //
-  // Doing so would degrade the debugging experience and adversely affect the
-  // accuracy of profiling information.
+  // Doing so would degrade the debugging experience.
+  //
+  // FIXME: Issue #152767: debug info should also be the same as the
+  // original branch, **if** the user explicitly indicated that (for sampling
+  // PGO)
   //
   // Currently, when hoisting the instructions, we take the following actions:
   // - Remove their debug intrinsic instructions.
@@ -3362,7 +3367,10 @@ void llvm::hoistAllInstructionsInto(BasicBlock *DomBlock, Instruction *InsertPt,
 
   for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
     Instruction *I = &*II;
-    I->dropUBImplyingAttrsAndMetadata();
+    if (ProfcheckDisableMetadataFixes)
+      I->dropUBImplyingAttrsAndMetadata();
+    else
+      I->dropUBImplyingAttrsAndMetadata({LLVMContext::MD_prof});
     if (I->isUsedByMetadata())
       dropDebugUsers(*I);
     // RemoveDIs: drop debug-info too as the following code does.
diff --git a/llvm/test/Transforms/SimplifyCFG/PhiBlockMerge.ll b/llvm/test/Transforms/SimplifyCFG/PhiBlockMerge.ll
index 2c5889a981db2..08397b5755a3f 100644
--- a/llvm/test/Transforms/SimplifyCFG/PhiBlockMerge.ll
+++ b/llvm/test/Transforms/SimplifyCFG/PhiBlockMerge.ll
@@ -1,20 +1,21 @@
-; NOTE: Assertions have been autogenerated by update_test_checks.py
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 5
 ; Test merging of blocks that only have PHI nodes in them
 ;
 ; RUN: opt < %s -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s
 ;
 
 define i32 @test(i1 %a, i1 %b) {
-; CHECK-LABEL: @test(
-; CHECK:       M:
-; CHECK-NEXT:    [[DOT:%.*]] = select i1 %b, i32 0, i32 1
-; CHECK-NEXT:    [[W:%.*]] = select i1 %a, i32 2, i32 [[DOT]]
+; CHECK-LABEL: define i32 @test(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]]) {
+; CHECK-NEXT:  [[M:.*:]]
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[B]], i32 0, i32 1, !prof [[PROF0:![0-9]+]]
+; CHECK-NEXT:    [[W:%.*]] = select i1 [[A]], i32 2, i32 [[SPEC_SELECT]], !prof [[PROF1:![0-9]+]]
 ; CHECK-NEXT:    [[R:%.*]] = add i32 [[W]], 1
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
-  br i1 %a, label %M, label %O
+  br i1 %a, label %M, label %O, !prof !0
 O:              ; preds = %0
-  br i1 %b, label %N, label %Q
+  br i1 %b, label %N, label %Q, !prof !1
 Q:              ; preds = %O
   br label %N
 N:              ; preds = %Q, %O
@@ -27,3 +28,9 @@ M:              ; preds = %N, %0
   ret i32 %R
 }
 
+!0 = !{!"branch_weights", i32 11, i32 7}
+!1 = !{!"branch_weights", i32 3, i32 5}
+;.
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 3, i32 5}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 11, i32 7}
+;.

[nikic]

LGTM

[nikic]

I think you don't need the --check-globals, just the new version is enough.

[nikic]

I think the Keep is no longer needed now, but I guess it doesn't hurt to keep it...