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