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,10 @@ } 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 @@ -1019,15 +1020,21 @@ or not intended to be run on the separate device. ``` - gpu.module @symbol_name { + gpu.module @symbol_name { + gpu.func {} + ... + gpu.module_end + } + gpu.module @symbol_name2 [#gpu.amdgpu] { gpu.func {} ... gpu.module_end } - ``` }]; - let builders = [OpBuilder<(ins "StringRef":$name)>]; + let builders = [ + OpBuilder<(ins "StringRef":$name, CArg<"ArrayAttr", "{}">:$targets)> + ]; 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 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -1456,18 +1456,35 @@ //===----------------------------------------------------------------------===// 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; } 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. @@ -1483,8 +1500,16 @@ 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); 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] { + } +}