[DFAJT][profcheck] Propagate select -> br profile metadata

mtrofin · mtrofin · commit bd0dbbbeb008 · 2025-10-07T07:27:39.000-07:00
diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
@@ -120,6 +120,9 @@ static cl::opt<unsigned>
     CostThreshold("dfa-cost-threshold",
                   cl::desc("Maximum cost accepted for the transformation"),
                   cl::Hidden, cl::init(50));
+
+extern cl::opt<bool> ProfcheckDisableMetadataFixes;
+
 } // namespace llvm
 
 namespace {
@@ -255,7 +258,11 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
 
     // Insert the real conditional branch based on the original condition.
     StartBlockTerm->eraseFromParent();
-    BranchInst::Create(EndBlock, NewBlock, SI->getCondition(), StartBlock);
+    auto *BI =
+        BranchInst::Create(EndBlock, NewBlock, SI->getCondition(), StartBlock);
+    if (!ProfcheckDisableMetadataFixes)
+      BI->setMetadata(LLVMContext::MD_prof,
+                      SI->getMetadata(LLVMContext::MD_prof));
     DTU->applyUpdates({{DominatorTree::Insert, StartBlock, EndBlock},
                        {DominatorTree::Insert, StartBlock, NewBlock}});
   } else {
@@ -290,7 +297,11 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
     //  (Use)
     BranchInst::Create(EndBlock, NewBlockF);
     // Insert the real conditional branch based on the original condition.
-    BranchInst::Create(EndBlock, NewBlockF, SI->getCondition(), NewBlockT);
+    auto *BI =
+        BranchInst::Create(EndBlock, NewBlockF, SI->getCondition(), NewBlockT);
+    if (!ProfcheckDisableMetadataFixes)
+      BI->setMetadata(LLVMContext::MD_prof,
+                      SI->getMetadata(LLVMContext::MD_prof));
     DTU->applyUpdates({{DominatorTree::Insert, NewBlockT, NewBlockF},
                        {DominatorTree::Insert, NewBlockT, EndBlock},
                        {DominatorTree::Insert, NewBlockF, EndBlock}});
diff --git a/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll b/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll
@@ -1,11 +1,12 @@
 ; REQUIRES: asserts
 ; RUN: opt -S -passes=dfa-jump-threading -debug-only=dfa-jump-threading -disable-output %s 2>&1 | FileCheck %s
+; RUN: opt -S -passes=dfa-jump-threading -print-prof-data %s -o - | FileCheck %s --check-prefix=PROFILE
 
 ; This test checks that the analysis identifies all threadable paths in a
 ; simple CFG. A threadable path includes a list of basic blocks, the exit
 ; state, and the block that determines the next state.
 ; < path of BBs that form a cycle > [ state, determinator ]
-define i32 @test1(i32 %num) {
+define i32 @test1(i32 %num) !prof !0{
 ; CHECK: < case2 for.inc for.body > [ 1, for.inc ]
 ; CHECK-NEXT: < for.inc for.body > [ 1, for.inc ]
 ; CHECK-NEXT: < case1 for.inc for.body > [ 2, for.inc ]
@@ -25,8 +26,11 @@ case1:
   br label %for.inc
 
 case2:
+  ; PROFILE-LABEL: @test1
+  ; PROFILE-LABEL: case2:
+  ; PROFILE: br i1 %cmp, label %for.inc.jt1, label %sel.si.unfold.false.jt2, !prof !1 ; !1 = !{!"branch_weights", i32 3, i32 5}
   %cmp = icmp eq i32 %count, 50
-  %sel = select i1 %cmp, i32 1, i32 2
+  %sel = select i1 %cmp, i32 1, i32 2, !prof !1
   br label %for.inc
 
 for.inc:
@@ -182,7 +186,7 @@ bb66:                                             ; preds = %bb59
 }
 
 ; Value %init is not predictable but it's okay since it is the value initial to the switch.
-define i32 @initial.value.positive1(i32 %init) {
+define i32 @initial.value.positive1(i32 %init) !prof !0 {
 ; CHECK: < loop.1.backedge loop.1 loop.2 loop.3 > [ 1, loop.1 ]
 ; CHECK-NEXT: < case4 loop.1.backedge state.1.be2.si.unfold.false loop.1 loop.2 loop.3 > [ 2, loop.1.backedge ]
 ; CHECK-NEXT: < case2 loop.1.backedge state.1.be2.si.unfold.false loop.1 loop.2 loop.3 > [ 4, loop.1.backedge ]
@@ -241,3 +245,6 @@ infloop.i:
 exit:
   ret i32 0
 }
+
+!0 = !{!"function_entry_count", i32 10}
+!1 = !{!"branch_weights", i32 3, i32 5}
diff --git a/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll b/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.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 -S -passes=dfa-jump-threading %s | FileCheck %s
 
 ; These tests check that the DFA jump threading transformation is applied
@@ -301,7 +301,7 @@ end:
   ret void
 }
 
-define void @pr106083_invalidBBarg_fold(i1 %cmp1, i1 %cmp2, i1 %not, ptr %d) {
+define void @pr106083_invalidBBarg_fold(i1 %cmp1, i1 %cmp2, i1 %not, ptr %d) !prof !0 {
 ; CHECK-LABEL: @pr106083_invalidBBarg_fold(
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    br label [[BB1:%.*]]
@@ -310,7 +310,7 @@ define void @pr106083_invalidBBarg_fold(i1 %cmp1, i1 %cmp2, i1 %not, ptr %d) {
 ; CHECK-NEXT:    br i1 [[NOT:%.*]], label [[BB7_JT0]], label [[BB2:%.*]]
 ; CHECK:       BB2:
 ; CHECK-NEXT:    store i16 0, ptr [[D:%.*]], align 2
-; CHECK-NEXT:    br i1 [[CMP2:%.*]], label [[BB7:%.*]], label [[SPEC_SELECT_SI_UNFOLD_FALSE_JT0:%.*]]
+; CHECK-NEXT:    br i1 [[CMP2:%.*]], label [[BB7:%.*]], label [[SPEC_SELECT_SI_UNFOLD_FALSE_JT0:%.*]], !prof [[PROF1:![0-9]+]]
 ; CHECK:       spec.select.si.unfold.false:
 ; CHECK-NEXT:    br label [[BB9]]
 ; CHECK:       spec.select.si.unfold.false.jt0:
@@ -357,7 +357,7 @@ BB1:                                              ; preds = %BB1.backedge, %BB7,
 
 BB2:                                              ; preds = %BB1
   store i16 0, ptr %d, align 2
-  %spec.select = select i1 %cmp2, i32 %sel, i32 0
+  %spec.select = select i1 %cmp2, i32 %sel, i32 0, !prof !1
   br label %BB7
 
 BB7:                                              ; preds = %BB2, %BB1
@@ -444,3 +444,10 @@ select.unfold:                                    ; preds = %bb1, %.loopexit6
 bb2:                                              ; preds = %select.unfold
   unreachable
 }
+
+!0 = !{!"function_entry_count", i32 10}
+!1 = !{!"branch_weights", i32 3, i32 5}
+;.
+; CHECK: [[META0:![0-9]+]] = !{!"function_entry_count", i32 10}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 3, i32 5}
+;.
diff --git a/llvm/utils/profcheck-xfail.txt b/llvm/utils/profcheck-xfail.txt
@@ -711,11 +711,6 @@ Transforms/CorrelatedValuePropagation/urem.ll
 Transforms/CrossDSOCFI/basic.ll
 Transforms/CrossDSOCFI/cfi_functions.ll
 Transforms/CrossDSOCFI/thumb.ll
-Transforms/DFAJumpThreading/dfa-jump-threading-analysis.ll
-Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll
-Transforms/DFAJumpThreading/dfa-unfold-select.ll
-Transforms/DFAJumpThreading/max-path-length.ll
-Transforms/DFAJumpThreading/negative.ll
 Transforms/ExpandFp/AMDGPU/frem-inf.ll
 Transforms/ExpandFp/AMDGPU/frem.ll
 Transforms/ExpandLargeDivRem/X86/sdiv129.ll
