[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 = [{
+    /// 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
----------------
ftynse wrote:

Ultra-nit: trailing dot, here and below.

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