From 5b4f2b906ba697c276b3e0940c24ae640cfde695 Mon Sep 17 00:00:00 2001 From: Fabian Mora Date: Mon, 15 Jan 2024 16:58:10 -0500 Subject: [PATCH] [mlir][gpu] Add an offloading handler attribute to `gpu.module` (#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. --- 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 71f6a2bc5fa2..955dd1e20d24 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -1194,7 +1194,9 @@ def GPU_BarrierOp : GPU_Op<"barrier"> { def GPU_GPUModuleOp : GPU_Op<"module", [ DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove, SymbolTable, Symbol, SingleBlockImplicitTerminator<"ModuleEndOp"> - ]>, Arguments<(ins OptionalAttr:$targets)> { + ]>, Arguments<(ins + OptionalAttr:$targets, + OptionalAttr:$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 @@ -1215,13 +1217,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] { + // Module with offloading handler and target attributes. + gpu.module @symbol_name2 <#gpu.select_object<1>> [ + #nvvm.target, + #rocdl.target] { gpu.func {} ... gpu.module_end @@ -1229,8 +1238,12 @@ def GPU_GPUModuleOp : GPU_Op<"module", [ ``` }]; let builders = [ - OpBuilder<(ins "StringRef":$name, CArg<"ArrayAttr", "{}">:$targets)>, - OpBuilder<(ins "StringRef":$name, "ArrayRef":$targets)> + OpBuilder<(ins "StringRef":$name, + CArg<"ArrayAttr", "{}">:$targets, + CArg<"Attribute", "{}">:$handler)>, + OpBuilder<(ins "StringRef":$name, + "ArrayRef":$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 020900934c9f..514b3e9a6e8a 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(); if (targets) - result.getOrAddProperties().targets = targets; + props.targets = targets; + props.offloadingHandler = offloadingHandler; } void GPUModuleOp::build(OpBuilder &builder, OperationState &result, - StringRef name, ArrayRef targets) { + StringRef name, ArrayRef 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(); + + // 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().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 70d36297e103..0527073da85b 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( + op.getOffloadingHandlerAttr()); + !handler && moduleHandler) + handler = moduleHandler; builder.setInsertionPointAfter(op); builder.create(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 4d3a898fdd15..273bc282b0b3 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@+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 05e368f7a642..c286c8bc9042 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, offload = "{{.*}}">, #gpu.object<#nvvm.target, offload = "{{.*}}">] + gpu.module @kernel_module3 <#gpu.select_object<1>> [ + #nvvm.target, + #nvvm.target] { + 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 5e60d91e4757..8d249c9e9b9b 100644 --- a/mlir/test/Dialect/GPU/ops.mlir +++ b/mlir/test/Dialect/GPU/ops.mlir @@ -423,3 +423,6 @@ gpu.module @module_with_two_target [#nvvm.target, #rocdl.target gpu.return } } + +gpu.module @module_with_offload_handler <#gpu.select_object<0>> [#nvvm.target] { +}