| //===- Function.cpp - MLIR Function Classes -------------------------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/IR/Function.h" |
| #include "mlir/IR/BlockAndValueMapping.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/FunctionImplementation.h" |
| #include "llvm/ADT/BitVector.h" |
| #include "llvm/ADT/MapVector.h" |
| #include "llvm/ADT/SmallString.h" |
| #include "llvm/ADT/Twine.h" |
| |
| using namespace mlir; |
| |
| //===----------------------------------------------------------------------===// |
| // Function Operation. |
| //===----------------------------------------------------------------------===// |
| |
| FuncOp FuncOp::create(Location location, StringRef name, FunctionType type, |
| ArrayRef<NamedAttribute> attrs) { |
| OperationState state(location, "func"); |
| OpBuilder builder(location->getContext()); |
| FuncOp::build(builder, state, name, type, attrs); |
| return cast<FuncOp>(Operation::create(state)); |
| } |
| FuncOp FuncOp::create(Location location, StringRef name, FunctionType type, |
| iterator_range<dialect_attr_iterator> attrs) { |
| SmallVector<NamedAttribute, 8> attrRef(attrs); |
| return create(location, name, type, llvm::makeArrayRef(attrRef)); |
| } |
| FuncOp FuncOp::create(Location location, StringRef name, FunctionType type, |
| ArrayRef<NamedAttribute> attrs, |
| ArrayRef<MutableDictionaryAttr> argAttrs) { |
| FuncOp func = create(location, name, type, attrs); |
| func.setAllArgAttrs(argAttrs); |
| return func; |
| } |
| |
| void FuncOp::build(OpBuilder &builder, OperationState &result, StringRef name, |
| FunctionType type, ArrayRef<NamedAttribute> attrs) { |
| result.addAttribute(SymbolTable::getSymbolAttrName(), |
| builder.getStringAttr(name)); |
| result.addAttribute(getTypeAttrName(), TypeAttr::get(type)); |
| result.attributes.append(attrs.begin(), attrs.end()); |
| result.addRegion(); |
| } |
| |
| void FuncOp::build(OpBuilder &builder, OperationState &result, StringRef name, |
| FunctionType type, ArrayRef<NamedAttribute> attrs, |
| ArrayRef<MutableDictionaryAttr> argAttrs) { |
| build(builder, result, name, type, attrs); |
| assert(type.getNumInputs() == argAttrs.size()); |
| SmallString<8> argAttrName; |
| for (unsigned i = 0, e = type.getNumInputs(); i != e; ++i) |
| if (auto argDict = argAttrs[i].getDictionary(builder.getContext())) |
| result.addAttribute(getArgAttrName(i, argAttrName), argDict); |
| } |
| |
| /// Parsing/Printing methods. |
| |
| ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) { |
| auto buildFuncType = [](Builder &builder, ArrayRef<Type> argTypes, |
| ArrayRef<Type> results, impl::VariadicFlag, |
| std::string &) { |
| return builder.getFunctionType(argTypes, results); |
| }; |
| |
| return impl::parseFunctionLikeOp(parser, result, /*allowVariadic=*/false, |
| buildFuncType); |
| } |
| |
| void FuncOp::print(OpAsmPrinter &p) { |
| FunctionType fnType = getType(); |
| impl::printFunctionLikeOp(p, *this, fnType.getInputs(), /*isVariadic=*/false, |
| fnType.getResults()); |
| } |
| |
| LogicalResult FuncOp::verify() { |
| // If this function is external there is nothing to do. |
| if (isExternal()) |
| return success(); |
| |
| // Verify that the argument list of the function and the arg list of the entry |
| // block line up. The trait already verified that the number of arguments is |
| // the same between the signature and the block. |
| auto fnInputTypes = getType().getInputs(); |
| Block &entryBlock = front(); |
| for (unsigned i = 0, e = entryBlock.getNumArguments(); i != e; ++i) |
| if (fnInputTypes[i] != entryBlock.getArgument(i).getType()) |
| return emitOpError("type of entry block argument #") |
| << i << '(' << entryBlock.getArgument(i).getType() |
| << ") must match the type of the corresponding argument in " |
| << "function signature(" << fnInputTypes[i] << ')'; |
| |
| return success(); |
| } |
| |
| void FuncOp::eraseArguments(ArrayRef<unsigned> argIndices) { |
| auto oldType = getType(); |
| int originalNumArgs = oldType.getNumInputs(); |
| llvm::BitVector eraseIndices(originalNumArgs); |
| for (auto index : argIndices) |
| eraseIndices.set(index); |
| auto shouldEraseArg = [&](int i) { return eraseIndices.test(i); }; |
| |
| // There are 3 things that need to be updated: |
| // - Function type. |
| // - Arg attrs. |
| // - Block arguments of entry block. |
| |
| // Update the function type and arg attrs. |
| SmallVector<Type, 4> newInputTypes; |
| SmallVector<MutableDictionaryAttr, 4> newArgAttrs; |
| for (int i = 0; i < originalNumArgs; i++) { |
| if (shouldEraseArg(i)) |
| continue; |
| newInputTypes.emplace_back(oldType.getInput(i)); |
| newArgAttrs.emplace_back(getArgAttrDict(i)); |
| } |
| setType(FunctionType::get(newInputTypes, oldType.getResults(), getContext())); |
| setAllArgAttrs(newArgAttrs); |
| |
| // Update the entry block's arguments. |
| // We do this in reverse so that we erase later indices before earlier |
| // indices, to avoid shifting the later indices. |
| Block &entry = front(); |
| for (int i = 0; i < originalNumArgs; i++) |
| if (shouldEraseArg(originalNumArgs - i - 1)) |
| entry.eraseArgument(originalNumArgs - i - 1); |
| } |
| |
| /// Clone the internal blocks from this function into dest and all attributes |
| /// from this function to dest. |
| void FuncOp::cloneInto(FuncOp dest, BlockAndValueMapping &mapper) { |
| // Add the attributes of this function to dest. |
| llvm::MapVector<Identifier, Attribute> newAttrs; |
| for (auto &attr : dest.getAttrs()) |
| newAttrs.insert(attr); |
| for (auto &attr : getAttrs()) |
| newAttrs.insert(attr); |
| dest.getOperation()->setAttrs( |
| DictionaryAttr::get(newAttrs.takeVector(), getContext())); |
| |
| // Clone the body. |
| getBody().cloneInto(&dest.getBody(), mapper); |
| } |
| |
| /// Create a deep copy of this function and all of its blocks, remapping |
| /// any operands that use values outside of the function using the map that is |
| /// provided (leaving them alone if no entry is present). Replaces references |
| /// to cloned sub-values with the corresponding value that is copied, and adds |
| /// those mappings to the mapper. |
| FuncOp FuncOp::clone(BlockAndValueMapping &mapper) { |
| FunctionType newType = getType(); |
| |
| // If the function has a body, then the user might be deleting arguments to |
| // the function by specifying them in the mapper. If so, we don't add the |
| // argument to the input type vector. |
| bool isExternalFn = isExternal(); |
| if (!isExternalFn) { |
| SmallVector<Type, 4> inputTypes; |
| inputTypes.reserve(newType.getNumInputs()); |
| for (unsigned i = 0, e = getNumArguments(); i != e; ++i) |
| if (!mapper.contains(getArgument(i))) |
| inputTypes.push_back(newType.getInput(i)); |
| newType = FunctionType::get(inputTypes, newType.getResults(), getContext()); |
| } |
| |
| // Create the new function. |
| FuncOp newFunc = cast<FuncOp>(getOperation()->cloneWithoutRegions()); |
| newFunc.setType(newType); |
| |
| /// Set the argument attributes for arguments that aren't being replaced. |
| for (unsigned i = 0, e = getNumArguments(), destI = 0; i != e; ++i) |
| if (isExternalFn || !mapper.contains(getArgument(i))) |
| newFunc.setArgAttrs(destI++, getArgAttrs(i)); |
| |
| /// Clone the current function into the new one and return it. |
| cloneInto(newFunc, mapper); |
| return newFunc; |
| } |
| FuncOp FuncOp::clone() { |
| BlockAndValueMapping mapper; |
| return clone(mapper); |
| } |