//===-- CompilationAttrs.td - GPU compilation attributes ---*- tablegen -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file defines GPU compilation related attributes.
//
//===----------------------------------------------------------------------===//
#ifndef GPU_COMPILATION_ATTRS
#define GPU_COMPILATION_ATTRS
include "mlir/Dialect/GPU/IR/GPUBase.td"
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
//===----------------------------------------------------------------------===//
// GPU kernel metadata attribute
//===----------------------------------------------------------------------===//
def GPU_KernelMetadataAttr : GPU_Attr<"KernelMetadata", "kernel_metadata"> {
let description = [{
GPU attribute for storing metadata related to a compiled kernel. The
attribute contains the name and arguments 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_metadata<@kernel1, (i32) -> (), arg_attrs = [...], metadata = {reg_count = 255, ...}>
#gpu.kernel_metadata<@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 KernelMetadataAttr& 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 KernelMetadataAttr that contains `attrs` in the metadata dictionary.
KernelMetadataAttr 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_metadata` 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_metadata<kernel0, () -> () >]>
// Table with multiple kernels.
#gpu.kernel_table<[
#gpu.kernel_metadata<"kernel0", (i32, f32) -> (), metadata = {sgpr_count = 255}>,
#gpu.kernel_metadata<"kernel1", (i32) -> ()>
]>
```
}];
let parameters = (ins
OptionalArrayRefParameter<"KernelMetadataAttr", "array of kernels">:$kernel_table
);
let assemblyFormat = [{
`<` (`[` qualified($kernel_table)^ `]`)? `>`
}];
let builders = [
AttrBuilder<(ins "ArrayRef<KernelMetadataAttr>":$kernels,
CArg<"bool", "false">:$isSorted)>
];
let skipDefaultBuilders = 1;
let genVerifyDecl = 1;
let extraClassDeclaration = [{
llvm::ArrayRef<KernelMetadataAttr>::iterator begin() const {
return getKernelTable().begin();
}
llvm::ArrayRef<KernelMetadataAttr>::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.
KernelMetadataAttr lookup(StringRef key) const;
KernelMetadataAttr lookup(StringAttr key) const;
}];
}
//===----------------------------------------------------------------------===//
// GPU object attribute.
//===----------------------------------------------------------------------===//
// For documentation on this enum cases, see the `GPU_ObjectAttr` docs.
def GPU_ObjectOffload : I32EnumAttrCase<"Offload", 1, "offload">;
def GPU_ObjectISA : I32EnumAttrCase<"Assembly", 2, "assembly">;
def GPU_ObjectBinary : I32EnumAttrCase<"Binary", 3, "bin">;
def GPU_ObjectFatbin : I32EnumAttrCase<"Fatbin", 4, "fatbin">;
def GPU_CompilationTargetEnum : GPU_I32Enum<
"CompilationTarget", "GPU compilation format", [
GPU_ObjectOffload,
GPU_ObjectISA,
GPU_ObjectBinary,
GPU_ObjectFatbin
]>;
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, 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
formats, and its meaning is target-dependent. For example, on the NVPTX and
AMDGPU targets, this format is associated with LLVM bitcode.
2. `Assembly`: represents GPU assembly code. For example, in the NVPTX
target, assembly is PTX code, which can be JITted at runtime.
3. `Binary`: represents executable code for a GPU single architecture. For
example, PTX code that was compiled for a specific compute capability. Note
that this format is likely to throw an error if there is an architecture
mismatch between the compiled and running architecture.
4. `Fatbin`: represents a GPU fat binary with executable code for multiple
architectures. This format is the default; thus, it gets elided inassembly
code.
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.
```
#gpu.object<#rocdl.target, offload = "..."> // An offload 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<"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^ `,`)?
(`kernels` `=` $kernels^ `,`)?
custom<Object>($format, $object)
`>`
}];
let genVerifyDecl = 1;
}
def GPUObjectArrayAttr :
TypedArrayAttrBase<GPU_ObjectAttr, "an array of GPU object attributes">;
//===----------------------------------------------------------------------===//
// GPU offloading LLVM translation handler attributes.
//===----------------------------------------------------------------------===//
def GPU_SelectObjectAttr : GPU_Attr<"SelectObject", "select_object", [
OffloadingTranslationAttrTrait
]> {
let description = [{
This GPU offloading handler selects a single GPU object for embedding. The
object is selected based on the `target` parameter, this parameter can be
either a number -i.e. selects the ith-target, or the target itself -i.e.
searches for the specified target in the object array.
The first object in a `gpu.binary` operation is selected if no target is
specified.
}];
let parameters = (ins
OptionalParameter<"Attribute", "Target to select for embedding.">:$target
);
let assemblyFormat = [{
(`<` $target^ `>`)?
}];
let genVerifyDecl = 1;
}
#endif // GPU_COMPILATION_ATTRS