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