[AMDGPU] Lazily init pal metadata on first function

[sebastian-ne]

Summary

Delay reading global metadata until the first function or the end of
the file is emitted. That way, earlier module passes can set metadata
that is emitted in the ELF.

emitStartOfAsmFile gets called when the passes are initialized,
which prevented earlier passes from changing the metadata.

This fixes issues encountered after converting
AMDGPUResourceUsageAnalysis to a Module pass in D117504.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

450–451

This feels like an abuse of the pass manager. You're not supposed to depend on passes on other functions from a function pass. What pass is setting this metadata? Is it a module pass?

[sebastian-ne]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

450–451

A module pass is setting the metadata.
emitStartOfAsmFile is called before the module pass is run (when AMDGPUResourceUsageAnalysis is a ModulePass, it works fine with the current main state for some reason).

[arsenm]

LGTM. Apparently emitStartOfAsmFile is called as part of AsmPrinter's doInitialization

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h

private:
 
   uint16_t getAmdhsaKernelCodeProperties(
       const MachineFunction &MF) const;
 
   amdhsa::kernel_descriptor_t getAmdhsaKernelDescriptor(
       const MachineFunction &MF,
       const SIProgramInfo &PI) const;
 
+  void initTargetStreamer(Module &M);
+
 public:
   explicit AMDGPUAsmPrinter(TargetMachine &TM,
                             std::unique_ptr<MCStreamer> Streamer);
 
   StringRef getPassName() const override;
 
   const MCSubtargetInfo* getGlobalSTI() const;
 
public:
   bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
                        const char *ExtraCode, raw_ostream &O) override;
 
 protected:
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   std::vector<std::string> DisasmLines, HexLines;
   size_t DisasmLineMaxLen;
+  bool IsTargetStreamerInitialized;
 };
 
 } // end namespace llvm
 
 #endif // LLVM_LIB_TARGET_AMDGPU_AMDGPUASMPRINTER_H

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

 
 AMDGPUTargetStreamer* AMDGPUAsmPrinter::getTargetStreamer() const {
   if (!OutStreamer)
     return nullptr;
   return static_cast<AMDGPUTargetStreamer*>(OutStreamer->getTargetStreamer());
 }
 
 void AMDGPUAsmPrinter::emitStartOfAsmFile(Module &M) {
+  IsTargetStreamerInitialized = false;
+}
+
+void AMDGPUAsmPrinter::initTargetStreamer(Module &M) {
+  IsTargetStreamerInitialized = true;
+
   // TODO: Which one is called first, emitStartOfAsmFile or
   // emitFunctionBodyStart?
   if (getTargetStreamer() && !getTargetStreamer()->getTargetID())
     initializeTargetID(M);
 
   if (TM.getTargetTriple().getOS() != Triple::AMDHSA &&
       TM.getTargetTriple().getOS() != Triple::AMDPAL)
     return;
void AMDGPUAsmPrinter::initTargetStreamer(Module &M) {
 
   // HSA and PAL emit NT_AMD_HSA_ISA_VERSION for code objects v2.
   IsaVersion Version = getIsaVersion(getGlobalSTI()->getCPU());
   getTargetStreamer()->EmitDirectiveHSACodeObjectISAV2(
       Version.Major, Version.Minor, Version.Stepping, "AMD", "AMDGPU");
 }
 
 void AMDGPUAsmPrinter::emitEndOfAsmFile(Module &M) {
+  // Init target streamer if it has not yet happened
+  if (!IsTargetStreamerInitialized)
+    initTargetStreamer(M);
+
   // Following code requires TargetStreamer to be present.
   if (!getTargetStreamer())
     return;
 
   if (TM.getTargetTriple().getOS() != Triple::AMDHSA ||
       isHsaAbiVersion2(getGlobalSTI()))
     getTargetStreamer()->EmitISAVersion();
 
amdhsa::kernel_descriptor_t AMDGPUAsmPrinter::getAmdhsaKernelDescriptor(
   if (STM.hasGFX90AInsts())
     KernelDescriptor.compute_pgm_rsrc3 =
       CurrentProgramInfo.ComputePGMRSrc3GFX90A;
 
   return KernelDescriptor;
 }
 
 bool AMDGPUAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
+  // Init target streamer lazily on the first function so that previous passes
+  // can set metadata.
+  if (!IsTargetStreamerInitialized)
+    initTargetStreamer(*MF.getFunction().getParent());
+
   ResourceUsage = &getAnalysis<AMDGPUResourceUsageAnalysis>();
   CurrentProgramInfo = SIProgramInfo();
 
   const AMDGPUMachineFunction *MFI = MF.getInfo<AMDGPUMachineFunction>();
 
   // The starting address of all shader programs must be 256 bytes aligned.
   // Regular functions just need the basic required instruction alignment.
   MF.setAlignment(MFI->isEntryFunction() ? Align(256) : Align(4));

llvm/unittests/CodeGen/AMDGPUMetadataTest.cpp

+//===- llvm/unittest/CodeGen/AMDGPUMetadataTest.cpp -----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Test that amdgpu metadata that is added in a pass is read by the asm emitter
+/// and stored in the ELF.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "gtest/gtest.h"
+
+namespace llvm {
+
+namespace {
+// Pass that adds global metadata
+struct AddMetadataPass : public ModulePass {
+  std::string PalMDString;
+
+public:
+  static char ID;
+  AddMetadataPass(std::string PalMDString)
+      : ModulePass(ID), PalMDString(PalMDString) {}
+  bool runOnModule(Module &M) override {
+    auto &Ctx = M.getContext();
+    auto *MD = M.getOrInsertNamedMetadata("amdgpu.pal.metadata.msgpack");
+    auto *PalMD = MDString::get(Ctx, PalMDString);
+    auto *TMD = MDTuple::get(Ctx, {PalMD});
+    MD->addOperand(TMD);
+    return true;
+  }
+};
+char AddMetadataPass::ID = 0;
+} // end anonymous namespace
+
+class AMDGPUSelectionDAGTest : public testing::Test {
+protected:
+  static void SetUpTestCase() {
+    InitializeAllTargets();
+    InitializeAllTargetMCs();
+  }
+
+  void SetUp() override {
+    std::string Error;
+    const Target *T = TargetRegistry::lookupTarget("amdgcn--amdpal", Error);
+    if (!T)
+      GTEST_SKIP();
+
+    TargetOptions Options;
+    TM = std::unique_ptr<LLVMTargetMachine>(
+        static_cast<LLVMTargetMachine *>(T->createTargetMachine(
+            "amdgcn--amdpal", "gfx1010", "", Options, None)));
+    if (!TM)
+      GTEST_SKIP();
+
+    LLVMContext Context;
+    std::unique_ptr<Module> M(new Module("TestModule", Context));
+    M->setDataLayout(TM->createDataLayout());
+
+    legacy::PassManager PM;
+    PM.add(new AddMetadataPass(PalMDString));
+    raw_svector_ostream OutStream(Elf);
+    if (TM->addPassesToEmitFile(PM, OutStream, nullptr,
+                                CodeGenFileType::CGFT_ObjectFile))
+      report_fatal_error("Target machine cannot emit a file of this type");
+
+    PM.run(*M);
+  }
+
+  static std::string PalMDString;
+
+  LLVMContext Context;
+  std::unique_ptr<LLVMTargetMachine> TM;
+  std::unique_ptr<Module> M;
+  SmallString<1024> Elf;
+};
+std::string AMDGPUSelectionDAGTest::PalMDString =
+    "\x81\xB0"
+    "amdpal.pipelines\x91\x81\xA4.api\xA6Vulkan";
+
+TEST_F(AMDGPUSelectionDAGTest, checkMetadata) {
+  // Check that the string is contained in the ELF
+  EXPECT_NE(Elf.find("Vulkan"), std::string::npos);
+}
+
+} // end namespace llvm

llvm/unittests/CodeGen/CMakeLists.txt

set(LLVM_LINK_COMPONENTS
   SelectionDAG
   Support
   Target
   )
 
 add_llvm_unittest(CodeGenTests
   AArch64SelectionDAGTest.cpp
   AllocationOrderTest.cpp
+  AMDGPUMetadataTest.cpp
   AsmPrinterDwarfTest.cpp
   DIEHashTest.cpp
   DIETest.cpp
   InstrRefLDVTest.cpp
   LowLevelTypeTest.cpp
   LexicalScopesTest.cpp
   MachineInstrBundleIteratorTest.cpp
   MachineInstrTest.cpp