1 //===- FunctionSupport.h - Utility types for function-like ops --*- C++ -*-===//
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 // This file defines support types for Operations that represent function-like
10 // constructs to use.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_IR_FUNCTIONSUPPORT_H
15 #define MLIR_IR_FUNCTIONSUPPORT_H
16 
17 #include "mlir/IR/BuiltinTypes.h"
18 #include "mlir/IR/OpDefinition.h"
19 #include "llvm/ADT/SmallString.h"
20 
21 namespace mlir {
22 
23 namespace impl {
24 
25 /// Return the name of the attribute used for function types.
getTypeAttrName()26 inline StringRef getTypeAttrName() { return "type"; }
27 
28 /// Return the name of the attribute used for function arguments.
getArgAttrName(unsigned arg,SmallVectorImpl<char> & out)29 inline StringRef getArgAttrName(unsigned arg, SmallVectorImpl<char> &out) {
30   out.clear();
31   return ("arg" + Twine(arg)).toStringRef(out);
32 }
33 
34 /// Returns true if the given name is a valid argument attribute name.
isArgAttrName(StringRef name)35 inline bool isArgAttrName(StringRef name) {
36   APInt unused;
37   return name.startswith("arg") &&
38          !name.drop_front(3).getAsInteger(/*Radix=*/10, unused);
39 }
40 
41 /// Return the name of the attribute used for function results.
getResultAttrName(unsigned arg,SmallVectorImpl<char> & out)42 inline StringRef getResultAttrName(unsigned arg, SmallVectorImpl<char> &out) {
43   out.clear();
44   return ("result" + Twine(arg)).toStringRef(out);
45 }
46 
47 /// Returns the dictionary attribute corresponding to the argument at 'index'.
48 /// If there are no argument attributes at 'index', a null attribute is
49 /// returned.
getArgAttrDict(Operation * op,unsigned index)50 inline DictionaryAttr getArgAttrDict(Operation *op, unsigned index) {
51   SmallString<8> nameOut;
52   return op->getAttrOfType<DictionaryAttr>(getArgAttrName(index, nameOut));
53 }
54 
55 /// Returns the dictionary attribute corresponding to the result at 'index'.
56 /// If there are no result attributes at 'index', a null attribute is
57 /// returned.
getResultAttrDict(Operation * op,unsigned index)58 inline DictionaryAttr getResultAttrDict(Operation *op, unsigned index) {
59   SmallString<8> nameOut;
60   return op->getAttrOfType<DictionaryAttr>(getResultAttrName(index, nameOut));
61 }
62 
63 /// Return all of the attributes for the argument at 'index'.
getArgAttrs(Operation * op,unsigned index)64 inline ArrayRef<NamedAttribute> getArgAttrs(Operation *op, unsigned index) {
65   auto argDict = getArgAttrDict(op, index);
66   return argDict ? argDict.getValue() : llvm::None;
67 }
68 
69 /// Return all of the attributes for the result at 'index'.
getResultAttrs(Operation * op,unsigned index)70 inline ArrayRef<NamedAttribute> getResultAttrs(Operation *op, unsigned index) {
71   auto resultDict = getResultAttrDict(op, index);
72   return resultDict ? resultDict.getValue() : llvm::None;
73 }
74 
75 /// Erase the specified arguments and update the function type attribute.
76 void eraseFunctionArguments(Operation *op, ArrayRef<unsigned> argIndices,
77                             unsigned originalNumArgs, Type newType);
78 
79 /// Erase the specified results and update the function type attribute.
80 void eraseFunctionResults(Operation *op, ArrayRef<unsigned> resultIndices,
81                           unsigned originalNumResults, Type newType);
82 
83 } // namespace impl
84 
85 namespace OpTrait {
86 
87 /// This trait provides APIs for Ops that behave like functions.  In particular:
88 /// - Ops must be symbols, i.e. also have the `Symbol` trait;
89 /// - Ops have a single region with multiple blocks that corresponds to the body
90 ///   of the function;
91 /// - An op with a single empty region corresponds to an external function;
92 /// - leading arguments of the first block of the region are treated as function
93 ///   arguments;
94 /// - they can have argument attributes that are stored in a dictionary
95 ///   attribute on the Op itself.
96 ///
97 /// This trait provides limited type support for the declared or defined
98 /// functions. The convenience function `getTypeAttrName()` returns the name of
99 /// an attribute that can be used to store the function type. In addition, this
100 /// trait provides `getType` and `setType` helpers to store a `FunctionType` in
101 /// the attribute named by `getTypeAttrName()`.
102 ///
103 /// In general, this trait assumes concrete ops use `FunctionType` under the
104 /// hood. If this is not the case, in order to use the function type support,
105 /// concrete ops must define the following methods, using the same name, to hide
106 /// the ones defined for `FunctionType`: `addBodyBlock`, `getType`,
107 /// `getTypeWithoutArgsAndResults` and `setType`.
108 ///
109 /// Besides the requirements above, concrete ops must interact with this trait
110 /// using the following functions:
111 /// - Concrete ops *must* define a member function `getNumFuncArguments()` that
112 ///   returns the number of function arguments based exclusively on type (so
113 ///   that it can be called on function declarations).
114 /// - Concrete ops *must* define a member function `getNumFuncResults()` that
115 ///   returns the number of function results based exclusively on type (so that
116 ///   it can be called on function declarations).
117 /// - To verify that the type respects op-specific invariants, concrete ops may
118 ///   redefine the `verifyType()` hook that will be called after verifying the
119 ///   presence of the `type` attribute and before any call to
120 ///   `getNumFuncArguments`/`getNumFuncResults` from the verifier.
121 /// - To verify that the body respects op-specific invariants, concrete ops may
122 ///   redefine the `verifyBody()` hook that will be called after verifying the
123 ///   function type and the presence of the (potentially empty) body region.
124 template <typename ConcreteType>
125 class FunctionLike : public OpTrait::TraitBase<ConcreteType, FunctionLike> {
126 public:
127   /// Verify that all of the argument attributes are dialect attributes.
128   static LogicalResult verifyTrait(Operation *op);
129 
130   //===--------------------------------------------------------------------===//
131   // Body Handling
132   //===--------------------------------------------------------------------===//
133 
134   /// Returns true if this function is external, i.e. it has no body.
isExternal()135   bool isExternal() { return empty(); }
136 
getBody()137   Region &getBody() { return this->getOperation()->getRegion(0); }
138 
139   /// Delete all blocks from this function.
eraseBody()140   void eraseBody() {
141     getBody().dropAllReferences();
142     getBody().getBlocks().clear();
143   }
144 
145   /// This is the list of blocks in the function.
146   using BlockListType = Region::BlockListType;
getBlocks()147   BlockListType &getBlocks() { return getBody().getBlocks(); }
148 
149   // Iteration over the block in the function.
150   using iterator = BlockListType::iterator;
151   using reverse_iterator = BlockListType::reverse_iterator;
152 
begin()153   iterator begin() { return getBody().begin(); }
end()154   iterator end() { return getBody().end(); }
rbegin()155   reverse_iterator rbegin() { return getBody().rbegin(); }
rend()156   reverse_iterator rend() { return getBody().rend(); }
157 
empty()158   bool empty() { return getBody().empty(); }
push_back(Block * block)159   void push_back(Block *block) { getBody().push_back(block); }
push_front(Block * block)160   void push_front(Block *block) { getBody().push_front(block); }
161 
back()162   Block &back() { return getBody().back(); }
front()163   Block &front() { return getBody().front(); }
164 
165   /// Add an entry block to an empty function, and set up the block arguments
166   /// to match the signature of the function. The newly inserted entry block
167   /// is returned.
168   ///
169   /// Note that the concrete class must define a method with the same name to
170   /// hide this one if the concrete class does not use FunctionType for the
171   /// function type under the hood.
172   Block *addEntryBlock();
173 
174   /// Add a normal block to the end of the function's block list. The function
175   /// should at least already have an entry block.
176   Block *addBlock();
177 
178   /// Hook for concrete ops to verify the contents of the body. Called as a
179   /// part of trait verification, after type verification and ensuring that a
180   /// region exists.
181   LogicalResult verifyBody();
182 
183   //===--------------------------------------------------------------------===//
184   // Type Attribute Handling
185   //===--------------------------------------------------------------------===//
186 
187   /// Return the name of the attribute used for function types.
getTypeAttrName()188   static StringRef getTypeAttrName() { return ::mlir::impl::getTypeAttrName(); }
189 
getTypeAttr()190   TypeAttr getTypeAttr() {
191     return this->getOperation()->template getAttrOfType<TypeAttr>(
192         getTypeAttrName());
193   }
194 
195   /// Return the type of this function.
196   ///
197   /// Note that the concrete class must define a method with the same name to
198   /// hide this one if the concrete class does not use FunctionType for the
199   /// function type under the hood.
getType()200   FunctionType getType() {
201     return getTypeAttr().getValue().template cast<FunctionType>();
202   }
203 
204   /// Return the type of this function without the specified arguments and
205   /// results. This is used to update the function's signature in the
206   /// `eraseArguments` and `eraseResults` methods. The arrays of indices are
207   /// allowed to have duplicates and can be in any order.
208   ///
209   /// Note that the concrete class must define a method with the same name to
210   /// hide this one if the concrete class does not use FunctionType for the
211   /// function type under the hood.
getTypeWithoutArgsAndResults(ArrayRef<unsigned> argIndices,ArrayRef<unsigned> resultIndices)212   FunctionType getTypeWithoutArgsAndResults(ArrayRef<unsigned> argIndices,
213                                             ArrayRef<unsigned> resultIndices) {
214     return getType().getWithoutArgsAndResults(argIndices, resultIndices);
215   }
216 
isTypeAttrValid()217   bool isTypeAttrValid() {
218     auto typeAttr = getTypeAttr();
219     if (!typeAttr)
220       return false;
221     return typeAttr.getValue() != Type{};
222   }
223 
224   /// Change the type of this function in place. This is an extremely dangerous
225   /// operation and it is up to the caller to ensure that this is legal for this
226   /// function, and to restore invariants:
227   ///  - the entry block args must be updated to match the function params.
228   ///  - the argument/result attributes may need an update: if the new type
229   ///    has less parameters we drop the extra attributes, if there are more
230   ///    parameters they won't have any attributes.
231   ///
232   /// Note that the concrete class must define a method with the same name to
233   /// hide this one if the concrete class does not use FunctionType for the
234   /// function type under the hood.
235   void setType(FunctionType newType);
236 
237   //===--------------------------------------------------------------------===//
238   // Argument and Result Handling
239   //===--------------------------------------------------------------------===//
240   using BlockArgListType = Region::BlockArgListType;
241 
getNumArguments()242   unsigned getNumArguments() {
243     return static_cast<ConcreteType *>(this)->getNumFuncArguments();
244   }
245 
getNumResults()246   unsigned getNumResults() {
247     return static_cast<ConcreteType *>(this)->getNumFuncResults();
248   }
249 
250   /// Gets argument.
getArgument(unsigned idx)251   BlockArgument getArgument(unsigned idx) { return getBody().getArgument(idx); }
252 
253   /// Support argument iteration.
254   using args_iterator = Region::args_iterator;
args_begin()255   args_iterator args_begin() { return getBody().args_begin(); }
args_end()256   args_iterator args_end() { return getBody().args_end(); }
getArguments()257   Block::BlockArgListType getArguments() { return getBody().getArguments(); }
258 
getArgumentTypes()259   ValueTypeRange<BlockArgListType> getArgumentTypes() {
260     return getBody().getArgumentTypes();
261   }
262 
263   /// Erase a single argument at `argIndex`.
eraseArgument(unsigned argIndex)264   void eraseArgument(unsigned argIndex) { eraseArguments({argIndex}); }
265 
266   /// Erases the arguments listed in `argIndices`.
267   /// `argIndices` is allowed to have duplicates and can be in any order.
eraseArguments(ArrayRef<unsigned> argIndices)268   void eraseArguments(ArrayRef<unsigned> argIndices) {
269     unsigned originalNumArgs = getNumArguments();
270     Type newType = getTypeWithoutArgsAndResults(argIndices, {});
271     ::mlir::impl::eraseFunctionArguments(this->getOperation(), argIndices,
272                                          originalNumArgs, newType);
273   }
274 
275   /// Erase a single result at `resultIndex`.
eraseResult(unsigned resultIndex)276   void eraseResult(unsigned resultIndex) { eraseResults({resultIndex}); }
277 
278   /// Erases the results listed in `resultIndices`.
279   /// `resultIndices` is allowed to have duplicates and can be in any order.
eraseResults(ArrayRef<unsigned> resultIndices)280   void eraseResults(ArrayRef<unsigned> resultIndices) {
281     unsigned originalNumResults = getNumResults();
282     Type newType = getTypeWithoutArgsAndResults({}, resultIndices);
283     ::mlir::impl::eraseFunctionResults(this->getOperation(), resultIndices,
284                                        originalNumResults, newType);
285   }
286 
287   //===--------------------------------------------------------------------===//
288   // Argument Attributes
289   //===--------------------------------------------------------------------===//
290 
291   /// FunctionLike operations allow for attaching attributes to each of the
292   /// respective function arguments. These argument attributes are stored as
293   /// DictionaryAttrs in the main operation attribute dictionary. The name of
294   /// these entries is `arg` followed by the index of the argument. These
295   /// argument attribute dictionaries are optional, and will generally only
296   /// exist if they are non-empty.
297 
298   /// Return all of the attributes for the argument at 'index'.
getArgAttrs(unsigned index)299   ArrayRef<NamedAttribute> getArgAttrs(unsigned index) {
300     return ::mlir::impl::getArgAttrs(this->getOperation(), index);
301   }
302 
303   /// Return all argument attributes of this function.
getAllArgAttrs(SmallVectorImpl<MutableDictionaryAttr> & result)304   void getAllArgAttrs(SmallVectorImpl<MutableDictionaryAttr> &result) {
305     for (unsigned i = 0, e = getNumArguments(); i != e; ++i)
306       result.emplace_back(getArgAttrDict(i));
307   }
308 
309   /// Return the specified attribute, if present, for the argument at 'index',
310   /// null otherwise.
getArgAttr(unsigned index,Identifier name)311   Attribute getArgAttr(unsigned index, Identifier name) {
312     auto argDict = getArgAttrDict(index);
313     return argDict ? argDict.get(name) : nullptr;
314   }
getArgAttr(unsigned index,StringRef name)315   Attribute getArgAttr(unsigned index, StringRef name) {
316     auto argDict = getArgAttrDict(index);
317     return argDict ? argDict.get(name) : nullptr;
318   }
319 
320   template <typename AttrClass>
getArgAttrOfType(unsigned index,Identifier name)321   AttrClass getArgAttrOfType(unsigned index, Identifier name) {
322     return getArgAttr(index, name).template dyn_cast_or_null<AttrClass>();
323   }
324   template <typename AttrClass>
getArgAttrOfType(unsigned index,StringRef name)325   AttrClass getArgAttrOfType(unsigned index, StringRef name) {
326     return getArgAttr(index, name).template dyn_cast_or_null<AttrClass>();
327   }
328 
329   /// Set the attributes held by the argument at 'index'.
330   void setArgAttrs(unsigned index, ArrayRef<NamedAttribute> attributes);
331   void setArgAttrs(unsigned index, MutableDictionaryAttr attributes);
setAllArgAttrs(ArrayRef<MutableDictionaryAttr> attributes)332   void setAllArgAttrs(ArrayRef<MutableDictionaryAttr> attributes) {
333     assert(attributes.size() == getNumArguments());
334     for (unsigned i = 0, e = attributes.size(); i != e; ++i)
335       setArgAttrs(i, attributes[i]);
336   }
337 
338   /// If the an attribute exists with the specified name, change it to the new
339   /// value. Otherwise, add a new attribute with the specified name/value.
340   void setArgAttr(unsigned index, Identifier name, Attribute value);
setArgAttr(unsigned index,StringRef name,Attribute value)341   void setArgAttr(unsigned index, StringRef name, Attribute value) {
342     setArgAttr(index, Identifier::get(name, this->getOperation()->getContext()),
343                value);
344   }
345 
346   /// Remove the attribute 'name' from the argument at 'index'.
347   MutableDictionaryAttr::RemoveResult removeArgAttr(unsigned index,
348                                                     Identifier name);
349 
350   //===--------------------------------------------------------------------===//
351   // Result Attributes
352   //===--------------------------------------------------------------------===//
353 
354   /// FunctionLike operations allow for attaching attributes to each of the
355   /// respective function results. These result attributes are stored as
356   /// DictionaryAttrs in the main operation attribute dictionary. The name of
357   /// these entries is `result` followed by the index of the result. These
358   /// result attribute dictionaries are optional, and will generally only
359   /// exist if they are non-empty.
360 
361   /// Return all of the attributes for the result at 'index'.
getResultAttrs(unsigned index)362   ArrayRef<NamedAttribute> getResultAttrs(unsigned index) {
363     return ::mlir::impl::getResultAttrs(this->getOperation(), index);
364   }
365 
366   /// Return all result attributes of this function.
getAllResultAttrs(SmallVectorImpl<MutableDictionaryAttr> & result)367   void getAllResultAttrs(SmallVectorImpl<MutableDictionaryAttr> &result) {
368     for (unsigned i = 0, e = getNumResults(); i != e; ++i)
369       result.emplace_back(getResultAttrDict(i));
370   }
371 
372   /// Return the specified attribute, if present, for the result at 'index',
373   /// null otherwise.
getResultAttr(unsigned index,Identifier name)374   Attribute getResultAttr(unsigned index, Identifier name) {
375     auto argDict = getResultAttrDict(index);
376     return argDict ? argDict.get(name) : nullptr;
377   }
getResultAttr(unsigned index,StringRef name)378   Attribute getResultAttr(unsigned index, StringRef name) {
379     auto argDict = getResultAttrDict(index);
380     return argDict ? argDict.get(name) : nullptr;
381   }
382 
383   template <typename AttrClass>
getResultAttrOfType(unsigned index,Identifier name)384   AttrClass getResultAttrOfType(unsigned index, Identifier name) {
385     return getResultAttr(index, name).template dyn_cast_or_null<AttrClass>();
386   }
387   template <typename AttrClass>
getResultAttrOfType(unsigned index,StringRef name)388   AttrClass getResultAttrOfType(unsigned index, StringRef name) {
389     return getResultAttr(index, name).template dyn_cast_or_null<AttrClass>();
390   }
391 
392   /// Set the attributes held by the result at 'index'.
393   void setResultAttrs(unsigned index, ArrayRef<NamedAttribute> attributes);
394   void setResultAttrs(unsigned index, MutableDictionaryAttr attributes);
setAllResultAttrs(ArrayRef<MutableDictionaryAttr> attributes)395   void setAllResultAttrs(ArrayRef<MutableDictionaryAttr> attributes) {
396     assert(attributes.size() == getNumResults());
397     for (unsigned i = 0, e = attributes.size(); i != e; ++i)
398       setResultAttrs(i, attributes[i]);
399   }
400 
401   /// If the an attribute exists with the specified name, change it to the new
402   /// value. Otherwise, add a new attribute with the specified name/value.
403   void setResultAttr(unsigned index, Identifier name, Attribute value);
setResultAttr(unsigned index,StringRef name,Attribute value)404   void setResultAttr(unsigned index, StringRef name, Attribute value) {
405     setResultAttr(index,
406                   Identifier::get(name, this->getOperation()->getContext()),
407                   value);
408   }
409 
410   /// Remove the attribute 'name' from the result at 'index'.
411   MutableDictionaryAttr::RemoveResult removeResultAttr(unsigned index,
412                                                        Identifier name);
413 
414 protected:
415   /// Returns the attribute entry name for the set of argument attributes at
416   /// 'index'.
getArgAttrName(unsigned index,SmallVectorImpl<char> & out)417   static StringRef getArgAttrName(unsigned index, SmallVectorImpl<char> &out) {
418     return ::mlir::impl::getArgAttrName(index, out);
419   }
420 
421   /// Returns the dictionary attribute corresponding to the argument at 'index'.
422   /// If there are no argument attributes at 'index', a null attribute is
423   /// returned.
getArgAttrDict(unsigned index)424   DictionaryAttr getArgAttrDict(unsigned index) {
425     assert(index < getNumArguments() && "invalid argument number");
426     return ::mlir::impl::getArgAttrDict(this->getOperation(), index);
427   }
428 
429   /// Returns the attribute entry name for the set of result attributes at
430   /// 'index'.
getResultAttrName(unsigned index,SmallVectorImpl<char> & out)431   static StringRef getResultAttrName(unsigned index,
432                                      SmallVectorImpl<char> &out) {
433     return ::mlir::impl::getResultAttrName(index, out);
434   }
435 
436   /// Returns the dictionary attribute corresponding to the result at 'index'.
437   /// If there are no result attributes at 'index', a null attribute is
438   /// returned.
getResultAttrDict(unsigned index)439   DictionaryAttr getResultAttrDict(unsigned index) {
440     assert(index < getNumResults() && "invalid result number");
441     return ::mlir::impl::getResultAttrDict(this->getOperation(), index);
442   }
443 
444   /// Hook for concrete classes to verify that the type attribute respects
445   /// op-specific invariants.  Default implementation always succeeds.
verifyType()446   LogicalResult verifyType() { return success(); }
447 };
448 
449 /// Default verifier checks that if the entry block exists, it has the same
450 /// number of arguments as the function-like operation.
451 template <typename ConcreteType>
verifyBody()452 LogicalResult FunctionLike<ConcreteType>::verifyBody() {
453   auto funcOp = cast<ConcreteType>(this->getOperation());
454 
455   if (funcOp.isExternal())
456     return success();
457 
458   unsigned numArguments = funcOp.getNumArguments();
459   if (funcOp.front().getNumArguments() != numArguments)
460     return funcOp.emitOpError("entry block must have ")
461            << numArguments << " arguments to match function signature";
462 
463   return success();
464 }
465 
466 template <typename ConcreteType>
verifyTrait(Operation * op)467 LogicalResult FunctionLike<ConcreteType>::verifyTrait(Operation *op) {
468   MLIRContext *ctx = op->getContext();
469   auto funcOp = cast<ConcreteType>(op);
470 
471   if (!funcOp.isTypeAttrValid())
472     return funcOp.emitOpError("requires a type attribute '")
473            << getTypeAttrName() << '\'';
474 
475   if (failed(funcOp.verifyType()))
476     return failure();
477 
478   for (unsigned i = 0, e = funcOp.getNumArguments(); i != e; ++i) {
479     // Verify that all of the argument attributes are dialect attributes, i.e.
480     // that they contain a dialect prefix in their name.  Call the dialect, if
481     // registered, to verify the attributes themselves.
482     for (auto attr : funcOp.getArgAttrs(i)) {
483       if (!attr.first.strref().contains('.'))
484         return funcOp.emitOpError("arguments may only have dialect attributes");
485       auto dialectNamePair = attr.first.strref().split('.');
486       if (auto *dialect = ctx->getLoadedDialect(dialectNamePair.first)) {
487         if (failed(dialect->verifyRegionArgAttribute(op, /*regionIndex=*/0,
488                                                      /*argIndex=*/i, attr)))
489           return failure();
490       }
491     }
492   }
493 
494   for (unsigned i = 0, e = funcOp.getNumResults(); i != e; ++i) {
495     // Verify that all of the result attributes are dialect attributes, i.e.
496     // that they contain a dialect prefix in their name.  Call the dialect, if
497     // registered, to verify the attributes themselves.
498     for (auto attr : funcOp.getResultAttrs(i)) {
499       if (!attr.first.strref().contains('.'))
500         return funcOp.emitOpError("results may only have dialect attributes");
501       auto dialectNamePair = attr.first.strref().split('.');
502       if (auto *dialect = ctx->getLoadedDialect(dialectNamePair.first)) {
503         if (failed(dialect->verifyRegionResultAttribute(op, /*regionIndex=*/0,
504                                                         /*resultIndex=*/i,
505                                                         attr)))
506           return failure();
507       }
508     }
509   }
510 
511   // Check that the op has exactly one region for the body.
512   if (op->getNumRegions() != 1)
513     return funcOp.emitOpError("expects one region");
514 
515   return funcOp.verifyBody();
516 }
517 
518 //===----------------------------------------------------------------------===//
519 // Function Body.
520 //===----------------------------------------------------------------------===//
521 
522 template <typename ConcreteType>
addEntryBlock()523 Block *FunctionLike<ConcreteType>::addEntryBlock() {
524   assert(empty() && "function already has an entry block");
525   auto *entry = new Block();
526   push_back(entry);
527   entry->addArguments(getType().getInputs());
528   return entry;
529 }
530 
531 template <typename ConcreteType>
addBlock()532 Block *FunctionLike<ConcreteType>::addBlock() {
533   assert(!empty() && "function should at least have an entry block");
534   push_back(new Block());
535   return &back();
536 }
537 
538 //===----------------------------------------------------------------------===//
539 // Function Type Attribute.
540 //===----------------------------------------------------------------------===//
541 
542 template <typename ConcreteType>
setType(FunctionType newType)543 void FunctionLike<ConcreteType>::setType(FunctionType newType) {
544   SmallVector<char, 16> nameBuf;
545   auto oldType = getType();
546   auto *concreteOp = static_cast<ConcreteType *>(this);
547 
548   for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
549     concreteOp->removeAttr(getArgAttrName(i, nameBuf));
550   for (int i = newType.getNumResults(), e = oldType.getNumResults(); i < e; i++)
551     concreteOp->removeAttr(getResultAttrName(i, nameBuf));
552   concreteOp->setAttr(getTypeAttrName(), TypeAttr::get(newType));
553 }
554 
555 //===----------------------------------------------------------------------===//
556 // Function Argument Attribute.
557 //===----------------------------------------------------------------------===//
558 
559 /// Set the attributes held by the argument at 'index'.
560 template <typename ConcreteType>
setArgAttrs(unsigned index,ArrayRef<NamedAttribute> attributes)561 void FunctionLike<ConcreteType>::setArgAttrs(
562     unsigned index, ArrayRef<NamedAttribute> attributes) {
563   assert(index < getNumArguments() && "invalid argument number");
564   SmallString<8> nameOut;
565   getArgAttrName(index, nameOut);
566 
567   if (attributes.empty())
568     return (void)static_cast<ConcreteType *>(this)->removeAttr(nameOut);
569   Operation *op = this->getOperation();
570   op->setAttr(nameOut, DictionaryAttr::get(attributes, op->getContext()));
571 }
572 
573 template <typename ConcreteType>
setArgAttrs(unsigned index,MutableDictionaryAttr attributes)574 void FunctionLike<ConcreteType>::setArgAttrs(unsigned index,
575                                              MutableDictionaryAttr attributes) {
576   assert(index < getNumArguments() && "invalid argument number");
577   SmallString<8> nameOut;
578   if (attributes.getAttrs().empty()) {
579     this->getOperation()->removeAttr(getArgAttrName(index, nameOut));
580   } else {
581     auto newAttr = attributes.getDictionary(
582         attributes.getAttrs().front().second.getContext());
583     return this->getOperation()->setAttr(getArgAttrName(index, nameOut),
584                                          newAttr);
585   }
586 }
587 
588 /// If the an attribute exists with the specified name, change it to the new
589 /// value. Otherwise, add a new attribute with the specified name/value.
590 template <typename ConcreteType>
setArgAttr(unsigned index,Identifier name,Attribute value)591 void FunctionLike<ConcreteType>::setArgAttr(unsigned index, Identifier name,
592                                             Attribute value) {
593   auto curAttr = getArgAttrDict(index);
594   MutableDictionaryAttr attrDict(curAttr);
595   attrDict.set(name, value);
596 
597   // If the attribute changed, then set the new arg attribute list.
598   if (curAttr != attrDict.getDictionary(value.getContext()))
599     setArgAttrs(index, attrDict);
600 }
601 
602 /// Remove the attribute 'name' from the argument at 'index'.
603 template <typename ConcreteType>
604 MutableDictionaryAttr::RemoveResult
removeArgAttr(unsigned index,Identifier name)605 FunctionLike<ConcreteType>::removeArgAttr(unsigned index, Identifier name) {
606   // Build an attribute list and remove the attribute at 'name'.
607   MutableDictionaryAttr attrDict(getArgAttrDict(index));
608   auto result = attrDict.remove(name);
609 
610   // If the attribute was removed, then update the argument dictionary.
611   if (result == MutableDictionaryAttr::RemoveResult::Removed)
612     setArgAttrs(index, attrDict);
613   return result;
614 }
615 
616 //===----------------------------------------------------------------------===//
617 // Function Result Attribute.
618 //===----------------------------------------------------------------------===//
619 
620 /// Set the attributes held by the result at 'index'.
621 template <typename ConcreteType>
setResultAttrs(unsigned index,ArrayRef<NamedAttribute> attributes)622 void FunctionLike<ConcreteType>::setResultAttrs(
623     unsigned index, ArrayRef<NamedAttribute> attributes) {
624   assert(index < getNumResults() && "invalid result number");
625   SmallString<8> nameOut;
626   getResultAttrName(index, nameOut);
627 
628   if (attributes.empty())
629     return (void)this->getOperation()->removeAttr(nameOut);
630   Operation *op = this->getOperation();
631   op->setAttr(nameOut, DictionaryAttr::get(attributes, op->getContext()));
632 }
633 
634 template <typename ConcreteType>
setResultAttrs(unsigned index,MutableDictionaryAttr attributes)635 void FunctionLike<ConcreteType>::setResultAttrs(
636     unsigned index, MutableDictionaryAttr attributes) {
637   assert(index < getNumResults() && "invalid result number");
638   SmallString<8> nameOut;
639   if (attributes.empty()) {
640     this->getOperation()->removeAttr(getResultAttrName(index, nameOut));
641   } else {
642     auto newAttr = attributes.getDictionary(this->getOperation()->getContext());
643     return this->getOperation()->setAttr(getResultAttrName(index, nameOut),
644                                          newAttr);
645   }
646 }
647 
648 /// If the an attribute exists with the specified name, change it to the new
649 /// value. Otherwise, add a new attribute with the specified name/value.
650 template <typename ConcreteType>
setResultAttr(unsigned index,Identifier name,Attribute value)651 void FunctionLike<ConcreteType>::setResultAttr(unsigned index, Identifier name,
652                                                Attribute value) {
653   auto curAttr = getResultAttrDict(index);
654   MutableDictionaryAttr attrDict(curAttr);
655   attrDict.set(name, value);
656 
657   // If the attribute changed, then set the new arg attribute list.
658   if (curAttr != attrDict.getDictionary(value.getContext()))
659     setResultAttrs(index, attrDict);
660 }
661 
662 /// Remove the attribute 'name' from the result at 'index'.
663 template <typename ConcreteType>
664 MutableDictionaryAttr::RemoveResult
removeResultAttr(unsigned index,Identifier name)665 FunctionLike<ConcreteType>::removeResultAttr(unsigned index, Identifier name) {
666   // Build an attribute list and remove the attribute at 'name'.
667   MutableDictionaryAttr attrDict(getResultAttrDict(index));
668   auto result = attrDict.remove(name);
669 
670   // If the attribute was removed, then update the result dictionary.
671   if (result == MutableDictionaryAttr::RemoveResult::Removed)
672     setResultAttrs(index, attrDict);
673   return result;
674 }
675 
676 } // end namespace OpTrait
677 
678 } // end namespace mlir
679 
680 #endif // MLIR_IR_FUNCTIONSUPPORT_H
681