453 lines · cpp
1//===- Async.cpp - MLIR Async Operations ----------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Async/IR/Async.h"10 11#include "mlir/IR/DialectImplementation.h"12#include "mlir/Interfaces/FunctionImplementation.h"13#include "llvm/ADT/TypeSwitch.h"14 15using namespace mlir;16using namespace mlir::async;17 18#include "mlir/Dialect/Async/IR/AsyncOpsDialect.cpp.inc"19 20void AsyncDialect::initialize() {21 addOperations<22#define GET_OP_LIST23#include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"24 >();25 addTypes<26#define GET_TYPEDEF_LIST27#include "mlir/Dialect/Async/IR/AsyncOpsTypes.cpp.inc"28 >();29}30 31//===----------------------------------------------------------------------===//32/// ExecuteOp33//===----------------------------------------------------------------------===//34 35constexpr char kOperandSegmentSizesAttr[] = "operandSegmentSizes";36 37OperandRange ExecuteOp::getEntrySuccessorOperands(RegionSuccessor successor) {38 assert(successor.getSuccessor() == &getBodyRegion() &&39 "invalid region index");40 return getBodyOperands();41}42 43bool ExecuteOp::areTypesCompatible(Type lhs, Type rhs) {44 const auto getValueOrTokenType = [](Type type) {45 if (auto value = llvm::dyn_cast<ValueType>(type))46 return value.getValueType();47 return type;48 };49 return getValueOrTokenType(lhs) == getValueOrTokenType(rhs);50}51 52void ExecuteOp::getSuccessorRegions(RegionBranchPoint point,53 SmallVectorImpl<RegionSuccessor> ®ions) {54 // The `body` region branch back to the parent operation.55 if (!point.isParent() &&56 point.getTerminatorPredecessorOrNull()->getParentRegion() ==57 &getBodyRegion()) {58 regions.push_back(RegionSuccessor(getOperation(), getBodyResults()));59 return;60 }61 62 // Otherwise the successor is the body region.63 regions.push_back(64 RegionSuccessor(&getBodyRegion(), getBodyRegion().getArguments()));65}66 67void ExecuteOp::build(OpBuilder &builder, OperationState &result,68 TypeRange resultTypes, ValueRange dependencies,69 ValueRange operands, BodyBuilderFn bodyBuilder) {70 OpBuilder::InsertionGuard guard(builder);71 result.addOperands(dependencies);72 result.addOperands(operands);73 74 // Add derived `operandSegmentSizes` attribute based on parsed operands.75 int32_t numDependencies = dependencies.size();76 int32_t numOperands = operands.size();77 auto operandSegmentSizes =78 builder.getDenseI32ArrayAttr({numDependencies, numOperands});79 result.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);80 81 // First result is always a token, and then `resultTypes` wrapped into82 // `async.value`.83 result.addTypes({TokenType::get(result.getContext())});84 for (Type type : resultTypes)85 result.addTypes(ValueType::get(type));86 87 // Add a body region with block arguments as unwrapped async value operands.88 Region *bodyRegion = result.addRegion();89 Block *bodyBlock = builder.createBlock(bodyRegion);90 for (Value operand : operands) {91 auto valueType = llvm::dyn_cast<ValueType>(operand.getType());92 bodyBlock->addArgument(valueType ? valueType.getValueType()93 : operand.getType(),94 operand.getLoc());95 }96 97 // Create the default terminator if the builder is not provided and if the98 // expected result is empty. Otherwise, leave this to the caller99 // because we don't know which values to return from the execute op.100 if (resultTypes.empty() && !bodyBuilder) {101 async::YieldOp::create(builder, result.location, ValueRange());102 } else if (bodyBuilder) {103 bodyBuilder(builder, result.location, bodyBlock->getArguments());104 }105}106 107void ExecuteOp::print(OpAsmPrinter &p) {108 // [%tokens,...]109 if (!getDependencies().empty())110 p << " [" << getDependencies() << "]";111 112 // (%value as %unwrapped: !async.value<!arg.type>, ...)113 if (!getBodyOperands().empty()) {114 p << " (";115 Block *entry = getBodyRegion().empty() ? nullptr : &getBodyRegion().front();116 llvm::interleaveComma(117 getBodyOperands(), p, [&, n = 0](Value operand) mutable {118 Value argument = entry ? entry->getArgument(n++) : Value();119 p << operand << " as " << argument << ": " << operand.getType();120 });121 p << ")";122 }123 124 // -> (!async.value<!return.type>, ...)125 p.printOptionalArrowTypeList(llvm::drop_begin(getResultTypes()));126 p.printOptionalAttrDictWithKeyword((*this)->getAttrs(),127 {kOperandSegmentSizesAttr});128 p << ' ';129 p.printRegion(getBodyRegion(), /*printEntryBlockArgs=*/false);130}131 132ParseResult ExecuteOp::parse(OpAsmParser &parser, OperationState &result) {133 MLIRContext *ctx = result.getContext();134 135 // Sizes of parsed variadic operands, will be updated below after parsing.136 int32_t numDependencies = 0;137 138 auto tokenTy = TokenType::get(ctx);139 140 // Parse dependency tokens.141 if (succeeded(parser.parseOptionalLSquare())) {142 SmallVector<OpAsmParser::UnresolvedOperand, 4> tokenArgs;143 if (parser.parseOperandList(tokenArgs) ||144 parser.resolveOperands(tokenArgs, tokenTy, result.operands) ||145 parser.parseRSquare())146 return failure();147 148 numDependencies = tokenArgs.size();149 }150 151 // Parse async value operands (%value as %unwrapped : !async.value<!type>).152 SmallVector<OpAsmParser::UnresolvedOperand, 4> valueArgs;153 SmallVector<OpAsmParser::Argument, 4> unwrappedArgs;154 SmallVector<Type, 4> valueTypes;155 156 // Parse a single instance of `%value as %unwrapped : !async.value<!type>`.157 auto parseAsyncValueArg = [&]() -> ParseResult {158 if (parser.parseOperand(valueArgs.emplace_back()) ||159 parser.parseKeyword("as") ||160 parser.parseArgument(unwrappedArgs.emplace_back()) ||161 parser.parseColonType(valueTypes.emplace_back()))162 return failure();163 164 auto valueTy = llvm::dyn_cast<ValueType>(valueTypes.back());165 unwrappedArgs.back().type = valueTy ? valueTy.getValueType() : Type();166 return success();167 };168 169 auto argsLoc = parser.getCurrentLocation();170 if (parser.parseCommaSeparatedList(OpAsmParser::Delimiter::OptionalParen,171 parseAsyncValueArg) ||172 parser.resolveOperands(valueArgs, valueTypes, argsLoc, result.operands))173 return failure();174 175 int32_t numOperands = valueArgs.size();176 177 // Add derived `operandSegmentSizes` attribute based on parsed operands.178 auto operandSegmentSizes =179 parser.getBuilder().getDenseI32ArrayAttr({numDependencies, numOperands});180 result.addAttribute(kOperandSegmentSizesAttr, operandSegmentSizes);181 182 // Parse the types of results returned from the async execute op.183 SmallVector<Type, 4> resultTypes;184 NamedAttrList attrs;185 if (parser.parseOptionalArrowTypeList(resultTypes) ||186 // Async execute first result is always a completion token.187 parser.addTypeToList(tokenTy, result.types) ||188 parser.addTypesToList(resultTypes, result.types) ||189 // Parse operation attributes.190 parser.parseOptionalAttrDictWithKeyword(attrs))191 return failure();192 193 result.addAttributes(attrs);194 195 // Parse asynchronous region.196 Region *body = result.addRegion();197 return parser.parseRegion(*body, /*arguments=*/unwrappedArgs);198}199 200LogicalResult ExecuteOp::verifyRegions() {201 // Unwrap async.execute value operands types.202 auto unwrappedTypes = llvm::map_range(getBodyOperands(), [](Value operand) {203 return llvm::cast<ValueType>(operand.getType()).getValueType();204 });205 206 // Verify that unwrapped argument types matches the body region arguments.207 if (getBodyRegion().getArgumentTypes() != unwrappedTypes)208 return emitOpError("async body region argument types do not match the "209 "execute operation arguments types");210 211 return success();212}213 214//===----------------------------------------------------------------------===//215/// CreateGroupOp216//===----------------------------------------------------------------------===//217 218LogicalResult CreateGroupOp::canonicalize(CreateGroupOp op,219 PatternRewriter &rewriter) {220 // Find all `await_all` users of the group.221 llvm::SmallVector<AwaitAllOp> awaitAllUsers;222 223 auto isAwaitAll = [&](Operation *op) -> bool {224 if (AwaitAllOp awaitAll = dyn_cast<AwaitAllOp>(op)) {225 awaitAllUsers.push_back(awaitAll);226 return true;227 }228 return false;229 };230 231 // Check if all users of the group are `await_all` operations.232 if (!llvm::all_of(op->getUsers(), isAwaitAll))233 return failure();234 235 // If group is only awaited without adding anything to it, we can safely erase236 // the create operation and all users.237 for (AwaitAllOp awaitAll : awaitAllUsers)238 rewriter.eraseOp(awaitAll);239 rewriter.eraseOp(op);240 241 return success();242}243 244//===----------------------------------------------------------------------===//245/// AwaitOp246//===----------------------------------------------------------------------===//247 248void AwaitOp::build(OpBuilder &builder, OperationState &result, Value operand,249 ArrayRef<NamedAttribute> attrs) {250 result.addOperands({operand});251 result.attributes.append(attrs.begin(), attrs.end());252 253 // Add unwrapped async.value type to the returned values types.254 if (auto valueType = llvm::dyn_cast<ValueType>(operand.getType()))255 result.addTypes(valueType.getValueType());256}257 258static ParseResult parseAwaitResultType(OpAsmParser &parser, Type &operandType,259 Type &resultType) {260 if (parser.parseType(operandType))261 return failure();262 263 // Add unwrapped async.value type to the returned values types.264 if (auto valueType = llvm::dyn_cast<ValueType>(operandType))265 resultType = valueType.getValueType();266 267 return success();268}269 270static void printAwaitResultType(OpAsmPrinter &p, Operation *op,271 Type operandType, Type resultType) {272 p << operandType;273}274 275LogicalResult AwaitOp::verify() {276 Type argType = getOperand().getType();277 278 // Awaiting on a token does not have any results.279 if (llvm::isa<TokenType>(argType) && !getResultTypes().empty())280 return emitOpError("awaiting on a token must have empty result");281 282 // Awaiting on a value unwraps the async value type.283 if (auto value = llvm::dyn_cast<ValueType>(argType)) {284 if (*getResultType() != value.getValueType())285 return emitOpError() << "result type " << *getResultType()286 << " does not match async value type "287 << value.getValueType();288 }289 290 return success();291}292 293//===----------------------------------------------------------------------===//294// FuncOp295//===----------------------------------------------------------------------===//296 297void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,298 FunctionType type, ArrayRef<NamedAttribute> attrs,299 ArrayRef<DictionaryAttr> argAttrs) {300 state.addAttribute(SymbolTable::getSymbolAttrName(),301 builder.getStringAttr(name));302 state.addAttribute(getFunctionTypeAttrName(state.name), TypeAttr::get(type));303 304 state.attributes.append(attrs.begin(), attrs.end());305 state.addRegion();306 307 if (argAttrs.empty())308 return;309 assert(type.getNumInputs() == argAttrs.size());310 call_interface_impl::addArgAndResultAttrs(311 builder, state, argAttrs, /*resultAttrs=*/{},312 getArgAttrsAttrName(state.name), getResAttrsAttrName(state.name));313}314 315ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) {316 auto buildFuncType =317 [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,318 function_interface_impl::VariadicFlag,319 std::string &) { return builder.getFunctionType(argTypes, results); };320 321 return function_interface_impl::parseFunctionOp(322 parser, result, /*allowVariadic=*/false,323 getFunctionTypeAttrName(result.name), buildFuncType,324 getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));325}326 327void FuncOp::print(OpAsmPrinter &p) {328 function_interface_impl::printFunctionOp(329 p, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),330 getArgAttrsAttrName(), getResAttrsAttrName());331}332 333/// Check that the result type of async.func is not void and must be334/// some async token or async values.335LogicalResult FuncOp::verify() {336 auto resultTypes = getResultTypes();337 if (resultTypes.empty())338 return emitOpError()339 << "result is expected to be at least of size 1, but got "340 << resultTypes.size();341 342 for (unsigned i = 0, e = resultTypes.size(); i != e; ++i) {343 auto type = resultTypes[i];344 if (!llvm::isa<TokenType>(type) && !llvm::isa<ValueType>(type))345 return emitOpError() << "result type must be async value type or async "346 "token type, but got "347 << type;348 // We only allow AsyncToken appear as the first return value349 if (llvm::isa<TokenType>(type) && i != 0) {350 return emitOpError()351 << " results' (optional) async token type is expected "352 "to appear as the 1st return value, but got "353 << i + 1;354 }355 }356 357 return success();358}359 360//===----------------------------------------------------------------------===//361/// CallOp362//===----------------------------------------------------------------------===//363 364LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {365 // Check that the callee attribute was specified.366 auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");367 if (!fnAttr)368 return emitOpError("requires a 'callee' symbol reference attribute");369 FuncOp fn = symbolTable.lookupNearestSymbolFrom<FuncOp>(*this, fnAttr);370 if (!fn)371 return emitOpError() << "'" << fnAttr.getValue()372 << "' does not reference a valid async function";373 374 // Verify that the operand and result types match the callee.375 auto fnType = fn.getFunctionType();376 if (fnType.getNumInputs() != getNumOperands())377 return emitOpError("incorrect number of operands for callee");378 379 for (unsigned i = 0, e = fnType.getNumInputs(); i != e; ++i)380 if (getOperand(i).getType() != fnType.getInput(i))381 return emitOpError("operand type mismatch: expected operand type ")382 << fnType.getInput(i) << ", but provided "383 << getOperand(i).getType() << " for operand number " << i;384 385 if (fnType.getNumResults() != getNumResults())386 return emitOpError("incorrect number of results for callee");387 388 for (unsigned i = 0, e = fnType.getNumResults(); i != e; ++i)389 if (getResult(i).getType() != fnType.getResult(i)) {390 auto diag = emitOpError("result type mismatch at index ") << i;391 diag.attachNote() << " op result types: " << getResultTypes();392 diag.attachNote() << "function result types: " << fnType.getResults();393 return diag;394 }395 396 return success();397}398 399FunctionType CallOp::getCalleeType() {400 return FunctionType::get(getContext(), getOperandTypes(), getResultTypes());401}402 403//===----------------------------------------------------------------------===//404/// ReturnOp405//===----------------------------------------------------------------------===//406 407LogicalResult ReturnOp::verify() {408 auto funcOp = (*this)->getParentOfType<FuncOp>();409 ArrayRef<Type> resultTypes = funcOp.isStateful()410 ? funcOp.getResultTypes().drop_front()411 : funcOp.getResultTypes();412 // Get the underlying value types from async types returned from the413 // parent `async.func` operation.414 auto types = llvm::map_range(resultTypes, [](const Type &result) {415 return llvm::cast<ValueType>(result).getValueType();416 });417 418 if (getOperandTypes() != types)419 return emitOpError("operand types do not match the types returned from "420 "the parent FuncOp");421 422 return success();423}424 425//===----------------------------------------------------------------------===//426// TableGen'd op method definitions427//===----------------------------------------------------------------------===//428 429#define GET_OP_CLASSES430#include "mlir/Dialect/Async/IR/AsyncOps.cpp.inc"431 432//===----------------------------------------------------------------------===//433// TableGen'd type method definitions434//===----------------------------------------------------------------------===//435 436#define GET_TYPEDEF_CLASSES437#include "mlir/Dialect/Async/IR/AsyncOpsTypes.cpp.inc"438 439void ValueType::print(AsmPrinter &printer) const {440 printer << "<";441 printer.printType(getValueType());442 printer << '>';443}444 445Type ValueType::parse(mlir::AsmParser &parser) {446 Type ty;447 if (parser.parseLess() || parser.parseType(ty) || parser.parseGreater()) {448 parser.emitError(parser.getNameLoc(), "failed to parse async value type");449 return Type();450 }451 return ValueType::get(ty);452}453