[Mlir-commits] [mlir] [mlir][gpu] Add metadata attributes for storing kernel metadata in GPU objects (PR #95292)

[Joseph Huber]

================
@@ -0,0 +1,230 @@
+//===- Utils.cpp - MLIR ROCDL target utils ----------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This files defines ROCDL target related utility classes and functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVM/ROCDL/Utils.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+
+#include "llvm/ADT/StringMap.h"
+#include "llvm/BinaryFormat/MsgPackDocument.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/AMDGPUMetadata.h"
+
+using namespace mlir;
+using namespace mlir::ROCDL;
+
+/// Search the ELF object and return an object containing the `amdhsa.kernels`
+/// metadata note. Function adapted from:
+/// llvm-project/llvm/tools/llvm-readobj/ELFDumper.cpp Also see
+/// `amdhsa.kernels`:
+/// https://llvm.org/docs/AMDGPUUsage.html#code-object-v3-metadata
+template <typename ELFT>
+static std::unique_ptr<llvm::msgpack::Document>
+getAMDHSANote(llvm::object::ELFObjectFile<ELFT> &elfObj) {
+  using namespace llvm;
+  using namespace llvm::object;
+  using namespace llvm::ELF;
+  const ELFFile<ELFT> &elf = elfObj.getELFFile();
+  Expected<typename ELFT::ShdrRange> secOrErr = elf.sections();
+  if (!secOrErr) {
+    consumeError(secOrErr.takeError());
+    return nullptr;
+  }
+  ArrayRef<typename ELFT::Shdr> sections = *secOrErr;
+  for (const typename ELFT::Shdr &section : sections) {
+    if (section.sh_type != ELF::SHT_NOTE)
+      continue;
+    size_t align = std::max(static_cast<unsigned>(section.sh_addralign), 4u);
+    Error err = Error::success();
+    for (const typename ELFT::Note note : elf.notes(section, err)) {
+      StringRef name = note.getName();
+      if (name != "AMDGPU")
+        continue;
+      uint32_t type = note.getType();
+      if (type != ELF::NT_AMDGPU_METADATA)
+        continue;
+      ArrayRef<uint8_t> desc = note.getDesc(align);
+      StringRef msgPackString =
+          StringRef(reinterpret_cast<const char *>(desc.data()), desc.size());
+      std::unique_ptr<llvm::msgpack::Document> msgPackDoc(
+          new llvm::msgpack::Document());
+      if (!msgPackDoc->readFromBlob(msgPackString, /*Multi=*/false))
+        return nullptr;
+      if (msgPackDoc->getRoot().isScalar())
+        return nullptr;
+      return msgPackDoc;
+    }
+  }
+  return nullptr;
+}
+
+/// Return the `amdhsa.kernels` metadata in the ELF object or nullptr on
+/// failure. This is a helper function that casts a generic `ObjectFile` to the
+/// appropiate `ELFObjectFile`.
+static std::unique_ptr<llvm::msgpack::Document>
+getAMDHSANote(ArrayRef<char> elfData) {
+  using namespace llvm;
+  using namespace llvm::object;
+  if (elfData.empty())
+    return nullptr;
+  MemoryBufferRef buffer(StringRef(elfData.data(), elfData.size()), "buffer");
+  Expected<std::unique_ptr<ObjectFile>> objOrErr =
+      ObjectFile::createELFObjectFile(buffer);
+  if (!objOrErr || !objOrErr.get()) {
+    // Drop the error.
+    llvm::consumeError(objOrErr.takeError());
+    return nullptr;
+  }
+  ObjectFile &elf = *(objOrErr.get());
+  if (auto *obj = dyn_cast<ELF32LEObjectFile>(&elf))
+    return getAMDHSANote(*obj);
+  else if (auto *obj = dyn_cast<ELF32BEObjectFile>(&elf))
+    return getAMDHSANote(*obj);
+  else if (auto *obj = dyn_cast<ELF64LEObjectFile>(&elf))
+    return getAMDHSANote(*obj);
+  else if (auto *obj = dyn_cast<ELF64BEObjectFile>(&elf))
+    return getAMDHSANote(*obj);
+  return nullptr;
+}
----------------
jhuber6 wrote:

I don't know if we're in a position where `mlir` can depend on `offload` here, possibly we could just move that to some LLVM utility. I don't know much of anything about MLIR unfortunately.

https://github.com/llvm/llvm-project/pull/95292