bb1d976feb
As we start defining more complex Ops, we increasingly see the need for Ops-with-regions to be able to construct Ops within their regions in their ::build methods. However, these methods only have access to Builder, and not OpBuilder. Creating a local instance of OpBuilder inside ::build and using it fails to trigger the operation creation hooks in derived builders (e.g., ConversionPatternRewriter). In this case, we risk breaking the logic of the derived builder. At the same time, OpBuilder::create, which is by far the largest user of ::build already passes "this" as the first argument, so an OpBuilder instance is already available. Update all ::build methods in all Ops in MLIR and Flang to take "OpBuilder &" instead of "Builder *". Note the change from pointer and to reference to comply with the common style in MLIR, this also ensures all other users must change their ::build methods. Differential Revision: https://reviews.llvm.org/D78713
789 lines
31 KiB
C++
789 lines
31 KiB
C++
//===- GPUDialect.cpp - MLIR Dialect for GPU Kernels implementation -------===//
|
|
//
|
|
// Part of the LLVM 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/IR/Ops.h"
|
|
#include "mlir/IR/Attributes.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
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
bool GPUDialect::isKernel(Operation *op) {
|
|
UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
|
|
return static_cast<bool>(isKernelAttr);
|
|
}
|
|
|
|
GPUDialect::GPUDialect(MLIRContext *context)
|
|
: Dialect(getDialectNamespace(), 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 != 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<SymbolRefAttr>(
|
|
LaunchFuncOp::getKernelAttrName()))
|
|
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.
|
|
Operation *kernelFunc = module.lookupSymbol(launchOp.kernel());
|
|
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 '")
|
|
<< launchOp.kernel() << "' is undefined";
|
|
if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
|
|
GPUDialect::getKernelFuncAttrName()))
|
|
return launchOp.emitOpError("kernel function is missing the '")
|
|
<< GPUDialect::getKernelFuncAttrName() << "' attribute";
|
|
|
|
// TODO(ntv,zinenko,herhut): if the kernel function has been converted to
|
|
// the LLVM dialect but the caller hasn't (which happens during the
|
|
// separate compilation), do not check type correspondence as it would
|
|
// require the verifier to be aware of the LLVM type conversion.
|
|
if (kernelLLVMFunction)
|
|
return success();
|
|
|
|
unsigned actualNumArguments = launchOp.getNumKernelOperands();
|
|
unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
|
|
if (expectedNumArguments != actualNumArguments)
|
|
return launchOp.emitOpError("got ")
|
|
<< actualNumArguments << " kernel operands but expected "
|
|
<< expectedNumArguments;
|
|
|
|
auto functionType = kernelGPUFunction.getType();
|
|
for (unsigned i = 0; i < expectedNumArguments; ++i) {
|
|
if (launchOp.getKernelOperand(i).getType() != functionType.getInput(i)) {
|
|
return launchOp.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");
|
|
} else {
|
|
StringRef opName = *allReduce.op();
|
|
if ((opName == "and" || opName == "or" || opName == "xor") &&
|
|
!allReduce.getType().isa<IntegerType>()) {
|
|
return allReduce.emitError()
|
|
<< '`' << opName << '`'
|
|
<< " accumulator is only compatible with Integer type";
|
|
}
|
|
}
|
|
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.isSignlessIntOrFloat() || 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(OpBuilder &builder, OperationState &result,
|
|
Value gridSizeX, Value gridSizeY, Value gridSizeZ,
|
|
Value blockSizeX, Value blockSizeY, Value blockSizeZ) {
|
|
// Add grid and block sizes as op operands, followed by the data operands.
|
|
result.addOperands(
|
|
{gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
|
|
|
|
// 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()));
|
|
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]};
|
|
}
|
|
|
|
KernelDim3 LaunchOp::getGridSizeOperandValues() {
|
|
return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
|
|
}
|
|
|
|
KernelDim3 LaunchOp::getBlockSizeOperandValues() {
|
|
return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
|
|
}
|
|
|
|
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.terminator`.
|
|
for (Block &block : op.body()) {
|
|
if (block.empty())
|
|
continue;
|
|
if (block.back().getNumSuccessors() != 0)
|
|
continue;
|
|
if (!isa<gpu::TerminatorOp>(&block.back())) {
|
|
return block.back()
|
|
.emitError()
|
|
.append("expected '", gpu::TerminatorOp::getOperationName(),
|
|
"' or a terminator with successors")
|
|
.attachNote(op.getLoc())
|
|
.append("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());
|
|
|
|
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
|
|
// 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();
|
|
|
|
// Introduce the body region and parse it. The region has
|
|
// kNumConfigRegionAttributes arguments that correspond to
|
|
// block/thread identifiers and grid/block sizes, all of the `index` type.
|
|
Type index = parser.getBuilder().getIndexType();
|
|
SmallVector<Type, LaunchOp::kNumConfigRegionAttributes> dataTypes(
|
|
LaunchOp::kNumConfigRegionAttributes, index);
|
|
Region *body = result.addRegion();
|
|
return failure(parser.parseRegion(*body, regionArgs, dataTypes) ||
|
|
parser.parseOptionalAttrDict(result.attributes));
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// LaunchFuncOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void LaunchFuncOp::build(OpBuilder &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);
|
|
auto kernelModule = kernelFunc.getParentOfType<GPUModuleOp>();
|
|
auto kernelSymbol = builder.getSymbolRefAttr(
|
|
kernelModule.getName(), {builder.getSymbolRefAttr(kernelFunc.getName())});
|
|
result.addAttribute(getKernelAttrName(), kernelSymbol);
|
|
}
|
|
|
|
void LaunchFuncOp::build(OpBuilder &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);
|
|
}
|
|
|
|
SymbolRefAttr LaunchFuncOp::kernel() {
|
|
return getAttrOfType<SymbolRefAttr>(getKernelAttrName());
|
|
}
|
|
|
|
unsigned LaunchFuncOp::getNumKernelOperands() {
|
|
return getNumOperands() - kNumConfigOperands;
|
|
}
|
|
|
|
StringRef LaunchFuncOp::getKernelModuleName() {
|
|
return kernel().getRootReference();
|
|
}
|
|
|
|
StringRef LaunchFuncOp::getKernelName() { return kernel().getLeafReference(); }
|
|
|
|
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<SymbolRefAttr>(op.getKernelAttrName());
|
|
if (!kernelAttr)
|
|
return op.emitOpError("symbol reference attribute '" +
|
|
op.getKernelAttrName() + "' 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(OpBuilder &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 function 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 << '(';
|
|
llvm::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");
|
|
|
|
if (isKernel() && getType().getNumResults() != 0)
|
|
return emitOpError() << "expected void return type for kernel function";
|
|
|
|
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();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// ReturnOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
static ParseResult parseReturnOp(OpAsmParser &parser, OperationState &result) {
|
|
llvm::SmallVector<OpAsmParser::OperandType, 4> operands;
|
|
llvm::SmallVector<Type, 4> types;
|
|
if (parser.parseOperandList(operands) ||
|
|
parser.parseOptionalColonTypeList(types) ||
|
|
parser.resolveOperands(operands, types, parser.getCurrentLocation(),
|
|
result.operands))
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verify(gpu::ReturnOp returnOp) {
|
|
GPUFuncOp function = returnOp.getParentOfType<GPUFuncOp>();
|
|
|
|
FunctionType funType = function.getType();
|
|
|
|
if (funType.getNumResults() != returnOp.operands().size())
|
|
return returnOp.emitOpError()
|
|
.append("expected ", funType.getNumResults(), " result operands")
|
|
.attachNote(function.getLoc())
|
|
.append("return type declared here");
|
|
|
|
for (auto pair : llvm::enumerate(
|
|
llvm::zip(function.getType().getResults(), returnOp.operands()))) {
|
|
Type type;
|
|
Value operand;
|
|
std::tie(type, operand) = pair.value();
|
|
if (type != operand.getType())
|
|
return returnOp.emitOpError() << "unexpected type `" << operand.getType()
|
|
<< "' for operand #" << pair.index();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// GPUModuleOp
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void GPUModuleOp::build(OpBuilder &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
|