[llvm] Check for all frame instructions in finalize isel. (PR #85945)

[Jonas Paulsson via llvm-commits]

https://github.com/JonPsson1 updated https://github.com/llvm/llvm-project/pull/85945

>From c9b89f76d2ef9d32f6fc73f38e3b9b61a8dc88df Mon Sep 17 00:00:00 2001
From: Jonas Paulsson <paulson1 at linux.ibm.com>
Date: Wed, 20 Mar 2024 11:03:39 -0400
Subject: [PATCH] Check for all frame instructions in finalize isel.

---
 llvm/lib/CodeGen/FinalizeISel.cpp           |  3 +--
 llvm/lib/CodeGen/PrologEpilogInserter.cpp   |  2 ++
 llvm/test/CodeGen/SystemZ/frame-adjstack.ll | 16 ++++++++++++++++
 3 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/SystemZ/frame-adjstack.ll

diff --git a/llvm/lib/CodeGen/FinalizeISel.cpp b/llvm/lib/CodeGen/FinalizeISel.cpp
index 978355f8eb1bbf..bf967eac22f177 100644
--- a/llvm/lib/CodeGen/FinalizeISel.cpp
+++ b/llvm/lib/CodeGen/FinalizeISel.cpp
@@ -59,8 +59,7 @@ bool FinalizeISel::runOnMachineFunction(MachineFunction &MF) {
 
       // Set AdjustsStack to true if the instruction selector emits a stack
       // frame setup instruction or a stack aligning inlineasm.
-      if (MI.getOpcode() == TII->getCallFrameSetupOpcode() ||
-          MI.isStackAligningInlineAsm())
+      if (TII->isFrameInstr(MI) || MI.isStackAligningInlineAsm())
         MF.getFrameInfo().setAdjustsStack(true);
 
       // If MI is a pseudo, expand it.
diff --git a/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/llvm/lib/CodeGen/PrologEpilogInserter.cpp
index c942b8a3e26880..eaf96ec5cbde8c 100644
--- a/llvm/lib/CodeGen/PrologEpilogInserter.cpp
+++ b/llvm/lib/CodeGen/PrologEpilogInserter.cpp
@@ -372,6 +372,8 @@ void PEI::calculateCallFrameInfo(MachineFunction &MF) {
   MFI.computeMaxCallFrameSize(MF, &FrameSDOps);
   assert(MFI.getMaxCallFrameSize() <= MaxCFSIn &&
          "Recomputing MaxCFS gave a larger value.");
+  assert((FrameSDOps.empty() || MF.getFrameInfo().adjustsStack()) &&
+         "AdjustsStack not set in presence of a frame pseudo instruction.");
 
   if (TFI->canSimplifyCallFramePseudos(MF)) {
     // If call frames are not being included as part of the stack frame, and
diff --git a/llvm/test/CodeGen/SystemZ/frame-adjstack.ll b/llvm/test/CodeGen/SystemZ/frame-adjstack.ll
new file mode 100644
index 00000000000000..7edacaa3d7d714
--- /dev/null
+++ b/llvm/test/CodeGen/SystemZ/frame-adjstack.ll
@@ -0,0 +1,16 @@
+; RUN: llc < %s -mtriple=s390x-linux-gnu -verify-machineinstrs | FileCheck %s
+;
+; Test that inserting a new MBB near a call during finalize isel custom
+; insertion does not cause all frame instructions to be missed. That would
+; result in a missing to set the AdjustsStack flag.
+
+; CHECK-LABEL: fun
+define void @fun(i1 %cc) {
+  %sel = select i1 %cc, i32 5, i32 0
+  tail call void @input_report_abs(i32 %sel)
+  %sel2 = select i1 %cc, i32 6, i32 1
+  tail call void @input_report_abs(i32 %sel2)
+  ret void
+}
+
+declare void @input_report_abs(i32)