[CodeGen][NPM] Port BranchFolder to NPM (llvm#128858)

EnableTailMerge is false by default and is handled by the pass builder.
Passes are independent of target pipeline options.

This completes the generic `MachineLateOptimization` passes for the NPM
pipeline.

Author: optimisan

llvm/include/llvm/CodeGen/BranchFoldingPass.h

@@ -0,0 +1,31 @@
+//===- llvm/CodeGen/BranchFoldingPass.h -------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_CODEGEN_BRANCHFOLDINGPASS_H
+#define LLVM_CODEGEN_BRANCHFOLDINGPASS_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+class BranchFolderPass : public PassInfoMixin<BranchFolderPass> {
+  bool EnableTailMerge;
+
+public:
+  BranchFolderPass(bool EnableTailMerge) : EnableTailMerge(EnableTailMerge) {}
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+
+  MachineFunctionProperties getRequiredProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_BRANCHFOLDINGPASS_H

llvm/include/llvm/InitializePasses.h

@@ -59,7 +59,7 @@ void initializeBasicBlockSectionsPass(PassRegistry &);
 void initializeBarrierNoopPass(PassRegistry &);
 void initializeBasicAAWrapperPassPass(PassRegistry &);
 void initializeBlockFrequencyInfoWrapperPassPass(PassRegistry &);
-void initializeBranchFolderPassPass(PassRegistry &);
+void initializeBranchFolderLegacyPass(PassRegistry &);
 void initializeBranchProbabilityInfoWrapperPassPass(PassRegistry &);
 void initializeBranchRelaxationPass(PassRegistry &);
 void initializeBreakCriticalEdgesPass(PassRegistry &);

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -23,6 +23,7 @@
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/TypeBasedAliasAnalysis.h"
 #include "llvm/CodeGen/AssignmentTrackingAnalysis.h"
+#include "llvm/CodeGen/BranchFoldingPass.h"
 #include "llvm/CodeGen/CallBrPrepare.h"
 #include "llvm/CodeGen/CodeGenPrepare.h"
 #include "llvm/CodeGen/DeadMachineInstructionElim.h"
@@ -1205,7 +1206,7 @@ template <typename Derived, typename TargetMachineT>
 void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineLateOptimization(
     AddMachinePass &addPass) const {
   // Branch folding must be run after regalloc and prolog/epilog insertion.
-  addPass(BranchFolderPass());
+  addPass(BranchFolderPass(Opt.EnableTailMerge));
 
   // Tail duplication.
   // Note that duplicating tail just increases code size and degrades

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -195,6 +195,15 @@ MACHINE_FUNCTION_PASS("verify<machine-trace-metrics>", MachineTraceMetricsVerifi
 #define MACHINE_FUNCTION_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER,    \
                                           PARAMS)
 #endif
+MACHINE_FUNCTION_PASS_WITH_PARAMS(
+    "branch-folder", "BranchFolderPass",
+    [](bool EnableTailMerge) { return BranchFolderPass(EnableTailMerge); },
+    [](StringRef Params) {
+      return parseSinglePassOption(Params, "enable-tail-merge",
+                                   "BranchFolderPass");
+    },
+    "enable-tail-merge")
+
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
     "machine-sink", "MachineSinkingPass",
     [](bool EnableSinkAndFold) {
@@ -246,7 +255,6 @@ DUMMY_MACHINE_FUNCTION_PASS("bbsections-prepare", BasicBlockSectionsPass)
 DUMMY_MACHINE_FUNCTION_PASS("bbsections-profile-reader", BasicBlockSectionsProfileReaderPass)
 DUMMY_MACHINE_FUNCTION_PASS("block-placement", MachineBlockPlacementPass)
 DUMMY_MACHINE_FUNCTION_PASS("block-placement-stats", MachineBlockPlacementStatsPass)
-DUMMY_MACHINE_FUNCTION_PASS("branch-folder", BranchFolderPass)
 DUMMY_MACHINE_FUNCTION_PASS("break-false-deps", BreakFalseDepsPass)
 DUMMY_MACHINE_FUNCTION_PASS("cfguard-longjmp", CFGuardLongjmpPass)
 DUMMY_MACHINE_FUNCTION_PASS("cfi-fixup", CFIFixupPass)

llvm/include/llvm/Target/CGPassBuilderOption.h

@@ -50,6 +50,7 @@ struct CGPassBuilderOption {
   bool EnableGlobalMergeFunc = false;
   bool EnableMachineFunctionSplitter = false;
   bool EnableSinkAndFold = false;
+  bool EnableTailMerge = true;
   bool MISchedPostRA = false;
   bool EarlyLiveIntervals = false;
   bool GCEmptyBlocks = false;

llvm/lib/CodeGen/BranchFolding.cpp

@@ -24,6 +24,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/Analysis.h"
+#include "llvm/CodeGen/BranchFoldingPass.h"
 #include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
@@ -88,38 +89,62 @@ TailMergeSize("tail-merge-size",
 namespace {
 
   /// BranchFolderPass - Wrap branch folder in a machine function pass.
-  class BranchFolderPass : public MachineFunctionPass {
-  public:
-    static char ID;
+class BranchFolderLegacy : public MachineFunctionPass {
+public:
+  static char ID;
 
-    explicit BranchFolderPass(): MachineFunctionPass(ID) {}
+  explicit BranchFolderLegacy() : MachineFunctionPass(ID) {}
 
-    bool runOnMachineFunction(MachineFunction &MF) override;
+  bool runOnMachineFunction(MachineFunction &MF) override;
 
-    void getAnalysisUsage(AnalysisUsage &AU) const override {
-      AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
-      AU.addRequired<MachineBranchProbabilityInfoWrapperPass>();
-      AU.addRequired<ProfileSummaryInfoWrapperPass>();
-      AU.addRequired<TargetPassConfig>();
-      MachineFunctionPass::getAnalysisUsage(AU);
-    }
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
+    AU.addRequired<MachineBranchProbabilityInfoWrapperPass>();
+    AU.addRequired<ProfileSummaryInfoWrapperPass>();
+    AU.addRequired<TargetPassConfig>();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
 
-    MachineFunctionProperties getRequiredProperties() const override {
-      return MachineFunctionProperties().set(
-          MachineFunctionProperties::Property::NoPHIs);
-    }
-  };
+  MachineFunctionProperties getRequiredProperties() const override {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+};
 
 } // end anonymous namespace
 
-char BranchFolderPass::ID = 0;
-
-char &llvm::BranchFolderPassID = BranchFolderPass::ID;
-
-INITIALIZE_PASS(BranchFolderPass, DEBUG_TYPE,
-                "Control Flow Optimizer", false, false)
+char BranchFolderLegacy::ID = 0;
+
+char &llvm::BranchFolderPassID = BranchFolderLegacy::ID;
+
+INITIALIZE_PASS(BranchFolderLegacy, DEBUG_TYPE, "Control Flow Optimizer", false,
+                false)
+
+PreservedAnalyses BranchFolderPass::run(MachineFunction &MF,
+                                        MachineFunctionAnalysisManager &MFAM) {
+  MFPropsModifier _(*this, MF);
+  bool EnableTailMerge =
+      !MF.getTarget().requiresStructuredCFG() && this->EnableTailMerge;
+
+  auto &MBPI = MFAM.getResult<MachineBranchProbabilityAnalysis>(MF);
+  auto *PSI = MFAM.getResult<ModuleAnalysisManagerMachineFunctionProxy>(MF)
+                  .getCachedResult<ProfileSummaryAnalysis>(
+                      *MF.getFunction().getParent());
+  if (!PSI)
+    report_fatal_error(
+        "ProfileSummaryAnalysis is required for BranchFoldingPass", false);
+
+  auto &MBFI = MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+  MBFIWrapper MBBFreqInfo(MBFI);
+  BranchFolder Folder(EnableTailMerge, /*CommonHoist=*/true, MBBFreqInfo, MBPI,
+                      PSI);
+  if (!Folder.OptimizeFunction(MF, MF.getSubtarget().getInstrInfo(),
+                               MF.getSubtarget().getRegisterInfo()))
+    return PreservedAnalyses::all();
+  return getMachineFunctionPassPreservedAnalyses();
+}
 
-bool BranchFolderPass::runOnMachineFunction(MachineFunction &MF) {
+bool BranchFolderLegacy::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 

llvm/lib/CodeGen/CodeGen.cpp

@@ -22,7 +22,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeAtomicExpandLegacyPass(Registry);
   initializeBasicBlockPathCloningPass(Registry);
   initializeBasicBlockSectionsPass(Registry);
-  initializeBranchFolderPassPass(Registry);
+  initializeBranchFolderLegacyPass(Registry);
   initializeBranchRelaxationPass(Registry);
   initializeBreakFalseDepsPass(Registry);
   initializeCallBrPreparePass(Registry);

llvm/lib/Passes/PassBuilder.cpp

@@ -80,6 +80,7 @@
 #include "llvm/CodeGen/AssignmentTrackingAnalysis.h"
 #include "llvm/CodeGen/AtomicExpand.h"
 #include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/BranchFoldingPass.h"
 #include "llvm/CodeGen/CallBrPrepare.h"
 #include "llvm/CodeGen/CodeGenPrepare.h"
 #include "llvm/CodeGen/ComplexDeinterleavingPass.h"

llvm/test/CodeGen/AArch64/branch-folder-oneinst.mir

@@ -1,4 +1,5 @@
 # RUN: llc -o - %s -mtriple=aarch64 -run-pass branch-folder -verify-machineinstrs | FileCheck %s
+# RUN: llc -o - %s -mtriple=aarch64 -passes="require<profile-summary>,function(machine-function(branch-folder<enable-tail-merge>))" -verify-machineinstrs | FileCheck %s
 # Check that BranchFolding pass is able to hoist a common instruction into a block with a single branch instruction.
 name: func
 tracksRegLiveness: true

llvm/test/CodeGen/AMDGPU/branch-folder-requires-no-phis.mir

@@ -1,9 +1,11 @@
 # REQUIRES: asserts
-# RUN: not --crash llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=branch-folder -o /dev/null %s 2>&1 | FileCheck %s
+# RUN: not --crash llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=branch-folder -o /dev/null %s 2>&1 | FileCheck %s --check-prefixes=LEGACY-CHECK,CHECK
+# RUN: not --crash llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -passes="require<profile-summary>,function(machine-function(branch-folder<enable-tail-merge>))" -o /dev/null %s 2>&1 | FileCheck %s --check-prefixes=NPM-CHECK,CHECK
 
 # BranchFolding breaks this function due to phis
 
-# CHECK: MachineFunctionProperties required by Control Flow Optimizer pass are not met by function func.
+# LEGACY-CHECK: MachineFunctionProperties required by Control Flow Optimizer pass are not met by function func.
+# NPM-CHECK: MachineFunctionProperties required by BranchFolderPass pass are not met by function func.
 # CHECK-NEXT: Required properties: NoPHIs
 # CHECK-NEXT: Current properties: IsSSA, TracksLiveness{{$}}
 ---