1 //===- TestFunctionLike.cpp - Pass to test helpers on FunctionLike --------===//
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/IR/BuiltinOps.h"
10 #include "mlir/Pass/Pass.h"
11
12 using namespace mlir;
13
14 namespace {
15 /// This is a test pass for verifying FuncOp's eraseArgument method.
16 struct TestFuncEraseArg
17 : public PassWrapper<TestFuncEraseArg, OperationPass<ModuleOp>> {
runOnOperation__anonfb1dd6050111::TestFuncEraseArg18 void runOnOperation() override {
19 auto module = getOperation();
20
21 for (FuncOp func : module.getOps<FuncOp>()) {
22 SmallVector<unsigned, 4> indicesToErase;
23 for (auto argIndex : llvm::seq<int>(0, func.getNumArguments())) {
24 if (func.getArgAttr(argIndex, "test.erase_this_arg")) {
25 // Push back twice to test that duplicate arg indices are handled
26 // correctly.
27 indicesToErase.push_back(argIndex);
28 indicesToErase.push_back(argIndex);
29 }
30 }
31 // Reverse the order to test that unsorted index lists are handled
32 // correctly.
33 std::reverse(indicesToErase.begin(), indicesToErase.end());
34 func.eraseArguments(indicesToErase);
35 }
36 }
37 };
38
39 /// This is a test pass for verifying FuncOp's eraseResult method.
40 struct TestFuncEraseResult
41 : public PassWrapper<TestFuncEraseResult, OperationPass<ModuleOp>> {
runOnOperation__anonfb1dd6050111::TestFuncEraseResult42 void runOnOperation() override {
43 auto module = getOperation();
44
45 for (FuncOp func : module.getOps<FuncOp>()) {
46 SmallVector<unsigned, 4> indicesToErase;
47 for (auto resultIndex : llvm::seq<int>(0, func.getNumResults())) {
48 if (func.getResultAttr(resultIndex, "test.erase_this_result")) {
49 // Push back twice to test that duplicate indices are handled
50 // correctly.
51 indicesToErase.push_back(resultIndex);
52 indicesToErase.push_back(resultIndex);
53 }
54 }
55 // Reverse the order to test that unsorted index lists are handled
56 // correctly.
57 std::reverse(indicesToErase.begin(), indicesToErase.end());
58 func.eraseResults(indicesToErase);
59 }
60 }
61 };
62
63 /// This is a test pass for verifying FuncOp's setType method.
64 struct TestFuncSetType
65 : public PassWrapper<TestFuncSetType, OperationPass<ModuleOp>> {
runOnOperation__anonfb1dd6050111::TestFuncSetType66 void runOnOperation() override {
67 auto module = getOperation();
68 SymbolTable symbolTable(module);
69
70 for (FuncOp func : module.getOps<FuncOp>()) {
71 auto sym = func->getAttrOfType<FlatSymbolRefAttr>("test.set_type_from");
72 if (!sym)
73 continue;
74 func.setType(symbolTable.lookup<FuncOp>(sym.getValue()).getType());
75 }
76 }
77 };
78 } // end anonymous namespace
79
80 namespace mlir {
registerTestFunc()81 void registerTestFunc() {
82 PassRegistration<TestFuncEraseArg>("test-func-erase-arg",
83 "Test erasing func args.");
84
85 PassRegistration<TestFuncEraseResult>("test-func-erase-result",
86 "Test erasing func results.");
87
88 PassRegistration<TestFuncSetType>("test-func-set-type",
89 "Test FuncOp::setType.");
90 }
91 } // namespace mlir
92