Page MenuHomePhabricator
Differential D114482 Diff 394698 mlir/test/lib/IR/TestVisitors.cpp

Changeset View

Standalone View

View Options

mlir/test/lib/IR/TestVisitors.cpp

//===- TestIRVisitors.cpp - Pass to test the IR visitors ------------------===// //===- TestIRVisitors.cpp - Pass to test the IR visitors ------------------===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
using namespace mlir; using namespace mlir;
static const char dummyOpName[] = "dummy";
static void printRegion(Region *region) { static void printRegion(Region *region) {
llvm::outs() << "region " << region->getRegionNumber() << " from operation '" llvm::outs() << "region " << region->getRegionNumber() << " from operation '"
<< region->getParentOp()->getName() << "'"; << region->getParentOp()->getName() << "'";
} }
static void printBlock(Block *block) { static void printBlock(Block *block) {
llvm::outs() << "block "; llvm::outs() << "block ";
block->printAsOperand(llvm::outs(), /*printType=*/false); block->printAsOperand(llvm::outs(), /*printType=*/false);
llvm::outs() << " from "; llvm::outs() << " from ";
printRegion(block->getParent()); printRegion(block->getParent());
} }
static void printOperation(Operation *op) { static void printOperation(Operation *op) {
llvm::outs() << "op '" << op->getName() << "'"; llvm::outs() << "op '" << op->getName() << "'";
} }
/// Tests pure callbacks. /// Tests pure callbacks.
static void testPureCallbacks(Operation *op) { static void testPureCallbacks(Operation *op) {
auto opPure = [](Operation *op) { auto opPure = [](Operation *op) {
llvm::outs() << "Visiting "; llvm::outs() << "Visiting ";
printOperation(op); printOperation(op);
llvm::outs() << "\n"; llvm::outs() << "\n";
}; };
auto opName = [](UnregisteredOp<dummyOpName> op) {
llvm::outs() << "Visiting ";
printOperation(op);
llvm::outs() << " with " << op->getNumOperands() << " operand\n";
};
auto blockPure = [](Block *block) { auto blockPure = [](Block *block) {
llvm::outs() << "Visiting "; llvm::outs() << "Visiting ";
printBlock(block); printBlock(block);
llvm::outs() << "\n"; llvm::outs() << "\n";
}; };
auto regionPure = [](Region *region) { auto regionPure = [](Region *region) {
llvm::outs() << "Visiting "; llvm::outs() << "Visiting ";
printRegion(region); printRegion(region);
llvm::outs() << "\n"; llvm::outs() << "\n";
}; };
llvm::outs() << "Op pre-order visits" llvm::outs() << "Op pre-order visits"
<< "\n"; << "\n";
op->walk<WalkOrder::PreOrder>(opPure); op->walk<WalkOrder::PreOrder>(opPure);
llvm::outs() << "Op name pre-order visits"
<< "\n";
op->walk<WalkOrder::PreOrder>(opName);
llvm::outs() << "Block pre-order visits" llvm::outs() << "Block pre-order visits"
<< "\n"; << "\n";
op->walk<WalkOrder::PreOrder>(blockPure); op->walk<WalkOrder::PreOrder>(blockPure);
llvm::outs() << "Region pre-order visits" llvm::outs() << "Region pre-order visits"
<< "\n"; << "\n";
op->walk<WalkOrder::PreOrder>(regionPure); op->walk<WalkOrder::PreOrder>(regionPure);
llvm::outs() << "Op post-order visits" llvm::outs() << "Op post-order visits"
<< "\n"; << "\n";
op->walk<WalkOrder::PostOrder>(opPure); op->walk<WalkOrder::PostOrder>(opPure);
llvm::outs() << "Op name post-order visits"
<< "\n";
op->walk<WalkOrder::PostOrder>(opName);
llvm::outs() << "Block post-order visits" llvm::outs() << "Block post-order visits"
<< "\n"; << "\n";
op->walk<WalkOrder::PostOrder>(blockPure); op->walk<WalkOrder::PostOrder>(blockPure);
llvm::outs() << "Region post-order visits" llvm::outs() << "Region post-order visits"
<< "\n"; << "\n";
op->walk<WalkOrder::PostOrder>(regionPure); op->walk<WalkOrder::PostOrder>(regionPure);
} }
/// Tests erasure callbacks that skip the walk. /// Tests erasure callbacks that skip the walk.
static void testSkipErasureCallbacks(Operation *op) { static void testSkipErasureCallbacks(Operation *op) {
auto skipOpErasure = [](Operation *op) { auto skipOpErasure = [](Operation *op) {
// Do not erase module and function op. Otherwise there wouldn't be too // Do not erase module and function op. Otherwise there wouldn't be too
// much to test in pre-order. // much to test in pre-order.
if (isa<ModuleOp>(op) || isa<FuncOp>(op)) if (isa<ModuleOp>(op) || isa<FuncOp>(op))
return WalkResult::advance(); return WalkResult::advance();
llvm::outs() << "Erasing "; llvm::outs() << "Erasing ";
printOperation(op); printOperation(op);
llvm::outs() << "\n"; llvm::outs() << "\n";
op->dropAllUses(); op->dropAllUses();
op->erase(); op->erase();
return WalkResult::skip(); return WalkResult::skip();
}; };
auto skipOpNameErasure = [](UnregisteredOp<dummyOpName> op) {
llvm::outs() << "Erasing ";
printOperation(op);
llvm::outs() << " with " << op->getNumOperands() << " operand\n";
op->dropAllUses();
op->erase();
return WalkResult::skip();
};
auto skipBlockErasure = [](Block *block) { auto skipBlockErasure = [](Block *block) {
// Do not erase module and function blocks. Otherwise there wouldn't be // Do not erase module and function blocks. Otherwise there wouldn't be
// too much to test in pre-order. // too much to test in pre-order.
Operation *parentOp = block->getParentOp(); Operation *parentOp = block->getParentOp();
if (isa<ModuleOp>(parentOp) || isa<FuncOp>(parentOp)) if (isa<ModuleOp>(parentOp) || isa<FuncOp>(parentOp))
return WalkResult::advance(); return WalkResult::advance();
llvm::outs() << "Erasing "; llvm::outs() << "Erasing ";
printBlock(block); printBlock(block);
llvm::outs() << "\n"; llvm::outs() << "\n";
block->erase(); block->erase();
return WalkResult::skip(); return WalkResult::skip();
}; };
llvm::outs() << "Op pre-order erasures (skip)" llvm::outs() << "Op pre-order erasures (skip)"
<< "\n"; << "\n";
Operation *cloned = op->clone(); Operation *cloned = op->clone();
cloned->walk<WalkOrder::PreOrder>(skipOpErasure); cloned->walk<WalkOrder::PreOrder>(skipOpErasure);
cloned->erase(); cloned->erase();
llvm::outs() << "Op name pre-order erasures (skip)"
<< "\n";
cloned = op->clone();
cloned->walk<WalkOrder::PreOrder>(skipOpNameErasure);
cloned->erase();
llvm::outs() << "Block pre-order erasures (skip)" llvm::outs() << "Block pre-order erasures (skip)"
<< "\n"; << "\n";
cloned = op->clone(); cloned = op->clone();
cloned->walk<WalkOrder::PreOrder>(skipBlockErasure); cloned->walk<WalkOrder::PreOrder>(skipBlockErasure);
cloned->erase(); cloned->erase();
llvm::outs() << "Op post-order erasures (skip)" llvm::outs() << "Op post-order erasures (skip)"
<< "\n"; << "\n";
cloned = op->clone(); cloned = op->clone();
cloned->walk<WalkOrder::PostOrder>(skipOpErasure); cloned->walk<WalkOrder::PostOrder>(skipOpErasure);
cloned->erase(); cloned->erase();
llvm::outs() << "Op name post-order erasures (skip)"
<< "\n";
cloned = op->clone();
cloned->walk<WalkOrder::PostOrder>(skipOpNameErasure);
cloned->erase();
llvm::outs() << "Block post-order erasures (skip)" llvm::outs() << "Block post-order erasures (skip)"
<< "\n"; << "\n";
cloned = op->clone(); cloned = op->clone();
cloned->walk<WalkOrder::PostOrder>(skipBlockErasure); cloned->walk<WalkOrder::PostOrder>(skipBlockErasure);
cloned->erase(); cloned->erase();
} }
/// Tests callbacks that erase the op or block but don't return 'Skip'. This /// Tests callbacks that erase the op or block but don't return 'Skip'. This
/// callbacks are only valid in post-order. /// callbacks are only valid in post-order.
static void testNoSkipErasureCallbacks(Operation *op) { static void testNoSkipErasureCallbacks(Operation *op) {
auto noSkipOpErasure = [](Operation *op) { auto noSkipOpErasure = [](Operation *op) {
llvm::outs() << "Erasing "; llvm::outs() << "Erasing ";
printOperation(op); printOperation(op);
llvm::outs() << "\n"; llvm::outs() << "\n";
op->dropAllUses(); op->dropAllUses();
op->erase(); op->erase();
}; };
auto noSkipOpNameErasure = [](UnregisteredOp<dummyOpName> op) {
llvm::outs() << "Erasing ";
printOperation(op);
llvm::outs() << " with " << op->getNumOperands() << " operand\n";
op->dropAllUses();
op->erase();
};
auto noSkipBlockErasure = [](Block *block) { auto noSkipBlockErasure = [](Block *block) {
llvm::outs() << "Erasing "; llvm::outs() << "Erasing ";
printBlock(block); printBlock(block);
llvm::outs() << "\n"; llvm::outs() << "\n";
block->erase(); block->erase();
}; };
llvm::outs() << "Op post-order erasures (no skip)" llvm::outs() << "Op post-order erasures (no skip)"
<< "\n"; << "\n";
Operation *cloned = op->clone(); Operation *cloned = op->clone();
cloned->walk<WalkOrder::PostOrder>(noSkipOpErasure); cloned->walk<WalkOrder::PostOrder>(noSkipOpErasure);
llvm::outs() << "Op name post-order erasures (no skip)"
<< "\n";
cloned = op->clone();
cloned->walk<WalkOrder::PostOrder>(noSkipOpNameErasure);
llvm::outs() << "Block post-order erasures (no skip)" llvm::outs() << "Block post-order erasures (no skip)"
<< "\n"; << "\n";
cloned = op->clone(); cloned = op->clone();
cloned->walk<WalkOrder::PostOrder>(noSkipBlockErasure); cloned->walk<WalkOrder::PostOrder>(noSkipBlockErasure);
cloned->erase(); cloned->erase();
} }
namespace { namespace {
Show All 19 Lines