1 //===- FunctionImplementation.cpp - Utilities for function-like ops -------===//
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/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
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 ///
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.
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 
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
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.
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.
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.
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.
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