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