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