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