diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -15,6 +15,7 @@ include "mlir/Dialect/DLTI/DLTIBase.td" include "mlir/Dialect/GPU/IR/GPUBase.td" +include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td" include "mlir/Dialect/GPU/IR/ParallelLoopMapperAttr.td" include "mlir/Dialect/GPU/TransformOps/GPUDeviceMappingAttr.td" include "mlir/IR/EnumAttr.td" @@ -998,10 +999,9 @@ } def GPU_GPUModuleOp : GPU_Op<"module", [ - DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove, - SymbolTable, Symbol, - SingleBlockImplicitTerminator<"ModuleEndOp"> -]> { + DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove, + SymbolTable, Symbol, SingleBlockImplicitTerminator<"ModuleEndOp"> + ]>, Arguments<(ins OptionalAttr:$targets)> { 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 @@ -1018,22 +1018,41 @@ allows filtering of code regions to execute passes on only code intended to or not intended to be run on the separate device. + Modules can contain zero or more target attributes. These attributes encode + how to transform modules into binary strings and are used by the + `gpu-module-to-binary` pass to transform modules into GPU binaries. + ``` - gpu.module @symbol_name { + gpu.module @symbol_name { + gpu.func {} + ... + gpu.module_end + } + gpu.module @symbol_name2 [#nvvm.target, #rocdl.target] { gpu.func {} ... gpu.module_end } - ``` }]; - let builders = [OpBuilder<(ins "StringRef":$name)>]; + let builders = [ + OpBuilder<(ins "StringRef":$name, CArg<"ArrayAttr", "{}">:$targets)>, + OpBuilder<(ins "StringRef":$name, "ArrayRef":$targets)> + ]; let regions = (region SizedRegion<1>:$bodyRegion); let hasCustomAssemblyFormat = 1; // We need to ensure the block inside the region is properly terminated; // the auto-generated builders do not guarantee that. let skipDefaultBuilders = 1; + + let extraClassDeclaration = [{ + /// Checks if `target` is in the `targets` list. + bool hasTarget(Attribute target); + + /// Sets the targets of the module. + void setTargets(ArrayRef targets); + }]; } def GPU_ModuleEndOp : GPU_Op<"module_end", [ diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -1500,18 +1500,41 @@ //===----------------------------------------------------------------------===// void GPUModuleOp::build(OpBuilder &builder, OperationState &result, - StringRef name) { + StringRef name, ArrayAttr targets) { ensureTerminator(*result.addRegion(), builder, result.location); result.attributes.push_back(builder.getNamedAttr( ::mlir::SymbolTable::getSymbolAttrName(), builder.getStringAttr(name))); + + if (targets) + result.getOrAddProperties().targets = targets; +} + +void GPUModuleOp::build(OpBuilder &builder, OperationState &result, + StringRef name, ArrayRef targets) { + build(builder, result, name, + targets.size() > 0 ? builder.getArrayAttr(targets) : ArrayAttr()); } ParseResult GPUModuleOp::parse(OpAsmParser &parser, OperationState &result) { StringAttr nameAttr; + ArrayAttr targetsAttr; + if (parser.parseSymbolName(nameAttr, mlir::SymbolTable::getSymbolAttrName(), - result.attributes) || - // If module attributes are present, parse them. - parser.parseOptionalAttrDictWithKeyword(result.attributes)) + result.attributes)) + return failure(); + + // Parse the optional array of target attributes. + OptionalParseResult targetsAttrResult = + parser.parseOptionalAttribute(targetsAttr, Type{}); + if (targetsAttrResult.has_value()) { + if (failed(*targetsAttrResult)) { + return failure(); + } + result.getOrAddProperties().targets = targetsAttr; + } + + // If module attributes are present, parse them. + if (parser.parseOptionalAttrDictWithKeyword(result.attributes)) return failure(); // Parse the module body. @@ -1527,13 +1550,33 @@ void GPUModuleOp::print(OpAsmPrinter &p) { p << ' '; p.printSymbolName(getName()); - p.printOptionalAttrDictWithKeyword((*this)->getAttrs(), - {mlir::SymbolTable::getSymbolAttrName()}); + + if (Attribute attr = getTargetsAttr()) { + p << ' '; + p.printAttribute(attr); + p << ' '; + } + + p.printOptionalAttrDictWithKeyword( + (*this)->getAttrs(), + {mlir::SymbolTable::getSymbolAttrName(), getTargetsAttrName()}); p << ' '; p.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false); } +bool GPUModuleOp::hasTarget(Attribute target) { + if (ArrayAttr targets = getTargetsAttr()) + return llvm::count(targets.getValue(), target); + return false; +} + +void GPUModuleOp::setTargets(ArrayRef targets) { + ArrayAttr &targetsAttr = getProperties().targets; + SmallVector targetsVector(targets); + targetsAttr = ArrayAttr::get(getContext(), targetsVector); +} + //===----------------------------------------------------------------------===// // GPUMemcpyOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir --- a/mlir/test/Dialect/GPU/invalid.mlir +++ b/mlir/test/Dialect/GPU/invalid.mlir @@ -610,3 +610,19 @@ } } } + +// ----- + +module { + // expected-error @+1 {{'gpu.module' op attribute 'targets' failed to satisfy constraint: array of GPU target attributes with at least 1 elements}} + gpu.module @gpu_funcs [] { + } +} + +// ----- + +module { + // expected-error @+1 {{'gpu.module' op attribute 'targets' failed to satisfy constraint: array of GPU target attributes with at least 1 elements}} + gpu.module @gpu_funcs [1] { + } +}