[llvm] [StackFrameLayoutAnalysis] Add basic Scalable stack slot output (PR #99883)

[David Green via llvm-commits]

https://github.com/davemgreen created https://github.com/llvm/llvm-project/pull/99883

The existing StackFrameLayoutAnalysis details do not do well with Scalable vector stack slots, which are not marked as scalable and intertwined with the other fixed-size slots. This patch adds some very basic support, marking them as scalable and sorting them to the end of the list. The slot addresses are not really correct (for fixed as well as scalable), but this prints something a little better with the limited information curently available.

>From 152c3679c5c0c7adc2119da3422bec2123bf2fa7 Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Mon, 22 Jul 2024 15:29:18 +0100
Subject: [PATCH] [StackFrameLayoutAnalysis] Add basic Scalable stack slot
 output

The existing StackFrameLayoutAnalysis details do not do well with Scalable
vector stack slots, which are not marked as scalable and intertwined with the
other fixed-size slots. This patch adds some very basic support, marking them
as scalable and sorting them to the end of the list. The slot addresses are not
really correct (for fixed as well as scalable), but this prints something a
little better with the limited information curently available.
---
 .../lib/CodeGen/StackFrameLayoutAnalysisPass.cpp | 16 +++++++++++-----
 .../CodeGen/AArch64/sve-stack-frame-layout.ll    |  8 ++++----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/CodeGen/StackFrameLayoutAnalysisPass.cpp b/llvm/lib/CodeGen/StackFrameLayoutAnalysisPass.cpp
index 5d3903ed84ce8..940aecd1cb363 100644
--- a/llvm/lib/CodeGen/StackFrameLayoutAnalysisPass.cpp
+++ b/llvm/lib/CodeGen/StackFrameLayoutAnalysisPass.cpp
@@ -62,11 +62,14 @@ struct StackFrameLayoutAnalysisPass : public MachineFunctionPass {
     int Align;
     int Offset;
     SlotType SlotTy;
+    bool Scalable;
 
     SlotData(const MachineFrameInfo &MFI, const int ValOffset, const int Idx)
         : Slot(Idx), Size(MFI.getObjectSize(Idx)),
           Align(MFI.getObjectAlign(Idx).value()),
-          Offset(MFI.getObjectOffset(Idx) - ValOffset), SlotTy(Invalid) {
+          Offset(MFI.getObjectOffset(Idx) - ValOffset), SlotTy(Invalid),
+          Scalable(false) {
+      Scalable = MFI.getStackID(Idx) == TargetStackID::ScalableVector;
       if (MFI.isSpillSlotObjectIndex(Idx))
         SlotTy = SlotType::Spill;
       else if (Idx == MFI.getStackProtectorIndex())
@@ -75,9 +78,12 @@ struct StackFrameLayoutAnalysisPass : public MachineFunctionPass {
         SlotTy = SlotType::Variable;
     }
 
-    // we use this to sort in reverse order, so that the layout is displayed
-    // correctly
-    bool operator<(const SlotData &Rhs) const { return Offset > Rhs.Offset; }
+    // We use this to sort in reverse order, so that the layout is displayed
+    // correctly. Scalable slots are sorted to the end of the list.
+    bool operator<(const SlotData &Rhs) const {
+      return std::make_tuple(!Scalable, Offset) >
+             std::make_tuple(!Rhs.Scalable, Rhs.Offset);
+    }
   };
 
   StackFrameLayoutAnalysisPass() : MachineFunctionPass(ID) {}
@@ -153,7 +159,7 @@ struct StackFrameLayoutAnalysisPass : public MachineFunctionPass {
     Rem << Prefix << ore::NV("Offset", D.Offset)
         << "], Type: " << ore::NV("Type", getTypeString(D.SlotTy))
         << ", Align: " << ore::NV("Align", D.Align)
-        << ", Size: " << ore::NV("Size", D.Size);
+        << ", Size: " << ore::NV("Size", ElementCount::get(D.Size, D.Scalable));
   }
 
   void emitSourceLocRemark(const MachineFunction &MF, const DILocalVariable *N,
diff --git a/llvm/test/CodeGen/AArch64/sve-stack-frame-layout.ll b/llvm/test/CodeGen/AArch64/sve-stack-frame-layout.ll
index 67c1486b7a8c0..34d85d1f76086 100644
--- a/llvm/test/CodeGen/AArch64/sve-stack-frame-layout.ll
+++ b/llvm/test/CodeGen/AArch64/sve-stack-frame-layout.ll
@@ -4,10 +4,10 @@
 
 ; CHECK-FRAMELAYOUT-LABEL: Function: csr_d8_allocnxv4i32i32f64
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-8], Type: Spill, Align: 8, Size: 8
-; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Variable, Align: 16, Size: 16
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Spill, Align: 8, Size: 8
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-20], Type: Variable, Align: 4, Size: 4
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-32], Type: Variable, Align: 8, Size: 8
+; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Variable, Align: 16, Size: vscale x 16
 
 define i32 @csr_d8_allocnxv4i32i32f64(double %d) "aarch64_pstate_sm_compatible" {
 ; CHECK-LABEL: csr_d8_allocnxv4i32i32f64:
@@ -46,11 +46,11 @@ entry:
 
 ; CHECK-FRAMELAYOUT-LABEL: Function: csr_d8_allocnxv4i32i32f64_fp
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-8], Type: Spill, Align: 8, Size: 8
-; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Variable, Align: 16, Size: 16
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Spill, Align: 8, Size: 8
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-20], Type: Variable, Align: 4, Size: 4
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-32], Type: Spill, Align: 16, Size: 8
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-40], Type: Variable, Align: 8, Size: 8
+; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Variable, Align: 16, Size: vscale x 16
 
 define i32 @csr_d8_allocnxv4i32i32f64_fp(double %d) "aarch64_pstate_sm_compatible" "frame-pointer"="all" {
 ; CHECK-LABEL: csr_d8_allocnxv4i32i32f64_fp:
@@ -92,11 +92,11 @@ entry:
 
 ; CHECK-FRAMELAYOUT-LABEL: Function: svecc_z8_allocnxv4i32i32f64_fp
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-8], Type: Spill, Align: 8, Size: 8
-; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Spill, Align: 16, Size: 16
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Spill, Align: 8, Size: 8
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-20], Type: Variable, Align: 4, Size: 4
-; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-32], Type: Variable, Align: 16, Size: 16
 ; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-32], Type: Variable, Align: 8, Size: 8
+; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-16], Type: Spill, Align: 16, Size: vscale x 16
+; CHECK-FRAMELAYOUT-NEXT: Offset: [SP-32], Type: Variable, Align: 16, Size: vscale x 16
 
 define i32 @svecc_z8_allocnxv4i32i32f64_fp(double %d, <vscale x 4 x i32> %v) "aarch64_pstate_sm_compatible" "frame-pointer"="all" {
 ; CHECK-LABEL: svecc_z8_allocnxv4i32i32f64_fp: