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