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