1 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 9 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" 10 11 #include "mlir/Conversion/LLVMCommon/VectorPattern.h" 12 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h" 13 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 14 #include "mlir/Dialect/MemRef/IR/MemRef.h" 15 #include "mlir/Dialect/StandardOps/IR/Ops.h" 16 #include "mlir/Dialect/Vector/VectorOps.h" 17 #include "mlir/IR/BuiltinTypes.h" 18 #include "mlir/Support/MathExtras.h" 19 #include "mlir/Target/LLVMIR/TypeToLLVM.h" 20 #include "mlir/Transforms/DialectConversion.h" 21 22 using namespace mlir; 23 using namespace mlir::vector; 24 25 // Helper to reduce vector type by one rank at front. 26 static VectorType reducedVectorTypeFront(VectorType tp) { 27 assert((tp.getRank() > 1) && "unlowerable vector type"); 28 return VectorType::get(tp.getShape().drop_front(), tp.getElementType()); 29 } 30 31 // Helper to reduce vector type by *all* but one rank at back. 32 static VectorType reducedVectorTypeBack(VectorType tp) { 33 assert((tp.getRank() > 1) && "unlowerable vector type"); 34 return VectorType::get(tp.getShape().take_back(), tp.getElementType()); 35 } 36 37 // Helper that picks the proper sequence for inserting. 38 static Value insertOne(ConversionPatternRewriter &rewriter, 39 LLVMTypeConverter &typeConverter, Location loc, 40 Value val1, Value val2, Type llvmType, int64_t rank, 41 int64_t pos) { 42 if (rank == 1) { 43 auto idxType = rewriter.getIndexType(); 44 auto constant = rewriter.create<LLVM::ConstantOp>( 45 loc, typeConverter.convertType(idxType), 46 rewriter.getIntegerAttr(idxType, pos)); 47 return rewriter.create<LLVM::InsertElementOp>(loc, llvmType, val1, val2, 48 constant); 49 } 50 return rewriter.create<LLVM::InsertValueOp>(loc, llvmType, val1, val2, 51 rewriter.getI64ArrayAttr(pos)); 52 } 53 54 // Helper that picks the proper sequence for inserting. 55 static Value insertOne(PatternRewriter &rewriter, Location loc, Value from, 56 Value into, int64_t offset) { 57 auto vectorType = into.getType().cast<VectorType>(); 58 if (vectorType.getRank() > 1) 59 return rewriter.create<InsertOp>(loc, from, into, offset); 60 return rewriter.create<vector::InsertElementOp>( 61 loc, vectorType, from, into, 62 rewriter.create<ConstantIndexOp>(loc, offset)); 63 } 64 65 // Helper that picks the proper sequence for extracting. 66 static Value extractOne(ConversionPatternRewriter &rewriter, 67 LLVMTypeConverter &typeConverter, Location loc, 68 Value val, Type llvmType, int64_t rank, int64_t pos) { 69 if (rank == 1) { 70 auto idxType = rewriter.getIndexType(); 71 auto constant = rewriter.create<LLVM::ConstantOp>( 72 loc, typeConverter.convertType(idxType), 73 rewriter.getIntegerAttr(idxType, pos)); 74 return rewriter.create<LLVM::ExtractElementOp>(loc, llvmType, val, 75 constant); 76 } 77 return rewriter.create<LLVM::ExtractValueOp>(loc, llvmType, val, 78 rewriter.getI64ArrayAttr(pos)); 79 } 80 81 // Helper that picks the proper sequence for extracting. 82 static Value extractOne(PatternRewriter &rewriter, Location loc, Value vector, 83 int64_t offset) { 84 auto vectorType = vector.getType().cast<VectorType>(); 85 if (vectorType.getRank() > 1) 86 return rewriter.create<ExtractOp>(loc, vector, offset); 87 return rewriter.create<vector::ExtractElementOp>( 88 loc, vectorType.getElementType(), vector, 89 rewriter.create<ConstantIndexOp>(loc, offset)); 90 } 91 92 // Helper that returns a subset of `arrayAttr` as a vector of int64_t. 93 // TODO: Better support for attribute subtype forwarding + slicing. 94 static SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr, 95 unsigned dropFront = 0, 96 unsigned dropBack = 0) { 97 assert(arrayAttr.size() > dropFront + dropBack && "Out of bounds"); 98 auto range = arrayAttr.getAsRange<IntegerAttr>(); 99 SmallVector<int64_t, 4> res; 100 res.reserve(arrayAttr.size() - dropFront - dropBack); 101 for (auto it = range.begin() + dropFront, eit = range.end() - dropBack; 102 it != eit; ++it) 103 res.push_back((*it).getValue().getSExtValue()); 104 return res; 105 } 106 107 // Helper that returns data layout alignment of a memref. 108 LogicalResult getMemRefAlignment(LLVMTypeConverter &typeConverter, 109 MemRefType memrefType, unsigned &align) { 110 Type elementTy = typeConverter.convertType(memrefType.getElementType()); 111 if (!elementTy) 112 return failure(); 113 114 // TODO: this should use the MLIR data layout when it becomes available and 115 // stop depending on translation. 116 llvm::LLVMContext llvmContext; 117 align = LLVM::TypeToLLVMIRTranslator(llvmContext) 118 .getPreferredAlignment(elementTy, typeConverter.getDataLayout()); 119 return success(); 120 } 121 122 // Return the minimal alignment value that satisfies all the AssumeAlignment 123 // uses of `value`. If no such uses exist, return 1. 124 static unsigned getAssumedAlignment(Value value) { 125 unsigned align = 1; 126 for (auto &u : value.getUses()) { 127 Operation *owner = u.getOwner(); 128 if (auto op = dyn_cast<memref::AssumeAlignmentOp>(owner)) 129 align = mlir::lcm(align, op.alignment()); 130 } 131 return align; 132 } 133 134 // Helper that returns data layout alignment of a memref associated with a 135 // load, store, scatter, or gather op, including additional information from 136 // assume_alignment calls on the source of the transfer 137 template <class OpAdaptor> 138 LogicalResult getMemRefOpAlignment(LLVMTypeConverter &typeConverter, 139 OpAdaptor op, unsigned &align) { 140 if (failed(getMemRefAlignment(typeConverter, op.getMemRefType(), align))) 141 return failure(); 142 align = std::max(align, getAssumedAlignment(op.base())); 143 return success(); 144 } 145 146 // Add an index vector component to a base pointer. This almost always succeeds 147 // unless the last stride is non-unit or the memory space is not zero. 148 static LogicalResult getIndexedPtrs(ConversionPatternRewriter &rewriter, 149 Location loc, Value memref, Value base, 150 Value index, MemRefType memRefType, 151 VectorType vType, Value &ptrs) { 152 int64_t offset; 153 SmallVector<int64_t, 4> strides; 154 auto successStrides = getStridesAndOffset(memRefType, strides, offset); 155 if (failed(successStrides) || strides.back() != 1 || 156 memRefType.getMemorySpaceAsInt() != 0) 157 return failure(); 158 auto pType = MemRefDescriptor(memref).getElementPtrType(); 159 auto ptrsType = LLVM::getFixedVectorType(pType, vType.getDimSize(0)); 160 ptrs = rewriter.create<LLVM::GEPOp>(loc, ptrsType, base, index); 161 return success(); 162 } 163 164 // Casts a strided element pointer to a vector pointer. The vector pointer 165 // will be in the same address space as the incoming memref type. 166 static Value castDataPtr(ConversionPatternRewriter &rewriter, Location loc, 167 Value ptr, MemRefType memRefType, Type vt) { 168 auto pType = LLVM::LLVMPointerType::get(vt, memRefType.getMemorySpaceAsInt()); 169 return rewriter.create<LLVM::BitcastOp>(loc, pType, ptr); 170 } 171 172 namespace { 173 174 /// Conversion pattern for a vector.bitcast. 175 class VectorBitCastOpConversion 176 : public ConvertOpToLLVMPattern<vector::BitCastOp> { 177 public: 178 using ConvertOpToLLVMPattern<vector::BitCastOp>::ConvertOpToLLVMPattern; 179 180 LogicalResult 181 matchAndRewrite(vector::BitCastOp bitCastOp, OpAdaptor adaptor, 182 ConversionPatternRewriter &rewriter) const override { 183 // Only 1-D vectors can be lowered to LLVM. 184 VectorType resultTy = bitCastOp.getType(); 185 if (resultTy.getRank() != 1) 186 return failure(); 187 Type newResultTy = typeConverter->convertType(resultTy); 188 rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(bitCastOp, newResultTy, 189 adaptor.getOperands()[0]); 190 return success(); 191 } 192 }; 193 194 /// Conversion pattern for a vector.matrix_multiply. 195 /// This is lowered directly to the proper llvm.intr.matrix.multiply. 196 class VectorMatmulOpConversion 197 : public ConvertOpToLLVMPattern<vector::MatmulOp> { 198 public: 199 using ConvertOpToLLVMPattern<vector::MatmulOp>::ConvertOpToLLVMPattern; 200 201 LogicalResult 202 matchAndRewrite(vector::MatmulOp matmulOp, OpAdaptor adaptor, 203 ConversionPatternRewriter &rewriter) const override { 204 rewriter.replaceOpWithNewOp<LLVM::MatrixMultiplyOp>( 205 matmulOp, typeConverter->convertType(matmulOp.res().getType()), 206 adaptor.lhs(), adaptor.rhs(), matmulOp.lhs_rows(), 207 matmulOp.lhs_columns(), matmulOp.rhs_columns()); 208 return success(); 209 } 210 }; 211 212 /// Conversion pattern for a vector.flat_transpose. 213 /// This is lowered directly to the proper llvm.intr.matrix.transpose. 214 class VectorFlatTransposeOpConversion 215 : public ConvertOpToLLVMPattern<vector::FlatTransposeOp> { 216 public: 217 using ConvertOpToLLVMPattern<vector::FlatTransposeOp>::ConvertOpToLLVMPattern; 218 219 LogicalResult 220 matchAndRewrite(vector::FlatTransposeOp transOp, OpAdaptor adaptor, 221 ConversionPatternRewriter &rewriter) const override { 222 rewriter.replaceOpWithNewOp<LLVM::MatrixTransposeOp>( 223 transOp, typeConverter->convertType(transOp.res().getType()), 224 adaptor.matrix(), transOp.rows(), transOp.columns()); 225 return success(); 226 } 227 }; 228 229 /// Overloaded utility that replaces a vector.load, vector.store, 230 /// vector.maskedload and vector.maskedstore with their respective LLVM 231 /// couterparts. 232 static void replaceLoadOrStoreOp(vector::LoadOp loadOp, 233 vector::LoadOpAdaptor adaptor, 234 VectorType vectorTy, Value ptr, unsigned align, 235 ConversionPatternRewriter &rewriter) { 236 rewriter.replaceOpWithNewOp<LLVM::LoadOp>(loadOp, ptr, align); 237 } 238 239 static void replaceLoadOrStoreOp(vector::MaskedLoadOp loadOp, 240 vector::MaskedLoadOpAdaptor adaptor, 241 VectorType vectorTy, Value ptr, unsigned align, 242 ConversionPatternRewriter &rewriter) { 243 rewriter.replaceOpWithNewOp<LLVM::MaskedLoadOp>( 244 loadOp, vectorTy, ptr, adaptor.mask(), adaptor.pass_thru(), align); 245 } 246 247 static void replaceLoadOrStoreOp(vector::StoreOp storeOp, 248 vector::StoreOpAdaptor adaptor, 249 VectorType vectorTy, Value ptr, unsigned align, 250 ConversionPatternRewriter &rewriter) { 251 rewriter.replaceOpWithNewOp<LLVM::StoreOp>(storeOp, adaptor.valueToStore(), 252 ptr, align); 253 } 254 255 static void replaceLoadOrStoreOp(vector::MaskedStoreOp storeOp, 256 vector::MaskedStoreOpAdaptor adaptor, 257 VectorType vectorTy, Value ptr, unsigned align, 258 ConversionPatternRewriter &rewriter) { 259 rewriter.replaceOpWithNewOp<LLVM::MaskedStoreOp>( 260 storeOp, adaptor.valueToStore(), ptr, adaptor.mask(), align); 261 } 262 263 /// Conversion pattern for a vector.load, vector.store, vector.maskedload, and 264 /// vector.maskedstore. 265 template <class LoadOrStoreOp, class LoadOrStoreOpAdaptor> 266 class VectorLoadStoreConversion : public ConvertOpToLLVMPattern<LoadOrStoreOp> { 267 public: 268 using ConvertOpToLLVMPattern<LoadOrStoreOp>::ConvertOpToLLVMPattern; 269 270 LogicalResult 271 matchAndRewrite(LoadOrStoreOp loadOrStoreOp, 272 typename LoadOrStoreOp::Adaptor adaptor, 273 ConversionPatternRewriter &rewriter) const override { 274 // Only 1-D vectors can be lowered to LLVM. 275 VectorType vectorTy = loadOrStoreOp.getVectorType(); 276 if (vectorTy.getRank() > 1) 277 return failure(); 278 279 auto loc = loadOrStoreOp->getLoc(); 280 MemRefType memRefTy = loadOrStoreOp.getMemRefType(); 281 282 // Resolve alignment. 283 unsigned align; 284 if (failed(getMemRefOpAlignment(*this->getTypeConverter(), loadOrStoreOp, 285 align))) 286 return failure(); 287 288 // Resolve address. 289 auto vtype = this->typeConverter->convertType(loadOrStoreOp.getVectorType()) 290 .template cast<VectorType>(); 291 Value dataPtr = this->getStridedElementPtr(loc, memRefTy, adaptor.base(), 292 adaptor.indices(), rewriter); 293 Value ptr = castDataPtr(rewriter, loc, dataPtr, memRefTy, vtype); 294 295 replaceLoadOrStoreOp(loadOrStoreOp, adaptor, vtype, ptr, align, rewriter); 296 return success(); 297 } 298 }; 299 300 /// Conversion pattern for a vector.gather. 301 class VectorGatherOpConversion 302 : public ConvertOpToLLVMPattern<vector::GatherOp> { 303 public: 304 using ConvertOpToLLVMPattern<vector::GatherOp>::ConvertOpToLLVMPattern; 305 306 LogicalResult 307 matchAndRewrite(vector::GatherOp gather, OpAdaptor adaptor, 308 ConversionPatternRewriter &rewriter) const override { 309 auto loc = gather->getLoc(); 310 MemRefType memRefType = gather.getMemRefType(); 311 312 // Resolve alignment. 313 unsigned align; 314 if (failed(getMemRefOpAlignment(*getTypeConverter(), gather, align))) 315 return failure(); 316 317 // Resolve address. 318 Value ptrs; 319 VectorType vType = gather.getVectorType(); 320 Value ptr = getStridedElementPtr(loc, memRefType, adaptor.base(), 321 adaptor.indices(), rewriter); 322 if (failed(getIndexedPtrs(rewriter, loc, adaptor.base(), ptr, 323 adaptor.index_vec(), memRefType, vType, ptrs))) 324 return failure(); 325 326 // Replace with the gather intrinsic. 327 rewriter.replaceOpWithNewOp<LLVM::masked_gather>( 328 gather, typeConverter->convertType(vType), ptrs, adaptor.mask(), 329 adaptor.pass_thru(), rewriter.getI32IntegerAttr(align)); 330 return success(); 331 } 332 }; 333 334 /// Conversion pattern for a vector.scatter. 335 class VectorScatterOpConversion 336 : public ConvertOpToLLVMPattern<vector::ScatterOp> { 337 public: 338 using ConvertOpToLLVMPattern<vector::ScatterOp>::ConvertOpToLLVMPattern; 339 340 LogicalResult 341 matchAndRewrite(vector::ScatterOp scatter, OpAdaptor adaptor, 342 ConversionPatternRewriter &rewriter) const override { 343 auto loc = scatter->getLoc(); 344 MemRefType memRefType = scatter.getMemRefType(); 345 346 // Resolve alignment. 347 unsigned align; 348 if (failed(getMemRefOpAlignment(*getTypeConverter(), scatter, align))) 349 return failure(); 350 351 // Resolve address. 352 Value ptrs; 353 VectorType vType = scatter.getVectorType(); 354 Value ptr = getStridedElementPtr(loc, memRefType, adaptor.base(), 355 adaptor.indices(), rewriter); 356 if (failed(getIndexedPtrs(rewriter, loc, adaptor.base(), ptr, 357 adaptor.index_vec(), memRefType, vType, ptrs))) 358 return failure(); 359 360 // Replace with the scatter intrinsic. 361 rewriter.replaceOpWithNewOp<LLVM::masked_scatter>( 362 scatter, adaptor.valueToStore(), ptrs, adaptor.mask(), 363 rewriter.getI32IntegerAttr(align)); 364 return success(); 365 } 366 }; 367 368 /// Conversion pattern for a vector.expandload. 369 class VectorExpandLoadOpConversion 370 : public ConvertOpToLLVMPattern<vector::ExpandLoadOp> { 371 public: 372 using ConvertOpToLLVMPattern<vector::ExpandLoadOp>::ConvertOpToLLVMPattern; 373 374 LogicalResult 375 matchAndRewrite(vector::ExpandLoadOp expand, OpAdaptor adaptor, 376 ConversionPatternRewriter &rewriter) const override { 377 auto loc = expand->getLoc(); 378 MemRefType memRefType = expand.getMemRefType(); 379 380 // Resolve address. 381 auto vtype = typeConverter->convertType(expand.getVectorType()); 382 Value ptr = getStridedElementPtr(loc, memRefType, adaptor.base(), 383 adaptor.indices(), rewriter); 384 385 rewriter.replaceOpWithNewOp<LLVM::masked_expandload>( 386 expand, vtype, ptr, adaptor.mask(), adaptor.pass_thru()); 387 return success(); 388 } 389 }; 390 391 /// Conversion pattern for a vector.compressstore. 392 class VectorCompressStoreOpConversion 393 : public ConvertOpToLLVMPattern<vector::CompressStoreOp> { 394 public: 395 using ConvertOpToLLVMPattern<vector::CompressStoreOp>::ConvertOpToLLVMPattern; 396 397 LogicalResult 398 matchAndRewrite(vector::CompressStoreOp compress, OpAdaptor adaptor, 399 ConversionPatternRewriter &rewriter) const override { 400 auto loc = compress->getLoc(); 401 MemRefType memRefType = compress.getMemRefType(); 402 403 // Resolve address. 404 Value ptr = getStridedElementPtr(loc, memRefType, adaptor.base(), 405 adaptor.indices(), rewriter); 406 407 rewriter.replaceOpWithNewOp<LLVM::masked_compressstore>( 408 compress, adaptor.valueToStore(), ptr, adaptor.mask()); 409 return success(); 410 } 411 }; 412 413 /// Conversion pattern for all vector reductions. 414 class VectorReductionOpConversion 415 : public ConvertOpToLLVMPattern<vector::ReductionOp> { 416 public: 417 explicit VectorReductionOpConversion(LLVMTypeConverter &typeConv, 418 bool reassociateFPRed) 419 : ConvertOpToLLVMPattern<vector::ReductionOp>(typeConv), 420 reassociateFPReductions(reassociateFPRed) {} 421 422 LogicalResult 423 matchAndRewrite(vector::ReductionOp reductionOp, OpAdaptor adaptor, 424 ConversionPatternRewriter &rewriter) const override { 425 auto kind = reductionOp.kind(); 426 Type eltType = reductionOp.dest().getType(); 427 Type llvmType = typeConverter->convertType(eltType); 428 Value operand = adaptor.getOperands()[0]; 429 if (eltType.isIntOrIndex()) { 430 // Integer reductions: add/mul/min/max/and/or/xor. 431 if (kind == "add") 432 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_add>(reductionOp, 433 llvmType, operand); 434 else if (kind == "mul") 435 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_mul>(reductionOp, 436 llvmType, operand); 437 else if (kind == "min" && 438 (eltType.isIndex() || eltType.isUnsignedInteger())) 439 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_umin>( 440 reductionOp, llvmType, operand); 441 else if (kind == "min") 442 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_smin>( 443 reductionOp, llvmType, operand); 444 else if (kind == "max" && 445 (eltType.isIndex() || eltType.isUnsignedInteger())) 446 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_umax>( 447 reductionOp, llvmType, operand); 448 else if (kind == "max") 449 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_smax>( 450 reductionOp, llvmType, operand); 451 else if (kind == "and") 452 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_and>(reductionOp, 453 llvmType, operand); 454 else if (kind == "or") 455 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_or>(reductionOp, 456 llvmType, operand); 457 else if (kind == "xor") 458 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_xor>(reductionOp, 459 llvmType, operand); 460 else 461 return failure(); 462 return success(); 463 } 464 465 if (!eltType.isa<FloatType>()) 466 return failure(); 467 468 // Floating-point reductions: add/mul/min/max 469 if (kind == "add") { 470 // Optional accumulator (or zero). 471 Value acc = adaptor.getOperands().size() > 1 472 ? adaptor.getOperands()[1] 473 : rewriter.create<LLVM::ConstantOp>( 474 reductionOp->getLoc(), llvmType, 475 rewriter.getZeroAttr(eltType)); 476 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_fadd>( 477 reductionOp, llvmType, acc, operand, 478 rewriter.getBoolAttr(reassociateFPReductions)); 479 } else if (kind == "mul") { 480 // Optional accumulator (or one). 481 Value acc = adaptor.getOperands().size() > 1 482 ? adaptor.getOperands()[1] 483 : rewriter.create<LLVM::ConstantOp>( 484 reductionOp->getLoc(), llvmType, 485 rewriter.getFloatAttr(eltType, 1.0)); 486 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_fmul>( 487 reductionOp, llvmType, acc, operand, 488 rewriter.getBoolAttr(reassociateFPReductions)); 489 } else if (kind == "min") 490 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_fmin>(reductionOp, 491 llvmType, operand); 492 else if (kind == "max") 493 rewriter.replaceOpWithNewOp<LLVM::vector_reduce_fmax>(reductionOp, 494 llvmType, operand); 495 else 496 return failure(); 497 return success(); 498 } 499 500 private: 501 const bool reassociateFPReductions; 502 }; 503 504 class VectorShuffleOpConversion 505 : public ConvertOpToLLVMPattern<vector::ShuffleOp> { 506 public: 507 using ConvertOpToLLVMPattern<vector::ShuffleOp>::ConvertOpToLLVMPattern; 508 509 LogicalResult 510 matchAndRewrite(vector::ShuffleOp shuffleOp, OpAdaptor adaptor, 511 ConversionPatternRewriter &rewriter) const override { 512 auto loc = shuffleOp->getLoc(); 513 auto v1Type = shuffleOp.getV1VectorType(); 514 auto v2Type = shuffleOp.getV2VectorType(); 515 auto vectorType = shuffleOp.getVectorType(); 516 Type llvmType = typeConverter->convertType(vectorType); 517 auto maskArrayAttr = shuffleOp.mask(); 518 519 // Bail if result type cannot be lowered. 520 if (!llvmType) 521 return failure(); 522 523 // Get rank and dimension sizes. 524 int64_t rank = vectorType.getRank(); 525 assert(v1Type.getRank() == rank); 526 assert(v2Type.getRank() == rank); 527 int64_t v1Dim = v1Type.getDimSize(0); 528 529 // For rank 1, where both operands have *exactly* the same vector type, 530 // there is direct shuffle support in LLVM. Use it! 531 if (rank == 1 && v1Type == v2Type) { 532 Value llvmShuffleOp = rewriter.create<LLVM::ShuffleVectorOp>( 533 loc, adaptor.v1(), adaptor.v2(), maskArrayAttr); 534 rewriter.replaceOp(shuffleOp, llvmShuffleOp); 535 return success(); 536 } 537 538 // For all other cases, insert the individual values individually. 539 Type eltType; 540 llvm::errs() << llvmType << "\n"; 541 if (auto arrayType = llvmType.dyn_cast<LLVM::LLVMArrayType>()) 542 eltType = arrayType.getElementType(); 543 else 544 eltType = llvmType.cast<VectorType>().getElementType(); 545 Value insert = rewriter.create<LLVM::UndefOp>(loc, llvmType); 546 int64_t insPos = 0; 547 for (auto en : llvm::enumerate(maskArrayAttr)) { 548 int64_t extPos = en.value().cast<IntegerAttr>().getInt(); 549 Value value = adaptor.v1(); 550 if (extPos >= v1Dim) { 551 extPos -= v1Dim; 552 value = adaptor.v2(); 553 } 554 Value extract = extractOne(rewriter, *getTypeConverter(), loc, value, 555 eltType, rank, extPos); 556 insert = insertOne(rewriter, *getTypeConverter(), loc, insert, extract, 557 llvmType, rank, insPos++); 558 } 559 rewriter.replaceOp(shuffleOp, insert); 560 return success(); 561 } 562 }; 563 564 class VectorExtractElementOpConversion 565 : public ConvertOpToLLVMPattern<vector::ExtractElementOp> { 566 public: 567 using ConvertOpToLLVMPattern< 568 vector::ExtractElementOp>::ConvertOpToLLVMPattern; 569 570 LogicalResult 571 matchAndRewrite(vector::ExtractElementOp extractEltOp, OpAdaptor adaptor, 572 ConversionPatternRewriter &rewriter) const override { 573 auto vectorType = extractEltOp.getVectorType(); 574 auto llvmType = typeConverter->convertType(vectorType.getElementType()); 575 576 // Bail if result type cannot be lowered. 577 if (!llvmType) 578 return failure(); 579 580 rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>( 581 extractEltOp, llvmType, adaptor.vector(), adaptor.position()); 582 return success(); 583 } 584 }; 585 586 class VectorExtractOpConversion 587 : public ConvertOpToLLVMPattern<vector::ExtractOp> { 588 public: 589 using ConvertOpToLLVMPattern<vector::ExtractOp>::ConvertOpToLLVMPattern; 590 591 LogicalResult 592 matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor, 593 ConversionPatternRewriter &rewriter) const override { 594 auto loc = extractOp->getLoc(); 595 auto vectorType = extractOp.getVectorType(); 596 auto resultType = extractOp.getResult().getType(); 597 auto llvmResultType = typeConverter->convertType(resultType); 598 auto positionArrayAttr = extractOp.position(); 599 600 // Bail if result type cannot be lowered. 601 if (!llvmResultType) 602 return failure(); 603 604 // Extract entire vector. Should be handled by folder, but just to be safe. 605 if (positionArrayAttr.empty()) { 606 rewriter.replaceOp(extractOp, adaptor.vector()); 607 return success(); 608 } 609 610 // One-shot extraction of vector from array (only requires extractvalue). 611 if (resultType.isa<VectorType>()) { 612 Value extracted = rewriter.create<LLVM::ExtractValueOp>( 613 loc, llvmResultType, adaptor.vector(), positionArrayAttr); 614 rewriter.replaceOp(extractOp, extracted); 615 return success(); 616 } 617 618 // Potential extraction of 1-D vector from array. 619 auto *context = extractOp->getContext(); 620 Value extracted = adaptor.vector(); 621 auto positionAttrs = positionArrayAttr.getValue(); 622 if (positionAttrs.size() > 1) { 623 auto oneDVectorType = reducedVectorTypeBack(vectorType); 624 auto nMinusOnePositionAttrs = 625 ArrayAttr::get(context, positionAttrs.drop_back()); 626 extracted = rewriter.create<LLVM::ExtractValueOp>( 627 loc, typeConverter->convertType(oneDVectorType), extracted, 628 nMinusOnePositionAttrs); 629 } 630 631 // Remaining extraction of element from 1-D LLVM vector 632 auto position = positionAttrs.back().cast<IntegerAttr>(); 633 auto i64Type = IntegerType::get(rewriter.getContext(), 64); 634 auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 635 extracted = 636 rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant); 637 rewriter.replaceOp(extractOp, extracted); 638 639 return success(); 640 } 641 }; 642 643 /// Conversion pattern that turns a vector.fma on a 1-D vector 644 /// into an llvm.intr.fmuladd. This is a trivial 1-1 conversion. 645 /// This does not match vectors of n >= 2 rank. 646 /// 647 /// Example: 648 /// ``` 649 /// vector.fma %a, %a, %a : vector<8xf32> 650 /// ``` 651 /// is converted to: 652 /// ``` 653 /// llvm.intr.fmuladd %va, %va, %va: 654 /// (!llvm."<8 x f32>">, !llvm<"<8 x f32>">, !llvm<"<8 x f32>">) 655 /// -> !llvm."<8 x f32>"> 656 /// ``` 657 class VectorFMAOp1DConversion : public ConvertOpToLLVMPattern<vector::FMAOp> { 658 public: 659 using ConvertOpToLLVMPattern<vector::FMAOp>::ConvertOpToLLVMPattern; 660 661 LogicalResult 662 matchAndRewrite(vector::FMAOp fmaOp, OpAdaptor adaptor, 663 ConversionPatternRewriter &rewriter) const override { 664 VectorType vType = fmaOp.getVectorType(); 665 if (vType.getRank() != 1) 666 return failure(); 667 rewriter.replaceOpWithNewOp<LLVM::FMulAddOp>(fmaOp, adaptor.lhs(), 668 adaptor.rhs(), adaptor.acc()); 669 return success(); 670 } 671 }; 672 673 class VectorInsertElementOpConversion 674 : public ConvertOpToLLVMPattern<vector::InsertElementOp> { 675 public: 676 using ConvertOpToLLVMPattern<vector::InsertElementOp>::ConvertOpToLLVMPattern; 677 678 LogicalResult 679 matchAndRewrite(vector::InsertElementOp insertEltOp, OpAdaptor adaptor, 680 ConversionPatternRewriter &rewriter) const override { 681 auto vectorType = insertEltOp.getDestVectorType(); 682 auto llvmType = typeConverter->convertType(vectorType); 683 684 // Bail if result type cannot be lowered. 685 if (!llvmType) 686 return failure(); 687 688 rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>( 689 insertEltOp, llvmType, adaptor.dest(), adaptor.source(), 690 adaptor.position()); 691 return success(); 692 } 693 }; 694 695 class VectorInsertOpConversion 696 : public ConvertOpToLLVMPattern<vector::InsertOp> { 697 public: 698 using ConvertOpToLLVMPattern<vector::InsertOp>::ConvertOpToLLVMPattern; 699 700 LogicalResult 701 matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor, 702 ConversionPatternRewriter &rewriter) const override { 703 auto loc = insertOp->getLoc(); 704 auto sourceType = insertOp.getSourceType(); 705 auto destVectorType = insertOp.getDestVectorType(); 706 auto llvmResultType = typeConverter->convertType(destVectorType); 707 auto positionArrayAttr = insertOp.position(); 708 709 // Bail if result type cannot be lowered. 710 if (!llvmResultType) 711 return failure(); 712 713 // Overwrite entire vector with value. Should be handled by folder, but 714 // just to be safe. 715 if (positionArrayAttr.empty()) { 716 rewriter.replaceOp(insertOp, adaptor.source()); 717 return success(); 718 } 719 720 // One-shot insertion of a vector into an array (only requires insertvalue). 721 if (sourceType.isa<VectorType>()) { 722 Value inserted = rewriter.create<LLVM::InsertValueOp>( 723 loc, llvmResultType, adaptor.dest(), adaptor.source(), 724 positionArrayAttr); 725 rewriter.replaceOp(insertOp, inserted); 726 return success(); 727 } 728 729 // Potential extraction of 1-D vector from array. 730 auto *context = insertOp->getContext(); 731 Value extracted = adaptor.dest(); 732 auto positionAttrs = positionArrayAttr.getValue(); 733 auto position = positionAttrs.back().cast<IntegerAttr>(); 734 auto oneDVectorType = destVectorType; 735 if (positionAttrs.size() > 1) { 736 oneDVectorType = reducedVectorTypeBack(destVectorType); 737 auto nMinusOnePositionAttrs = 738 ArrayAttr::get(context, positionAttrs.drop_back()); 739 extracted = rewriter.create<LLVM::ExtractValueOp>( 740 loc, typeConverter->convertType(oneDVectorType), extracted, 741 nMinusOnePositionAttrs); 742 } 743 744 // Insertion of an element into a 1-D LLVM vector. 745 auto i64Type = IntegerType::get(rewriter.getContext(), 64); 746 auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position); 747 Value inserted = rewriter.create<LLVM::InsertElementOp>( 748 loc, typeConverter->convertType(oneDVectorType), extracted, 749 adaptor.source(), constant); 750 751 // Potential insertion of resulting 1-D vector into array. 752 if (positionAttrs.size() > 1) { 753 auto nMinusOnePositionAttrs = 754 ArrayAttr::get(context, positionAttrs.drop_back()); 755 inserted = rewriter.create<LLVM::InsertValueOp>(loc, llvmResultType, 756 adaptor.dest(), inserted, 757 nMinusOnePositionAttrs); 758 } 759 760 rewriter.replaceOp(insertOp, inserted); 761 return success(); 762 } 763 }; 764 765 /// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1. 766 /// 767 /// Example: 768 /// ``` 769 /// %d = vector.fma %a, %b, %c : vector<2x4xf32> 770 /// ``` 771 /// is rewritten into: 772 /// ``` 773 /// %r = splat %f0: vector<2x4xf32> 774 /// %va = vector.extractvalue %a[0] : vector<2x4xf32> 775 /// %vb = vector.extractvalue %b[0] : vector<2x4xf32> 776 /// %vc = vector.extractvalue %c[0] : vector<2x4xf32> 777 /// %vd = vector.fma %va, %vb, %vc : vector<4xf32> 778 /// %r2 = vector.insertvalue %vd, %r[0] : vector<4xf32> into vector<2x4xf32> 779 /// %va2 = vector.extractvalue %a2[1] : vector<2x4xf32> 780 /// %vb2 = vector.extractvalue %b2[1] : vector<2x4xf32> 781 /// %vc2 = vector.extractvalue %c2[1] : vector<2x4xf32> 782 /// %vd2 = vector.fma %va2, %vb2, %vc2 : vector<4xf32> 783 /// %r3 = vector.insertvalue %vd2, %r2[1] : vector<4xf32> into vector<2x4xf32> 784 /// // %r3 holds the final value. 785 /// ``` 786 class VectorFMAOpNDRewritePattern : public OpRewritePattern<FMAOp> { 787 public: 788 using OpRewritePattern<FMAOp>::OpRewritePattern; 789 790 LogicalResult matchAndRewrite(FMAOp op, 791 PatternRewriter &rewriter) const override { 792 auto vType = op.getVectorType(); 793 if (vType.getRank() < 2) 794 return failure(); 795 796 auto loc = op.getLoc(); 797 auto elemType = vType.getElementType(); 798 Value zero = rewriter.create<ConstantOp>(loc, elemType, 799 rewriter.getZeroAttr(elemType)); 800 Value desc = rewriter.create<SplatOp>(loc, vType, zero); 801 for (int64_t i = 0, e = vType.getShape().front(); i != e; ++i) { 802 Value extrLHS = rewriter.create<ExtractOp>(loc, op.lhs(), i); 803 Value extrRHS = rewriter.create<ExtractOp>(loc, op.rhs(), i); 804 Value extrACC = rewriter.create<ExtractOp>(loc, op.acc(), i); 805 Value fma = rewriter.create<FMAOp>(loc, extrLHS, extrRHS, extrACC); 806 desc = rewriter.create<InsertOp>(loc, fma, desc, i); 807 } 808 rewriter.replaceOp(op, desc); 809 return success(); 810 } 811 }; 812 813 // When ranks are different, InsertStridedSlice needs to extract a properly 814 // ranked vector from the destination vector into which to insert. This pattern 815 // only takes care of this part and forwards the rest of the conversion to 816 // another pattern that converts InsertStridedSlice for operands of the same 817 // rank. 818 // 819 // RewritePattern for InsertStridedSliceOp where source and destination vectors 820 // have different ranks. In this case: 821 // 1. the proper subvector is extracted from the destination vector 822 // 2. a new InsertStridedSlice op is created to insert the source in the 823 // destination subvector 824 // 3. the destination subvector is inserted back in the proper place 825 // 4. the op is replaced by the result of step 3. 826 // The new InsertStridedSlice from step 2. will be picked up by a 827 // `VectorInsertStridedSliceOpSameRankRewritePattern`. 828 class VectorInsertStridedSliceOpDifferentRankRewritePattern 829 : public OpRewritePattern<InsertStridedSliceOp> { 830 public: 831 using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 832 833 LogicalResult matchAndRewrite(InsertStridedSliceOp op, 834 PatternRewriter &rewriter) const override { 835 auto srcType = op.getSourceVectorType(); 836 auto dstType = op.getDestVectorType(); 837 838 if (op.offsets().getValue().empty()) 839 return failure(); 840 841 auto loc = op.getLoc(); 842 int64_t rankDiff = dstType.getRank() - srcType.getRank(); 843 assert(rankDiff >= 0); 844 if (rankDiff == 0) 845 return failure(); 846 847 int64_t rankRest = dstType.getRank() - rankDiff; 848 // Extract / insert the subvector of matching rank and InsertStridedSlice 849 // on it. 850 Value extracted = 851 rewriter.create<ExtractOp>(loc, op.dest(), 852 getI64SubArray(op.offsets(), /*dropFront=*/0, 853 /*dropBack=*/rankRest)); 854 // A different pattern will kick in for InsertStridedSlice with matching 855 // ranks. 856 auto stridedSliceInnerOp = rewriter.create<InsertStridedSliceOp>( 857 loc, op.source(), extracted, 858 getI64SubArray(op.offsets(), /*dropFront=*/rankDiff), 859 getI64SubArray(op.strides(), /*dropFront=*/0)); 860 rewriter.replaceOpWithNewOp<InsertOp>( 861 op, stridedSliceInnerOp.getResult(), op.dest(), 862 getI64SubArray(op.offsets(), /*dropFront=*/0, 863 /*dropBack=*/rankRest)); 864 return success(); 865 } 866 }; 867 868 // RewritePattern for InsertStridedSliceOp where source and destination vectors 869 // have the same rank. In this case, we reduce 870 // 1. the proper subvector is extracted from the destination vector 871 // 2. a new InsertStridedSlice op is created to insert the source in the 872 // destination subvector 873 // 3. the destination subvector is inserted back in the proper place 874 // 4. the op is replaced by the result of step 3. 875 // The new InsertStridedSlice from step 2. will be picked up by a 876 // `VectorInsertStridedSliceOpSameRankRewritePattern`. 877 class VectorInsertStridedSliceOpSameRankRewritePattern 878 : public OpRewritePattern<InsertStridedSliceOp> { 879 public: 880 using OpRewritePattern<InsertStridedSliceOp>::OpRewritePattern; 881 882 void initialize() { 883 // This pattern creates recursive InsertStridedSliceOp, but the recursion is 884 // bounded as the rank is strictly decreasing. 885 setHasBoundedRewriteRecursion(); 886 } 887 888 LogicalResult matchAndRewrite(InsertStridedSliceOp op, 889 PatternRewriter &rewriter) const override { 890 auto srcType = op.getSourceVectorType(); 891 auto dstType = op.getDestVectorType(); 892 893 if (op.offsets().getValue().empty()) 894 return failure(); 895 896 int64_t rankDiff = dstType.getRank() - srcType.getRank(); 897 assert(rankDiff >= 0); 898 if (rankDiff != 0) 899 return failure(); 900 901 if (srcType == dstType) { 902 rewriter.replaceOp(op, op.source()); 903 return success(); 904 } 905 906 int64_t offset = 907 op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 908 int64_t size = srcType.getShape().front(); 909 int64_t stride = 910 op.strides().getValue().front().cast<IntegerAttr>().getInt(); 911 912 auto loc = op.getLoc(); 913 Value res = op.dest(); 914 // For each slice of the source vector along the most major dimension. 915 for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 916 off += stride, ++idx) { 917 // 1. extract the proper subvector (or element) from source 918 Value extractedSource = extractOne(rewriter, loc, op.source(), idx); 919 if (extractedSource.getType().isa<VectorType>()) { 920 // 2. If we have a vector, extract the proper subvector from destination 921 // Otherwise we are at the element level and no need to recurse. 922 Value extractedDest = extractOne(rewriter, loc, op.dest(), off); 923 // 3. Reduce the problem to lowering a new InsertStridedSlice op with 924 // smaller rank. 925 extractedSource = rewriter.create<InsertStridedSliceOp>( 926 loc, extractedSource, extractedDest, 927 getI64SubArray(op.offsets(), /* dropFront=*/1), 928 getI64SubArray(op.strides(), /* dropFront=*/1)); 929 } 930 // 4. Insert the extractedSource into the res vector. 931 res = insertOne(rewriter, loc, extractedSource, res, off); 932 } 933 934 rewriter.replaceOp(op, res); 935 return success(); 936 } 937 }; 938 939 /// Returns the strides if the memory underlying `memRefType` has a contiguous 940 /// static layout. 941 static llvm::Optional<SmallVector<int64_t, 4>> 942 computeContiguousStrides(MemRefType memRefType) { 943 int64_t offset; 944 SmallVector<int64_t, 4> strides; 945 if (failed(getStridesAndOffset(memRefType, strides, offset))) 946 return None; 947 if (!strides.empty() && strides.back() != 1) 948 return None; 949 // If no layout or identity layout, this is contiguous by definition. 950 if (memRefType.getAffineMaps().empty() || 951 memRefType.getAffineMaps().front().isIdentity()) 952 return strides; 953 954 // Otherwise, we must determine contiguity form shapes. This can only ever 955 // work in static cases because MemRefType is underspecified to represent 956 // contiguous dynamic shapes in other ways than with just empty/identity 957 // layout. 958 auto sizes = memRefType.getShape(); 959 for (int index = 0, e = strides.size() - 1; index < e; ++index) { 960 if (ShapedType::isDynamic(sizes[index + 1]) || 961 ShapedType::isDynamicStrideOrOffset(strides[index]) || 962 ShapedType::isDynamicStrideOrOffset(strides[index + 1])) 963 return None; 964 if (strides[index] != strides[index + 1] * sizes[index + 1]) 965 return None; 966 } 967 return strides; 968 } 969 970 class VectorTypeCastOpConversion 971 : public ConvertOpToLLVMPattern<vector::TypeCastOp> { 972 public: 973 using ConvertOpToLLVMPattern<vector::TypeCastOp>::ConvertOpToLLVMPattern; 974 975 LogicalResult 976 matchAndRewrite(vector::TypeCastOp castOp, OpAdaptor adaptor, 977 ConversionPatternRewriter &rewriter) const override { 978 auto loc = castOp->getLoc(); 979 MemRefType sourceMemRefType = 980 castOp.getOperand().getType().cast<MemRefType>(); 981 MemRefType targetMemRefType = castOp.getType(); 982 983 // Only static shape casts supported atm. 984 if (!sourceMemRefType.hasStaticShape() || 985 !targetMemRefType.hasStaticShape()) 986 return failure(); 987 988 auto llvmSourceDescriptorTy = 989 adaptor.getOperands()[0].getType().dyn_cast<LLVM::LLVMStructType>(); 990 if (!llvmSourceDescriptorTy) 991 return failure(); 992 MemRefDescriptor sourceMemRef(adaptor.getOperands()[0]); 993 994 auto llvmTargetDescriptorTy = typeConverter->convertType(targetMemRefType) 995 .dyn_cast_or_null<LLVM::LLVMStructType>(); 996 if (!llvmTargetDescriptorTy) 997 return failure(); 998 999 // Only contiguous source buffers supported atm. 1000 auto sourceStrides = computeContiguousStrides(sourceMemRefType); 1001 if (!sourceStrides) 1002 return failure(); 1003 auto targetStrides = computeContiguousStrides(targetMemRefType); 1004 if (!targetStrides) 1005 return failure(); 1006 // Only support static strides for now, regardless of contiguity. 1007 if (llvm::any_of(*targetStrides, [](int64_t stride) { 1008 return ShapedType::isDynamicStrideOrOffset(stride); 1009 })) 1010 return failure(); 1011 1012 auto int64Ty = IntegerType::get(rewriter.getContext(), 64); 1013 1014 // Create descriptor. 1015 auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy); 1016 Type llvmTargetElementTy = desc.getElementPtrType(); 1017 // Set allocated ptr. 1018 Value allocated = sourceMemRef.allocatedPtr(rewriter, loc); 1019 allocated = 1020 rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated); 1021 desc.setAllocatedPtr(rewriter, loc, allocated); 1022 // Set aligned ptr. 1023 Value ptr = sourceMemRef.alignedPtr(rewriter, loc); 1024 ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr); 1025 desc.setAlignedPtr(rewriter, loc, ptr); 1026 // Fill offset 0. 1027 auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0); 1028 auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr); 1029 desc.setOffset(rewriter, loc, zero); 1030 1031 // Fill size and stride descriptors in memref. 1032 for (auto indexedSize : llvm::enumerate(targetMemRefType.getShape())) { 1033 int64_t index = indexedSize.index(); 1034 auto sizeAttr = 1035 rewriter.getIntegerAttr(rewriter.getIndexType(), indexedSize.value()); 1036 auto size = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, sizeAttr); 1037 desc.setSize(rewriter, loc, index, size); 1038 auto strideAttr = rewriter.getIntegerAttr(rewriter.getIndexType(), 1039 (*targetStrides)[index]); 1040 auto stride = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, strideAttr); 1041 desc.setStride(rewriter, loc, index, stride); 1042 } 1043 1044 rewriter.replaceOp(castOp, {desc}); 1045 return success(); 1046 } 1047 }; 1048 1049 class VectorPrintOpConversion : public ConvertOpToLLVMPattern<vector::PrintOp> { 1050 public: 1051 using ConvertOpToLLVMPattern<vector::PrintOp>::ConvertOpToLLVMPattern; 1052 1053 // Proof-of-concept lowering implementation that relies on a small 1054 // runtime support library, which only needs to provide a few 1055 // printing methods (single value for all data types, opening/closing 1056 // bracket, comma, newline). The lowering fully unrolls a vector 1057 // in terms of these elementary printing operations. The advantage 1058 // of this approach is that the library can remain unaware of all 1059 // low-level implementation details of vectors while still supporting 1060 // output of any shaped and dimensioned vector. Due to full unrolling, 1061 // this approach is less suited for very large vectors though. 1062 // 1063 // TODO: rely solely on libc in future? something else? 1064 // 1065 LogicalResult 1066 matchAndRewrite(vector::PrintOp printOp, OpAdaptor adaptor, 1067 ConversionPatternRewriter &rewriter) const override { 1068 Type printType = printOp.getPrintType(); 1069 1070 if (typeConverter->convertType(printType) == nullptr) 1071 return failure(); 1072 1073 // Make sure element type has runtime support. 1074 PrintConversion conversion = PrintConversion::None; 1075 VectorType vectorType = printType.dyn_cast<VectorType>(); 1076 Type eltType = vectorType ? vectorType.getElementType() : printType; 1077 Operation *printer; 1078 if (eltType.isF32()) { 1079 printer = 1080 LLVM::lookupOrCreatePrintF32Fn(printOp->getParentOfType<ModuleOp>()); 1081 } else if (eltType.isF64()) { 1082 printer = 1083 LLVM::lookupOrCreatePrintF64Fn(printOp->getParentOfType<ModuleOp>()); 1084 } else if (eltType.isIndex()) { 1085 printer = 1086 LLVM::lookupOrCreatePrintU64Fn(printOp->getParentOfType<ModuleOp>()); 1087 } else if (auto intTy = eltType.dyn_cast<IntegerType>()) { 1088 // Integers need a zero or sign extension on the operand 1089 // (depending on the source type) as well as a signed or 1090 // unsigned print method. Up to 64-bit is supported. 1091 unsigned width = intTy.getWidth(); 1092 if (intTy.isUnsigned()) { 1093 if (width <= 64) { 1094 if (width < 64) 1095 conversion = PrintConversion::ZeroExt64; 1096 printer = LLVM::lookupOrCreatePrintU64Fn( 1097 printOp->getParentOfType<ModuleOp>()); 1098 } else { 1099 return failure(); 1100 } 1101 } else { 1102 assert(intTy.isSignless() || intTy.isSigned()); 1103 if (width <= 64) { 1104 // Note that we *always* zero extend booleans (1-bit integers), 1105 // so that true/false is printed as 1/0 rather than -1/0. 1106 if (width == 1) 1107 conversion = PrintConversion::ZeroExt64; 1108 else if (width < 64) 1109 conversion = PrintConversion::SignExt64; 1110 printer = LLVM::lookupOrCreatePrintI64Fn( 1111 printOp->getParentOfType<ModuleOp>()); 1112 } else { 1113 return failure(); 1114 } 1115 } 1116 } else { 1117 return failure(); 1118 } 1119 1120 // Unroll vector into elementary print calls. 1121 int64_t rank = vectorType ? vectorType.getRank() : 0; 1122 emitRanks(rewriter, printOp, adaptor.source(), vectorType, printer, rank, 1123 conversion); 1124 emitCall(rewriter, printOp->getLoc(), 1125 LLVM::lookupOrCreatePrintNewlineFn( 1126 printOp->getParentOfType<ModuleOp>())); 1127 rewriter.eraseOp(printOp); 1128 return success(); 1129 } 1130 1131 private: 1132 enum class PrintConversion { 1133 // clang-format off 1134 None, 1135 ZeroExt64, 1136 SignExt64 1137 // clang-format on 1138 }; 1139 1140 void emitRanks(ConversionPatternRewriter &rewriter, Operation *op, 1141 Value value, VectorType vectorType, Operation *printer, 1142 int64_t rank, PrintConversion conversion) const { 1143 Location loc = op->getLoc(); 1144 if (rank == 0) { 1145 switch (conversion) { 1146 case PrintConversion::ZeroExt64: 1147 value = rewriter.create<ZeroExtendIOp>( 1148 loc, value, IntegerType::get(rewriter.getContext(), 64)); 1149 break; 1150 case PrintConversion::SignExt64: 1151 value = rewriter.create<SignExtendIOp>( 1152 loc, value, IntegerType::get(rewriter.getContext(), 64)); 1153 break; 1154 case PrintConversion::None: 1155 break; 1156 } 1157 emitCall(rewriter, loc, printer, value); 1158 return; 1159 } 1160 1161 emitCall(rewriter, loc, 1162 LLVM::lookupOrCreatePrintOpenFn(op->getParentOfType<ModuleOp>())); 1163 Operation *printComma = 1164 LLVM::lookupOrCreatePrintCommaFn(op->getParentOfType<ModuleOp>()); 1165 int64_t dim = vectorType.getDimSize(0); 1166 for (int64_t d = 0; d < dim; ++d) { 1167 auto reducedType = 1168 rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr; 1169 auto llvmType = typeConverter->convertType( 1170 rank > 1 ? reducedType : vectorType.getElementType()); 1171 Value nestedVal = extractOne(rewriter, *getTypeConverter(), loc, value, 1172 llvmType, rank, d); 1173 emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1, 1174 conversion); 1175 if (d != dim - 1) 1176 emitCall(rewriter, loc, printComma); 1177 } 1178 emitCall(rewriter, loc, 1179 LLVM::lookupOrCreatePrintCloseFn(op->getParentOfType<ModuleOp>())); 1180 } 1181 1182 // Helper to emit a call. 1183 static void emitCall(ConversionPatternRewriter &rewriter, Location loc, 1184 Operation *ref, ValueRange params = ValueRange()) { 1185 rewriter.create<LLVM::CallOp>(loc, TypeRange(), SymbolRefAttr::get(ref), 1186 params); 1187 } 1188 }; 1189 1190 /// Progressive lowering of ExtractStridedSliceOp to either: 1191 /// 1. express single offset extract as a direct shuffle. 1192 /// 2. extract + lower rank strided_slice + insert for the n-D case. 1193 class VectorExtractStridedSliceOpConversion 1194 : public OpRewritePattern<ExtractStridedSliceOp> { 1195 public: 1196 using OpRewritePattern<ExtractStridedSliceOp>::OpRewritePattern; 1197 1198 void initialize() { 1199 // This pattern creates recursive ExtractStridedSliceOp, but the recursion 1200 // is bounded as the rank is strictly decreasing. 1201 setHasBoundedRewriteRecursion(); 1202 } 1203 1204 LogicalResult matchAndRewrite(ExtractStridedSliceOp op, 1205 PatternRewriter &rewriter) const override { 1206 auto dstType = op.getType(); 1207 1208 assert(!op.offsets().getValue().empty() && "Unexpected empty offsets"); 1209 1210 int64_t offset = 1211 op.offsets().getValue().front().cast<IntegerAttr>().getInt(); 1212 int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt(); 1213 int64_t stride = 1214 op.strides().getValue().front().cast<IntegerAttr>().getInt(); 1215 1216 auto loc = op.getLoc(); 1217 auto elemType = dstType.getElementType(); 1218 assert(elemType.isSignlessIntOrIndexOrFloat()); 1219 1220 // Single offset can be more efficiently shuffled. 1221 if (op.offsets().getValue().size() == 1) { 1222 SmallVector<int64_t, 4> offsets; 1223 offsets.reserve(size); 1224 for (int64_t off = offset, e = offset + size * stride; off < e; 1225 off += stride) 1226 offsets.push_back(off); 1227 rewriter.replaceOpWithNewOp<ShuffleOp>(op, dstType, op.vector(), 1228 op.vector(), 1229 rewriter.getI64ArrayAttr(offsets)); 1230 return success(); 1231 } 1232 1233 // Extract/insert on a lower ranked extract strided slice op. 1234 Value zero = rewriter.create<ConstantOp>(loc, elemType, 1235 rewriter.getZeroAttr(elemType)); 1236 Value res = rewriter.create<SplatOp>(loc, dstType, zero); 1237 for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e; 1238 off += stride, ++idx) { 1239 Value one = extractOne(rewriter, loc, op.vector(), off); 1240 Value extracted = rewriter.create<ExtractStridedSliceOp>( 1241 loc, one, getI64SubArray(op.offsets(), /* dropFront=*/1), 1242 getI64SubArray(op.sizes(), /* dropFront=*/1), 1243 getI64SubArray(op.strides(), /* dropFront=*/1)); 1244 res = insertOne(rewriter, loc, extracted, res, idx); 1245 } 1246 rewriter.replaceOp(op, res); 1247 return success(); 1248 } 1249 }; 1250 1251 } // namespace 1252 1253 /// Populate the given list with patterns that convert from Vector to LLVM. 1254 void mlir::populateVectorToLLVMConversionPatterns( 1255 LLVMTypeConverter &converter, RewritePatternSet &patterns, 1256 bool reassociateFPReductions) { 1257 MLIRContext *ctx = converter.getDialect()->getContext(); 1258 patterns.add<VectorFMAOpNDRewritePattern, 1259 VectorInsertStridedSliceOpDifferentRankRewritePattern, 1260 VectorInsertStridedSliceOpSameRankRewritePattern, 1261 VectorExtractStridedSliceOpConversion>(ctx); 1262 patterns.add<VectorReductionOpConversion>(converter, reassociateFPReductions); 1263 patterns 1264 .add<VectorBitCastOpConversion, VectorShuffleOpConversion, 1265 VectorExtractElementOpConversion, VectorExtractOpConversion, 1266 VectorFMAOp1DConversion, VectorInsertElementOpConversion, 1267 VectorInsertOpConversion, VectorPrintOpConversion, 1268 VectorTypeCastOpConversion, 1269 VectorLoadStoreConversion<vector::LoadOp, vector::LoadOpAdaptor>, 1270 VectorLoadStoreConversion<vector::MaskedLoadOp, 1271 vector::MaskedLoadOpAdaptor>, 1272 VectorLoadStoreConversion<vector::StoreOp, vector::StoreOpAdaptor>, 1273 VectorLoadStoreConversion<vector::MaskedStoreOp, 1274 vector::MaskedStoreOpAdaptor>, 1275 VectorGatherOpConversion, VectorScatterOpConversion, 1276 VectorExpandLoadOpConversion, VectorCompressStoreOpConversion>( 1277 converter); 1278 // Transfer ops with rank > 1 are handled by VectorToSCF. 1279 populateVectorTransferLoweringPatterns(patterns, /*maxTransferRank=*/1); 1280 } 1281 1282 void mlir::populateVectorToLLVMMatrixConversionPatterns( 1283 LLVMTypeConverter &converter, RewritePatternSet &patterns) { 1284 patterns.add<VectorMatmulOpConversion>(converter); 1285 patterns.add<VectorFlatTransposeOpConversion>(converter); 1286 } 1287