1 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===// 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/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" 10 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h" 11 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" 12 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 13 #include "mlir/Dialect/StandardOps/Ops.h" 14 #include "mlir/Dialect/VectorOps/VectorOps.h" 15 #include "mlir/IR/Attributes.h" 16 #include "mlir/IR/Builders.h" 17 #include "mlir/IR/MLIRContext.h" 18 #include "mlir/IR/Module.h" 19 #include "mlir/IR/Operation.h" 20 #include "mlir/IR/PatternMatch.h" 21 #include "mlir/IR/StandardTypes.h" 22 #include "mlir/IR/Types.h" 23 #include "mlir/Pass/Pass.h" 24 #include "mlir/Pass/PassManager.h" 25 #include "mlir/Transforms/DialectConversion.h" 26 #include "mlir/Transforms/Passes.h" 27 28 #include "llvm/IR/DerivedTypes.h" 29 #include "llvm/IR/Module.h" 30 #include "llvm/IR/Type.h" 31 #include "llvm/Support/Allocator.h" 32 #include "llvm/Support/ErrorHandling.h" 33 34 using namespace mlir; 35 using namespace mlir::vector; 36 37 template <typename T> 38 static LLVM::LLVMType getPtrToElementType(T containerType, 39 LLVMTypeConverter &lowering) { 40 return lowering.convertType(containerType.getElementType()) 41 .template cast<LLVM::LLVMType>() 42 .getPointerTo(); 43 } 44 45 // Helper to reduce vector type by one rank at front. 46 static VectorType reducedVectorTypeFront(VectorType tp) { 47 assert((tp.getRank() > 1) && "unlowerable vector type"); 48 return VectorType::get(tp.getShape().drop_front(), tp.getElementType()); 49 } 50 51 // Helper to reduce vector type by *all* but one rank at back. 52 static VectorType reducedVectorTypeBack(VectorType tp) { 53 assert((tp.getRank() > 1) && "unlowerable vector type"); 54 return VectorType::get(tp.getShape().take_back(), tp.getElementType()); 55 } 56 57 // Helper that picks the proper sequence for inserting. 58 static Value insertOne(ConversionPatternRewriter &rewriter, 59 LLVMTypeConverter &lowering, Location loc, Value val1, 60 Value val2, Type llvmType, int64_t rank, int64_t pos) { 61 if (rank == 1) { 62 auto idxType = rewriter.getIndexType(); 63 auto constant = rewriter.create<LLVM::ConstantOp>( 64 loc, lowering.convertType(idxType), 65 rewriter.getIntegerAttr(idxType, pos)); 66 return rewriter.create<LLVM::InsertElementOp>(loc, llvmType, val1, val2, 67 constant); 68 } 69 return rewriter.create<LLVM::InsertValueOp>(loc, llvmType, val1, val2, 70 rewriter.getI64ArrayAttr(pos)); 71 } 72 73 // Helper that picks the proper sequence for inserting. 74 static Value insertOne(PatternRewriter &rewriter, Location loc, Value from, 75 Value into, int64_t offset) { 76 auto vectorType = into.getType().cast<VectorType>(); 77 if (vectorType.getRank() > 1) 78 return rewriter.create<InsertOp>(loc, from, into, offset); 79 return rewriter.create<vector::InsertElementOp>( 80 loc, vectorType, from, into, 81 rewriter.create<ConstantIndexOp>(loc, offset)); 82 } 83 84 // Helper that picks the proper sequence for extracting. 85 static Value extractOne(ConversionPatternRewriter &rewriter, 86 LLVMTypeConverter &lowering, Location loc, Value val, 87 Type llvmType, int64_t rank, int64_t pos) { 88 if (rank == 1) { 89 auto idxType = rewriter.getIndexType(); 90 auto constant = rewriter.create<LLVM::ConstantOp>( 91 loc, lowering.convertType(idxType), 92 rewriter.getIntegerAttr(idxType, pos)); 93 return rewriter.create<LLVM::ExtractElementOp>(loc, llvmType, val, 94 constant); 95 } 96 return rewriter.create<LLVM::ExtractValueOp>(loc, llvmType, val, 97 rewriter.getI64ArrayAttr(pos)); 98 } 99 100 // Helper that picks the proper sequence for extracting. 101 static Value extractOne(PatternRewriter &rewriter, Location loc, Value vector, 102 int64_t offset) { 103 auto vectorType = vector.getType().cast<VectorType>(); 104 if (vectorType.getRank() > 1) 105 return rewriter.create<ExtractOp>(loc, vector, offset); 106 return rewriter.create<vector::ExtractElementOp>( 107 loc, vectorType.getElementType(), vector, 108 rewriter.create<ConstantIndexOp>(loc, offset)); 109 } 110 111 // Helper that returns a subset of `arrayAttr` as a vector of int64_t. 112 // TODO(rriddle): Better support for attribute subtype forwarding + slicing. 113 static SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr, 114 unsigned dropFront = 0, 115 unsigned dropBack = 0) { 116 assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds"); 117 auto range = arrayAttr.getAsRange<IntegerAttr>(); 118 SmallVector<int64_t, 4> res; 119 res.reserve(arrayAttr.size() - dropFront - dropBack); 120 for (auto it = range.begin() + dropFront, eit = range.end() - dropBack; 121 it != eit; ++it) 122 res.push_back((*it).getValue().getSExtValue()); 123 return res; 124 } 125 126 class VectorBroadcastOpConversion : public LLVMOpLowering { 127 public: 128 explicit VectorBroadcastOpConversion(MLIRContext *context, 129 LLVMTypeConverter &typeConverter) 130 : LLVMOpLowering(vector::BroadcastOp::getOperationName(), context, 131 typeConverter) {} 132 133 PatternMatchResult 134 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 135 ConversionPatternRewriter &rewriter) const override { 136 auto broadcastOp = cast<vector::BroadcastOp>(op); 137 VectorType dstVectorType = broadcastOp.getVectorType(); 138 if (lowering.convertType(dstVectorType) == nullptr) 139 return matchFailure(); 140 // Rewrite when the full vector type can be lowered (which 141 // implies all 'reduced' types can be lowered too). 142 auto adaptor = vector::BroadcastOpOperandAdaptor(operands); 143 VectorType srcVectorType = 144 broadcastOp.getSourceType().dyn_cast<VectorType>(); 145 rewriter.replaceOp( 146 op, expandRanks(adaptor.source(), // source value to be expanded 147 op->getLoc(), // location of original broadcast 148 srcVectorType, dstVectorType, rewriter)); 149 return matchSuccess(); 150 } 151 152 private: 153 // Expands the given source value over all the ranks, as defined 154 // by the source and destination type (a null source type denotes 155 // expansion from a scalar value into a vector). 156 // 157 // TODO(ajcbik): consider replacing this one-pattern lowering 158 // with a two-pattern lowering using other vector 159 // ops once all insert/extract/shuffle operations 160 // are available with lowering implemention. 161 // 162 Value expandRanks(Value value, Location loc, VectorType srcVectorType, 163 VectorType dstVectorType, 164 ConversionPatternRewriter &rewriter) const { 165 assert((dstVectorType != nullptr) && "invalid result type in broadcast"); 166 // Determine rank of source and destination. 167 int64_t srcRank = srcVectorType ? srcVectorType.getRank() : 0; 168 int64_t dstRank = dstVectorType.getRank(); 169 int64_t curDim = dstVectorType.getDimSize(0); 170 if (srcRank < dstRank) 171 // Duplicate this rank. 172 return duplicateOneRank(value, loc, srcVectorType, dstVectorType, dstRank, 173 curDim, rewriter); 174 // If all trailing dimensions are the same, the broadcast consists of 175 // simply passing through the source value and we are done. Otherwise, 176 // any non-matching dimension forces a stretch along this rank. 177 assert((srcVectorType != nullptr) && (srcRank > 0) && 178 (srcRank == dstRank) && "invalid rank in broadcast"); 179 for (int64_t r = 0; r < dstRank; r++) { 180 if (srcVectorType.getDimSize(r) != dstVectorType.getDimSize(r)) { 181 return stretchOneRank(value, loc, srcVectorType, dstVectorType, dstRank, 182 curDim, rewriter); 183 } 184 } 185 return value; 186 } 187 188 // Picks the best way to duplicate a single rank. For the 1-D case, a 189 // single insert-elt/shuffle is the most efficient expansion. For higher 190 // dimensions, however, we need dim x insert-values on a new broadcast 191 // with one less leading dimension, which will be lowered "recursively" 192 // to matching LLVM IR. 193 // For example: 194 // v = broadcast s : f32 to vector<4x2xf32> 195 // becomes: 196 // x = broadcast s : f32 to vector<2xf32> 197 // v = [x,x,x,x] 198 // becomes: 199 // x = [s,s] 200 // v = [x,x,x,x] 201 Value duplicateOneRank(Value value, Location loc, VectorType srcVectorType, 202 VectorType dstVectorType, int64_t rank, int64_t dim, 203 ConversionPatternRewriter &rewriter) const { 204 Type llvmType = lowering.convertType(dstVectorType); 205 assert((llvmType != nullptr) && "unlowerable vector type"); 206 if (rank == 1) { 207 Value undef = rewriter.create<LLVM::UndefOp>(loc, llvmType); 208 Value expand = 209 insertOne(rewriter, lowering, loc, undef, value, llvmType, rank, 0); 210 SmallVector<int32_t, 4> zeroValues(dim, 0); 211 return rewriter.create<LLVM::ShuffleVectorOp>( 212 loc, expand, undef, rewriter.getI32ArrayAttr(zeroValues)); 213 } 214 Value expand = expandRanks(value, loc, srcVectorType, 215 reducedVectorTypeFront(dstVectorType), rewriter); 216 Value result = rewriter.create<LLVM::UndefOp>(loc, llvmType); 217 for (int64_t d = 0; d < dim; ++d) { 218 result = 219 insertOne(rewriter, lowering, loc, result, expand, llvmType, rank, d); 220 } 221 return result; 222 } 223 224 // Picks the best way to stretch a single rank. For the 1-D case, a 225 // single insert-elt/shuffle is the most efficient expansion when at 226 // a stretch. Otherwise, every dimension needs to be expanded 227 // individually and individually inserted in the resulting vector. 228 // For example: 229 // v = broadcast w : vector<4x1x2xf32> to vector<4x2x2xf32> 230 // becomes: 231 // a = broadcast w[0] : vector<1x2xf32> to vector<2x2xf32> 232 // b = broadcast w[1] : vector<1x2xf32> to vector<2x2xf32> 233 // c = broadcast w[2] : vector<1x2xf32> to vector<2x2xf32> 234 // d = broadcast w[3] : vector<1x2xf32> to vector<2x2xf32> 235 // v = [a,b,c,d] 236 // becomes: 237 // x = broadcast w[0][0] : vector<2xf32> to vector <2x2xf32> 238 // y = broadcast w[1][0] : vector<2xf32> to vector <2x2xf32> 239 // a = [x, y] 240 // etc. 241 Value stretchOneRank(Value value, Location loc, VectorType srcVectorType, 242 VectorType dstVectorType, int64_t rank, int64_t dim, 243 ConversionPatternRewriter &rewriter) const { 244 Type llvmType = lowering.convertType(dstVectorType); 245 assert((llvmType != nullptr) && "unlowerable vector type"); 246 Value result = rewriter.create<LLVM::UndefOp>(loc, llvmType); 247 bool atStretch = dim != srcVectorType.getDimSize(0); 248 if (rank == 1) { 249 assert(atStretch); 250 Type redLlvmType = lowering.convertType(dstVectorType.getElementType()); 251 Value one = 252 extractOne(rewriter, lowering, loc, value, redLlvmType, rank, 0); 253 Value expand = 254 insertOne(rewriter, lowering, loc, result, one, llvmType, rank, 0); 255 SmallVector<int32_t, 4> zeroValues(dim, 0); 256 return rewriter.create<LLVM::ShuffleVectorOp>( 257 loc, expand, result, rewriter.getI32ArrayAttr(zeroValues)); 258 } 259 VectorType redSrcType = reducedVectorTypeFront(srcVectorType); 260 VectorType redDstType = reducedVectorTypeFront(dstVectorType); 261 Type redLlvmType = lowering.convertType(redSrcType); 262 for (int64_t d = 0; d < dim; ++d) { 263 int64_t pos = atStretch ? 0 : d; 264 Value one = 265 extractOne(rewriter, lowering, loc, value, redLlvmType, rank, pos); 266 Value expand = expandRanks(one, loc, redSrcType, redDstType, rewriter); 267 result = 268 insertOne(rewriter, lowering, loc, result, expand, llvmType, rank, d); 269 } 270 return result; 271 } 272 }; 273 274 class VectorShuffleOpConversion : public LLVMOpLowering { 275 public: 276 explicit VectorShuffleOpConversion(MLIRContext *context, 277 LLVMTypeConverter &typeConverter) 278 : LLVMOpLowering(vector::ShuffleOp::getOperationName(), context, 279 typeConverter) {} 280 281 PatternMatchResult 282 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 283 ConversionPatternRewriter &rewriter) const override { 284 auto loc = op->getLoc(); 285 auto adaptor = vector::ShuffleOpOperandAdaptor(operands); 286 auto shuffleOp = cast<vector::ShuffleOp>(op); 287 auto v1Type = shuffleOp.getV1VectorType(); 288 auto v2Type = shuffleOp.getV2VectorType(); 289 auto vectorType = shuffleOp.getVectorType(); 290 Type llvmType = lowering.convertType(vectorType); 291 auto maskArrayAttr = shuffleOp.mask(); 292 293 // Bail if result type cannot be lowered. 294 if (!llvmType) 295 return matchFailure(); 296 297 // Get rank and dimension sizes. 298 int64_t rank = vectorType.getRank(); 299 assert(v1Type.getRank() == rank); 300 assert(v2Type.getRank() == rank); 301 int64_t v1Dim = v1Type.getDimSize(0); 302 303 // For rank 1, where both operands have *exactly* the same vector type, 304 // there is direct shuffle support in LLVM. Use it! 305 if (rank == 1 && v1Type == v2Type) { 306 Value shuffle = rewriter.create<LLVM::ShuffleVectorOp>( 307 loc, adaptor.v1(), adaptor.v2(), maskArrayAttr); 308 rewriter.replaceOp(op, shuffle); 309 return matchSuccess(); 310 } 311 312 // For all other cases, insert the individual values individually. 313 Value insert = rewriter.create<LLVM::UndefOp>(loc, llvmType); 314 int64_t insPos = 0; 315 for (auto en : llvm::enumerate(maskArrayAttr)) { 316 int64_t extPos = en.value().cast<IntegerAttr>().getInt(); 317 Value value = adaptor.v1(); 318 if (extPos >= v1Dim) { 319 extPos -= v1Dim; 320 value = adaptor.v2(); 321 } 322 Value extract = 323 extractOne(rewriter, lowering, loc, value, llvmType, rank, extPos); 324 insert = insertOne(rewriter, lowering, loc, insert, extract, llvmType, 325 rank, insPos++); 326 } 327 rewriter.replaceOp(op, insert); 328 return matchSuccess(); 329 } 330 }; 331 332 class VectorExtractElementOpConversion : public LLVMOpLowering { 333 public: 334 explicit VectorExtractElementOpConversion(MLIRContext *context, 335 LLVMTypeConverter &typeConverter) 336 : LLVMOpLowering(vector::ExtractElementOp::getOperationName(), context, 337 typeConverter) {} 338 339 PatternMatchResult 340 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 341 ConversionPatternRewriter &rewriter) const override { 342 auto adaptor = vector::ExtractElementOpOperandAdaptor(operands); 343 auto extractEltOp = cast<vector::ExtractElementOp>(op); 344 auto vectorType = extractEltOp.getVectorType(); 345 auto llvmType = lowering.convertType(vectorType.getElementType()); 346 347 // Bail if result type cannot be lowered. 348 if (!llvmType) 349 return matchFailure(); 350 351 rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>( 352 op, llvmType, adaptor.vector(), adaptor.position()); 353 return matchSuccess(); 354 } 355 }; 356 357 class VectorExtractOpConversion : public LLVMOpLowering { 358 public: 359 explicit VectorExtractOpConversion(MLIRContext *context, 360 LLVMTypeConverter &typeConverter) 361 : LLVMOpLowering(vector::ExtractOp::getOperationName(), context, 362 typeConverter) {} 363 364 PatternMatchResult 365 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 366 ConversionPatternRewriter &rewriter) const override { 367 auto loc = op->getLoc(); 368 auto adaptor = vector::ExtractOpOperandAdaptor(operands); 369 auto extractOp = cast<vector::ExtractOp>(op); 370 auto vectorType = extractOp.getVectorType(); 371 auto resultType = extractOp.getResult()->getType(); 372 auto llvmResultType = lowering.convertType(resultType); 373 auto positionArrayAttr = extractOp.position(); 374 375 // Bail if result type cannot be lowered. 376 if (!llvmResultType) 377 return matchFailure(); 378 379 // One-shot extraction of vector from array (only requires extractvalue). 380 if (resultType.isa<VectorType>()) { 381 Value extracted = rewriter.create<LLVM::ExtractValueOp>( 382 loc, llvmResultType, adaptor.vector(), positionArrayAttr); 383 rewriter.replaceOp(op, extracted); 384 return matchSuccess(); 385 } 386 387 // Potential extraction of 1-D vector from array. 388 auto *context = op->getContext(); 389 Value extracted = adaptor.vector(); 390 auto positionAttrs = positionArrayAttr.getValue(); 391 if (positionAttrs.size() > 1) { 392 auto oneDVectorType = reducedVectorTypeBack(vectorType); 393 auto nMinusOnePositionAttrs = 394 ArrayAttr::get(positionAttrs.drop_back(), context); 395 extracted = rewriter.create<LLVM::ExtractValueOp>( 396 loc, lowering.convertType(oneDVectorType), extracted, 397 nMinusOnePositionAttrs); 398 } 399 400 // Remaining extraction of element from 1-D LLVM vector 401 auto position = positionAttrs.back().cast<IntegerAttr>(); 402 auto i64Type = LLVM::LLVMType::getInt64Ty(lowering.getDialect()); 403 auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 404 extracted = 405 rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant); 406 rewriter.replaceOp(op, extracted); 407 408 return matchSuccess(); 409 } 410 }; 411 412 class VectorInsertElementOpConversion : public LLVMOpLowering { 413 public: 414 explicit VectorInsertElementOpConversion(MLIRContext *context, 415 LLVMTypeConverter &typeConverter) 416 : LLVMOpLowering(vector::InsertElementOp::getOperationName(), context, 417 typeConverter) {} 418 419 PatternMatchResult 420 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 421 ConversionPatternRewriter &rewriter) const override { 422 auto adaptor = vector::InsertElementOpOperandAdaptor(operands); 423 auto insertEltOp = cast<vector::InsertElementOp>(op); 424 auto vectorType = insertEltOp.getDestVectorType(); 425 auto llvmType = lowering.convertType(vectorType); 426 427 // Bail if result type cannot be lowered. 428 if (!llvmType) 429 return matchFailure(); 430 431 rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>( 432 op, llvmType, adaptor.dest(), adaptor.source(), adaptor.position()); 433 return matchSuccess(); 434 } 435 }; 436 437 class VectorInsertOpConversion : public LLVMOpLowering { 438 public: 439 explicit VectorInsertOpConversion(MLIRContext *context, 440 LLVMTypeConverter &typeConverter) 441 : LLVMOpLowering(vector::InsertOp::getOperationName(), context, 442 typeConverter) {} 443 444 PatternMatchResult 445 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 446 ConversionPatternRewriter &rewriter) const override { 447 auto loc = op->getLoc(); 448 auto adaptor = vector::InsertOpOperandAdaptor(operands); 449 auto insertOp = cast<vector::InsertOp>(op); 450 auto sourceType = insertOp.getSourceType(); 451 auto destVectorType = insertOp.getDestVectorType(); 452 auto llvmResultType = lowering.convertType(destVectorType); 453 auto positionArrayAttr = insertOp.position(); 454 455 // Bail if result type cannot be lowered. 456 if (!llvmResultType) 457 return matchFailure(); 458 459 // One-shot insertion of a vector into an array (only requires insertvalue). 460 if (sourceType.isa<VectorType>()) { 461 Value inserted = rewriter.create<LLVM::InsertValueOp>( 462 loc, llvmResultType, adaptor.dest(), adaptor.source(), 463 positionArrayAttr); 464 rewriter.replaceOp(op, inserted); 465 return matchSuccess(); 466 } 467 468 // Potential extraction of 1-D vector from array. 469 auto *context = op->getContext(); 470 Value extracted = adaptor.dest(); 471 auto positionAttrs = positionArrayAttr.getValue(); 472 auto position = positionAttrs.back().cast<IntegerAttr>(); 473 auto oneDVectorType = destVectorType; 474 if (positionAttrs.size() > 1) { 475 oneDVectorType = reducedVectorTypeBack(destVectorType); 476 auto nMinusOnePositionAttrs = 477 ArrayAttr::get(positionAttrs.drop_back(), context); 478 extracted = rewriter.create<LLVM::ExtractValueOp>( 479 loc, lowering.convertType(oneDVectorType), extracted, 480 nMinusOnePositionAttrs); 481 } 482 483 // Insertion of an element into a 1-D LLVM vector. 484 auto i64Type = LLVM::LLVMType::getInt64Ty(lowering.getDialect()); 485 auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 486 Value inserted = rewriter.create<LLVM::InsertElementOp>( 487 loc, lowering.convertType(oneDVectorType), extracted, adaptor.source(), 488 constant); 489 490 // Potential insertion of resulting 1-D vector into array. 491 if (positionAttrs.size() > 1) { 492 auto nMinusOnePositionAttrs = 493 ArrayAttr::get(positionAttrs.drop_back(), context); 494 inserted = rewriter.create<LLVM::InsertValueOp>(loc, llvmResultType, 495 adaptor.dest(), inserted, 496 nMinusOnePositionAttrs); 497 } 498 499 rewriter.replaceOp(op, inserted); 500 return matchSuccess(); 501 } 502 }; 503 504 // When ranks are different, InsertStridedSlice needs to extract a properly 505 // ranked vector from the destination vector into which to insert. This pattern 506 // only takes care of this part and forwards the rest of the conversion to 507 // another pattern that converts InsertStridedSlice for operands of the same 508 // rank. 509 // 510 // RewritePattern for InsertStridedSliceOp where source and destination vectors 511 // have different ranks. In this case: 512 // 1. the proper subvector is extracted from the destination vector 513 // 2. a new InsertStridedSlice op is created to insert the source in the 514 // destination subvector 515 // 3. the destination subvector is inserted back in the proper place 516 // 4. the op is replaced by the result of step 3. 517 // The new InsertStridedSlice from step 2. will be picked up by a 518 // `VectorInsertStridedSliceOpSameRankRewritePattern`. 519 class VectorInsertStridedSliceOpDifferentRankRewritePattern 520 : public OpRewritePattern<InsertStridedSliceOp> { 521 public: 522 using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 523 524 PatternMatchResult matchAndRewrite(InsertStridedSliceOp op, 525 PatternRewriter &rewriter) const override { 526 auto srcType = op.getSourceVectorType(); 527 auto dstType = op.getDestVectorType(); 528 529 if (op.offsets().getValue().empty()) 530 return matchFailure(); 531 532 auto loc = op.getLoc(); 533 int64_t rankDiff = dstType.getRank() - srcType.getRank(); 534 assert(rankDiff >= 0); 535 if (rankDiff == 0) 536 return matchFailure(); 537 538 int64_t rankRest = dstType.getRank() - rankDiff; 539 // Extract / insert the subvector of matching rank and InsertStridedSlice 540 // on it. 541 Value extracted = 542 rewriter.create<ExtractOp>(loc, op.dest(), 543 getI64SubArray(op.offsets(), /*dropFront=*/0, 544 /*dropFront=*/rankRest)); 545 // A different pattern will kick in for InsertStridedSlice with matching 546 // ranks. 547 auto stridedSliceInnerOp = rewriter.create<InsertStridedSliceOp>( 548 loc, op.source(), extracted, 549 getI64SubArray(op.offsets(), /*dropFront=*/rankDiff), 550 getI64SubArray(op.strides(), /*dropFront=*/rankDiff)); 551 rewriter.replaceOpWithNewOp<InsertOp>( 552 op, stridedSliceInnerOp.getResult(), op.dest(), 553 getI64SubArray(op.offsets(), /*dropFront=*/0, 554 /*dropFront=*/rankRest)); 555 return matchSuccess(); 556 } 557 }; 558 559 // RewritePattern for InsertStridedSliceOp where source and destination vectors 560 // have the same rank. In this case, we reduce 561 // 1. the proper subvector is extracted from the destination vector 562 // 2. a new InsertStridedSlice op is created to insert the source in the 563 // destination subvector 564 // 3. the destination subvector is inserted back in the proper place 565 // 4. the op is replaced by the result of step 3. 566 // The new InsertStridedSlice from step 2. will be picked up by a 567 // `VectorInsertStridedSliceOpSameRankRewritePattern`. 568 class VectorInsertStridedSliceOpSameRankRewritePattern 569 : public OpRewritePattern<InsertStridedSliceOp> { 570 public: 571 using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 572 573 PatternMatchResult matchAndRewrite(InsertStridedSliceOp op, 574 PatternRewriter &rewriter) const override { 575 auto srcType = op.getSourceVectorType(); 576 auto dstType = op.getDestVectorType(); 577 578 if (op.offsets().getValue().empty()) 579 return matchFailure(); 580 581 int64_t rankDiff = dstType.getRank() - srcType.getRank(); 582 assert(rankDiff >= 0); 583 if (rankDiff != 0) 584 return matchFailure(); 585 586 if (srcType == dstType) { 587 rewriter.replaceOp(op, op.source()); 588 return matchSuccess(); 589 } 590 591 int64_t offset = 592 op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 593 int64_t size = srcType.getShape().front(); 594 int64_t stride = 595 op.strides().getValue().front().cast<IntegerAttr>().getInt(); 596 597 auto loc = op.getLoc(); 598 Value res = op.dest(); 599 // For each slice of the source vector along the most major dimension. 600 for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 601 off += stride, ++idx) { 602 // 1. extract the proper subvector (or element) from source 603 Value extractedSource = extractOne(rewriter, loc, op.source(), idx); 604 if (extractedSource.getType().isa<VectorType>()) { 605 // 2. If we have a vector, extract the proper subvector from destination 606 // Otherwise we are at the element level and no need to recurse. 607 Value extractedDest = extractOne(rewriter, loc, op.dest(), off); 608 // 3. Reduce the problem to lowering a new InsertStridedSlice op with 609 // smaller rank. 610 InsertStridedSliceOp insertStridedSliceOp = 611 rewriter.create<InsertStridedSliceOp>( 612 loc, extractedSource, extractedDest, 613 getI64SubArray(op.offsets(), /* dropFront=*/1), 614 getI64SubArray(op.strides(), /* dropFront=*/1)); 615 // Call matchAndRewrite recursively from within the pattern. This 616 // circumvents the current limitation that a given pattern cannot 617 // be called multiple times by the PatternRewrite infrastructure (to 618 // avoid infinite recursion, but in this case, infinite recursion 619 // cannot happen because the rank is strictly decreasing). 620 // TODO(rriddle, nicolasvasilache) Implement something like a hook for 621 // a potential function that must decrease and allow the same pattern 622 // multiple times. 623 auto success = matchAndRewrite(insertStridedSliceOp, rewriter); 624 (void)success; 625 assert(success && "Unexpected failure"); 626 extractedSource = insertStridedSliceOp; 627 } 628 // 4. Insert the extractedSource into the res vector. 629 res = insertOne(rewriter, loc, extractedSource, res, off); 630 } 631 632 rewriter.replaceOp(op, res); 633 return matchSuccess(); 634 } 635 }; 636 637 class VectorOuterProductOpConversion : public LLVMOpLowering { 638 public: 639 explicit VectorOuterProductOpConversion(MLIRContext *context, 640 LLVMTypeConverter &typeConverter) 641 : LLVMOpLowering(vector::OuterProductOp::getOperationName(), context, 642 typeConverter) {} 643 644 PatternMatchResult 645 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 646 ConversionPatternRewriter &rewriter) const override { 647 auto loc = op->getLoc(); 648 auto adaptor = vector::OuterProductOpOperandAdaptor(operands); 649 auto *ctx = op->getContext(); 650 auto vLHS = adaptor.lhs()->getType().cast<LLVM::LLVMType>(); 651 auto vRHS = adaptor.rhs()->getType().cast<LLVM::LLVMType>(); 652 auto rankLHS = vLHS.getUnderlyingType()->getVectorNumElements(); 653 auto rankRHS = vRHS.getUnderlyingType()->getVectorNumElements(); 654 auto llvmArrayOfVectType = lowering.convertType( 655 cast<vector::OuterProductOp>(op).getResult()->getType()); 656 Value desc = rewriter.create<LLVM::UndefOp>(loc, llvmArrayOfVectType); 657 Value a = adaptor.lhs(), b = adaptor.rhs(); 658 Value acc = adaptor.acc().empty() ? nullptr : adaptor.acc().front(); 659 SmallVector<Value, 8> lhs, accs; 660 lhs.reserve(rankLHS); 661 accs.reserve(rankLHS); 662 for (unsigned d = 0, e = rankLHS; d < e; ++d) { 663 // shufflevector explicitly requires i32. 664 auto attr = rewriter.getI32IntegerAttr(d); 665 SmallVector<Attribute, 4> bcastAttr(rankRHS, attr); 666 auto bcastArrayAttr = ArrayAttr::get(bcastAttr, ctx); 667 Value aD = nullptr, accD = nullptr; 668 // 1. Broadcast the element a[d] into vector aD. 669 aD = rewriter.create<LLVM::ShuffleVectorOp>(loc, a, a, bcastArrayAttr); 670 // 2. If acc is present, extract 1-d vector acc[d] into accD. 671 if (acc) 672 accD = rewriter.create<LLVM::ExtractValueOp>( 673 loc, vRHS, acc, rewriter.getI64ArrayAttr(d)); 674 // 3. Compute aD outer b (plus accD, if relevant). 675 Value aOuterbD = 676 accD ? rewriter.create<LLVM::FMulAddOp>(loc, vRHS, aD, b, accD) 677 .getResult() 678 : rewriter.create<LLVM::FMulOp>(loc, aD, b).getResult(); 679 // 4. Insert as value `d` in the descriptor. 680 desc = rewriter.create<LLVM::InsertValueOp>(loc, llvmArrayOfVectType, 681 desc, aOuterbD, 682 rewriter.getI64ArrayAttr(d)); 683 } 684 rewriter.replaceOp(op, desc); 685 return matchSuccess(); 686 } 687 }; 688 689 class VectorTypeCastOpConversion : public LLVMOpLowering { 690 public: 691 explicit VectorTypeCastOpConversion(MLIRContext *context, 692 LLVMTypeConverter &typeConverter) 693 : LLVMOpLowering(vector::TypeCastOp::getOperationName(), context, 694 typeConverter) {} 695 696 PatternMatchResult 697 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 698 ConversionPatternRewriter &rewriter) const override { 699 auto loc = op->getLoc(); 700 vector::TypeCastOp castOp = cast<vector::TypeCastOp>(op); 701 MemRefType sourceMemRefType = 702 castOp.getOperand()->getType().cast<MemRefType>(); 703 MemRefType targetMemRefType = 704 castOp.getResult()->getType().cast<MemRefType>(); 705 706 // Only static shape casts supported atm. 707 if (!sourceMemRefType.hasStaticShape() || 708 !targetMemRefType.hasStaticShape()) 709 return matchFailure(); 710 711 auto llvmSourceDescriptorTy = 712 operands[0]->getType().dyn_cast<LLVM::LLVMType>(); 713 if (!llvmSourceDescriptorTy || !llvmSourceDescriptorTy.isStructTy()) 714 return matchFailure(); 715 MemRefDescriptor sourceMemRef(operands[0]); 716 717 auto llvmTargetDescriptorTy = lowering.convertType(targetMemRefType) 718 .dyn_cast_or_null<LLVM::LLVMType>(); 719 if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy()) 720 return matchFailure(); 721 722 int64_t offset; 723 SmallVector<int64_t, 4> strides; 724 auto successStrides = 725 getStridesAndOffset(sourceMemRefType, strides, offset); 726 bool isContiguous = (strides.back() == 1); 727 if (isContiguous) { 728 auto sizes = sourceMemRefType.getShape(); 729 for (int index = 0, e = strides.size() - 2; index < e; ++index) { 730 if (strides[index] != strides[index + 1] * sizes[index + 1]) { 731 isContiguous = false; 732 break; 733 } 734 } 735 } 736 // Only contiguous source tensors supported atm. 737 if (failed(successStrides) || !isContiguous) 738 return matchFailure(); 739 740 auto int64Ty = LLVM::LLVMType::getInt64Ty(lowering.getDialect()); 741 742 // Create descriptor. 743 auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy); 744 Type llvmTargetElementTy = desc.getElementType(); 745 // Set allocated ptr. 746 Value allocated = sourceMemRef.allocatedPtr(rewriter, loc); 747 allocated = 748 rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated); 749 desc.setAllocatedPtr(rewriter, loc, allocated); 750 // Set aligned ptr. 751 Value ptr = sourceMemRef.alignedPtr(rewriter, loc); 752 ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr); 753 desc.setAlignedPtr(rewriter, loc, ptr); 754 // Fill offset 0. 755 auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0); 756 auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr); 757 desc.setOffset(rewriter, loc, zero); 758 759 // Fill size and stride descriptors in memref. 760 for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) { 761 int64_t index = indexedSize.index(); 762 auto sizeAttr = 763 rewriter.getIntegerAttr(rewriter.getIndexType(), indexedSize.value()); 764 auto size = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, sizeAttr); 765 desc.setSize(rewriter, loc, index, size); 766 auto strideAttr = 767 rewriter.getIntegerAttr(rewriter.getIndexType(), strides[index]); 768 auto stride = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, strideAttr); 769 desc.setStride(rewriter, loc, index, stride); 770 } 771 772 rewriter.replaceOp(op, {desc}); 773 return matchSuccess(); 774 } 775 }; 776 777 class VectorPrintOpConversion : public LLVMOpLowering { 778 public: 779 explicit VectorPrintOpConversion(MLIRContext *context, 780 LLVMTypeConverter &typeConverter) 781 : LLVMOpLowering(vector::PrintOp::getOperationName(), context, 782 typeConverter) {} 783 784 // Proof-of-concept lowering implementation that relies on a small 785 // runtime support library, which only needs to provide a few 786 // printing methods (single value for all data types, opening/closing 787 // bracket, comma, newline). The lowering fully unrolls a vector 788 // in terms of these elementary printing operations. The advantage 789 // of this approach is that the library can remain unaware of all 790 // low-level implementation details of vectors while still supporting 791 // output of any shaped and dimensioned vector. Due to full unrolling, 792 // this approach is less suited for very large vectors though. 793 // 794 // TODO(ajcbik): rely solely on libc in future? something else? 795 // 796 PatternMatchResult 797 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 798 ConversionPatternRewriter &rewriter) const override { 799 auto printOp = cast<vector::PrintOp>(op); 800 auto adaptor = vector::PrintOpOperandAdaptor(operands); 801 Type printType = printOp.getPrintType(); 802 803 if (lowering.convertType(printType) == nullptr) 804 return matchFailure(); 805 806 // Make sure element type has runtime support (currently just Float/Double). 807 VectorType vectorType = printType.dyn_cast<VectorType>(); 808 Type eltType = vectorType ? vectorType.getElementType() : printType; 809 int64_t rank = vectorType ? vectorType.getRank() : 0; 810 Operation *printer; 811 if (eltType.isF32()) 812 printer = getPrintFloat(op); 813 else if (eltType.isF64()) 814 printer = getPrintDouble(op); 815 else 816 return matchFailure(); 817 818 // Unroll vector into elementary print calls. 819 emitRanks(rewriter, op, adaptor.source(), vectorType, printer, rank); 820 emitCall(rewriter, op->getLoc(), getPrintNewline(op)); 821 rewriter.eraseOp(op); 822 return matchSuccess(); 823 } 824 825 private: 826 void emitRanks(ConversionPatternRewriter &rewriter, Operation *op, 827 Value value, VectorType vectorType, Operation *printer, 828 int64_t rank) const { 829 Location loc = op->getLoc(); 830 if (rank == 0) { 831 emitCall(rewriter, loc, printer, value); 832 return; 833 } 834 835 emitCall(rewriter, loc, getPrintOpen(op)); 836 Operation *printComma = getPrintComma(op); 837 int64_t dim = vectorType.getDimSize(0); 838 for (int64_t d = 0; d < dim; ++d) { 839 auto reducedType = 840 rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr; 841 auto llvmType = lowering.convertType( 842 rank > 1 ? reducedType : vectorType.getElementType()); 843 Value nestedVal = 844 extractOne(rewriter, lowering, loc, value, llvmType, rank, d); 845 emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1); 846 if (d != dim - 1) 847 emitCall(rewriter, loc, printComma); 848 } 849 emitCall(rewriter, loc, getPrintClose(op)); 850 } 851 852 // Helper to emit a call. 853 static void emitCall(ConversionPatternRewriter &rewriter, Location loc, 854 Operation *ref, ValueRange params = ValueRange()) { 855 rewriter.create<LLVM::CallOp>(loc, ArrayRef<Type>{}, 856 rewriter.getSymbolRefAttr(ref), params); 857 } 858 859 // Helper for printer method declaration (first hit) and lookup. 860 static Operation *getPrint(Operation *op, LLVM::LLVMDialect *dialect, 861 StringRef name, ArrayRef<LLVM::LLVMType> params) { 862 auto module = op->getParentOfType<ModuleOp>(); 863 auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(name); 864 if (func) 865 return func; 866 OpBuilder moduleBuilder(module.getBodyRegion()); 867 return moduleBuilder.create<LLVM::LLVMFuncOp>( 868 op->getLoc(), name, 869 LLVM::LLVMType::getFunctionTy(LLVM::LLVMType::getVoidTy(dialect), 870 params, /*isVarArg=*/false)); 871 } 872 873 // Helpers for method names. 874 Operation *getPrintFloat(Operation *op) const { 875 LLVM::LLVMDialect *dialect = lowering.getDialect(); 876 return getPrint(op, dialect, "print_f32", 877 LLVM::LLVMType::getFloatTy(dialect)); 878 } 879 Operation *getPrintDouble(Operation *op) const { 880 LLVM::LLVMDialect *dialect = lowering.getDialect(); 881 return getPrint(op, dialect, "print_f64", 882 LLVM::LLVMType::getDoubleTy(dialect)); 883 } 884 Operation *getPrintOpen(Operation *op) const { 885 return getPrint(op, lowering.getDialect(), "print_open", {}); 886 } 887 Operation *getPrintClose(Operation *op) const { 888 return getPrint(op, lowering.getDialect(), "print_close", {}); 889 } 890 Operation *getPrintComma(Operation *op) const { 891 return getPrint(op, lowering.getDialect(), "print_comma", {}); 892 } 893 Operation *getPrintNewline(Operation *op) const { 894 return getPrint(op, lowering.getDialect(), "print_newline", {}); 895 } 896 }; 897 898 /// Progressive lowering of StridedSliceOp to either: 899 /// 1. extractelement + insertelement for the 1-D case 900 /// 2. extract + optional strided_slice + insert for the n-D case. 901 class VectorStridedSliceOpConversion : public OpRewritePattern<StridedSliceOp> { 902 public: 903 using OpRewritePattern<StridedSliceOp>::OpRewritePattern; 904 905 PatternMatchResult matchAndRewrite(StridedSliceOp op, 906 PatternRewriter &rewriter) const override { 907 auto dstType = op.getResult().getType().cast<VectorType>(); 908 909 assert(!op.offsets().getValue().empty() && "Unexpected empty offsets"); 910 911 int64_t offset = 912 op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 913 int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt(); 914 int64_t stride = 915 op.strides().getValue().front().cast<IntegerAttr>().getInt(); 916 917 auto loc = op.getLoc(); 918 auto elemType = dstType.getElementType(); 919 assert(elemType.isIntOrIndexOrFloat()); 920 Value zero = rewriter.create<ConstantOp>(loc, elemType, 921 rewriter.getZeroAttr(elemType)); 922 Value res = rewriter.create<SplatOp>(loc, dstType, zero); 923 for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 924 off += stride, ++idx) { 925 Value extracted = extractOne(rewriter, loc, op.vector(), off); 926 if (op.offsets().getValue().size() > 1) { 927 StridedSliceOp stridedSliceOp = rewriter.create<StridedSliceOp>( 928 loc, extracted, getI64SubArray(op.offsets(), /* dropFront=*/1), 929 getI64SubArray(op.sizes(), /* dropFront=*/1), 930 getI64SubArray(op.strides(), /* dropFront=*/1)); 931 // Call matchAndRewrite recursively from within the pattern. This 932 // circumvents the current limitation that a given pattern cannot 933 // be called multiple times by the PatternRewrite infrastructure (to 934 // avoid infinite recursion, but in this case, infinite recursion 935 // cannot happen because the rank is strictly decreasing). 936 // TODO(rriddle, nicolasvasilache) Implement something like a hook for 937 // a potential function that must decrease and allow the same pattern 938 // multiple times. 939 auto success = matchAndRewrite(stridedSliceOp, rewriter); 940 (void)success; 941 assert(success && "Unexpected failure"); 942 extracted = stridedSliceOp; 943 } 944 res = insertOne(rewriter, loc, extracted, res, idx); 945 } 946 rewriter.replaceOp(op, {res}); 947 return matchSuccess(); 948 } 949 }; 950 951 /// Populate the given list with patterns that convert from Vector to LLVM. 952 void mlir::populateVectorToLLVMConversionPatterns( 953 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 954 MLIRContext *ctx = converter.getDialect()->getContext(); 955 patterns.insert<VectorInsertStridedSliceOpDifferentRankRewritePattern, 956 VectorInsertStridedSliceOpSameRankRewritePattern, 957 VectorStridedSliceOpConversion>(ctx); 958 patterns.insert<VectorBroadcastOpConversion, VectorShuffleOpConversion, 959 VectorExtractElementOpConversion, VectorExtractOpConversion, 960 VectorInsertElementOpConversion, VectorInsertOpConversion, 961 VectorOuterProductOpConversion, VectorTypeCastOpConversion, 962 VectorPrintOpConversion>(ctx, converter); 963 } 964 965 namespace { 966 struct LowerVectorToLLVMPass : public ModulePass<LowerVectorToLLVMPass> { 967 void runOnModule() override; 968 }; 969 } // namespace 970 971 void LowerVectorToLLVMPass::runOnModule() { 972 // Convert to the LLVM IR dialect using the converter defined above. 973 OwningRewritePatternList patterns; 974 LLVMTypeConverter converter(&getContext()); 975 populateVectorToLLVMConversionPatterns(converter, patterns); 976 populateStdToLLVMConversionPatterns(converter, patterns); 977 978 ConversionTarget target(getContext()); 979 target.addLegalDialect<LLVM::LLVMDialect>(); 980 target.addDynamicallyLegalOp<FuncOp>( 981 [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); 982 if (failed( 983 applyPartialConversion(getModule(), target, patterns, &converter))) { 984 signalPassFailure(); 985 } 986 } 987 988 OpPassBase<ModuleOp> *mlir::createLowerVectorToLLVMPass() { 989 return new LowerVectorToLLVMPass(); 990 } 991 992 static PassRegistration<LowerVectorToLLVMPass> 993 pass("convert-vector-to-llvm", 994 "Lower the operations from the vector dialect into the LLVM dialect"); 995