[Mlir-commits] [mlir] [mlir][linalg] Hoist masked vector transfer pairs (PR #221481)

[Andrzej Warzyński]

================
@@ -233,13 +238,25 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
                         << "\n");
 
       SetVector<Operation *> forwardSlice;
-      getForwardSlice(transferRead.getOperation(), &forwardSlice);
+      getForwardSlice(readUnit, &forwardSlice);
+      // A forward slice does not follow a value out of a `vector.mask` region,
+      // so re-slice from each enclosing mask op to reach the write.
+      for (unsigned i = 0; i < forwardSlice.size(); ++i) {
+        auto mask = forwardSlice[i]->getParentOfType<vector::MaskOp>();
+        if (mask && forwardSlice.insert(mask.getOperation()))
+          getForwardSlice(mask.getOperation(), &forwardSlice);
+      }
 
       // Look for the last TransferWriteOp in the forwardSlice of
-      // `transferRead` that operates on the same memref.
+      // `transferRead` that operates on the same memref. A masked write appears
+      // as the `vector.mask` op wrapping it.
       vector::TransferWriteOp transferWrite;
       for (auto *sliceOp : llvm::reverse(forwardSlice)) {
         auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
----------------
banach-space wrote:

[nit] 
```suggestion
        auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
        // If this is not an xfer_write Op, check whether it's masked xfer_write
```

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