[NewPM][CodeGen] Port localstackalloc to new pass manager (#94303)

There are two AArch64 tests use `-start-before` and `-print-after`. Rest tests uses `--passes` to test this pass.

Author: paperchalice

llvm/include/llvm/CodeGen/LocalStackSlotAllocation.h

@@ -0,0 +1,23 @@
+//===- LocalStackSlotAllocation.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_LOCALSTACKSLOTALLOCATION_H
+#define LLVM_CODEGEN_LOCALSTACKSLOTALLOCATION_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+class LocalStackSlotAllocationPass
+    : public PassInfoMixin<LocalStackSlotAllocationPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF, MachineFunctionAnalysisManager &);
+};
+
+} // namespace llvm
+#endif // LLVM_CODEGEN_LOCALSTACKSLOTALLOCATION_H

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -36,6 +36,7 @@
 #include "llvm/CodeGen/InterleavedAccess.h"
 #include "llvm/CodeGen/InterleavedLoadCombine.h"
 #include "llvm/CodeGen/JMCInstrumenter.h"
+#include "llvm/CodeGen/LocalStackSlotAllocation.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
 #include "llvm/CodeGen/MIRPrinter.h"
 #include "llvm/CodeGen/MachineFunctionAnalysis.h"
@@ -881,7 +882,7 @@ Error CodeGenPassBuilder<Derived, TargetMachineT>::addMachinePasses(
   } else {
     // If the target requests it, assign local variables to stack slots relative
     // to one another and simplify frame index references where possible.
-    addPass(LocalStackSlotPass());
+    addPass(LocalStackSlotAllocationPass());
   }
 
   if (TM.Options.EnableIPRA)
@@ -995,7 +996,7 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
 
   // If the target requests it, assign local variables to stack slots relative
   // to one another and simplify frame index references where possible.
-  addPass(LocalStackSlotPass());
+  addPass(LocalStackSlotAllocationPass());
 
   // With optimization, dead code should already be eliminated. However
   // there is one known exception: lowered code for arguments that are only

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -125,6 +125,7 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PI
 #endif
 MACHINE_FUNCTION_PASS("dead-mi-elimination", DeadMachineInstructionElimPass())
 MACHINE_FUNCTION_PASS("finalize-isel", FinalizeISelPass())
+MACHINE_FUNCTION_PASS("localstackalloc", LocalStackSlotAllocationPass())
 MACHINE_FUNCTION_PASS("no-op-machine-function", NoOpMachineFunctionPass())
 MACHINE_FUNCTION_PASS("print", PrintMIRPass())
 MACHINE_FUNCTION_PASS("require-all-machine-function-properties",
@@ -182,7 +183,6 @@ DUMMY_MACHINE_FUNCTION_PASS("kcfi", MachineKCFIPass)
 DUMMY_MACHINE_FUNCTION_PASS("legalizer", LegalizerPass)
 DUMMY_MACHINE_FUNCTION_PASS("livedebugvalues", LiveDebugValuesPass)
 DUMMY_MACHINE_FUNCTION_PASS("liveintervals", LiveIntervalsPass)
-DUMMY_MACHINE_FUNCTION_PASS("localstackalloc", LocalStackSlotPass)
 DUMMY_MACHINE_FUNCTION_PASS("lrshrink", LiveRangeShrinkPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-combiner", MachineCombinerPass)
 DUMMY_MACHINE_FUNCTION_PASS("machine-cp", MachineCopyPropagationPass)

llvm/lib/CodeGen/LocalStackSlotAllocation.cpp

@@ -13,6 +13,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/LocalStackSlotAllocation.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
@@ -71,7 +72,7 @@ namespace {
     int getFrameIndex() const { return FrameIdx; }
   };
 
-  class LocalStackSlotPass: public MachineFunctionPass {
+  class LocalStackSlotImpl {
     SmallVector<int64_t, 16> LocalOffsets;
 
     /// StackObjSet - A set of stack object indexes
@@ -86,14 +87,21 @@ namespace {
     void calculateFrameObjectOffsets(MachineFunction &Fn);
     bool insertFrameReferenceRegisters(MachineFunction &Fn);
 
+  public:
+    bool runOnMachineFunction(MachineFunction &MF);
+  };
+
+  class LocalStackSlotPass : public MachineFunctionPass {
   public:
     static char ID; // Pass identification, replacement for typeid
 
     explicit LocalStackSlotPass() : MachineFunctionPass(ID) {
       initializeLocalStackSlotPassPass(*PassRegistry::getPassRegistry());
     }
 
-    bool runOnMachineFunction(MachineFunction &MF) override;
+    bool runOnMachineFunction(MachineFunction &MF) override {
+      return LocalStackSlotImpl().runOnMachineFunction(MF);
+    }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.setPreservesCFG();
@@ -103,13 +111,24 @@ namespace {
 
 } // end anonymous namespace
 
+PreservedAnalyses
+LocalStackSlotAllocationPass::run(MachineFunction &MF,
+                                  MachineFunctionAnalysisManager &) {
+  bool Changed = LocalStackSlotImpl().runOnMachineFunction(MF);
+  if (!Changed)
+    return PreservedAnalyses::all();
+  auto PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}
+
 char LocalStackSlotPass::ID = 0;
 
 char &llvm::LocalStackSlotAllocationID = LocalStackSlotPass::ID;
 INITIALIZE_PASS(LocalStackSlotPass, DEBUG_TYPE,
                 "Local Stack Slot Allocation", false, false)
 
-bool LocalStackSlotPass::runOnMachineFunction(MachineFunction &MF) {
+bool LocalStackSlotImpl::runOnMachineFunction(MachineFunction &MF) {
   MachineFrameInfo &MFI = MF.getFrameInfo();
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   unsigned LocalObjectCount = MFI.getObjectIndexEnd();
@@ -139,7 +158,7 @@ bool LocalStackSlotPass::runOnMachineFunction(MachineFunction &MF) {
 }
 
 /// AdjustStackOffset - Helper function used to adjust the stack frame offset.
-void LocalStackSlotPass::AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx,
+void LocalStackSlotImpl::AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx,
                                            int64_t &Offset, bool StackGrowsDown,
                                            Align &MaxAlign) {
   // If the stack grows down, add the object size to find the lowest address.
@@ -171,7 +190,7 @@ void LocalStackSlotPass::AdjustStackOffset(MachineFrameInfo &MFI, int FrameIdx,
 
 /// AssignProtectedObjSet - Helper function to assign large stack objects (i.e.,
 /// those required to be close to the Stack Protector) to stack offsets.
-void LocalStackSlotPass::AssignProtectedObjSet(
+void LocalStackSlotImpl::AssignProtectedObjSet(
     const StackObjSet &UnassignedObjs, SmallSet<int, 16> &ProtectedObjs,
     MachineFrameInfo &MFI, bool StackGrowsDown, int64_t &Offset,
     Align &MaxAlign) {
@@ -183,7 +202,7 @@ void LocalStackSlotPass::AssignProtectedObjSet(
 
 /// calculateFrameObjectOffsets - Calculate actual frame offsets for all of the
 /// abstract stack objects.
-void LocalStackSlotPass::calculateFrameObjectOffsets(MachineFunction &Fn) {
+void LocalStackSlotImpl::calculateFrameObjectOffsets(MachineFunction &Fn) {
   // Loop over all of the stack objects, assigning sequential addresses...
   MachineFrameInfo &MFI = Fn.getFrameInfo();
   const TargetFrameLowering &TFI = *Fn.getSubtarget().getFrameLowering();
@@ -281,7 +300,7 @@ lookupCandidateBaseReg(unsigned BaseReg,
   return TRI->isFrameOffsetLegal(&MI, BaseReg, Offset);
 }
 
-bool LocalStackSlotPass::insertFrameReferenceRegisters(MachineFunction &Fn) {
+bool LocalStackSlotImpl::insertFrameReferenceRegisters(MachineFunction &Fn) {
   // Scan the function's instructions looking for frame index references.
   // For each, ask the target if it wants a virtual base register for it
   // based on what we can tell it about where the local will end up in the

llvm/lib/Passes/PassBuilder.cpp

@@ -90,6 +90,7 @@
 #include "llvm/CodeGen/InterleavedAccess.h"
 #include "llvm/CodeGen/InterleavedLoadCombine.h"
 #include "llvm/CodeGen/JMCInstrumenter.h"
+#include "llvm/CodeGen/LocalStackSlotAllocation.h"
 #include "llvm/CodeGen/LowerEmuTLS.h"
 #include "llvm/CodeGen/MIRPrinter.h"
 #include "llvm/CodeGen/MachineFunctionAnalysis.h"

llvm/test/CodeGen/AArch64/aarch64st1.mir

@@ -1,5 +1,6 @@
 # Check that it doesn't crash with unhandled opcode error, see pr52249
 # RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass localstackalloc -o - %s | FileCheck %s
+# RUN: llc -mtriple=aarch64-none-linux-gnu -passes=localstackalloc -o - %s | FileCheck %s
 --- |
 
   define void @test_st1_to_sp(<2 x i32> %a, <4 x i16> %b, <8 x i8> %c, <2 x i64> %d) gc "statepoint-example" { entry: ret void }

llvm/test/CodeGen/AArch64/sve-localstackalloc.mir

@@ -1,4 +1,5 @@
 # RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve -run-pass=localstackalloc -o - %s | FileCheck %s
+# RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve -passes=localstackalloc -o - %s | FileCheck %s
 
 --- |
   ; ModuleID = '<stdin>'