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

This patch adds the `#gpu.kernel_metadata` and `#gpu.kernel_table`
attributes. The `#gpu.kernel_metadata` attribute allows storing metadata
related to a compiled kernel, for example, the number of scalar
registers used by the kernel. The attribute only has 2 required
parameters, the name and function type. It also has 2 optional
parameters, the arguments attributes and generic dictionary for storing
all other metadata.

The `#gpu.kernel_table` stores a table of `#gpu.kernel_metadata`,
mapping the name of the kernel to the metadata.

Finally, the function `ROCDL::getAMDHSAKernelsELFMetadata` was added to
collect ELF metadata from a binary, and to test the class methods in
both attributes.

Example:
```mlir
gpu.binary binary [#gpu.object<#rocdl.target<chip = "gfx900">, kernels = #gpu.kernel_table<[
    #gpu.kernel_metadata<"kernel0", (i32) -> (), metadata = {sgpr_count = 255}>,
    #gpu.kernel_metadata<"kernel1", (i32, f32) -> (), arg_attrs = [{llvm.read_only}, {}]>
  ]> , bin = "BLOB">]

```
The motivation behind these attributes is to provide useful information
for things like tunning.

---------

Co-authored-by: Mehdi Amini <[email protected]>

Author: fabianmcg

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/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());
+}