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