[mlir][Linalg] Tile sizes for Conv ops vectorization added as pass arguments

[limo1996]

Summary

Current setup for conv op vectorization does not enable user to specify tile
sizes as well as dimensions for vectorization. In this commit we change that by
adding tile sizes as pass arguments. Every dimension with corresponding tile
size > 1 is automatically vectorized.

[ftynse]

In general, I wouldn't couple vectorization flags to tiling flags (well, maybe have tiling as a precondition) because we may want to tile several times. But this is already an improvement over the current state.

[limo1996]

In D88533#2302544, @ftynse wrote:

In general, I wouldn't couple vectorization flags to tiling flags (well, maybe have tiling as a precondition) because we may want to tile several times. But this is already an improvement over the current state.

I agree. We can tile for free before the pass though.. You can see it in the tests. But yeah, right now we cannot vectorize without tiling (tiling is in fact needed only when dimension is dynamic)..

[nicolasvasilache]

Same comment as Alex.
Still this is an improvement so let's land it.

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

 
 struct TiledLinalgOp {
   LinalgOp op;
   SmallVector<Operation *, 8> loops;
 };
 
 /// Populates patterns for vectorization of all ConvN-D ops.
 void populateConvVectorizationPatterns(
-    MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns);
+    MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns,
+    ArrayRef<int64_t> tileSizes);
 
 /// Performs standalone tiling of a single LinalgOp by `tileSizes`.
 /// and permute the loop nest according to `interchangeVector`
 /// The permutation is expressed as a list of integers that specify
 /// the new ordering of the loop nest. The length of `interchangeVector`
 /// must be equal to the length of `tileSizes`.
 /// An empty vector is interpreted as the identity permutation and the
 /// transformation returns early.

 };
 
 /// Converts Convolution op into vector contraction.
 ///
 /// Conversion expects ConvOp to have dimensions marked in the *mask* as
 /// false of size 1. This ensures that the ConvOp can be lowered to vector
 /// contraction of dimensions marked in the *mask* as true.
 ///
-/// A good example is ConvNHWCOp which is 2D Conv op with channels as the last
-/// dimension. For this op we contract last 3 dimensions.
+/// A good example for vectorization is ConvNHWCOp which is 2D Conv op
+/// with channels as the last dimension. Let's vectorize last 3 dimensions.
 /// The initial op definition looks like this:
 /// ```
 /// linalg.conv_2d_nhwc  %arg0, %arg1, %arg2 :
 ///   (memref<1x3x3x3xf32>, memref<1x3x3x3xf32>, memref<?x?x?x?xf32>)
 /// ```
 /// This op can be expressed as a dot product between %arg0 (input) and
 /// %arg1 (kernel) which is written into first entry of %arg2 (output). This is
 /// the ConvOp this pass expects and converts into:
public:
   ConvOpVectorization(MLIRContext *context, SmallVector<bool, 4> msk)
       : OpRewritePattern<ConvOp>(context) {
     assert(msk.size() == N && "Mask size does not match rank");
     this->mask = msk;
   }
 
   LogicalResult matchAndRewrite(ConvOp minOp,
                                 PatternRewriter &rewriter) const override;
-
-  // TODO: Make these pass arguments.
-  static const int tileSize = 3;
-  static const int noTile = 1;
 };
 
 //===----------------------------------------------------------------------===//
 // 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

mlir/integration_test/Dialect/Linalg/CPU/test-conv-1d-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=4" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=4" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns a 1-D buffer of size %s1 filled with the value %f
 func @alloc_1d_filled_f32(%s1 : index, %f : f32) -> memref<?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-1d-ncw-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,0,4" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,0,4" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f
 func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-1d-nwc-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,4" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,4" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f
 func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-2d-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,2" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,3,3"  -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,2" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,3,3"  -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns a 2-D buffer of size (%s1, %s2) filled with the value %f
 func @alloc_2d_filled_f32(%s1 : index, %s2 : index, %f : f32) -> memref<?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-2d-nchw-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,0,4,4" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,1,3,3,3"  -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,0,4,4" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,1,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 4-D buffer of size (%s1, %s2, %s3, %s4) filled with the value %f
 func @alloc_4d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %f : f32) -> memref<?x?x?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-2d-nhwc-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,3,2" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,1,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,3,2" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,1,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 4-D buffer of size (%s1, %s2, %s3, %s4) filled with the value %f
 func @alloc_4d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %f : f32) -> memref<?x?x?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-3d-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,2,2" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,2,2" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f
 func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-3d-ncdhw-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,0,5,5,5" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,1,1,3,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,0,5,5,5" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,1,1,3,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 5-D buffer of size (%s1, %s2, %s3, %s4, %s5) filled with the value %f
 func @alloc_5d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %s5 : index, %f : f32) -> memref<?x?x?x?x?xf32> {

mlir/integration_test/Dialect/Linalg/CPU/test-conv-3d-ndhwc-call.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,5,5,5" -convert-linalg-to-loops \
 // RUN:   -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
-// RUN: mlir-opt %s -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,1,1,1,1,3,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,5,5,5" \
-// RUN:   -test-conv-vectorization -convert-linalg-to-llvm | \
+// RUN:   -test-conv-vectorization="tile-sizes=1,1,1,1,1,3,3,3,3" -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @print_memref_f32(memref<*xf32>)
 
 // Creates and returns 5-D buffer of size (%s1, %s2, %s3, %s4, %s5) filled with the value %f
 func @alloc_5d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %s5 : index, %f : f32) -> memref<?x?x?x?x?xf32> {

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

LogicalResult ConvOpVectorization<ConvOp, N>::matchAndRewrite(
 
   ArrayRef<int64_t> inShape = inShapeType.getShape();
   ArrayRef<int64_t> kShape = kShapeType.getShape();
 
   if (!inShapeType.hasStaticShape() || !kShapeType.hasStaticShape())
     return failure();
 
   SmallVector<AffineExpr, 4> mapping;
-  // Fail to apply when the size of not vectorized dimension is not 1 or
-  // when the size of vectorized dimension is not dimSize.
+  SmallVector<int64_t, 4> vectorDims;
+  // Fail to apply when the size of not vectorized dimension is not 1.
   for (unsigned i = 0; i < N; i++) {
     if (!mask[i] && (inShape[i] != 1 || kShape[i] != 1))
       return failure();
-    if (mask[i] && (inShape[i] != tileSize || kShape[i] != tileSize))
+
+    if (mask[i] && inShape[i] != kShape[i])
       return failure();
 
-    if (mask[i])
+    if (mask[i]) {
       mapping.push_back(getAffineDimExpr(i, context));
+      vectorDims.push_back(inShape[i]);
+    }
   }
 
   Value input = op.getInput(0);
   Value kernel = op.getInput(1);
   Value output = op.getOutputBuffer(0);
 
   unsigned rank = inShapeType.getRank();
   unsigned numDims = mapping.size();
   Type elemType = inShapeType.getElementType();
 
   auto map = AffineMap::get(rank, 0, mapping, context);
   SmallVector<Value, 4> zeros(rank, std_constant_index(0));
-  auto vecType =
-      VectorType::get(SmallVector<int64_t, 4>(numDims, tileSize), elemType);
+  auto vecType = VectorType::get(vectorDims, elemType);
 
   auto inputVec = vector_transfer_read(vecType, input, zeros, map);
   auto kernelVec = vector_transfer_read(vecType, kernel, zeros, map);
 
   auto acc = std_constant(elemType, rewriter.getZeroAttr(elemType));
 
   std::array<AffineMap, 3> indexingMaps{
       AffineMap::getMultiDimIdentityMap(numDims, context),

 /// conversion into corresponding pattern lists.
 template <typename ConvOp, unsigned N>
 static void
 populateVectorizationPatterns(OwningRewritePatternList &tilingPatterns,
                               OwningRewritePatternList &promotionPatterns,
                               OwningRewritePatternList &vectorizationPatterns,
                               ArrayRef<int64_t> tileSizes,
                               MLIRContext *context) {
+  if (tileSizes.size() < N)
+    return;
+
   constexpr static StringRef kTiledMarker = "TILED";
   constexpr static StringRef kPromotedMarker = "PROMOTED";
   tilingPatterns.insert<LinalgTilingPattern<ConvOp>>(
       context, LinalgTilingOptions().setTileSizes(tileSizes),
       LinalgMarker({}, Identifier::get(kTiledMarker, context)));
 
   promotionPatterns.insert<LinalgPromotionPattern<ConvOp>>(
       context, LinalgPromotionOptions().setUseFullTileBuffersByDefault(true),
       LinalgMarker(Identifier::get(kTiledMarker, context),
                    Identifier::get(kPromotedMarker, context)));
 
   SmallVector<bool, 4> mask(N);
   int offset = tileSizes.size() - N;
   std::transform(tileSizes.begin() + offset, tileSizes.end(), mask.begin(),
-                 [](int64_t i) -> bool { return i != ConvOpConst::noTile; });
+                 [](int64_t i) -> bool { return i > 1; });
 
   vectorizationPatterns.insert<ConvOpVectorization<ConvOp, N>>(context, mask);
 }
 
 void mlir::linalg::populateConvVectorizationPatterns(
-    MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns) {
-  const int64_t tileSize = ConvOpConst::tileSize;
-  const int64_t noTile = ConvOpConst::noTile;
-  auto makeTileSizes = [&](unsigned numNoTile, unsigned numTile) {
-    SmallVector<int64_t, 10> result(numNoTile, noTile);
-    result.append(numTile, tileSize);
-    return result;
-  };
-
+    MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns,
+    ArrayRef<int64_t> tileSizes) {
   OwningRewritePatternList tiling, promotion, vectorization;
-  populateVectorizationPatterns<ConvWOp, 1>(
-      tiling, promotion, vectorization,
-      makeTileSizes(/*numNoTile=*/1, /*numTile*/ 1), context);
+  populateVectorizationPatterns<ConvWOp, 1>(tiling, promotion, vectorization,
+                                            tileSizes, context);
 
   populateVectorizationPatterns<ConvNWCOp, 3>(tiling, promotion, vectorization,
-                                              makeTileSizes(3, 2), context);
+                                              tileSizes, context);
 
   populateVectorizationPatterns<ConvNCWOp, 3>(tiling, promotion, vectorization,
-                                              makeTileSizes(3, 2), context);
+                                              tileSizes, context);
 
   populateVectorizationPatterns<ConvHWOp, 2>(tiling, promotion, vectorization,
-                                             makeTileSizes(2, 2), context);
+                                             tileSizes, context);
 
   populateVectorizationPatterns<ConvNHWCOp, 4>(tiling, promotion, vectorization,
-                                               makeTileSizes(4, 3), context);
+                                               tileSizes, context);
 
   populateVectorizationPatterns<ConvNCHWOp, 4>(tiling, promotion, vectorization,
-                                               makeTileSizes(4, 3), context);
+                                               tileSizes, context);
 
   populateVectorizationPatterns<ConvDHWOp, 3>(tiling, promotion, vectorization,
-                                              makeTileSizes(3, 3), context);
+                                              tileSizes, context);
 
   populateVectorizationPatterns<ConvNDHWCOp, 5>(
-      tiling, promotion, vectorization, makeTileSizes(5, 4), context);
+      tiling, promotion, vectorization, tileSizes, context);
 
   populateVectorizationPatterns<ConvNCDHWOp, 5>(
-      tiling, promotion, vectorization, makeTileSizes(5, 4), context);
+      tiling, promotion, vectorization, tileSizes, context);
 
   patterns.push_back(std::move(tiling));
   patterns.push_back(std::move(promotion));
   patterns.push_back(std::move(vectorization));
 }

mlir/test/Conversion/LinalgToVector/linalg-to-vector.mlir

-// RUN: mlir-opt %s -test-conv-vectorization --cse | FileCheck %s
+// RUN: mlir-opt %s -test-conv-vectorization="tile-sizes=1,3" --cse | FileCheck %s
 
 // CHECK-DAG:  #[[$map0:.*]] = affine_map<(d0)[s0] -> (1, -d0 + s0)>
 // CHECK-DAG:  #[[$map1:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECK-DAG:  #[[$map2:.*]] = affine_map<(d0, d1) -> (d0 + d1)>
 // CHECK-DAG:  #[[$map3:.*]] = affine_map<(d0, d1)[s0] -> (3, -d0 - d1 + s0)>
 // CHECK-DAG:  #[[$map4:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)>
 // CHECK-DAG:  #[[$map5:.*]] = affine_map<(d0) -> (d0)>
 

mlir/test/lib/Transforms/TestConvVectorization.cpp

 
 using namespace mlir;
 using namespace vector;
 
 namespace {
 /// A pass converting MLIR Linalg ops into Vector ops.
 class TestConvVectorization
     : public PassWrapper<TestConvVectorization, OperationPass<ModuleOp>> {
+public:
+  TestConvVectorization() = default;
+  TestConvVectorization(const TestConvVectorization &) {}
+  explicit TestConvVectorization(ArrayRef<int64_t> tileSizesParam) {
+    tileSizes = tileSizesParam;
+  }
+
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<VectorDialect>();
     registry.insert<linalg::LinalgDialect>();
     registry.insert<scf::SCFDialect>();
     registry.insert<AffineDialect>();
     registry.insert<StandardOpsDialect>();
   }
+
+  ListOption<int64_t> tileSizes{
+      *this, "tile-sizes", llvm::cl::desc("Vectorization sizes."),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
 };
 } // namespace
 
 void TestConvVectorization::runOnOperation() {
   MLIRContext *context = &getContext();
   ModuleOp module = getOperation();
 
   ConversionTarget target(*context);
   target.addLegalDialect<AffineDialect, scf::SCFDialect, StandardOpsDialect,
                          VectorDialect>();
   target.addLegalOp<ModuleOp, FuncOp, ModuleTerminatorOp, ReturnOp>();
   target.addLegalOp<linalg::FillOp, linalg::YieldOp>();
 
   SmallVector<OwningRewritePatternList, 4> stage1Patterns;
-  linalg::populateConvVectorizationPatterns(context, stage1Patterns);
+  linalg::populateConvVectorizationPatterns(context, stage1Patterns, tileSizes);
 
   OwningRewritePatternList stage2Patterns =
       linalg::getLinalgTilingCanonicalizationPatterns(context);
   stage2Patterns.insert<linalg::AffineMinSCFCanonicalizationPattern>(context);
 
   auto stage3Transforms = [](Operation *op) {
     PassManager pm(op->getContext());
     pm.addPass(createLoopInvariantCodeMotionPass());