[CodeGen][NPM] Port MIRFSDiscriminator to NPM

Author: optimisan

llvm/include/llvm/CodeGen/MIRFSDiscriminator.h

@@ -20,6 +20,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachinePassManager.h"
 #include "llvm/Support/Discriminator.h"
 
 #include <cassert>
@@ -29,26 +30,21 @@ namespace llvm {
 class MachineFunction;
 
 using namespace sampleprof;
-class MIRAddFSDiscriminators : public MachineFunctionPass {
+class MIRAddFSDiscriminators {
   MachineFunction *MF = nullptr;
   FSDiscriminatorPass Pass;
   unsigned LowBit;
   unsigned HighBit;
 
 public:
-  static char ID;
   /// PassNum is the sequence number this pass is called, start from 1.
   MIRAddFSDiscriminators(FSDiscriminatorPass P = FSDiscriminatorPass::Pass1)
-      : MachineFunctionPass(ID), Pass(P) {
+      : Pass(P) {
     LowBit = getFSPassBitBegin(P);
     HighBit = getFSPassBitEnd(P);
     assert(LowBit < HighBit && "HighBit needs to be greater than Lowbit");
   }
 
-  StringRef getPassName() const override {
-    return "Add FS discriminators in MIR";
-  }
-
   /// getNumFSBBs() - Return the number of machine BBs that have FS samples.
   unsigned getNumFSBBs();
 
@@ -59,8 +55,33 @@ class MIRAddFSDiscriminators : public MachineFunctionPass {
   /// getMachineFunction - Return the current machine function.
   const MachineFunction *getMachineFunction() const { return MF; }
 
+  bool runOnMachineFunction(MachineFunction &);
+};
+
+class MIRAddFSDiscriminatorsLegacy : public MachineFunctionPass {
+  FSDiscriminatorPass Pass;
+
+public:
+  static char ID;
+  MIRAddFSDiscriminatorsLegacy(
+      FSDiscriminatorPass P = FSDiscriminatorPass::Pass1)
+      : MachineFunctionPass(ID), Pass(P) {}
+  StringRef getPassName() const override {
+    return "Add FS discriminators in MIR";
+  }
+
 private:
-  bool runOnMachineFunction(MachineFunction &) override;
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+
+class MIRAddFSDiscriminatorsPass
+    : public PassInfoMixin<MIRAddFSDiscriminatorsPass> {
+  FSDiscriminatorPass Pass;
+
+public:
+  MIRAddFSDiscriminatorsPass(FSDiscriminatorPass P = FSDiscriminatorPass::Pass1)
+      : Pass(P) {}
+  PreservedAnalyses run(MachineFunction &F, MachineFunctionAnalysisManager &AM);
 };
 
 } // namespace llvm

llvm/include/llvm/InitializePasses.h

@@ -179,7 +179,7 @@ void initializeLowerIntrinsicsPass(PassRegistry &);
 void initializeLowerInvokeLegacyPassPass(PassRegistry &);
 void initializeLowerSwitchLegacyPassPass(PassRegistry &);
 void initializeKCFIPass(PassRegistry &);
-void initializeMIRAddFSDiscriminatorsPass(PassRegistry &);
+void initializeMIRAddFSDiscriminatorsLegacyPass(PassRegistry &);
 void initializeMIRCanonicalizerPass(PassRegistry &);
 void initializeMIRNamerPass(PassRegistry &);
 void initializeMIRPrintingPassPass(PassRegistry &);

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -191,6 +191,14 @@ 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(
+    "mirfs-discriminators", "MIRAddFSDiscriminatorsPass",
+    [](FSDiscriminatorPass Pass) { return MIRAddFSDiscriminatorsPass(Pass); },
+    [PB = this](StringRef Params) {
+      return parseFSDiscriminatorPassOptions(*PB, Params);
+    },
+    "fs-discr-pass")
+
 MACHINE_FUNCTION_PASS_WITH_PARAMS(
     "regallocfast", "RegAllocFastPass",
     [](RegAllocFastPass::Options Opts) { return RegAllocFastPass(Opts); },
@@ -262,7 +270,6 @@ DUMMY_MACHINE_FUNCTION_PASS("machine-sanmd", MachineSanitizerBinaryMetadata)
 DUMMY_MACHINE_FUNCTION_PASS("machine-sink", MachineSinkingPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-uniformity", MachineUniformityInfoWrapperPass)
 DUMMY_MACHINE_FUNCTION_PASS("machineinstr-printer", MachineFunctionPrinterPass)
-DUMMY_MACHINE_FUNCTION_PASS("mirfs-discriminators", MIRAddFSDiscriminatorsPass)
 DUMMY_MACHINE_FUNCTION_PASS("patchable-function", PatchableFunctionPass)
 DUMMY_MACHINE_FUNCTION_PASS("postra-machine-sink", PostRAMachineSinkingPass)
 DUMMY_MACHINE_FUNCTION_PASS("print-machine-uniformity", MachineUniformityInfoPrinterPass)

llvm/lib/CodeGen/CodeGen.cpp

@@ -67,7 +67,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeLocalStackSlotPassPass(Registry);
   initializeLowerGlobalDtorsLegacyPassPass(Registry);
   initializeLowerIntrinsicsPass(Registry);
-  initializeMIRAddFSDiscriminatorsPass(Registry);
+  initializeMIRAddFSDiscriminatorsLegacyPass(Registry);
   initializeMIRCanonicalizerPass(Registry);
   initializeMIRNamerPass(Registry);
   initializeMIRProfileLoaderPassPass(Registry);

llvm/lib/CodeGen/MIRFSDiscriminator.cpp

@@ -42,16 +42,16 @@ cl::opt<bool> ImprovedFSDiscriminator(
              "encoding)"));
 }
 
-char MIRAddFSDiscriminators::ID = 0;
+char MIRAddFSDiscriminatorsLegacy::ID = 0;
 
-INITIALIZE_PASS(MIRAddFSDiscriminators, DEBUG_TYPE,
+INITIALIZE_PASS(MIRAddFSDiscriminatorsLegacy, DEBUG_TYPE,
                 "Add MIR Flow Sensitive Discriminators",
                 /* cfg = */ false, /* is_analysis = */ false)
 
-char &llvm::MIRAddFSDiscriminatorsID = MIRAddFSDiscriminators::ID;
+char &llvm::MIRAddFSDiscriminatorsID = MIRAddFSDiscriminatorsLegacy::ID;
 
 FunctionPass *llvm::createMIRAddFSDiscriminatorsPass(FSDiscriminatorPass P) {
-  return new MIRAddFSDiscriminators(P);
+  return new MIRAddFSDiscriminatorsLegacy(P);
 }
 
 // TODO(xur): Remove this once we switch to ImprovedFSDiscriminator.
@@ -88,6 +88,18 @@ static uint64_t getCallStackHash(const DILocation *DIL) {
   return Ret;
 }
 
+PreservedAnalyses
+MIRAddFSDiscriminatorsPass::run(MachineFunction &MF,
+                                MachineFunctionAnalysisManager &) {
+  if (MIRAddFSDiscriminators(Pass).runOnMachineFunction(MF))
+    return PreservedAnalyses::none();
+  return getMachineFunctionPassPreservedAnalyses();
+}
+
+bool MIRAddFSDiscriminatorsLegacy::runOnMachineFunction(MachineFunction &MF) {
+  return MIRAddFSDiscriminators(Pass).runOnMachineFunction(MF);
+}
+
 // Traverse the CFG and assign FD discriminators. If two instructions
 // have the same lineno and discriminator, but residing in different BBs,
 // the latter instruction will get a new discriminator value. The new

llvm/lib/Passes/PassBuilder.cpp

@@ -107,6 +107,7 @@
 #include "llvm/CodeGen/LiveVariables.h"
 #include "llvm/CodeGen/LocalStackSlotAllocation.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
+#include "llvm/CodeGen/MIRFSDiscriminator.h"
 #include "llvm/CodeGen/MIRPrinter.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
@@ -1431,6 +1432,32 @@ parseRegAllocGreedyFilterFunc(PassBuilder &PB, StringRef Params) {
       inconvertibleErrorCode());
 }
 
+Expected<FSDiscriminatorPass>
+parseFSDiscriminatorPassOptions(PassBuilder &PB, StringRef Params) {
+  if (Params == "Base")
+    return FSDiscriminatorPass::Base;
+  if (Params.consume_front("Pass")) {
+    switch (Params.front()) {
+    case '0':
+      return FSDiscriminatorPass::Pass0;
+    case '1':
+      return FSDiscriminatorPass::Pass1;
+    case '2':
+      return FSDiscriminatorPass::Pass2;
+    case '3':
+      return FSDiscriminatorPass::Pass3;
+    case '4':
+      return FSDiscriminatorPass::Pass4;
+    default:
+      if (Params == "Last")
+        return FSDiscriminatorPass::PassLast;
+    }
+  }
+  return make_error<StringError>(
+      formatv("invalid FSDiscriminator pass parameter '{0}' ", Params).str(),
+      inconvertibleErrorCode());
+}
+
 } // namespace
 
 /// Tests whether a pass name starts with a valid prefix for a default pipeline