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