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 #include "mlir/Transforms/Passes.h" 12 #include "mlir/Transforms/Utils.h" 13 14 using namespace mlir; 15 16 namespace { 17 /// Simple constant folding pass. 18 struct TestConstantFold : public PassWrapper<TestConstantFold, FunctionPass> { 19 // All constants in the function post folding. 20 SmallVector<Operation *, 8> existingConstants; 21 22 void foldOperation(Operation *op, OperationFolder &helper); 23 void runOnFunction() override; 24 }; 25 } // end anonymous namespace 26 27 void TestConstantFold::foldOperation(Operation *op, OperationFolder &helper) { 28 auto processGeneratedConstants = [this](Operation *op) { 29 existingConstants.push_back(op); 30 }; 31 32 // Attempt to fold the specified operation, including handling unused or 33 // duplicated constants. 34 (void)helper.tryToFold(op, processGeneratedConstants); 35 } 36 37 // For now, we do a simple top-down pass over a function folding constants. We 38 // don't handle conditional control flow, block arguments, folding conditional 39 // branches, or anything else fancy. 40 void TestConstantFold::runOnFunction() { 41 existingConstants.clear(); 42 43 // Collect and fold the operations within the function. 44 SmallVector<Operation *, 8> ops; 45 getFunction().walk([&](Operation *op) { ops.push_back(op); }); 46 47 // Fold the constants in reverse so that the last generated constants from 48 // folding are at the beginning. This creates somewhat of a linear ordering to 49 // the newly generated constants that matches the operation order and improves 50 // the readability of test cases. 51 OperationFolder helper(&getContext()); 52 for (Operation *op : llvm::reverse(ops)) 53 foldOperation(op, helper); 54 55 // By the time we are done, we may have simplified a bunch of code, leaving 56 // around dead constants. Check for them now and remove them. 57 for (auto *cst : existingConstants) { 58 if (cst->use_empty()) 59 cst->erase(); 60 } 61 } 62 63 namespace mlir { 64 namespace test { 65 void registerTestConstantFold() { 66 PassRegistration<TestConstantFold>("test-constant-fold", 67 "Test operation constant folding"); 68 } 69 } // namespace test 70 } // namespace mlir 71