1 //===- VectorToSCF.cpp - Conversion from Vector to mix of SCF and Std -----===// 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 // This file implements target-dependent lowering of vector transfer operations. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include <type_traits> 14 15 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h" 16 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h" 17 #include "mlir/Dialect/SCF/EDSC/Builders.h" 18 #include "mlir/Dialect/SCF/EDSC/Intrinsics.h" 19 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h" 20 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h" 21 #include "mlir/Dialect/Vector/VectorOps.h" 22 #include "mlir/IR/AffineExpr.h" 23 #include "mlir/IR/AffineMap.h" 24 #include "mlir/IR/Attributes.h" 25 #include "mlir/IR/Builders.h" 26 #include "mlir/IR/Location.h" 27 #include "mlir/IR/Matchers.h" 28 #include "mlir/IR/OperationSupport.h" 29 #include "mlir/IR/PatternMatch.h" 30 #include "mlir/IR/Types.h" 31 32 using namespace mlir; 33 using namespace mlir::edsc; 34 using namespace mlir::edsc::intrinsics; 35 using vector::TransferReadOp; 36 using vector::TransferWriteOp; 37 38 /// Helper class captures the common information needed to lower N>1-D vector 39 /// transfer operations (read and write). 40 /// On construction, this class opens an edsc::ScopedContext for simpler IR 41 /// manipulation. 42 /// In pseudo-IR, for an n-D vector_transfer_read such as: 43 /// 44 /// ``` 45 /// vector_transfer_read(%m, %offsets, identity_map, %fill) : 46 /// memref<(leading_dims) x (major_dims) x (minor_dims) x type>, 47 /// vector<(major_dims) x (minor_dims) x type> 48 /// ``` 49 /// 50 /// where rank(minor_dims) is the lower-level vector rank (e.g. 1 for LLVM or 51 /// higher). 52 /// 53 /// This is the entry point to emitting pseudo-IR resembling: 54 /// 55 /// ``` 56 /// %tmp = alloc(): memref<(major_dims) x vector<minor_dim x type>> 57 /// for (%ivs_major, {0}, {vector_shape}, {1}) { // (N-1)-D loop nest 58 /// if (any_of(%ivs_major + %offsets, <, major_dims)) { 59 /// %v = vector_transfer_read( 60 /// {%offsets_leading, %ivs_major + %offsets_major, %offsets_minor}, 61 /// %ivs_minor): 62 /// memref<(leading_dims) x (major_dims) x (minor_dims) x type>, 63 /// vector<(minor_dims) x type>; 64 /// store(%v, %tmp); 65 /// } else { 66 /// %v = splat(vector<(minor_dims) x type>, %fill) 67 /// store(%v, %tmp, %ivs_major); 68 /// } 69 /// } 70 /// %res = load(%tmp, %0): memref<(major_dims) x vector<minor_dim x type>>): 71 // vector<(major_dims) x (minor_dims) x type> 72 /// ``` 73 /// 74 template <typename ConcreteOp> 75 class NDTransferOpHelper { 76 public: 77 NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp) 78 : rewriter(rewriter), loc(xferOp.getLoc()), 79 scope(std::make_unique<ScopedContext>(rewriter, loc)), xferOp(xferOp), 80 op(xferOp.getOperation()) { 81 vectorType = xferOp.getVectorType(); 82 // TODO(ntv, ajcbik): when we go to k > 1-D vectors adapt minorRank. 83 minorRank = 1; 84 majorRank = vectorType.getRank() - minorRank; 85 leadingRank = xferOp.getMemRefType().getRank() - (majorRank + minorRank); 86 majorVectorType = 87 VectorType::get(vectorType.getShape().take_front(majorRank), 88 vectorType.getElementType()); 89 minorVectorType = 90 VectorType::get(vectorType.getShape().take_back(minorRank), 91 vectorType.getElementType()); 92 /// Memref of minor vector type is used for individual transfers. 93 memRefMinorVectorType = 94 MemRefType::get(majorVectorType.getShape(), minorVectorType, {}, 95 xferOp.getMemRefType().getMemorySpace()); 96 } 97 98 LogicalResult doReplace(); 99 100 private: 101 /// Creates the loop nest on the "major" dimensions and calls the 102 /// `loopBodyBuilder` lambda in the context of the loop nest. 103 template <typename Lambda> 104 void emitLoops(Lambda loopBodyBuilder); 105 106 /// Operate within the body of `emitLoops` to: 107 /// 1. Compute the indexings `majorIvs + majorOffsets`. 108 /// 2. Compute a boolean that determines whether the first `majorIvs.rank()` 109 /// dimensions `majorIvs + majorOffsets` are all within `memrefBounds`. 110 /// 3. Create an IfOp conditioned on the boolean in step 2. 111 /// 4. Call a `thenBlockBuilder` and an `elseBlockBuilder` to append 112 /// operations to the IfOp blocks as appropriate. 113 template <typename LambdaThen, typename LambdaElse> 114 void emitInBounds(ValueRange majorIvs, ValueRange majorOffsets, 115 MemRefBoundsCapture &memrefBounds, 116 LambdaThen thenBlockBuilder, LambdaElse elseBlockBuilder); 117 118 /// Common state to lower vector transfer ops. 119 PatternRewriter &rewriter; 120 Location loc; 121 std::unique_ptr<ScopedContext> scope; 122 ConcreteOp xferOp; 123 Operation *op; 124 // A vector transfer copies data between: 125 // - memref<(leading_dims) x (major_dims) x (minor_dims) x type> 126 // - vector<(major_dims) x (minor_dims) x type> 127 unsigned minorRank; // for now always 1 128 unsigned majorRank; // vector rank - minorRank 129 unsigned leadingRank; // memref rank - vector rank 130 VectorType vectorType; // vector<(major_dims) x (minor_dims) x type> 131 VectorType majorVectorType; // vector<(major_dims) x type> 132 VectorType minorVectorType; // vector<(minor_dims) x type> 133 MemRefType memRefMinorVectorType; // memref<vector<(minor_dims) x type>> 134 }; 135 136 template <typename ConcreteOp> 137 template <typename Lambda> 138 void NDTransferOpHelper<ConcreteOp>::emitLoops(Lambda loopBodyBuilder) { 139 /// Loop nest operates on the major dimensions 140 MemRefBoundsCapture memrefBoundsCapture(xferOp.memref()); 141 VectorBoundsCapture vectorBoundsCapture(majorVectorType); 142 auto majorLbs = vectorBoundsCapture.getLbs(); 143 auto majorUbs = vectorBoundsCapture.getUbs(); 144 auto majorSteps = vectorBoundsCapture.getSteps(); 145 SmallVector<Value, 8> majorIvs(vectorBoundsCapture.rank()); 146 AffineLoopNestBuilder(majorIvs, majorLbs, majorUbs, majorSteps)([&] { 147 ValueRange indices(xferOp.indices()); 148 loopBodyBuilder(majorIvs, indices.take_front(leadingRank), 149 indices.drop_front(leadingRank).take_front(majorRank), 150 indices.take_back(minorRank), memrefBoundsCapture); 151 }); 152 } 153 154 template <typename ConcreteOp> 155 template <typename LambdaThen, typename LambdaElse> 156 void NDTransferOpHelper<ConcreteOp>::emitInBounds( 157 ValueRange majorIvs, ValueRange majorOffsets, 158 MemRefBoundsCapture &memrefBounds, LambdaThen thenBlockBuilder, 159 LambdaElse elseBlockBuilder) { 160 Value inBounds = std_constant_int(/*value=*/1, /*width=*/1); 161 SmallVector<Value, 4> majorIvsPlusOffsets; 162 majorIvsPlusOffsets.reserve(majorIvs.size()); 163 for (auto it : llvm::zip(majorIvs, majorOffsets, memrefBounds.getUbs())) { 164 Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it); 165 using namespace mlir::edsc::op; 166 majorIvsPlusOffsets.push_back(iv + off); 167 Value inBounds2 = majorIvsPlusOffsets.back() < ub; 168 inBounds = inBounds && inBounds2; 169 } 170 171 auto ifOp = ScopedContext::getBuilderRef().create<scf::IfOp>( 172 ScopedContext::getLocation(), TypeRange{}, inBounds, 173 /*withElseRegion=*/std::is_same<ConcreteOp, TransferReadOp>()); 174 BlockBuilder(&ifOp.thenRegion().front(), 175 Append())([&] { thenBlockBuilder(majorIvsPlusOffsets); }); 176 if (std::is_same<ConcreteOp, TransferReadOp>()) 177 BlockBuilder(&ifOp.elseRegion().front(), 178 Append())([&] { elseBlockBuilder(majorIvsPlusOffsets); }); 179 } 180 181 template <> 182 LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() { 183 Value alloc = std_alloc(memRefMinorVectorType); 184 185 emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets, 186 ValueRange majorOffsets, ValueRange minorOffsets, 187 MemRefBoundsCapture &memrefBounds) { 188 // If in-bounds, index into memref and lower to 1-D transfer read. 189 auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) { 190 auto map = AffineMap::getMinorIdentityMap( 191 xferOp.getMemRefType().getRank(), minorRank, xferOp.getContext()); 192 // Lower to 1-D vector_transfer_read and let recursion handle it. 193 Value memref = xferOp.memref(); 194 SmallVector<Value, 8> indexing; 195 indexing.reserve(leadingRank + majorRank + minorRank); 196 indexing.append(leadingOffsets.begin(), leadingOffsets.end()); 197 indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end()); 198 indexing.append(minorOffsets.begin(), minorOffsets.end()); 199 auto loaded1D = 200 vector_transfer_read(minorVectorType, memref, indexing, 201 AffineMapAttr::get(map), xferOp.padding()); 202 // Store the 1-D vector. 203 std_store(loaded1D, alloc, majorIvs); 204 }; 205 // If out-of-bounds, just store a splatted vector. 206 auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) { 207 auto vector = std_splat(minorVectorType, xferOp.padding()); 208 std_store(vector, alloc, majorIvs); 209 }; 210 emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder, 211 elseBlockBuilder); 212 }); 213 214 Value loaded = 215 std_load(vector_type_cast(MemRefType::get({}, vectorType), alloc)); 216 rewriter.replaceOp(op, loaded); 217 218 return success(); 219 } 220 221 template <> 222 LogicalResult NDTransferOpHelper<TransferWriteOp>::doReplace() { 223 Value alloc = std_alloc(memRefMinorVectorType); 224 225 std_store(xferOp.vector(), 226 vector_type_cast(MemRefType::get({}, vectorType), alloc)); 227 228 emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets, 229 ValueRange majorOffsets, ValueRange minorOffsets, 230 MemRefBoundsCapture &memrefBounds) { 231 auto thenBlockBuilder = [&](ValueRange majorIvsPlusOffsets) { 232 // Lower to 1-D vector_transfer_write and let recursion handle it. 233 Value loaded1D = std_load(alloc, majorIvs); 234 auto map = AffineMap::getMinorIdentityMap( 235 xferOp.getMemRefType().getRank(), minorRank, xferOp.getContext()); 236 SmallVector<Value, 8> indexing; 237 indexing.reserve(leadingRank + majorRank + minorRank); 238 indexing.append(leadingOffsets.begin(), leadingOffsets.end()); 239 indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end()); 240 indexing.append(minorOffsets.begin(), minorOffsets.end()); 241 vector_transfer_write(loaded1D, xferOp.memref(), indexing, 242 AffineMapAttr::get(map)); 243 }; 244 // Don't write anything when out of bounds. 245 auto elseBlockBuilder = [&](ValueRange majorIvsPlusOffsets) {}; 246 emitInBounds(majorIvs, majorOffsets, memrefBounds, thenBlockBuilder, 247 elseBlockBuilder); 248 }); 249 250 rewriter.eraseOp(op); 251 252 return success(); 253 } 254 255 /// Analyzes the `transfer` to find an access dimension along the fastest remote 256 /// MemRef dimension. If such a dimension with coalescing properties is found, 257 /// `pivs` and `vectorBoundsCapture` are swapped so that the invocation of 258 /// LoopNestBuilder captures it in the innermost loop. 259 template <typename TransferOpTy> 260 static int computeCoalescedIndex(TransferOpTy transfer) { 261 // rank of the remote memory access, coalescing behavior occurs on the 262 // innermost memory dimension. 263 auto remoteRank = transfer.getMemRefType().getRank(); 264 // Iterate over the results expressions of the permutation map to determine 265 // the loop order for creating pointwise copies between remote and local 266 // memories. 267 int coalescedIdx = -1; 268 auto exprs = transfer.permutation_map().getResults(); 269 for (auto en : llvm::enumerate(exprs)) { 270 auto dim = en.value().template dyn_cast<AffineDimExpr>(); 271 if (!dim) { 272 continue; 273 } 274 auto memRefDim = dim.getPosition(); 275 if (memRefDim == remoteRank - 1) { 276 // memRefDim has coalescing properties, it should be swapped in the last 277 // position. 278 assert(coalescedIdx == -1 && "Unexpected > 1 coalesced indices"); 279 coalescedIdx = en.index(); 280 } 281 } 282 return coalescedIdx; 283 } 284 285 /// Emits remote memory accesses that are clipped to the boundaries of the 286 /// MemRef. 287 template <typename TransferOpTy> 288 static SmallVector<Value, 8> 289 clip(TransferOpTy transfer, MemRefBoundsCapture &bounds, ArrayRef<Value> ivs) { 290 using namespace mlir::edsc; 291 292 Value zero(std_constant_index(0)), one(std_constant_index(1)); 293 SmallVector<Value, 8> memRefAccess(transfer.indices()); 294 SmallVector<Value, 8> clippedScalarAccessExprs(memRefAccess.size()); 295 // Indices accessing to remote memory are clipped and their expressions are 296 // returned in clippedScalarAccessExprs. 297 for (unsigned memRefDim = 0; memRefDim < clippedScalarAccessExprs.size(); 298 ++memRefDim) { 299 // Linear search on a small number of entries. 300 int loopIndex = -1; 301 auto exprs = transfer.permutation_map().getResults(); 302 for (auto en : llvm::enumerate(exprs)) { 303 auto expr = en.value(); 304 auto dim = expr.template dyn_cast<AffineDimExpr>(); 305 // Sanity check. 306 assert( 307 (dim || expr.template cast<AffineConstantExpr>().getValue() == 0) && 308 "Expected dim or 0 in permutationMap"); 309 if (dim && memRefDim == dim.getPosition()) { 310 loopIndex = en.index(); 311 break; 312 } 313 } 314 315 // We cannot distinguish atm between unrolled dimensions that implement 316 // the "always full" tile abstraction and need clipping from the other 317 // ones. So we conservatively clip everything. 318 using namespace edsc::op; 319 auto N = bounds.ub(memRefDim); 320 auto i = memRefAccess[memRefDim]; 321 if (loopIndex < 0) { 322 auto N_minus_1 = N - one; 323 auto select_1 = std_select(i < N, i, N_minus_1); 324 clippedScalarAccessExprs[memRefDim] = 325 std_select(i < zero, zero, select_1); 326 } else { 327 auto ii = ivs[loopIndex]; 328 auto i_plus_ii = i + ii; 329 auto N_minus_1 = N - one; 330 auto select_1 = std_select(i_plus_ii < N, i_plus_ii, N_minus_1); 331 clippedScalarAccessExprs[memRefDim] = 332 std_select(i_plus_ii < zero, zero, select_1); 333 } 334 } 335 336 return clippedScalarAccessExprs; 337 } 338 339 namespace { 340 341 /// Implements lowering of TransferReadOp and TransferWriteOp to a 342 /// proper abstraction for the hardware. 343 /// 344 /// For now, we only emit a simple loop nest that performs clipped pointwise 345 /// copies from a remote to a locally allocated memory. 346 /// 347 /// Consider the case: 348 /// 349 /// ```mlir 350 /// // Read the slice `%A[%i0, %i1:%i1+256, %i2:%i2+32]` into 351 /// // vector<32x256xf32> and pad with %f0 to handle the boundary case: 352 /// %f0 = constant 0.0f : f32 353 /// scf.for %i0 = 0 to %0 { 354 /// scf.for %i1 = 0 to %1 step %c256 { 355 /// scf.for %i2 = 0 to %2 step %c32 { 356 /// %v = vector.transfer_read %A[%i0, %i1, %i2], %f0 357 /// {permutation_map: (d0, d1, d2) -> (d2, d1)} : 358 /// memref<?x?x?xf32>, vector<32x256xf32> 359 /// }}} 360 /// ``` 361 /// 362 /// The rewriters construct loop and indices that access MemRef A in a pattern 363 /// resembling the following (while guaranteeing an always full-tile 364 /// abstraction): 365 /// 366 /// ```mlir 367 /// scf.for %d2 = 0 to %c256 { 368 /// scf.for %d1 = 0 to %c32 { 369 /// %s = %A[%i0, %i1 + %d1, %i2 + %d2] : f32 370 /// %tmp[%d2, %d1] = %s 371 /// } 372 /// } 373 /// ``` 374 /// 375 /// In the current state, only a clipping transfer is implemented by `clip`, 376 /// which creates individual indexing expressions of the form: 377 /// 378 /// ```mlir-dsc 379 /// auto condMax = i + ii < N; 380 /// auto max = std_select(condMax, i + ii, N - one) 381 /// auto cond = i + ii < zero; 382 /// std_select(cond, zero, max); 383 /// ``` 384 /// 385 /// In the future, clipping should not be the only way and instead we should 386 /// load vectors + mask them. Similarly on the write side, load/mask/store for 387 /// implementing RMW behavior. 388 /// 389 /// Lowers TransferOp into a combination of: 390 /// 1. local memory allocation; 391 /// 2. perfect loop nest over: 392 /// a. scalar load/stores from local buffers (viewed as a scalar memref); 393 /// a. scalar store/load to original memref (with clipping). 394 /// 3. vector_load/store 395 /// 4. local memory deallocation. 396 /// Minor variations occur depending on whether a TransferReadOp or 397 /// a TransferWriteOp is rewritten. 398 template <typename TransferOpTy> 399 struct VectorTransferRewriter : public RewritePattern { 400 explicit VectorTransferRewriter(MLIRContext *context) 401 : RewritePattern(TransferOpTy::getOperationName(), 1, context) {} 402 403 /// Used for staging the transfer in a local scalar buffer. 404 MemRefType tmpMemRefType(TransferOpTy transfer) const { 405 auto vectorType = transfer.getVectorType(); 406 return MemRefType::get(vectorType.getShape(), vectorType.getElementType(), 407 {}, 0); 408 } 409 410 /// Performs the rewrite. 411 LogicalResult matchAndRewrite(Operation *op, 412 PatternRewriter &rewriter) const override; 413 }; 414 415 /// Lowers TransferReadOp into a combination of: 416 /// 1. local memory allocation; 417 /// 2. perfect loop nest over: 418 /// a. scalar load from local buffers (viewed as a scalar memref); 419 /// a. scalar store to original memref (with clipping). 420 /// 3. vector_load from local buffer (viewed as a memref<1 x vector>); 421 /// 4. local memory deallocation. 422 /// 423 /// Lowers the data transfer part of a TransferReadOp while ensuring no 424 /// out-of-bounds accesses are possible. Out-of-bounds behavior is handled by 425 /// clipping. This means that a given value in memory can be read multiple 426 /// times and concurrently. 427 /// 428 /// Important notes about clipping and "full-tiles only" abstraction: 429 /// ================================================================= 430 /// When using clipping for dealing with boundary conditions, the same edge 431 /// value will appear multiple times (a.k.a edge padding). This is fine if the 432 /// subsequent vector operations are all data-parallel but **is generally 433 /// incorrect** in the presence of reductions or extract operations. 434 /// 435 /// More generally, clipping is a scalar abstraction that is expected to work 436 /// fine as a baseline for CPUs and GPUs but not for vector_load and DMAs. 437 /// To deal with real vector_load and DMAs, a "padded allocation + view" 438 /// abstraction with the ability to read out-of-memref-bounds (but still within 439 /// the allocated region) is necessary. 440 /// 441 /// Whether using scalar loops or vector_load/DMAs to perform the transfer, 442 /// junk values will be materialized in the vectors and generally need to be 443 /// filtered out and replaced by the "neutral element". This neutral element is 444 /// op-dependent so, in the future, we expect to create a vector filter and 445 /// apply it to a splatted constant vector with the proper neutral element at 446 /// each ssa-use. This filtering is not necessary for pure data-parallel 447 /// operations. 448 /// 449 /// In the case of vector_store/DMAs, Read-Modify-Write will be required, which 450 /// also have concurrency implications. Note that by using clipped scalar stores 451 /// in the presence of data-parallel only operations, we generate code that 452 /// writes the same value multiple time on the edge locations. 453 /// 454 /// TODO(ntv): implement alternatives to clipping. 455 /// TODO(ntv): support non-data-parallel operations. 456 457 /// Performs the rewrite. 458 template <> 459 LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite( 460 Operation *op, PatternRewriter &rewriter) const { 461 using namespace mlir::edsc::op; 462 463 TransferReadOp transfer = cast<TransferReadOp>(op); 464 if (AffineMap::isMinorIdentity(transfer.permutation_map())) { 465 // If > 1D, emit a bunch of loops around 1-D vector transfers. 466 if (transfer.getVectorType().getRank() > 1) 467 return NDTransferOpHelper<TransferReadOp>(rewriter, transfer).doReplace(); 468 // If 1-D this is now handled by the target-specific lowering. 469 if (transfer.getVectorType().getRank() == 1) 470 return failure(); 471 } 472 473 // Conservative lowering to scalar load / stores. 474 // 1. Setup all the captures. 475 ScopedContext scope(rewriter, transfer.getLoc()); 476 StdIndexedValue remote(transfer.memref()); 477 MemRefBoundsCapture memRefBoundsCapture(transfer.memref()); 478 VectorBoundsCapture vectorBoundsCapture(transfer.vector()); 479 int coalescedIdx = computeCoalescedIndex(transfer); 480 // Swap the vectorBoundsCapture which will reorder loop bounds. 481 if (coalescedIdx >= 0) 482 vectorBoundsCapture.swapRanges(vectorBoundsCapture.rank() - 1, 483 coalescedIdx); 484 485 auto lbs = vectorBoundsCapture.getLbs(); 486 auto ubs = vectorBoundsCapture.getUbs(); 487 SmallVector<Value, 8> steps; 488 steps.reserve(vectorBoundsCapture.getSteps().size()); 489 for (auto step : vectorBoundsCapture.getSteps()) 490 steps.push_back(std_constant_index(step)); 491 492 // 2. Emit alloc-copy-load-dealloc. 493 Value tmp = std_alloc(tmpMemRefType(transfer)); 494 StdIndexedValue local(tmp); 495 Value vec = vector_type_cast(tmp); 496 SmallVector<Value, 8> ivs(lbs.size()); 497 LoopNestBuilder(ivs, lbs, ubs, steps)([&] { 498 // Swap the ivs which will reorder memory accesses. 499 if (coalescedIdx >= 0) 500 std::swap(ivs.back(), ivs[coalescedIdx]); 501 // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist). 502 local(ivs) = remote(clip(transfer, memRefBoundsCapture, ivs)); 503 }); 504 Value vectorValue = std_load(vec); 505 (std_dealloc(tmp)); // vexing parse 506 507 // 3. Propagate. 508 rewriter.replaceOp(op, vectorValue); 509 return success(); 510 } 511 512 /// Lowers TransferWriteOp into a combination of: 513 /// 1. local memory allocation; 514 /// 2. vector_store to local buffer (viewed as a memref<1 x vector>); 515 /// 3. perfect loop nest over: 516 /// a. scalar load from local buffers (viewed as a scalar memref); 517 /// a. scalar store to original memref (with clipping). 518 /// 4. local memory deallocation. 519 /// 520 /// More specifically, lowers the data transfer part while ensuring no 521 /// out-of-bounds accesses are possible. Out-of-bounds behavior is handled by 522 /// clipping. This means that a given value in memory can be written to multiple 523 /// times and concurrently. 524 /// 525 /// See `Important notes about clipping and full-tiles only abstraction` in the 526 /// description of `readClipped` above. 527 /// 528 /// TODO(ntv): implement alternatives to clipping. 529 /// TODO(ntv): support non-data-parallel operations. 530 template <> 531 LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite( 532 Operation *op, PatternRewriter &rewriter) const { 533 using namespace edsc::op; 534 535 TransferWriteOp transfer = cast<TransferWriteOp>(op); 536 if (AffineMap::isMinorIdentity(transfer.permutation_map())) { 537 // If > 1D, emit a bunch of loops around 1-D vector transfers. 538 if (transfer.getVectorType().getRank() > 1) 539 return NDTransferOpHelper<TransferWriteOp>(rewriter, transfer) 540 .doReplace(); 541 // If 1-D this is now handled by the target-specific lowering. 542 if (transfer.getVectorType().getRank() == 1) 543 return failure(); 544 } 545 546 // 1. Setup all the captures. 547 ScopedContext scope(rewriter, transfer.getLoc()); 548 StdIndexedValue remote(transfer.memref()); 549 MemRefBoundsCapture memRefBoundsCapture(transfer.memref()); 550 Value vectorValue(transfer.vector()); 551 VectorBoundsCapture vectorBoundsCapture(transfer.vector()); 552 int coalescedIdx = computeCoalescedIndex(transfer); 553 // Swap the vectorBoundsCapture which will reorder loop bounds. 554 if (coalescedIdx >= 0) 555 vectorBoundsCapture.swapRanges(vectorBoundsCapture.rank() - 1, 556 coalescedIdx); 557 558 auto lbs = vectorBoundsCapture.getLbs(); 559 auto ubs = vectorBoundsCapture.getUbs(); 560 SmallVector<Value, 8> steps; 561 steps.reserve(vectorBoundsCapture.getSteps().size()); 562 for (auto step : vectorBoundsCapture.getSteps()) 563 steps.push_back(std_constant_index(step)); 564 565 // 2. Emit alloc-store-copy-dealloc. 566 Value tmp = std_alloc(tmpMemRefType(transfer)); 567 StdIndexedValue local(tmp); 568 Value vec = vector_type_cast(tmp); 569 std_store(vectorValue, vec); 570 SmallVector<Value, 8> ivs(lbs.size()); 571 LoopNestBuilder(ivs, lbs, ubs, steps)([&] { 572 // Swap the ivs which will reorder memory accesses. 573 if (coalescedIdx >= 0) 574 std::swap(ivs.back(), ivs[coalescedIdx]); 575 // Computes clippedScalarAccessExprs in the loop nest scope (ivs exist). 576 remote(clip(transfer, memRefBoundsCapture, ivs)) = local(ivs); 577 }); 578 (std_dealloc(tmp)); // vexing parse... 579 580 rewriter.eraseOp(op); 581 return success(); 582 } 583 584 } // namespace 585 586 void mlir::populateVectorToSCFConversionPatterns( 587 OwningRewritePatternList &patterns, MLIRContext *context) { 588 patterns.insert<VectorTransferRewriter<vector::TransferReadOp>, 589 VectorTransferRewriter<vector::TransferWriteOp>>(context); 590 } 591