[llvm] [CodeGen] Add skipFunction() check to MachineFunctionSplitter (PR #166260)

[Grigory Pastukhov via llvm-commits]

https://github.com/grigorypas updated https://github.com/llvm/llvm-project/pull/166260

>From 5064384e8f62afce434165bc26642f2f35780ac7 Mon Sep 17 00:00:00 2001
From: Grigory Pastukhov <gpastukhov at meta.com>
Date: Mon, 3 Nov 2025 15:11:07 -0800
Subject: [PATCH 1/2] Add skipFunction() check to MachineFunctionSplitter pass

---
 llvm/lib/CodeGen/MachineFunctionSplitter.cpp  |  2 +-
 .../machine-function-splitter-optnone.ll      | 50 +++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll

diff --git a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
index c31454a8affda..8ee546d8f038d 100644
--- a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
+++ b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
@@ -136,7 +136,7 @@ bool MachineFunctionSplitter::runOnMachineFunction(MachineFunction &MF) {
   // of exception handling code may be split to cold if user passes the
   // mfs-split-ehcode flag.
   bool UseProfileData = MF.getFunction().hasProfileData();
-  if (!UseProfileData && !SplitAllEHCode)
+  if (!skipFunction(MF.getFunction()) && !UseProfileData && !SplitAllEHCode)
     return false;
 
   const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
diff --git a/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll b/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll
new file mode 100644
index 0000000000000..d2318c777bcce
--- /dev/null
+++ b/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll
@@ -0,0 +1,50 @@
+; REQUIRES: x86-registered-target
+
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -split-machine-functions | FileCheck %s
+
+;; Check that functions with optnone attribute are not split.
+; CHECK-LABEL: foo_optnone:
+; CHECK-NOT:   .section .text.split.foo_optnone
+; CHECK-NOT:   foo_optnone.cold:
+; CHECK:       .LBB0_2:
+; CHECK:       .size   foo_optnone
+
+define void @foo_optnone(i1 zeroext %0) nounwind optnone noinline !prof !14 !section_prefix !15 {
+entry:
+  br i1 %0, label %hot, label %cold, !prof !17
+
+hot:
+  %1 = call i32 @bar()
+  br label %exit
+
+cold:
+  %2 = call i32 @baz()
+  br label %exit
+
+exit:
+  %3 = tail call i32 @qux()
+  ret void
+}
+
+declare i32 @bar()
+declare i32 @baz()
+declare i32 @qux()
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 10000}
+!4 = !{!"MaxCount", i64 10}
+!5 = !{!"MaxInternalCount", i64 1}
+!6 = !{!"MaxFunctionCount", i64 1000}
+!7 = !{!"NumCounts", i64 3}
+!8 = !{!"NumFunctions", i64 5}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 100, i32 1}
+!12 = !{i32 999900, i64 100, i32 1}
+!13 = !{i32 999999, i64 1, i32 2}
+!14 = !{!"function_entry_count", i64 7000}
+!15 = !{!"function_section_prefix", !"hot"}
+!17 = !{!"branch_weights", i32 7000, i32 0}

>From 67eb58e1eca706a100df5a32414de9c079d74fb5 Mon Sep 17 00:00:00 2001
From: Grigory Pastukhov <gpastukhov at meta.com>
Date: Mon, 3 Nov 2025 15:50:44 -0800
Subject: [PATCH 2/2] Fix the bug and address comment

---
 llvm/lib/CodeGen/MachineFunctionSplitter.cpp                 | 5 ++++-
 .../CodeGen/Generic/machine-function-splitter-optnone.ll     | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
index 8ee546d8f038d..b5d3092ee84d8 100644
--- a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
+++ b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
@@ -129,6 +129,9 @@ static bool isColdBlock(const MachineBasicBlock &MBB,
 }
 
 bool MachineFunctionSplitter::runOnMachineFunction(MachineFunction &MF) {
+  if (skipFunction(MF.getFunction()))
+    return false;
+
   // Do not split functions when -basic-block-sections=all is specified.
   if (MF.getTarget().getBBSectionsType() == llvm::BasicBlockSection::All)
     return false;
@@ -136,7 +139,7 @@ bool MachineFunctionSplitter::runOnMachineFunction(MachineFunction &MF) {
   // of exception handling code may be split to cold if user passes the
   // mfs-split-ehcode flag.
   bool UseProfileData = MF.getFunction().hasProfileData();
-  if (!skipFunction(MF.getFunction()) && !UseProfileData && !SplitAllEHCode)
+  if (!UseProfileData && !SplitAllEHCode)
     return false;
 
   const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
diff --git a/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll b/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll
index d2318c777bcce..67d2ad72ee2f4 100644
--- a/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll
+++ b/llvm/test/CodeGen/Generic/machine-function-splitter-optnone.ll
@@ -1,6 +1,6 @@
 ; REQUIRES: x86-registered-target
 
-; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -split-machine-functions | FileCheck %s
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -split-machine-functions -O0 -mfs-psi-cutoff=0 -mfs-count-threshold=10000 | FileCheck %s
 
 ;; Check that functions with optnone attribute are not split.
 ; CHECK-LABEL: foo_optnone: