1 //===- Builders.cpp - Helpers for constructing MLIR Classes ---------------===//
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/Builders.h"
10 #include "mlir/IR/AffineExpr.h"
11 #include "mlir/IR/AffineMap.h"
12 #include "mlir/IR/BlockAndValueMapping.h"
13 #include "mlir/IR/BuiltinTypes.h"
14 #include "mlir/IR/Dialect.h"
15 #include "mlir/IR/IntegerSet.h"
16 #include "mlir/IR/Matchers.h"
17 #include "mlir/IR/SymbolTable.h"
18 #include "llvm/Support/raw_ostream.h"
19 
20 using namespace mlir;
21 
22 Identifier Builder::getIdentifier(StringRef str) {
23   return Identifier::get(str, context);
24 }
25 
26 //===----------------------------------------------------------------------===//
27 // Locations.
28 //===----------------------------------------------------------------------===//
29 
30 Location Builder::getUnknownLoc() { return UnknownLoc::get(context); }
31 
32 Location Builder::getFusedLoc(ArrayRef<Location> locs, Attribute metadata) {
33   return FusedLoc::get(locs, metadata, context);
34 }
35 
36 //===----------------------------------------------------------------------===//
37 // Types.
38 //===----------------------------------------------------------------------===//
39 
40 FloatType Builder::getBF16Type() { return FloatType::getBF16(context); }
41 
42 FloatType Builder::getF16Type() { return FloatType::getF16(context); }
43 
44 FloatType Builder::getF32Type() { return FloatType::getF32(context); }
45 
46 FloatType Builder::getF64Type() { return FloatType::getF64(context); }
47 
48 FloatType Builder::getF80Type() { return FloatType::getF80(context); }
49 
50 FloatType Builder::getF128Type() { return FloatType::getF128(context); }
51 
52 IndexType Builder::getIndexType() { return IndexType::get(context); }
53 
54 IntegerType Builder::getI1Type() { return IntegerType::get(context, 1); }
55 
56 IntegerType Builder::getI32Type() { return IntegerType::get(context, 32); }
57 
58 IntegerType Builder::getI64Type() { return IntegerType::get(context, 64); }
59 
60 IntegerType Builder::getIntegerType(unsigned width) {
61   return IntegerType::get(context, width);
62 }
63 
64 IntegerType Builder::getIntegerType(unsigned width, bool isSigned) {
65   return IntegerType::get(
66       context, width, isSigned ? IntegerType::Signed : IntegerType::Unsigned);
67 }
68 
69 FunctionType Builder::getFunctionType(TypeRange inputs, TypeRange results) {
70   return FunctionType::get(context, inputs, results);
71 }
72 
73 TupleType Builder::getTupleType(TypeRange elementTypes) {
74   return TupleType::get(context, elementTypes);
75 }
76 
77 NoneType Builder::getNoneType() { return NoneType::get(context); }
78 
79 //===----------------------------------------------------------------------===//
80 // Attributes.
81 //===----------------------------------------------------------------------===//
82 
83 NamedAttribute Builder::getNamedAttr(StringRef name, Attribute val) {
84   return NamedAttribute(getIdentifier(name), val);
85 }
86 
87 UnitAttr Builder::getUnitAttr() { return UnitAttr::get(context); }
88 
89 BoolAttr Builder::getBoolAttr(bool value) {
90   return BoolAttr::get(context, value);
91 }
92 
93 DictionaryAttr Builder::getDictionaryAttr(ArrayRef<NamedAttribute> value) {
94   return DictionaryAttr::get(context, value);
95 }
96 
97 IntegerAttr Builder::getIndexAttr(int64_t value) {
98   return IntegerAttr::get(getIndexType(), APInt(64, value));
99 }
100 
101 IntegerAttr Builder::getI64IntegerAttr(int64_t value) {
102   return IntegerAttr::get(getIntegerType(64), APInt(64, value));
103 }
104 
105 DenseIntElementsAttr Builder::getBoolVectorAttr(ArrayRef<bool> values) {
106   return DenseIntElementsAttr::get(
107       VectorType::get(static_cast<int64_t>(values.size()), getI1Type()),
108       values);
109 }
110 
111 DenseIntElementsAttr Builder::getI32VectorAttr(ArrayRef<int32_t> values) {
112   return DenseIntElementsAttr::get(
113       VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(32)),
114       values);
115 }
116 
117 DenseIntElementsAttr Builder::getI64VectorAttr(ArrayRef<int64_t> values) {
118   return DenseIntElementsAttr::get(
119       VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(64)),
120       values);
121 }
122 
123 DenseIntElementsAttr Builder::getI32TensorAttr(ArrayRef<int32_t> values) {
124   return DenseIntElementsAttr::get(
125       RankedTensorType::get(static_cast<int64_t>(values.size()),
126                             getIntegerType(32)),
127       values);
128 }
129 
130 DenseIntElementsAttr Builder::getI64TensorAttr(ArrayRef<int64_t> values) {
131   return DenseIntElementsAttr::get(
132       RankedTensorType::get(static_cast<int64_t>(values.size()),
133                             getIntegerType(64)),
134       values);
135 }
136 
137 DenseIntElementsAttr Builder::getIndexTensorAttr(ArrayRef<int64_t> values) {
138   return DenseIntElementsAttr::get(
139       RankedTensorType::get(static_cast<int64_t>(values.size()),
140                             getIndexType()),
141       values);
142 }
143 
144 IntegerAttr Builder::getI32IntegerAttr(int32_t value) {
145   return IntegerAttr::get(getIntegerType(32), APInt(32, value));
146 }
147 
148 IntegerAttr Builder::getSI32IntegerAttr(int32_t value) {
149   return IntegerAttr::get(getIntegerType(32, /*isSigned=*/true),
150                           APInt(32, value, /*isSigned=*/true));
151 }
152 
153 IntegerAttr Builder::getUI32IntegerAttr(uint32_t value) {
154   return IntegerAttr::get(getIntegerType(32, /*isSigned=*/false),
155                           APInt(32, (uint64_t)value, /*isSigned=*/false));
156 }
157 
158 IntegerAttr Builder::getI16IntegerAttr(int16_t value) {
159   return IntegerAttr::get(getIntegerType(16), APInt(16, value));
160 }
161 
162 IntegerAttr Builder::getI8IntegerAttr(int8_t value) {
163   return IntegerAttr::get(getIntegerType(8), APInt(8, value));
164 }
165 
166 IntegerAttr Builder::getIntegerAttr(Type type, int64_t value) {
167   if (type.isIndex())
168     return IntegerAttr::get(type, APInt(64, value));
169   return IntegerAttr::get(
170       type, APInt(type.getIntOrFloatBitWidth(), value, type.isSignedInteger()));
171 }
172 
173 IntegerAttr Builder::getIntegerAttr(Type type, const APInt &value) {
174   return IntegerAttr::get(type, value);
175 }
176 
177 FloatAttr Builder::getF64FloatAttr(double value) {
178   return FloatAttr::get(getF64Type(), APFloat(value));
179 }
180 
181 FloatAttr Builder::getF32FloatAttr(float value) {
182   return FloatAttr::get(getF32Type(), APFloat(value));
183 }
184 
185 FloatAttr Builder::getF16FloatAttr(float value) {
186   return FloatAttr::get(getF16Type(), value);
187 }
188 
189 FloatAttr Builder::getFloatAttr(Type type, double value) {
190   return FloatAttr::get(type, value);
191 }
192 
193 FloatAttr Builder::getFloatAttr(Type type, const APFloat &value) {
194   return FloatAttr::get(type, value);
195 }
196 
197 StringAttr Builder::getStringAttr(StringRef bytes) {
198   return StringAttr::get(context, bytes);
199 }
200 
201 ArrayAttr Builder::getArrayAttr(ArrayRef<Attribute> value) {
202   return ArrayAttr::get(context, value);
203 }
204 
205 FlatSymbolRefAttr Builder::getSymbolRefAttr(Operation *value) {
206   auto symName =
207       value->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
208   assert(symName && "value does not have a valid symbol name");
209   return getSymbolRefAttr(symName.getValue());
210 }
211 FlatSymbolRefAttr Builder::getSymbolRefAttr(StringRef value) {
212   return SymbolRefAttr::get(getContext(), value);
213 }
214 SymbolRefAttr
215 Builder::getSymbolRefAttr(StringRef value,
216                           ArrayRef<FlatSymbolRefAttr> nestedReferences) {
217   return SymbolRefAttr::get(getContext(), value, nestedReferences);
218 }
219 
220 ArrayAttr Builder::getBoolArrayAttr(ArrayRef<bool> values) {
221   auto attrs = llvm::to_vector<8>(llvm::map_range(
222       values, [this](bool v) -> Attribute { return getBoolAttr(v); }));
223   return getArrayAttr(attrs);
224 }
225 
226 ArrayAttr Builder::getI32ArrayAttr(ArrayRef<int32_t> values) {
227   auto attrs = llvm::to_vector<8>(llvm::map_range(
228       values, [this](int32_t v) -> Attribute { return getI32IntegerAttr(v); }));
229   return getArrayAttr(attrs);
230 }
231 ArrayAttr Builder::getI64ArrayAttr(ArrayRef<int64_t> values) {
232   auto attrs = llvm::to_vector<8>(llvm::map_range(
233       values, [this](int64_t v) -> Attribute { return getI64IntegerAttr(v); }));
234   return getArrayAttr(attrs);
235 }
236 
237 ArrayAttr Builder::getIndexArrayAttr(ArrayRef<int64_t> values) {
238   auto attrs = llvm::to_vector<8>(
239       llvm::map_range(values, [this](int64_t v) -> Attribute {
240         return getIntegerAttr(IndexType::get(getContext()), v);
241       }));
242   return getArrayAttr(attrs);
243 }
244 
245 ArrayAttr Builder::getF32ArrayAttr(ArrayRef<float> values) {
246   auto attrs = llvm::to_vector<8>(llvm::map_range(
247       values, [this](float v) -> Attribute { return getF32FloatAttr(v); }));
248   return getArrayAttr(attrs);
249 }
250 
251 ArrayAttr Builder::getF64ArrayAttr(ArrayRef<double> values) {
252   auto attrs = llvm::to_vector<8>(llvm::map_range(
253       values, [this](double v) -> Attribute { return getF64FloatAttr(v); }));
254   return getArrayAttr(attrs);
255 }
256 
257 ArrayAttr Builder::getStrArrayAttr(ArrayRef<StringRef> values) {
258   auto attrs = llvm::to_vector<8>(llvm::map_range(
259       values, [this](StringRef v) -> Attribute { return getStringAttr(v); }));
260   return getArrayAttr(attrs);
261 }
262 
263 ArrayAttr Builder::getTypeArrayAttr(TypeRange values) {
264   auto attrs = llvm::to_vector<8>(llvm::map_range(
265       values, [](Type v) -> Attribute { return TypeAttr::get(v); }));
266   return getArrayAttr(attrs);
267 }
268 
269 ArrayAttr Builder::getAffineMapArrayAttr(ArrayRef<AffineMap> values) {
270   auto attrs = llvm::to_vector<8>(llvm::map_range(
271       values, [](AffineMap v) -> Attribute { return AffineMapAttr::get(v); }));
272   return getArrayAttr(attrs);
273 }
274 
275 Attribute Builder::getZeroAttr(Type type) {
276   if (type.isa<FloatType>())
277     return getFloatAttr(type, 0.0);
278   if (type.isa<IndexType>())
279     return getIndexAttr(0);
280   if (auto integerType = type.dyn_cast<IntegerType>())
281     return getIntegerAttr(type, APInt(type.cast<IntegerType>().getWidth(), 0));
282   if (type.isa<RankedTensorType, VectorType>()) {
283     auto vtType = type.cast<ShapedType>();
284     auto element = getZeroAttr(vtType.getElementType());
285     if (!element)
286       return {};
287     return DenseElementsAttr::get(vtType, element);
288   }
289   return {};
290 }
291 
292 //===----------------------------------------------------------------------===//
293 // Affine Expressions, Affine Maps, and Integer Sets.
294 //===----------------------------------------------------------------------===//
295 
296 AffineExpr Builder::getAffineDimExpr(unsigned position) {
297   return mlir::getAffineDimExpr(position, context);
298 }
299 
300 AffineExpr Builder::getAffineSymbolExpr(unsigned position) {
301   return mlir::getAffineSymbolExpr(position, context);
302 }
303 
304 AffineExpr Builder::getAffineConstantExpr(int64_t constant) {
305   return mlir::getAffineConstantExpr(constant, context);
306 }
307 
308 AffineMap Builder::getEmptyAffineMap() { return AffineMap::get(context); }
309 
310 AffineMap Builder::getConstantAffineMap(int64_t val) {
311   return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
312                         getAffineConstantExpr(val));
313 }
314 
315 AffineMap Builder::getDimIdentityMap() {
316   return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, getAffineDimExpr(0));
317 }
318 
319 AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
320   SmallVector<AffineExpr, 4> dimExprs;
321   dimExprs.reserve(rank);
322   for (unsigned i = 0; i < rank; ++i)
323     dimExprs.push_back(getAffineDimExpr(i));
324   return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, dimExprs,
325                         context);
326 }
327 
328 AffineMap Builder::getSymbolIdentityMap() {
329   return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
330                         getAffineSymbolExpr(0));
331 }
332 
333 AffineMap Builder::getSingleDimShiftAffineMap(int64_t shift) {
334   // expr = d0 + shift.
335   auto expr = getAffineDimExpr(0) + shift;
336   return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
337 }
338 
339 AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
340   SmallVector<AffineExpr, 4> shiftedResults;
341   shiftedResults.reserve(map.getNumResults());
342   for (auto resultExpr : map.getResults())
343     shiftedResults.push_back(resultExpr + shift);
344   return AffineMap::get(map.getNumDims(), map.getNumSymbols(), shiftedResults,
345                         context);
346 }
347 
348 //===----------------------------------------------------------------------===//
349 // OpBuilder
350 //===----------------------------------------------------------------------===//
351 
352 OpBuilder::Listener::~Listener() {}
353 
354 /// Insert the given operation at the current insertion point and return it.
355 Operation *OpBuilder::insert(Operation *op) {
356   if (block)
357     block->getOperations().insert(insertPoint, op);
358 
359   if (listener)
360     listener->notifyOperationInserted(op);
361   return op;
362 }
363 
364 /// Add new block with 'argTypes' arguments and set the insertion point to the
365 /// end of it. The block is inserted at the provided insertion point of
366 /// 'parent'.
367 Block *OpBuilder::createBlock(Region *parent, Region::iterator insertPt,
368                               TypeRange argTypes) {
369   assert(parent && "expected valid parent region");
370   if (insertPt == Region::iterator())
371     insertPt = parent->end();
372 
373   Block *b = new Block();
374   b->addArguments(argTypes);
375   parent->getBlocks().insert(insertPt, b);
376   setInsertionPointToEnd(b);
377 
378   if (listener)
379     listener->notifyBlockCreated(b);
380   return b;
381 }
382 
383 /// Add new block with 'argTypes' arguments and set the insertion point to the
384 /// end of it.  The block is placed before 'insertBefore'.
385 Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes) {
386   assert(insertBefore && "expected valid insertion block");
387   return createBlock(insertBefore->getParent(), Region::iterator(insertBefore),
388                      argTypes);
389 }
390 
391 /// Create an operation given the fields represented as an OperationState.
392 Operation *OpBuilder::createOperation(const OperationState &state) {
393   return insert(Operation::create(state));
394 }
395 
396 /// Attempts to fold the given operation and places new results within
397 /// 'results'. Returns success if the operation was folded, failure otherwise.
398 /// Note: This function does not erase the operation on a successful fold.
399 LogicalResult OpBuilder::tryFold(Operation *op,
400                                  SmallVectorImpl<Value> &results) {
401   results.reserve(op->getNumResults());
402   auto cleanupFailure = [&] {
403     results.assign(op->result_begin(), op->result_end());
404     return failure();
405   };
406 
407   // If this operation is already a constant, there is nothing to do.
408   if (matchPattern(op, m_Constant()))
409     return cleanupFailure();
410 
411   // Check to see if any operands to the operation is constant and whether
412   // the operation knows how to constant fold itself.
413   SmallVector<Attribute, 4> constOperands(op->getNumOperands());
414   for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
415     matchPattern(op->getOperand(i), m_Constant(&constOperands[i]));
416 
417   // Try to fold the operation.
418   SmallVector<OpFoldResult, 4> foldResults;
419   if (failed(op->fold(constOperands, foldResults)) || foldResults.empty())
420     return cleanupFailure();
421 
422   // A temporary builder used for creating constants during folding.
423   OpBuilder cstBuilder(context);
424   SmallVector<Operation *, 1> generatedConstants;
425 
426   // Populate the results with the folded results.
427   Dialect *dialect = op->getDialect();
428   for (auto &it : llvm::enumerate(foldResults)) {
429     // Normal values get pushed back directly.
430     if (auto value = it.value().dyn_cast<Value>()) {
431       results.push_back(value);
432       continue;
433     }
434 
435     // Otherwise, try to materialize a constant operation.
436     if (!dialect)
437       return cleanupFailure();
438 
439     // Ask the dialect to materialize a constant operation for this value.
440     Attribute attr = it.value().get<Attribute>();
441     auto *constOp = dialect->materializeConstant(
442         cstBuilder, attr, op->getResult(it.index()).getType(), op->getLoc());
443     if (!constOp) {
444       // Erase any generated constants.
445       for (Operation *cst : generatedConstants)
446         cst->erase();
447       return cleanupFailure();
448     }
449     assert(matchPattern(constOp, m_Constant()));
450 
451     generatedConstants.push_back(constOp);
452     results.push_back(constOp->getResult(0));
453   }
454 
455   // If we were successful, insert any generated constants.
456   for (Operation *cst : generatedConstants)
457     insert(cst);
458 
459   return success();
460 }
461 
462 Operation *OpBuilder::clone(Operation &op, BlockAndValueMapping &mapper) {
463   Operation *newOp = op.clone(mapper);
464   // The `insert` call below handles the notification for inserting `newOp`
465   // itself. But if `newOp` has any regions, we need to notify the listener
466   // about any ops that got inserted inside those regions as part of cloning.
467   if (listener) {
468     auto walkFn = [&](Operation *walkedOp) {
469       listener->notifyOperationInserted(walkedOp);
470     };
471     for (Region &region : newOp->getRegions())
472       region.walk(walkFn);
473   }
474   return insert(newOp);
475 }
476 
477 Operation *OpBuilder::clone(Operation &op) {
478   BlockAndValueMapping mapper;
479   return clone(op, mapper);
480 }
481