532 lines · cpp
1//===- ControlFlowInterfaces.cpp - ControlFlow Interfaces -----------------===//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 <utility>10 11#include "mlir/IR/BuiltinTypes.h"12#include "mlir/IR/Operation.h"13#include "mlir/Interfaces/ControlFlowInterfaces.h"14#include "llvm/Support/DebugLog.h"15 16using namespace mlir;17 18//===----------------------------------------------------------------------===//19// ControlFlowInterfaces20//===----------------------------------------------------------------------===//21 22#include "mlir/Interfaces/ControlFlowInterfaces.cpp.inc"23 24SuccessorOperands::SuccessorOperands(MutableOperandRange forwardedOperands)25 : producedOperandCount(0), forwardedOperands(std::move(forwardedOperands)) {26}27 28SuccessorOperands::SuccessorOperands(unsigned int producedOperandCount,29 MutableOperandRange forwardedOperands)30 : producedOperandCount(producedOperandCount),31 forwardedOperands(std::move(forwardedOperands)) {}32 33//===----------------------------------------------------------------------===//34// BranchOpInterface35//===----------------------------------------------------------------------===//36 37/// Returns the `BlockArgument` corresponding to operand `operandIndex` in some38/// successor if 'operandIndex' is within the range of 'operands', or39/// std::nullopt if `operandIndex` isn't a successor operand index.40std::optional<BlockArgument>41detail::getBranchSuccessorArgument(const SuccessorOperands &operands,42 unsigned operandIndex, Block *successor) {43 LDBG() << "Getting branch successor argument for operand index "44 << operandIndex << " in successor block";45 46 OperandRange forwardedOperands = operands.getForwardedOperands();47 // Check that the operands are valid.48 if (forwardedOperands.empty()) {49 LDBG() << "No forwarded operands, returning nullopt";50 return std::nullopt;51 }52 53 // Check to ensure that this operand is within the range.54 unsigned operandsStart = forwardedOperands.getBeginOperandIndex();55 if (operandIndex < operandsStart ||56 operandIndex >= (operandsStart + forwardedOperands.size())) {57 LDBG() << "Operand index " << operandIndex << " out of range ["58 << operandsStart << ", "59 << (operandsStart + forwardedOperands.size())60 << "), returning nullopt";61 return std::nullopt;62 }63 64 // Index the successor.65 unsigned argIndex =66 operands.getProducedOperandCount() + operandIndex - operandsStart;67 LDBG() << "Computed argument index " << argIndex << " for successor block";68 return successor->getArgument(argIndex);69}70 71/// Verify that the given operands match those of the given successor block.72LogicalResult73detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,74 const SuccessorOperands &operands) {75 LDBG() << "Verifying branch successor operands for successor #" << succNo76 << " in operation " << op->getName();77 78 // Check the count.79 unsigned operandCount = operands.size();80 Block *destBB = op->getSuccessor(succNo);81 LDBG() << "Branch has " << operandCount << " operands, target block has "82 << destBB->getNumArguments() << " arguments";83 84 if (operandCount != destBB->getNumArguments())85 return op->emitError() << "branch has " << operandCount86 << " operands for successor #" << succNo87 << ", but target block has "88 << destBB->getNumArguments();89 90 // Check the types.91 LDBG() << "Checking type compatibility for "92 << (operandCount - operands.getProducedOperandCount())93 << " forwarded operands";94 for (unsigned i = operands.getProducedOperandCount(); i != operandCount;95 ++i) {96 Type operandType = operands[i].getType();97 Type argType = destBB->getArgument(i).getType();98 LDBG() << "Checking type compatibility: operand type " << operandType99 << " vs argument type " << argType;100 101 if (!cast<BranchOpInterface>(op).areTypesCompatible(operandType, argType))102 return op->emitError() << "type mismatch for bb argument #" << i103 << " of successor #" << succNo;104 }105 106 LDBG() << "Branch successor operand verification successful";107 return success();108}109 110//===----------------------------------------------------------------------===//111// WeightedBranchOpInterface112//===----------------------------------------------------------------------===//113 114static LogicalResult verifyWeights(Operation *op,115 llvm::ArrayRef<int32_t> weights,116 std::size_t expectedWeightsNum,117 llvm::StringRef weightAnchorName,118 llvm::StringRef weightRefName) {119 if (weights.empty())120 return success();121 122 if (weights.size() != expectedWeightsNum)123 return op->emitError() << "expects number of " << weightAnchorName124 << " weights to match number of " << weightRefName125 << ": " << weights.size() << " vs "126 << expectedWeightsNum;127 128 if (llvm::all_of(weights, [](int32_t value) { return value == 0; }))129 return op->emitError() << "branch weights cannot all be zero";130 131 return success();132}133 134LogicalResult detail::verifyBranchWeights(Operation *op) {135 llvm::ArrayRef<int32_t> weights =136 cast<WeightedBranchOpInterface>(op).getWeights();137 return verifyWeights(op, weights, op->getNumSuccessors(), "branch",138 "successors");139}140 141//===----------------------------------------------------------------------===//142// WeightedRegionBranchOpInterface143//===----------------------------------------------------------------------===//144 145LogicalResult detail::verifyRegionBranchWeights(Operation *op) {146 llvm::ArrayRef<int32_t> weights =147 cast<WeightedRegionBranchOpInterface>(op).getWeights();148 return verifyWeights(op, weights, op->getNumRegions(), "region", "regions");149}150 151//===----------------------------------------------------------------------===//152// RegionBranchOpInterface153//===----------------------------------------------------------------------===//154 155static InFlightDiagnostic &printRegionEdgeName(InFlightDiagnostic &diag,156 RegionBranchPoint sourceNo,157 RegionSuccessor succRegionNo) {158 diag << "from ";159 if (Operation *op = sourceNo.getTerminatorPredecessorOrNull())160 diag << "Operation " << op->getName();161 else162 diag << "parent operands";163 164 diag << " to ";165 if (Region *region = succRegionNo.getSuccessor())166 diag << "Region #" << region->getRegionNumber();167 else168 diag << "parent results";169 return diag;170}171 172/// Verify that types match along all region control flow edges originating from173/// `sourcePoint`. `getInputsTypesForRegion` is a function that returns the174/// types of the inputs that flow to a successor region.175static LogicalResult176verifyTypesAlongAllEdges(RegionBranchOpInterface branchOp,177 RegionBranchPoint sourcePoint,178 function_ref<FailureOr<TypeRange>(RegionSuccessor)>179 getInputsTypesForRegion) {180 SmallVector<RegionSuccessor, 2> successors;181 branchOp.getSuccessorRegions(sourcePoint, successors);182 183 for (RegionSuccessor &succ : successors) {184 FailureOr<TypeRange> sourceTypes = getInputsTypesForRegion(succ);185 if (failed(sourceTypes))186 return failure();187 188 TypeRange succInputsTypes = succ.getSuccessorInputs().getTypes();189 if (sourceTypes->size() != succInputsTypes.size()) {190 InFlightDiagnostic diag =191 branchOp->emitOpError("region control flow edge ");192 std::string succStr;193 llvm::raw_string_ostream os(succStr);194 os << succ;195 return printRegionEdgeName(diag, sourcePoint, succ)196 << ": source has " << sourceTypes->size()197 << " operands, but target successor " << os.str() << " needs "198 << succInputsTypes.size();199 }200 201 for (const auto &typesIdx :202 llvm::enumerate(llvm::zip(*sourceTypes, succInputsTypes))) {203 Type sourceType = std::get<0>(typesIdx.value());204 Type inputType = std::get<1>(typesIdx.value());205 206 if (!branchOp.areTypesCompatible(sourceType, inputType)) {207 InFlightDiagnostic diag =208 branchOp->emitOpError("along control flow edge ");209 return printRegionEdgeName(diag, sourcePoint, succ)210 << ": source type #" << typesIdx.index() << " " << sourceType211 << " should match input type #" << typesIdx.index() << " "212 << inputType;213 }214 }215 }216 217 return success();218}219 220/// Verify that types match along control flow edges described the given op.221LogicalResult detail::verifyTypesAlongControlFlowEdges(Operation *op) {222 auto regionInterface = cast<RegionBranchOpInterface>(op);223 224 auto inputTypesFromParent = [&](RegionSuccessor successor) -> TypeRange {225 return regionInterface.getEntrySuccessorOperands(successor).getTypes();226 };227 228 // Verify types along control flow edges originating from the parent.229 if (failed(verifyTypesAlongAllEdges(230 regionInterface, RegionBranchPoint::parent(), inputTypesFromParent)))231 return failure();232 233 // Verify types along control flow edges originating from each region.234 for (Region ®ion : op->getRegions()) {235 // Collect all return-like terminators in the region.236 SmallVector<RegionBranchTerminatorOpInterface> regionReturnOps;237 for (Block &block : region)238 if (!block.empty())239 if (auto terminator =240 dyn_cast<RegionBranchTerminatorOpInterface>(block.back()))241 regionReturnOps.push_back(terminator);242 243 // If there is no return-like terminator, the op itself should verify244 // type consistency.245 if (regionReturnOps.empty())246 continue;247 248 // Verify types along control flow edges originating from each return-like249 // terminator.250 for (RegionBranchTerminatorOpInterface regionReturnOp : regionReturnOps) {251 252 auto inputTypesForRegion =253 [&](RegionSuccessor successor) -> FailureOr<TypeRange> {254 OperandRange terminatorOperands =255 regionReturnOp.getSuccessorOperands(successor);256 return TypeRange(terminatorOperands.getTypes());257 };258 if (failed(verifyTypesAlongAllEdges(regionInterface, regionReturnOp,259 inputTypesForRegion)))260 return failure();261 }262 }263 264 return success();265}266 267/// Stop condition for `traverseRegionGraph`. The traversal is interrupted if268/// this function returns "true" for a successor region. The first parameter is269/// the successor region. The second parameter indicates all already visited270/// regions.271using StopConditionFn = function_ref<bool(Region *, ArrayRef<bool> visited)>;272 273/// Traverse the region graph starting at `begin`. The traversal is interrupted274/// if `stopCondition` evaluates to "true" for a successor region. In that case,275/// this function returns "true". Otherwise, if the traversal was not276/// interrupted, this function returns "false".277static bool traverseRegionGraph(Region *begin,278 StopConditionFn stopConditionFn) {279 auto op = cast<RegionBranchOpInterface>(begin->getParentOp());280 LDBG() << "Starting region graph traversal from region #"281 << begin->getRegionNumber() << " in operation " << op->getName();282 283 SmallVector<bool> visited(op->getNumRegions(), false);284 visited[begin->getRegionNumber()] = true;285 LDBG() << "Initialized visited array with " << op->getNumRegions()286 << " regions";287 288 // Retrieve all successors of the region and enqueue them in the worklist.289 SmallVector<Region *> worklist;290 auto enqueueAllSuccessors = [&](Region *region) {291 LDBG() << "Enqueuing successors for region #" << region->getRegionNumber();292 SmallVector<Attribute> operandAttributes(op->getNumOperands());293 for (Block &block : *region) {294 if (block.empty())295 continue;296 auto terminator =297 dyn_cast<RegionBranchTerminatorOpInterface>(block.back());298 if (!terminator)299 continue;300 SmallVector<RegionSuccessor> successors;301 operandAttributes.resize(terminator->getNumOperands());302 terminator.getSuccessorRegions(operandAttributes, successors);303 LDBG() << "Found " << successors.size()304 << " successors from terminator in block";305 for (RegionSuccessor successor : successors) {306 if (!successor.isParent()) {307 worklist.push_back(successor.getSuccessor());308 LDBG() << "Added region #"309 << successor.getSuccessor()->getRegionNumber()310 << " to worklist";311 } else {312 LDBG() << "Skipping parent successor";313 }314 }315 }316 };317 enqueueAllSuccessors(begin);318 LDBG() << "Initial worklist size: " << worklist.size();319 320 // Process all regions in the worklist via DFS.321 while (!worklist.empty()) {322 Region *nextRegion = worklist.pop_back_val();323 LDBG() << "Processing region #" << nextRegion->getRegionNumber()324 << " from worklist (remaining: " << worklist.size() << ")";325 326 if (stopConditionFn(nextRegion, visited)) {327 LDBG() << "Stop condition met for region #"328 << nextRegion->getRegionNumber() << ", returning true";329 return true;330 }331 if (!nextRegion->getParentOp()) {332 llvm::errs() << "Region " << *nextRegion << " has no parent op\n";333 return false;334 }335 if (visited[nextRegion->getRegionNumber()]) {336 LDBG() << "Region #" << nextRegion->getRegionNumber()337 << " already visited, skipping";338 continue;339 }340 visited[nextRegion->getRegionNumber()] = true;341 LDBG() << "Marking region #" << nextRegion->getRegionNumber()342 << " as visited";343 enqueueAllSuccessors(nextRegion);344 }345 346 LDBG() << "Traversal completed, returning false";347 return false;348}349 350/// Return `true` if region `r` is reachable from region `begin` according to351/// the RegionBranchOpInterface (by taking a branch).352static bool isRegionReachable(Region *begin, Region *r) {353 assert(begin->getParentOp() == r->getParentOp() &&354 "expected that both regions belong to the same op");355 return traverseRegionGraph(begin,356 [&](Region *nextRegion, ArrayRef<bool> visited) {357 // Interrupt traversal if `r` was reached.358 return nextRegion == r;359 });360}361 362/// Return `true` if `a` and `b` are in mutually exclusive regions.363///364/// 1. Find the first common of `a` and `b` (ancestor) that implements365/// RegionBranchOpInterface.366/// 2. Determine the regions `regionA` and `regionB` in which `a` and `b` are367/// contained.368/// 3. Check if `regionA` and `regionB` are mutually exclusive. They are369/// mutually exclusive if they are not reachable from each other as per370/// RegionBranchOpInterface::getSuccessorRegions.371bool mlir::insideMutuallyExclusiveRegions(Operation *a, Operation *b) {372 LDBG() << "Checking if operations are in mutually exclusive regions: "373 << a->getName() << " and " << b->getName();374 375 assert(a && "expected non-empty operation");376 assert(b && "expected non-empty operation");377 378 auto branchOp = a->getParentOfType<RegionBranchOpInterface>();379 while (branchOp) {380 LDBG() << "Checking branch operation " << branchOp->getName();381 382 // Check if b is inside branchOp. (We already know that a is.)383 if (!branchOp->isProperAncestor(b)) {384 LDBG() << "Operation b is not inside branchOp, checking next ancestor";385 // Check next enclosing RegionBranchOpInterface.386 branchOp = branchOp->getParentOfType<RegionBranchOpInterface>();387 continue;388 }389 390 LDBG() << "Both operations are inside branchOp, finding their regions";391 392 // b is contained in branchOp. Retrieve the regions in which `a` and `b`393 // are contained.394 Region *regionA = nullptr, *regionB = nullptr;395 for (Region &r : branchOp->getRegions()) {396 if (r.findAncestorOpInRegion(*a)) {397 assert(!regionA && "already found a region for a");398 regionA = &r;399 LDBG() << "Found region #" << r.getRegionNumber() << " for operation a";400 }401 if (r.findAncestorOpInRegion(*b)) {402 assert(!regionB && "already found a region for b");403 regionB = &r;404 LDBG() << "Found region #" << r.getRegionNumber() << " for operation b";405 }406 }407 assert(regionA && regionB && "could not find region of op");408 409 LDBG() << "Region A: #" << regionA->getRegionNumber() << ", Region B: #"410 << regionB->getRegionNumber();411 412 // `a` and `b` are in mutually exclusive regions if both regions are413 // distinct and neither region is reachable from the other region.414 bool regionsAreDistinct = (regionA != regionB);415 bool aNotReachableFromB = !isRegionReachable(regionA, regionB);416 bool bNotReachableFromA = !isRegionReachable(regionB, regionA);417 418 LDBG() << "Regions distinct: " << regionsAreDistinct419 << ", A not reachable from B: " << aNotReachableFromB420 << ", B not reachable from A: " << bNotReachableFromA;421 422 bool mutuallyExclusive =423 regionsAreDistinct && aNotReachableFromB && bNotReachableFromA;424 LDBG() << "Operations are mutually exclusive: " << mutuallyExclusive;425 426 return mutuallyExclusive;427 }428 429 // Could not find a common RegionBranchOpInterface among a's and b's430 // ancestors.431 LDBG() << "No common RegionBranchOpInterface found, operations are not "432 "mutually exclusive";433 return false;434}435 436bool RegionBranchOpInterface::isRepetitiveRegion(unsigned index) {437 LDBG() << "Checking if region #" << index << " is repetitive in operation "438 << getOperation()->getName();439 440 Region *region = &getOperation()->getRegion(index);441 bool isRepetitive = isRegionReachable(region, region);442 443 LDBG() << "Region #" << index << " is repetitive: " << isRepetitive;444 return isRepetitive;445}446 447bool RegionBranchOpInterface::hasLoop() {448 LDBG() << "Checking if operation " << getOperation()->getName()449 << " has loops";450 451 SmallVector<RegionSuccessor> entryRegions;452 getSuccessorRegions(RegionBranchPoint::parent(), entryRegions);453 LDBG() << "Found " << entryRegions.size() << " entry regions";454 455 for (RegionSuccessor successor : entryRegions) {456 if (!successor.isParent()) {457 LDBG() << "Checking entry region #"458 << successor.getSuccessor()->getRegionNumber() << " for loops";459 460 bool hasLoop =461 traverseRegionGraph(successor.getSuccessor(),462 [](Region *nextRegion, ArrayRef<bool> visited) {463 // Interrupt traversal if the region was already464 // visited.465 return visited[nextRegion->getRegionNumber()];466 });467 468 if (hasLoop) {469 LDBG() << "Found loop in entry region #"470 << successor.getSuccessor()->getRegionNumber();471 return true;472 }473 } else {474 LDBG() << "Skipping parent successor";475 }476 }477 478 LDBG() << "No loops found in operation";479 return false;480}481 482Region *mlir::getEnclosingRepetitiveRegion(Operation *op) {483 LDBG() << "Finding enclosing repetitive region for operation "484 << op->getName();485 486 while (Region *region = op->getParentRegion()) {487 LDBG() << "Checking region #" << region->getRegionNumber()488 << " in operation " << region->getParentOp()->getName();489 490 op = region->getParentOp();491 if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op)) {492 LDBG()493 << "Found RegionBranchOpInterface, checking if region is repetitive";494 if (branchOp.isRepetitiveRegion(region->getRegionNumber())) {495 LDBG() << "Found repetitive region #" << region->getRegionNumber();496 return region;497 }498 } else {499 LDBG() << "Parent operation does not implement RegionBranchOpInterface";500 }501 }502 503 LDBG() << "No enclosing repetitive region found";504 return nullptr;505}506 507Region *mlir::getEnclosingRepetitiveRegion(Value value) {508 LDBG() << "Finding enclosing repetitive region for value";509 510 Region *region = value.getParentRegion();511 while (region) {512 LDBG() << "Checking region #" << region->getRegionNumber()513 << " in operation " << region->getParentOp()->getName();514 515 Operation *op = region->getParentOp();516 if (auto branchOp = dyn_cast<RegionBranchOpInterface>(op)) {517 LDBG()518 << "Found RegionBranchOpInterface, checking if region is repetitive";519 if (branchOp.isRepetitiveRegion(region->getRegionNumber())) {520 LDBG() << "Found repetitive region #" << region->getRegionNumber();521 return region;522 }523 } else {524 LDBG() << "Parent operation does not implement RegionBranchOpInterface";525 }526 region = op->getParentRegion();527 }528 529 LDBG() << "No enclosing repetitive region found for value";530 return nullptr;531}532