1 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===// 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/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" 10 11 #include "../PassDetail.h" 12 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h" 13 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" 14 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 15 #include "mlir/Dialect/StandardOps/IR/Ops.h" 16 #include "mlir/Dialect/Vector/VectorOps.h" 17 #include "mlir/IR/Attributes.h" 18 #include "mlir/IR/Builders.h" 19 #include "mlir/IR/MLIRContext.h" 20 #include "mlir/IR/Module.h" 21 #include "mlir/IR/Operation.h" 22 #include "mlir/IR/PatternMatch.h" 23 #include "mlir/IR/StandardTypes.h" 24 #include "mlir/IR/Types.h" 25 #include "mlir/Transforms/DialectConversion.h" 26 #include "mlir/Transforms/Passes.h" 27 #include "llvm/IR/DerivedTypes.h" 28 #include "llvm/IR/Module.h" 29 #include "llvm/IR/Type.h" 30 #include "llvm/Support/Allocator.h" 31 #include "llvm/Support/ErrorHandling.h" 32 33 using namespace mlir; 34 using namespace mlir::vector; 35 36 template <typename T> 37 static LLVM::LLVMType getPtrToElementType(T containerType, 38 LLVMTypeConverter &typeConverter) { 39 return typeConverter.convertType(containerType.getElementType()) 40 .template cast<LLVM::LLVMType>() 41 .getPointerTo(); 42 } 43 44 // Helper to reduce vector type by one rank at front. 45 static VectorType reducedVectorTypeFront(VectorType tp) { 46 assert((tp.getRank() > 1) && "unlowerable vector type"); 47 return VectorType::get(tp.getShape().drop_front(), tp.getElementType()); 48 } 49 50 // Helper to reduce vector type by *all* but one rank at back. 51 static VectorType reducedVectorTypeBack(VectorType tp) { 52 assert((tp.getRank() > 1) && "unlowerable vector type"); 53 return VectorType::get(tp.getShape().take_back(), tp.getElementType()); 54 } 55 56 // Helper that picks the proper sequence for inserting. 57 static Value insertOne(ConversionPatternRewriter &rewriter, 58 LLVMTypeConverter &typeConverter, Location loc, 59 Value val1, Value val2, Type llvmType, int64_t rank, 60 int64_t pos) { 61 if (rank == 1) { 62 auto idxType = rewriter.getIndexType(); 63 auto constant = rewriter.create<LLVM::ConstantOp>( 64 loc, typeConverter.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 &typeConverter, Location loc, 87 Value val, 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, typeConverter.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 namespace { 127 128 class VectorBroadcastOpConversion : public ConvertToLLVMPattern { 129 public: 130 explicit VectorBroadcastOpConversion(MLIRContext *context, 131 LLVMTypeConverter &typeConverter) 132 : ConvertToLLVMPattern(vector::BroadcastOp::getOperationName(), context, 133 typeConverter) {} 134 135 LogicalResult 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 (typeConverter.convertType(dstVectorType) == nullptr) 141 return failure(); 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 success(); 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 implementation. 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 = typeConverter.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 = insertOne(rewriter, typeConverter, loc, undef, value, 211 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 = insertOne(rewriter, typeConverter, loc, result, expand, llvmType, 221 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 = typeConverter.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 = 253 typeConverter.convertType(dstVectorType.getElementType()); 254 Value one = 255 extractOne(rewriter, typeConverter, loc, value, redLlvmType, rank, 0); 256 Value expand = insertOne(rewriter, typeConverter, loc, result, one, 257 llvmType, rank, 0); 258 SmallVector<int32_t, 4> zeroValues(dim, 0); 259 return rewriter.create<LLVM::ShuffleVectorOp>( 260 loc, expand, result, rewriter.getI32ArrayAttr(zeroValues)); 261 } 262 VectorType redSrcType = reducedVectorTypeFront(srcVectorType); 263 VectorType redDstType = reducedVectorTypeFront(dstVectorType); 264 Type redLlvmType = typeConverter.convertType(redSrcType); 265 for (int64_t d = 0; d < dim; ++d) { 266 int64_t pos = atStretch ? 0 : d; 267 Value one = extractOne(rewriter, typeConverter, loc, value, redLlvmType, 268 rank, pos); 269 Value expand = expandRanks(one, loc, redSrcType, redDstType, rewriter); 270 result = insertOne(rewriter, typeConverter, loc, result, expand, llvmType, 271 rank, d); 272 } 273 return result; 274 } 275 }; 276 277 /// Conversion pattern for a vector.matrix_multiply. 278 /// This is lowered directly to the proper llvm.intr.matrix.multiply. 279 class VectorMatmulOpConversion : public ConvertToLLVMPattern { 280 public: 281 explicit VectorMatmulOpConversion(MLIRContext *context, 282 LLVMTypeConverter &typeConverter) 283 : ConvertToLLVMPattern(vector::MatmulOp::getOperationName(), context, 284 typeConverter) {} 285 286 LogicalResult 287 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 288 ConversionPatternRewriter &rewriter) const override { 289 auto matmulOp = cast<vector::MatmulOp>(op); 290 auto adaptor = vector::MatmulOpOperandAdaptor(operands); 291 rewriter.replaceOpWithNewOp<LLVM::MatrixMultiplyOp>( 292 op, typeConverter.convertType(matmulOp.res().getType()), adaptor.lhs(), 293 adaptor.rhs(), matmulOp.lhs_rows(), matmulOp.lhs_columns(), 294 matmulOp.rhs_columns()); 295 return success(); 296 } 297 }; 298 299 class VectorReductionOpConversion : public ConvertToLLVMPattern { 300 public: 301 explicit VectorReductionOpConversion(MLIRContext *context, 302 LLVMTypeConverter &typeConverter) 303 : ConvertToLLVMPattern(vector::ReductionOp::getOperationName(), context, 304 typeConverter) {} 305 306 LogicalResult 307 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 308 ConversionPatternRewriter &rewriter) const override { 309 auto reductionOp = cast<vector::ReductionOp>(op); 310 auto kind = reductionOp.kind(); 311 Type eltType = reductionOp.dest().getType(); 312 Type llvmType = typeConverter.convertType(eltType); 313 if (eltType.isSignlessInteger(32) || eltType.isSignlessInteger(64)) { 314 // Integer reductions: add/mul/min/max/and/or/xor. 315 if (kind == "add") 316 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_add>( 317 op, llvmType, operands[0]); 318 else if (kind == "mul") 319 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_mul>( 320 op, llvmType, operands[0]); 321 else if (kind == "min") 322 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smin>( 323 op, llvmType, operands[0]); 324 else if (kind == "max") 325 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smax>( 326 op, llvmType, operands[0]); 327 else if (kind == "and") 328 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_and>( 329 op, llvmType, operands[0]); 330 else if (kind == "or") 331 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_or>( 332 op, llvmType, operands[0]); 333 else if (kind == "xor") 334 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_xor>( 335 op, llvmType, operands[0]); 336 else 337 return failure(); 338 return success(); 339 340 } else if (eltType.isF32() || eltType.isF64()) { 341 // Floating-point reductions: add/mul/min/max 342 if (kind == "add") { 343 // Optional accumulator (or zero). 344 Value acc = operands.size() > 1 ? operands[1] 345 : rewriter.create<LLVM::ConstantOp>( 346 op->getLoc(), llvmType, 347 rewriter.getZeroAttr(eltType)); 348 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>( 349 op, llvmType, acc, operands[0]); 350 } else if (kind == "mul") { 351 // Optional accumulator (or one). 352 Value acc = operands.size() > 1 353 ? operands[1] 354 : rewriter.create<LLVM::ConstantOp>( 355 op->getLoc(), llvmType, 356 rewriter.getFloatAttr(eltType, 1.0)); 357 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>( 358 op, llvmType, acc, operands[0]); 359 } else if (kind == "min") 360 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmin>( 361 op, llvmType, operands[0]); 362 else if (kind == "max") 363 rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmax>( 364 op, llvmType, operands[0]); 365 else 366 return failure(); 367 return success(); 368 } 369 return failure(); 370 } 371 }; 372 373 class VectorShuffleOpConversion : public ConvertToLLVMPattern { 374 public: 375 explicit VectorShuffleOpConversion(MLIRContext *context, 376 LLVMTypeConverter &typeConverter) 377 : ConvertToLLVMPattern(vector::ShuffleOp::getOperationName(), context, 378 typeConverter) {} 379 380 LogicalResult 381 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 382 ConversionPatternRewriter &rewriter) const override { 383 auto loc = op->getLoc(); 384 auto adaptor = vector::ShuffleOpOperandAdaptor(operands); 385 auto shuffleOp = cast<vector::ShuffleOp>(op); 386 auto v1Type = shuffleOp.getV1VectorType(); 387 auto v2Type = shuffleOp.getV2VectorType(); 388 auto vectorType = shuffleOp.getVectorType(); 389 Type llvmType = typeConverter.convertType(vectorType); 390 auto maskArrayAttr = shuffleOp.mask(); 391 392 // Bail if result type cannot be lowered. 393 if (!llvmType) 394 return failure(); 395 396 // Get rank and dimension sizes. 397 int64_t rank = vectorType.getRank(); 398 assert(v1Type.getRank() == rank); 399 assert(v2Type.getRank() == rank); 400 int64_t v1Dim = v1Type.getDimSize(0); 401 402 // For rank 1, where both operands have *exactly* the same vector type, 403 // there is direct shuffle support in LLVM. Use it! 404 if (rank == 1 && v1Type == v2Type) { 405 Value shuffle = rewriter.create<LLVM::ShuffleVectorOp>( 406 loc, adaptor.v1(), adaptor.v2(), maskArrayAttr); 407 rewriter.replaceOp(op, shuffle); 408 return success(); 409 } 410 411 // For all other cases, insert the individual values individually. 412 Value insert = rewriter.create<LLVM::UndefOp>(loc, llvmType); 413 int64_t insPos = 0; 414 for (auto en : llvm::enumerate(maskArrayAttr)) { 415 int64_t extPos = en.value().cast<IntegerAttr>().getInt(); 416 Value value = adaptor.v1(); 417 if (extPos >= v1Dim) { 418 extPos -= v1Dim; 419 value = adaptor.v2(); 420 } 421 Value extract = extractOne(rewriter, typeConverter, loc, value, llvmType, 422 rank, extPos); 423 insert = insertOne(rewriter, typeConverter, loc, insert, extract, 424 llvmType, rank, insPos++); 425 } 426 rewriter.replaceOp(op, insert); 427 return success(); 428 } 429 }; 430 431 class VectorExtractElementOpConversion : public ConvertToLLVMPattern { 432 public: 433 explicit VectorExtractElementOpConversion(MLIRContext *context, 434 LLVMTypeConverter &typeConverter) 435 : ConvertToLLVMPattern(vector::ExtractElementOp::getOperationName(), 436 context, typeConverter) {} 437 438 LogicalResult 439 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 440 ConversionPatternRewriter &rewriter) const override { 441 auto adaptor = vector::ExtractElementOpOperandAdaptor(operands); 442 auto extractEltOp = cast<vector::ExtractElementOp>(op); 443 auto vectorType = extractEltOp.getVectorType(); 444 auto llvmType = typeConverter.convertType(vectorType.getElementType()); 445 446 // Bail if result type cannot be lowered. 447 if (!llvmType) 448 return failure(); 449 450 rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>( 451 op, llvmType, adaptor.vector(), adaptor.position()); 452 return success(); 453 } 454 }; 455 456 class VectorExtractOpConversion : public ConvertToLLVMPattern { 457 public: 458 explicit VectorExtractOpConversion(MLIRContext *context, 459 LLVMTypeConverter &typeConverter) 460 : ConvertToLLVMPattern(vector::ExtractOp::getOperationName(), context, 461 typeConverter) {} 462 463 LogicalResult 464 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 465 ConversionPatternRewriter &rewriter) const override { 466 auto loc = op->getLoc(); 467 auto adaptor = vector::ExtractOpOperandAdaptor(operands); 468 auto extractOp = cast<vector::ExtractOp>(op); 469 auto vectorType = extractOp.getVectorType(); 470 auto resultType = extractOp.getResult().getType(); 471 auto llvmResultType = typeConverter.convertType(resultType); 472 auto positionArrayAttr = extractOp.position(); 473 474 // Bail if result type cannot be lowered. 475 if (!llvmResultType) 476 return failure(); 477 478 // One-shot extraction of vector from array (only requires extractvalue). 479 if (resultType.isa<VectorType>()) { 480 Value extracted = rewriter.create<LLVM::ExtractValueOp>( 481 loc, llvmResultType, adaptor.vector(), positionArrayAttr); 482 rewriter.replaceOp(op, extracted); 483 return success(); 484 } 485 486 // Potential extraction of 1-D vector from array. 487 auto *context = op->getContext(); 488 Value extracted = adaptor.vector(); 489 auto positionAttrs = positionArrayAttr.getValue(); 490 if (positionAttrs.size() > 1) { 491 auto oneDVectorType = reducedVectorTypeBack(vectorType); 492 auto nMinusOnePositionAttrs = 493 ArrayAttr::get(positionAttrs.drop_back(), context); 494 extracted = rewriter.create<LLVM::ExtractValueOp>( 495 loc, typeConverter.convertType(oneDVectorType), extracted, 496 nMinusOnePositionAttrs); 497 } 498 499 // Remaining extraction of element from 1-D LLVM vector 500 auto position = positionAttrs.back().cast<IntegerAttr>(); 501 auto i64Type = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 502 auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 503 extracted = 504 rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant); 505 rewriter.replaceOp(op, extracted); 506 507 return success(); 508 } 509 }; 510 511 /// Conversion pattern that turns a vector.fma on a 1-D vector 512 /// into an llvm.intr.fmuladd. This is a trivial 1-1 conversion. 513 /// This does not match vectors of n >= 2 rank. 514 /// 515 /// Example: 516 /// ``` 517 /// vector.fma %a, %a, %a : vector<8xf32> 518 /// ``` 519 /// is converted to: 520 /// ``` 521 /// llvm.intr.fma %va, %va, %va: 522 /// (!llvm<"<8 x float>">, !llvm<"<8 x float>">, !llvm<"<8 x float>">) 523 /// -> !llvm<"<8 x float>"> 524 /// ``` 525 class VectorFMAOp1DConversion : public ConvertToLLVMPattern { 526 public: 527 explicit VectorFMAOp1DConversion(MLIRContext *context, 528 LLVMTypeConverter &typeConverter) 529 : ConvertToLLVMPattern(vector::FMAOp::getOperationName(), context, 530 typeConverter) {} 531 532 LogicalResult 533 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 534 ConversionPatternRewriter &rewriter) const override { 535 auto adaptor = vector::FMAOpOperandAdaptor(operands); 536 vector::FMAOp fmaOp = cast<vector::FMAOp>(op); 537 VectorType vType = fmaOp.getVectorType(); 538 if (vType.getRank() != 1) 539 return failure(); 540 rewriter.replaceOpWithNewOp<LLVM::FMAOp>(op, adaptor.lhs(), adaptor.rhs(), 541 adaptor.acc()); 542 return success(); 543 } 544 }; 545 546 class VectorInsertElementOpConversion : public ConvertToLLVMPattern { 547 public: 548 explicit VectorInsertElementOpConversion(MLIRContext *context, 549 LLVMTypeConverter &typeConverter) 550 : ConvertToLLVMPattern(vector::InsertElementOp::getOperationName(), 551 context, typeConverter) {} 552 553 LogicalResult 554 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 555 ConversionPatternRewriter &rewriter) const override { 556 auto adaptor = vector::InsertElementOpOperandAdaptor(operands); 557 auto insertEltOp = cast<vector::InsertElementOp>(op); 558 auto vectorType = insertEltOp.getDestVectorType(); 559 auto llvmType = typeConverter.convertType(vectorType); 560 561 // Bail if result type cannot be lowered. 562 if (!llvmType) 563 return failure(); 564 565 rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>( 566 op, llvmType, adaptor.dest(), adaptor.source(), adaptor.position()); 567 return success(); 568 } 569 }; 570 571 class VectorInsertOpConversion : public ConvertToLLVMPattern { 572 public: 573 explicit VectorInsertOpConversion(MLIRContext *context, 574 LLVMTypeConverter &typeConverter) 575 : ConvertToLLVMPattern(vector::InsertOp::getOperationName(), context, 576 typeConverter) {} 577 578 LogicalResult 579 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 580 ConversionPatternRewriter &rewriter) const override { 581 auto loc = op->getLoc(); 582 auto adaptor = vector::InsertOpOperandAdaptor(operands); 583 auto insertOp = cast<vector::InsertOp>(op); 584 auto sourceType = insertOp.getSourceType(); 585 auto destVectorType = insertOp.getDestVectorType(); 586 auto llvmResultType = typeConverter.convertType(destVectorType); 587 auto positionArrayAttr = insertOp.position(); 588 589 // Bail if result type cannot be lowered. 590 if (!llvmResultType) 591 return failure(); 592 593 // One-shot insertion of a vector into an array (only requires insertvalue). 594 if (sourceType.isa<VectorType>()) { 595 Value inserted = rewriter.create<LLVM::InsertValueOp>( 596 loc, llvmResultType, adaptor.dest(), adaptor.source(), 597 positionArrayAttr); 598 rewriter.replaceOp(op, inserted); 599 return success(); 600 } 601 602 // Potential extraction of 1-D vector from array. 603 auto *context = op->getContext(); 604 Value extracted = adaptor.dest(); 605 auto positionAttrs = positionArrayAttr.getValue(); 606 auto position = positionAttrs.back().cast<IntegerAttr>(); 607 auto oneDVectorType = destVectorType; 608 if (positionAttrs.size() > 1) { 609 oneDVectorType = reducedVectorTypeBack(destVectorType); 610 auto nMinusOnePositionAttrs = 611 ArrayAttr::get(positionAttrs.drop_back(), context); 612 extracted = rewriter.create<LLVM::ExtractValueOp>( 613 loc, typeConverter.convertType(oneDVectorType), extracted, 614 nMinusOnePositionAttrs); 615 } 616 617 // Insertion of an element into a 1-D LLVM vector. 618 auto i64Type = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 619 auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 620 Value inserted = rewriter.create<LLVM::InsertElementOp>( 621 loc, typeConverter.convertType(oneDVectorType), extracted, 622 adaptor.source(), constant); 623 624 // Potential insertion of resulting 1-D vector into array. 625 if (positionAttrs.size() > 1) { 626 auto nMinusOnePositionAttrs = 627 ArrayAttr::get(positionAttrs.drop_back(), context); 628 inserted = rewriter.create<LLVM::InsertValueOp>(loc, llvmResultType, 629 adaptor.dest(), inserted, 630 nMinusOnePositionAttrs); 631 } 632 633 rewriter.replaceOp(op, inserted); 634 return success(); 635 } 636 }; 637 638 /// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1. 639 /// 640 /// Example: 641 /// ``` 642 /// %d = vector.fma %a, %b, %c : vector<2x4xf32> 643 /// ``` 644 /// is rewritten into: 645 /// ``` 646 /// %r = splat %f0: vector<2x4xf32> 647 /// %va = vector.extractvalue %a[0] : vector<2x4xf32> 648 /// %vb = vector.extractvalue %b[0] : vector<2x4xf32> 649 /// %vc = vector.extractvalue %c[0] : vector<2x4xf32> 650 /// %vd = vector.fma %va, %vb, %vc : vector<4xf32> 651 /// %r2 = vector.insertvalue %vd, %r[0] : vector<4xf32> into vector<2x4xf32> 652 /// %va2 = vector.extractvalue %a2[1] : vector<2x4xf32> 653 /// %vb2 = vector.extractvalue %b2[1] : vector<2x4xf32> 654 /// %vc2 = vector.extractvalue %c2[1] : vector<2x4xf32> 655 /// %vd2 = vector.fma %va2, %vb2, %vc2 : vector<4xf32> 656 /// %r3 = vector.insertvalue %vd2, %r2[1] : vector<4xf32> into vector<2x4xf32> 657 /// // %r3 holds the final value. 658 /// ``` 659 class VectorFMAOpNDRewritePattern : public OpRewritePattern<FMAOp> { 660 public: 661 using OpRewritePattern<FMAOp>::OpRewritePattern; 662 663 LogicalResult matchAndRewrite(FMAOp op, 664 PatternRewriter &rewriter) const override { 665 auto vType = op.getVectorType(); 666 if (vType.getRank() < 2) 667 return failure(); 668 669 auto loc = op.getLoc(); 670 auto elemType = vType.getElementType(); 671 Value zero = rewriter.create<ConstantOp>(loc, elemType, 672 rewriter.getZeroAttr(elemType)); 673 Value desc = rewriter.create<SplatOp>(loc, vType, zero); 674 for (int64_t i = 0, e = vType.getShape().front(); i != e; ++i) { 675 Value extrLHS = rewriter.create<ExtractOp>(loc, op.lhs(), i); 676 Value extrRHS = rewriter.create<ExtractOp>(loc, op.rhs(), i); 677 Value extrACC = rewriter.create<ExtractOp>(loc, op.acc(), i); 678 Value fma = rewriter.create<FMAOp>(loc, extrLHS, extrRHS, extrACC); 679 desc = rewriter.create<InsertOp>(loc, fma, desc, i); 680 } 681 rewriter.replaceOp(op, desc); 682 return success(); 683 } 684 }; 685 686 // When ranks are different, InsertStridedSlice needs to extract a properly 687 // ranked vector from the destination vector into which to insert. This pattern 688 // only takes care of this part and forwards the rest of the conversion to 689 // another pattern that converts InsertStridedSlice for operands of the same 690 // rank. 691 // 692 // RewritePattern for InsertStridedSliceOp where source and destination vectors 693 // have different ranks. In this case: 694 // 1. the proper subvector is extracted from the destination vector 695 // 2. a new InsertStridedSlice op is created to insert the source in the 696 // destination subvector 697 // 3. the destination subvector is inserted back in the proper place 698 // 4. the op is replaced by the result of step 3. 699 // The new InsertStridedSlice from step 2. will be picked up by a 700 // `VectorInsertStridedSliceOpSameRankRewritePattern`. 701 class VectorInsertStridedSliceOpDifferentRankRewritePattern 702 : public OpRewritePattern<InsertStridedSliceOp> { 703 public: 704 using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 705 706 LogicalResult matchAndRewrite(InsertStridedSliceOp op, 707 PatternRewriter &rewriter) const override { 708 auto srcType = op.getSourceVectorType(); 709 auto dstType = op.getDestVectorType(); 710 711 if (op.offsets().getValue().empty()) 712 return failure(); 713 714 auto loc = op.getLoc(); 715 int64_t rankDiff = dstType.getRank() - srcType.getRank(); 716 assert(rankDiff >= 0); 717 if (rankDiff == 0) 718 return failure(); 719 720 int64_t rankRest = dstType.getRank() - rankDiff; 721 // Extract / insert the subvector of matching rank and InsertStridedSlice 722 // on it. 723 Value extracted = 724 rewriter.create<ExtractOp>(loc, op.dest(), 725 getI64SubArray(op.offsets(), /*dropFront=*/0, 726 /*dropFront=*/rankRest)); 727 // A different pattern will kick in for InsertStridedSlice with matching 728 // ranks. 729 auto stridedSliceInnerOp = rewriter.create<InsertStridedSliceOp>( 730 loc, op.source(), extracted, 731 getI64SubArray(op.offsets(), /*dropFront=*/rankDiff), 732 getI64SubArray(op.strides(), /*dropFront=*/0)); 733 rewriter.replaceOpWithNewOp<InsertOp>( 734 op, stridedSliceInnerOp.getResult(), op.dest(), 735 getI64SubArray(op.offsets(), /*dropFront=*/0, 736 /*dropFront=*/rankRest)); 737 return success(); 738 } 739 }; 740 741 // RewritePattern for InsertStridedSliceOp where source and destination vectors 742 // have the same rank. In this case, we reduce 743 // 1. the proper subvector is extracted from the destination vector 744 // 2. a new InsertStridedSlice op is created to insert the source in the 745 // destination subvector 746 // 3. the destination subvector is inserted back in the proper place 747 // 4. the op is replaced by the result of step 3. 748 // The new InsertStridedSlice from step 2. will be picked up by a 749 // `VectorInsertStridedSliceOpSameRankRewritePattern`. 750 class VectorInsertStridedSliceOpSameRankRewritePattern 751 : public OpRewritePattern<InsertStridedSliceOp> { 752 public: 753 using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 754 755 LogicalResult matchAndRewrite(InsertStridedSliceOp op, 756 PatternRewriter &rewriter) const override { 757 auto srcType = op.getSourceVectorType(); 758 auto dstType = op.getDestVectorType(); 759 760 if (op.offsets().getValue().empty()) 761 return failure(); 762 763 int64_t rankDiff = dstType.getRank() - srcType.getRank(); 764 assert(rankDiff >= 0); 765 if (rankDiff != 0) 766 return failure(); 767 768 if (srcType == dstType) { 769 rewriter.replaceOp(op, op.source()); 770 return success(); 771 } 772 773 int64_t offset = 774 op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 775 int64_t size = srcType.getShape().front(); 776 int64_t stride = 777 op.strides().getValue().front().cast<IntegerAttr>().getInt(); 778 779 auto loc = op.getLoc(); 780 Value res = op.dest(); 781 // For each slice of the source vector along the most major dimension. 782 for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 783 off += stride, ++idx) { 784 // 1. extract the proper subvector (or element) from source 785 Value extractedSource = extractOne(rewriter, loc, op.source(), idx); 786 if (extractedSource.getType().isa<VectorType>()) { 787 // 2. If we have a vector, extract the proper subvector from destination 788 // Otherwise we are at the element level and no need to recurse. 789 Value extractedDest = extractOne(rewriter, loc, op.dest(), off); 790 // 3. Reduce the problem to lowering a new InsertStridedSlice op with 791 // smaller rank. 792 InsertStridedSliceOp insertStridedSliceOp = 793 rewriter.create<InsertStridedSliceOp>( 794 loc, extractedSource, extractedDest, 795 getI64SubArray(op.offsets(), /* dropFront=*/1), 796 getI64SubArray(op.strides(), /* dropFront=*/1)); 797 // Call matchAndRewrite recursively from within the pattern. This 798 // circumvents the current limitation that a given pattern cannot 799 // be called multiple times by the PatternRewrite infrastructure (to 800 // avoid infinite recursion, but in this case, infinite recursion 801 // cannot happen because the rank is strictly decreasing). 802 // TODO(rriddle, nicolasvasilache) Implement something like a hook for 803 // a potential function that must decrease and allow the same pattern 804 // multiple times. 805 auto success = matchAndRewrite(insertStridedSliceOp, rewriter); 806 (void)success; 807 assert(succeeded(success) && "Unexpected failure"); 808 extractedSource = insertStridedSliceOp; 809 } 810 // 4. Insert the extractedSource into the res vector. 811 res = insertOne(rewriter, loc, extractedSource, res, off); 812 } 813 814 rewriter.replaceOp(op, res); 815 return success(); 816 } 817 }; 818 819 class VectorTypeCastOpConversion : public ConvertToLLVMPattern { 820 public: 821 explicit VectorTypeCastOpConversion(MLIRContext *context, 822 LLVMTypeConverter &typeConverter) 823 : ConvertToLLVMPattern(vector::TypeCastOp::getOperationName(), context, 824 typeConverter) {} 825 826 LogicalResult 827 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 828 ConversionPatternRewriter &rewriter) const override { 829 auto loc = op->getLoc(); 830 vector::TypeCastOp castOp = cast<vector::TypeCastOp>(op); 831 MemRefType sourceMemRefType = 832 castOp.getOperand().getType().cast<MemRefType>(); 833 MemRefType targetMemRefType = 834 castOp.getResult().getType().cast<MemRefType>(); 835 836 // Only static shape casts supported atm. 837 if (!sourceMemRefType.hasStaticShape() || 838 !targetMemRefType.hasStaticShape()) 839 return failure(); 840 841 auto llvmSourceDescriptorTy = 842 operands[0].getType().dyn_cast<LLVM::LLVMType>(); 843 if (!llvmSourceDescriptorTy || !llvmSourceDescriptorTy.isStructTy()) 844 return failure(); 845 MemRefDescriptor sourceMemRef(operands[0]); 846 847 auto llvmTargetDescriptorTy = typeConverter.convertType(targetMemRefType) 848 .dyn_cast_or_null<LLVM::LLVMType>(); 849 if (!llvmTargetDescriptorTy || !llvmTargetDescriptorTy.isStructTy()) 850 return failure(); 851 852 int64_t offset; 853 SmallVector<int64_t, 4> strides; 854 auto successStrides = 855 getStridesAndOffset(sourceMemRefType, strides, offset); 856 bool isContiguous = (strides.back() == 1); 857 if (isContiguous) { 858 auto sizes = sourceMemRefType.getShape(); 859 for (int index = 0, e = strides.size() - 2; index < e; ++index) { 860 if (strides[index] != strides[index + 1] * sizes[index + 1]) { 861 isContiguous = false; 862 break; 863 } 864 } 865 } 866 // Only contiguous source tensors supported atm. 867 if (failed(successStrides) || !isContiguous) 868 return failure(); 869 870 auto int64Ty = LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()); 871 872 // Create descriptor. 873 auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy); 874 Type llvmTargetElementTy = desc.getElementType(); 875 // Set allocated ptr. 876 Value allocated = sourceMemRef.allocatedPtr(rewriter, loc); 877 allocated = 878 rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated); 879 desc.setAllocatedPtr(rewriter, loc, allocated); 880 // Set aligned ptr. 881 Value ptr = sourceMemRef.alignedPtr(rewriter, loc); 882 ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr); 883 desc.setAlignedPtr(rewriter, loc, ptr); 884 // Fill offset 0. 885 auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0); 886 auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr); 887 desc.setOffset(rewriter, loc, zero); 888 889 // Fill size and stride descriptors in memref. 890 for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) { 891 int64_t index = indexedSize.index(); 892 auto sizeAttr = 893 rewriter.getIntegerAttr(rewriter.getIndexType(), indexedSize.value()); 894 auto size = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, sizeAttr); 895 desc.setSize(rewriter, loc, index, size); 896 auto strideAttr = 897 rewriter.getIntegerAttr(rewriter.getIndexType(), strides[index]); 898 auto stride = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, strideAttr); 899 desc.setStride(rewriter, loc, index, stride); 900 } 901 902 rewriter.replaceOp(op, {desc}); 903 return success(); 904 } 905 }; 906 907 class VectorPrintOpConversion : public ConvertToLLVMPattern { 908 public: 909 explicit VectorPrintOpConversion(MLIRContext *context, 910 LLVMTypeConverter &typeConverter) 911 : ConvertToLLVMPattern(vector::PrintOp::getOperationName(), context, 912 typeConverter) {} 913 914 // Proof-of-concept lowering implementation that relies on a small 915 // runtime support library, which only needs to provide a few 916 // printing methods (single value for all data types, opening/closing 917 // bracket, comma, newline). The lowering fully unrolls a vector 918 // in terms of these elementary printing operations. The advantage 919 // of this approach is that the library can remain unaware of all 920 // low-level implementation details of vectors while still supporting 921 // output of any shaped and dimensioned vector. Due to full unrolling, 922 // this approach is less suited for very large vectors though. 923 // 924 // TODO(ajcbik): rely solely on libc in future? something else? 925 // 926 LogicalResult 927 matchAndRewrite(Operation *op, ArrayRef<Value> operands, 928 ConversionPatternRewriter &rewriter) const override { 929 auto printOp = cast<vector::PrintOp>(op); 930 auto adaptor = vector::PrintOpOperandAdaptor(operands); 931 Type printType = printOp.getPrintType(); 932 933 if (typeConverter.convertType(printType) == nullptr) 934 return failure(); 935 936 // Make sure element type has runtime support (currently just Float/Double). 937 VectorType vectorType = printType.dyn_cast<VectorType>(); 938 Type eltType = vectorType ? vectorType.getElementType() : printType; 939 int64_t rank = vectorType ? vectorType.getRank() : 0; 940 Operation *printer; 941 if (eltType.isSignlessInteger(32)) 942 printer = getPrintI32(op); 943 else if (eltType.isSignlessInteger(64)) 944 printer = getPrintI64(op); 945 else if (eltType.isF32()) 946 printer = getPrintFloat(op); 947 else if (eltType.isF64()) 948 printer = getPrintDouble(op); 949 else 950 return failure(); 951 952 // Unroll vector into elementary print calls. 953 emitRanks(rewriter, op, adaptor.source(), vectorType, printer, rank); 954 emitCall(rewriter, op->getLoc(), getPrintNewline(op)); 955 rewriter.eraseOp(op); 956 return success(); 957 } 958 959 private: 960 void emitRanks(ConversionPatternRewriter &rewriter, Operation *op, 961 Value value, VectorType vectorType, Operation *printer, 962 int64_t rank) const { 963 Location loc = op->getLoc(); 964 if (rank == 0) { 965 emitCall(rewriter, loc, printer, value); 966 return; 967 } 968 969 emitCall(rewriter, loc, getPrintOpen(op)); 970 Operation *printComma = getPrintComma(op); 971 int64_t dim = vectorType.getDimSize(0); 972 for (int64_t d = 0; d < dim; ++d) { 973 auto reducedType = 974 rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr; 975 auto llvmType = typeConverter.convertType( 976 rank > 1 ? reducedType : vectorType.getElementType()); 977 Value nestedVal = 978 extractOne(rewriter, typeConverter, loc, value, llvmType, rank, d); 979 emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1); 980 if (d != dim - 1) 981 emitCall(rewriter, loc, printComma); 982 } 983 emitCall(rewriter, loc, getPrintClose(op)); 984 } 985 986 // Helper to emit a call. 987 static void emitCall(ConversionPatternRewriter &rewriter, Location loc, 988 Operation *ref, ValueRange params = ValueRange()) { 989 rewriter.create<LLVM::CallOp>(loc, ArrayRef<Type>{}, 990 rewriter.getSymbolRefAttr(ref), params); 991 } 992 993 // Helper for printer method declaration (first hit) and lookup. 994 static Operation *getPrint(Operation *op, LLVM::LLVMDialect *dialect, 995 StringRef name, ArrayRef<LLVM::LLVMType> params) { 996 auto module = op->getParentOfType<ModuleOp>(); 997 auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(name); 998 if (func) 999 return func; 1000 OpBuilder moduleBuilder(module.getBodyRegion()); 1001 return moduleBuilder.create<LLVM::LLVMFuncOp>( 1002 op->getLoc(), name, 1003 LLVM::LLVMType::getFunctionTy(LLVM::LLVMType::getVoidTy(dialect), 1004 params, /*isVarArg=*/false)); 1005 } 1006 1007 // Helpers for method names. 1008 Operation *getPrintI32(Operation *op) const { 1009 LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1010 return getPrint(op, dialect, "print_i32", 1011 LLVM::LLVMType::getInt32Ty(dialect)); 1012 } 1013 Operation *getPrintI64(Operation *op) const { 1014 LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1015 return getPrint(op, dialect, "print_i64", 1016 LLVM::LLVMType::getInt64Ty(dialect)); 1017 } 1018 Operation *getPrintFloat(Operation *op) const { 1019 LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1020 return getPrint(op, dialect, "print_f32", 1021 LLVM::LLVMType::getFloatTy(dialect)); 1022 } 1023 Operation *getPrintDouble(Operation *op) const { 1024 LLVM::LLVMDialect *dialect = typeConverter.getDialect(); 1025 return getPrint(op, dialect, "print_f64", 1026 LLVM::LLVMType::getDoubleTy(dialect)); 1027 } 1028 Operation *getPrintOpen(Operation *op) const { 1029 return getPrint(op, typeConverter.getDialect(), "print_open", {}); 1030 } 1031 Operation *getPrintClose(Operation *op) const { 1032 return getPrint(op, typeConverter.getDialect(), "print_close", {}); 1033 } 1034 Operation *getPrintComma(Operation *op) const { 1035 return getPrint(op, typeConverter.getDialect(), "print_comma", {}); 1036 } 1037 Operation *getPrintNewline(Operation *op) const { 1038 return getPrint(op, typeConverter.getDialect(), "print_newline", {}); 1039 } 1040 }; 1041 1042 /// Progressive lowering of StridedSliceOp to either: 1043 /// 1. extractelement + insertelement for the 1-D case 1044 /// 2. extract + optional strided_slice + insert for the n-D case. 1045 class VectorStridedSliceOpConversion : public OpRewritePattern<StridedSliceOp> { 1046 public: 1047 using OpRewritePattern<StridedSliceOp>::OpRewritePattern; 1048 1049 LogicalResult matchAndRewrite(StridedSliceOp op, 1050 PatternRewriter &rewriter) const override { 1051 auto dstType = op.getResult().getType().cast<VectorType>(); 1052 1053 assert(!op.offsets().getValue().empty() && "Unexpected empty offsets"); 1054 1055 int64_t offset = 1056 op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 1057 int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt(); 1058 int64_t stride = 1059 op.strides().getValue().front().cast<IntegerAttr>().getInt(); 1060 1061 auto loc = op.getLoc(); 1062 auto elemType = dstType.getElementType(); 1063 assert(elemType.isSignlessIntOrIndexOrFloat()); 1064 Value zero = rewriter.create<ConstantOp>(loc, elemType, 1065 rewriter.getZeroAttr(elemType)); 1066 Value res = rewriter.create<SplatOp>(loc, dstType, zero); 1067 for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 1068 off += stride, ++idx) { 1069 Value extracted = extractOne(rewriter, loc, op.vector(), off); 1070 if (op.offsets().getValue().size() > 1) { 1071 StridedSliceOp stridedSliceOp = rewriter.create<StridedSliceOp>( 1072 loc, extracted, getI64SubArray(op.offsets(), /* dropFront=*/1), 1073 getI64SubArray(op.sizes(), /* dropFront=*/1), 1074 getI64SubArray(op.strides(), /* dropFront=*/1)); 1075 // Call matchAndRewrite recursively from within the pattern. This 1076 // circumvents the current limitation that a given pattern cannot 1077 // be called multiple times by the PatternRewrite infrastructure (to 1078 // avoid infinite recursion, but in this case, infinite recursion 1079 // cannot happen because the rank is strictly decreasing). 1080 // TODO(rriddle, nicolasvasilache) Implement something like a hook for 1081 // a potential function that must decrease and allow the same pattern 1082 // multiple times. 1083 auto success = matchAndRewrite(stridedSliceOp, rewriter); 1084 (void)success; 1085 assert(succeeded(success) && "Unexpected failure"); 1086 extracted = stridedSliceOp; 1087 } 1088 res = insertOne(rewriter, loc, extracted, res, idx); 1089 } 1090 rewriter.replaceOp(op, {res}); 1091 return success(); 1092 } 1093 }; 1094 1095 } // namespace 1096 1097 /// Populate the given list with patterns that convert from Vector to LLVM. 1098 void mlir::populateVectorToLLVMConversionPatterns( 1099 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 1100 MLIRContext *ctx = converter.getDialect()->getContext(); 1101 patterns.insert<VectorFMAOpNDRewritePattern, 1102 VectorInsertStridedSliceOpDifferentRankRewritePattern, 1103 VectorInsertStridedSliceOpSameRankRewritePattern, 1104 VectorStridedSliceOpConversion>(ctx); 1105 patterns.insert<VectorBroadcastOpConversion, VectorReductionOpConversion, 1106 VectorShuffleOpConversion, VectorExtractElementOpConversion, 1107 VectorExtractOpConversion, VectorFMAOp1DConversion, 1108 VectorInsertElementOpConversion, VectorInsertOpConversion, 1109 VectorTypeCastOpConversion, VectorPrintOpConversion>( 1110 ctx, converter); 1111 } 1112 1113 void mlir::populateVectorToLLVMMatrixConversionPatterns( 1114 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 1115 MLIRContext *ctx = converter.getDialect()->getContext(); 1116 patterns.insert<VectorMatmulOpConversion>(ctx, converter); 1117 } 1118 1119 namespace { 1120 struct LowerVectorToLLVMPass 1121 : public ConvertVectorToLLVMBase<LowerVectorToLLVMPass> { 1122 void runOnOperation() override; 1123 }; 1124 } // namespace 1125 1126 void LowerVectorToLLVMPass::runOnOperation() { 1127 // Perform progressive lowering of operations on slices and 1128 // all contraction operations. Also applies folding and DCE. 1129 { 1130 OwningRewritePatternList patterns; 1131 populateVectorSlicesLoweringPatterns(patterns, &getContext()); 1132 populateVectorContractLoweringPatterns(patterns, &getContext()); 1133 applyPatternsGreedily(getOperation(), patterns); 1134 } 1135 1136 // Convert to the LLVM IR dialect. 1137 LLVMTypeConverter converter(&getContext()); 1138 OwningRewritePatternList patterns; 1139 populateVectorToLLVMMatrixConversionPatterns(converter, patterns); 1140 populateVectorToLLVMConversionPatterns(converter, patterns); 1141 populateVectorToLLVMMatrixConversionPatterns(converter, patterns); 1142 populateStdToLLVMConversionPatterns(converter, patterns); 1143 1144 LLVMConversionTarget target(getContext()); 1145 target.addDynamicallyLegalOp<FuncOp>( 1146 [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); 1147 if (failed(applyPartialConversion(getOperation(), target, patterns, 1148 &converter))) { 1149 signalPassFailure(); 1150 } 1151 } 1152 1153 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertVectorToLLVMPass() { 1154 return std::make_unique<LowerVectorToLLVMPass>(); 1155 } 1156