[CodeGen] Support custom format of stack maps

[cherry]

Summary

Add a hook to the GCMetadataPrinter for emitting stack maps in
custom format. The hook will be called at stack map generation
time. The default stack map format is used if there is no hook.

For this to be useful a few data structures and accessors are
exposed from the StackMaps class, so the custom printer can
access the stack map data.

[apilipenko]

include/llvm/CodeGen/StackMaps.h

267–281 ↗

(On Diff #171779)

The code move part of this change can be landed separately without further review so as to reduce the size of this diff.

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3005–3007

I'm not sure that early termination is the right thing to do. There might be more than one GC strategy in use in the current module, and all of them might want to emit stack maps.

[cherry]

Thanks for the review!

include/llvm/CodeGen/StackMaps.h

267–281 ↗

(On Diff #171779)

Split to https://reviews.llvm.org/D54224.

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3005–3007

In theory there could be more than one GC strategies, although I couldn't see how it can be used.

What about changing it to running all the GCMetadataPrinters, then return if any of the emitStackMaps returns true?

[reames]

A few drive by comments for the moment.

include/llvm/CodeGen/GCMetadataPrinter.h

68

This new call is arguably duplicating the existing finishModule callback. It'd be nice to find a way to merge the two honestly.

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3005–3007

The structure of the existing code is that a single model can contain functions compiled with two different gc strategies. I agree it's a bit odd, but it should continue to work.

I'd also like to move towards being able to emit *multiple* formats for a single GC. (Primarily for testing purposes.) That can be out of scope for your change if you desire.

[apilipenko]

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3005–3007

You can make the value returned from emitStackMaps indicate whether the default stackmap is needed. You would print the default format (SM.serializeToStackMapSection ) if at least one of the MPs requested default format serialization.

If a strategy doesn't have a MP than by default it requests default serialization printing.

[cherry]

include/llvm/CodeGen/GCMetadataPrinter.h

68

I agree that this would be nice. However, given the signature of finishAssembly, assuming we don't want to change that, there are some problems:

• we need to pass a StackMaps object. One possibility is to link it to the GCModuleInfo object.

• finishAssembly doesn't return anything, but we need a return value to decide whether to use the default format.

Any idea of how to solve these?

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3005–3007

Done as @apilipenko suggested.

For emitting multiple formats, the hook GCMetadataPrinter::emitStackMaps itself can do it. It can also call SM.serializeToStackMapSection to emit the default format (along with a custom format).

[reames]

LGTM

We can continue the discussion around finalizeModule and emitStackMaps separately. It's definitely not a blocker for this patch.

include/llvm/CodeGen/GCMetadataPrinter.h

68

In case my tone was unclear, this was not a must fix. It was a please think about it, and share any ideas you might have. If we can't find the right abstraction, I'm fine with this landing with the new API.

[cherry]

Thank you for the review!

include/llvm/CodeGen/AsmPrinter.h

 class MCSection;
 class MCStreamer;
 class MCSubtargetInfo;
 class MCSymbol;
 class MCTargetOptions;
 class MDNode;
 class Module;
 class raw_ostream;
+class StackMaps;
 class TargetLoweringObjectFile;
 class TargetMachine;
 
 /// This class is intended to be used as a driving class for all asm writers.
 class AsmPrinter : public MachineFunctionPass {
 public:
   /// Target machine description.
   TargetMachine &TM;
public:
   /// accesses to GOT entries.
   void computeGlobalGOTEquivs(Module &M);
 
   /// Constant expressions using GOT equivalent globals may not be
   /// eligible for PC relative GOT entry conversion, in such cases we need to
   /// emit the proxies we previously omitted in EmitGlobalVariable.
   void emitGlobalGOTEquivs();
 
+  /// Emit the stack maps.
+  void emitStackMaps(StackMaps &SM);
+
   //===------------------------------------------------------------------===//
   // Overridable Hooks
   //===------------------------------------------------------------------===//
 
   // Targets can, or in the case of EmitInstruction, must implement these to
   // customize output.
 
   /// This virtual method can be overridden by targets that want to emit

include/llvm/CodeGen/GCMetadataPrinter.h

 
 namespace llvm {
 
 class AsmPrinter;
 class GCMetadataPrinter;
 class GCModuleInfo;
 class GCStrategy;
 class Module;
+class StackMaps;
 
 /// GCMetadataPrinterRegistry - The GC assembly printer registry uses all the
 /// defaults from Registry.
 using GCMetadataPrinterRegistry = Registry<GCMetadataPrinter>;
 
 /// GCMetadataPrinter - Emits GC metadata as assembly code.  Instances are
 /// created, managed, and owned by the AsmPrinter.
 class GCMetadataPrinter {
public:
 
   /// Called before the assembly for the module is generated by
   /// the AsmPrinter (but after target specific hooks.)
   virtual void beginAssembly(Module &M, GCModuleInfo &Info, AsmPrinter &AP) {}
 
   /// Called after the assembly for the module is generated by
   /// the AsmPrinter (but before target specific hooks)
   virtual void finishAssembly(Module &M, GCModuleInfo &Info, AsmPrinter &AP) {}
+
+  /// Called when the stack maps are generated. Return true if
+  /// stack maps with a custom format are generated. Otherwise
+  /// returns false and the default format will be used.
+  virtual bool emitStackMaps(StackMaps &SM, AsmPrinter &AP) { return false; }
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_GCMETADATAPRINTER_H

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

 #include "llvm/CodeGen/MachineInstrBundle.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineModuleInfoImpls.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
+#include "llvm/CodeGen/StackMaps.h"
 #include "llvm/CodeGen/TargetFrameLowering.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Comdat.h"
isBlockOnlyReachableByFallthrough(const MachineBasicBlock *MBB) const {
 
   return true;
 }
 
 GCMetadataPrinter *AsmPrinter::GetOrCreateGCPrinter(GCStrategy &S) {
   if (!S.usesMetadata())
     return nullptr;
 
-  assert(!S.useStatepoints() && "statepoints do not currently support custom"
-         " stackmap formats, please see the documentation for a description of"
-         " the default format.  If you really need a custom serialized format,"
-         " please file a bug");
-
   gcp_map_type &GCMap = getGCMap(GCMetadataPrinters);
   gcp_map_type::iterator GCPI = GCMap.find(&S);
   if (GCPI != GCMap.end())
     return GCPI->second.get();
 
   auto Name = S.getName();
 
   for (GCMetadataPrinterRegistry::iterator
          I = GCMetadataPrinterRegistry::begin(),
          E = GCMetadataPrinterRegistry::end(); I != E; ++I)
     if (Name == I->getName()) {
       std::unique_ptr<GCMetadataPrinter> GMP = I->instantiate();
       GMP->S = &S;
       auto IterBool = GCMap.insert(std::make_pair(&S, std::move(GMP)));
       return IterBool.first->second.get();
     }
 
   report_fatal_error("no GCMetadataPrinter registered for GC: " + Twine(Name));
 }
 
+void AsmPrinter::emitStackMaps(StackMaps &SM) {
+  GCModuleInfo *MI = getAnalysisIfAvailable<GCModuleInfo>();
+  assert(MI && "AsmPrinter didn't require GCModuleInfo?");
+  for (auto &I : *MI)
+    if (GCMetadataPrinter *MP = GetOrCreateGCPrinter(*I))
+      if (MP->emitStackMaps(SM, *this))
+        // We've emitted a custom format, done.
+        return;
+
+  // Use the default serialization if no hooks.
+  SM.serializeToStackMapSection();
+}
+
 /// Pin vtable to this file.
 AsmPrinterHandler::~AsmPrinterHandler() = default;
 
 void AsmPrinterHandler::markFunctionEnd() {}
 
 // In the binary's "xray_instr_map" section, an array of these function entries
 // describes each instrumentation point.  When XRay patches your code, the index
 // into this table will be given to your handler as a patch point identifier.

lib/Target/AArch64/AArch64AsmPrinter.cpp

void AArch64AsmPrinter::EmitEndOfAsmFile(Module &M) {
   const Triple &TT = TM.getTargetTriple();
   if (TT.isOSBinFormatMachO()) {
     // Funny Darwin hack: This flag tells the linker that no global symbols
     // contain code that falls through to other global symbols (e.g. the obvious
     // implementation of multiple entry points).  If this doesn't occur, the
     // linker can safely perform dead code stripping.  Since LLVM never
     // generates code that does this, it is always safe to set.
     OutStreamer->EmitAssemblerFlag(MCAF_SubsectionsViaSymbols);
-    SM.serializeToStackMapSection();
+    emitStackMaps(SM);
   }
 }
 
 void AArch64AsmPrinter::EmitLOHs() {
   SmallVector<MCSymbol *, 3> MCArgs;
 
   for (const auto &D : AArch64FI->getLOHContainer()) {
     for (const MachineInstr *MI : D.getArgs()) {

lib/Target/PowerPC/PPCAsmPrinter.cpp

 MCSymbol *PPCAsmPrinter::lookUpOrCreateTOCEntry(MCSymbol *Sym) {
   MCSymbol *&TOCEntry = TOC[Sym];
   if (!TOCEntry)
     TOCEntry = createTempSymbol("C");
   return TOCEntry;
 }
 
 void PPCAsmPrinter::EmitEndOfAsmFile(Module &M) {
-  SM.serializeToStackMapSection();
+  emitStackMaps(SM);
 }
 
 void PPCAsmPrinter::LowerSTACKMAP(StackMaps &SM, const MachineInstr &MI) {
   unsigned NumNOPBytes = MI.getOperand(1).getImm();
 
   SM.recordStackMap(MI);
   assert(NumNOPBytes % 4 == 0 && "Invalid number of NOP bytes requested!");
 

lib/Target/SystemZ/SystemZAsmPrinter.cpp

bool SystemZAsmPrinter::PrintAsmMemoryOperand(const MachineInstr *MI,
                                               raw_ostream &OS) {
   SystemZInstPrinter::printAddress(MI->getOperand(OpNo).getReg(),
                                    MI->getOperand(OpNo + 1).getImm(),
                                    MI->getOperand(OpNo + 2).getReg(), OS);
   return false;
 }
 
 void SystemZAsmPrinter::EmitEndOfAsmFile(Module &M) {
-  SM.serializeToStackMapSection();
+  emitStackMaps(SM);
 }
 
 // Force static initialization.
 extern "C" void LLVMInitializeSystemZAsmPrinter() {
   RegisterAsmPrinter<SystemZAsmPrinter> X(getTheSystemZTarget());
 }

lib/Target/X86/X86AsmPrinter.cpp

void X86AsmPrinter::EmitEndOfAsmFile(Module &M) {
   const Triple &TT = TM.getTargetTriple();
 
   if (TT.isOSBinFormatMachO()) {
     // Mach-O uses non-lazy symbol stubs to encode per-TU information into
     // global table for symbol lookup.
     emitNonLazyStubs(MMI, *OutStreamer);
 
     // Emit stack and fault map information.
-    SM.serializeToStackMapSection();
+    emitStackMaps(SM);
     FM.serializeToFaultMapSection();
 
     // This flag tells the linker that no global symbols contain code that fall
     // through to other global symbols (e.g. an implementation of multiple entry
     // points). If this doesn't occur, the linker can safely perform dead code
     // stripping. Since LLVM never generates code that does this, it is always
     // safe to set.
     OutStreamer->EmitAssemblerFlag(MCAF_SubsectionsViaSymbols);
     return;
   }
 
   if (TT.isKnownWindowsMSVCEnvironment() && MMI->usesVAFloatArgument()) {
     StringRef SymbolName =
         (TT.getArch() == Triple::x86_64) ? "_fltused" : "__fltused";
     MCSymbol *S = MMI->getContext().getOrCreateSymbol(SymbolName);
     OutStreamer->EmitSymbolAttribute(S, MCSA_Global);
     return;
   }
 
   if (TT.isOSBinFormatCOFF()) {
-    SM.serializeToStackMapSection();
+    emitStackMaps(SM);
     return;
   }
 
   if (TT.isOSBinFormatELF()) {
-    SM.serializeToStackMapSection();
+    emitStackMaps(SM);
     FM.serializeToFaultMapSection();
     return;
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Target Registry Stuff
 //===----------------------------------------------------------------------===//
 
 // Force static initialization.
 extern "C" void LLVMInitializeX86AsmPrinter() {
   RegisterAsmPrinter<X86AsmPrinter> X(getTheX86_32Target());
   RegisterAsmPrinter<X86AsmPrinter> Y(getTheX86_64Target());
 }