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