[CodeGen][NewPM] Port "ShrinkWrap" pass to NPM (#129880)

Author: vikramRH

llvm/include/llvm/CodeGen/ShrinkWrap.h

@@ -0,0 +1,29 @@
+//===- llvm/CodeGen/ShrinkWrap.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_SHRINKWRAP_H
+#define LLVM_CODEGEN_SHRINKWRAP_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+class ShrinkWrapPass : public PassInfoMixin<ShrinkWrapPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+
+  MachineFunctionProperties getRequiredProperties() const {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoVRegs);
+  }
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_SHRINKWRAP_H

llvm/include/llvm/InitializePasses.h

@@ -285,7 +285,7 @@ void initializeScavengerTestPass(PassRegistry &);
 void initializeScopedNoAliasAAWrapperPassPass(PassRegistry &);
 void initializeSeparateConstOffsetFromGEPLegacyPassPass(PassRegistry &);
 void initializeShadowStackGCLoweringPass(PassRegistry &);
-void initializeShrinkWrapPass(PassRegistry &);
+void initializeShrinkWrapLegacyPass(PassRegistry &);
 void initializeSingleLoopExtractorPass(PassRegistry &);
 void initializeSinkingLegacyPassPass(PassRegistry &);
 void initializeSjLjEHPreparePass(PassRegistry &);

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -84,6 +84,7 @@
 #include "llvm/CodeGen/SanitizerBinaryMetadata.h"
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
+#include "llvm/CodeGen/ShrinkWrap.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
 #include "llvm/CodeGen/StackColoring.h"
 #include "llvm/CodeGen/StackFrameLayoutAnalysisPass.h"

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -192,6 +192,7 @@ MACHINE_FUNCTION_PASS("remove-loads-into-fake-uses", RemoveLoadsIntoFakeUsesPass
 MACHINE_FUNCTION_PASS("remove-redundant-debug-values", RemoveRedundantDebugValuesPass())
 MACHINE_FUNCTION_PASS("require-all-machine-function-properties",
                       RequireAllMachineFunctionPropertiesPass())
+MACHINE_FUNCTION_PASS("shrink-wrap", ShrinkWrapPass())
 MACHINE_FUNCTION_PASS("stack-coloring", StackColoringPass())
 MACHINE_FUNCTION_PASS("stack-frame-layout", StackFrameLayoutAnalysisPass())
 MACHINE_FUNCTION_PASS("stack-slot-coloring", StackSlotColoringPass())
@@ -317,7 +318,6 @@ DUMMY_MACHINE_FUNCTION_PASS("regalloc", RegAllocPass)
 DUMMY_MACHINE_FUNCTION_PASS("regallocscoringpass", RegAllocScoringPass)
 DUMMY_MACHINE_FUNCTION_PASS("regbankselect", RegBankSelectPass)
 DUMMY_MACHINE_FUNCTION_PASS("reset-machine-function", ResetMachineFunctionPass)
-DUMMY_MACHINE_FUNCTION_PASS("shrink-wrap", ShrinkWrapPass)
 DUMMY_MACHINE_FUNCTION_PASS("stackmap-liveness", StackMapLivenessPass)
 DUMMY_MACHINE_FUNCTION_PASS("unpack-mi-bundles", UnpackMachineBundlesPass)
 DUMMY_MACHINE_FUNCTION_PASS("virtregrewriter", VirtRegRewriterPass)

llvm/lib/CodeGen/CodeGen.cpp

@@ -123,7 +123,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeSafeStackLegacyPassPass(Registry);
   initializeSelectOptimizePass(Registry);
   initializeShadowStackGCLoweringPass(Registry);
-  initializeShrinkWrapPass(Registry);
+  initializeShrinkWrapLegacyPass(Registry);
   initializeSjLjEHPreparePass(Registry);
   initializeSlotIndexesWrapperPassPass(Registry);
   initializeStackColoringLegacyPass(Registry);

llvm/lib/CodeGen/ShrinkWrap.cpp

@@ -47,6 +47,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/ShrinkWrap.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SetVector.h"
@@ -110,7 +111,7 @@ namespace {
 /// does not rely on expensive data-flow analysis. Instead we use the
 /// dominance properties and loop information to decide which point
 /// are safe for such insertion.
-class ShrinkWrap : public MachineFunctionPass {
+class ShrinkWrapImpl {
   /// Hold callee-saved information.
   RegisterClassInfo RCI;
   MachineDominatorTree *MDT = nullptr;
@@ -224,13 +225,8 @@ class ShrinkWrap : public MachineFunctionPass {
   /// Initialize the pass for \p MF.
   void init(MachineFunction &MF) {
     RCI.runOnMachineFunction(MF);
-    MDT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
-    MPDT = &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
     Save = nullptr;
     Restore = nullptr;
-    MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
-    MLI = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
-    ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
     EntryFreq = MBFI->getEntryFreq();
     const TargetSubtargetInfo &Subtarget = MF.getSubtarget();
     const TargetInstrInfo &TII = *Subtarget.getInstrInfo();
@@ -248,14 +244,24 @@ class ShrinkWrap : public MachineFunctionPass {
   /// shrink-wrapping.
   bool ArePointsInteresting() const { return Save != Entry && Save && Restore; }
 
+public:
+  ShrinkWrapImpl(MachineDominatorTree *MDT, MachinePostDominatorTree *MPDT,
+                 MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *MLI,
+                 MachineOptimizationRemarkEmitter *ORE)
+      : MDT(MDT), MPDT(MPDT), MBFI(MBFI), MLI(MLI), ORE(ORE) {}
+
   /// Check if shrink wrapping is enabled for this target and function.
   static bool isShrinkWrapEnabled(const MachineFunction &MF);
 
+  bool run(MachineFunction &MF);
+};
+
+class ShrinkWrapLegacy : public MachineFunctionPass {
 public:
   static char ID;
 
-  ShrinkWrap() : MachineFunctionPass(ID) {
-    initializeShrinkWrapPass(*PassRegistry::getPassRegistry());
+  ShrinkWrapLegacy() : MachineFunctionPass(ID) {
+    initializeShrinkWrapLegacyPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -282,20 +288,22 @@ class ShrinkWrap : public MachineFunctionPass {
 
 } // end anonymous namespace
 
-char ShrinkWrap::ID = 0;
+char ShrinkWrapLegacy::ID = 0;
 
-char &llvm::ShrinkWrapID = ShrinkWrap::ID;
+char &llvm::ShrinkWrapID = ShrinkWrapLegacy::ID;
 
-INITIALIZE_PASS_BEGIN(ShrinkWrap, DEBUG_TYPE, "Shrink Wrap Pass", false, false)
+INITIALIZE_PASS_BEGIN(ShrinkWrapLegacy, DEBUG_TYPE, "Shrink Wrap Pass", false,
+                      false)
 INITIALIZE_PASS_DEPENDENCY(MachineBlockFrequencyInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachinePostDominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineLoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
-INITIALIZE_PASS_END(ShrinkWrap, DEBUG_TYPE, "Shrink Wrap Pass", false, false)
+INITIALIZE_PASS_END(ShrinkWrapLegacy, DEBUG_TYPE, "Shrink Wrap Pass", false,
+                    false)
 
-bool ShrinkWrap::useOrDefCSROrFI(const MachineInstr &MI, RegScavenger *RS,
-                                 bool StackAddressUsed) const {
+bool ShrinkWrapImpl::useOrDefCSROrFI(const MachineInstr &MI, RegScavenger *RS,
+                                     bool StackAddressUsed) const {
   /// Check if \p Op is known to access an address not on the function's stack .
   /// At the moment, accesses where the underlying object is a global, function
   /// argument, or jump table are considered non-stack accesses. Note that the
@@ -549,7 +557,7 @@ static void rollbackRestoreSplit(MachineFunction &MF, MachineBasicBlock *NMBB,
 // A block is deemed fit for restore point split iff there exist
 // 1. DirtyPreds - preds of CurRestore reachable from use or def of CSR/FI
 // 2. CleanPreds - preds of CurRestore that arent DirtyPreds
-bool ShrinkWrap::checkIfRestoreSplittable(
+bool ShrinkWrapImpl::checkIfRestoreSplittable(
     const MachineBasicBlock *CurRestore,
     const DenseSet<const MachineBasicBlock *> &ReachableByDirty,
     SmallVectorImpl<MachineBasicBlock *> &DirtyPreds,
@@ -572,8 +580,8 @@ bool ShrinkWrap::checkIfRestoreSplittable(
   return !(CleanPreds.empty() || DirtyPreds.empty());
 }
 
-bool ShrinkWrap::postShrinkWrapping(bool HasCandidate, MachineFunction &MF,
-                                    RegScavenger *RS) {
+bool ShrinkWrapImpl::postShrinkWrapping(bool HasCandidate, MachineFunction &MF,
+                                        RegScavenger *RS) {
   if (!EnablePostShrinkWrapOpt)
     return false;
 
@@ -679,8 +687,8 @@ bool ShrinkWrap::postShrinkWrapping(bool HasCandidate, MachineFunction &MF,
   return true;
 }
 
-void ShrinkWrap::updateSaveRestorePoints(MachineBasicBlock &MBB,
-                                         RegScavenger *RS) {
+void ShrinkWrapImpl::updateSaveRestorePoints(MachineBasicBlock &MBB,
+                                             RegScavenger *RS) {
   // Get rid of the easy cases first.
   if (!Save)
     Save = &MBB;
@@ -810,7 +818,7 @@ static bool giveUpWithRemarks(MachineOptimizationRemarkEmitter *ORE,
   return false;
 }
 
-bool ShrinkWrap::performShrinkWrapping(
+bool ShrinkWrapImpl::performShrinkWrapping(
     const ReversePostOrderTraversal<MachineBasicBlock *> &RPOT,
     RegScavenger *RS) {
   for (MachineBasicBlock *MBB : RPOT) {
@@ -919,10 +927,7 @@ bool ShrinkWrap::performShrinkWrapping(
   return true;
 }
 
-bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
-  if (skipFunction(MF.getFunction()) || MF.empty() || !isShrinkWrapEnabled(MF))
-    return false;
-
+bool ShrinkWrapImpl::run(MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "**** Analysing " << MF.getName() << '\n');
 
   init(MF);
@@ -969,7 +974,44 @@ bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
   return Changed;
 }
 
-bool ShrinkWrap::isShrinkWrapEnabled(const MachineFunction &MF) {
+bool ShrinkWrapLegacy::runOnMachineFunction(MachineFunction &MF) {
+  if (skipFunction(MF.getFunction()) || MF.empty() ||
+      !ShrinkWrapImpl::isShrinkWrapEnabled(MF))
+    return false;
+
+  MachineDominatorTree *MDT =
+      &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+  MachinePostDominatorTree *MPDT =
+      &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
+  MachineBlockFrequencyInfo *MBFI =
+      &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+  MachineLoopInfo *MLI = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+  MachineOptimizationRemarkEmitter *ORE =
+      &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
+
+  return ShrinkWrapImpl(MDT, MPDT, MBFI, MLI, ORE).run(MF);
+}
+
+PreservedAnalyses ShrinkWrapPass::run(MachineFunction &MF,
+                                      MachineFunctionAnalysisManager &MFAM) {
+  MFPropsModifier _(*this, MF);
+  if (MF.empty() || !ShrinkWrapImpl::isShrinkWrapEnabled(MF))
+    return PreservedAnalyses::all();
+
+  MachineDominatorTree &MDT = MFAM.getResult<MachineDominatorTreeAnalysis>(MF);
+  MachinePostDominatorTree &MPDT =
+      MFAM.getResult<MachinePostDominatorTreeAnalysis>(MF);
+  MachineBlockFrequencyInfo &MBFI =
+      MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+  MachineLoopInfo &MLI = MFAM.getResult<MachineLoopAnalysis>(MF);
+  MachineOptimizationRemarkEmitter &ORE =
+      MFAM.getResult<MachineOptimizationRemarkEmitterAnalysis>(MF);
+
+  ShrinkWrapImpl(&MDT, &MPDT, &MBFI, &MLI, &ORE).run(MF);
+  return PreservedAnalyses::all();
+}
+
+bool ShrinkWrapImpl::isShrinkWrapEnabled(const MachineFunction &MF) {
   const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
 
   switch (EnableShrinkWrapOpt) {

llvm/lib/Passes/PassBuilder.cpp

@@ -158,6 +158,7 @@
 #include "llvm/CodeGen/SanitizerBinaryMetadata.h"
 #include "llvm/CodeGen/SelectOptimize.h"
 #include "llvm/CodeGen/ShadowStackGCLowering.h"
+#include "llvm/CodeGen/ShrinkWrap.h"
 #include "llvm/CodeGen/SjLjEHPrepare.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/SpillPlacement.h"

llvm/test/CodeGen/AArch64/dont-shrink-wrap-stack-mayloadorstore.mir

@@ -4,6 +4,7 @@
  ; to a position where the stack might still be accessed by a load or store
 
  ; RUN: llc -x=mir -simplify-mir -run-pass=shrink-wrap -o - %s | FileCheck %s
+ ; RUN: llc -x=mir -simplify-mir -passes='shrink-wrap' -o - %s | FileCheck %s
  ; CHECK:      name:            compiler_pop_stack
  ; CHECK:      frameInfo:
  ; CHECK:      savePoint:       '%bb.1'

llvm/test/CodeGen/AArch64/shrinkwrap-split-restore-point.mir

@@ -1,5 +1,6 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
 # RUN: llc -mtriple=aarch64 -run-pass=shrink-wrap -o - %s | FileCheck %s
+# RUN: llc -mtriple=aarch64 -passes='shrink-wrap' -o - %s | FileCheck %s
 
 --- |
   define void @shrink_test1(i32 %a) {

llvm/test/CodeGen/PowerPC/shrink-wrap.mir

@@ -1,10 +1,13 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
 # RUN: llc -verify-machineinstrs -mcpu=pwr9 -mtriple powerpc64le-unknown-linux-gnu \
 # RUN:   -run-pass=shrink-wrap -o - %s | FileCheck %s
+# RUN: llc -mcpu=pwr9 -mtriple powerpc64le-unknown-linux-gnu -passes='shrink-wrap' -o - %s | FileCheck %s
 # RUN: llc -verify-machineinstrs -mcpu=pwr9 -mtriple powerpc-ibm-aix-xcoff \
 # RUN:   -run-pass=shrink-wrap -mattr=-altivec -o - %s | FileCheck %s
+# RUN: llc -mcpu=pwr9 -mtriple powerpc-ibm-aix-xcoff -passes='shrink-wrap' -mattr=-altivec -o - %s | FileCheck %s
 # RUN: llc -verify-machineinstrs -mcpu=pwr9 -mtriple powerpc64-ibm-aix-xcoff \
 # RUN:   -run-pass=shrink-wrap -mattr=-altivec -o - %s | FileCheck %s
+# RUN: llc -mcpu=pwr9 -mtriple powerpc64-ibm-aix-xcoff -passes='shrink-wrap' -mattr=-altivec -o - %s | FileCheck %s
 --- |
   ; ModuleID = 'test.ll'
   source_filename = "test.ll"