1 //===- FunctionSupport.cpp - Utility types 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/FunctionInterfaces.h"
10 #include "mlir/Support/LLVM.h"
11 #include "llvm/ADT/BitVector.h"
12 
13 using namespace mlir;
14 
15 /// Helper to call a callback once on each index in the range
16 /// [0, `totalIndices`), *except* for the indices given in `indices`.
17 /// `indices` is allowed to have duplicates and can be in any order.
18 inline static void iterateIndicesExcept(unsigned totalIndices,
19                                         ArrayRef<unsigned> indices,
20                                         function_ref<void(unsigned)> callback) {
21   llvm::BitVector skipIndices(totalIndices);
22   for (unsigned i : indices)
23     skipIndices.set(i);
24 
25   for (unsigned i = 0; i < totalIndices; ++i)
26     if (!skipIndices.test(i))
27       callback(i);
28 }
29 
30 //===----------------------------------------------------------------------===//
31 // Tablegen Interface Definitions
32 //===----------------------------------------------------------------------===//
33 
34 #include "mlir/IR/FunctionOpInterfaces.cpp.inc"
35 
36 //===----------------------------------------------------------------------===//
37 // Function Arguments and Results.
38 //===----------------------------------------------------------------------===//
39 
40 static bool isEmptyAttrDict(Attribute attr) {
41   return attr.cast<DictionaryAttr>().empty();
42 }
43 
44 DictionaryAttr mlir::function_interface_impl::getArgAttrDict(Operation *op,
45                                                              unsigned index) {
46   ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
47   DictionaryAttr argAttrs =
48       attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
49   return argAttrs;
50 }
51 
52 DictionaryAttr
53 mlir::function_interface_impl::getResultAttrDict(Operation *op,
54                                                  unsigned index) {
55   ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
56   DictionaryAttr resAttrs =
57       attrs ? attrs[index].cast<DictionaryAttr>() : DictionaryAttr();
58   return resAttrs;
59 }
60 
61 void mlir::function_interface_impl::detail::setArgResAttrDict(
62     Operation *op, StringRef attrName, unsigned numTotalIndices, unsigned index,
63     DictionaryAttr attrs) {
64   ArrayAttr allAttrs = op->getAttrOfType<ArrayAttr>(attrName);
65   if (!allAttrs) {
66     if (attrs.empty())
67       return;
68 
69     // If this attribute is not empty, we need to create a new attribute array.
70     SmallVector<Attribute, 8> newAttrs(numTotalIndices,
71                                        DictionaryAttr::get(op->getContext()));
72     newAttrs[index] = attrs;
73     op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
74     return;
75   }
76   // Check to see if the attribute is different from what we already have.
77   if (allAttrs[index] == attrs)
78     return;
79 
80   // If it is, check to see if the attribute array would now contain only empty
81   // dictionaries.
82   ArrayRef<Attribute> rawAttrArray = allAttrs.getValue();
83   if (attrs.empty() &&
84       llvm::all_of(rawAttrArray.take_front(index), isEmptyAttrDict) &&
85       llvm::all_of(rawAttrArray.drop_front(index + 1), isEmptyAttrDict)) {
86     op->removeAttr(attrName);
87     return;
88   }
89 
90   // Otherwise, create a new attribute array with the updated dictionary.
91   SmallVector<Attribute, 8> newAttrs(rawAttrArray.begin(), rawAttrArray.end());
92   newAttrs[index] = attrs;
93   op->setAttr(attrName, ArrayAttr::get(op->getContext(), newAttrs));
94 }
95 
96 /// Set all of the argument or result attribute dictionaries for a function.
97 static void setAllArgResAttrDicts(Operation *op, StringRef attrName,
98                                   ArrayRef<Attribute> attrs) {
99   if (llvm::all_of(attrs, isEmptyAttrDict))
100     op->removeAttr(attrName);
101   else
102     op->setAttr(attrName, ArrayAttr::get(op->getContext(), attrs));
103 }
104 
105 void mlir::function_interface_impl::setAllArgAttrDicts(
106     Operation *op, ArrayRef<DictionaryAttr> attrs) {
107   setAllArgAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
108 }
109 void mlir::function_interface_impl::setAllArgAttrDicts(
110     Operation *op, ArrayRef<Attribute> attrs) {
111   auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
112     return !attr ? DictionaryAttr::get(op->getContext()) : attr;
113   });
114   setAllArgResAttrDicts(op, getArgDictAttrName(),
115                         llvm::to_vector<8>(wrappedAttrs));
116 }
117 
118 void mlir::function_interface_impl::setAllResultAttrDicts(
119     Operation *op, ArrayRef<DictionaryAttr> attrs) {
120   setAllResultAttrDicts(op, ArrayRef<Attribute>(attrs.data(), attrs.size()));
121 }
122 void mlir::function_interface_impl::setAllResultAttrDicts(
123     Operation *op, ArrayRef<Attribute> attrs) {
124   auto wrappedAttrs = llvm::map_range(attrs, [op](Attribute attr) -> Attribute {
125     return !attr ? DictionaryAttr::get(op->getContext()) : attr;
126   });
127   setAllArgResAttrDicts(op, getResultDictAttrName(),
128                         llvm::to_vector<8>(wrappedAttrs));
129 }
130 
131 void mlir::function_interface_impl::insertFunctionArguments(
132     Operation *op, ArrayRef<unsigned> argIndices, TypeRange argTypes,
133     ArrayRef<DictionaryAttr> argAttrs, ArrayRef<Location> argLocs,
134     unsigned originalNumArgs, Type newType) {
135   assert(argIndices.size() == argTypes.size());
136   assert(argIndices.size() == argAttrs.size() || argAttrs.empty());
137   assert(argIndices.size() == argLocs.size());
138   if (argIndices.empty())
139     return;
140 
141   // There are 3 things that need to be updated:
142   // - Function type.
143   // - Arg attrs.
144   // - Block arguments of entry block.
145   Block &entry = op->getRegion(0).front();
146 
147   // Update the argument attributes of the function.
148   auto oldArgAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
149   if (oldArgAttrs || !argAttrs.empty()) {
150     SmallVector<DictionaryAttr, 4> newArgAttrs;
151     newArgAttrs.reserve(originalNumArgs + argIndices.size());
152     unsigned oldIdx = 0;
153     auto migrate = [&](unsigned untilIdx) {
154       if (!oldArgAttrs) {
155         newArgAttrs.resize(newArgAttrs.size() + untilIdx - oldIdx);
156       } else {
157         auto oldArgAttrRange = oldArgAttrs.getAsRange<DictionaryAttr>();
158         newArgAttrs.append(oldArgAttrRange.begin() + oldIdx,
159                            oldArgAttrRange.begin() + untilIdx);
160       }
161       oldIdx = untilIdx;
162     };
163     for (unsigned i = 0, e = argIndices.size(); i < e; ++i) {
164       migrate(argIndices[i]);
165       newArgAttrs.push_back(argAttrs.empty() ? DictionaryAttr{} : argAttrs[i]);
166     }
167     migrate(originalNumArgs);
168     setAllArgAttrDicts(op, newArgAttrs);
169   }
170 
171   // Update the function type and any entry block arguments.
172   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
173   for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
174     entry.insertArgument(argIndices[i] + i, argTypes[i], argLocs[i]);
175 }
176 
177 void mlir::function_interface_impl::insertFunctionResults(
178     Operation *op, ArrayRef<unsigned> resultIndices, TypeRange resultTypes,
179     ArrayRef<DictionaryAttr> resultAttrs, unsigned originalNumResults,
180     Type newType) {
181   assert(resultIndices.size() == resultTypes.size());
182   assert(resultIndices.size() == resultAttrs.size() || resultAttrs.empty());
183   if (resultIndices.empty())
184     return;
185 
186   // There are 2 things that need to be updated:
187   // - Function type.
188   // - Result attrs.
189 
190   // Update the result attributes of the function.
191   auto oldResultAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName());
192   if (oldResultAttrs || !resultAttrs.empty()) {
193     SmallVector<DictionaryAttr, 4> newResultAttrs;
194     newResultAttrs.reserve(originalNumResults + resultIndices.size());
195     unsigned oldIdx = 0;
196     auto migrate = [&](unsigned untilIdx) {
197       if (!oldResultAttrs) {
198         newResultAttrs.resize(newResultAttrs.size() + untilIdx - oldIdx);
199       } else {
200         auto oldResultAttrsRange = oldResultAttrs.getAsRange<DictionaryAttr>();
201         newResultAttrs.append(oldResultAttrsRange.begin() + oldIdx,
202                               oldResultAttrsRange.begin() + untilIdx);
203       }
204       oldIdx = untilIdx;
205     };
206     for (unsigned i = 0, e = resultIndices.size(); i < e; ++i) {
207       migrate(resultIndices[i]);
208       newResultAttrs.push_back(resultAttrs.empty() ? DictionaryAttr{}
209                                                    : resultAttrs[i]);
210     }
211     migrate(originalNumResults);
212     setAllResultAttrDicts(op, newResultAttrs);
213   }
214 
215   // Update the function type.
216   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
217 }
218 
219 void mlir::function_interface_impl::eraseFunctionArguments(
220     Operation *op, ArrayRef<unsigned> argIndices, unsigned originalNumArgs,
221     Type newType) {
222   // There are 3 things that need to be updated:
223   // - Function type.
224   // - Arg attrs.
225   // - Block arguments of entry block.
226   Block &entry = op->getRegion(0).front();
227 
228   // Update the argument attributes of the function.
229   if (auto argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName())) {
230     SmallVector<DictionaryAttr, 4> newArgAttrs;
231     newArgAttrs.reserve(argAttrs.size());
232     iterateIndicesExcept(originalNumArgs, argIndices, [&](unsigned i) {
233       newArgAttrs.emplace_back(argAttrs[i].cast<DictionaryAttr>());
234     });
235     setAllArgAttrDicts(op, newArgAttrs);
236   }
237 
238   // Update the function type and any entry block arguments.
239   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
240   entry.eraseArguments(argIndices);
241 }
242 
243 void mlir::function_interface_impl::eraseFunctionResults(
244     Operation *op, ArrayRef<unsigned> resultIndices,
245     unsigned originalNumResults, Type newType) {
246   // There are 2 things that need to be updated:
247   // - Function type.
248   // - Result attrs.
249 
250   // Update the result attributes of the function.
251   if (auto resAttrs = op->getAttrOfType<ArrayAttr>(getResultDictAttrName())) {
252     SmallVector<DictionaryAttr, 4> newResultAttrs;
253     newResultAttrs.reserve(resAttrs.size());
254     iterateIndicesExcept(originalNumResults, resultIndices, [&](unsigned i) {
255       newResultAttrs.emplace_back(resAttrs[i].cast<DictionaryAttr>());
256     });
257     setAllResultAttrDicts(op, newResultAttrs);
258   }
259 
260   // Update the function type.
261   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
262 }
263 
264 TypeRange mlir::function_interface_impl::insertTypesInto(
265     TypeRange oldTypes, ArrayRef<unsigned> indices, TypeRange newTypes,
266     SmallVectorImpl<Type> &storage) {
267   assert(indices.size() == newTypes.size() &&
268          "mismatch between indice and type count");
269   if (indices.empty())
270     return oldTypes;
271 
272   auto fromIt = oldTypes.begin();
273   for (auto it : llvm::zip(indices, newTypes)) {
274     const auto toIt = oldTypes.begin() + std::get<0>(it);
275     storage.append(fromIt, toIt);
276     storage.push_back(std::get<1>(it));
277     fromIt = toIt;
278   }
279   storage.append(fromIt, oldTypes.end());
280   return storage;
281 }
282 
283 TypeRange
284 mlir::function_interface_impl::filterTypesOut(TypeRange types,
285                                               ArrayRef<unsigned> indices,
286                                               SmallVectorImpl<Type> &storage) {
287   if (indices.empty())
288     return types;
289   iterateIndicesExcept(types.size(), indices,
290                        [&](unsigned i) { storage.emplace_back(types[i]); });
291   return storage;
292 }
293 
294 //===----------------------------------------------------------------------===//
295 // Function type signature.
296 //===----------------------------------------------------------------------===//
297 
298 void mlir::function_interface_impl::setFunctionType(Operation *op,
299                                                     Type newType) {
300   FunctionOpInterface funcOp = cast<FunctionOpInterface>(op);
301   unsigned oldNumArgs = funcOp.getNumArguments();
302   unsigned oldNumResults = funcOp.getNumResults();
303   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
304   unsigned newNumArgs = funcOp.getNumArguments();
305   unsigned newNumResults = funcOp.getNumResults();
306 
307   // Functor used to update the argument and result attributes of the function.
308   auto updateAttrFn = [&](StringRef attrName, unsigned oldCount,
309                           unsigned newCount, auto setAttrFn) {
310     if (oldCount == newCount)
311       return;
312     // The new type has no arguments/results, just drop the attribute.
313     if (newCount == 0) {
314       op->removeAttr(attrName);
315       return;
316     }
317     ArrayAttr attrs = op->getAttrOfType<ArrayAttr>(attrName);
318     if (!attrs)
319       return;
320 
321     // The new type has less arguments/results, take the first N attributes.
322     if (newCount < oldCount)
323       return setAttrFn(op, attrs.getValue().take_front(newCount));
324 
325     // Otherwise, the new type has more arguments/results. Initialize the new
326     // arguments/results with empty attributes.
327     SmallVector<Attribute> newAttrs(attrs.begin(), attrs.end());
328     newAttrs.resize(newCount);
329     setAttrFn(op, newAttrs);
330   };
331 
332   // Update the argument and result attributes.
333   updateAttrFn(
334       getArgDictAttrName(), oldNumArgs, newNumArgs,
335       [&](Operation *op, auto &&attrs) { setAllArgAttrDicts(op, attrs); });
336   updateAttrFn(
337       getResultDictAttrName(), oldNumResults, newNumResults,
338       [&](Operation *op, auto &&attrs) { setAllResultAttrDicts(op, attrs); });
339 }
340