[AMDGPU] Lazily init pal metadata on first function

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.

Differential Revision: https://reviews.llvm.org/D118492

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

@@ -111,6 +111,12 @@ AMDGPUTargetStreamer* AMDGPUAsmPrinter::getTargetStreamer() const {
 }
 
 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())
@@ -143,6 +149,10 @@ void AMDGPUAsmPrinter::emitStartOfAsmFile(Module &M) {
 }
 
 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;
@@ -437,6 +447,11 @@ amdhsa::kernel_descriptor_t AMDGPUAsmPrinter::getAmdhsaKernelDescriptor(
 }
 
 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();
 

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h

@@ -77,6 +77,8 @@ private:
       const MachineFunction &MF,
       const SIProgramInfo &PI) const;
 
+  void initTargetStreamer(Module &M);
+
 public:
   explicit AMDGPUAsmPrinter(TargetMachine &TM,
                             std::unique_ptr<MCStreamer> Streamer);
@@ -132,6 +134,7 @@ protected:
 
   std::vector<std::string> DisasmLines, HexLines;
   size_t DisasmLineMaxLen;
+  bool IsTargetStreamerInitialized;
 };
 
 } // end namespace llvm

llvm/unittests/CodeGen/AMDGPUMetadataTest.cpp

@@ -0,0 +1,94 @@
+//===- 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

@@ -17,6 +17,7 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_unittest(CodeGenTests
   AArch64SelectionDAGTest.cpp
   AllocationOrderTest.cpp
+  AMDGPUMetadataTest.cpp
   AsmPrinterDwarfTest.cpp
   DIEHashTest.cpp
   DIETest.cpp