[Mlir-commits] [mlir] [mlir][linalg] Enable masked vectorisation for depthwise convolutions (PR #81625)

[Andrzej Warzyński]

================
@@ -3371,7 +3465,17 @@ vectorizeConvolution(RewriterBase &rewriter, LinalgOp op,
   res = e.generateNcwPooling();
   if (succeeded(res))
     return res;
-  return e.generateDilatedConv(flatten1DDepthwiseConv);
+
+  uint64_t vecChDimSize = ShapedType::kDynamic;
+  bool vecChDimScalableFlag = false;
+  if (!inputVecSizes.empty()) {
+    // Only use the input vector size corresponding to the channel dim. Other
+    // vector dims will be inferred from the Ops.
+    vecChDimSize = inputVecSizes[2];
+    vecChDimScalableFlag = inputScalableVecDims[2];
----------------
banach-space wrote:

Yes. Note that I am inserting this just before calling `generateDilatedConv` and: https://github.com/llvm/llvm-project/blob/982e9022ca4abaad58c693d2b0aba0e908ee2d39/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp#L3400-L3416

Specifically:
```cpp:
      return depthwiseConv(flatten);
```
and
https://github.com/llvm/llvm-project/blob/982e9022ca4abaad58c693d2b0aba0e908ee2d39/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp#L3125-L3134

I can add a comment to clarify the assumptions made.

https://github.com/llvm/llvm-project/pull/81625