Skip to content

Commit 7be061e

Browse files
fabianmcgfabianmcg
committed
move to gpu
1 parent ca33796 commit 7be061e

File tree

12 files changed

+481
-11
lines changed

12 files changed

+481
-11
lines changed

‎mlir/include/mlir-c/Dialect/GPU.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ MLIR_CAPI_EXPORTED bool mlirAttributeIsAGPUObjectAttr(MlirAttribute attr);
2727

2828
MLIR_CAPI_EXPORTED MlirAttribute
2929
mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target, uint32_t format,
30-
MlirStringRef objectStrRef, MlirAttribute mlirObjectProps);
30+
MlirStringRef objectStrRef, MlirAttribute mlirObjectProps,
31+
MlirAttribute mlirKernelsAttr);
3132

3233
MLIR_CAPI_EXPORTED MlirAttribute
3334
mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr);

‎mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td

Lines changed: 145 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,136 @@
1616
include "mlir/Dialect/GPU/IR/GPUBase.td"
1717
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
1818

19+
//===----------------------------------------------------------------------===//
20+
// GPU kernel attribute
21+
//===----------------------------------------------------------------------===//
22+
23+
def GPU_KernelAttr : GPU_Attr<"Kernel", "kernel"> {
24+
let description = [{
25+
GPU attribute for storing metadata related to a compiled kernel. It
26+
contains the attribute dictionary of the LLVM function used to generate the
27+
kernel, as well as an optional dictionary for additional metadata, like
28+
occupancy information.
29+
30+
Examples:
31+
```mlir
32+
#gpu.kernel<{sym_name = "test_fusion__part_0", ...},
33+
metadata = {reg_count = 255, ...}>
34+
```
35+
}];
36+
let parameters = (ins
37+
"DictionaryAttr":$func_attrs,
38+
OptionalParameter<"DictionaryAttr", "metadata dictionary">:$metadata
39+
);
40+
let assemblyFormat = [{
41+
`<` $func_attrs (`,` `metadata` `=` $metadata^ )? `>`
42+
}];
43+
let builders = [
44+
AttrBuilderWithInferredContext<(ins "DictionaryAttr":$funcAttrs,
45+
CArg<"DictionaryAttr",
46+
"nullptr">:$metadata), [{
47+
assert(funcAttrs && "invalid function attributes dictionary");
48+
return $_get(funcAttrs.getContext(), funcAttrs, metadata);
49+
}]>,
50+
AttrBuilderWithInferredContext<(ins "Operation*":$kernel,
51+
CArg<"DictionaryAttr",
52+
"nullptr">:$metadata)>
53+
];
54+
let extraClassDeclaration = [{
55+
/// Returns the function attribute corresponding to key or nullptr if missing.
56+
Attribute getAttr(StringRef key) const {
57+
return getFuncAttrs().get(key);
58+
}
59+
template <typename ConcreteAttr>
60+
ConcreteAttr getAttr(StringRef key) const {
61+
return llvm::dyn_cast_or_null<ConcreteAttr>(getAttr(key));
62+
}
63+
Attribute getAttr(StringAttr key) const;
64+
template <typename ConcreteAttr>
65+
ConcreteAttr getAttr(StringAttr key) const {
66+
return llvm::dyn_cast_or_null<ConcreteAttr>(getAttr(key));
67+
}
68+
69+
/// Returns the name of the kernel.
70+
StringAttr getName() const {
71+
return getAttr<StringAttr>("sym_name");
72+
}
73+
74+
/// Returns the metadta attribute corresponding to key or nullptr if missing.
75+
Attribute getMDAttr(StringRef key) const {
76+
if (DictionaryAttr attrs = getMetadata())
77+
return attrs.get(key);
78+
return nullptr;
79+
}
80+
template <typename ConcreteAttr>
81+
ConcreteAttr getMDAttr(StringRef key) const {
82+
return llvm::dyn_cast_or_null<ConcreteAttr>(getMDAttr(key));
83+
}
84+
Attribute getMDAttr(StringAttr key) const;
85+
template <typename ConcreteAttr>
86+
ConcreteAttr getMDAttr(StringAttr key) const {
87+
return llvm::dyn_cast_or_null<ConcreteAttr>(getMDAttr(key));
88+
}
89+
90+
/// Helper function for appending metadata to a kernel attribute.
91+
KernelAttr appendMetadata(ArrayRef<NamedAttribute> attrs) const;
92+
}];
93+
}
94+
95+
//===----------------------------------------------------------------------===//
96+
// GPU kernel table attribute
97+
//===----------------------------------------------------------------------===//
98+
99+
def GPU_KernelTableAttr : GPU_Attr<"KernelTable", "kernel_table"> {
100+
let description = [{
101+
GPU attribute representing a table of kernels metadata. All the attributes
102+
in the dictionary must be of type `#gpu.kernel`.
103+
104+
Examples:
105+
```mlir
106+
#gpu.kernel_table<{kernel0 = #gpu.kernel<...>}>
107+
```
108+
}];
109+
let parameters = (ins
110+
"DictionaryAttr":$kernel_table
111+
);
112+
let assemblyFormat = [{
113+
`<` $kernel_table `>`
114+
}];
115+
let builders = [
116+
AttrBuilderWithInferredContext<(ins "DictionaryAttr":$kernel_table), [{
117+
assert(kernel_table && "invalid kernel table");
118+
return $_get(kernel_table.getContext(), kernel_table);
119+
}]>
120+
];
121+
let skipDefaultBuilders = 1;
122+
let genVerifyDecl = 1;
123+
let extraClassDeclaration = [{
124+
/// Helper iterator class for traversing the kernel table.
125+
struct KernelIterator
126+
: llvm::mapped_iterator_base<KernelIterator,
127+
llvm::ArrayRef<NamedAttribute>::iterator,
128+
std::pair<StringAttr, KernelAttr>> {
129+
using llvm::mapped_iterator_base<
130+
KernelIterator, llvm::ArrayRef<NamedAttribute>::iterator,
131+
std::pair<StringAttr, KernelAttr>>::mapped_iterator_base;
132+
/// Map the iterator to the kernel name and a KernelAttribute.
133+
std::pair<StringAttr, KernelAttr> mapElement(NamedAttribute attr) const {
134+
return {attr.getName(), llvm::cast<KernelAttr>(attr.getValue())};
135+
}
136+
};
137+
auto begin() const {
138+
return KernelIterator(getKernelTable().begin());
139+
}
140+
auto end() const {
141+
return KernelIterator(getKernelTable().end());
142+
}
143+
size_t size() const {
144+
return getKernelTable().size();
145+
}
146+
}];
147+
}
148+
19149
//===----------------------------------------------------------------------===//
20150
// GPU object attribute.
21151
//===----------------------------------------------------------------------===//
@@ -63,16 +193,29 @@ def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
63193
#gpu.object<#nvvm.target, properties = {O = 3 : i32}, assembly = "..."> // An assembly object with additional properties.
64194
#gpu.object<#rocdl.target, bin = "..."> // A binary object.
65195
#gpu.object<#nvvm.target, "..."> // A fatbin object.
196+
#gpu.object<#nvvm.target, kernels = #gpu.kernel_table<...>, "..."> // An object with a kernel table.
66197
```
67198
}];
68199
let parameters = (ins
69200
"Attribute":$target,
70201
DefaultValuedParameter<"CompilationTarget", "CompilationTarget::Fatbin">:$format,
71202
"StringAttr":$object,
72-
OptionalParameter<"DictionaryAttr">:$properties
203+
OptionalParameter<"DictionaryAttr">:$properties,
204+
OptionalParameter<"KernelTableAttr">:$kernels
73205
);
206+
let builders = [
207+
AttrBuilderWithInferredContext<(ins "Attribute":$target,
208+
"CompilationTarget":$format,
209+
"StringAttr":$object,
210+
CArg<"DictionaryAttr", "nullptr">:$properties,
211+
CArg<"KernelTableAttr", "nullptr">:$kernels), [{
212+
assert(target && "invalid target");
213+
return $_get(target.getContext(), target, format, object, properties, kernels);
214+
}]>
215+
];
74216
let assemblyFormat = [{ `<`
75-
$target `,` (`properties` `=` $properties ^ `,`)?
217+
$target `,` (`properties` `=` $properties^ `,`)?
218+
(`kernels` `=` $kernels^ `,`)?
76219
custom<Object>($format, $object)
77220
`>`
78221
}];

‎mlir/include/mlir/Target/LLVM/ROCDL/Utils.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#define MLIR_TARGET_LLVM_ROCDL_UTILS_H
1515

1616
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
17+
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
1718
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
1819
#include "mlir/Target/LLVM/ModuleToObject.h"
1920

@@ -85,6 +86,12 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
8586
/// List of LLVM bitcode files to link to.
8687
SmallVector<std::string> fileList;
8788
};
89+
90+
/// Returns a `#gpu.kernel_table` containing kernel metadata for each of the
91+
/// kernels in `gpuModule`. If `elfData` is valid, then the `amdhsa.kernels` ELF
92+
/// metadata will be added to the `#gpu.kernel_table`.
93+
gpu::KernelTableAttr getAMDHSAKernelsMetadata(Operation *gpuModule,
94+
ArrayRef<char> elfData = {});
8895
} // namespace ROCDL
8996
} // namespace mlir
9097

‎mlir/lib/CAPI/Dialect/GPU.cpp

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,19 @@ bool mlirAttributeIsAGPUObjectAttr(MlirAttribute attr) {
2525

2626
MlirAttribute mlirGPUObjectAttrGet(MlirContext mlirCtx, MlirAttribute target,
2727
uint32_t format, MlirStringRef objectStrRef,
28-
MlirAttribute mlirObjectProps) {
28+
MlirAttribute mlirObjectProps,
29+
MlirAttribute mlirKernelsAttr) {
2930
MLIRContext *ctx = unwrap(mlirCtx);
3031
llvm::StringRef object = unwrap(objectStrRef);
3132
DictionaryAttr objectProps;
3233
if (mlirObjectProps.ptr != nullptr)
3334
objectProps = llvm::cast<DictionaryAttr>(unwrap(mlirObjectProps));
34-
return wrap(gpu::ObjectAttr::get(ctx, unwrap(target),
35-
static_cast<gpu::CompilationTarget>(format),
36-
StringAttr::get(ctx, object), objectProps));
35+
gpu::KernelTableAttr kernels;
36+
if (mlirKernelsAttr.ptr != nullptr)
37+
kernels = llvm::cast<gpu::KernelTableAttr>(unwrap(mlirKernelsAttr));
38+
return wrap(gpu::ObjectAttr::get(
39+
ctx, unwrap(target), static_cast<gpu::CompilationTarget>(format),
40+
StringAttr::get(ctx, object), objectProps, kernels));
3741
}
3842

3943
MlirAttribute mlirGPUObjectAttrGetTarget(MlirAttribute mlirObjectAttr) {

‎mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2140,7 +2140,8 @@ void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
21402140

21412141
LogicalResult ObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
21422142
Attribute target, CompilationTarget format,
2143-
StringAttr object, DictionaryAttr properties) {
2143+
StringAttr object, DictionaryAttr properties,
2144+
KernelTableAttr kernels) {
21442145
if (!target)
21452146
return emitError() << "the target attribute cannot be null";
21462147
if (target.hasPromiseOrImplementsInterface<TargetAttrInterface>())
@@ -2226,6 +2227,40 @@ LogicalResult gpu::DynamicSharedMemoryOp::verify() {
22262227
return success();
22272228
}
22282229

2230+
//===----------------------------------------------------------------------===//
2231+
// GPU KernelAttr
2232+
//===----------------------------------------------------------------------===//
2233+
2234+
KernelAttr KernelAttr::get(Operation *kernelOp, DictionaryAttr metadata) {
2235+
assert(kernelOp && "invalid kernel");
2236+
return get(kernelOp->getAttrDictionary(), metadata);
2237+
}
2238+
2239+
KernelAttr KernelAttr::appendMetadata(ArrayRef<NamedAttribute> attrs) const {
2240+
if (attrs.empty())
2241+
return *this;
2242+
NamedAttrList attrList(attrs);
2243+
attrList.append(getMetadata());
2244+
return KernelAttr::get(getFuncAttrs(), attrList.getDictionary(getContext()));
2245+
}
2246+
2247+
//===----------------------------------------------------------------------===//
2248+
// GPU KernelTableAttr
2249+
//===----------------------------------------------------------------------===//
2250+
2251+
LogicalResult
2252+
KernelTableAttr::verify(function_ref<InFlightDiagnostic()> emitError,
2253+
DictionaryAttr dict) {
2254+
if (!dict)
2255+
return emitError() << "table cannot be null";
2256+
if (llvm::any_of(dict, [](NamedAttribute attr) {
2257+
return !llvm::isa<KernelAttr>(attr.getValue());
2258+
}))
2259+
return emitError()
2260+
<< "all the dictionary values must be `#gpu.kernel` attributes";
2261+
return success();
2262+
}
2263+
22292264
//===----------------------------------------------------------------------===//
22302265
// GPU target options
22312266
//===----------------------------------------------------------------------===//

‎mlir/lib/Target/LLVM/CMakeLists.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,9 +108,11 @@ endif()
108108

109109
add_mlir_dialect_library(MLIRROCDLTarget
110110
ROCDL/Target.cpp
111+
ROCDL/Utils.cpp
111112

112113
LINK_COMPONENTS
113114
MCParser
115+
Object
114116
${AMDGPU_LIBS}
115117

116118
LINK_LIBS PUBLIC

‎mlir/lib/Target/LLVM/NVVM/Target.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -606,5 +606,5 @@ NVVMTargetAttrImpl::createObject(Attribute attribute,
606606
return builder.getAttr<gpu::ObjectAttr>(
607607
attribute, format,
608608
builder.getStringAttr(StringRef(object.data(), object.size())),
609-
objectProps);
609+
objectProps, nullptr);
610610
}

‎mlir/lib/Target/LLVM/ROCDL/Target.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -482,5 +482,6 @@ ROCDLTargetAttrImpl::createObject(Attribute attribute,
482482
attribute,
483483
format > gpu::CompilationTarget::Binary ? gpu::CompilationTarget::Binary
484484
: format,
485-
builder.getStringAttr(StringRef(object.data(), object.size())), nullptr);
485+
builder.getStringAttr(StringRef(object.data(), object.size())), nullptr,
486+
nullptr);
486487
}

0 commit comments

Comments
 (0)