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>> {
18   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>> {
42   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>> {
66   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 {
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