[Mlir-commits] [mlir] [mlir][gpu] Add an offloading handler attribute to `gpu.module` (PR #78047)
Fabian Mora
llvmlistbot at llvm.org
Sat Jan 13 08:23:07 PST 2024
https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/78047
This patch adds an optional offloading handler attribute to the`gpu.module` op. This attribute will be used during `gpu-module-to-binary` pass to override the offloading handler used in the `gpu.binary` op.
>From b45c496015a9c6e1ceb25c5b556eb2daa60b8bd2 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sat, 13 Jan 2024 16:12:14 +0000
Subject: [PATCH] [mlir][gpu] Add an offloading handler attribute to
`gpu.module`
This patch adds an optional offloading handler attribute to the`gpu.module` op.
This attribute will be used during `gpu-module-to-binary` pass to override the
offloading handler used in the `gpu.binary` op.
---
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 21 ++++++++--
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp | 38 +++++++++++++++----
.../Dialect/GPU/Transforms/ModuleToBinary.cpp | 5 +++
mlir/test/Dialect/GPU/invalid.mlir | 7 ++++
.../Dialect/GPU/module-to-binary-nvvm.mlir | 12 ++++++
mlir/test/Dialect/GPU/ops.mlir | 3 ++
6 files changed, 74 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 8d4a110ee801f0..228aad3d84629c 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -1191,7 +1191,9 @@ def GPU_BarrierOp : GPU_Op<"barrier"> {
def GPU_GPUModuleOp : GPU_Op<"module", [
DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove,
SymbolTable, Symbol, SingleBlockImplicitTerminator<"ModuleEndOp">
- ]>, Arguments<(ins OptionalAttr<GPUNonEmptyTargetArrayAttr>:$targets)> {
+ ]>, Arguments<(ins
+ OptionalAttr<GPUNonEmptyTargetArrayAttr>:$targets,
+ OptionalAttr<OffloadingTranslationAttr>:$offloadingHandler)> {
let summary = "A top level compilation unit containing code to be run on a GPU.";
let description = [{
GPU module contains code that is intended to be run on a GPU. A host device
@@ -1212,13 +1214,20 @@ def GPU_GPUModuleOp : GPU_Op<"module", [
how to transform modules into binary strings and are used by the
`gpu-module-to-binary` pass to transform modules into GPU binaries.
+ Modules can contain an optional `OffloadingTranslationAttr` attribute. This
+ attribute will be used during the `gpu-module-to-binary` pass to specify the
+ `OffloadingTranslationAttr` used when creating the `gpu.binary` operation.
+
```
gpu.module @symbol_name {
gpu.func {}
...
gpu.module_end
}
- gpu.module @symbol_name2 [#nvvm.target, #rocdl.target<chip = "gfx90a">] {
+ // Module with offloading handler and target attributes.
+ gpu.module @symbol_name2 <#gpu.select_object<1>> [
+ #nvvm.target,
+ #rocdl.target<chip = "gfx90a">] {
gpu.func {}
...
gpu.module_end
@@ -1226,8 +1235,12 @@ def GPU_GPUModuleOp : GPU_Op<"module", [
```
}];
let builders = [
- OpBuilder<(ins "StringRef":$name, CArg<"ArrayAttr", "{}">:$targets)>,
- OpBuilder<(ins "StringRef":$name, "ArrayRef<Attribute>":$targets)>
+ OpBuilder<(ins "StringRef":$name,
+ CArg<"ArrayAttr", "{}">:$targets,
+ CArg<"Attribute", "{}">:$handler)>,
+ OpBuilder<(ins "StringRef":$name,
+ "ArrayRef<Attribute>":$targets,
+ CArg<"Attribute", "{}">:$handler)>
];
let regions = (region SizedRegion<1>:$bodyRegion);
let hasCustomAssemblyFormat = 1;
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 020900934c9f72..514b3e9a6e8a56 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -1724,19 +1724,24 @@ LogicalResult gpu::ReturnOp::verify() {
//===----------------------------------------------------------------------===//
void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
- StringRef name, ArrayAttr targets) {
+ StringRef name, ArrayAttr targets,
+ Attribute offloadingHandler) {
ensureTerminator(*result.addRegion(), builder, result.location);
result.attributes.push_back(builder.getNamedAttr(
::mlir::SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));
+ Properties &props = result.getOrAddProperties<Properties>();
if (targets)
- result.getOrAddProperties<Properties>().targets = targets;
+ props.targets = targets;
+ props.offloadingHandler = offloadingHandler;
}
void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
- StringRef name, ArrayRef<Attribute> targets) {
+ StringRef name, ArrayRef<Attribute> targets,
+ Attribute offloadingHandler) {
build(builder, result, name,
- targets.empty() ? ArrayAttr() : builder.getArrayAttr(targets));
+ targets.empty() ? ArrayAttr() : builder.getArrayAttr(targets),
+ offloadingHandler);
}
ParseResult GPUModuleOp::parse(OpAsmParser &parser, OperationState &result) {
@@ -1747,6 +1752,16 @@ ParseResult GPUModuleOp::parse(OpAsmParser &parser, OperationState &result) {
result.attributes))
return failure();
+ Properties &props = result.getOrAddProperties<Properties>();
+
+ // Parse the optional offloadingHandler
+ if (succeeded(parser.parseOptionalLess())) {
+ if (parser.parseAttribute(props.offloadingHandler))
+ return failure();
+ if (parser.parseGreater())
+ return failure();
+ }
+
// Parse the optional array of target attributes.
OptionalParseResult targetsAttrResult =
parser.parseOptionalAttribute(targetsAttr, Type{});
@@ -1754,7 +1769,7 @@ ParseResult GPUModuleOp::parse(OpAsmParser &parser, OperationState &result) {
if (failed(*targetsAttrResult)) {
return failure();
}
- result.getOrAddProperties<Properties>().targets = targetsAttr;
+ props.targets = targetsAttr;
}
// If module attributes are present, parse them.
@@ -1775,15 +1790,22 @@ void GPUModuleOp::print(OpAsmPrinter &p) {
p << ' ';
p.printSymbolName(getName());
+ if (Attribute attr = getOffloadingHandlerAttr()) {
+ p << " <";
+ p.printAttribute(attr);
+ p << ">";
+ }
+
if (Attribute attr = getTargetsAttr()) {
p << ' ';
p.printAttribute(attr);
p << ' ';
}
- p.printOptionalAttrDictWithKeyword(
- (*this)->getAttrs(),
- {mlir::SymbolTable::getSymbolAttrName(), getTargetsAttrName()});
+ p.printOptionalAttrDictWithKeyword((*this)->getAttrs(),
+ {mlir::SymbolTable::getSymbolAttrName(),
+ getTargetsAttrName(),
+ getOffloadingHandlerAttrName()});
p << ' ';
p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
/*printBlockTerminators=*/false);
diff --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
index 70d36297e103f3..0527073da85b69 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
@@ -124,6 +124,11 @@ LogicalResult moduleSerializer(GPUModuleOp op,
}
objects.push_back(object);
}
+ if (auto moduleHandler =
+ dyn_cast_or_null<OffloadingLLVMTranslationAttrInterface>(
+ op.getOffloadingHandlerAttr());
+ !handler && moduleHandler)
+ handler = moduleHandler;
builder.setInsertionPointAfter(op);
builder.create<gpu::BinaryOp>(op.getLoc(), op.getName(), handler,
builder.getArrayAttr(objects));
diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
index 4d3a898fdd1565..273bc282b0b3b0 100644
--- a/mlir/test/Dialect/GPU/invalid.mlir
+++ b/mlir/test/Dialect/GPU/invalid.mlir
@@ -818,3 +818,10 @@ func.func @main(%arg0 : index) {
return
}
+// -----
+
+module attributes {gpu.container_module} {
+ // expected-error at +1 {{expected attribute value}}
+ gpu.module @kernel <> {
+ }
+}
diff --git a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir
index 05e368f7a642e6..c286c8bc9042ff 100644
--- a/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir
+++ b/mlir/test/Dialect/GPU/module-to-binary-nvvm.mlir
@@ -22,4 +22,16 @@ module attributes {gpu.container_module} {
llvm.return
}
}
+
+ // CHECK-LABEL:gpu.binary @kernel_module3 <#gpu.select_object<1 : i64>>
+ // CHECK:[#gpu.object<#nvvm.target<chip = "sm_70">, offload = "{{.*}}">, #gpu.object<#nvvm.target<chip = "sm_80">, offload = "{{.*}}">]
+ gpu.module @kernel_module3 <#gpu.select_object<1>> [
+ #nvvm.target<chip = "sm_70">,
+ #nvvm.target<chip = "sm_80">] {
+ llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr,
+ %arg2: !llvm.ptr, %arg3: i64, %arg4: i64,
+ %arg5: i64) attributes {gpu.kernel} {
+ llvm.return
+ }
+ }
}
diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir
index 60512424383052..488fa7aaf6adca 100644
--- a/mlir/test/Dialect/GPU/ops.mlir
+++ b/mlir/test/Dialect/GPU/ops.mlir
@@ -412,3 +412,6 @@ gpu.module @module_with_two_target [#nvvm.target, #rocdl.target<chip = "gfx90a">
gpu.return
}
}
+
+gpu.module @module_with_offload_handler <#gpu.select_object<0>> [#nvvm.target] {
+}
More information about the Mlir-commits
mailing list