94 lines · cpp
1//===- TestSingleFold.cpp - Pass to test single-pass 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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Pass/Pass.h"10#include "mlir/Transforms/FoldUtils.h"11 12using namespace mlir;13 14namespace {15/// Test pass for single-pass constant folding.16///17/// This pass tests the behavior of operations when folded exactly once. Unlike18/// canonicalization passes that may apply multiple rounds of folding, this pass19/// ensures that each operation is folded at most once, which is useful for20/// testing scenarios where the fold implementation should handle complex cases21/// without requiring multiple iterations.22///23/// The pass also removes dead constants after folding to clean up unused24/// intermediate results.25struct TestSingleFold : public PassWrapper<TestSingleFold, OperationPass<>>,26 public RewriterBase::Listener {27 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSingleFold)28 29 TestSingleFold() = default;30 TestSingleFold(const TestSingleFold &pass) : PassWrapper(pass) {}31 32 StringRef getArgument() const final { return "test-single-fold"; }33 StringRef getDescription() const final {34 return "Test single-pass operation folding and dead constant elimination";35 }36 // All constants in the operation post folding.37 SmallVector<Operation *> existingConstants;38 39 void foldOperation(Operation *op, OperationFolder &helper);40 void runOnOperation() override;41 42 void notifyOperationInserted(Operation *op,43 OpBuilder::InsertPoint previous) override {44 existingConstants.push_back(op);45 }46 void notifyOperationErased(Operation *op) override {47 auto *it = llvm::find(existingConstants, op);48 if (it != existingConstants.end())49 existingConstants.erase(it);50 }51 52 Option<int> maxIterations{*this, "max-iterations",53 llvm::cl::desc("Max iterations in the tryToFold"),54 llvm::cl::init(1)};55};56} // namespace57 58void TestSingleFold::foldOperation(Operation *op, OperationFolder &helper) {59 // Attempt to fold the specified operation, including handling unused or60 // duplicated constants.61 bool inPlaceUpdate = false;62 (void)helper.tryToFold(op, &inPlaceUpdate, maxIterations);63}64 65void TestSingleFold::runOnOperation() {66 existingConstants.clear();67 68 // Collect and fold the operations within the operation.69 SmallVector<Operation *, 8> ops;70 getOperation()->walk<mlir::WalkOrder::PreOrder>(71 [&](Operation *op) { ops.push_back(op); });72 73 // Fold the constants in reverse so that the last generated constants from74 // folding are at the beginning. This creates somewhat of a linear ordering to75 // the newly generated constants that matches the operation order and improves76 // the readability of test cases.77 OperationFolder helper(&getContext(), /*listener=*/this);78 for (Operation *op : llvm::reverse(ops))79 foldOperation(op, helper);80 81 // By the time we are done, we may have simplified a bunch of code, leaving82 // around dead constants. Check for them now and remove them.83 for (auto *cst : existingConstants) {84 if (cst->use_empty())85 cst->erase();86 }87}88 89namespace mlir {90namespace test {91void registerTestSingleFold() { PassRegistration<TestSingleFold>(); }92} // namespace test93} // namespace mlir94