1 //===- TestFunctionLike.cpp - Pass to test helpers on FunctionLike --------===// 2 // 3 // Part of the MLIR 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/Function.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 : public ModulePass<TestFuncEraseArg> { 17 void runOnModule() override { 18 auto module = getModule(); 19 20 for (FuncOp func : module.getOps<FuncOp>()) { 21 SmallVector<unsigned, 4> indicesToErase; 22 for (auto argIndex : llvm::seq<int>(0, func.getNumArguments())) { 23 if (func.getArgAttr(argIndex, "test.erase_this_arg")) { 24 // Push back twice to test that duplicate arg indices are handled 25 // correctly. 26 indicesToErase.push_back(argIndex); 27 indicesToErase.push_back(argIndex); 28 } 29 } 30 // Reverse the order to test that unsorted index lists are handled 31 // correctly. 32 std::reverse(indicesToErase.begin(), indicesToErase.end()); 33 func.eraseArguments(indicesToErase); 34 } 35 } 36 }; 37 38 /// This is a test pass for verifying FuncOp's setType method. 39 struct TestFuncSetType : public ModulePass<TestFuncSetType> { 40 void runOnModule() override { 41 auto module = getModule(); 42 SymbolTable symbolTable(module); 43 44 for (FuncOp func : module.getOps<FuncOp>()) { 45 auto sym = func.getAttrOfType<FlatSymbolRefAttr>("test.set_type_from"); 46 if (!sym) 47 continue; 48 func.setType(symbolTable.lookup<FuncOp>(sym.getValue()).getType()); 49 } 50 } 51 }; 52 } // end anonymous namespace 53 54 static PassRegistration<TestFuncEraseArg> pass("test-func-erase-arg", 55 "Test erasing func args."); 56 57 static PassRegistration<TestFuncSetType> pass2("test-func-set-type", 58 "Test FuncOp::setType."); 59