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