1 //===- FunctionImplementation.cpp - Utilities for function-like ops -------===//
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/FunctionImplementation.h"
10 #include "mlir/IR/Builders.h"
11 #include "mlir/IR/FunctionSupport.h"
12 #include "mlir/IR/SymbolTable.h"
13
14 using namespace mlir;
15
16 static ParseResult
parseArgumentList(OpAsmParser & parser,bool allowVariadic,SmallVectorImpl<Type> & argTypes,SmallVectorImpl<OpAsmParser::OperandType> & argNames,SmallVectorImpl<SmallVector<NamedAttribute,2>> & argAttrs,bool & isVariadic)17 parseArgumentList(OpAsmParser &parser, bool allowVariadic,
18 SmallVectorImpl<Type> &argTypes,
19 SmallVectorImpl<OpAsmParser::OperandType> &argNames,
20 SmallVectorImpl<SmallVector<NamedAttribute, 2>> &argAttrs,
21 bool &isVariadic) {
22 if (parser.parseLParen())
23 return failure();
24
25 // The argument list either has to consistently have ssa-id's followed by
26 // types, or just be a type list. It isn't ok to sometimes have SSA ID's and
27 // sometimes not.
28 auto parseArgument = [&]() -> ParseResult {
29 llvm::SMLoc loc = parser.getCurrentLocation();
30
31 // Parse argument name if present.
32 OpAsmParser::OperandType argument;
33 Type argumentType;
34 if (succeeded(parser.parseOptionalRegionArgument(argument)) &&
35 !argument.name.empty()) {
36 // Reject this if the preceding argument was missing a name.
37 if (argNames.empty() && !argTypes.empty())
38 return parser.emitError(loc, "expected type instead of SSA identifier");
39 argNames.push_back(argument);
40
41 if (parser.parseColonType(argumentType))
42 return failure();
43 } else if (allowVariadic && succeeded(parser.parseOptionalEllipsis())) {
44 isVariadic = true;
45 return success();
46 } else if (!argNames.empty()) {
47 // Reject this if the preceding argument had a name.
48 return parser.emitError(loc, "expected SSA identifier");
49 } else if (parser.parseType(argumentType)) {
50 return failure();
51 }
52
53 // Add the argument type.
54 argTypes.push_back(argumentType);
55
56 // Parse any argument attributes.
57 SmallVector<NamedAttribute, 2> attrs;
58 if (parser.parseOptionalAttrDict(attrs))
59 return failure();
60 argAttrs.push_back(attrs);
61 return success();
62 };
63
64 // Parse the function arguments.
65 isVariadic = false;
66 if (failed(parser.parseOptionalRParen())) {
67 do {
68 unsigned numTypedArguments = argTypes.size();
69 if (parseArgument())
70 return failure();
71
72 llvm::SMLoc loc = parser.getCurrentLocation();
73 if (argTypes.size() == numTypedArguments &&
74 succeeded(parser.parseOptionalComma()))
75 return parser.emitError(
76 loc, "variadic arguments must be in the end of the argument list");
77 } while (succeeded(parser.parseOptionalComma()));
78 parser.parseRParen();
79 }
80
81 return success();
82 }
83
84 /// Parse a function result list.
85 ///
86 /// function-result-list ::= function-result-list-parens
87 /// | non-function-type
88 /// function-result-list-parens ::= `(` `)`
89 /// | `(` function-result-list-no-parens `)`
90 /// function-result-list-no-parens ::= function-result (`,` function-result)*
91 /// function-result ::= type attribute-dict?
92 ///
parseFunctionResultList(OpAsmParser & parser,SmallVectorImpl<Type> & resultTypes,SmallVectorImpl<SmallVector<NamedAttribute,2>> & resultAttrs)93 static ParseResult parseFunctionResultList(
94 OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes,
95 SmallVectorImpl<SmallVector<NamedAttribute, 2>> &resultAttrs) {
96 if (failed(parser.parseOptionalLParen())) {
97 // We already know that there is no `(`, so parse a type.
98 // Because there is no `(`, it cannot be a function type.
99 Type ty;
100 if (parser.parseType(ty))
101 return failure();
102 resultTypes.push_back(ty);
103 resultAttrs.emplace_back();
104 return success();
105 }
106
107 // Special case for an empty set of parens.
108 if (succeeded(parser.parseOptionalRParen()))
109 return success();
110
111 // Parse individual function results.
112 do {
113 resultTypes.emplace_back();
114 resultAttrs.emplace_back();
115 if (parser.parseType(resultTypes.back()) ||
116 parser.parseOptionalAttrDict(resultAttrs.back())) {
117 return failure();
118 }
119 } while (succeeded(parser.parseOptionalComma()));
120 return parser.parseRParen();
121 }
122
123 /// Parses a function signature using `parser`. The `allowVariadic` argument
124 /// indicates whether functions with variadic arguments are supported. The
125 /// trailing arguments are populated by this function with names, types and
126 /// attributes of the arguments and those of the results.
parseFunctionSignature(OpAsmParser & parser,bool allowVariadic,SmallVectorImpl<OpAsmParser::OperandType> & argNames,SmallVectorImpl<Type> & argTypes,SmallVectorImpl<SmallVector<NamedAttribute,2>> & argAttrs,bool & isVariadic,SmallVectorImpl<Type> & resultTypes,SmallVectorImpl<SmallVector<NamedAttribute,2>> & resultAttrs)127 ParseResult mlir::impl::parseFunctionSignature(
128 OpAsmParser &parser, bool allowVariadic,
129 SmallVectorImpl<OpAsmParser::OperandType> &argNames,
130 SmallVectorImpl<Type> &argTypes,
131 SmallVectorImpl<SmallVector<NamedAttribute, 2>> &argAttrs, bool &isVariadic,
132 SmallVectorImpl<Type> &resultTypes,
133 SmallVectorImpl<SmallVector<NamedAttribute, 2>> &resultAttrs) {
134 if (parseArgumentList(parser, allowVariadic, argTypes, argNames, argAttrs,
135 isVariadic))
136 return failure();
137 if (succeeded(parser.parseOptionalArrow()))
138 return parseFunctionResultList(parser, resultTypes, resultAttrs);
139 return success();
140 }
141
addArgAndResultAttrs(Builder & builder,OperationState & result,ArrayRef<SmallVector<NamedAttribute,2>> argAttrs,ArrayRef<SmallVector<NamedAttribute,2>> resultAttrs)142 void mlir::impl::addArgAndResultAttrs(
143 Builder &builder, OperationState &result,
144 ArrayRef<SmallVector<NamedAttribute, 2>> argAttrs,
145 ArrayRef<SmallVector<NamedAttribute, 2>> resultAttrs) {
146 // Add the attributes to the function arguments.
147 SmallString<8> attrNameBuf;
148 for (unsigned i = 0, e = argAttrs.size(); i != e; ++i)
149 if (!argAttrs[i].empty())
150 result.addAttribute(getArgAttrName(i, attrNameBuf),
151 builder.getDictionaryAttr(argAttrs[i]));
152
153 // Add the attributes to the function results.
154 for (unsigned i = 0, e = resultAttrs.size(); i != e; ++i)
155 if (!resultAttrs[i].empty())
156 result.addAttribute(getResultAttrName(i, attrNameBuf),
157 builder.getDictionaryAttr(resultAttrs[i]));
158 }
159
160 /// Parser implementation for function-like operations. Uses `funcTypeBuilder`
161 /// to construct the custom function type given lists of input and output types.
162 ParseResult
parseFunctionLikeOp(OpAsmParser & parser,OperationState & result,bool allowVariadic,mlir::impl::FuncTypeBuilder funcTypeBuilder)163 mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,
164 bool allowVariadic,
165 mlir::impl::FuncTypeBuilder funcTypeBuilder) {
166 SmallVector<OpAsmParser::OperandType, 4> entryArgs;
167 SmallVector<SmallVector<NamedAttribute, 2>, 4> argAttrs;
168 SmallVector<SmallVector<NamedAttribute, 2>, 4> resultAttrs;
169 SmallVector<Type, 4> argTypes;
170 SmallVector<Type, 4> resultTypes;
171 auto &builder = parser.getBuilder();
172
173 // Parse the name as a symbol.
174 StringAttr nameAttr;
175 if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
176 result.attributes))
177 return failure();
178
179 // Parse the function signature.
180 auto signatureLocation = parser.getCurrentLocation();
181 bool isVariadic = false;
182 if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
183 argAttrs, isVariadic, resultTypes, resultAttrs))
184 return failure();
185
186 std::string errorMessage;
187 if (auto type = funcTypeBuilder(builder, argTypes, resultTypes,
188 impl::VariadicFlag(isVariadic), errorMessage))
189 result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
190 else
191 return parser.emitError(signatureLocation)
192 << "failed to construct function type"
193 << (errorMessage.empty() ? "" : ": ") << errorMessage;
194
195 // If function attributes are present, parse them.
196 if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
197 return failure();
198
199 // Add the attributes to the function arguments.
200 assert(argAttrs.size() == argTypes.size());
201 assert(resultAttrs.size() == resultTypes.size());
202 addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);
203
204 // Parse the optional function body.
205 auto *body = result.addRegion();
206 return parser.parseOptionalRegion(
207 *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes);
208 }
209
210 // Print a function result list.
printFunctionResultList(OpAsmPrinter & p,ArrayRef<Type> types,ArrayRef<ArrayRef<NamedAttribute>> attrs)211 static void printFunctionResultList(OpAsmPrinter &p, ArrayRef<Type> types,
212 ArrayRef<ArrayRef<NamedAttribute>> attrs) {
213 assert(!types.empty() && "Should not be called for empty result list.");
214 auto &os = p.getStream();
215 bool needsParens =
216 types.size() > 1 || types[0].isa<FunctionType>() || !attrs[0].empty();
217 if (needsParens)
218 os << '(';
219 interleaveComma(llvm::zip(types, attrs), os,
220 [&](const std::tuple<Type, ArrayRef<NamedAttribute>> &t) {
221 p.printType(std::get<0>(t));
222 p.printOptionalAttrDict(std::get<1>(t));
223 });
224 if (needsParens)
225 os << ')';
226 }
227
228 /// Print the signature of the function-like operation `op`. Assumes `op` has
229 /// the FunctionLike trait and passed the verification.
printFunctionSignature(OpAsmPrinter & p,Operation * op,ArrayRef<Type> argTypes,bool isVariadic,ArrayRef<Type> resultTypes)230 void mlir::impl::printFunctionSignature(OpAsmPrinter &p, Operation *op,
231 ArrayRef<Type> argTypes,
232 bool isVariadic,
233 ArrayRef<Type> resultTypes) {
234 Region &body = op->getRegion(0);
235 bool isExternal = body.empty();
236
237 p << '(';
238 for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {
239 if (i > 0)
240 p << ", ";
241
242 if (!isExternal) {
243 p.printOperand(body.front().getArgument(i));
244 p << ": ";
245 }
246
247 p.printType(argTypes[i]);
248 p.printOptionalAttrDict(::mlir::impl::getArgAttrs(op, i));
249 }
250
251 if (isVariadic) {
252 if (!argTypes.empty())
253 p << ", ";
254 p << "...";
255 }
256
257 p << ')';
258
259 if (!resultTypes.empty()) {
260 p.getStream() << " -> ";
261 SmallVector<ArrayRef<NamedAttribute>, 4> resultAttrs;
262 for (int i = 0, e = resultTypes.size(); i < e; ++i)
263 resultAttrs.push_back(::mlir::impl::getResultAttrs(op, i));
264 printFunctionResultList(p, resultTypes, resultAttrs);
265 }
266 }
267
268 /// Prints the list of function prefixed with the "attributes" keyword. The
269 /// attributes with names listed in "elided" as well as those used by the
270 /// function-like operation internally are not printed. Nothing is printed
271 /// if all attributes are elided. Assumes `op` has the `FunctionLike` trait and
272 /// passed the verification.
printFunctionAttributes(OpAsmPrinter & p,Operation * op,unsigned numInputs,unsigned numResults,ArrayRef<StringRef> elided)273 void mlir::impl::printFunctionAttributes(OpAsmPrinter &p, Operation *op,
274 unsigned numInputs,
275 unsigned numResults,
276 ArrayRef<StringRef> elided) {
277 // Print out function attributes, if present.
278 SmallVector<StringRef, 2> ignoredAttrs = {
279 ::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName()};
280 ignoredAttrs.append(elided.begin(), elided.end());
281
282 SmallString<8> attrNameBuf;
283
284 // Ignore any argument attributes.
285 std::vector<SmallString<8>> argAttrStorage;
286 for (unsigned i = 0; i != numInputs; ++i)
287 if (op->getAttr(getArgAttrName(i, attrNameBuf)))
288 argAttrStorage.emplace_back(attrNameBuf);
289 ignoredAttrs.append(argAttrStorage.begin(), argAttrStorage.end());
290
291 // Ignore any result attributes.
292 std::vector<SmallString<8>> resultAttrStorage;
293 for (unsigned i = 0; i != numResults; ++i)
294 if (op->getAttr(getResultAttrName(i, attrNameBuf)))
295 resultAttrStorage.emplace_back(attrNameBuf);
296 ignoredAttrs.append(resultAttrStorage.begin(), resultAttrStorage.end());
297
298 p.printOptionalAttrDictWithKeyword(op->getAttrs(), ignoredAttrs);
299 }
300
301 /// Printer implementation for function-like operations. Accepts lists of
302 /// argument and result types to use while printing.
printFunctionLikeOp(OpAsmPrinter & p,Operation * op,ArrayRef<Type> argTypes,bool isVariadic,ArrayRef<Type> resultTypes)303 void mlir::impl::printFunctionLikeOp(OpAsmPrinter &p, Operation *op,
304 ArrayRef<Type> argTypes, bool isVariadic,
305 ArrayRef<Type> resultTypes) {
306 // Print the operation and the function name.
307 auto funcName =
308 op->getAttrOfType<StringAttr>(::mlir::SymbolTable::getSymbolAttrName())
309 .getValue();
310 p << op->getName() << ' ';
311 p.printSymbolName(funcName);
312
313 printFunctionSignature(p, op, argTypes, isVariadic, resultTypes);
314 printFunctionAttributes(p, op, argTypes.size(), resultTypes.size());
315
316 // Print the body if this is not an external function.
317 Region &body = op->getRegion(0);
318 if (!body.empty())
319 p.printRegion(body, /*printEntryBlockArgs=*/false,
320 /*printBlockTerminators=*/true);
321 }
322