|
| 1 | +//===-- AddContextArg.cpp - Add context argument ----------------*- C++ -*-===// |
| 2 | +// |
| 3 | +// This file is licensed 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 | +#include "mlir/Conversion/Passes.h" |
| 9 | +#include "mlir/Dialect/Func/IR/FuncOps.h" |
| 10 | + |
| 11 | +namespace mlir::gc { |
| 12 | +#define GEN_PASS_DECL_ADDCONTEXTARG |
| 13 | +#define GEN_PASS_DEF_ADDCONTEXTARG |
| 14 | +#include "gc/Transforms/Passes.h.inc" |
| 15 | +} // namespace mlir::gc |
| 16 | + |
| 17 | +using namespace mlir; |
| 18 | + |
| 19 | +namespace { |
| 20 | +struct AddContextArg final : gc::impl::AddContextArgBase<AddContextArg> { |
| 21 | + void runOnOperation() override { |
| 22 | + auto func = getOperation(); |
| 23 | + if (func.isExternal()) { |
| 24 | + return; |
| 25 | + } |
| 26 | + |
| 27 | + auto funcType = func.getFunctionType(); |
| 28 | + auto argTypes = llvm::to_vector<8>(funcType.getInputs()); |
| 29 | + auto resultTypes = llvm::to_vector<1>(funcType.getResults()); |
| 30 | + auto ctx = func->getContext(); |
| 31 | + auto newArgType = MemRefType::get({}, IntegerType::get(ctx, 8)); |
| 32 | + argTypes.emplace_back(newArgType); |
| 33 | + auto newFuncType = FunctionType::get(ctx, argTypes, resultTypes); |
| 34 | + func.setType(newFuncType); |
| 35 | + func.getBody().front().addArgument(newArgType, func.getLoc()); |
| 36 | + |
| 37 | + // Find all function calls and append the last argument of the current |
| 38 | + // function to the call. |
| 39 | + auto module = func->getParentOfType<ModuleOp>(); |
| 40 | + func.walk([&](func::CallOp call) { |
| 41 | + // If the function to be called is defined in the current module, then the |
| 42 | + // context arg will be added to this function signature either and, thus, |
| 43 | + // wee need add the context arg to the function call. |
| 44 | + if (auto callee = module.lookupSymbol<func::FuncOp>(call.getCallee()); |
| 45 | + !callee || callee.isExternal()) { |
| 46 | + return; |
| 47 | + } |
| 48 | + auto args = llvm::to_vector<8>(call.getOperands()); |
| 49 | + args.emplace_back(func.getArgument(func.getNumArguments() - 1)); |
| 50 | + call->setOperands(args); |
| 51 | + }); |
| 52 | + } |
| 53 | +}; |
| 54 | +} // namespace |
0 commit comments