[llvm] [CodeGen] Don't check attrs for stack realign (PR #92564)

[via llvm-commits]

https://github.com/aengelke created https://github.com/llvm/llvm-project/pull/92564

shouldRealignStack/canRealignStack are repeatedly called in PEI (through hasStackRealignment). Checking function attributes is expensive, so cache this data in the MachineFrameInfo, which had most data already.

(Draft PR for now, I want to see if the CI passes first.)

>From 59d4a5aecb1bd0de14802535d4f04ab53cbc23f8 Mon Sep 17 00:00:00 2001
From: Alexis Engelke <engelke at in.tum.de>
Date: Fri, 17 May 2024 17:44:45 +0200
Subject: [PATCH] [CodeGen] Don't check attrs for stack realign

shouldRealignStack/canRealignStack are repeatedly called in PEI (through
hasStackRealignment). Checking function attributes is expensive, so
cache this data in the MachineFrameInfo, which had most data already.
---
 llvm/include/llvm/CodeGen/MachineFrameInfo.h | 8 ++++++++
 llvm/lib/CodeGen/MachineFunction.cpp         | 5 +++--
 llvm/lib/CodeGen/TargetRegisterInfo.cpp      | 9 ++-------
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/MachineFrameInfo.h b/llvm/include/llvm/CodeGen/MachineFrameInfo.h
index a2c78e9e093d0..466fed7fb3a29 100644
--- a/llvm/include/llvm/CodeGen/MachineFrameInfo.h
+++ b/llvm/include/llvm/CodeGen/MachineFrameInfo.h
@@ -346,6 +346,8 @@ class MachineFrameInfo {
 
   MachineFrameInfo(const MachineFrameInfo &) = delete;
 
+  bool isStackRealignable() const { return StackRealignable; }
+
   /// Return true if there are any stack objects in this function.
   bool hasStackObjects() const { return !Objects.empty(); }
 
@@ -603,6 +605,12 @@ class MachineFrameInfo {
   /// Make sure the function is at least Align bytes aligned.
   void ensureMaxAlignment(Align Alignment);
 
+  /// Return true if stack realignment is forced by function attributes or if
+  /// the stack alignment.
+  bool shouldRealignStack() const {
+    return ForcedRealign || MaxAlignment > StackAlignment;
+  }
+
   /// Return true if this function adjusts the stack -- e.g.,
   /// when calling another function. This is only valid during and after
   /// prolog/epilog code insertion.
diff --git a/llvm/lib/CodeGen/MachineFunction.cpp b/llvm/lib/CodeGen/MachineFunction.cpp
index 8366ad2859069..4182e75354125 100644
--- a/llvm/lib/CodeGen/MachineFunction.cpp
+++ b/llvm/lib/CodeGen/MachineFunction.cpp
@@ -200,10 +200,11 @@ void MachineFunction::init() {
   // explicitly asked us not to.
   bool CanRealignSP = STI->getFrameLowering()->isStackRealignable() &&
                       !F.hasFnAttribute("no-realign-stack");
+  bool ForceRealignSP = F.hasFnAttribute(Attribute::StackAlignment) ||
+                        F.hasFnAttribute("stackrealign");
   FrameInfo = new (Allocator) MachineFrameInfo(
       getFnStackAlignment(STI, F), /*StackRealignable=*/CanRealignSP,
-      /*ForcedRealign=*/CanRealignSP &&
-          F.hasFnAttribute(Attribute::StackAlignment));
+      /*ForcedRealign=*/ForceRealignSP && CanRealignSP);
 
   setUnsafeStackSize(F, *FrameInfo);
 
diff --git a/llvm/lib/CodeGen/TargetRegisterInfo.cpp b/llvm/lib/CodeGen/TargetRegisterInfo.cpp
index 4e06393f4cc1d..ffc8055dd27e8 100644
--- a/llvm/lib/CodeGen/TargetRegisterInfo.cpp
+++ b/llvm/lib/CodeGen/TargetRegisterInfo.cpp
@@ -474,16 +474,11 @@ bool TargetRegisterInfo::isCalleeSavedPhysReg(
 }
 
 bool TargetRegisterInfo::canRealignStack(const MachineFunction &MF) const {
-  return !MF.getFunction().hasFnAttribute("no-realign-stack");
+  return MF.getFrameInfo().isStackRealignable();
 }
 
 bool TargetRegisterInfo::shouldRealignStack(const MachineFunction &MF) const {
-  const MachineFrameInfo &MFI = MF.getFrameInfo();
-  const TargetFrameLowering *TFI = MF.getSubtarget().getFrameLowering();
-  const Function &F = MF.getFunction();
-  return F.hasFnAttribute("stackrealign") ||
-         (MFI.getMaxAlign() > TFI->getStackAlign()) ||
-         F.hasFnAttribute(Attribute::StackAlignment);
+  return MF.getFrameInfo().shouldRealignStack();
 }
 
 bool TargetRegisterInfo::regmaskSubsetEqual(const uint32_t *mask0,