688 lines · cpp
1//===- ControlFlowOps.cpp - MLIR SPIR-V Control Flow Ops -----------------===//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// Defines the control flow operations in the SPIR-V dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"14#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"15#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"16#include "mlir/Interfaces/CallInterfaces.h"17 18#include "llvm/Support/InterleavedRange.h"19 20#include "SPIRVOpUtils.h"21#include "SPIRVParsingUtils.h"22 23using namespace mlir::spirv::AttrNames;24 25namespace mlir::spirv {26 27/// Parses Function, Selection and Loop control attributes. If no control is28/// specified, "None" is used as a default.29template <typename EnumAttrClass, typename EnumClass>30static ParseResult31parseControlAttribute(OpAsmParser &parser, OperationState &state,32 StringRef attrName = spirv::attributeName<EnumClass>()) {33 if (succeeded(parser.parseOptionalKeyword(kControl))) {34 EnumClass control;35 if (parser.parseLParen() ||36 spirv::parseEnumKeywordAttr<EnumAttrClass>(control, parser, state) ||37 parser.parseRParen())38 return failure();39 return success();40 }41 // Set control to "None" otherwise.42 Builder builder = parser.getBuilder();43 state.addAttribute(attrName,44 builder.getAttr<EnumAttrClass>(static_cast<EnumClass>(0)));45 return success();46}47 48//===----------------------------------------------------------------------===//49// spirv.BranchOp50//===----------------------------------------------------------------------===//51 52SuccessorOperands BranchOp::getSuccessorOperands(unsigned index) {53 assert(index == 0 && "invalid successor index");54 return SuccessorOperands(0, getTargetOperandsMutable());55}56 57//===----------------------------------------------------------------------===//58// spirv.BranchConditionalOp59//===----------------------------------------------------------------------===//60 61SuccessorOperands BranchConditionalOp::getSuccessorOperands(unsigned index) {62 assert(index < 2 && "invalid successor index");63 return SuccessorOperands(index == kTrueIndex64 ? getTrueTargetOperandsMutable()65 : getFalseTargetOperandsMutable());66}67 68ParseResult BranchConditionalOp::parse(OpAsmParser &parser,69 OperationState &result) {70 auto &builder = parser.getBuilder();71 OpAsmParser::UnresolvedOperand condInfo;72 Block *dest;73 74 // Parse the condition.75 Type boolTy = builder.getI1Type();76 if (parser.parseOperand(condInfo) ||77 parser.resolveOperand(condInfo, boolTy, result.operands))78 return failure();79 80 // Parse the optional branch weights.81 if (succeeded(parser.parseOptionalLSquare())) {82 IntegerAttr trueWeight, falseWeight;83 NamedAttrList weights;84 85 auto i32Type = builder.getIntegerType(32);86 if (parser.parseAttribute(trueWeight, i32Type, "weight", weights) ||87 parser.parseComma() ||88 parser.parseAttribute(falseWeight, i32Type, "weight", weights) ||89 parser.parseRSquare())90 return failure();91 92 StringAttr branchWeightsAttrName =93 BranchConditionalOp::getBranchWeightsAttrName(result.name);94 result.addAttribute(branchWeightsAttrName,95 builder.getArrayAttr({trueWeight, falseWeight}));96 }97 98 // Parse the true branch.99 SmallVector<Value, 4> trueOperands;100 if (parser.parseComma() ||101 parser.parseSuccessorAndUseList(dest, trueOperands))102 return failure();103 result.addSuccessors(dest);104 result.addOperands(trueOperands);105 106 // Parse the false branch.107 SmallVector<Value, 4> falseOperands;108 if (parser.parseComma() ||109 parser.parseSuccessorAndUseList(dest, falseOperands))110 return failure();111 result.addSuccessors(dest);112 result.addOperands(falseOperands);113 result.addAttribute(spirv::BranchConditionalOp::getOperandSegmentSizeAttr(),114 builder.getDenseI32ArrayAttr(115 {1, static_cast<int32_t>(trueOperands.size()),116 static_cast<int32_t>(falseOperands.size())}));117 118 return success();119}120 121void BranchConditionalOp::print(OpAsmPrinter &printer) {122 printer << ' ' << getCondition();123 124 if (std::optional<ArrayAttr> weights = getBranchWeights()) {125 printer << ' '126 << llvm::interleaved_array(weights->getAsValueRange<IntegerAttr>());127 }128 129 printer << ", ";130 printer.printSuccessorAndUseList(getTrueBlock(), getTrueBlockArguments());131 printer << ", ";132 printer.printSuccessorAndUseList(getFalseBlock(), getFalseBlockArguments());133}134 135LogicalResult BranchConditionalOp::verify() {136 if (auto weights = getBranchWeights()) {137 if (weights->getValue().size() != 2) {138 return emitOpError("must have exactly two branch weights");139 }140 if (llvm::all_of(*weights, [](Attribute attr) {141 return llvm::cast<IntegerAttr>(attr).getValue().isZero();142 }))143 return emitOpError("branch weights cannot both be zero");144 }145 146 return success();147}148 149//===----------------------------------------------------------------------===//150// spirv.FunctionCall151//===----------------------------------------------------------------------===//152 153LogicalResult FunctionCallOp::verify() {154 if (getNumResults() > 1) {155 return emitOpError(156 "expected callee function to have 0 or 1 result, but provided ")157 << getNumResults();158 }159 return success();160}161 162LogicalResult163FunctionCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {164 auto fnName = getCalleeAttr();165 166 auto funcOp =167 symbolTable.lookupNearestSymbolFrom<spirv::FuncOp>(*this, fnName);168 if (!funcOp) {169 return emitOpError("callee function '")170 << fnName.getValue() << "' not found in nearest symbol table";171 }172 173 auto functionType = funcOp.getFunctionType();174 175 if (functionType.getNumInputs() != getNumOperands()) {176 return emitOpError("has incorrect number of operands for callee: expected ")177 << functionType.getNumInputs() << ", but provided "178 << getNumOperands();179 }180 181 for (uint32_t i = 0, e = functionType.getNumInputs(); i != e; ++i) {182 if (getOperand(i).getType() != functionType.getInput(i)) {183 return emitOpError("operand type mismatch: expected operand type ")184 << functionType.getInput(i) << ", but provided "185 << getOperand(i).getType() << " for operand number " << i;186 }187 }188 189 if (functionType.getNumResults() != getNumResults()) {190 return emitOpError(191 "has incorrect number of results has for callee: expected ")192 << functionType.getNumResults() << ", but provided "193 << getNumResults();194 }195 196 if (getNumResults() &&197 (getResult(0).getType() != functionType.getResult(0))) {198 return emitOpError("result type mismatch: expected ")199 << functionType.getResult(0) << ", but provided "200 << getResult(0).getType();201 }202 203 return success();204}205 206CallInterfaceCallable FunctionCallOp::getCallableForCallee() {207 return (*this)->getAttrOfType<SymbolRefAttr>(getCalleeAttrName());208}209 210void FunctionCallOp::setCalleeFromCallable(CallInterfaceCallable callee) {211 (*this)->setAttr(getCalleeAttrName(), cast<SymbolRefAttr>(callee));212}213 214Operation::operand_range FunctionCallOp::getArgOperands() {215 return getArguments();216}217 218MutableOperandRange FunctionCallOp::getArgOperandsMutable() {219 return getArgumentsMutable();220}221 222//===----------------------------------------------------------------------===//223// spirv.Switch224//===----------------------------------------------------------------------===//225 226void SwitchOp::build(OpBuilder &builder, OperationState &result, Value selector,227 Block *defaultTarget, ValueRange defaultOperands,228 DenseIntElementsAttr literals, BlockRange targets,229 ArrayRef<ValueRange> targetOperands) {230 build(builder, result, selector, defaultOperands, targetOperands, literals,231 defaultTarget, targets);232}233 234void SwitchOp::build(OpBuilder &builder, OperationState &result, Value selector,235 Block *defaultTarget, ValueRange defaultOperands,236 ArrayRef<APInt> literals, BlockRange targets,237 ArrayRef<ValueRange> targetOperands) {238 DenseIntElementsAttr literalsAttr;239 if (!literals.empty()) {240 ShapedType literalType = VectorType::get(241 static_cast<int64_t>(literals.size()), selector.getType());242 literalsAttr = DenseIntElementsAttr::get(literalType, literals);243 }244 build(builder, result, selector, defaultTarget, defaultOperands, literalsAttr,245 targets, targetOperands);246}247 248void SwitchOp::build(OpBuilder &builder, OperationState &result, Value selector,249 Block *defaultTarget, ValueRange defaultOperands,250 ArrayRef<int32_t> literals, BlockRange targets,251 ArrayRef<ValueRange> targetOperands) {252 DenseIntElementsAttr literalsAttr;253 if (!literals.empty()) {254 ShapedType literalType = VectorType::get(255 static_cast<int64_t>(literals.size()), selector.getType());256 literalsAttr = DenseIntElementsAttr::get(literalType, literals);257 }258 build(builder, result, selector, defaultTarget, defaultOperands, literalsAttr,259 targets, targetOperands);260}261 262LogicalResult SwitchOp::verify() {263 std::optional<DenseIntElementsAttr> literals = getLiterals();264 BlockRange targets = getTargets();265 266 if (!literals && targets.empty())267 return success();268 269 Type selectorType = getSelector().getType();270 Type literalType = literals->getType().getElementType();271 if (literalType != selectorType)272 return emitOpError() << "'selector' type (" << selectorType273 << ") should match literals type (" << literalType274 << ")";275 276 if (literals && literals->size() != static_cast<int64_t>(targets.size()))277 return emitOpError() << "number of literals (" << literals->size()278 << ") should match number of targets ("279 << targets.size() << ")";280 return success();281}282 283SuccessorOperands SwitchOp::getSuccessorOperands(unsigned index) {284 assert(index < getNumSuccessors() && "invalid successor index");285 return SuccessorOperands(index == 0 ? getDefaultOperandsMutable()286 : getTargetOperandsMutable(index - 1));287}288 289Block *SwitchOp::getSuccessorForOperands(ArrayRef<Attribute> operands) {290 std::optional<DenseIntElementsAttr> literals = getLiterals();291 292 if (!literals)293 return getDefaultTarget();294 295 SuccessorRange targets = getTargets();296 if (auto value = dyn_cast_or_null<IntegerAttr>(operands.front())) {297 for (auto [index, literal] : llvm::enumerate(literals->getValues<APInt>()))298 if (literal == value.getValue())299 return targets[index];300 return getDefaultTarget();301 }302 return nullptr;303}304 305//===----------------------------------------------------------------------===//306// spirv.mlir.loop307//===----------------------------------------------------------------------===//308 309void LoopOp::build(OpBuilder &builder, OperationState &state) {310 state.addAttribute("loop_control", builder.getAttr<spirv::LoopControlAttr>(311 spirv::LoopControl::None));312 state.addRegion();313}314 315ParseResult LoopOp::parse(OpAsmParser &parser, OperationState &result) {316 if (parseControlAttribute<spirv::LoopControlAttr, spirv::LoopControl>(parser,317 result))318 return failure();319 320 if (succeeded(parser.parseOptionalArrow()))321 if (parser.parseTypeList(result.types))322 return failure();323 324 return parser.parseRegion(*result.addRegion(), /*arguments=*/{});325}326 327void LoopOp::print(OpAsmPrinter &printer) {328 auto control = getLoopControl();329 if (control != spirv::LoopControl::None)330 printer << " control(" << spirv::stringifyLoopControl(control) << ")";331 if (getNumResults() > 0) {332 printer << " -> ";333 printer << getResultTypes();334 }335 printer << ' ';336 printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false,337 /*printBlockTerminators=*/true);338}339 340/// Returns true if the given `srcBlock` contains only one `spirv.Branch` to the341/// given `dstBlock`.342static bool hasOneBranchOpTo(Block &srcBlock, Block &dstBlock) {343 // Check that there is only one op in the `srcBlock`.344 if (!llvm::hasSingleElement(srcBlock))345 return false;346 347 auto branchOp = dyn_cast<spirv::BranchOp>(srcBlock.back());348 return branchOp && branchOp.getSuccessor() == &dstBlock;349}350 351/// Returns true if the given `block` only contains one `spirv.mlir.merge` op.352static bool isMergeBlock(Block &block) {353 return llvm::hasSingleElement(block) && isa<spirv::MergeOp>(block.front());354}355 356/// Returns true if a `spirv.mlir.merge` op outside the merge block.357static bool hasOtherMerge(Region ®ion) {358 return !region.empty() && llvm::any_of(region.getOps(), [&](Operation &op) {359 return isa<spirv::MergeOp>(op) && op.getBlock() != ®ion.back();360 });361}362 363LogicalResult LoopOp::verifyRegions() {364 auto *op = getOperation();365 366 // We need to verify that the blocks follow the following layout:367 //368 // +-------------+369 // | entry block |370 // +-------------+371 // |372 // v373 // +-------------+374 // | loop header | <-----+375 // +-------------+ |376 // |377 // ... |378 // \ | / |379 // v |380 // +---------------+ |381 // | loop continue | -----+382 // +---------------+383 //384 // ...385 // \ | /386 // v387 // +-------------+388 // | merge block |389 // +-------------+390 391 auto ®ion = op->getRegion(0);392 // Allow empty region as a degenerated case, which can come from393 // optimizations.394 if (region.empty())395 return success();396 397 // The last block is the merge block.398 Block &merge = region.back();399 if (!isMergeBlock(merge))400 return emitOpError("last block must be the merge block with only one "401 "'spirv.mlir.merge' op");402 if (hasOtherMerge(region))403 return emitOpError(404 "should not have 'spirv.mlir.merge' op outside the merge block");405 406 if (region.hasOneBlock())407 return emitOpError(408 "must have an entry block branching to the loop header block");409 // The first block is the entry block.410 Block &entry = region.front();411 412 if (std::next(region.begin(), 2) == region.end())413 return emitOpError(414 "must have a loop header block branched from the entry block");415 // The second block is the loop header block.416 Block &header = *std::next(region.begin(), 1);417 418 if (!hasOneBranchOpTo(entry, header))419 return emitOpError(420 "entry block must only have one 'spirv.Branch' op to the second block");421 422 if (std::next(region.begin(), 3) == region.end())423 return emitOpError(424 "requires a loop continue block branching to the loop header block");425 // The second to last block is the loop continue block.426 Block &cont = *std::prev(region.end(), 2);427 428 // Make sure that we have a branch from the loop continue block to the loop429 // header block.430 if (llvm::none_of(431 llvm::seq<unsigned>(0, cont.getNumSuccessors()),432 [&](unsigned index) { return cont.getSuccessor(index) == &header; }))433 return emitOpError("second to last block must be the loop continue "434 "block that branches to the loop header block");435 436 // Make sure that no other blocks (except the entry and loop continue block)437 // branches to the loop header block.438 for (auto &block : llvm::make_range(std::next(region.begin(), 2),439 std::prev(region.end(), 2))) {440 for (auto i : llvm::seq<unsigned>(0, block.getNumSuccessors())) {441 if (block.getSuccessor(i) == &header) {442 return emitOpError("can only have the entry and loop continue "443 "block branching to the loop header block");444 }445 }446 }447 448 return success();449}450 451Block *LoopOp::getEntryBlock() {452 assert(!getBody().empty() && "op region should not be empty!");453 return &getBody().front();454}455 456Block *LoopOp::getHeaderBlock() {457 assert(!getBody().empty() && "op region should not be empty!");458 // The second block is the loop header block.459 return &*std::next(getBody().begin());460}461 462Block *LoopOp::getContinueBlock() {463 assert(!getBody().empty() && "op region should not be empty!");464 // The second to last block is the loop continue block.465 return &*std::prev(getBody().end(), 2);466}467 468Block *LoopOp::getMergeBlock() {469 assert(!getBody().empty() && "op region should not be empty!");470 // The last block is the loop merge block.471 return &getBody().back();472}473 474void LoopOp::addEntryAndMergeBlock(OpBuilder &builder) {475 assert(getBody().empty() && "entry and merge block already exist");476 OpBuilder::InsertionGuard g(builder);477 builder.createBlock(&getBody());478 builder.createBlock(&getBody());479 480 // Add a spirv.mlir.merge op into the merge block.481 spirv::MergeOp::create(builder, getLoc());482}483 484//===----------------------------------------------------------------------===//485// spirv.Return486//===----------------------------------------------------------------------===//487 488LogicalResult ReturnOp::verify() {489 // Verification is performed in spirv.func op.490 return success();491}492 493//===----------------------------------------------------------------------===//494// spirv.ReturnValue495//===----------------------------------------------------------------------===//496 497LogicalResult ReturnValueOp::verify() {498 // Verification is performed in spirv.func op.499 return success();500}501 502//===----------------------------------------------------------------------===//503// spirv.Select504//===----------------------------------------------------------------------===//505 506LogicalResult SelectOp::verify() {507 if (auto conditionTy = llvm::dyn_cast<VectorType>(getCondition().getType())) {508 auto resultVectorTy = llvm::dyn_cast<VectorType>(getResult().getType());509 if (!resultVectorTy) {510 return emitOpError("result expected to be of vector type when "511 "condition is of vector type");512 }513 if (resultVectorTy.getNumElements() != conditionTy.getNumElements()) {514 return emitOpError("result should have the same number of elements as "515 "the condition when condition is of vector type");516 }517 }518 return success();519}520 521// Custom availability implementation is needed for spirv.Select given the522// syntax changes starting v1.4.523SmallVector<ArrayRef<spirv::Extension>, 1> SelectOp::getExtensions() {524 return {};525}526SmallVector<ArrayRef<spirv::Capability>, 1> SelectOp::getCapabilities() {527 return {};528}529std::optional<spirv::Version> SelectOp::getMinVersion() {530 // Per the spec, "Before version 1.4, results are only computed per531 // component."532 if (isa<spirv::ScalarType>(getCondition().getType()) &&533 isa<spirv::CompositeType>(getType()))534 return Version::V_1_4;535 536 return Version::V_1_0;537}538std::optional<spirv::Version> SelectOp::getMaxVersion() {539 return Version::V_1_6;540}541 542//===----------------------------------------------------------------------===//543// spirv.mlir.selection544//===----------------------------------------------------------------------===//545 546ParseResult SelectionOp::parse(OpAsmParser &parser, OperationState &result) {547 if (parseControlAttribute<spirv::SelectionControlAttr,548 spirv::SelectionControl>(parser, result))549 return failure();550 551 if (succeeded(parser.parseOptionalArrow()))552 if (parser.parseTypeList(result.types))553 return failure();554 555 return parser.parseRegion(*result.addRegion(), /*arguments=*/{});556}557 558void SelectionOp::print(OpAsmPrinter &printer) {559 auto control = getSelectionControl();560 if (control != spirv::SelectionControl::None)561 printer << " control(" << spirv::stringifySelectionControl(control) << ")";562 if (getNumResults() > 0) {563 printer << " -> ";564 printer << getResultTypes();565 }566 printer << ' ';567 printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false,568 /*printBlockTerminators=*/true);569}570 571LogicalResult SelectionOp::verifyRegions() {572 auto *op = getOperation();573 574 // We need to verify that the blocks follow the following layout:575 //576 // +--------------+577 // | header block |578 // +--------------+579 // / | \580 // ...581 //582 //583 // +---------+ +---------+ +---------+584 // | case #0 | | case #1 | | case #2 | ...585 // +---------+ +---------+ +---------+586 //587 //588 // ...589 // \ | /590 // v591 // +-------------+592 // | merge block |593 // +-------------+594 595 auto ®ion = op->getRegion(0);596 // Allow empty region as a degenerated case, which can come from597 // optimizations.598 if (region.empty())599 return success();600 601 // The last block is the merge block.602 if (!isMergeBlock(region.back()))603 return emitOpError("last block must be the merge block with only one "604 "'spirv.mlir.merge' op");605 if (hasOtherMerge(region))606 return emitOpError(607 "should not have 'spirv.mlir.merge' op outside the merge block");608 609 if (region.hasOneBlock())610 return emitOpError("must have a selection header block");611 612 return success();613}614 615Block *SelectionOp::getHeaderBlock() {616 assert(!getBody().empty() && "op region should not be empty!");617 // The first block is the loop header block.618 return &getBody().front();619}620 621Block *SelectionOp::getMergeBlock() {622 assert(!getBody().empty() && "op region should not be empty!");623 // The last block is the loop merge block.624 return &getBody().back();625}626 627void SelectionOp::addMergeBlock(OpBuilder &builder) {628 assert(getBody().empty() && "entry and merge block already exist");629 OpBuilder::InsertionGuard guard(builder);630 builder.createBlock(&getBody());631 632 // Add a spirv.mlir.merge op into the merge block.633 spirv::MergeOp::create(builder, getLoc());634}635 636SelectionOp637SelectionOp::createIfThen(Location loc, Value condition,638 function_ref<void(OpBuilder &builder)> thenBody,639 OpBuilder &builder) {640 auto selectionOp =641 spirv::SelectionOp::create(builder, loc, spirv::SelectionControl::None);642 643 selectionOp.addMergeBlock(builder);644 Block *mergeBlock = selectionOp.getMergeBlock();645 Block *thenBlock = nullptr;646 647 // Build the "then" block.648 {649 OpBuilder::InsertionGuard guard(builder);650 thenBlock = builder.createBlock(mergeBlock);651 thenBody(builder);652 spirv::BranchOp::create(builder, loc, mergeBlock);653 }654 655 // Build the header block.656 {657 OpBuilder::InsertionGuard guard(builder);658 builder.createBlock(thenBlock);659 spirv::BranchConditionalOp::create(builder, loc, condition, thenBlock,660 /*trueArguments=*/ArrayRef<Value>(),661 mergeBlock,662 /*falseArguments=*/ArrayRef<Value>());663 }664 665 return selectionOp;666}667 668//===----------------------------------------------------------------------===//669// spirv.Unreachable670//===----------------------------------------------------------------------===//671 672LogicalResult spirv::UnreachableOp::verify() {673 auto *block = (*this)->getBlock();674 // Fast track: if this is in entry block, its invalid. Otherwise, if no675 // predecessors, it's valid.676 if (block->isEntryBlock())677 return emitOpError("cannot be used in reachable block");678 if (block->hasNoPredecessors())679 return success();680 681 // TODO: further verification needs to analyze reachability from682 // the entry block.683 684 return success();685}686 687} // namespace mlir::spirv688