[mlir][gpu] Add metadata attributes for storing kernel metadata in GPU objects

mlir/include/mlir-c/Dialect/GPU.h

@@ -37,6 +37,11 @@ MLIR_CAPI_EXPORTED MlirAttribute
 mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target, uint32_t format,
                      MlirStringRef objectStrRef, MlirAttribute mlirObjectProps);
 
+MLIR_CAPI_EXPORTED MlirAttribute mlirGPUObjectAttrGetWithKernels(
+    MlirContext mlirCtx, MlirAttribute target, uint32_t format,
+    MlirStringRef objectStrRef, MlirAttribute mlirObjectProps,
+    MlirAttribute mlirKernelsAttr);
+
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr);
 
@@ -52,6 +57,12 @@ mlirGPUObjectAttrHasProperties(MlirAttribute mlirObjectAttr);
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirGPUObjectAttrGetProperties(MlirAttribute mlirObjectAttr);
 
+MLIR_CAPI_EXPORTED bool
+mlirGPUObjectAttrHasKernels(MlirAttribute mlirObjectAttr);
+
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirGPUObjectAttrGetKernels(MlirAttribute mlirObjectAttr);
+
 #ifdef __cplusplus
 }
 #endif

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td

@@ -16,6 +16,155 @@
 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. The
+    attribute contains the name and function type of the kernel.
+
+    The attribute also contains optional parameters for storing the arguments
+    attributes as well as a dictionary for additional metadata, like occupancy
+    information or other function attributes.
+
+    Note: The `arg_attrs` parameter is expected to follow all the constraints
+    imposed by the `mlir::FunctionOpInterface` interface.
+
+    Examples:
+    ```mlir
+      #gpu.kernel<@kernel1, (i32) -> (), arg_attrs = [...],  metadata = {reg_count = 255, ...}>
+      #gpu.kernel<@kernel2, (i32, f64) -> ()>
+    ```
+  }];
+  let parameters = (ins
+    "StringAttr":$name,
+    "Type":$function_type,
+    OptionalParameter<"ArrayAttr", "arguments attributes">:$arg_attrs,
+    OptionalParameter<"DictionaryAttr", "metadata dictionary">:$metadata
+  );
+  let assemblyFormat = [{
+    `<` $name `,` $function_type (`,` struct($arg_attrs, $metadata)^)? `>`
+  }];
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "StringAttr":$name,
+                                        "Type":$functionType,
+                                        CArg<"ArrayAttr", "nullptr">:$argAttrs,
+                                        CArg<"DictionaryAttr",
+                                             "nullptr">:$metadata), [{
+      assert(name && "invalid name");
+      return $_get(name.getContext(), name, functionType, argAttrs, metadata);
+    }]>,
+    AttrBuilderWithInferredContext<(ins "FunctionOpInterface":$kernel,
+                                         CArg<"DictionaryAttr",
+                                              "nullptr">:$metadata)>
+  ];
+  let genVerifyDecl = 1;
+  let extraClassDeclaration = [{
+    /// Compare two kernels based on the name.
+    bool operator<(const KernelAttr& other) const {
+      return getName().getValue() < other.getName().getValue();
+    }
+
+    /// Returns the metadata attribute corresponding to `key` or `nullptr`
+    /// if missing.
+    Attribute getAttr(StringRef key) const {
+      DictionaryAttr attrs = getMetadata();
+      return attrs ? attrs.get(key) : nullptr;
+    }
+    template <typename ConcreteAttr>
+    ConcreteAttr getAttr(StringRef key) const {
+      return llvm::dyn_cast_or_null<ConcreteAttr>(getAttr(key));
+    }
+    Attribute getAttr(StringAttr key) const {
+      DictionaryAttr attrs = getMetadata();
+      return attrs ? attrs.get(key) : nullptr;
+    }
+    template <typename ConcreteAttr>
+    ConcreteAttr getAttr(StringAttr key) const {
+      return llvm::dyn_cast_or_null<ConcreteAttr>(getAttr(key));
+    }
+
+    /// Returns the attribute dictionary at position `index`.
+    DictionaryAttr getArgAttrDict(unsigned index) {
+      ArrayAttr argArray = getArgAttrs();
+      return argArray ? llvm::cast<DictionaryAttr>(argArray[index]) : nullptr;
+    }
+
+    /// Return the specified attribute, if present, for the argument at 'index',
+    /// null otherwise.
+    Attribute getArgAttr(unsigned index, StringAttr name) {
+      DictionaryAttr argDict = getArgAttrDict(index);
+      return argDict ? argDict.get(name) : nullptr;
+    }
+    Attribute getArgAttr(unsigned index, StringRef name) {
+      DictionaryAttr argDict = getArgAttrDict(index);
+      return argDict ? argDict.get(name) : nullptr;
+    }
+
+    /// Returns a new KernelAttr that contains `attrs` in the metadata dictionary.
+    KernelAttr appendMetadata(ArrayRef<NamedAttribute> attrs) const;
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// GPU kernel table attribute
+//===----------------------------------------------------------------------===//
+
+def GPU_KernelTableAttr : GPU_Attr<"KernelTable", "kernel_table"> {
+  let description = [{
+    GPU attribute representing a list of `#gpu.kernel` attributes. This
+    attribute supports searching kernels by name. All kernels in the table must
+    have an unique name.
+
+    Examples:
+    ```mlir
+      // Empty table.
+      #gpu.kernel_table<>
+
+      // Table with a single kernel.
+      #gpu.kernel_table<[#gpu.kernel<kernel0, () -> () >]>
+
+      // Table with multiple kernels.
+      #gpu.kernel_table<[
+        #gpu.kernel<"kernel0", (i32, f32) -> (), metadata = {sgpr_count = 255}>,
+        #gpu.kernel<"kernel1", (i32) -> ()>
+      ]>
+    ```
+  }];
+  let parameters = (ins
+    OptionalArrayRefParameter<"KernelAttr", "array of kernels">:$kernel_table
+  );
+  let assemblyFormat = [{
+    `<` (`[` qualified($kernel_table)^ `]`)? `>`
+  }];
+  let builders = [
+    AttrBuilder<(ins "ArrayRef<KernelAttr>":$kernels,
+                     CArg<"bool", "false">:$isSorted)>
+  ];
+  let skipDefaultBuilders = 1;
+  let genVerifyDecl = 1;
+  let extraClassDeclaration = [{
+    llvm::ArrayRef<KernelAttr>::iterator begin() const {
+      return getKernelTable().begin();
+    }
+    llvm::ArrayRef<KernelAttr>::iterator end() const {
+      return getKernelTable().end();
+    }
+    size_t size() const {
+      return getKernelTable().size();
+    }
+    bool empty() const {
+      return getKernelTable().empty();
+    }
+
+    /// Returns the kernel with name `key` or `nullptr` if not present.
+    KernelAttr lookup(StringRef key) const;
+    KernelAttr lookup(StringAttr key) const;
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // GPU object attribute.
 //===----------------------------------------------------------------------===//
@@ -36,8 +185,9 @@ def GPU_CompilationTargetEnum : GPU_I32Enum<
 def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
   let description = [{
     A GPU object attribute glues together a GPU target, the object kind, a
-    binary string with the object, and the object properties, encapsulating how
-    the object was generated and its properties with the object itself.
+    binary string with the object, the object properties, and kernel metadata,
+    encapsulating how the object was generated and its properties with the
+    object itself.
 
     There are four object formats:
     1. `Offload`: represents generic objects not described by the other three
@@ -55,6 +205,10 @@ def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
 
     Object properties are specified through the `properties` dictionary
     attribute and can be used to define additional information.
+
+    Kernel metadata is specified through the `kernels` parameter, and can be
+    used to specify additional information on a kernel by kernel basis.
+
     The target attribute must implement or promise the `TargetAttrInterface`
     interface.
 
@@ -63,16 +217,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)
     `>`
   }];

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/Support/LLVM.h"
 #include "mlir/Target/LLVM/ModuleToObject.h"
@@ -107,6 +108,20 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
   /// AMD GCN libraries to use when linking, the default is using none.
   AMDGCNLibraries deviceLibs = AMDGCNLibraries::None;
 };
+
+/// Returns a map containing the `amdhsa.kernels` ELF metadata for each of the
+/// kernels in the binary, or `std::nullopt` if the metadata couldn't be
+/// retrieved. The map associates the name of the kernel with the list of named
+/// attributes found in `amdhsa.kernels`. For more information on the ELF
+/// metadata see: https://llvm.org/docs/AMDGPUUsage.html#amdhsa
+std::optional<DenseMap<StringAttr, NamedAttrList>>
+getAMDHSAKernelsELFMetadata(Builder &builder, ArrayRef<char> elfData);
+
+/// 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 getKernelMetadata(Operation *gpuModule,
+                                       ArrayRef<char> elfData = {});
 } // namespace ROCDL
 } // namespace mlir
 

mlir/lib/Bindings/Python/DialectGPU.cpp

@@ -48,17 +48,21 @@ PYBIND11_MODULE(_mlirDialectsGPU, m) {
       .def_classmethod(
           "get",
           [](py::object cls, MlirAttribute target, uint32_t format,
-             py::bytes object, std::optional<MlirAttribute> mlirObjectProps) {
+             py::bytes object, std::optional<MlirAttribute> mlirObjectProps,
+             std::optional<MlirAttribute> mlirKernelsAttr) {
             py::buffer_info info(py::buffer(object).request());
             MlirStringRef objectStrRef =
                 mlirStringRefCreate(static_cast<char *>(info.ptr), info.size);
-            return cls(mlirGPUObjectAttrGet(
+            return cls(mlirGPUObjectAttrGetWithKernels(
                 mlirAttributeGetContext(target), target, format, objectStrRef,
                 mlirObjectProps.has_value() ? *mlirObjectProps
+                                            : MlirAttribute{nullptr},
+                mlirKernelsAttr.has_value() ? *mlirKernelsAttr
                                             : MlirAttribute{nullptr}));
           },
           "cls"_a, "target"_a, "format"_a, "object"_a,
-          "properties"_a = py::none(), "Gets a gpu.object from parameters.")
+          "properties"_a = py::none(), "kernels"_a = py::none(),
+          "Gets a gpu.object from parameters.")
       .def_property_readonly(
           "target",
           [](MlirAttribute self) { return mlirGPUObjectAttrGetTarget(self); })
@@ -71,9 +75,16 @@ PYBIND11_MODULE(_mlirDialectsGPU, m) {
             MlirStringRef stringRef = mlirGPUObjectAttrGetObject(self);
             return py::bytes(stringRef.data, stringRef.length);
           })
-      .def_property_readonly("properties", [](MlirAttribute self) {
-        if (mlirGPUObjectAttrHasProperties(self))
-          return py::cast(mlirGPUObjectAttrGetProperties(self));
+      .def_property_readonly("properties",
+                             [](MlirAttribute self) {
+                               if (mlirGPUObjectAttrHasProperties(self))
+                                 return py::cast(
+                                     mlirGPUObjectAttrGetProperties(self));
+                               return py::none().cast<py::object>();
+                             })
+      .def_property_readonly("kernels", [](MlirAttribute self) {
+        if (mlirGPUObjectAttrHasKernels(self))
+          return py::cast(mlirGPUObjectAttrGetKernels(self));
         return py::none().cast<py::object>();
       });
 }

mlir/lib/CAPI/Dialect/GPU.cpp

@@ -43,9 +43,28 @@ MlirAttribute mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target,
   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));
+  return wrap(gpu::ObjectAttr::get(
+      ctx, unwrap(target), static_cast<gpu::CompilationTarget>(format),
+      StringAttr::get(ctx, object), objectProps, nullptr));
+}
+
+MlirAttribute mlirGPUObjectAttrGetWithKernels(MlirContext mlirCtx,
+                                              MlirAttribute target,
+                                              uint32_t format,
+                                              MlirStringRef objectStrRef,
+                                              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));
+  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) {
@@ -78,3 +97,15 @@ MlirAttribute mlirGPUObjectAttrGetProperties(MlirAttribute mlirObjectAttr) {
       llvm::cast<gpu::ObjectAttr>(unwrap(mlirObjectAttr));
   return wrap(objectAttr.getProperties());
 }
+
+bool mlirGPUObjectAttrHasKernels(MlirAttribute mlirObjectAttr) {
+  gpu::ObjectAttr objectAttr =
+      llvm::cast<gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  return objectAttr.getKernels() != nullptr;
+}
+
+MlirAttribute mlirGPUObjectAttrGetKernels(MlirAttribute mlirObjectAttr) {
+  gpu::ObjectAttr objectAttr =
+      llvm::cast<gpu::ObjectAttr>(unwrap(mlirObjectAttr));
+  return wrap(objectAttr.getKernels());
+}