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/BuiltinTypes.h"
13#include "mlir/IR/Dialect.h"
14#include "mlir/IR/IRMapping.h"
15#include "mlir/IR/IntegerSet.h"
16#include "mlir/IR/Matchers.h"
17#include "mlir/IR/SymbolTable.h"
18#include "llvm/ADT/SmallVectorExtras.h"
19#include "llvm/Support/raw_ostream.h"
20
21using namespace mlir;
22
23//===----------------------------------------------------------------------===//
24// Locations.
25//===----------------------------------------------------------------------===//
26
27Location Builder::getUnknownLoc() { return UnknownLoc::get(context); }
28
29Location Builder::getFusedLoc(ArrayRef<Location> locs, Attribute metadata) {
30 return FusedLoc::get(locs, metadata, context);
31}
32
33//===----------------------------------------------------------------------===//
34// Types.
35//===----------------------------------------------------------------------===//
36
37FloatType Builder::getFloat8E5M2Type() {
38 return FloatType::getFloat8E5M2(ctx: context);
39}
40
41FloatType Builder::getFloat8E4M3FNType() {
42 return FloatType::getFloat8E4M3FN(ctx: context);
43}
44
45FloatType Builder::getFloat8E5M2FNUZType() {
46 return FloatType::getFloat8E5M2FNUZ(ctx: context);
47}
48
49FloatType Builder::getFloat8E4M3FNUZType() {
50 return FloatType::getFloat8E4M3FNUZ(ctx: context);
51}
52
53FloatType Builder::getFloat8E4M3B11FNUZType() {
54 return FloatType::getFloat8E4M3B11FNUZ(ctx: context);
55}
56
57FloatType Builder::getBF16Type() { return FloatType::getBF16(ctx: context); }
58
59FloatType Builder::getF16Type() { return FloatType::getF16(ctx: context); }
60
61FloatType Builder::getTF32Type() { return FloatType::getTF32(ctx: context); }
62
63FloatType Builder::getF32Type() { return FloatType::getF32(ctx: context); }
64
65FloatType Builder::getF64Type() { return FloatType::getF64(ctx: context); }
66
67FloatType Builder::getF80Type() { return FloatType::getF80(ctx: context); }
68
69FloatType Builder::getF128Type() { return FloatType::getF128(ctx: context); }
70
71IndexType Builder::getIndexType() { return IndexType::get(context); }
72
73IntegerType Builder::getI1Type() { return IntegerType::get(context, 1); }
74
75IntegerType Builder::getI2Type() { return IntegerType::get(context, 2); }
76
77IntegerType Builder::getI4Type() { return IntegerType::get(context, 4); }
78
79IntegerType Builder::getI8Type() { return IntegerType::get(context, 8); }
80
81IntegerType Builder::getI16Type() { return IntegerType::get(context, 16); }
82
83IntegerType Builder::getI32Type() { return IntegerType::get(context, 32); }
84
85IntegerType Builder::getI64Type() { return IntegerType::get(context, 64); }
86
87IntegerType Builder::getIntegerType(unsigned width) {
88 return IntegerType::get(context, width);
89}
90
91IntegerType Builder::getIntegerType(unsigned width, bool isSigned) {
92 return IntegerType::get(
93 context, width, isSigned ? IntegerType::Signed : IntegerType::Unsigned);
94}
95
96FunctionType Builder::getFunctionType(TypeRange inputs, TypeRange results) {
97 return FunctionType::get(context, inputs, results);
98}
99
100TupleType Builder::getTupleType(TypeRange elementTypes) {
101 return TupleType::get(context, elementTypes);
102}
103
104NoneType Builder::getNoneType() { return NoneType::get(context); }
105
106//===----------------------------------------------------------------------===//
107// Attributes.
108//===----------------------------------------------------------------------===//
109
110NamedAttribute Builder::getNamedAttr(StringRef name, Attribute val) {
111 return NamedAttribute(getStringAttr(name), val);
112}
113
114UnitAttr Builder::getUnitAttr() { return UnitAttr::get(context); }
115
116BoolAttr Builder::getBoolAttr(bool value) {
117 return BoolAttr::get(context, value);
118}
119
120DictionaryAttr Builder::getDictionaryAttr(ArrayRef<NamedAttribute> value) {
121 return DictionaryAttr::get(context, value);
122}
123
124IntegerAttr Builder::getIndexAttr(int64_t value) {
125 return IntegerAttr::get(getIndexType(), APInt(64, value));
126}
127
128IntegerAttr Builder::getI64IntegerAttr(int64_t value) {
129 return IntegerAttr::get(getIntegerType(64), APInt(64, value));
130}
131
132DenseIntElementsAttr Builder::getBoolVectorAttr(ArrayRef<bool> values) {
133 return DenseIntElementsAttr::get(
134 VectorType::get(static_cast<int64_t>(values.size()), getI1Type()),
135 values);
136}
137
138DenseIntElementsAttr Builder::getI32VectorAttr(ArrayRef<int32_t> values) {
139 return DenseIntElementsAttr::get(
140 VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(32)),
141 values);
142}
143
144DenseIntElementsAttr Builder::getI64VectorAttr(ArrayRef<int64_t> values) {
145 return DenseIntElementsAttr::get(
146 VectorType::get(static_cast<int64_t>(values.size()), getIntegerType(64)),
147 values);
148}
149
150DenseIntElementsAttr Builder::getIndexVectorAttr(ArrayRef<int64_t> values) {
151 return DenseIntElementsAttr::get(
152 VectorType::get(static_cast<int64_t>(values.size()), getIndexType()),
153 values);
154}
155
156DenseFPElementsAttr Builder::getF32VectorAttr(ArrayRef<float> values) {
157 return DenseFPElementsAttr::get(
158 VectorType::get(static_cast<float>(values.size()), getF32Type()), values);
159}
160
161DenseFPElementsAttr Builder::getF64VectorAttr(ArrayRef<double> values) {
162 return DenseFPElementsAttr::get(
163 VectorType::get(static_cast<double>(values.size()), getF64Type()),
164 values);
165}
166
167DenseBoolArrayAttr Builder::getDenseBoolArrayAttr(ArrayRef<bool> values) {
168 return DenseBoolArrayAttr::get(context, values);
169}
170
171DenseI8ArrayAttr Builder::getDenseI8ArrayAttr(ArrayRef<int8_t> values) {
172 return DenseI8ArrayAttr::get(context, values);
173}
174
175DenseI16ArrayAttr Builder::getDenseI16ArrayAttr(ArrayRef<int16_t> values) {
176 return DenseI16ArrayAttr::get(context, values);
177}
178
179DenseI32ArrayAttr Builder::getDenseI32ArrayAttr(ArrayRef<int32_t> values) {
180 return DenseI32ArrayAttr::get(context, values);
181}
182
183DenseI64ArrayAttr Builder::getDenseI64ArrayAttr(ArrayRef<int64_t> values) {
184 return DenseI64ArrayAttr::get(context, values);
185}
186
187DenseF32ArrayAttr Builder::getDenseF32ArrayAttr(ArrayRef<float> values) {
188 return DenseF32ArrayAttr::get(context, values);
189}
190
191DenseF64ArrayAttr Builder::getDenseF64ArrayAttr(ArrayRef<double> values) {
192 return DenseF64ArrayAttr::get(context, values);
193}
194
195DenseIntElementsAttr Builder::getI32TensorAttr(ArrayRef<int32_t> values) {
196 return DenseIntElementsAttr::get(
197 RankedTensorType::get(static_cast<int64_t>(values.size()),
198 getIntegerType(32)),
199 values);
200}
201
202DenseIntElementsAttr Builder::getI64TensorAttr(ArrayRef<int64_t> values) {
203 return DenseIntElementsAttr::get(
204 RankedTensorType::get(static_cast<int64_t>(values.size()),
205 getIntegerType(64)),
206 values);
207}
208
209DenseIntElementsAttr Builder::getIndexTensorAttr(ArrayRef<int64_t> values) {
210 return DenseIntElementsAttr::get(
211 RankedTensorType::get(static_cast<int64_t>(values.size()),
212 getIndexType()),
213 values);
214}
215
216IntegerAttr Builder::getI32IntegerAttr(int32_t value) {
217 return IntegerAttr::get(getIntegerType(32), APInt(32, value));
218}
219
220IntegerAttr Builder::getSI32IntegerAttr(int32_t value) {
221 return IntegerAttr::get(getIntegerType(32, /*isSigned=*/true),
222 APInt(32, value, /*isSigned=*/true));
223}
224
225IntegerAttr Builder::getUI32IntegerAttr(uint32_t value) {
226 return IntegerAttr::get(getIntegerType(32, /*isSigned=*/false),
227 APInt(32, (uint64_t)value, /*isSigned=*/false));
228}
229
230IntegerAttr Builder::getI16IntegerAttr(int16_t value) {
231 return IntegerAttr::get(getIntegerType(16), APInt(16, value));
232}
233
234IntegerAttr Builder::getI8IntegerAttr(int8_t value) {
235 return IntegerAttr::get(getIntegerType(8), APInt(8, value));
236}
237
238IntegerAttr Builder::getIntegerAttr(Type type, int64_t value) {
239 if (type.isIndex())
240 return IntegerAttr::get(type, APInt(64, value));
241 return IntegerAttr::get(
242 type, APInt(type.getIntOrFloatBitWidth(), value, type.isSignedInteger()));
243}
244
245IntegerAttr Builder::getIntegerAttr(Type type, const APInt &value) {
246 return IntegerAttr::get(type, value);
247}
248
249FloatAttr Builder::getF64FloatAttr(double value) {
250 return FloatAttr::get(getF64Type(), APFloat(value));
251}
252
253FloatAttr Builder::getF32FloatAttr(float value) {
254 return FloatAttr::get(getF32Type(), APFloat(value));
255}
256
257FloatAttr Builder::getF16FloatAttr(float value) {
258 return FloatAttr::get(getF16Type(), value);
259}
260
261FloatAttr Builder::getFloatAttr(Type type, double value) {
262 return FloatAttr::get(type, value);
263}
264
265FloatAttr Builder::getFloatAttr(Type type, const APFloat &value) {
266 return FloatAttr::get(type, value);
267}
268
269StringAttr Builder::getStringAttr(const Twine &bytes) {
270 return StringAttr::get(context, bytes);
271}
272
273ArrayAttr Builder::getArrayAttr(ArrayRef<Attribute> value) {
274 return ArrayAttr::get(context, value);
275}
276
277ArrayAttr Builder::getBoolArrayAttr(ArrayRef<bool> values) {
278 auto attrs = llvm::map_to_vector<8>(
279 C&: values, F: [this](bool v) -> Attribute { return getBoolAttr(value: v); });
280 return getArrayAttr(attrs);
281}
282
283ArrayAttr Builder::getI32ArrayAttr(ArrayRef<int32_t> values) {
284 auto attrs = llvm::map_to_vector<8>(
285 C&: values, F: [this](int32_t v) -> Attribute { return getI32IntegerAttr(v); });
286 return getArrayAttr(attrs);
287}
288ArrayAttr Builder::getI64ArrayAttr(ArrayRef<int64_t> values) {
289 auto attrs = llvm::map_to_vector<8>(
290 C&: values, F: [this](int64_t v) -> Attribute { return getI64IntegerAttr(v); });
291 return getArrayAttr(attrs);
292}
293
294ArrayAttr Builder::getIndexArrayAttr(ArrayRef<int64_t> values) {
295 auto attrs = llvm::map_to_vector<8>(C&: values, F: [this](int64_t v) -> Attribute {
296 return getIntegerAttr(IndexType::get(getContext()), v);
297 });
298 return getArrayAttr(attrs);
299}
300
301ArrayAttr Builder::getF32ArrayAttr(ArrayRef<float> values) {
302 auto attrs = llvm::map_to_vector<8>(
303 C&: values, F: [this](float v) -> Attribute { return getF32FloatAttr(v); });
304 return getArrayAttr(attrs);
305}
306
307ArrayAttr Builder::getF64ArrayAttr(ArrayRef<double> values) {
308 auto attrs = llvm::map_to_vector<8>(
309 C&: values, F: [this](double v) -> Attribute { return getF64FloatAttr(v); });
310 return getArrayAttr(attrs);
311}
312
313ArrayAttr Builder::getStrArrayAttr(ArrayRef<StringRef> values) {
314 auto attrs = llvm::map_to_vector<8>(
315 C&: values, F: [this](StringRef v) -> Attribute { return getStringAttr(v); });
316 return getArrayAttr(attrs);
317}
318
319ArrayAttr Builder::getTypeArrayAttr(TypeRange values) {
320 auto attrs = llvm::map_to_vector<8>(
321 C&: values, F: [](Type v) -> Attribute { return TypeAttr::get(v); });
322 return getArrayAttr(attrs);
323}
324
325ArrayAttr Builder::getAffineMapArrayAttr(ArrayRef<AffineMap> values) {
326 auto attrs = llvm::map_to_vector<8>(
327 C&: values, F: [](AffineMap v) -> Attribute { return AffineMapAttr::get(v); });
328 return getArrayAttr(attrs);
329}
330
331TypedAttr Builder::getZeroAttr(Type type) {
332 if (llvm::isa<FloatType>(Val: type))
333 return getFloatAttr(type, 0.0);
334 if (llvm::isa<IndexType>(type))
335 return getIndexAttr(0);
336 if (llvm::dyn_cast<IntegerType>(type))
337 return getIntegerAttr(type,
338 APInt(llvm::cast<IntegerType>(type).getWidth(), 0));
339 if (llvm::isa<RankedTensorType, VectorType>(Val: type)) {
340 auto vtType = llvm::cast<ShapedType>(type);
341 auto element = getZeroAttr(vtType.getElementType());
342 if (!element)
343 return {};
344 return DenseElementsAttr::get(vtType, element);
345 }
346 return {};
347}
348
349TypedAttr Builder::getOneAttr(Type type) {
350 if (llvm::isa<FloatType>(Val: type))
351 return getFloatAttr(type, 1.0);
352 if (llvm::isa<IndexType>(type))
353 return getIndexAttr(1);
354 if (llvm::dyn_cast<IntegerType>(type))
355 return getIntegerAttr(type,
356 APInt(llvm::cast<IntegerType>(type).getWidth(), 1));
357 if (llvm::isa<RankedTensorType, VectorType>(Val: type)) {
358 auto vtType = llvm::cast<ShapedType>(type);
359 auto element = getOneAttr(vtType.getElementType());
360 if (!element)
361 return {};
362 return DenseElementsAttr::get(vtType, element);
363 }
364 return {};
365}
366
367//===----------------------------------------------------------------------===//
368// Affine Expressions, Affine Maps, and Integer Sets.
369//===----------------------------------------------------------------------===//
370
371AffineExpr Builder::getAffineDimExpr(unsigned position) {
372 return mlir::getAffineDimExpr(position, context);
373}
374
375AffineExpr Builder::getAffineSymbolExpr(unsigned position) {
376 return mlir::getAffineSymbolExpr(position, context);
377}
378
379AffineExpr Builder::getAffineConstantExpr(int64_t constant) {
380 return mlir::getAffineConstantExpr(constant, context);
381}
382
383AffineMap Builder::getEmptyAffineMap() { return AffineMap::get(context); }
384
385AffineMap Builder::getConstantAffineMap(int64_t val) {
386 return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
387 result: getAffineConstantExpr(constant: val));
388}
389
390AffineMap Builder::getDimIdentityMap() {
391 return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, result: getAffineDimExpr(position: 0));
392}
393
394AffineMap Builder::getMultiDimIdentityMap(unsigned rank) {
395 SmallVector<AffineExpr, 4> dimExprs;
396 dimExprs.reserve(N: rank);
397 for (unsigned i = 0; i < rank; ++i)
398 dimExprs.push_back(Elt: getAffineDimExpr(position: i));
399 return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/0, results: dimExprs,
400 context);
401}
402
403AffineMap Builder::getSymbolIdentityMap() {
404 return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
405 result: getAffineSymbolExpr(position: 0));
406}
407
408AffineMap Builder::getSingleDimShiftAffineMap(int64_t shift) {
409 // expr = d0 + shift.
410 auto expr = getAffineDimExpr(position: 0) + shift;
411 return AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, result: expr);
412}
413
414AffineMap Builder::getShiftedAffineMap(AffineMap map, int64_t shift) {
415 SmallVector<AffineExpr, 4> shiftedResults;
416 shiftedResults.reserve(N: map.getNumResults());
417 for (auto resultExpr : map.getResults())
418 shiftedResults.push_back(Elt: resultExpr + shift);
419 return AffineMap::get(dimCount: map.getNumDims(), symbolCount: map.getNumSymbols(), results: shiftedResults,
420 context);
421}
422
423//===----------------------------------------------------------------------===//
424// OpBuilder
425//===----------------------------------------------------------------------===//
426
427/// Insert the given operation at the current insertion point and return it.
428Operation *OpBuilder::insert(Operation *op) {
429 if (block) {
430 block->getOperations().insert(where: insertPoint, New: op);
431 if (listener)
432 listener->notifyOperationInserted(op, /*previous=*/{});
433 }
434 return op;
435}
436
437Block *OpBuilder::createBlock(Region *parent, Region::iterator insertPt,
438 TypeRange argTypes, ArrayRef<Location> locs) {
439 assert(parent && "expected valid parent region");
440 assert(argTypes.size() == locs.size() && "argument location mismatch");
441 if (insertPt == Region::iterator())
442 insertPt = parent->end();
443
444 Block *b = new Block();
445 b->addArguments(types: argTypes, locs);
446 parent->getBlocks().insert(where: insertPt, New: b);
447 setInsertionPointToEnd(b);
448
449 if (listener)
450 listener->notifyBlockInserted(block: b, /*previous=*/nullptr, /*previousIt=*/{});
451 return b;
452}
453
454/// Add new block with 'argTypes' arguments and set the insertion point to the
455/// end of it. The block is placed before 'insertBefore'.
456Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes,
457 ArrayRef<Location> locs) {
458 assert(insertBefore && "expected valid insertion block");
459 return createBlock(parent: insertBefore->getParent(), insertPt: Region::iterator(insertBefore),
460 argTypes, locs);
461}
462
463/// Create an operation given the fields represented as an OperationState.
464Operation *OpBuilder::create(const OperationState &state) {
465 return insert(op: Operation::create(state));
466}
467
468/// Creates an operation with the given fields.
469Operation *OpBuilder::create(Location loc, StringAttr opName,
470 ValueRange operands, TypeRange types,
471 ArrayRef<NamedAttribute> attributes,
472 BlockRange successors,
473 MutableArrayRef<std::unique_ptr<Region>> regions) {
474 OperationState state(loc, opName, operands, types, attributes, successors,
475 regions);
476 return create(state);
477}
478
479LogicalResult OpBuilder::tryFold(Operation *op,
480 SmallVectorImpl<Value> &results) {
481 assert(results.empty() && "expected empty results");
482 ResultRange opResults = op->getResults();
483
484 results.reserve(N: opResults.size());
485 auto cleanupFailure = [&] {
486 results.clear();
487 return failure();
488 };
489
490 // If this operation is already a constant, there is nothing to do.
491 if (matchPattern(op, pattern: m_Constant()))
492 return cleanupFailure();
493
494 // Try to fold the operation.
495 SmallVector<OpFoldResult, 4> foldResults;
496 if (failed(result: op->fold(results&: foldResults)))
497 return cleanupFailure();
498
499 // An in-place fold does not require generation of any constants.
500 if (foldResults.empty())
501 return success();
502
503 // A temporary builder used for creating constants during folding.
504 OpBuilder cstBuilder(context);
505 SmallVector<Operation *, 1> generatedConstants;
506
507 // Populate the results with the folded results.
508 Dialect *dialect = op->getDialect();
509 for (auto [foldResult, expectedType] :
510 llvm::zip_equal(t&: foldResults, u: opResults.getTypes())) {
511
512 // Normal values get pushed back directly.
513 if (auto value = llvm::dyn_cast_if_present<Value>(Val&: foldResult)) {
514 results.push_back(Elt: value);
515 continue;
516 }
517
518 // Otherwise, try to materialize a constant operation.
519 if (!dialect)
520 return cleanupFailure();
521
522 // Ask the dialect to materialize a constant operation for this value.
523 Attribute attr = foldResult.get<Attribute>();
524 auto *constOp = dialect->materializeConstant(builder&: cstBuilder, value: attr, type: expectedType,
525 loc: op->getLoc());
526 if (!constOp) {
527 // Erase any generated constants.
528 for (Operation *cst : generatedConstants)
529 cst->erase();
530 return cleanupFailure();
531 }
532 assert(matchPattern(constOp, m_Constant()));
533
534 generatedConstants.push_back(Elt: constOp);
535 results.push_back(Elt: constOp->getResult(idx: 0));
536 }
537
538 // If we were successful, insert any generated constants.
539 for (Operation *cst : generatedConstants)
540 insert(op: cst);
541
542 return success();
543}
544
545/// Helper function that sends block insertion notifications for every block
546/// that is directly nested in the given op.
547static void notifyBlockInsertions(Operation *op,
548 OpBuilder::Listener *listener) {
549 for (Region &r : op->getRegions())
550 for (Block &b : r.getBlocks())
551 listener->notifyBlockInserted(block: &b, /*previous=*/nullptr,
552 /*previousIt=*/{});
553}
554
555Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
556 Operation *newOp = op.clone(mapper);
557 newOp = insert(op: newOp);
558
559 // The `insert` call above handles the notification for inserting `newOp`
560 // itself. But if `newOp` has any regions, we need to notify the listener
561 // about any ops that got inserted inside those regions as part of cloning.
562 if (listener) {
563 // The `insert` call above notifies about op insertion, but not about block
564 // insertion.
565 notifyBlockInsertions(op: newOp, listener);
566 auto walkFn = [&](Operation *walkedOp) {
567 listener->notifyOperationInserted(op: walkedOp, /*previous=*/{});
568 notifyBlockInsertions(op: walkedOp, listener);
569 };
570 for (Region &region : newOp->getRegions())
571 region.walk<WalkOrder::PreOrder>(callback&: walkFn);
572 }
573
574 return newOp;
575}
576
577Operation *OpBuilder::clone(Operation &op) {
578 IRMapping mapper;
579 return clone(op, mapper);
580}
581
582void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
583 Region::iterator before, IRMapping &mapping) {
584 region.cloneInto(dest: &parent, destPos: before, mapper&: mapping);
585
586 // Fast path: If no listener is attached, there is no more work to do.
587 if (!listener)
588 return;
589
590 // Notify about op/block insertion.
591 for (auto it = mapping.lookup(from: &region.front())->getIterator(); it != before;
592 ++it) {
593 listener->notifyBlockInserted(block: &*it, /*previous=*/nullptr,
594 /*previousIt=*/{});
595 it->walk<WalkOrder::PreOrder>(callback: [&](Operation *walkedOp) {
596 listener->notifyOperationInserted(op: walkedOp, /*previous=*/{});
597 notifyBlockInsertions(op: walkedOp, listener);
598 });
599 }
600}
601
602void OpBuilder::cloneRegionBefore(Region &region, Region &parent,
603 Region::iterator before) {
604 IRMapping mapping;
605 cloneRegionBefore(region, parent, before, mapping);
606}
607
608void OpBuilder::cloneRegionBefore(Region &region, Block *before) {
609 cloneRegionBefore(region, parent&: *before->getParent(), before: before->getIterator());
610}
611

source code of mlir/lib/IR/Builders.cpp