| //===- GPUDialect.cpp - MLIR Dialect for GPU Kernels implementation -------===// |
| // |
| // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file implements the GPU kernel-related dialect and its operations. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/GPU/GPUDialect.h" |
| #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
| #include "mlir/Dialect/StandardOps/Ops.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/Function.h" |
| #include "mlir/IR/FunctionImplementation.h" |
| #include "mlir/IR/Module.h" |
| #include "mlir/IR/OpImplementation.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/IR/StandardTypes.h" |
| |
| using namespace mlir; |
| using namespace mlir::gpu; |
| |
| //===----------------------------------------------------------------------===// |
| // GPUDialect |
| //===----------------------------------------------------------------------===// |
| |
| StringRef GPUDialect::getDialectName() { return "gpu"; } |
| |
| bool GPUDialect::isKernel(Operation *op) { |
| UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName()); |
| return static_cast<bool>(isKernelAttr); |
| } |
| |
| GPUDialect::GPUDialect(MLIRContext *context) |
| : Dialect(getDialectName(), context) { |
| addOperations< |
| #define GET_OP_LIST |
| #include "mlir/Dialect/GPU/GPUOps.cpp.inc" |
| >(); |
| } |
| |
| LogicalResult GPUDialect::verifyOperationAttribute(Operation *op, |
| NamedAttribute attr) { |
| if (!attr.second.isa<UnitAttr>() || |
| !attr.first.is(getContainerModuleAttrName())) |
| return success(); |
| |
| auto module = dyn_cast<ModuleOp>(op); |
| if (!module) |
| return op->emitError("expected '") |
| << getContainerModuleAttrName() << "' attribute to be attached to '" |
| << ModuleOp::getOperationName() << '\''; |
| |
| auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult { |
| // Ignore launches that are nested more or less deep than functions in the |
| // module we are currently checking. |
| if (!launchOp.getParentOp() || |
| launchOp.getParentOp()->getParentOp() != module) |
| return success(); |
| |
| // Ignore launch ops with missing attributes here. The errors will be |
| // reported by the verifiers of those ops. |
| if (!launchOp.getAttrOfType<StringAttr>( |
| LaunchFuncOp::getKernelAttrName()) || |
| !launchOp.getAttrOfType<SymbolRefAttr>( |
| LaunchFuncOp::getKernelModuleAttrName())) |
| return success(); |
| |
| // Check that `launch_func` refers to a well-formed GPU kernel module. |
| StringRef kernelModuleName = launchOp.getKernelModuleName(); |
| auto kernelModule = module.lookupSymbol<GPUModuleOp>(kernelModuleName); |
| if (!kernelModule) |
| return launchOp.emitOpError() |
| << "kernel module '" << kernelModuleName << "' is undefined"; |
| |
| // Check that `launch_func` refers to a well-formed kernel function. |
| StringRef kernelName = launchOp.kernel(); |
| Operation *kernelFunc = kernelModule.lookupSymbol(kernelName); |
| auto kernelGPUFunction = dyn_cast_or_null<gpu::GPUFuncOp>(kernelFunc); |
| auto kernelLLVMFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(kernelFunc); |
| if (!kernelGPUFunction && !kernelLLVMFunction) |
| return launchOp.emitOpError("kernel function '") |
| << kernelName << "' is undefined"; |
| if (!kernelFunc->getAttrOfType<mlir::UnitAttr>( |
| GPUDialect::getKernelFuncAttrName())) |
| return launchOp.emitOpError("kernel function is missing the '") |
| << GPUDialect::getKernelFuncAttrName() << "' attribute"; |
| |
| unsigned actualNumArguments = launchOp.getNumKernelOperands(); |
| unsigned expectedNumArguments = kernelLLVMFunction |
| ? kernelLLVMFunction.getNumArguments() |
| : kernelGPUFunction.getNumArguments(); |
| if (expectedNumArguments != actualNumArguments) |
| return launchOp.emitOpError("got ") |
| << actualNumArguments << " kernel operands but expected " |
| << expectedNumArguments; |
| |
| // Due to the ordering of the current impl of lowering and LLVMLowering, |
| // type checks need to be temporarily disabled. |
| // TODO(ntv,zinenko,herhut): reactivate checks once "changing gpu.launchFunc |
| // to encode target module" has landed. |
| // auto functionType = kernelFunc.getType(); |
| // for (unsigned i = 0; i < numKernelFuncArgs; ++i) { |
| // if (getKernelOperand(i).getType() != functionType.getInput(i)) { |
| // return emitOpError("type of function argument ") |
| // << i << " does not match"; |
| // } |
| // } |
| |
| return success(); |
| }); |
| |
| return walkResult.wasInterrupted() ? failure() : success(); |
| } |
| |
| template <typename T> static LogicalResult verifyIndexOp(T op) { |
| auto dimension = op.dimension(); |
| if (dimension != "x" && dimension != "y" && dimension != "z") |
| return op.emitError("dimension \"") << dimension << "\" is invalid"; |
| return success(); |
| } |
| |
| static LogicalResult verifyAllReduce(gpu::AllReduceOp allReduce) { |
| if (allReduce.body().empty() != allReduce.op().hasValue()) |
| return allReduce.emitError( |
| "expected either an op attribute or a non-empty body"); |
| if (!allReduce.body().empty()) { |
| if (allReduce.body().front().getNumArguments() != 2) |
| return allReduce.emitError("expected two region arguments"); |
| for (auto argument : allReduce.body().front().getArguments()) { |
| if (argument.getType() != allReduce.getType()) |
| return allReduce.emitError("incorrect region argument type"); |
| } |
| unsigned yieldCount = 0; |
| for (Block &block : allReduce.body()) { |
| if (auto yield = dyn_cast<gpu::YieldOp>(block.getTerminator())) { |
| if (yield.getNumOperands() != 1) |
| return allReduce.emitError("expected one gpu.yield operand"); |
| if (yield.getOperand(0).getType() != allReduce.getType()) |
| return allReduce.emitError("incorrect gpu.yield type"); |
| ++yieldCount; |
| } |
| } |
| if (yieldCount == 0) |
| return allReduce.emitError("expected gpu.yield op in region"); |
| } |
| return success(); |
| } |
| |
| static LogicalResult verifyShuffleOp(gpu::ShuffleOp shuffleOp) { |
| auto type = shuffleOp.value().getType(); |
| if (shuffleOp.result().getType() != type) { |
| return shuffleOp.emitOpError() |
| << "requires the same type for value operand and result"; |
| } |
| if (!type.isIntOrFloat() || type.getIntOrFloatBitWidth() != 32) { |
| return shuffleOp.emitOpError() |
| << "requires value operand type to be f32 or i32"; |
| } |
| return success(); |
| } |
| |
| static void printShuffleOp(OpAsmPrinter &p, ShuffleOp op) { |
| p << ShuffleOp::getOperationName() << ' ' << op.getOperands() << ' ' |
| << op.mode() << " : " << op.value().getType(); |
| } |
| |
| static ParseResult parseShuffleOp(OpAsmParser &parser, OperationState &state) { |
| SmallVector<OpAsmParser::OperandType, 3> operandInfo; |
| if (parser.parseOperandList(operandInfo, 3)) |
| return failure(); |
| |
| StringRef mode; |
| if (parser.parseKeyword(&mode)) |
| return failure(); |
| state.addAttribute("mode", parser.getBuilder().getStringAttr(mode)); |
| |
| Type valueType; |
| Type int32Type = parser.getBuilder().getIntegerType(32); |
| Type int1Type = parser.getBuilder().getI1Type(); |
| if (parser.parseColonType(valueType) || |
| parser.resolveOperands(operandInfo, {valueType, int32Type, int32Type}, |
| parser.getCurrentLocation(), state.operands) || |
| parser.addTypesToList({valueType, int1Type}, state.types)) |
| return failure(); |
| return success(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // LaunchOp |
| //===----------------------------------------------------------------------===// |
| |
| void LaunchOp::build(Builder *builder, OperationState &result, Value gridSizeX, |
| Value gridSizeY, Value gridSizeZ, Value blockSizeX, |
| Value blockSizeY, Value blockSizeZ, ValueRange operands) { |
| // Add grid and block sizes as op operands, followed by the data operands. |
| result.addOperands( |
| {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ}); |
| result.addOperands(operands); |
| |
| // Create a kernel body region with kNumConfigRegionAttributes + N arguments, |
| // where the first kNumConfigRegionAttributes arguments have `index` type and |
| // the rest have the same types as the data operands. |
| Region *kernelRegion = result.addRegion(); |
| Block *body = new Block(); |
| body->addArguments( |
| std::vector<Type>(kNumConfigRegionAttributes, builder->getIndexType())); |
| body->addArguments(llvm::to_vector<4>(operands.getTypes())); |
| kernelRegion->push_back(body); |
| } |
| |
| KernelDim3 LaunchOp::getBlockIds() { |
| assert(!body().getBlocks().empty() && "FuncOp body must not be empty."); |
| auto args = body().getBlocks().front().getArguments(); |
| return KernelDim3{args[0], args[1], args[2]}; |
| } |
| |
| KernelDim3 LaunchOp::getThreadIds() { |
| assert(!body().getBlocks().empty() && "FuncOp body must not be empty."); |
| auto args = body().getBlocks().front().getArguments(); |
| return KernelDim3{args[3], args[4], args[5]}; |
| } |
| |
| KernelDim3 LaunchOp::getGridSize() { |
| assert(!body().getBlocks().empty() && "FuncOp body must not be empty."); |
| auto args = body().getBlocks().front().getArguments(); |
| return KernelDim3{args[6], args[7], args[8]}; |
| } |
| |
| KernelDim3 LaunchOp::getBlockSize() { |
| assert(!body().getBlocks().empty() && "FuncOp body must not be empty."); |
| auto args = body().getBlocks().front().getArguments(); |
| return KernelDim3{args[9], args[10], args[11]}; |
| } |
| |
| LaunchOp::operand_range LaunchOp::getKernelOperandValues() { |
| return llvm::drop_begin(getOperands(), kNumConfigOperands); |
| } |
| |
| LaunchOp::operand_type_range LaunchOp::getKernelOperandTypes() { |
| return llvm::drop_begin(getOperandTypes(), kNumConfigOperands); |
| } |
| |
| KernelDim3 LaunchOp::getGridSizeOperandValues() { |
| return KernelDim3{getOperand(0), getOperand(1), getOperand(2)}; |
| } |
| |
| KernelDim3 LaunchOp::getBlockSizeOperandValues() { |
| return KernelDim3{getOperand(3), getOperand(4), getOperand(5)}; |
| } |
| |
| iterator_range<Block::args_iterator> LaunchOp::getKernelArguments() { |
| auto args = body().getBlocks().front().getArguments(); |
| return llvm::drop_begin(args, LaunchOp::kNumConfigRegionAttributes); |
| } |
| |
| static LogicalResult verify(LaunchOp op) { |
| // Kernel launch takes kNumConfigOperands leading operands for grid/block |
| // sizes and transforms them into kNumConfigRegionAttributes region arguments |
| // for block/thread identifiers and grid/block sizes. |
| if (!op.body().empty()) { |
| Block &entryBlock = op.body().front(); |
| if (entryBlock.getNumArguments() != |
| LaunchOp::kNumConfigOperands + op.getNumOperands()) |
| return op.emitOpError("unexpected number of region arguments"); |
| } |
| |
| // Block terminators without successors are expected to exit the kernel region |
| // and must be `gpu.launch`. |
| for (Block &block : op.body()) { |
| if (block.empty()) |
| continue; |
| if (block.back().getNumSuccessors() != 0) |
| continue; |
| if (!isa<gpu::ReturnOp>(&block.back())) { |
| return block.back() |
| .emitError("expected 'gpu.terminator' or a terminator with " |
| "successors") |
| .attachNote(op.getLoc()) |
| << "in '" << LaunchOp::getOperationName() << "' body region"; |
| } |
| } |
| |
| return success(); |
| } |
| |
| // Pretty-print the kernel grid/block size assignment as |
| // (%iter-x, %iter-y, %iter-z) in |
| // (%size-x = %ssa-use, %size-y = %ssa-use, %size-z = %ssa-use) |
| // where %size-* and %iter-* will correspond to the body region arguments. |
| static void printSizeAssignment(OpAsmPrinter &p, KernelDim3 size, |
| ValueRange operands, KernelDim3 ids) { |
| p << '(' << ids.x << ", " << ids.y << ", " << ids.z << ") in ("; |
| p << size.x << " = " << operands[0] << ", "; |
| p << size.y << " = " << operands[1] << ", "; |
| p << size.z << " = " << operands[2] << ')'; |
| } |
| |
| static void printLaunchOp(OpAsmPrinter &p, LaunchOp op) { |
| ValueRange operands = op.getOperands(); |
| |
| // Print the launch configuration. |
| p << LaunchOp::getOperationName() << ' ' << op.getBlocksKeyword(); |
| printSizeAssignment(p, op.getGridSize(), operands.take_front(3), |
| op.getBlockIds()); |
| p << ' ' << op.getThreadsKeyword(); |
| printSizeAssignment(p, op.getBlockSize(), operands.slice(3, 3), |
| op.getThreadIds()); |
| |
| // From now on, the first kNumConfigOperands operands corresponding to grid |
| // and block sizes are irrelevant, so we can drop them. |
| operands = operands.drop_front(LaunchOp::kNumConfigOperands); |
| |
| // Print the data argument remapping. |
| if (!op.body().empty() && !operands.empty()) { |
| p << ' ' << op.getArgsKeyword() << '('; |
| Block *entryBlock = &op.body().front(); |
| interleaveComma(llvm::seq<int>(0, operands.size()), p, [&](int i) { |
| p << entryBlock->getArgument(LaunchOp::kNumConfigRegionAttributes + i) |
| << " = " << operands[i]; |
| }); |
| p << ") "; |
| } |
| |
| // Print the types of data arguments. |
| if (!operands.empty()) |
| p << ": " << operands.getTypes(); |
| |
| p.printRegion(op.body(), /*printEntryBlockArgs=*/false); |
| p.printOptionalAttrDict(op.getAttrs()); |
| } |
| |
| // Parse the size assignment blocks for blocks and threads. These have the form |
| // (%region_arg, %region_arg, %region_arg) in |
| // (%region_arg = %operand, %region_arg = %operand, %region_arg = %operand) |
| // where %region_arg are percent-identifiers for the region arguments to be |
| // introduced further (SSA defs), and %operand are percent-identifiers for the |
| // SSA value uses. |
| static ParseResult |
| parseSizeAssignment(OpAsmParser &parser, |
| MutableArrayRef<OpAsmParser::OperandType> sizes, |
| MutableArrayRef<OpAsmParser::OperandType> regionSizes, |
| MutableArrayRef<OpAsmParser::OperandType> indices) { |
| assert(indices.size() == 3 && "space for three indices expected"); |
| SmallVector<OpAsmParser::OperandType, 3> args; |
| if (parser.parseRegionArgumentList(args, /*requiredOperandCount=*/3, |
| OpAsmParser::Delimiter::Paren) || |
| parser.parseKeyword("in") || parser.parseLParen()) |
| return failure(); |
| std::move(args.begin(), args.end(), indices.begin()); |
| |
| for (int i = 0; i < 3; ++i) { |
| if (i != 0 && parser.parseComma()) |
| return failure(); |
| if (parser.parseRegionArgument(regionSizes[i]) || parser.parseEqual() || |
| parser.parseOperand(sizes[i])) |
| return failure(); |
| } |
| |
| return parser.parseRParen(); |
| } |
| |
| // Parses a Launch operation. |
| // operation ::= `gpu.launch` `blocks` `(` ssa-id-list `)` `in` ssa-reassignment |
| // `threads` `(` ssa-id-list `)` `in` ssa-reassignment |
| // (`args` ssa-reassignment `:` type-list)? |
| // region attr-dict? |
| // ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)` |
| static ParseResult parseLaunchOp(OpAsmParser &parser, OperationState &result) { |
| // Sizes of the grid and block. |
| SmallVector<OpAsmParser::OperandType, LaunchOp::kNumConfigOperands> sizes( |
| LaunchOp::kNumConfigOperands); |
| MutableArrayRef<OpAsmParser::OperandType> sizesRef(sizes); |
| |
| // Actual (data) operands passed to the kernel. |
| SmallVector<OpAsmParser::OperandType, 4> dataOperands; |
| |
| // Region arguments to be created. |
| SmallVector<OpAsmParser::OperandType, 16> regionArgs( |
| LaunchOp::kNumConfigRegionAttributes); |
| MutableArrayRef<OpAsmParser::OperandType> regionArgsRef(regionArgs); |
| |
| // Parse the size assignment segments: the first segment assigns grid sizes |
| // and defines values for block identifiers; the second segment assigns block |
| // sizes and defines values for thread identifiers. In the region argument |
| // list, identifiers precede sizes, and block-related values precede |
| // thread-related values. |
| if (parser.parseKeyword(LaunchOp::getBlocksKeyword().data()) || |
| parseSizeAssignment(parser, sizesRef.take_front(3), |
| regionArgsRef.slice(6, 3), |
| regionArgsRef.slice(0, 3)) || |
| parser.parseKeyword(LaunchOp::getThreadsKeyword().data()) || |
| parseSizeAssignment(parser, sizesRef.drop_front(3), |
| regionArgsRef.slice(9, 3), |
| regionArgsRef.slice(3, 3)) || |
| parser.resolveOperands(sizes, parser.getBuilder().getIndexType(), |
| result.operands)) |
| return failure(); |
| |
| // If kernel argument renaming segment is present, parse it. When present, |
| // the segment should have at least one element. If this segment is present, |
| // so is the trailing type list. Parse it as well and use the parsed types |
| // to resolve the operands passed to the kernel arguments. |
| SmallVector<Type, 4> dataTypes; |
| if (!parser.parseOptionalKeyword(LaunchOp::getArgsKeyword())) { |
| llvm::SMLoc argsLoc = parser.getCurrentLocation(); |
| |
| regionArgs.push_back({}); |
| dataOperands.push_back({}); |
| if (parser.parseLParen() || parser.parseRegionArgument(regionArgs.back()) || |
| parser.parseEqual() || parser.parseOperand(dataOperands.back())) |
| return failure(); |
| |
| while (!parser.parseOptionalComma()) { |
| regionArgs.push_back({}); |
| dataOperands.push_back({}); |
| if (parser.parseRegionArgument(regionArgs.back()) || |
| parser.parseEqual() || parser.parseOperand(dataOperands.back())) |
| return failure(); |
| } |
| |
| if (parser.parseRParen() || parser.parseColonTypeList(dataTypes) || |
| parser.resolveOperands(dataOperands, dataTypes, argsLoc, |
| result.operands)) |
| return failure(); |
| } |
| |
| // Introduce the body region and parse it. The region has |
| // kNumConfigRegionAttributes leading arguments that correspond to |
| // block/thread identifiers and grid/block sizes, all of the `index` type. |
| // Follow the actual kernel arguments. |
| Type index = parser.getBuilder().getIndexType(); |
| dataTypes.insert(dataTypes.begin(), LaunchOp::kNumConfigRegionAttributes, |
| index); |
| Region *body = result.addRegion(); |
| return failure(parser.parseRegion(*body, regionArgs, dataTypes) || |
| parser.parseOptionalAttrDict(result.attributes)); |
| } |
| |
| void LaunchOp::eraseKernelArgument(unsigned index) { |
| Block &entryBlock = body().front(); |
| assert(index < entryBlock.getNumArguments() - kNumConfigRegionAttributes && |
| "kernel argument index overflow"); |
| entryBlock.eraseArgument(kNumConfigRegionAttributes + index); |
| getOperation()->eraseOperand(kNumConfigOperands + index); |
| } |
| |
| namespace { |
| // Clone any known constants passed as operands to the kernel into its body. |
| class PropagateConstantBounds : public OpRewritePattern<LaunchOp> { |
| using OpRewritePattern<LaunchOp>::OpRewritePattern; |
| |
| PatternMatchResult matchAndRewrite(LaunchOp launchOp, |
| PatternRewriter &rewriter) const override { |
| rewriter.startRootUpdate(launchOp); |
| PatternRewriter::InsertionGuard guard(rewriter); |
| rewriter.setInsertionPointToStart(&launchOp.body().front()); |
| |
| // Traverse operands passed to kernel and check if some of them are known |
| // constants. If so, clone the constant operation inside the kernel region |
| // and use it instead of passing the value from the parent region. Perform |
| // the traversal in the inverse order to simplify index arithmetics when |
| // dropping arguments. |
| auto operands = launchOp.getKernelOperandValues(); |
| auto kernelArgs = launchOp.getKernelArguments(); |
| bool found = false; |
| for (unsigned i = operands.size(); i > 0; --i) { |
| unsigned index = i - 1; |
| Value operand = operands[index]; |
| if (!isa_and_nonnull<ConstantOp>(operand.getDefiningOp())) |
| continue; |
| |
| found = true; |
| Value internalConstant = |
| rewriter.clone(*operand.getDefiningOp())->getResult(0); |
| Value kernelArg = *std::next(kernelArgs.begin(), index); |
| kernelArg.replaceAllUsesWith(internalConstant); |
| launchOp.eraseKernelArgument(index); |
| } |
| |
| if (!found) { |
| rewriter.cancelRootUpdate(launchOp); |
| return matchFailure(); |
| } |
| |
| rewriter.finalizeRootUpdate(launchOp); |
| return matchSuccess(); |
| } |
| }; |
| } // end namespace |
| |
| void LaunchOp::getCanonicalizationPatterns(OwningRewritePatternList &results, |
| MLIRContext *context) { |
| results.insert<PropagateConstantBounds>(context); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // LaunchFuncOp |
| //===----------------------------------------------------------------------===// |
| |
| void LaunchFuncOp::build(Builder *builder, OperationState &result, |
| GPUFuncOp kernelFunc, Value gridSizeX, Value gridSizeY, |
| Value gridSizeZ, Value blockSizeX, Value blockSizeY, |
| Value blockSizeZ, ValueRange kernelOperands) { |
| // Add grid and block sizes as op operands, followed by the data operands. |
| result.addOperands( |
| {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ}); |
| result.addOperands(kernelOperands); |
| result.addAttribute(getKernelAttrName(), |
| builder->getStringAttr(kernelFunc.getName())); |
| auto kernelModule = kernelFunc.getParentOfType<GPUModuleOp>(); |
| result.addAttribute(getKernelModuleAttrName(), |
| builder->getSymbolRefAttr(kernelModule.getName())); |
| } |
| |
| void LaunchFuncOp::build(Builder *builder, OperationState &result, |
| GPUFuncOp kernelFunc, KernelDim3 gridSize, |
| KernelDim3 blockSize, ValueRange kernelOperands) { |
| build(builder, result, kernelFunc, gridSize.x, gridSize.y, gridSize.z, |
| blockSize.x, blockSize.y, blockSize.z, kernelOperands); |
| } |
| |
| StringRef LaunchFuncOp::kernel() { |
| return getAttrOfType<StringAttr>(getKernelAttrName()).getValue(); |
| } |
| |
| unsigned LaunchFuncOp::getNumKernelOperands() { |
| return getNumOperands() - kNumConfigOperands; |
| } |
| |
| StringRef LaunchFuncOp::getKernelModuleName() { |
| return getAttrOfType<SymbolRefAttr>(getKernelModuleAttrName()) |
| .getRootReference(); |
| } |
| |
| Value LaunchFuncOp::getKernelOperand(unsigned i) { |
| return getOperation()->getOperand(i + kNumConfigOperands); |
| } |
| |
| KernelDim3 LaunchFuncOp::getGridSizeOperandValues() { |
| return KernelDim3{getOperand(0), getOperand(1), getOperand(2)}; |
| } |
| |
| KernelDim3 LaunchFuncOp::getBlockSizeOperandValues() { |
| return KernelDim3{getOperand(3), getOperand(4), getOperand(5)}; |
| } |
| |
| static LogicalResult verify(LaunchFuncOp op) { |
| auto module = op.getParentOfType<ModuleOp>(); |
| if (!module) |
| return op.emitOpError("expected to belong to a module"); |
| |
| if (!module.getAttrOfType<UnitAttr>(GPUDialect::getContainerModuleAttrName())) |
| return op.emitOpError( |
| "expected the closest surrounding module to have the '" + |
| GPUDialect::getContainerModuleAttrName() + "' attribute"); |
| |
| auto kernelAttr = op.getAttrOfType<StringAttr>(op.getKernelAttrName()); |
| if (!kernelAttr) |
| return op.emitOpError("string attribute '" + op.getKernelAttrName() + |
| "' must be specified"); |
| |
| auto kernelModuleAttr = |
| op.getAttrOfType<SymbolRefAttr>(op.getKernelModuleAttrName()); |
| if (!kernelModuleAttr) |
| return op.emitOpError("symbol reference attribute '" + |
| op.getKernelModuleAttrName() + "' must be specified"); |
| |
| return success(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // GPUFuncOp |
| //===----------------------------------------------------------------------===// |
| |
| /// Adds a workgroup attribution to "op" of the MemRef type with the given shape |
| /// and element type. |
| Value GPUFuncOp::addWorkgroupAttribution(ArrayRef<int64_t> shape, |
| Type elementType) { |
| unsigned pos = getNumFuncArguments() + getNumWorkgroupAttributions(); |
| Block &bodyBlock = body().front(); |
| Value attribution = bodyBlock.insertArgument( |
| std::next(bodyBlock.args_begin(), pos), |
| MemRefType::get(shape, elementType, /*affineMapComposition=*/{}, |
| GPUDialect::getWorkgroupAddressSpace())); |
| auto numWorkgroupBuffersAttr = |
| getAttrOfType<IntegerAttr>(getNumWorkgroupAttributionsAttrName()); |
| setAttr(getNumWorkgroupAttributionsAttrName(), |
| IntegerAttr::get(numWorkgroupBuffersAttr.getType(), |
| numWorkgroupBuffersAttr.getValue() + 1)); |
| return attribution; |
| } |
| |
| void GPUFuncOp::build(Builder *builder, OperationState &result, StringRef name, |
| FunctionType type, ArrayRef<Type> workgroupAttributions, |
| ArrayRef<Type> privateAttributions, |
| ArrayRef<NamedAttribute> attrs) { |
| result.addAttribute(SymbolTable::getSymbolAttrName(), |
| builder->getStringAttr(name)); |
| result.addAttribute(getTypeAttrName(), TypeAttr::get(type)); |
| result.addAttribute(getNumWorkgroupAttributionsAttrName(), |
| builder->getI64IntegerAttr(workgroupAttributions.size())); |
| result.addAttributes(attrs); |
| Region *body = result.addRegion(); |
| Block *entryBlock = new Block; |
| entryBlock->addArguments(type.getInputs()); |
| entryBlock->addArguments(workgroupAttributions); |
| entryBlock->addArguments(privateAttributions); |
| |
| body->getBlocks().push_back(entryBlock); |
| } |
| |
| /// Parses a GPU function memory attribution. |
| /// |
| /// memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)? |
| /// (`private` `(` ssa-id-and-type-list `)`)? |
| /// |
| /// Note that this function parses only one of the two similar parts, with the |
| /// keyword provided as argument. |
| static ParseResult |
| parseAttributions(OpAsmParser &parser, StringRef keyword, |
| SmallVectorImpl<OpAsmParser::OperandType> &args, |
| SmallVectorImpl<Type> &argTypes) { |
| // If we could not parse the keyword, just assume empty list and succeed. |
| if (failed(parser.parseOptionalKeyword(keyword))) |
| return success(); |
| |
| if (failed(parser.parseLParen())) |
| return failure(); |
| |
| // Early exit for an empty list. |
| if (succeeded(parser.parseOptionalRParen())) |
| return success(); |
| |
| do { |
| OpAsmParser::OperandType arg; |
| Type type; |
| |
| if (parser.parseRegionArgument(arg) || parser.parseColonType(type)) |
| return failure(); |
| |
| args.push_back(arg); |
| argTypes.push_back(type); |
| } while (succeeded(parser.parseOptionalComma())); |
| |
| return parser.parseRParen(); |
| } |
| |
| /// Parses a GPU function. |
| /// |
| /// <operation> ::= `gpu.func` symbol-ref-id `(` argument-list `)` |
| /// (`->` function-result-list)? memory-attribution `kernel`? |
| /// function-attributes? region |
| static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result) { |
| SmallVector<OpAsmParser::OperandType, 8> entryArgs; |
| SmallVector<SmallVector<NamedAttribute, 2>, 1> argAttrs; |
| SmallVector<SmallVector<NamedAttribute, 2>, 1> resultAttrs; |
| SmallVector<Type, 8> argTypes; |
| SmallVector<Type, 4> resultTypes; |
| bool isVariadic; |
| |
| // Parse the function name. |
| StringAttr nameAttr; |
| if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(), |
| result.attributes)) |
| return failure(); |
| |
| auto signatureLocation = parser.getCurrentLocation(); |
| if (failed(impl::parseFunctionSignature( |
| parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs, |
| isVariadic, resultTypes, resultAttrs))) |
| return failure(); |
| |
| if (entryArgs.empty() && !argTypes.empty()) |
| return parser.emitError(signatureLocation) |
| << "gpu.func requires named arguments"; |
| |
| // Construct the function type. More types will be added to the region, but |
| // not to the functiont type. |
| Builder &builder = parser.getBuilder(); |
| auto type = builder.getFunctionType(argTypes, resultTypes); |
| result.addAttribute(GPUFuncOp::getTypeAttrName(), TypeAttr::get(type)); |
| |
| // Parse workgroup memory attributions. |
| if (failed(parseAttributions(parser, GPUFuncOp::getWorkgroupKeyword(), |
| entryArgs, argTypes))) |
| return failure(); |
| |
| // Store the number of operands we just parsed as the number of workgroup |
| // memory attributions. |
| unsigned numWorkgroupAttrs = argTypes.size() - type.getNumInputs(); |
| result.addAttribute(GPUFuncOp::getNumWorkgroupAttributionsAttrName(), |
| builder.getI64IntegerAttr(numWorkgroupAttrs)); |
| |
| // Parse private memory attributions. |
| if (failed(parseAttributions(parser, GPUFuncOp::getPrivateKeyword(), |
| entryArgs, argTypes))) |
| return failure(); |
| |
| // Parse the kernel attribute if present. |
| if (succeeded(parser.parseOptionalKeyword(GPUFuncOp::getKernelKeyword()))) |
| result.addAttribute(GPUDialect::getKernelFuncAttrName(), |
| builder.getUnitAttr()); |
| |
| // Parse attributes. |
| if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes))) |
| return failure(); |
| mlir::impl::addArgAndResultAttrs(builder, result, argAttrs, resultAttrs); |
| |
| // Parse the region. If no argument names were provided, take all names |
| // (including those of attributions) from the entry block. |
| auto *body = result.addRegion(); |
| return parser.parseRegion(*body, entryArgs, argTypes); |
| } |
| |
| static void printAttributions(OpAsmPrinter &p, StringRef keyword, |
| ArrayRef<BlockArgument> values) { |
| if (values.empty()) |
| return; |
| |
| p << ' ' << keyword << '('; |
| interleaveComma(values, p, |
| [&p](BlockArgument v) { p << v << " : " << v.getType(); }); |
| p << ')'; |
| } |
| |
| /// Prints a GPU Func op. |
| static void printGPUFuncOp(OpAsmPrinter &p, GPUFuncOp op) { |
| p << GPUFuncOp::getOperationName() << ' '; |
| p.printSymbolName(op.getName()); |
| |
| FunctionType type = op.getType(); |
| impl::printFunctionSignature(p, op.getOperation(), type.getInputs(), |
| /*isVariadic=*/false, type.getResults()); |
| |
| printAttributions(p, op.getWorkgroupKeyword(), op.getWorkgroupAttributions()); |
| printAttributions(p, op.getPrivateKeyword(), op.getPrivateAttributions()); |
| if (op.isKernel()) |
| p << ' ' << op.getKernelKeyword(); |
| |
| impl::printFunctionAttributes(p, op.getOperation(), type.getNumInputs(), |
| type.getNumResults(), |
| {op.getNumWorkgroupAttributionsAttrName(), |
| GPUDialect::getKernelFuncAttrName()}); |
| p.printRegion(op.getBody(), /*printEntryBlockArgs=*/false); |
| } |
| |
| void GPUFuncOp::setType(FunctionType newType) { |
| auto oldType = getType(); |
| assert(newType.getNumResults() == oldType.getNumResults() && |
| "unimplemented: changes to the number of results"); |
| |
| SmallVector<char, 16> nameBuf; |
| for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++) |
| removeAttr(getArgAttrName(i, nameBuf)); |
| |
| setAttr(getTypeAttrName(), TypeAttr::get(newType)); |
| } |
| |
| /// Hook for FunctionLike verifier. |
| LogicalResult GPUFuncOp::verifyType() { |
| Type type = getTypeAttr().getValue(); |
| if (!type.isa<FunctionType>()) |
| return emitOpError("requires '" + getTypeAttrName() + |
| "' attribute of function type"); |
| return success(); |
| } |
| |
| static LogicalResult verifyAttributions(Operation *op, |
| ArrayRef<BlockArgument> attributions, |
| unsigned memorySpace) { |
| for (Value v : attributions) { |
| auto type = v.getType().dyn_cast<MemRefType>(); |
| if (!type) |
| return op->emitOpError() << "expected memref type in attribution"; |
| |
| if (type.getMemorySpace() != memorySpace) { |
| return op->emitOpError() |
| << "expected memory space " << memorySpace << " in attribution"; |
| } |
| } |
| return success(); |
| } |
| |
| /// Verifies the body of the function. |
| LogicalResult GPUFuncOp::verifyBody() { |
| unsigned numFuncArguments = getNumArguments(); |
| unsigned numWorkgroupAttributions = getNumWorkgroupAttributions(); |
| unsigned numBlockArguments = front().getNumArguments(); |
| if (numBlockArguments < numFuncArguments + numWorkgroupAttributions) |
| return emitOpError() << "expected at least " |
| << numFuncArguments + numWorkgroupAttributions |
| << " arguments to body region"; |
| |
| ArrayRef<Type> funcArgTypes = getType().getInputs(); |
| for (unsigned i = 0; i < numFuncArguments; ++i) { |
| Type blockArgType = front().getArgument(i).getType(); |
| if (funcArgTypes[i] != blockArgType) |
| return emitOpError() << "expected body region argument #" << i |
| << " to be of type " << funcArgTypes[i] << ", got " |
| << blockArgType; |
| } |
| |
| if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(), |
| GPUDialect::getWorkgroupAddressSpace())) || |
| failed(verifyAttributions(getOperation(), getPrivateAttributions(), |
| GPUDialect::getPrivateAddressSpace()))) |
| return failure(); |
| |
| return success(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // GPUModuleOp |
| //===----------------------------------------------------------------------===// |
| |
| void GPUModuleOp::build(Builder *builder, OperationState &result, |
| StringRef name) { |
| ensureTerminator(*result.addRegion(), *builder, result.location); |
| result.attributes.push_back(builder->getNamedAttr( |
| ::mlir::SymbolTable::getSymbolAttrName(), builder->getStringAttr(name))); |
| } |
| |
| static ParseResult parseGPUModuleOp(OpAsmParser &parser, |
| OperationState &result) { |
| StringAttr nameAttr; |
| if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(), |
| result.attributes)) |
| return failure(); |
| |
| // If module attributes are present, parse them. |
| if (parser.parseOptionalAttrDictWithKeyword(result.attributes)) |
| return failure(); |
| |
| // Parse the module body. |
| auto *body = result.addRegion(); |
| if (parser.parseRegion(*body, None, None)) |
| return failure(); |
| |
| // Ensure that this module has a valid terminator. |
| GPUModuleOp::ensureTerminator(*body, parser.getBuilder(), result.location); |
| return success(); |
| } |
| |
| static void print(OpAsmPrinter &p, GPUModuleOp op) { |
| p << op.getOperationName() << ' '; |
| p.printSymbolName(op.getName()); |
| p.printOptionalAttrDictWithKeyword(op.getAttrs(), |
| {SymbolTable::getSymbolAttrName()}); |
| p.printRegion(op.getOperation()->getRegion(0), /*printEntryBlockArgs=*/false, |
| /*printBlockTerminators=*/false); |
| } |
| |
| // Namespace avoids ambiguous ReturnOpOperandAdaptor. |
| namespace mlir { |
| namespace gpu { |
| #define GET_OP_CLASSES |
| #include "mlir/Dialect/GPU/GPUOps.cpp.inc" |
| } // namespace gpu |
| } // namespace mlir |