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

[Oleksandr Alex Zinenko]

================
@@ -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 = [{
+    /// Returns the metadata attribute corresponding to `key` or `nullptr`
+    /// if missing.
+    Attribute getAttr(StringRef key) const {
+      auto 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 {
+      auto 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) {
+      auto 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) {
+      auto argDict = getArgAttrDict(index);
+      return argDict ? argDict.get(name) : nullptr;
+    }
+    Attribute getArgAttr(unsigned index, StringRef name) {
+      auto 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 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 {
----------------
ftynse wrote:

Nit: `KernelIterator` should work as return type here.

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