[mlir][Linalg][Vector] Add forwarding patterns between linalg.copy and vector.transfer

[nicolasvasilache]

Summary

This revision adds custom rewrites for patterns that arise during linalg structured
ops vectorization. These patterns allow the composition of linalg promotion,
vectorization and removal of redundant copies.

The patterns are voluntarily limited and restrictive atm.
More robust behavior will be implemented once more powerful side effect modeling and
analyses are available on view/subview.

On the transfer_read side, the following pattern is rewritten:

%alloc = ...
[optional] %view = std.view %alloc ...
%subView = subview %allocOrView ...
[optional] linalg.fill(%allocOrView, %cst) ...
...
linalg.copy(%in, %subView) ...
vector.transfer_read %allocOrView[...], %cst ...

into

[unchanged] %alloc = ...
[unchanged] [optional] %view = std.view %alloc ...
[unchanged] [unchanged] %subView = subview %allocOrView ...
...
vector.transfer_read %in[...], %cst ...

On the transfer_write side, the following pattern is rewriten:

%alloc = ...
[optional] %view = std.view %alloc ...
%subView = subview %allocOrView...
...
vector.transfer_write %..., %allocOrView[...]
linalg.copy(%subView, %out)

[rriddle]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

162

I don't think this namespace is necessary.

184

nit: Drop else after return.

[aartbik]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

459

and missing

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

133

perhaps give it a comment:

Test for interleaved uses, conservatively erring on their existence

or something to show the conservative assumption (the may part)

193

without interleaved .... ?

229

comment reads a bit awkward...

261

as requested above by riddle

[ftynse]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

454

Nit: %A is not defined here, did you mean {%alloc or %view}. Same in the commit description.

458

Could you also mention what is this pattern rewritten _to_ ?

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

106

I wonder if we want to write some macro wrapper along the lines of

#ifdef DEBUG_TYPE
#define DBGS() \
  (dbgs() << '[' << DEBUG_TYPE << "] ")
#else
#define DBGS() \
  (dbgs())
#endif

and propose it for LLVM.

140

Please document the precondition somwhere

168

I don't understand this comment

261

Maybe this can be factored out into a function?

[nicolasvasilache]

Address review + cleanups.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

106

It hasn't bugged me enough to go and do it but it would make sense.

[nicolasvasilache]

Address review.

[NathanielMcVicar]

I believe this change broke the windows buildbot as well, just as an additional data point. http://lab.llvm.org:8011/builders/mlir-windows/builds/2485

1025.187 [461/64/2346] Building CXX object tools\mlir\lib\Dialect\Linalg\Transforms\CMakeFiles\obj.MLIRLinalgTransforms.dir\Loops.cpp.obj
FAILED: tools/mlir/lib/Dialect/Linalg/Transforms/CMakeFiles/obj.MLIRLinalgTransforms.dir/Loops.cpp.obj 
C:\PROGRA~2\MICROS~1\2017\COMMUN~1\VC\Tools\MSVC\1416~1.270\bin\Hostx64\x64\cl.exe  /nologo /TP -DBUILD_EXAMPLES -DGTEST_HAS_RTTI=0 -DMLIR_CUDA_CONVERSIONS_ENABLED=1 -DMLIR_ROCM_CONVERSIONS_ENABLED=1 -DUNICODE -D_CRT_NONSTDC_NO_DEPRECATE -D_CRT_NONSTDC_NO_WARNINGS -D_CRT_SECURE_NO_DEPRECATE -D_CRT_SECURE_NO_WARNINGS -D_HAS_EXCEPTIONS=0 -D_SCL_SECURE_NO_DEPRECATE -D_SCL_SECURE_NO_WARNINGS -D_UNICODE -D__STDC_CONSTANT_MACROS -D__STDC_FORMAT_MACROS -D__STDC_LIMIT_MACROS -Itools\mlir\lib\Dialect\Linalg\Transforms -IE:\build_slave\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\Linalg\Transforms -Iinclude -IE:\build_slave\mlir-x64-windows-ninja\llvm-project\llvm\include -IE:\build_slave\mlir-x64-windows-ninja\llvm-project\mlir\include -Itools\mlir\include /DWIN32 /D_WINDOWS   /Zc:inline /Zc:strictStrings /Oi /Zc:rvalueCast /W4 -wd4141 -wd4146 -wd4244 -wd4267 -wd4291 -wd4351 -wd4456 -wd4457 -wd4458 -wd4459 -wd4503 -wd4624 -wd4722 -wd4100 -wd4127 -wd4512 -wd4505 -wd4610 -wd4510 -wd4702 -wd4245 -wd4706 -wd4310 -wd4701 -wd4703 -wd4389 -wd4611 -wd4805 -wd4204 -wd4577 -wd4091 -wd4592 -wd4319 -wd4709 -wd4324 -w14062 -we4238 /Gw /MD /O2 /Ob2    /EHs-c- /GR- -UNDEBUG -std:c++14 /showIncludes /Fotools\mlir\lib\Dialect\Linalg\Transforms\CMakeFiles\obj.MLIRLinalgTransforms.dir\Loops.cpp.obj /Fdtools\mlir\lib\Dialect\Linalg\Transforms\CMakeFiles\obj.MLIRLinalgTransforms.dir\ /FS -c E:\build_slave\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\Linalg\Transforms\Loops.cpp
E:\build_slave\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/IR/PatternMatch.h(192): error C2664: 'mlir::LogicalResult mlir::OpRewritePattern<mlir::vector::TransferWriteOp>::match(SourceOp) const': cannot convert argument 1 from 'SourceOp' to 'mlir::Operation *'
        with
        [
            SourceOp=mlir::vector::TransferWriteOp
        ]
E:\build_slave\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/IR/PatternMatch.h(192): note: use of undefined type 'mlir::vector::TransferWriteOp'
E:\build_slave\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/Dialect/Linalg/Transforms/Transforms.h(20): note: see declaration of 'mlir::vector::TransferWriteOp'

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 //===- Transforms.h - Linalg transformations as patterns --------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_
 #define DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_
 
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 namespace mlir {
+namespace vector {
+
+class TransferReadOp;
+class TransferWriteOp;
+
+} // namespace vector
+
 namespace linalg {
 
 struct LinalgTilingOptions;
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
 using LinalgLoops = SmallVector<Operation *, 4>;
private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgMarker marker;
   /// Controls whether the pattern lowers to library calls, scf.for, affine.for
   /// or scf.parallel.
   LinalgLoweringType loweringType;
 };
 
 //===----------------------------------------------------------------------===//
+// Op-specific patterns.
+//===----------------------------------------------------------------------===//
+/// Match and rewrite for the pattern:
+/// ```
+///    %alloc = ...
+///    [optional] %view = std.view %alloc ...
+///    %subView = subview %allocOrView ...
+///    [optional] linalg.fill(%allocOrView, %cst) ...
+///    ...
+///    linalg.copy(%in, %subView) ...
+///    vector.transfer_read %allocOrView[...], %cst ...
+/// ```
+/// into
+/// ```
+///    [unchanged] %alloc = ...
+///    [unchanged] [optional] %view = std.view %alloc ...
+///    [unchanged] [unchanged] %subView = subview %allocOrView ...
+///    ...
+///    vector.transfer_read %in[...], %cst ...
+/// ```
+/// Where there is no interleaved use between linalg.copy and transfer_read as
+/// well as no interleaved use between linalg.fill and linalg.copy (if
+/// linalg.fill is specified).
+/// This is a custom rewrite to forward partial reads (with optional fills) to
+/// vector.transfer_read.
+struct LinalgCopyVTRForwardingPattern
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp xferOp,
+                                PatternRewriter &rewriter) const override;
+};
+
+/// Match and rewrite for the pattern:
+/// ```
+///    %alloc = ...
+///    [optional] %view = std.view %alloc ...
+///    %subView = subview %allocOrView...
+///    ...
+///    vector.transfer_write %..., %allocOrView[...]
+///    linalg.copy(%subView, %out)
+/// ```
+/// into
+/// ```
+///    [unchanged] %alloc = ...
+///    [unchanged] [optional] %view = std.view %alloc ...
+///    [unchanged] %subView = subview %allocOrView...
+///    ...
+///    vector.transfer_write %..., %out[...]
+/// ```
+/// Where there is no interleaved use between transfer_write and linalg.copy.
+/// This is a custom rewrite to forward partial writes to vector.transfer_write.
+struct LinalgCopyVTWForwardingPattern
+    : public OpRewritePattern<vector::TransferWriteOp> {
+  using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferWriteOp xferOp,
+                                PatternRewriter &rewriter) const override;
+};
+
+//===----------------------------------------------------------------------===//
 // Support for staged pattern application.
 //===----------------------------------------------------------------------===//
 /// Helper function to allow applying rewrite patterns, interleaved with more
 /// global transformations, in a staged fashion:
 ///   1. the first stage consists of a list of OwningRewritePatternList. Each
 ///   OwningRewritePatternList in this list is applied once, in order.
 ///   2. the second stage consists of a single OwningRewritePattern that is
 ///   applied greedily until convergence.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

void mlir::linalg::vectorizeLinalgOp(OpBuilder &builder, Operation *op) {
   assert(succeeded(vectorizeLinalgOpPrecondition(op)));
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op)) {
     // TODO: add a level of indirection to linalg.generic.
     if (convOp.padding())
       llvm_unreachable("Unexpected conv with padding");
   }
 
+  StringRef dbgPref = "\n[" DEBUG_TYPE "]: ";
+  (void)dbgPref;
   edsc::ScopedContext scope(builder, op->getLoc());
   if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
     // Vectorize fill as a vector.broadcast.
-    LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
-                         "]: Rewrite linalg.fill as vector.broadcast: "
-                      << *op << ":\n");
+    LLVM_DEBUG(dbgs() << dbgPref
+                      << "Rewrite linalg.fill as vector.broadcast: " << *op);
     Value memref = vector_type_cast(fillOp.getOutputBuffer(0));
     Value dst = std_load(memref);
     Value res = vector_broadcast(dst.getType(), fillOp.value());
     std_store(res, memref);
     return;
   }
 
   // Vectorize other ops as vector contraction (currently only matmul).
-  LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
-                       "]: Rewrite linalg op as vector.contract: "
-                    << *op << ":\n");
+  LLVM_DEBUG(dbgs() << dbgPref
+                    << "Rewrite linalg op as vector.contract: " << *op);
   auto linalgOp = cast<linalg::LinalgOp>(op);
   Value a = std_load(vector_type_cast(linalgOp.getInput(0)));
   Value b = std_load(vector_type_cast(linalgOp.getInput(1)));
   Value memref = vector_type_cast(linalgOp.getOutputBuffer(0));
   Value c = std_load(memref);
   Value res = vector_contract(a, b, c, linalgOp.indexing_maps(),
                               linalgOp.iterator_types());
   std_store(res, memref);
 }
+
+/// Check whether there is any interleaved use of any `values` between `firstOp`
+/// and `secondOp`. Conservatively return `true` if any op or value is in a
+/// different block.
+static bool mayExistInterleavedUses(Operation *firstOp, Operation *secondOp,
+                                    ValueRange values) {
+  StringRef dbgPref = "\n[" DEBUG_TYPE "]: ";
+  (void)dbgPref;
+  if (firstOp->getBlock() != secondOp->getBlock() ||
+      !firstOp->isBeforeInBlock(secondOp)) {
+    LLVM_DEBUG(llvm::dbgs()
+               << dbgPref << "interleavedUses precondition failed, firstOp: "
+               << *firstOp << ", second op: " << *secondOp);
+    return true;
+  }
+  for (auto v : values) {
+    for (auto &u : v.getUses()) {
+      Operation *owner = u.getOwner();
+      if (owner == firstOp || owner == secondOp)
+        continue;
+      // TODO: this is too conservative, use dominance info in the future.
+      if (owner->getBlock() == firstOp->getBlock() &&
+          (owner->isBeforeInBlock(firstOp) || secondOp->isBeforeInBlock(owner)))
+        continue;
+      LLVM_DEBUG(llvm::dbgs()
+                 << dbgPref << " found interleaved op " << *owner
+                 << ", firstOp: " << *firstOp << ", second op: " << *secondOp);
+      return true;
+    }
+  }
+  return false;
+}
+
+/// Return the unique subview use of `v` if it is indeed unique, null otherwise.
+static SubViewOp getSubViewUseIfUnique(Value v) {
+  SubViewOp subViewOp;
+  for (auto &u : v.getUses()) {
+    if (auto newSubViewOp = dyn_cast<SubViewOp>(u.getOwner())) {
+      if (subViewOp)
+        return SubViewOp();
+      subViewOp = newSubViewOp;
+    }
+  }
+  return subViewOp;
+}
+
+/// TODO: use interfaces, side-effects and aliasing analysis as appropriate,
+/// when available.
+LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
+    vector::TransferReadOp xferOp, PatternRewriter &rewriter) const {
+
+  // Transfer into `view`.
+  Value viewOrAlloc = xferOp.memref();
+  if (!viewOrAlloc.getDefiningOp<ViewOp>() &&
+      !viewOrAlloc.getDefiningOp<AllocOp>())
+    return failure();
+
+  StringRef dbgPref = "\n[" DEBUG_TYPE "]: VTRForwarding: ";
+  (void)dbgPref;
+  LLVM_DEBUG(llvm::dbgs() << dbgPref << viewOrAlloc);
+
+  // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`.
+  SubViewOp subViewOp = getSubViewUseIfUnique(viewOrAlloc);
+  if (!subViewOp)
+    return failure();
+  Value subView = subViewOp.getResult();
+  LLVM_DEBUG(llvm::dbgs() << dbgPref << "with subView " << subView);
+
+  // Find the copy into `subView` without interleaved uses.
+  CopyOp copyOp;
+  for (auto &u : subView.getUses()) {
+    if (auto newCopyOp = dyn_cast<CopyOp>(u.getOwner())) {
+      if (newCopyOp.getOutputBuffer(0) != subView)
+        continue;
+      LLVM_DEBUG(llvm::dbgs() << dbgPref << "copy candidate " << *newCopyOp);
+      if (mayExistInterleavedUses(newCopyOp, xferOp, {viewOrAlloc, subView}))
+        continue;
+      copyOp = newCopyOp;
+      break;
+    }
+  }
+  if (!copyOp)
+    return failure();
+  LLVM_DEBUG(llvm::dbgs() << dbgPref << "with copy " << *copyOp);
+
+  // Find the fill into `viewOrAlloc` without interleaved uses before the copy.
+  FillOp maybeFillOp;
+  for (auto &u : viewOrAlloc.getUses()) {
+    if (auto newFillOp = dyn_cast<FillOp>(u.getOwner())) {
+      if (newFillOp.getOutputBuffer(0) != viewOrAlloc)
+        continue;
+      LLVM_DEBUG(llvm::dbgs() << dbgPref << "fill candidate " << *newFillOp);
+      if (mayExistInterleavedUses(newFillOp, copyOp, {viewOrAlloc, subView}))
+        continue;
+      maybeFillOp = newFillOp;
+      break;
+    }
+  }
+  // Ensure padding matches.
+  if (maybeFillOp && xferOp.padding() != maybeFillOp.value())
+    return failure();
+  if (maybeFillOp)
+    LLVM_DEBUG(llvm::dbgs() << dbgPref << "with maybeFillOp " << *maybeFillOp);
+
+  // `in` is the subview that linalg.copy reads. Replace it.
+  Value in = copyOp.getInput(0);
+
+  Value res = rewriter.create<vector::TransferReadOp>(
+      xferOp.getLoc(), xferOp.getVectorType(), in, xferOp.indices(),
+      xferOp.permutation_map(), xferOp.padding(),
+      xferOp.masked() ? *xferOp.masked() : ArrayAttr());
+
+  if (maybeFillOp)
+    rewriter.eraseOp(maybeFillOp);
+  rewriter.eraseOp(copyOp);
+  rewriter.replaceOp(xferOp, res);
+
+  return success();
+}
+
+/// TODO: use interfaces, side-effects and aliasing analysis as appropriate,
+/// when available.
+LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
+    vector::TransferWriteOp xferOp, PatternRewriter &rewriter) const {
+  // Transfer into `viewOrAlloc`.
+  Value viewOrAlloc = xferOp.memref();
+  if (!viewOrAlloc.getDefiningOp<ViewOp>() &&
+      !viewOrAlloc.getDefiningOp<AllocOp>())
+    return failure();
+
+  // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`.
+  SubViewOp subViewOp = getSubViewUseIfUnique(viewOrAlloc);
+  if (!subViewOp)
+    return failure();
+  Value subView = subViewOp.getResult();
+
+  // Find the copy from `subView` without interleaved uses.
+  CopyOp copyOp;
+  for (auto &u : subViewOp.getResult().getUses()) {
+    if (auto newCopyOp = dyn_cast<CopyOp>(u.getOwner())) {
+      if (newCopyOp.getInput(0) != subView)
+        continue;
+      if (mayExistInterleavedUses(xferOp, newCopyOp, {viewOrAlloc, subView}))
+        continue;
+      copyOp = newCopyOp;
+      break;
+    }
+  }
+  if (!copyOp)
+    return failure();
+
+  // `out` is the subview copied into that we replace.
+  Value out = copyOp.getOutputBuffer(0);
+
+  // Forward vector.transfer into copy.
+  rewriter.create<vector::TransferWriteOp>(
+      xferOp.getLoc(), xferOp.vector(), out, xferOp.indices(),
+      xferOp.permutation_map(),
+      xferOp.masked() ? *xferOp.masked() : ArrayAttr());
+
+  rewriter.eraseOp(copyOp);
+  rewriter.eraseOp(xferOp);
+
+  return success();
+}

mlir/test/Dialect/Linalg/forward-vector-transfers.mlir

+// RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-vector-transfer-forwarding-patterns | FileCheck %s
+
+// CHECK-LABEL: testAllocRead
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: linalg.fill
+//   CHECK-NOT: linalg.copy
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_read %[[ARG0]]
+func @testAllocRead(%in: memref<? x f32>) -> vector<32 x f32> {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<32 x f32>
+  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
+  %0 = vector.transfer_read %alloc[%c0], %f0: memref<32 x f32>, vector<32 x f32>
+  dealloc %alloc : memref<32 x f32>
+  return %0: vector<32 x f32>
+}
+
+// CHECK-LABEL: testAllocFillRead
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: linalg.fill
+//   CHECK-NOT: linalg.copy
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_read %[[ARG0]]
+func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<32 x f32>
+  linalg.fill(%alloc, %f0): memref<32 x f32>, f32
+  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
+  %0 = vector.transfer_read %alloc[%c0], %f0: memref<32 x f32>, vector<32 x f32>
+  dealloc %alloc : memref<32 x f32>
+  return %0: vector<32 x f32>
+}
+
+// CHECK-LABEL: testViewRead
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: linalg.fill
+//   CHECK-NOT: linalg.copy
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_read %[[ARG0]]
+func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<128 x i8>
+  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
+  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
+  %0 = vector.transfer_read %view[%c0], %f0: memref<32 x f32>, vector<32 x f32>
+  dealloc %alloc : memref<128 x i8>
+  return %0: vector<32 x f32>
+}
+
+// CHECK-LABEL: testViewFillRead
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: linalg.fill
+//   CHECK-NOT: linalg.copy
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_read %[[ARG0]]
+func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<128 x i8>
+  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
+  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  linalg.fill(%view, %f0): memref<32 x f32>, f32
+  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
+  %0 = vector.transfer_read %view[%c0], %f0: memref<32 x f32>, vector<32 x f32>
+  dealloc %alloc : memref<128 x i8>
+  return %0: vector<32 x f32>
+}
+
+// CHECK-LABEL: testAllocWrite
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
+//  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: linalg.copy
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]
+func @testAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<32 x f32>
+  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  vector.transfer_write %vec, %alloc[%c0] : vector<32 x f32>, memref<32 x f32>
+  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
+  dealloc %alloc : memref<32 x f32>
+  return
+}
+
+// CHECK-LABEL: testViewWrite
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
+//  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: linalg.copy
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]
+func @testViewWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<128 x i8>
+  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
+  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  vector.transfer_write %vec, %view[%c0] : vector<32 x f32>, memref<32 x f32>
+  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
+  dealloc %alloc : memref<128 x i8>
+  return
+}
+
+///===--------------------------------------------------------------------===///
+// Negative tests
+///===--------------------------------------------------------------------===///
+
+// This should fail the rewrite due to mismatching fill and transfer read value.
+// CHECK-LABEL: failAllocFillRead
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: vector.transfer_read %[[ARG0]]
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: linalg.copy
+//       CHECK: vector.transfer_read %[[ALLOC]]
+func @failAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %f1 = constant 1.0: f32
+  %alloc = alloc() : memref<32 x f32>
+  linalg.fill(%alloc, %f0): memref<32 x f32>, f32
+  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
+  "some_interleaved_use"(%subview) : (memref<16 x f32>) -> ()
+  %0 = vector.transfer_read %alloc[%c0], %f1: memref<32 x f32>, vector<32 x f32>
+  dealloc %alloc : memref<32 x f32>
+  return %0: vector<32 x f32>
+}
+
+// This should fail the rewrite due to some interleaved use.
+// CHECK-LABEL: failAllocWrite
+//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
+//  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
+//   CHECK-NOT: vector.transfer_write %[[ARG0]], %[[ARG1]]
+//       CHECK: %[[ALLOC:.*]] = alloc
+//       CHECK: vector.transfer_write %[[ARG0]], %[[ALLOC]]
+//       CHECK: linalg.copy
+func @failAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
+  %c0 = constant 0: index
+  %f0 = constant 0.0: f32
+  %alloc = alloc() : memref<32 x f32>
+  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
+  vector.transfer_write %vec, %alloc[%c0] : vector<32 x f32>, memref<32 x f32>
+  "some_interleaved_use"(%subview) : (memref<16 x f32>) -> ()
+  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
+  dealloc %alloc : memref<32 x f32>
+  return
+}

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

 //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for testing Linalg transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
Option<bool> testMatmulToVectorPatterns2dTiling{
       *this, "test-matmul-to-vector-patterns-tile-2d",
       llvm::cl::desc(
           "Test a fused pass that applies patterns from matmul to vectors via "
           "2-d tiling"),
       llvm::cl::init(false)};
   Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
                                     llvm::cl::desc("Test promotion options"),
                                     llvm::cl::init(false)};
+  Option<bool> testVectorTransferForwardingPatterns{
+      *this, "test-vector-transfer-forwarding-patterns",
+      llvm::cl::desc(
+          "Test a fused pass that forwards linalg.copy to vector.transfer"),
+      llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//
static void applyPatterns(FuncOp funcOp) {
   applyPatternsAndFoldGreedily(funcOp, patterns);
 
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 }
 
-static OwningRewritePatternList
-getMatmulToVectorCanonicalizationPatterns(MLIRContext *context) {
-  OwningRewritePatternList patterns;
-  AffineApplyOp::getCanonicalizationPatterns(patterns, context);
-  AffineMinOp::getCanonicalizationPatterns(patterns, context);
-  AffineMaxOp::getCanonicalizationPatterns(patterns, context);
-  AllocOp::getCanonicalizationPatterns(patterns, context);
-  SubViewOp::getCanonicalizationPatterns(patterns, context);
-  ViewOp::getCanonicalizationPatterns(patterns, context);
-  MatmulOp::getCanonicalizationPatterns(patterns, context);
-  return patterns;
-}
-
 static void fillL1TilingAndMatmulToVectorPatterns(
     FuncOp funcOp, StringRef startMarker,
     SmallVectorImpl<OwningRewritePatternList> &patternsVector) {
   MLIRContext *context = funcOp.getContext();
   patternsVector.emplace_back(LinalgTilingPattern<MatmulOp>(
       context,
       LinalgTilingOptions().setTileSizes({8, 12, 16}).setInterchange({1, 0, 2}),
       LinalgMarker({startMarker}, "L1")));
patterns.insert<LinalgPromotionPattern<MatmulOp>>(
               },
               [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
                 copyCallBackFn(b, src, dst, true);
                 return success();
               }),
       LinalgMarker({"PROMOTE"}));
 }
 
-/// Apply transformations specified as patterns.
-void TestLinalgTransforms::runOnFunction() {
-  if (testPatterns) {
-    applyPatterns(getFunction());
-    return;
-  }
-  if (testPromotionOptions) {
-    OwningRewritePatternList patterns;
-    fillPromotionCallBackPatterns(&getContext(), patterns);
-    applyPatternsAndFoldGreedily(getFunction(), patterns);
-  } else {
-    SmallVector<OwningRewritePatternList, 4> stage1Patterns;
-    if (testMatmulToVectorPatterns1dTiling) {
-      fillL1TilingAndMatmulToVectorPatterns(getFunction(), "START",
-                                            stage1Patterns);
-    } else if (testMatmulToVectorPatterns2dTiling) {
-      stage1Patterns.emplace_back(
-          LinalgTilingPattern<MatmulOp>(&getContext(),
-                                        LinalgTilingOptions()
-                                            .setTileSizes({768, 264, 768})
-                                            .setInterchange({1, 2, 0}),
-                                        LinalgMarker({"START"}, "L2")));
-      fillL1TilingAndMatmulToVectorPatterns(getFunction(), "L2",
-                                            stage1Patterns);
-    }
-    OwningRewritePatternList stage2Patterns =
-        getMatmulToVectorCanonicalizationPatterns(&getContext());
-    applyStagedPatterns(getFunction(), stage1Patterns, stage2Patterns);
-  }
+static void
+applyMatmulToVectorPatterns(FuncOp funcOp,
+                            bool testMatmulToVectorPatterns1dTiling,
+                            bool testMatmulToVectorPatterns2dTiling) {
+  MLIRContext *ctx = funcOp.getContext();
+    fillL1TilingAndMatmulToVectorPatterns(funcOp, "START", stage1Patterns);
+        LinalgTilingPattern<MatmulOp>(ctx,
+    fillL1TilingAndMatmulToVectorPatterns(funcOp, "L2", stage1Patterns);
+      getLinalgTilingCanonicalizationPatterns(ctx);
+  applyStagedPatterns(funcOp, stage1Patterns, stage2Patterns);
 
-  // Drop the marker.
+static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
+  OwningRewritePatternList forwardPattern;
+  forwardPattern.insert<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
+  forwardPattern.insert<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
+  applyPatternsAndFoldGreedily(funcOp, forwardPattern);
+}
+
+/// Apply transformations specified as patterns.
+void TestLinalgTransforms::runOnFunction() {
+  auto lambda = [&](void *) {
     getFunction().walk([](LinalgOp op) {
       op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
+  };
+  std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
+
+  if (testPromotionOptions) {
+    OwningRewritePatternList patterns;
+    fillPromotionCallBackPatterns(&getContext(), patterns);
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+    return;
+  }
+  if (testPatterns)
+    return applyPatterns(getFunction());
+  if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
+    return applyMatmulToVectorPatterns(getFunction(),
+                                       testMatmulToVectorPatterns1dTiling,
+                                       testMatmulToVectorPatterns2dTiling);
+  if (testVectorTransferForwardingPatterns)
+    return applyVectorTransferForwardingPatterns(getFunction());
 }
 
 namespace mlir {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace mlir