| 1 | //===- TestConstantFold.cpp - Pass to test constant folding ---------------===// |
|---|
| 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-exception |
|---|
| 6 | // |
|---|
| 7 | //===----------------------------------------------------------------------===// |
|---|
| 8 | |
|---|
| 9 | #include "mlir/Pass/Pass.h" |
|---|
| 10 | #include "mlir/Transforms/FoldUtils.h" |
|---|
| 11 | |
|---|
| 12 | using namespace mlir; |
|---|
| 13 | |
|---|
| 14 | namespace { |
|---|
| 15 | /// Simple constant folding pass. |
|---|
| 16 | struct TestConstantFold : public PassWrapper<TestConstantFold, OperationPass<>>, |
|---|
| 17 | public RewriterBase::Listener { |
|---|
| 18 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestConstantFold) |
|---|
| 19 | |
|---|
| 20 | StringRef getArgument() const final { return "test-constant-fold" ; } |
|---|
| 21 | StringRef getDescription() const final { |
|---|
| 22 | return "Test operation constant folding" ; |
|---|
| 23 | } |
|---|
| 24 | // All constants in the operation post folding. |
|---|
| 25 | SmallVector<Operation *> existingConstants; |
|---|
| 26 | |
|---|
| 27 | void foldOperation(Operation *op, OperationFolder &helper); |
|---|
| 28 | void runOnOperation() override; |
|---|
| 29 | |
|---|
| 30 | void notifyOperationInserted(Operation *op, |
|---|
| 31 | OpBuilder::InsertPoint previous) override { |
|---|
| 32 | existingConstants.push_back(Elt: op); |
|---|
| 33 | } |
|---|
| 34 | void notifyOperationErased(Operation *op) override { |
|---|
| 35 | auto *it = llvm::find(Range&: existingConstants, Val: op); |
|---|
| 36 | if (it != existingConstants.end()) |
|---|
| 37 | existingConstants.erase(CI: it); |
|---|
| 38 | } |
|---|
| 39 | }; |
|---|
| 40 | } // namespace |
|---|
| 41 | |
|---|
| 42 | void TestConstantFold::foldOperation(Operation *op, OperationFolder &helper) { |
|---|
| 43 | // Attempt to fold the specified operation, including handling unused or |
|---|
| 44 | // duplicated constants. |
|---|
| 45 | (void)helper.tryToFold(op); |
|---|
| 46 | } |
|---|
| 47 | |
|---|
| 48 | void TestConstantFold::runOnOperation() { |
|---|
| 49 | existingConstants.clear(); |
|---|
| 50 | |
|---|
| 51 | // Collect and fold the operations within the operation. |
|---|
| 52 | SmallVector<Operation *, 8> ops; |
|---|
| 53 | getOperation()->walk<mlir::WalkOrder::PreOrder>(callback: [&](Operation *op) { ops.push_back(Elt: op); }); |
|---|
| 54 | |
|---|
| 55 | // Fold the constants in reverse so that the last generated constants from |
|---|
| 56 | // folding are at the beginning. This creates somewhat of a linear ordering to |
|---|
| 57 | // the newly generated constants that matches the operation order and improves |
|---|
| 58 | // the readability of test cases. |
|---|
| 59 | OperationFolder helper(&getContext(), /*listener=*/this); |
|---|
| 60 | for (Operation *op : llvm::reverse(C&: ops)) |
|---|
| 61 | foldOperation(op, helper); |
|---|
| 62 | |
|---|
| 63 | // By the time we are done, we may have simplified a bunch of code, leaving |
|---|
| 64 | // around dead constants. Check for them now and remove them. |
|---|
| 65 | for (auto *cst : existingConstants) { |
|---|
| 66 | if (cst->use_empty()) |
|---|
| 67 | cst->erase(); |
|---|
| 68 | } |
|---|
| 69 | } |
|---|
| 70 | |
|---|
| 71 | namespace mlir { |
|---|
| 72 | namespace test { |
|---|
| 73 | void registerTestConstantFold() { PassRegistration<TestConstantFold>(); } |
|---|
| 74 | } // namespace test |
|---|
| 75 | } // namespace mlir |
|---|
| 76 | |
|---|
