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