5cc0f76d34
The pattern rewriter documentation states that "*all* IR mutations [...] are required to be performed via the `PatternRewriter`." This commit adds two functions that were missing from the rewriter API: `moveOpBefore` and `moveOpAfter`. After an operation was moved, the `notifyOperationInserted` callback is triggered. This allows listeners such as the greedy pattern rewrite driver to react to IR changes. This commit narrows the discrepancy between the kind of IR modification that can be performed and the kind of IR modifications that can be listened to.
76 lines
2.6 KiB
C++
76 lines
2.6 KiB
C++
//===- TestConstantFold.cpp - Pass to test constant folding ---------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/FoldUtils.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
/// Simple constant folding pass.
|
|
struct TestConstantFold : public PassWrapper<TestConstantFold, OperationPass<>>,
|
|
public RewriterBase::Listener {
|
|
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestConstantFold)
|
|
|
|
StringRef getArgument() const final { return "test-constant-fold"; }
|
|
StringRef getDescription() const final {
|
|
return "Test operation constant folding";
|
|
}
|
|
// All constants in the operation post folding.
|
|
SmallVector<Operation *> existingConstants;
|
|
|
|
void foldOperation(Operation *op, OperationFolder &helper);
|
|
void runOnOperation() override;
|
|
|
|
void notifyOperationInserted(Operation *op,
|
|
OpBuilder::InsertPoint previous) override {
|
|
existingConstants.push_back(op);
|
|
}
|
|
void notifyOperationRemoved(Operation *op) override {
|
|
auto it = llvm::find(existingConstants, op);
|
|
if (it != existingConstants.end())
|
|
existingConstants.erase(it);
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void TestConstantFold::foldOperation(Operation *op, OperationFolder &helper) {
|
|
// Attempt to fold the specified operation, including handling unused or
|
|
// duplicated constants.
|
|
(void)helper.tryToFold(op);
|
|
}
|
|
|
|
void TestConstantFold::runOnOperation() {
|
|
existingConstants.clear();
|
|
|
|
// Collect and fold the operations within the operation.
|
|
SmallVector<Operation *, 8> ops;
|
|
getOperation()->walk<mlir::WalkOrder::PreOrder>([&](Operation *op) { ops.push_back(op); });
|
|
|
|
// Fold the constants in reverse so that the last generated constants from
|
|
// folding are at the beginning. This creates somewhat of a linear ordering to
|
|
// the newly generated constants that matches the operation order and improves
|
|
// the readability of test cases.
|
|
OperationFolder helper(&getContext(), /*listener=*/this);
|
|
for (Operation *op : llvm::reverse(ops))
|
|
foldOperation(op, helper);
|
|
|
|
// By the time we are done, we may have simplified a bunch of code, leaving
|
|
// around dead constants. Check for them now and remove them.
|
|
for (auto *cst : existingConstants) {
|
|
if (cst->use_empty())
|
|
cst->erase();
|
|
}
|
|
}
|
|
|
|
namespace mlir {
|
|
namespace test {
|
|
void registerTestConstantFold() { PassRegistration<TestConstantFold>(); }
|
|
} // namespace test
|
|
} // namespace mlir
|