[Mlir-commits] [mlir] [mlir][gpu] Add metadata attributes for storing kernel metadata in GPU objects (PR #95292)
Fabian Mora
llvmlistbot at llvm.org
Tue Jul 30 08:35:25 PDT 2024
================
@@ -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
----------------
fabianmcg wrote:
I changed to a sorted list.
https://github.com/llvm/llvm-project/pull/95292
More information about the Mlir-commits
mailing list