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

[Benjamin Maxwell]

================
@@ -3186,20 +3259,21 @@ struct Conv1DGenerator
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       for (int64_t w = 0; w < wSize; w += wSizeStep) {
         lhsVals.push_back(rewriter.create<vector::ExtractStridedSliceOp>(
-            loc, lhs,
+            loc, maybeMaskedLhs->getResult(0),
             /*offsets=*/ArrayRef<int64_t>{0, w * strideW + kw * dilationW, 0},
             inOutSliceSizes, inOutStrides));
       }
     }
     // Extract rhs slice of size {c} @ [kw].
     for (int64_t kw = 0; kw < kwSize; ++kw) {
       rhsVals.push_back(rewriter.create<vector::ExtractOp>(
-          loc, rhs, /*offsets=*/ArrayRef<int64_t>{kw}));
+          loc, maybeMaskedRhs->getResult(0),
+          /*offsets=*/ArrayRef<int64_t>{kw}));
     }
     // Extract res slice: {n, wSizeStep, c} @ [0, w, 0].
     for (int64_t w = 0; w < wSize; w += wSizeStep) {
       resVals.push_back(rewriter.create<vector::ExtractStridedSliceOp>(
-          loc, res,
+          loc, maybeMaskedRes->getResult(0),
----------------
MacDue wrote:

The logic below this (for when `flatten=true`) looks like it will drop the scalable dims: 
```
    auto inOutFlattenSliceSizes =
        SmallVector<int64_t>{nSize, wSizeStep * cSize};
    auto lhsCastType = VectorType::get(inOutFlattenSliceSizes, lhsEltType);
    auto resCastType = VectorType::get(inOutFlattenSliceSizes, resEltType);
    // Compute contraction: O{n, w, c} += I{n, sw * w + dw * kw, c} * F{c}
```

You may want to update that (or bail out if early if `flatten=true`).

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