[Mlir-commits] [mlir] [mlir][vector] Refactor `createWriteOrMaskedWrite` (PR #138137)

[Diego Caballero]

================
@@ -1506,84 +1506,86 @@ static SmallVector<int64_t> getTiledPackShape(linalg::PackOp packOp,
   return applyPermutation(destShape, linalg::getPackInverseDestPerm(packOp));
 }
 
-/// Creates a TransferWriteOp to write `input` into a newly initialized
-/// output tensor.
+/// Creates an optionally masked TransferWriteOp
 ///
-/// Given:
-/// - an input vector to write,
-/// - the mixed destination sizes for the output tensor,
-/// - and the vector sizes used for vectorization (i.e., the leading N dims,
-///   for some value of N),
-///
-/// this function generates the following sequence of ops:
-///
-///   %dest = tensor.empty(%destSizes)
-///   %res = vector.transfer_write %input into %dest
+/// Generates the following operation:
+///   %res = vector.transfer_write %vectorToStore into %dest
 ///
 /// If the leading N dimensions of the destination tensor do not match
-/// `inputVecSizesForLeadingDims` (where N =
-/// rank(`inputVecSizesForLeadingDims`)), masking is applied to ensure
-/// correctness:
+/// `inputVecSizesForLeadingDims`, where=
+///   * N = rank(`inputVecSizesForLeadingDims`)),
+/// masking is applied to ensure correctness:
 ///
-///   %dest = tensor.empty(%destSizes)
-///   %write = vector.transfer_write %input into %dest
-///   %mask = vector.create_mask(%destSizes)
+///   %write = vector.transfer_write %vectorToStore into %dest
+///   %mask = vector.create_mask(%destShape)
 ///   %res = vector.mask %mask { %write }
----------------
dcaballe wrote:

Wait, this is not a correct use of `vector.mask`. The masked operation should be nested within the `vector.mask` operation. I saw a similar example in the CSE discussion so let me take a look at the verifier...

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