[mlir][linalg] Add vectorization for element-wise linalg ops

[ThomasRaoux]

Summary

Add support for vectorization for linalg.generic representing element-wise ops. Those are converted to transfer_read + vector ops + transfer_write.
Also re-organize the vectorization tests to be together.

Implementation derived from @burmako's work.

[nicolasvasilache]

Very cool, thanks for pushing this @ThomasRaoux !

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

92

I thought we had some LLVM interface already for elementwise like (ElementsMappable something from @silvas' work?)

106

if (llvm::any_of(genericOp.getOutputIndexingMaps(), lambda)) return false;``` ?

110

same

135

Ah cool that the insertion point has a loc now, can you just modify the ScopedContext to just take a builder and simplify uses in this file?

140

I'd move up 2 lines to keep C++ together and meta-programmed IR together.

170

early return and drop the else to reduce indentation entropy?

207

The type-only logic here seems like it could be refactored into VectorType?

235

return llvm::TypeSwitch<Operation*, Operation*>(...)
.Case(...)
.Default(...);

[aartbik]

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

124

note that load and transfer-read are slightly different terms (see doc),
so perhaps call this transferLoadVector?

134

same, transferStoreVector?

150

"to reasds" and "to writes" at a quick glance reads like bad grammar, but then you realize it is used as a noun

so how about "to transfer reads" and "to transfer writes" to make that more clear right away

[aartbik]

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

124

Ooop, I meant transferRead and transferWrite obviously, not Load/Store ;-)

[ThomasRaoux]

Address review comments.

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

92

Well, I want to restrict it to operations supported below so I'm not sure there is a better solution?

106

There is no getOutputIndexingMaps in linalg op. I can get a list of all indexingmaps but I want to distinguish the input and outputs so I couldn't find a better way to do that. Any ideas?

207

Good point, changed it to just skip if none of the types are vector, it makes logic simpler and should be more generic.

[silvas]

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

92

Just left a comment below. This should fit nicely with ElementwiseMappable.

234

You can do this for all ops that have ElementwiseMappable trait using a single piece of code like we do for elementwise on tensors ->linalg.generic: https://github.com/llvm/llvm-project/blob/f24c642178a5a31aa1c44585537c08fdfc3fdfd5/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp#L53
You just copy all the attributes and convert all the types to vector types.

The invariants of ElementwiseMappable guarantee that this can always be done: https://github.com/llvm/llvm-project/blob/f24c642178a5a31aa1c44585537c08fdfc3fdfd5/mlir/include/mlir/IR/OpDefinition.h#L1243

[mravishankar]

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

107

In general this could be a permutation. Is that not supported?

111

This could also be a permuted projection. Is that a todo?

[ThomasRaoux]

Use ElementwiseMappable.

[ThomasRaoux]

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

107

Yes I think permutation could be supported, I think we can just propagate the affineMap into the transfer read/write. Do you think this needs to go into this patch?

111

I was trying to make it work with permuted projection but it isn't trivial as the type of the memref may not be compatible with the type of the vector we want out of the transfer_read.
For instance if we have a map like affine_map<(d0, d1) -> (d0)> with 1D memref I would have to load the vector as 1D and then broadcast it along the first dimension which isn't supported by the broadcast.
Maybe I'm missing something, I can add a TODO if you want.

234

Thanks Sean! I used ElementwiseMappable and it does clean up significantly.

[ThomasRaoux]

Adding todo to relax indexing map restrictions.

[pravnar]

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

206

Typo? Probably should be "broadcastIfNeeded" :)

[ThomasRaoux]

Fix typo

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

206

Oops, yes. Fixed it. Thanks.

[burmako]

Thanks a lot for your work! This will be really useful :)

I'd also like to share the credit that you kindly gave to me with fedelebron@ and agrue@.

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

100

Can you talk more about the context behind the definition of isElementwise, e.g. in the comment above?

I'm asking because in one of the internal tests we have affine_map<(d0, d1) -> (d0)> as one of the supported indexing maps, and, if I understand correctly, that doesn't seem to work with isElementwise.

This comment doesn't need to block this patch (thanks a lot for working on it!). I'm mostly curious about what use cases you're targeting, and what use cases we'll need to add support for later on.

206

Typo: broadcastIfNeeded.

mlir/test/Dialect/Linalg/vectorization.mlir

187

What do you think about spelling out the types for vector.transfer_read/transfer_write, since these types can affect semantics? (This is something that I regret not doing in the original code).

[ThomasRaoux]

Add comment and spell out vector types in test.

[ThomasRaoux]

In D92540#2432328, @burmako wrote:

Thanks a lot for your work! This will be really useful :)

I'd also like to share the credit that you kindly gave to me with fedelebron@ and agrue@.

Thanks for the review.

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

100

I added a comment explaining why I added this restriction. I probably need to talk to you to understand how you handle affine_map<(d0, d1) -> (d0)>. I limited it to minor identity so that I can convert the input vector to a larger vector using just a vector.broadcast. With such an affine map I believe I need to do broadcast follow by transpose and it can get more complicated which is why I limited the scope in this case.
Most likely this will need to be relaxed in the future. I'd be happy to learn about the use cases you have that would require that.

mlir/test/Dialect/Linalg/vectorization.mlir

187

Good point, added it.

mlir/include/mlir/EDSC/Builders.h

 namespace edsc {
 /// Helper class to transparently handle builder insertion points by RAII.
 /// As its name indicates, a ScopedContext is means to be used locally in a
 /// scoped fashion. This abstracts away all the boilerplate related to
 /// checking proper usage of captures, NestedBuilders as well as handling the
 /// setting and restoring of insertion points.
 class ScopedContext {
 public:
+  ScopedContext(OpBuilder &b);
   ScopedContext(OpBuilder &b, Location location);
 
   /// Sets the insertion point of the builder to 'newInsertPt' for the duration
   /// of the scope. The existing insertion point of the builder is restored on
   /// destruction.
   ScopedContext(OpBuilder &b, OpBuilder::InsertPoint newInsertPt,
                 Location location);
   ~ScopedContext();

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

static LogicalResult isContraction(Operation *op) {
   auto mapRange = genericOp.indexing_maps().getAsValueRange<AffineMapAttr>();
   return success(
       genericOp.getNumInputs() == 2 && genericOp.getNumOutputs() == 1 &&
       llvm::all_of(mapRange,
                    [](AffineMap m) { return m.isProjectedPermutation(); }) &&
       hasMultiplyAddBody(genericOp.region()));
 }
 
+static bool hasOnlyScalarElementwiseOp(Region &r) {
+  if (!llvm::hasSingleElement(r))
+    return false;
+  for (Operation &op : r.front()) {
+    if (!(isa<ConstantOp, linalg::YieldOp>(op) ||
+          op.hasTrait<OpTrait::ElementwiseMappable>()) ||
+        llvm::any_of(op.getResultTypes(),
+                     [](Type type) { return !type.isIntOrIndexOrFloat(); }))
+      return false;
+  }
+  return true;
+}
+
+// Return true if the op is an element-wise linalg op.
+static bool isElementwise(Operation *op) {
+  auto genericOp = dyn_cast<linalg::GenericOp>(op);
+  if (!genericOp)
+    return false;
+  if (genericOp.getNumLoops() != genericOp.getNumParallelLoops())
+    return false;
+  // TODO: relax the restrictions on indexing map.
+  for (unsigned i = 0, e = genericOp.getNumOutputs(); i < e; i++) {
+    if (!genericOp.getOutputIndexingMap(i).isIdentity())
+      return false;
+  }
+  // Currently limit the input indexing map to minor identity as other
+  // permutations might require adding transpose ops to convert the vector read
+  // to the right shape.
+  for (unsigned i = 0, e = genericOp.getNumInputs(); i < e; i++) {
+    if (!genericOp.getInputIndexingMap(i).isMinorIdentity())
+      return false;
+  }
+  return hasOnlyScalarElementwiseOp(genericOp.getRegion());
+}
+
+static VectorType extractVectorTypeFromScalarView(Value v) {
+  MemRefType mt = v.getType().cast<MemRefType>();
+  return mt.getShape().empty()
+             ? VectorType()
+             : VectorType::get(mt.getShape(), mt.getElementType());
+}
+
+static Value transferReadVector(OpBuilder &builder, Value memref) {
+  edsc::ScopedContext scope(builder);
+  auto memrefType = memref.getType().cast<MemRefType>();
+  if (VectorType vectorType = extractVectorTypeFromScalarView(memref)) {
+    SmallVector<Value, 4> indices(memrefType.getRank(), std_constant_index(0));
+    return vector_transfer_read(vectorType, memref, indices);
+  }
+  return std_load(memref);
+}
+
+static void transferWriteVector(OpBuilder &builder, Value value, Value memref) {
+  edsc::ScopedContext scope(builder);
+  auto memrefType = memref.getType().cast<MemRefType>();
+  if (VectorType vectorType = extractVectorTypeFromScalarView(memref)) {
+    SmallVector<Value, 4> indices(memrefType.getRank(), std_constant_index(0));
+    if (vectorType != value.getType())
+      value = vector_broadcast(vectorType, value);
+    vector_transfer_write(value, memref, indices);
+  } else {
+    std_store(value, memref);
+  }
+}
+
+namespace {
+// Transforms scalar operations into their vectorized counterparts,
+// while using the provided generic op to map:
+//   * Its arguments to transfer reads from the views of the generic op.
+//   * linalg.yield ops to transfer writes to the views of the generic op.
+class GenericVectorizer {
+public:
+  GenericVectorizer(OpBuilder &builder, linalg::GenericOp generic)
+      : builder(builder), generic(generic) {}
+
+  // Takes a scalar operation and builds its vectorized counterpart or
+  // counterparts using the underlying builder.
+  // If operands of the scalar operation are referring to previously vectorized
+  // operations, then in their vectorized form these operands will be referring
+  // to previous vectorization results.
+  void vectorize(Operation &scalarOp) {
+    auto yieldOp = dyn_cast<linalg::YieldOp>(scalarOp);
+    if (yieldOp) {
+      for (auto outputAndMemref :
+           llvm::zip(yieldOp.values(), generic.getOutputBuffers())) {
+        Value vectorValue = vectorize(std::get<0>(outputAndMemref));
+        transferWriteVector(builder, vectorValue, std::get<1>(outputAndMemref));
+      }
+      return;
+    }
+    Operation *vectorOp = uncachedVectorize(scalarOp);
+    assert(scalarOp.getNumResults() == vectorOp->getNumResults());
+    for (auto result :
+         llvm::zip(scalarOp.getResults(), vectorOp->getResults())) {
+      valueCache[std::get<0>(result)] = std::get<1>(result);
+    }
+  }
+
+private:
+  // Transforms a scalar value into its vectorized counterpart, recursively
+  // vectorizing operations as necessary using the underlying builder.
+  // Keeps track of previously vectorized values and reuses vectorization
+  // results if these values come up again.
+  Value vectorize(Value scalarValue) {
+    // Don't vectorize values coming from outside the region.
+    if (scalarValue.getParentRegion() != &generic.region())
+      return scalarValue;
+    auto vectorValueIt = valueCache.find(scalarValue);
+    if (vectorValueIt != valueCache.end())
+      return vectorValueIt->second;
+
+    // If the value is from the region but not in the cache it means it is a
+    // block argument.
+    auto scalarArg = scalarValue.cast<BlockArgument>();
+    assert(scalarArg.getOwner() == &generic.region().front());
+    Value vector_arg =
+        generic.getInputsAndOutputBuffers()[scalarArg.getArgNumber()];
+    Value vectorResult = transferReadVector(builder, vector_arg);
+    valueCache[scalarArg] = vectorResult;
+    return vectorResult;
+  }
+
+  // Return the largest shape of all the given values. Return an empty
+  // SmallVector if there are no vector value.
+  static SmallVector<int64_t, 4> getLargestShape(ArrayRef<Value> values) {
+    SmallVector<int64_t, 4> largestShape;
+    int64_t maxSize = 1;
+    for (Value value : values) {
+      auto vecType = value.getType().dyn_cast<VectorType>();
+      if (!vecType)
+        continue;
+      if (maxSize < vecType.getNumElements()) {
+        largestShape.assign(vecType.getShape().begin(),
+                            vecType.getShape().end());
+      }
+    }
+    return largestShape;
+  }
+
+  // If the value's type doesn't have the given shape broadcast it.
+  Value broadcastIfNeeded(Value value, ArrayRef<int64_t> shape) {
+    auto vecType = value.getType().dyn_cast<VectorType>();
+    if (shape.empty() || (vecType != nullptr && vecType.getShape() == shape))
+      return value;
+    auto newVecType = VectorType::get(shape, vecType ? vecType.getElementType()
+                                                     : value.getType());
+    return builder.create<vector::BroadcastOp>(
+        builder.getInsertionPoint()->getLoc(), newVecType, value);
+  }
+
+  // Takes a scalar operation and builds its vectorized counterpart or
+  // counterparts using underlying builder without involving any caches.
+  Operation *uncachedVectorize(Operation &base_scalarOp) {
+    SmallVector<Value, 4> vectorizedOperands;
+    for (Value operand : base_scalarOp.getOperands()) {
+      vectorizedOperands.push_back(vectorize(operand));
+    }
+    SmallVector<int64_t, 4> shape = getLargestShape(vectorizedOperands);
+    for (Value &operand : vectorizedOperands)
+      operand = broadcastIfNeeded(operand, shape);
+    OperationState state(base_scalarOp.getLoc(), base_scalarOp.getName());
+    state.addAttributes(base_scalarOp.getAttrs());
+    state.addOperands(vectorizedOperands);
+    if (shape.empty()) {
+      state.addTypes(base_scalarOp.getResultTypes());
+    } else {
+      SmallVector<VectorType, 4> vectorizedTypes;
+      for (auto Type : base_scalarOp.getResultTypes())
+        vectorizedTypes.push_back(VectorType::get(shape, Type));
+      state.addTypes(vectorizedTypes);
+    }
+    return builder.createOperation(state);
+  }
+
+  OpBuilder &builder;
+  linalg::GenericOp generic;
+  llvm::DenseMap<Value, Value> valueCache;
+};
+} // namespace
+
+// Replaces elementwise linalg.generic ops with their bodies with scalar
+// operations from these bodies promoted to vector operations.
+static void vectorizeElementwise(linalg::GenericOp op, OpBuilder &builder) {
+  GenericVectorizer vectorizer(builder, op);
+  for (Operation &scalarOp : op.region().front()) {
+    vectorizer.vectorize(scalarOp);
+  }
+}
+
 LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
   auto linalgOp = cast<linalg::LinalgOp>(op);
   // All types must be static shape to go to vector.
   for (Value operand : linalgOp.getInputsAndOutputBuffers())
     if (!operand.getType().cast<ShapedType>().hasStaticShape())
       return failure();
   for (Type outputTensorType : linalgOp.getOutputTensorTypes())
     if (!outputTensorType.cast<ShapedType>().hasStaticShape())
       return failure();
 
   if (isa<linalg::FillOp, linalg::CopyOp>(op))
     return success();
-
+  if (isElementwise(op))
+    return success();
   return isContraction(op);
 }
 
 void mlir::linalg::vectorizeLinalgOp(OpBuilder &builder, Operation *op) {
   assert(succeeded(vectorizeLinalgOpPrecondition(op)));
 
   StringRef dbgPref = "\n[" DEBUG_TYPE "]: ";
   (void)dbgPref;
   edsc::ScopedContext scope(builder, op->getLoc());
   // In the case of 0-D memrefs, return null and special case to scalar load or
   // store later.
-  auto extractVectorTypeFromScalarView = [](Value v) {
-    MemRefType mt = v.getType().cast<MemRefType>();
-    return mt.getShape().empty()
-               ? VectorType()
-               : VectorType::get(mt.getShape(), mt.getElementType());
-  };
   if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
     // Vectorize fill as a vector.broadcast.
     LLVM_DEBUG(dbgs() << dbgPref
                       << "Rewrite linalg.fill as vector.broadcast: " << *op);
-    Value viewOutput = fillOp.output();
-    if (VectorType outputType = extractVectorTypeFromScalarView(viewOutput)) {
-      auto vecType =
-          VectorType::get(fillOp.getOutputBufferType(0).getShape(),
-                          fillOp.getOutputBufferType(0).getElementType());
-      Value vector = vector_broadcast(vecType, fillOp.value());
-      Value zero = std_constant_index(0);
-      SmallVector<Value, 4> indicesOutput(outputType.getRank(), zero);
-      vector_transfer_write(vector, viewOutput, indicesOutput);
-    } else {
-      std_store(fillOp.value(), viewOutput);
-    }
+    transferWriteVector(builder, fillOp.value(), fillOp.output());
     return;
   }
   if (auto copyOp = dyn_cast<linalg::CopyOp>(op)) {
     // Vectorize copy as a vector.transfer_read+vector.transfer_write.
     LLVM_DEBUG(dbgs() << dbgPref
                       << "Rewrite linalg.copy as vector.transfer_read + "
                          "vector.transfer_write: "
                       << *op);
-    Value zero = std_constant_index(0);
-    Value viewInput = copyOp.input();
-    Value viewOutput = copyOp.output();
-    Value vector;
-    if (VectorType inputType = extractVectorTypeFromScalarView(viewInput)) {
-      SmallVector<Value, 4> indicesInput(inputType.getRank(), zero);
-      if (copyOp.inputPermutation())
-        vector = vector_transfer_read(
-            extractVectorTypeFromScalarView(viewInput), viewInput, indicesInput,
-            copyOp.inputPermutation().getValue());
-      else
-        vector =
-            vector_transfer_read(extractVectorTypeFromScalarView(viewInput),
-                                 viewInput, indicesInput);
-    } else {
-      vector = std_load(viewInput).value;
-    }
-    if (VectorType outputType = extractVectorTypeFromScalarView(viewOutput)) {
-      SmallVector<Value, 4> indicesOutput(outputType.getRank(), zero);
-      if (copyOp.outputPermutation())
-        vector_transfer_write(vector, viewOutput, indicesOutput,
-                              copyOp.outputPermutation().getValue());
-      else
-        vector_transfer_write(vector, viewOutput, indicesOutput);
-    } else {
-      std_store(vector, viewOutput);
-    }
+    Value vector = transferReadVector(builder, copyOp.input());
+    transferWriteVector(builder, vector, copyOp.output());
     return;
   }
 
+  if (isElementwise(op)) {
+    LLVM_DEBUG(dbgs() << dbgPref
+                      << "Rewrite linalg op as vector.transfer_read + "
+                         "vector_op + vector.transfer_write: "
+                      << *op);
+    return vectorizeElementwise(cast<linalg::GenericOp>(op), builder);
+  }
+
   assert(succeeded(isContraction(op)) && "Expected contraction");
 
   // Vectorize other ops as vector contraction.
   // TODO: interface.
   LLVM_DEBUG(dbgs() << dbgPref
                     << "Rewrite linalg op as vector.contract: " << *op);
   auto linalgOp = cast<linalg::LinalgOp>(op);
   Value viewA = linalgOp.getInput(0);

mlir/lib/EDSC/Builders.cpp

 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 
 #include "llvm/ADT/Optional.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 
+mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b)
+    : ScopedContext(b, b.getInsertionPoint()->getLoc()) {}
+
 mlir::edsc::ScopedContext::ScopedContext(OpBuilder &b, Location location)
     : builder(b), guard(builder), location(location),
       enclosingScopedContext(ScopedContext::getCurrentScopedContext()) {
   getCurrentScopedContext() = this;
 }
 
 /// Sets the insertion point of the builder to 'newInsertPt' for the duration
 /// of the scope. The existing insertion point of the builder is restored on

mlir/test/Dialect/Linalg/transform-patterns-matmul-to-vector.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-1d | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-2d | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-transform-patterns=test-contraction-to-vector-patterns | FileCheck %s --check-prefix=VECTOR-CONTRACTION
 
 func @matmul(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
   linalg.matmul {__internal_linalg_transform__ = "START"}
     ins(%A, %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                 memref<1584x1584xf32, offset: 0, strides: [1584, 1]>)
    outs(%C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>)

 //      CHECK: linalg.copy
 //      CHECK: linalg.copy
 //
 //      CHECK: vector.contract
 // CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"]
 // CHECK-SAME:   : vector<8x16xf32>, vector<16x12xf32> into vector<8x12xf32>
 //
 //      CHECK: linalg.copy
-
-// VECTOR-CONTRACTION-LABEL: contraction_dot
-func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
-  // VECTOR-CONTRACTION: vector.contract
-  // VECTOR-CONTRACTION-SAME: vector<1584xf32>, vector<1584xf32> into f32
-  linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>)
-            outs(%C: memref<f32>)
-  return
-}
-
-// VECTOR-CONTRACTION-LABEL: contraction_matvec
-func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) {
-  // VECTOR-CONTRACTION: vector.contract
-  // VECTOR-CONTRACTION-SAME: vector<1584x1584xf32>, vector<1584xf32> into vector<1584xf32>
-  linalg.matvec ins(%A, %B: memref<1584x1584xf32>, memref<1584xf32>)
-            outs(%C: memref<1584xf32>)
-  return
-}
-
-// VECTOR-CONTRACTION-LABEL: contraction_matmul
-func @contraction_matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
-  // VECTOR-CONTRACTION: vector.contract
-  // VECTOR-CONTRACTION-SAME: vector<1584x1584xf32>, vector<1584x1584xf32> into vector<1584x1584xf32>
-  linalg.matmul ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>)
-            outs(%C: memref<1584x1584xf32>)
-  return
-}
-
-// VECTOR-CONTRACTION-LABEL: contraction_batch_matmul
-func @contraction_batch_matmul(%A: memref<1584x1584x1584xf32>, %B: memref<1584x1584x1584xf32>, %C: memref<1584x1584x1584xf32>) {
-  // VECTOR-CONTRACTION: vector.contract
-  // VECTOR-CONTRACTION-SAME: vector<1584x1584x1584xf32>, vector<1584x1584x1584xf32> into vector<1584x1584x1584xf32>
-  linalg.batch_matmul
-    ins(%A, %B: memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>)
-   outs(%C: memref<1584x1584x1584xf32>)
-  return
-}

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-patterns | FileCheck %s
 
 // CHECK-DAG: #[[$STRIDED_1D:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1  + s0)>
 // Map corresponding to a 2D memory access where the stride along the last dim is known to be 1.
 // CHECK-DAG: #[[$STRIDED_2D_u_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // Map corresponding to a 2D memory access where the stride along all dims are unknown.
 // CHECK-DAG: #[[$STRIDED_2D:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
-// CHECK-DAG: #[[$mk:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
 // CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
-// CHECK-DAG: #[[$mn:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
 // CHECK-DAG: #[[$nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)>
 // CHECK-DAG: #[[$km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)>
 
 func @dot(%x: memref<?xf32, offset: ?, strides: [1]>,
           %y: memref<?xf32, offset: ?, strides: [1]>,
           %v: memref<f32>) {
   linalg.dot { __internal_linalg_transform__ = "MEM" }
     ins(%x, %y: memref<?xf32, offset: ?, strides: [1]>,

 // CHECK:                         scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
 // CHECK:                           scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2]] {
 // CHECK:                             scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3]] {
 // CHECK:                               scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4]] {
 // CHECK:                                 linalg.matmul
 // CHECK:                                   ins({{.*}}, {{.*}}: memref<?x?xf32, #[[$STRIDED_2D]]>, memref<?x?xf32, #[[$STRIDED_2D]]>)
 // CHECK:                                  outs({{.*}}: memref<?x?xf32, #[[$STRIDED_2D]]>)
 
-#matmul_trait = {
-  args_in = 2,
-  args_out = 1,
-  indexing_maps = [
-    affine_map<(m, n, k) -> (m, k)>,
-    affine_map<(m, n, k) -> (k, n)>,
-    affine_map<(m, n, k) -> (m, n)>
-  ],
-  iterator_types = ["parallel", "parallel", "reduction"],
-  __internal_linalg_transform__ = "VECTORIZE"
-}
-func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
-                         %C: memref<8x32xf32>) {
-  linalg.generic #matmul_trait
-    ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>)
-   outs(%C : memref<8x32xf32>) {
-    ^bb(%a: f32, %b: f32, %c: f32) :
-      %d = mulf %a, %b: f32
-      %e = addf %c, %d: f32
-      linalg.yield %e : f32
-  }
-  return
-}
-// CHECK-LABEL: func @vectorization_test
-//       CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x16xf32>
-//       CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<16x32xf32>
-//       CHECK: vector.transfer_read %{{.*}} : memref<8x32xf32>, vector<8x32xf32>
-//       CHECK: vector.contract {indexing_maps = [#[[$mk]], #[[$kn]], #[[$mn]]], iterator_types = ["parallel", "parallel", "reduction"]} %{{.*}}, %{{.*}}, %{{.*}} : vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
-//       CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<8x32xf32>
-
-func @vectorization_test_integer(%A: memref<8x16xi32>, %B: memref<16x32xi32>,
-                                 %C: memref<8x32xi32>) {
-  linalg.generic #matmul_trait
-    ins(%A, %B : memref<8x16xi32>, memref<16x32xi32>)
-   outs(%C : memref<8x32xi32>) {
-    ^bb(%a: i32, %b: i32, %c: i32) :
-      %d = muli %a, %b: i32
-      %e = addi %c, %d: i32
-      linalg.yield %e : i32
-  }
-  return
-}
-// CHECK-LABEL: func @vectorization_test_integer
-//       CHECK: vector.transfer_read %{{.*}} : memref<8x16xi32>, vector<8x16xi32>
-//       CHECK: vector.transfer_read %{{.*}} : memref<16x32xi32>, vector<16x32xi32>
-//       CHECK: vector.transfer_read %{{.*}} : memref<8x32xi32>, vector<8x32xi32>
-//       CHECK: vector.contract {indexing_maps = [#[[$mk]], #[[$kn]], #[[$mn]]], iterator_types = ["parallel", "parallel", "reduction"]} %{{.*}}, %{{.*}}, %{{.*}} : vector<8x16xi32>, vector<16x32xi32> into vector<8x32xi32>
-//       CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xi32>, memref<8x32xi32>
-
-func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
-                         %C: memref<8x32xf32>) {
-  linalg.matmul { __internal_linalg_transform__ = "VECTORIZE"}
-    ins(%A, %B: memref<8x16xf32>, memref<16x32xf32>)
-   outs(%C: memref<8x32xf32>)
-  return
-}
-// CHECK-LABEL: func @vectorization_test_2
-//       CHECK: vector.contract {{.*}} :
-//                vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
-
-func @test_vectorize_fill(%A : memref<8x16xf32>, %arg0 : f32) {
-  linalg.fill(%A, %arg0) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<8x16xf32>, f32
-  return
-}
-// CHECK-LABEL: func @test_vectorize_fill
-//       CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32>
-//       CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
-
-func @test_vectorize_fill_scalar(%A : memref<f32>, %arg0 : f32) {
-  linalg.fill(%A, %arg0) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<f32>, f32
-  return
-}
-// CHECK-LABEL: func @test_vectorize_fill
-//  CHECK-SAME: (%[[M:.*]]: memref<f32>, %[[V:.*]]: f32)
-//       CHECK:   store %[[V]], %[[M]][] : memref<f32>
-
-func @test_vectorize_copy(%A : memref<8x16xf32>, %B : memref<8x16xf32>) {
-  linalg.copy(%A, %B) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<8x16xf32>, memref<8x16xf32>
-  return
-}
-// CHECK-LABEL: func @test_vectorize_copy
-//       CHECK: %[[V:.*]] = vector.transfer_read {{.*}} : memref<8x16xf32>, vector<8x16xf32>
-//       CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
-
-func @test_vectorize_copy_scalar(%A : memref<f32>, %B : memref<f32>) {
-  linalg.copy(%A, %B) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<f32>, memref<f32>
-  return
-}
-// CHECK-LABEL: func @test_vectorize_copy_scalar
-//       CHECK: %[[V:.*]] = load {{.*}} : memref<f32>
-//       CHECK: store %[[V]], {{.*}} : memref<f32>
-
-
 #matmul_accesses = [
   affine_map<(m, n, k) -> (m, k)>,
   affine_map<(m, n, k) -> (k, n)>,
   affine_map<(m, n, k) -> (m, n)>
 ]
 #generic_matmul_trait = {
   args_in = 2,
   args_out = 1,

mlir/test/Dialect/Linalg/vectorization.mlir

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns | FileCheck %s
+
+// CHECK-DAG: #[[$mk:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK-DAG: #[[$mn:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: contraction_dot
+func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
+  // CHECK: vector.contract
+  // CHECK-SAME: vector<1584xf32>, vector<1584xf32> into f32
+  linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>)
+            outs(%C: memref<f32>)
+  return
+}
+
+// CHECK-LABEL: contraction_matvec
+func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) {
+  // CHECK: vector.contract
+  // CHECK-SAME: vector<1584x1584xf32>, vector<1584xf32> into vector<1584xf32>
+  linalg.matvec ins(%A, %B: memref<1584x1584xf32>, memref<1584xf32>)
+            outs(%C: memref<1584xf32>)
+  return
+}
+
+// CHECK-LABEL: contraction_matmul
+func @contraction_matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
+  // CHECK: vector.contract
+  // CHECK-SAME: vector<1584x1584xf32>, vector<1584x1584xf32> into vector<1584x1584xf32>
+  linalg.matmul ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>)
+            outs(%C: memref<1584x1584xf32>)
+  return
+}
+
+// CHECK-LABEL: contraction_batch_matmul
+func @contraction_batch_matmul(%A: memref<1584x1584x1584xf32>, %B: memref<1584x1584x1584xf32>, %C: memref<1584x1584x1584xf32>) {
+  // CHECK: vector.contract
+  // CHECK-SAME: vector<1584x1584x1584xf32>, vector<1584x1584x1584xf32> into vector<1584x1584x1584xf32>
+  linalg.batch_matmul
+    ins(%A, %B: memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>)
+   outs(%C: memref<1584x1584x1584xf32>)
+  return
+}
+
+#matmul_trait = {
+  args_in = 2,
+  args_out = 1,
+  indexing_maps = [
+    affine_map<(m, n, k) -> (m, k)>,
+    affine_map<(m, n, k) -> (k, n)>,
+    affine_map<(m, n, k) -> (m, n)>
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
+                         %C: memref<8x32xf32>) {
+  linalg.generic #matmul_trait
+    ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>)
+   outs(%C : memref<8x32xf32>) {
+    ^bb(%a: f32, %b: f32, %c: f32) :
+      %d = mulf %a, %b: f32
+      %e = addf %c, %d: f32
+      linalg.yield %e : f32
+  }
+  return
+}
+// CHECK-LABEL: func @vectorization_test
+//       CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x16xf32>
+//       CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<16x32xf32>
+//       CHECK: vector.transfer_read %{{.*}} : memref<8x32xf32>, vector<8x32xf32>
+//       CHECK: vector.contract {indexing_maps = [#[[$mk]], #[[$kn]], #[[$mn]]], iterator_types = ["parallel", "parallel", "reduction"]} %{{.*}}, %{{.*}}, %{{.*}} : vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
+//       CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<8x32xf32>
+
+func @vectorization_test_integer(%A: memref<8x16xi32>, %B: memref<16x32xi32>,
+                                 %C: memref<8x32xi32>) {
+  linalg.generic #matmul_trait
+    ins(%A, %B : memref<8x16xi32>, memref<16x32xi32>)
+   outs(%C : memref<8x32xi32>) {
+    ^bb(%a: i32, %b: i32, %c: i32) :
+      %d = muli %a, %b: i32
+      %e = addi %c, %d: i32
+      linalg.yield %e : i32
+  }
+  return
+}
+// CHECK-LABEL: func @vectorization_test_integer
+//       CHECK: vector.transfer_read %{{.*}} : memref<8x16xi32>, vector<8x16xi32>
+//       CHECK: vector.transfer_read %{{.*}} : memref<16x32xi32>, vector<16x32xi32>
+//       CHECK: vector.transfer_read %{{.*}} : memref<8x32xi32>, vector<8x32xi32>
+//       CHECK: vector.contract {indexing_maps = [#[[$mk]], #[[$kn]], #[[$mn]]], iterator_types = ["parallel", "parallel", "reduction"]} %{{.*}}, %{{.*}}, %{{.*}} : vector<8x16xi32>, vector<16x32xi32> into vector<8x32xi32>
+//       CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xi32>, memref<8x32xi32>
+
+func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
+                         %C: memref<8x32xf32>) {
+  linalg.matmul
+    ins(%A, %B: memref<8x16xf32>, memref<16x32xf32>)
+   outs(%C: memref<8x32xf32>)
+  return
+}
+// CHECK-LABEL: func @vectorization_test_2
+//       CHECK: vector.contract {{.*}} :
+//                vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
+
+func @test_vectorize_fill(%A : memref<8x16xf32>, %arg0 : f32) {
+  linalg.fill(%A, %arg0) :  memref<8x16xf32>, f32
+  return
+}
+// CHECK-LABEL: func @test_vectorize_fill
+//       CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32>
+//       CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
+
+func @test_vectorize_fill_scalar(%A : memref<f32>, %arg0 : f32) {
+  linalg.fill(%A, %arg0) :  memref<f32>, f32
+  return
+}
+// CHECK-LABEL: func @test_vectorize_fill
+//  CHECK-SAME: (%[[M:.*]]: memref<f32>, %[[V:.*]]: f32)
+//       CHECK:   store %[[V]], %[[M]][] : memref<f32>
+
+func @test_vectorize_copy(%A : memref<8x16xf32>, %B : memref<8x16xf32>) {
+  linalg.copy(%A, %B) :  memref<8x16xf32>, memref<8x16xf32>
+  return
+}
+// CHECK-LABEL: func @test_vectorize_copy
+//       CHECK: %[[V:.*]] = vector.transfer_read {{.*}} : memref<8x16xf32>, vector<8x16xf32>
+//       CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
+
+func @test_vectorize_copy_scalar(%A : memref<f32>, %B : memref<f32>) {
+  linalg.copy(%A, %B) :  memref<f32>, memref<f32>
+  return
+}
+// CHECK-LABEL: func @test_vectorize_copy_scalar
+//       CHECK: %[[V:.*]] = load {{.*}} : memref<f32>
+//       CHECK: store %[[V]], {{.*}} : memref<f32>
+
+func @generic_vectorize(%arg0: memref<4x256xf32>, %arg1: memref<4x256xf32>,
+  %arg2: memref<256xf32>, %i: f32) {
+  %c1_f32 = constant 1.0 : f32
+  linalg.generic {
+    args_in = 0 : i64,
+    args_out = 10 : i64,
+    indexing_maps = [
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>,
+      affine_map<(d0, d1) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel"]}
+  ins(%arg1, %arg2: memref<4x256xf32>, memref<256xf32>)
+  outs(
+    %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 :
+    memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>,
+    memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>,
+    memref<4x256xf32>, memref<4x256xf32>) {
+  ^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32,
+    %arg9 : f32, %arg10 : f32, %arg11 : f32, %arg12 : f32, %arg13 : f32,
+    %arg14 : f32):
+    %6 = addf %arg4, %arg6 : f32
+    %7 = cmpf "ogt", %arg3, %arg6 : f32
+    %8 = constant 2.0 : f32
+    %9 = divf %arg5, %i : f32
+    %10 = exp2 %arg5 : f32
+    %11 = mulf %arg5, %8 : f32
+    %12 = rsqrt %arg5 : f32
+    %13 = select %7, %arg5, %arg6 : f32
+    %14 = subf %arg5, %arg6 : f32
+    %15 = tanh %arg5 : f32
+    linalg.yield %6, %8, %c1_f32, %9, %10, %11, %12, %13, %14, %15 : f32, f32,
+      f32, f32, f32, f32, f32, f32, f32, f32
+  }
+  return
+}
+
+// CHECK-LABEL: func @generic_vectorize
+//  CHECK-SAME: (%[[ARG0:.*]]: memref<4x256xf32>, %[[ARG1:.*]]: memref<4x256xf32>,
+//  CHECK-SAME:  %[[ARG2:.*]]: memref<256xf32>, %[[ARG3:.*]]: f32)
+//   CHECK-DAG:   %[[CST0:.*]] = constant dense<2.000000e+00> : vector<4x256xf32>
+//   CHECK-DAG:   %[[CST1:.*]] = constant dense<1.000000e+00> : vector<4x256xf32>
+//   CHECK-DAG:   %[[C0:.*]] = constant 0 : index
+//       CHECK:   %[[V0:.*]] = vector.transfer_read %[[ARG2]][%[[C0]]], {{.*}} : memref<256xf32>, vector<256xf32>
+//       CHECK:   %[[V1:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32>
+//       CHECK:   %[[V0B:.*]] = vector.broadcast %[[V0]] : vector<256xf32> to vector<4x256xf32>
+//       CHECK:   %[[ADD:.*]] = addf %[[V0B]], %[[V1]] : vector<4x256xf32>
+//       CHECK:   %[[V2:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32>
+//       CHECK:   %[[CMP:.*]] = cmpf "ogt", %[[V2]], %[[V1]] : vector<4x256xf32>
+//       CHECK:   %[[V3:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> 
+//       CHECK:   %[[ARG3B:.*]] = vector.broadcast %[[ARG3]] : f32 to vector<4x256xf32>
+//       CHECK:   %[[DIV:.*]] = divf %[[V3]], %[[ARG3B]] : vector<4x256xf32>
+//       CHECK:   %[[EXP:.*]] = exp2 %[[V3]] : vector<4x256xf32>
+//       CHECK:   %[[MUL:.*]] = mulf %[[V3]], %[[CST0]] : vector<4x256xf32>
+//       CHECK:   %[[RSQRT:.*]] = rsqrt %[[V3]] : vector<4x256xf32>
+//       CHECK:   %[[SEL:.*]] = select %[[CMP]], %[[V3]], %[[V1]] : vector<4x256xi1>, vector<4x256xf32>
+//       CHECK:   %[[SUB:.*]] = subf %[[V3]], %[[V1]] : vector<4x256xf32>
+//       CHECK:   %[[TAN:.*]] = tanh %[[V3]] : vector<4x256xf32>
+//       CHECK:   vector.transfer_write %[[ADD]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[CST0]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[CST1]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[DIV]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[EXP]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[MUL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[RSQRT]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[SEL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[SUB]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>
+//       CHECK:   vector.transfer_write %[[TAN]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32>

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

Option<bool> testTileAndDistributionOptions{
       llvm::cl::desc("Test tile and distribute 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)};
   Option<bool> testGenericToVectorPattern{
-      *this, "test-contraction-to-vector-patterns",
+      *this, "test-linalg-to-vector-patterns",
       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
                      "in vector.contract form"),
       llvm::cl::init(false)};
   Option<bool> testAffineMinSCFCanonicalizationPatterns{
       *this, "test-affine-min-scf-canonicalization-patterns",
       llvm::cl::desc("Test affine-min + scf canonicalization patterns."),
       llvm::cl::init(false)};
 };

 
 static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
   OwningRewritePatternList forwardPattern;
   forwardPattern.insert<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
   forwardPattern.insert<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
   applyPatternsAndFoldGreedily(funcOp, std::move(forwardPattern));
 }
 
-static void applyContractionToVectorPatterns(FuncOp funcOp) {
+static void applyLinalgToVectorPatterns(FuncOp funcOp) {
   OwningRewritePatternList patterns;
-  patterns.insert<LinalgVectorizationPattern<BatchMatmulOp>,
-                  LinalgVectorizationPattern<MatmulOp>,
-                  LinalgVectorizationPattern<MatvecOp>,
-                  LinalgVectorizationPattern<VecmatOp>,
-                  LinalgVectorizationPattern<DotOp>,
-                  LinalgVectorizationPattern<GenericOp>>(funcOp.getContext());
+  patterns.insert<
+      LinalgVectorizationPattern<BatchMatmulOp>,
+      LinalgVectorizationPattern<VecmatOp>, LinalgVectorizationPattern<DotOp>,
+      LinalgVectorizationPattern<FillOp>, LinalgVectorizationPattern<CopyOp>,
   applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
 static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
   OwningRewritePatternList foldPattern;
   foldPattern.insert<AffineMinSCFCanonicalizationPattern>(funcOp.getContext());
   FrozenRewritePatternList frozenPatterns(std::move(foldPattern));
 
if (testPatterns)
     return applyPatterns(getFunction());
   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
     return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
   if (testGenericToVectorPattern)
-    return applyContractionToVectorPatterns(getFunction());
+    return applyLinalgToVectorPatterns(getFunction());
   if (testAffineMinSCFCanonicalizationPatterns)
     return applyAffineMinSCFCanonicalizationPatterns(getFunction());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace test
 } // namespace mlir