65 lines · cpp
1//===- Visitors.cpp - MLIR Visitor Utilities ------------------------------===//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/IR/Visitors.h"10#include "mlir/IR/Operation.h"11 12using namespace mlir;13 14WalkStage::WalkStage(Operation *op)15 : numRegions(op->getNumRegions()), nextRegion(0) {}16 17MutableArrayRef<Region> ForwardIterator::makeIterable(Operation &range) {18 return range.getRegions();19}20 21void detail::walk(Operation *op,22 function_ref<void(Operation *, const WalkStage &)> callback) {23 WalkStage stage(op);24 25 for (Region ®ion : op->getRegions()) {26 // Invoke callback on the parent op before visiting each child region.27 callback(op, stage);28 stage.advance();29 30 for (Block &block : region) {31 for (Operation &nestedOp : block)32 walk(&nestedOp, callback);33 }34 }35 36 // Invoke callback after all regions have been visited.37 callback(op, stage);38}39 40WalkResult detail::walk(41 Operation *op,42 function_ref<WalkResult(Operation *, const WalkStage &)> callback) {43 WalkStage stage(op);44 45 for (Region ®ion : op->getRegions()) {46 // Invoke callback on the parent op before visiting each child region.47 WalkResult result = callback(op, stage);48 49 if (result.wasSkipped())50 return WalkResult::advance();51 if (result.wasInterrupted())52 return WalkResult::interrupt();53 54 stage.advance();55 56 for (Block &block : region) {57 // Early increment here in the case where the operation is erased.58 for (Operation &nestedOp : llvm::make_early_inc_range(block))59 if (walk(&nestedOp, callback).wasInterrupted())60 return WalkResult::interrupt();61 }62 }63 return callback(op, stage);64}65