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/FunctionInterfaces.h"
12 #include "mlir/IR/SymbolTable.h"
13 
14 using namespace mlir;
15 
16 ParseResult mlir::function_interface_impl::parseFunctionArgumentList(
17     OpAsmParser &parser, bool allowAttributes, bool allowVariadic,
18     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
19     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
20     SmallVectorImpl<Optional<Location>> &argLocations, 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     if (!allowAttributes && !attrs.empty())
60       return parser.emitError(loc, "expected arguments without attributes");
61     argAttrs.push_back(attrs);
62 
63     // Parse a location if specified.
64     Optional<Location> explicitLoc;
65     if (!argument.name.empty() &&
66         parser.parseOptionalLocationSpecifier(explicitLoc))
67       return failure();
68     argLocations.push_back(explicitLoc);
69 
70     return success();
71   };
72 
73   // Parse the function arguments.
74   isVariadic = false;
75   if (failed(parser.parseOptionalRParen())) {
76     do {
77       unsigned numTypedArguments = argTypes.size();
78       if (parseArgument())
79         return failure();
80 
81       llvm::SMLoc loc = parser.getCurrentLocation();
82       if (argTypes.size() == numTypedArguments &&
83           succeeded(parser.parseOptionalComma()))
84         return parser.emitError(
85             loc, "variadic arguments must be in the end of the argument list");
86     } while (succeeded(parser.parseOptionalComma()));
87     parser.parseRParen();
88   }
89 
90   return success();
91 }
92 
93 /// Parse a function result list.
94 ///
95 ///   function-result-list ::= function-result-list-parens
96 ///                          | non-function-type
97 ///   function-result-list-parens ::= `(` `)`
98 ///                                 | `(` function-result-list-no-parens `)`
99 ///   function-result-list-no-parens ::= function-result (`,` function-result)*
100 ///   function-result ::= type attribute-dict?
101 ///
102 static ParseResult
103 parseFunctionResultList(OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes,
104                         SmallVectorImpl<NamedAttrList> &resultAttrs) {
105   if (failed(parser.parseOptionalLParen())) {
106     // We already know that there is no `(`, so parse a type.
107     // Because there is no `(`, it cannot be a function type.
108     Type ty;
109     if (parser.parseType(ty))
110       return failure();
111     resultTypes.push_back(ty);
112     resultAttrs.emplace_back();
113     return success();
114   }
115 
116   // Special case for an empty set of parens.
117   if (succeeded(parser.parseOptionalRParen()))
118     return success();
119 
120   // Parse individual function results.
121   do {
122     resultTypes.emplace_back();
123     resultAttrs.emplace_back();
124     if (parser.parseType(resultTypes.back()) ||
125         parser.parseOptionalAttrDict(resultAttrs.back())) {
126       return failure();
127     }
128   } while (succeeded(parser.parseOptionalComma()));
129   return parser.parseRParen();
130 }
131 
132 ParseResult mlir::function_interface_impl::parseFunctionSignature(
133     OpAsmParser &parser, bool allowVariadic,
134     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
135     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
136     SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic,
137     SmallVectorImpl<Type> &resultTypes,
138     SmallVectorImpl<NamedAttrList> &resultAttrs) {
139   bool allowArgAttrs = true;
140   if (parseFunctionArgumentList(parser, allowArgAttrs, allowVariadic, argNames,
141                                 argTypes, argAttrs, argLocations, isVariadic))
142     return failure();
143   if (succeeded(parser.parseOptionalArrow()))
144     return parseFunctionResultList(parser, resultTypes, resultAttrs);
145   return success();
146 }
147 
148 /// Implementation of `addArgAndResultAttrs` that is attribute list type
149 /// agnostic.
150 template <typename AttrListT, typename AttrArrayBuildFnT>
151 static void addArgAndResultAttrsImpl(Builder &builder, OperationState &result,
152                                      ArrayRef<AttrListT> argAttrs,
153                                      ArrayRef<AttrListT> resultAttrs,
154                                      AttrArrayBuildFnT &&buildAttrArrayFn) {
155   auto nonEmptyAttrsFn = [](const AttrListT &attrs) { return !attrs.empty(); };
156 
157   // Add the attributes to the function arguments.
158   if (!argAttrs.empty() && llvm::any_of(argAttrs, nonEmptyAttrsFn)) {
159     ArrayAttr attrDicts = builder.getArrayAttr(buildAttrArrayFn(argAttrs));
160     result.addAttribute(function_interface_impl::getArgDictAttrName(),
161                         attrDicts);
162   }
163   // Add the attributes to the function results.
164   if (!resultAttrs.empty() && llvm::any_of(resultAttrs, nonEmptyAttrsFn)) {
165     ArrayAttr attrDicts = builder.getArrayAttr(buildAttrArrayFn(resultAttrs));
166     result.addAttribute(function_interface_impl::getResultDictAttrName(),
167                         attrDicts);
168   }
169 }
170 
171 void mlir::function_interface_impl::addArgAndResultAttrs(
172     Builder &builder, OperationState &result, ArrayRef<DictionaryAttr> argAttrs,
173     ArrayRef<DictionaryAttr> resultAttrs) {
174   auto buildFn = [](ArrayRef<DictionaryAttr> attrs) {
175     return ArrayRef<Attribute>(attrs.data(), attrs.size());
176   };
177   addArgAndResultAttrsImpl(builder, result, argAttrs, resultAttrs, buildFn);
178 }
179 void mlir::function_interface_impl::addArgAndResultAttrs(
180     Builder &builder, OperationState &result, ArrayRef<NamedAttrList> argAttrs,
181     ArrayRef<NamedAttrList> resultAttrs) {
182   MLIRContext *context = builder.getContext();
183   auto buildFn = [=](ArrayRef<NamedAttrList> attrs) {
184     return llvm::to_vector<8>(
185         llvm::map_range(attrs, [=](const NamedAttrList &attrList) -> Attribute {
186           return attrList.getDictionary(context);
187         }));
188   };
189   addArgAndResultAttrsImpl(builder, result, argAttrs, resultAttrs, buildFn);
190 }
191 
192 ParseResult mlir::function_interface_impl::parseFunctionOp(
193     OpAsmParser &parser, OperationState &result, bool allowVariadic,
194     FuncTypeBuilder funcTypeBuilder) {
195   SmallVector<OpAsmParser::OperandType> entryArgs;
196   SmallVector<NamedAttrList> argAttrs;
197   SmallVector<NamedAttrList> resultAttrs;
198   SmallVector<Type> argTypes;
199   SmallVector<Type> resultTypes;
200   SmallVector<Optional<Location>> argLocations;
201   auto &builder = parser.getBuilder();
202 
203   // Parse visibility.
204   impl::parseOptionalVisibilityKeyword(parser, result.attributes);
205 
206   // Parse the name as a symbol.
207   StringAttr nameAttr;
208   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
209                              result.attributes))
210     return failure();
211 
212   // Parse the function signature.
213   llvm::SMLoc signatureLocation = parser.getCurrentLocation();
214   bool isVariadic = false;
215   if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
216                              argAttrs, argLocations, isVariadic, resultTypes,
217                              resultAttrs))
218     return failure();
219 
220   std::string errorMessage;
221   Type type = funcTypeBuilder(builder, argTypes, resultTypes,
222                               VariadicFlag(isVariadic), errorMessage);
223   if (!type) {
224     return parser.emitError(signatureLocation)
225            << "failed to construct function type"
226            << (errorMessage.empty() ? "" : ": ") << errorMessage;
227   }
228   result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
229 
230   // If function attributes are present, parse them.
231   NamedAttrList parsedAttributes;
232   llvm::SMLoc attributeDictLocation = parser.getCurrentLocation();
233   if (parser.parseOptionalAttrDictWithKeyword(parsedAttributes))
234     return failure();
235 
236   // Disallow attributes that are inferred from elsewhere in the attribute
237   // dictionary.
238   for (StringRef disallowed :
239        {SymbolTable::getVisibilityAttrName(), SymbolTable::getSymbolAttrName(),
240         getTypeAttrName()}) {
241     if (parsedAttributes.get(disallowed))
242       return parser.emitError(attributeDictLocation, "'")
243              << disallowed
244              << "' is an inferred attribute and should not be specified in the "
245                 "explicit attribute dictionary";
246   }
247   result.attributes.append(parsedAttributes);
248 
249   // Add the attributes to the function arguments.
250   assert(argAttrs.size() == argTypes.size());
251   assert(resultAttrs.size() == resultTypes.size());
252   addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);
253 
254   // Parse the optional function body. The printer will not print the body if
255   // its empty, so disallow parsing of empty body in the parser.
256   auto *body = result.addRegion();
257   llvm::SMLoc loc = parser.getCurrentLocation();
258   OptionalParseResult parseResult = parser.parseOptionalRegion(
259       *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes,
260       entryArgs.empty() ? ArrayRef<Optional<Location>>() : argLocations,
261       /*enableNameShadowing=*/false);
262   if (parseResult.hasValue()) {
263     if (failed(*parseResult))
264       return failure();
265     // Function body was parsed, make sure its not empty.
266     if (body->empty())
267       return parser.emitError(loc, "expected non-empty function body");
268   }
269   return success();
270 }
271 
272 /// Print a function result list. The provided `attrs` must either be null, or
273 /// contain a set of DictionaryAttrs of the same arity as `types`.
274 static void printFunctionResultList(OpAsmPrinter &p, ArrayRef<Type> types,
275                                     ArrayAttr attrs) {
276   assert(!types.empty() && "Should not be called for empty result list.");
277   assert((!attrs || attrs.size() == types.size()) &&
278          "Invalid number of attributes.");
279 
280   auto &os = p.getStream();
281   bool needsParens = types.size() > 1 || types[0].isa<FunctionType>() ||
282                      (attrs && !attrs[0].cast<DictionaryAttr>().empty());
283   if (needsParens)
284     os << '(';
285   llvm::interleaveComma(llvm::seq<size_t>(0, types.size()), os, [&](size_t i) {
286     p.printType(types[i]);
287     if (attrs)
288       p.printOptionalAttrDict(attrs[i].cast<DictionaryAttr>().getValue());
289   });
290   if (needsParens)
291     os << ')';
292 }
293 
294 void mlir::function_interface_impl::printFunctionSignature(
295     OpAsmPrinter &p, Operation *op, ArrayRef<Type> argTypes, bool isVariadic,
296     ArrayRef<Type> resultTypes) {
297   Region &body = op->getRegion(0);
298   bool isExternal = body.empty();
299 
300   p << '(';
301   ArrayAttr argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
302   for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {
303     if (i > 0)
304       p << ", ";
305 
306     if (!isExternal) {
307       ArrayRef<NamedAttribute> attrs;
308       if (argAttrs)
309         attrs = argAttrs[i].cast<DictionaryAttr>().getValue();
310       p.printRegionArgument(body.getArgument(i), attrs);
311     } else {
312       p.printType(argTypes[i]);
313       if (argAttrs)
314         p.printOptionalAttrDict(argAttrs[i].cast<DictionaryAttr>().getValue());
315     }
316   }
317 
318   if (isVariadic) {
319     if (!argTypes.empty())
320       p << ", ";
321     p << "...";
322   }
323 
324   p << ')';
325 
326   if (!resultTypes.empty()) {
327     p.getStream() << " -> ";
328     auto resultAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
329     printFunctionResultList(p, resultTypes, resultAttrs);
330   }
331 }
332 
333 void mlir::function_interface_impl::printFunctionAttributes(
334     OpAsmPrinter &p, Operation *op, unsigned numInputs, unsigned numResults,
335     ArrayRef<StringRef> elided) {
336   // Print out function attributes, if present.
337   SmallVector<StringRef, 2> ignoredAttrs = {
338       ::mlir::SymbolTable::getSymbolAttrName(), getTypeAttrName(),
339       getArgDictAttrName(), getResultDictAttrName()};
340   ignoredAttrs.append(elided.begin(), elided.end());
341 
342   p.printOptionalAttrDictWithKeyword(op->getAttrs(), ignoredAttrs);
343 }
344 
345 void mlir::function_interface_impl::printFunctionOp(
346     OpAsmPrinter &p, Operation *op, ArrayRef<Type> argTypes, bool isVariadic,
347     ArrayRef<Type> resultTypes) {
348   // Print the operation and the function name.
349   auto funcName =
350       op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
351           .getValue();
352   p << ' ';
353 
354   StringRef visibilityAttrName = SymbolTable::getVisibilityAttrName();
355   if (auto visibility = op->getAttrOfType<StringAttr>(visibilityAttrName))
356     p << visibility.getValue() << ' ';
357   p.printSymbolName(funcName);
358 
359   printFunctionSignature(p, op, argTypes, isVariadic, resultTypes);
360   printFunctionAttributes(p, op, argTypes.size(), resultTypes.size(),
361                           {visibilityAttrName});
362   // Print the body if this is not an external function.
363   Region &body = op->getRegion(0);
364   if (!body.empty()) {
365     p << ' ';
366     p.printRegion(body, /*printEntryBlockArgs=*/false,
367                   /*printBlockTerminators=*/true);
368   }
369 }
370