[StackMaps] Force section serialization for functions with no records

[yrouban]

Summary

This patch allows to get function frame size from the .llvm_stackmaps section.
This section was generated for those functions that have at least one stack map record. For the other functions we could not get frame size. (The .eh_frame section does not help. Any other way?).

The StackMap class collected functions when stack map records were added. This patch introduces method StackMap::registerFunction() that is called before adding stack map records. All registered functions will have records with their frame sizes in the stack map section if the new option -stackmap-force-serialization is set to true even if there is no stack map records. With -stackmap-force-serialization set to false (the default) the old behavior is used: only those functions that have at least one record get into the stack map section and the section is generated only if there is at least one such function.

[yrouban]

fixed nits reported by build bots

[reames]

I don't really have an objection to this, but have you looked at the -stack-size-section support? I think you might be able to get the same information that way. If that doesn't work for some reason, just let me know and we can move forward with the review of this.

[yrouban]

In D103887#2814177, @reames wrote:

I don't really have an objection to this, but have you looked at the -stack-size-section support? I think you might be able to get the same information that way. If that doesn't work for some reason, just let me know and we can move forward with the review of this.

Thank you, Philip. The .stack_sizes section seems to be exactly what we need.

llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp

1661

must call the base PPCAsmPrinter::emitFunctionBodyStart(); to register the function in StackMaps.

llvm/test/CodeGen/AArch64/stackmap.ll

63

remove extra '_' in the name

[yrouban]

In D103887#2826669, @yrouban wrote:

Thank you, Philip. The .stack_sizes section seems to be exactly what we need.

One minor issue I had to solve downstream: the .stack_sizes section does not have the ELF::SHF_ALLOC flag as the llvm_stackmaps does. So I had to set this flag in MCObjectFileInfo.cpp.

llvm/include/llvm/CodeGen/StackMaps.h

public:
   }
 
   using LocationVec = SmallVector<Location, 8>;
   using LiveOutVec = SmallVector<LiveOutReg, 8>;
   using ConstantPool = MapVector<uint64_t, uint64_t>;
 
   struct FunctionInfo {
     uint64_t StackSize = 0;
-    uint64_t RecordCount = 1;
+    uint64_t RecordCount = 0;
 
     FunctionInfo() = default;
     explicit FunctionInfo(uint64_t StackSize) : StackSize(StackSize) {}
   };
 
   struct CallsiteInfo {
     const MCExpr *CSOffsetExpr = nullptr;
     uint64_t ID = 0;
     LocationVec Locations;
     LiveOutVec LiveOuts;
 
     CallsiteInfo() = default;
     CallsiteInfo(const MCExpr *CSOffsetExpr, uint64_t ID,
                  LocationVec &&Locations, LiveOutVec &&LiveOuts)
         : CSOffsetExpr(CSOffsetExpr), ID(ID), Locations(std::move(Locations)),
           LiveOuts(std::move(LiveOuts)) {}
   };
 
   using FnInfoMap = MapVector<const MCSymbol *, FunctionInfo>;
   using CallsiteInfoList = std::vector<CallsiteInfo>;
 
+  /// Register function. Must be called before the other record methods for the
+  /// current function.
+  void recordFunction();
+
   /// Generate a stackmap record for a stackmap instruction.
   ///
   /// MI must be a raw STACKMAP, not a PATCHPOINT.
   void recordStackMap(const MCSymbol &L,
                       const MachineInstr &MI);
 
   /// Generate a stackmap record for a patchpoint instruction.
   void recordPatchPoint(const MCSymbol &L,

llvm/lib/CodeGen/StackMaps.cpp

 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <iterator>
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "stackmaps"
 
 static cl::opt<int> StackMapVersion(
     "stackmap-version", cl::init(3), cl::Hidden,
     cl::desc("Specify the stackmap encoding version (default = 3)"));
 
+static cl::opt<bool> ForceSerializationIfEmpty(
+    "stackmap-force-serialization", cl::init(false), cl::Hidden,
+    cl::desc("Forcibly serializes stack map section even if empty."));
+
 const char *StackMaps::WSMP = "Stack Maps: ";
 
 static uint64_t getConstMetaVal(const MachineInstr &MI, unsigned Idx) {
   assert(MI.getOperand(Idx).isImm() &&
          MI.getOperand(Idx).getImm() == StackMaps::ConstantOp);
   const auto &MO = MI.getOperand(Idx + 1);
   assert(MO.isImm());
   return MO.getImm();
void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
   // entry.
   const MCExpr *CSOffsetExpr = MCBinaryExpr::createSub(
       MCSymbolRefExpr::create(&MILabel, OutContext),
       MCSymbolRefExpr::create(AP.CurrentFnSymForSize, OutContext), OutContext);
 
   CSInfos.emplace_back(CSOffsetExpr, ID, std::move(Locations),
                        std::move(LiveOuts));
 
-  // Record the stack size of the current function and update callsite count.
+  // Update callsite count.
+  auto CurrentIt = FnInfos.find(AP.CurrentFnSym);
+  assert(CurrentIt != FnInfos.end() &&
+         "Current function must have been recorded");
+  CurrentIt->second.RecordCount++;
+}
+
+void StackMaps::recordFunction() {
+  // Record the stack size of the current function.
+  // Safely avoid doing it twice. It is parser's responsibility to report the
+  // redefinition error.
+  if (FnInfos.find(AP.CurrentFnSym) != FnInfos.end())
+    return;
+
   const MachineFrameInfo &MFI = AP.MF->getFrameInfo();
   const TargetRegisterInfo *RegInfo = AP.MF->getSubtarget().getRegisterInfo();
   bool HasDynamicFrameSize =
       MFI.hasVarSizedObjects() || RegInfo->hasStackRealignment(*(AP.MF));
   uint64_t FrameSize = HasDynamicFrameSize ? UINT64_MAX : MFI.getStackSize();
-
-  auto CurrentIt = FnInfos.find(AP.CurrentFnSym);
-  if (CurrentIt != FnInfos.end())
-    CurrentIt->second.RecordCount++;
-  else
-    FnInfos.insert(std::make_pair(AP.CurrentFnSym, FunctionInfo(FrameSize)));
 }
 
 void StackMaps::recordStackMap(const MCSymbol &L, const MachineInstr &MI) {
   assert(MI.getOpcode() == TargetOpcode::STACKMAP && "expected stackmap");
 
   StackMapOpers opers(&MI);
   const int64_t ID = MI.getOperand(PatchPointOpers::IDPos).getImm();
   recordStackMapOpers(L, MI, ID, std::next(MI.operands_begin(),

 /// uint32 : NumRecords
 void StackMaps::emitStackmapHeader(MCStreamer &OS) {
   // Header.
   OS.emitIntValue(StackMapVersion, 1); // Version.
   OS.emitIntValue(0, 1);               // Reserved.
   OS.emitInt16(0);                     // Reserved.
 
   // Num functions.
-  LLVM_DEBUG(dbgs() << WSMP << "#functions = " << FnInfos.size() << '\n');
-  OS.emitInt32(FnInfos.size());
+  // Count functions that will be emitted.
+  uint64_t NumOfFunctions =
+      ForceSerializationIfEmpty
+          ? FnInfos.size()
+          : std::count_if(FnInfos.begin(), FnInfos.end(),
+                          [](const FnInfoMap::value_type &FnInfo) {
+                            return FnInfo.second.RecordCount != 0;
+                          });
+
+  LLVM_DEBUG(dbgs() << WSMP << "#functions = " << NumOfFunctions << '\n');
+  OS.emitInt32(NumOfFunctions);
   // Num constants.
   LLVM_DEBUG(dbgs() << WSMP << "#constants = " << ConstPool.size() << '\n');
   OS.emitInt32(ConstPool.size());
   // Num callsites.
   LLVM_DEBUG(dbgs() << WSMP << "#callsites = " << CSInfos.size() << '\n');
   OS.emitInt32(CSInfos.size());
 }
 
 /// Emit the function frame record for each function.
 ///
 /// StkSizeRecord[NumFunctions] {
 ///   uint64 : Function Address
 ///   uint64 : Stack Size
 ///   uint64 : Record Count
 /// }
 void StackMaps::emitFunctionFrameRecords(MCStreamer &OS) {
   // Function Frame records.
   LLVM_DEBUG(dbgs() << WSMP << "functions:\n");
   for (auto const &FR : FnInfos) {
+    if (FR.second.RecordCount == 0 && !ForceSerializationIfEmpty)
+      continue;
     LLVM_DEBUG(dbgs() << WSMP << "function addr: " << FR.first
                       << " frame size: " << FR.second.StackSize
                       << " callsite count: " << FR.second.RecordCount << '\n');
+    assert(OS.getContext().getAsmInfo()->getCodePointerSize() == 8 &&
+           'StackMap supports only 64-bit platforms');
     OS.emitSymbolValue(FR.first, 8);
     OS.emitIntValue(FR.second.StackSize, 8);
     OS.emitIntValue(FR.second.RecordCount, 8);
   }
 }
 
 /// Emit the constant pool.
 ///

 }
 
 /// Serialize the stackmap data.
 void StackMaps::serializeToStackMapSection() {
   (void)WSMP;
   // Bail out if there's no stack map data.
   assert((!CSInfos.empty() || ConstPool.empty()) &&
          "Expected empty constant pool too!");
-  assert((!CSInfos.empty() || FnInfos.empty()) &&
-         "Expected empty function record too!");
-  if (CSInfos.empty())
+  if (CSInfos.empty() && !ForceSerializationIfEmpty)
     return;
 
   MCContext &OutContext = AP.OutStreamer->getContext();
   MCStreamer &OS = *AP.OutStreamer;
 
   // Create the section.
   MCSection *StackMapSection =
       OutContext.getObjectFileInfo()->getStackMapSection();

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

private:
 
   bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNum,
                        const char *ExtraCode, raw_ostream &O) override;
   bool PrintAsmMemoryOperand(const MachineInstr *MI, unsigned OpNum,
                              const char *ExtraCode, raw_ostream &O) override;
 
   void PrintDebugValueComment(const MachineInstr *MI, raw_ostream &OS);
 
+  void emitFunctionBodyStart() override { SM.recordFunction(); }
   void emitFunctionBodyEnd() override;
 
   MCSymbol *GetCPISymbol(unsigned CPID) const override;
   void emitEndOfAsmFile(Module &M) override;
 
   AArch64FunctionInfo *AArch64FI = nullptr;
 
   /// Emit the LOHs contained in AArch64FI.

llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp

MCSymbol *lookUpOrCreateTOCEntry(const MCSymbol *Sym,
                                        MCSymbolRefExpr::VariantKind::VK_None);
 
   bool doInitialization(Module &M) override {
     if (!TOC.empty())
       TOC.clear();
     return AsmPrinter::doInitialization(M);
   }
 
+  void emitFunctionBodyStart() override { SM.recordFunction(); }
   void emitInstruction(const MachineInstr *MI) override;
 
   /// This function is for PrintAsmOperand and PrintAsmMemoryOperand,
   /// invoked by EmitMSInlineAsmStr and EmitGCCInlineAsmStr only.
   /// The \p MI would be INLINEASM ONLY.
   void printOperand(const MachineInstr *MI, unsigned OpNo, raw_ostream &O);
 
   void PrintSymbolOperand(const MachineOperand &MO, raw_ostream &O) override;
if (!TOC.empty()) {
     }
   }
 
   PPCAsmPrinter::emitEndOfAsmFile(M);
 }
 
 /// EmitFunctionBodyStart - Emit a global entry point prefix for ELFv2.
 void PPCLinuxAsmPrinter::emitFunctionBodyStart() {
+  PPCAsmPrinter::emitFunctionBodyStart();
+
   // In the ELFv2 ABI, in functions that use the TOC register, we need to
   // provide two entry points.  The ABI guarantees that when calling the
   // local entry point, r2 is set up by the caller to contain the TOC base
   // for this function, and when calling the global entry point, r12 is set
   // up by the caller to hold the address of the global entry point.  We
   // thus emit a prefix sequence along the following lines:
   //
   // func:
   // .Lfunc_gepNN:

llvm/lib/Target/SystemZ/SystemZAsmPrinter.h

 private:
   StackMaps SM;
 
 public:
   SystemZAsmPrinter(TargetMachine &TM, std::unique_ptr<MCStreamer> Streamer)
       : AsmPrinter(TM, std::move(Streamer)), SM(*this) {}
 
   // Override AsmPrinter.
+  void emitFunctionBodyStart() override { SM.recordFunction(); }
   StringRef getPassName() const override { return "SystemZ Assembly Printer"; }
   void emitInstruction(const MachineInstr *MI) override;
   void emitMachineConstantPoolValue(MachineConstantPoolValue *MCPV) override;
   void emitEndOfAsmFile(Module &M) override;
   bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
                        const char *ExtraCode, raw_ostream &OS) override;
   bool PrintAsmMemoryOperand(const MachineInstr *MI, unsigned OpNo,
                              const char *ExtraCode, raw_ostream &OS) override;

llvm/lib/Target/X86/X86AsmPrinter.cpp

bool X86AsmPrinter::runOnMachineFunction(MachineFunction &MF) {
 
   EmitFPOData = false;
 
   // We didn't modify anything.
   return false;
 }
 
 void X86AsmPrinter::emitFunctionBodyStart() {
+  SM.recordFunction();
   if (EmitFPOData) {
     if (auto *XTS =
         static_cast<X86TargetStreamer *>(OutStreamer->getTargetStreamer()))
       XTS->emitFPOProc(
           CurrentFnSym,
           MF->getInfo<X86MachineFunctionInfo>()->getArgumentStackSize());
   }
 }

llvm/test/CodeGen/AArch64/stackmap.ll

-; RUN: llc < %s -mtriple=arm64-linux-gnu | FileCheck %s
+; RUN: llc < %s -mtriple=arm64-linux-gnu -stackmap-force-serialization=0 | FileCheck -check-prefixes=CHECK,CHECK-FORCED0 %s
+; RUN: llc < %s -mtriple=arm64-linux-gnu -stackmap-force-serialization=1 | FileCheck -check-prefixes=CHECK,CHECK-FORCED1 %s
 ;
 ; Note: Print verbose stackmaps using -debug-only=stackmaps.
 
 ; CHECK-LABEL:  .section  .llvm_stackmaps
 ; CHECK-NEXT:  __LLVM_StackMaps:
 ; Header
 ; CHECK-NEXT:   .byte 3
 ; CHECK-NEXT:   .byte 0
 ; CHECK-NEXT:   .hword 0
 ; Num Functions
-; CHECK-NEXT:   .word 14
+; CHECK-FORCED0-NEXT: .word 14
+; CHECK-FORCED1-NEXT: .word 15
 ; Num LargeConstants
 ; CHECK-NEXT:   .word 3
 ; Num Callsites
 ; CHECK-NEXT:   .word 18
 
 ; Functions and stack size
 ; CHECK-NEXT:   .xword constantargs
 ; CHECK-NEXT:   .xword 16

 ; CHECK-NEXT:   .xword 16
 ; CHECK-NEXT:   .xword 4
 ; CHECK-NEXT:   .xword clobberLR
 ; CHECK-NEXT:   .xword 112
 ; CHECK-NEXT:   .xword 1
 ; CHECK-NEXT:   .xword needsStackRealignment
 ; CHECK-NEXT:   .xword -1
 ; CHECK-NEXT:   .xword 1
+; CHECK-FORCED1-NEXT:   .xword noStackMapEntries
+; CHECK-FORCED1-NEXT:   .xword 0
+; CHECK-FORCED1-NEXT:   .xword 0
 
 ; Large Constants
 ; CHECK-NEXT:   .xword   2147483648
 ; CHECK-NEXT:   .xword   4294967295
 ; CHECK-NEXT:   .xword   4294967296
 
 ; Callsites
 ; Constant arguments
define void @needsStackRealignment() {
   %val = alloca i64, i32 3, align 128
   tail call void (...) @escape_values(i64* %val)
 ; Note: Adding any non-constant to the stackmap would fail because we
 ; expected to be able to address off the frame pointer.  In a realigned
 ; frame, we must use the stack pointer instead.  This is a separate bug.
   tail call void (i64, i32, ...) @llvm.experimental.stackmap(i64 0, i32 0)
   ret void
 }
+
+define void @noStackMapEntries() {
+  unreachable
+}
 declare void @escape_values(...)
 
 declare void @llvm.experimental.stackmap(i64, i32, ...)
 declare void @llvm.experimental.patchpoint.void(i64, i32, i8*, i32, ...)
 declare i64 @llvm.experimental.patchpoint.i64(i64, i32, i8*, i32, ...)

llvm/test/CodeGen/SystemZ/stackmap.ll

-; RUN: llc < %s -mtriple=s390x-linux-gnu | FileCheck %s
+; RUN: llc < %s -mtriple=s390x-linux-gnu -stackmap-force-serialization=0 | FileCheck -check-prefixes=CHECK,CHECK-FORCED0 %s
+; RUN: llc < %s -mtriple=s390x-linux-gnu -stackmap-force-serialization=1 | FileCheck -check-prefixes=CHECK,CHECK-FORCED1 %s
 ;
 ; Note: Print verbose stackmaps using -debug-only=stackmaps.
 
 ; CHECK:       .section .llvm_stackmaps
 ; CHECK-NEXT:  __LLVM_StackMaps:
 ; Header
 ; CHECK-NEXT:   .byte 3
 ; CHECK-NEXT:   .byte 0
 ; CHECK-NEXT:   .short 0
 ; Num Functions
-; CHECK-NEXT:   .long 15
+; CHECK-FORCED0-NEXT: .long 15
+; CHECK-FORCED1-NEXT: .long 16
 ; Num LargeConstants
 ; CHECK-NEXT:   .long 3
 ; Num Callsites
 ; CHECK-NEXT:   .long 19
 
 ; Functions and stack size
 ; CHECK-NEXT:   .quad constantargs
 ; CHECK-NEXT:   .quad 160

 ; CHECK-NEXT:   .quad 160
 ; CHECK-NEXT:   .quad 4
 ; CHECK-NEXT:   .quad clobberScratch
 ; CHECK-NEXT:   .quad 168
 ; CHECK-NEXT:   .quad 1
 ; CHECK-NEXT:   .quad needsStackRealignment
 ; CHECK-NEXT:   .quad -1
 ; CHECK-NEXT:   .quad 1
+; CHECK-FORCED1-NEXT:   .quad noStackMapEntries
+; CHECK-FORCED1-NEXT:   .quad 0
+; CHECK-FORCED1-NEXT:   .quad 0
 
 ; Large Constants
 ; CHECK-NEXT:   .quad   2147483648
 ; CHECK-NEXT:   .quad   4294967295
 ; CHECK-NEXT:   .quad   4294967296
 
 ; Callsites
 ; Constant arguments
define void @needsStackRealignment() {
   %val = alloca i64, i32 3, align 128
   tail call void (...) @escape_values(i64* %val)
 ; Note: Adding any non-constant to the stackmap would fail because we
 ; expected to be able to address off the frame pointer.  In a realigned
 ; frame, we must use the stack pointer instead.  This is a separate bug.
   tail call void (i64, i32, ...) @llvm.experimental.stackmap(i64 0, i32 0)
   ret void
 }
+
+define void @noStackMapEntries() {
+  unreachable
+}
 declare void @escape_values(...)
 
 declare void @llvm.experimental.stackmap(i64, i32, ...)
 declare void @llvm.experimental.patchpoint.void(i64, i32, i8*, i32, ...)
 declare i64 @llvm.experimental.patchpoint.i64(i64, i32, i8*, i32, ...)

llvm/test/CodeGen/X86/stackmap.ll

-; RUN: llc < %s -mtriple=x86_64-apple-darwin -mcpu=corei7 | FileCheck %s
+; RUN: llc < %s -mtriple=x86_64-apple-darwin -mcpu=corei7 -stackmap-force-serialization=0 | FileCheck -check-prefixes=CHECK,CHECK-FORCED0 %s
+; RUN: llc < %s -mtriple=x86_64-apple-darwin -mcpu=corei7 -stackmap-force-serialization=1 | FileCheck -check-prefixes=CHECK,CHECK-FORCED1 %s
 ;
 ; Note: Print verbose stackmaps using -debug-only=stackmaps.
 
 ; CHECK-LABEL:  .section  __LLVM_STACKMAPS,__llvm_stackmaps
 ; CHECK-NEXT:  __LLVM_StackMaps:
 ; Header
 ; CHECK-NEXT:   .byte 3
 ; CHECK-NEXT:   .byte 0
 ; CHECK-NEXT:   .short 0
 ; Num Functions
-; CHECK-NEXT:   .long 16
+; CHECK-FORCED0-NEXT: .long 16
+; CHECK-FORCED1-NEXT: .long 17
 ; Num LargeConstants
 ; CHECK-NEXT:   .long 3
 ; Num Callsites
 ; CHECK-NEXT:   .long 20
 
 ; Functions and stack size
 ; CHECK-NEXT:   .quad _constantargs
 ; CHECK-NEXT:   .quad 8

 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad 4
 ; CHECK-NEXT:   .quad _clobberScratch
 ; CHECK-NEXT:   .quad 56
 ; CHECK-NEXT:   .quad 1
 ; CHECK-NEXT:   .quad _needsStackRealignment
 ; CHECK-NEXT:   .quad -1
 ; CHECK-NEXT:   .quad 1
+; CHECK-FORCED1-NEXT:   .quad _noStackMapEntries
+; CHECK-FORCED1-NEXT:   .quad 0
+; CHECK-FORCED1-NEXT:   .quad 0
 
 ; Large Constants
 ; CHECK-NEXT:   .quad   2147483648
 ; CHECK-NEXT:   .quad   4294967295
 ; CHECK-NEXT:   .quad   4294967296
 
 ; Callsites
 ; Constant arguments
define void @needsStackRealignment() {
   %val = alloca i64, i32 3, align 128
   tail call void (...) @escape_values(i64* %val)
 ; Note: Adding any non-constant to the stackmap would fail because we
 ; expected to be able to address off the frame pointer.  In a realigned
 ; frame, we must use the stack pointer instead.  This is a separate bug.
   tail call void (i64, i32, ...) @llvm.experimental.stackmap(i64 0, i32 0)
   ret void
 }
+
+define void @noStackMapEntries() {
+  unreachable
+}
 declare void @escape_values(...)
 
 declare void @llvm.experimental.stackmap(i64, i32, ...)
 declare void @llvm.experimental.patchpoint.void(i64, i32, i8*, i32, ...)
 declare i64 @llvm.experimental.patchpoint.i64(i64, i32, i8*, i32, ...)