[llvm] 4a02562 - [AMDGPU] Lazily init pal metadata on first function

[Sebastian Neubauer via llvm-commits]

Author: Sebastian Neubauer
Date: 2022-02-04T18:39:35+01:00
New Revision: 4a02562275d460a4a2970ca3938cbca710c39268

URL: https://github.com/llvm/llvm-project/commit/4a02562275d460a4a2970ca3938cbca710c39268
DIFF: https://github.com/llvm/llvm-project/commit/4a02562275d460a4a2970ca3938cbca710c39268.diff

LOG: [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

Added: 
    llvm/unittests/CodeGen/AMDGPUMetadataTest.cpp

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
    llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
    llvm/unittests/CodeGen/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
index 6e2984f2a04fc..468da4c2c83d1 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp
+++ b/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();
 

diff  --git a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
index d5c60aa3be7d6..ddda2cf107b11 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.h
@@ -77,6 +77,8 @@ class AMDGPUAsmPrinter final : public AsmPrinter {
       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 @@ class AMDGPUAsmPrinter final : public AsmPrinter {
 
   std::vector<std::string> DisasmLines, HexLines;
   size_t DisasmLineMaxLen;
+  bool IsTargetStreamerInitialized;
 };
 
 } // end namespace llvm

diff  --git a/llvm/unittests/CodeGen/AMDGPUMetadataTest.cpp b/llvm/unittests/CodeGen/AMDGPUMetadataTest.cpp
new file mode 100644
index 0000000000000..d1829091a1c6f
--- /dev/null
+++ b/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

diff  --git a/llvm/unittests/CodeGen/CMakeLists.txt b/llvm/unittests/CodeGen/CMakeLists.txt
index 74346691bbcfe..a85e3c0d9921b 100644
--- a/llvm/unittests/CodeGen/CMakeLists.txt
+++ b/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