[Mlir-commits] [mlir] [mlir][ROCDL] Add metadata attributes for storing ELF object information (PR #95292)
Fabian Mora
llvmlistbot at llvm.org
Wed Jun 12 13:50:02 PDT 2024
https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/95292
>From 7be061ec5d33a84a7095ffa1c0d04cb5fb7b1882 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Wed, 12 Jun 2024 20:49:09 +0000
Subject: [PATCH] move to gpu
---
mlir/include/mlir-c/Dialect/GPU.h | 3 +-
.../mlir/Dialect/GPU/IR/CompilationAttrs.td | 147 +++++++++++-
mlir/include/mlir/Target/LLVM/ROCDL/Utils.h | 7 +
mlir/lib/CAPI/Dialect/GPU.cpp | 12 +-
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp | 37 ++-
mlir/lib/Target/LLVM/CMakeLists.txt | 2 +
mlir/lib/Target/LLVM/NVVM/Target.cpp | 2 +-
mlir/lib/Target/LLVM/ROCDL/Target.cpp | 3 +-
mlir/lib/Target/LLVM/ROCDL/Utils.cpp | 226 ++++++++++++++++++
mlir/lib/Target/SPIRV/Target.cpp | 2 +-
mlir/test/Dialect/GPU/ops.mlir | 6 +
.../Target/LLVM/SerializeROCDLTarget.cpp | 45 ++++
12 files changed, 481 insertions(+), 11 deletions(-)
create mode 100644 mlir/lib/Target/LLVM/ROCDL/Utils.cpp
diff --git a/mlir/include/mlir-c/Dialect/GPU.h b/mlir/include/mlir-c/Dialect/GPU.h
index 2adf73ddff6ea..74325bfb00e2e 100644
--- a/mlir/include/mlir-c/Dialect/GPU.h
+++ b/mlir/include/mlir-c/Dialect/GPU.h
@@ -27,7 +27,8 @@ MLIR_CAPI_EXPORTED bool mlirAttributeIsAGPUObjectAttr(MlirAttribute attr);
MLIR_CAPI_EXPORTED MlirAttribute
mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target, uint32_t format,
- MlirStringRef objectStrRef, MlirAttribute mlirObjectProps);
+ MlirStringRef objectStrRef, MlirAttribute mlirObjectProps,
+ MlirAttribute mlirKernelsAttr);
MLIR_CAPI_EXPORTED MlirAttribute
mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr);
diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td
index 6659f4a2c58e8..f4037b55c85b4 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td
@@ -16,6 +16,136 @@
include "mlir/Dialect/GPU/IR/GPUBase.td"
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
+//===----------------------------------------------------------------------===//
+// GPU kernel attribute
+//===----------------------------------------------------------------------===//
+
+def GPU_KernelAttr : GPU_Attr<"Kernel", "kernel"> {
+ let description = [{
+ GPU attribute for storing metadata related to a compiled kernel. It
+ contains the attribute dictionary of the LLVM function used to generate the
+ kernel, as well as an optional dictionary for additional metadata, like
+ occupancy information.
+
+ Examples:
+ ```mlir
+ #gpu.kernel<{sym_name = "test_fusion__part_0", ...},
+ metadata = {reg_count = 255, ...}>
+ ```
+ }];
+ let parameters = (ins
+ "DictionaryAttr":$func_attrs,
+ OptionalParameter<"DictionaryAttr", "metadata dictionary">:$metadata
+ );
+ let assemblyFormat = [{
+ `<` $func_attrs (`,` `metadata` `=` $metadata^ )? `>`
+ }];
+ let builders = [
+ AttrBuilderWithInferredContext<(ins "DictionaryAttr":$funcAttrs,
+ CArg<"DictionaryAttr",
+ "nullptr">:$metadata), [{
+ assert(funcAttrs && "invalid function attributes dictionary");
+ return $_get(funcAttrs.getContext(), funcAttrs, metadata);
+ }]>,
+ AttrBuilderWithInferredContext<(ins "Operation*":$kernel,
+ CArg<"DictionaryAttr",
+ "nullptr">:$metadata)>
+ ];
+ let extraClassDeclaration = [{
+ /// Returns the function attribute corresponding to key or nullptr if missing.
+ Attribute getAttr(StringRef key) const {
+ return getFuncAttrs().get(key);
+ }
+ template <typename ConcreteAttr>
+ ConcreteAttr getAttr(StringRef key) const {
+ return llvm::dyn_cast_or_null<ConcreteAttr>(getAttr(key));
+ }
+ Attribute getAttr(StringAttr key) const;
+ template <typename ConcreteAttr>
+ ConcreteAttr getAttr(StringAttr key) const {
+ return llvm::dyn_cast_or_null<ConcreteAttr>(getAttr(key));
+ }
+
+ /// Returns the name of the kernel.
+ StringAttr getName() const {
+ return getAttr<StringAttr>("sym_name");
+ }
+
+ /// Returns the metadta attribute corresponding to key or nullptr if missing.
+ Attribute getMDAttr(StringRef key) const {
+ if (DictionaryAttr attrs = getMetadata())
+ return attrs.get(key);
+ return nullptr;
+ }
+ template <typename ConcreteAttr>
+ ConcreteAttr getMDAttr(StringRef key) const {
+ return llvm::dyn_cast_or_null<ConcreteAttr>(getMDAttr(key));
+ }
+ Attribute getMDAttr(StringAttr key) const;
+ template <typename ConcreteAttr>
+ ConcreteAttr getMDAttr(StringAttr key) const {
+ return llvm::dyn_cast_or_null<ConcreteAttr>(getMDAttr(key));
+ }
+
+ /// Helper function for appending metadata to a kernel attribute.
+ KernelAttr appendMetadata(ArrayRef<NamedAttribute> attrs) const;
+ }];
+}
+
+//===----------------------------------------------------------------------===//
+// GPU kernel table attribute
+//===----------------------------------------------------------------------===//
+
+def GPU_KernelTableAttr : GPU_Attr<"KernelTable", "kernel_table"> {
+ let description = [{
+ GPU attribute representing a table of kernels metadata. All the attributes
+ in the dictionary must be of type `#gpu.kernel`.
+
+ Examples:
+ ```mlir
+ #gpu.kernel_table<{kernel0 = #gpu.kernel<...>}>
+ ```
+ }];
+ let parameters = (ins
+ "DictionaryAttr":$kernel_table
+ );
+ let assemblyFormat = [{
+ `<` $kernel_table `>`
+ }];
+ let builders = [
+ AttrBuilderWithInferredContext<(ins "DictionaryAttr":$kernel_table), [{
+ assert(kernel_table && "invalid kernel table");
+ return $_get(kernel_table.getContext(), kernel_table);
+ }]>
+ ];
+ let skipDefaultBuilders = 1;
+ let genVerifyDecl = 1;
+ let extraClassDeclaration = [{
+ /// Helper iterator class for traversing the kernel table.
+ struct KernelIterator
+ : llvm::mapped_iterator_base<KernelIterator,
+ llvm::ArrayRef<NamedAttribute>::iterator,
+ std::pair<StringAttr, KernelAttr>> {
+ using llvm::mapped_iterator_base<
+ KernelIterator, llvm::ArrayRef<NamedAttribute>::iterator,
+ std::pair<StringAttr, KernelAttr>>::mapped_iterator_base;
+ /// Map the iterator to the kernel name and a KernelAttribute.
+ std::pair<StringAttr, KernelAttr> mapElement(NamedAttribute attr) const {
+ return {attr.getName(), llvm::cast<KernelAttr>(attr.getValue())};
+ }
+ };
+ auto begin() const {
+ return KernelIterator(getKernelTable().begin());
+ }
+ auto end() const {
+ return KernelIterator(getKernelTable().end());
+ }
+ size_t size() const {
+ return getKernelTable().size();
+ }
+ }];
+}
+
//===----------------------------------------------------------------------===//
// GPU object attribute.
//===----------------------------------------------------------------------===//
@@ -63,16 +193,29 @@ def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
#gpu.object<#nvvm.target, properties = {O = 3 : i32}, assembly = "..."> // An assembly object with additional properties.
#gpu.object<#rocdl.target, bin = "..."> // A binary object.
#gpu.object<#nvvm.target, "..."> // A fatbin object.
+ #gpu.object<#nvvm.target, kernels = #gpu.kernel_table<...>, "..."> // An object with a kernel table.
```
}];
let parameters = (ins
"Attribute":$target,
DefaultValuedParameter<"CompilationTarget", "CompilationTarget::Fatbin">:$format,
"StringAttr":$object,
- OptionalParameter<"DictionaryAttr">:$properties
+ OptionalParameter<"DictionaryAttr">:$properties,
+ OptionalParameter<"KernelTableAttr">:$kernels
);
+ let builders = [
+ AttrBuilderWithInferredContext<(ins "Attribute":$target,
+ "CompilationTarget":$format,
+ "StringAttr":$object,
+ CArg<"DictionaryAttr", "nullptr">:$properties,
+ CArg<"KernelTableAttr", "nullptr">:$kernels), [{
+ assert(target && "invalid target");
+ return $_get(target.getContext(), target, format, object, properties, kernels);
+ }]>
+ ];
let assemblyFormat = [{ `<`
- $target `,` (`properties` `=` $properties ^ `,`)?
+ $target `,` (`properties` `=` $properties^ `,`)?
+ (`kernels` `=` $kernels^ `,`)?
custom<Object>($format, $object)
`>`
}];
diff --git a/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h b/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
index 374fa65bd02e3..015a853d79d31 100644
--- a/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
+++ b/mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
@@ -14,6 +14,7 @@
#define MLIR_TARGET_LLVM_ROCDL_UTILS_H
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/Target/LLVM/ModuleToObject.h"
@@ -85,6 +86,12 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
/// List of LLVM bitcode files to link to.
SmallVector<std::string> fileList;
};
+
+/// Returns a `#gpu.kernel_table` containing kernel metadata for each of the
+/// kernels in `gpuModule`. If `elfData` is valid, then the `amdhsa.kernels` ELF
+/// metadata will be added to the `#gpu.kernel_table`.
+gpu::KernelTableAttr getAMDHSAKernelsMetadata(Operation *gpuModule,
+ ArrayRef<char> elfData = {});
} // namespace ROCDL
} // namespace mlir
diff --git a/mlir/lib/CAPI/Dialect/GPU.cpp b/mlir/lib/CAPI/Dialect/GPU.cpp
index e471e8cd9588e..8a17d117cdd8c 100644
--- a/mlir/lib/CAPI/Dialect/GPU.cpp
+++ b/mlir/lib/CAPI/Dialect/GPU.cpp
@@ -25,15 +25,19 @@ bool mlirAttributeIsAGPUObjectAttr(MlirAttribute attr) {
MlirAttribute mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target,
uint32_t format, MlirStringRef objectStrRef,
- MlirAttribute mlirObjectProps) {
+ MlirAttribute mlirObjectProps,
+ MlirAttribute mlirKernelsAttr) {
MLIRContext *ctx = unwrap(mlirCtx);
llvm::StringRef object = unwrap(objectStrRef);
DictionaryAttr objectProps;
if (mlirObjectProps.ptr != nullptr)
objectProps = llvm::cast<DictionaryAttr>(unwrap(mlirObjectProps));
- return wrap(gpu::ObjectAttr::get(ctx, unwrap(target),
- static_cast<gpu::CompilationTarget>(format),
- StringAttr::get(ctx, object), objectProps));
+ gpu::KernelTableAttr kernels;
+ if (mlirKernelsAttr.ptr != nullptr)
+ kernels = llvm::cast<gpu::KernelTableAttr>(unwrap(mlirKernelsAttr));
+ return wrap(gpu::ObjectAttr::get(
+ ctx, unwrap(target), static_cast<gpu::CompilationTarget>(format),
+ StringAttr::get(ctx, object), objectProps, kernels));
}
MlirAttribute mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr) {
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index d8e29da6512d4..20383393295f5 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2140,7 +2140,8 @@ void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
LogicalResult ObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
Attribute target, CompilationTarget format,
- StringAttr object, DictionaryAttr properties) {
+ StringAttr object, DictionaryAttr properties,
+ KernelTableAttr kernels) {
if (!target)
return emitError() << "the target attribute cannot be null";
if (target.hasPromiseOrImplementsInterface<TargetAttrInterface>())
@@ -2226,6 +2227,40 @@ LogicalResult gpu::DynamicSharedMemoryOp::verify() {
return success();
}
+//===----------------------------------------------------------------------===//
+// GPU KernelAttr
+//===----------------------------------------------------------------------===//
+
+KernelAttr KernelAttr::get(Operation *kernelOp, DictionaryAttr metadata) {
+ assert(kernelOp && "invalid kernel");
+ return get(kernelOp->getAttrDictionary(), metadata);
+}
+
+KernelAttr KernelAttr::appendMetadata(ArrayRef<NamedAttribute> attrs) const {
+ if (attrs.empty())
+ return *this;
+ NamedAttrList attrList(attrs);
+ attrList.append(getMetadata());
+ return KernelAttr::get(getFuncAttrs(), attrList.getDictionary(getContext()));
+}
+
+//===----------------------------------------------------------------------===//
+// GPU KernelTableAttr
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+KernelTableAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+ DictionaryAttr dict) {
+ if (!dict)
+ return emitError() << "table cannot be null";
+ if (llvm::any_of(dict, [](NamedAttribute attr) {
+ return !llvm::isa<KernelAttr>(attr.getValue());
+ }))
+ return emitError()
+ << "all the dictionary values must be `#gpu.kernel` attributes";
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// GPU target options
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVM/CMakeLists.txt b/mlir/lib/Target/LLVM/CMakeLists.txt
index 5a3fa160850b4..08da4966e1499 100644
--- a/mlir/lib/Target/LLVM/CMakeLists.txt
+++ b/mlir/lib/Target/LLVM/CMakeLists.txt
@@ -108,9 +108,11 @@ endif()
add_mlir_dialect_library(MLIRROCDLTarget
ROCDL/Target.cpp
+ ROCDL/Utils.cpp
LINK_COMPONENTS
MCParser
+ Object
${AMDGPU_LIBS}
LINK_LIBS PUBLIC
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index acb903aa37caf..347b5af8d23c7 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -606,5 +606,5 @@ NVVMTargetAttrImpl::createObject(Attribute attribute,
return builder.getAttr<gpu::ObjectAttr>(
attribute, format,
builder.getStringAttr(StringRef(object.data(), object.size())),
- objectProps);
+ objectProps, nullptr);
}
diff --git a/mlir/lib/Target/LLVM/ROCDL/Target.cpp b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
index cc13e5b7436ea..f4907fe2904b9 100644
--- a/mlir/lib/Target/LLVM/ROCDL/Target.cpp
+++ b/mlir/lib/Target/LLVM/ROCDL/Target.cpp
@@ -482,5 +482,6 @@ ROCDLTargetAttrImpl::createObject(Attribute attribute,
attribute,
format > gpu::CompilationTarget::Binary ? gpu::CompilationTarget::Binary
: format,
- builder.getStringAttr(StringRef(object.data(), object.size())), nullptr);
+ builder.getStringAttr(StringRef(object.data(), object.size())), nullptr,
+ nullptr);
}
diff --git a/mlir/lib/Target/LLVM/ROCDL/Utils.cpp b/mlir/lib/Target/LLVM/ROCDL/Utils.cpp
new file mode 100644
index 0000000000000..5029293bae2a6
--- /dev/null
+++ b/mlir/lib/Target/LLVM/ROCDL/Utils.cpp
@@ -0,0 +1,226 @@
+//===- 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::optional<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();
+ auto secOrErr = elf.sections();
+ if (!secOrErr)
+ return std::nullopt;
+ ArrayRef<typename ELFT::Shdr> sections = *secOrErr;
+ for (auto 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());
+ msgpack::Document msgPackDoc;
+ if (!msgPackDoc.readFromBlob(msgPackString, /*Multi=*/false))
+ return std::nullopt;
+ if (msgPackDoc.getRoot().isScalar())
+ return std::nullopt;
+ return std::optional<llvm::msgpack::Document>(std::move(msgPackDoc));
+ }
+ }
+ return std::nullopt;
+}
+
+/// Return the `amdhsa.kernels` metadata in the ELF object or std::nullopt on
+/// failure. This is a helper function that casts a generic `ObjectFile` to the
+/// appropiate `ELFObjectFile`.
+static std::optional<llvm::msgpack::Document>
+getAMDHSANote(ArrayRef<char> elfData) {
+ using namespace llvm;
+ using namespace llvm::object;
+ if (elfData.empty())
+ return std::nullopt;
+ 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 std::nullopt;
+ }
+ ObjectFile &elf = *(objOrErr.get());
+ std::optional<llvm::msgpack::Document> metadata;
+ if (auto *obj = dyn_cast<ELF32LEObjectFile>(&elf))
+ metadata = getAMDHSANote(*obj);
+ else if (auto *obj = dyn_cast<ELF32BEObjectFile>(&elf))
+ metadata = getAMDHSANote(*obj);
+ else if (auto *obj = dyn_cast<ELF64LEObjectFile>(&elf))
+ metadata = getAMDHSANote(*obj);
+ else if (auto *obj = dyn_cast<ELF64BEObjectFile>(&elf))
+ metadata = getAMDHSANote(*obj);
+ return metadata;
+}
+
+/// Utility functions for converting `llvm::msgpack::DocNode` nodes.
+static Attribute convertNode(Builder &builder, llvm::msgpack::DocNode &node);
+static Attribute convertNode(Builder &builder,
+ llvm::msgpack::MapDocNode &node) {
+ NamedAttrList attrs;
+ for (auto kv : node) {
+ if (!kv.first.isString())
+ continue;
+ if (Attribute attr = convertNode(builder, kv.second)) {
+ auto key = kv.first.getString();
+ key.consume_front(".");
+ key.consume_back(".");
+ attrs.append(key, attr);
+ }
+ }
+ if (attrs.empty())
+ return nullptr;
+ return builder.getDictionaryAttr(attrs);
+}
+
+static Attribute convertNode(Builder &builder,
+ llvm::msgpack::ArrayDocNode &node) {
+ using NodeKind = llvm::msgpack::Type;
+ // Use `DenseIntAttr` if we know all the attrs are ints.
+ if (llvm::all_of(node, [](llvm::msgpack::DocNode &n) {
+ auto kind = n.getKind();
+ return kind == NodeKind::Int || kind == NodeKind::UInt;
+ })) {
+ SmallVector<int64_t> values;
+ for (llvm::msgpack::DocNode &n : node) {
+ auto kind = n.getKind();
+ if (kind == NodeKind::Int)
+ values.push_back(n.getInt());
+ else if (kind == NodeKind::UInt)
+ values.push_back(n.getUInt());
+ }
+ return builder.getDenseI64ArrayAttr(values);
+ }
+ // Convert the array.
+ SmallVector<Attribute> attrs;
+ for (llvm::msgpack::DocNode &n : node) {
+ if (Attribute attr = convertNode(builder, n))
+ attrs.push_back(attr);
+ }
+ if (attrs.empty())
+ return nullptr;
+ return builder.getArrayAttr(attrs);
+}
+
+static Attribute convertNode(Builder &builder, llvm::msgpack::DocNode &node) {
+ using namespace llvm::msgpack;
+ using NodeKind = llvm::msgpack::Type;
+ switch (node.getKind()) {
+ case NodeKind::Int:
+ return builder.getI64IntegerAttr(node.getInt());
+ case NodeKind::UInt:
+ return builder.getI64IntegerAttr(node.getUInt());
+ case NodeKind::Boolean:
+ return builder.getI64IntegerAttr(node.getBool());
+ case NodeKind::String:
+ return builder.getStringAttr(node.getString());
+ case NodeKind::Array:
+ return convertNode(builder, node.getArray());
+ case NodeKind::Map:
+ return convertNode(builder, node.getMap());
+ default:
+ return nullptr;
+ }
+}
+
+/// The following function should succeed for Code object V3 and above.
+static llvm::StringMap<DictionaryAttr> getELFMetadata(Builder &builder,
+ ArrayRef<char> elfData) {
+ std::optional<llvm::msgpack::Document> metadata = getAMDHSANote(elfData);
+ if (!metadata)
+ return {};
+ llvm::StringMap<DictionaryAttr> kernelMD;
+ llvm::msgpack::DocNode &root = (metadata)->getRoot();
+ // Fail if `root` is not a map -it should be for AMD Obj Ver 3.
+ if (!root.isMap())
+ return kernelMD;
+ auto &kernels = root.getMap()["amdhsa.kernels"];
+ // Fail if `amdhsa.kernels` is not an array.
+ if (!kernels.isArray())
+ return kernelMD;
+ // Convert each of the kernels.
+ for (auto &kernel : kernels.getArray()) {
+ if (!kernel.isMap())
+ continue;
+ auto &kernelMap = kernel.getMap();
+ auto &name = kernelMap[".name"];
+ if (!name.isString())
+ continue;
+ NamedAttrList attrList;
+ // Convert the kernel properties.
+ for (auto kv : kernelMap) {
+ if (!kv.first.isString())
+ continue;
+ StringRef key = kv.first.getString();
+ key.consume_front(".");
+ key.consume_back(".");
+ if (key == "name")
+ continue;
+ if (Attribute attr = convertNode(builder, kv.second))
+ attrList.append(key, attr);
+ }
+ if (!attrList.empty())
+ kernelMD[name.getString()] = builder.getDictionaryAttr(attrList);
+ }
+ return kernelMD;
+}
+
+gpu::KernelTableAttr
+mlir::ROCDL::getAMDHSAKernelsMetadata(Operation *gpuModule,
+ ArrayRef<char> elfData) {
+ auto module = cast<gpu::GPUModuleOp>(gpuModule);
+ Builder builder(module.getContext());
+ NamedAttrList moduleAttrs;
+ llvm::StringMap<DictionaryAttr> mdMap = getELFMetadata(builder, elfData);
+ for (auto funcOp : module.getBody()->getOps<LLVM::LLVMFuncOp>()) {
+ if (!funcOp->getDiscardableAttr("rocdl.kernel"))
+ continue;
+ moduleAttrs.append(
+ funcOp.getName(),
+ gpu::KernelAttr::get(funcOp, mdMap.lookup(funcOp.getName())));
+ }
+ return gpu::KernelTableAttr::get(
+ moduleAttrs.getDictionary(module.getContext()));
+}
diff --git a/mlir/lib/Target/SPIRV/Target.cpp b/mlir/lib/Target/SPIRV/Target.cpp
index 4c416abe71cac..e7651b5f3c767 100644
--- a/mlir/lib/Target/SPIRV/Target.cpp
+++ b/mlir/lib/Target/SPIRV/Target.cpp
@@ -98,5 +98,5 @@ SPIRVTargetAttrImpl::createObject(Attribute attribute,
return builder.getAttr<gpu::ObjectAttr>(
attribute, format,
builder.getStringAttr(StringRef(object.data(), object.size())),
- objectProps);
+ objectProps, nullptr);
}
diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir
index ba7897f4e80cb..692ef2a5e3bef 100644
--- a/mlir/test/Dialect/GPU/ops.mlir
+++ b/mlir/test/Dialect/GPU/ops.mlir
@@ -441,3 +441,9 @@ gpu.module @module_with_two_target [#nvvm.target, #rocdl.target<chip = "gfx90a">
gpu.module @module_with_offload_handler <#gpu.select_object<0>> [#nvvm.target] {
}
+
+
+gpu.binary @binary [#gpu.object<#rocdl.target<chip = "gfx900">, kernels = #gpu.kernel_table<{
+ kernel0 = #gpu.kernel<{sym_name = "kernel0"}, metadata = {sgpr_count = 255}>,
+ kernel1 = #gpu.kernel<{sym_name = "kernel1"}>
+ }> , bin = "BLOB">]
diff --git a/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp b/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp
index 33291bc4bcaed..3d5c84efb6f4f 100644
--- a/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeROCDLTarget.cpp
@@ -158,3 +158,48 @@ TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(SerializeROCDLToBinary)) {
ASSERT_FALSE(object->empty());
}
}
+
+// Test ROCDL metadata.
+TEST_F(MLIRTargetLLVMROCDL, SKIP_WITHOUT_AMDGPU(GetELFMetadata)) {
+ if (!hasROCMTools())
+ GTEST_SKIP() << "ROCm installation not found, skipping test.";
+
+ MLIRContext context(registry);
+
+ OwningOpRef<ModuleOp> module =
+ parseSourceString<ModuleOp>(moduleStr, &context);
+ ASSERT_TRUE(!!module);
+
+ // Create a ROCDL target.
+ ROCDL::ROCDLTargetAttr target = ROCDL::ROCDLTargetAttr::get(&context);
+
+ // Serialize the module.
+ auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
+ ASSERT_TRUE(!!serializer);
+ gpu::TargetOptions options("", {}, "", gpu::CompilationTarget::Binary);
+ for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
+ std::optional<SmallVector<char, 0>> object =
+ serializer.serializeToObject(gpuModule, options);
+ // Check that the serializer was successful.
+ ASSERT_TRUE(object != std::nullopt);
+ ASSERT_FALSE(object->empty());
+ if (!object)
+ continue;
+ // Get the metadata.
+ gpu::KernelTableAttr metadata =
+ ROCDL::getAMDHSAKernelsMetadata(gpuModule, *object);
+ ASSERT_TRUE(metadata != nullptr);
+ // There should be only a single kernel.
+ ASSERT_TRUE(metadata.size() == 1);
+ // Test the `ROCDLObjectMDAttr` iterators.
+ for (auto [name, kernel] : metadata) {
+ ASSERT_TRUE(name.getValue() == "rocdl_kernel");
+ // Check that the ELF metadata is present.
+ ASSERT_TRUE(kernel.getMetadata() != nullptr);
+ // Verify that `sgpr_count` is present and it is an integer attribute.
+ ASSERT_TRUE(kernel.getMDAttr<IntegerAttr>("sgpr_count") != nullptr);
+ // Verify that `vgpr_count` is present and it is an integer attribute.
+ ASSERT_TRUE(kernel.getMDAttr<IntegerAttr>("vgpr_count") != nullptr);
+ }
+ }
+}
More information about the Mlir-commits
mailing list