[llvm-branch-commits] [mlir] [mlir][linalg] Enable scalable vectorization of linalg.unpack (PR #149293)

[Han-Chung Wang via llvm-branch-commits]

================
@@ -1876,22 +1872,35 @@ static VectorType getCollapsedVecType(VectorType type,
   return VectorType::get(newShape, type.getElementType(), newScalableFlags);
 }
 
-/// Vectorize a `linalg::UnPackOp` to these 4 Ops:
-///   Vector::TransferReadOp - Reads a vector from the source tensor
-///   vector::TransposeOp - Transpose the Source tensor
-///   ShapeCastOp - Reshape the data based on the target.
-///   vector::TransferWriteOp. - Write the result vector back to the destination
-///   tensor.
-///   If the vector sizes are not provided:
-/// Vectorize `linalg.unpack %src into %dest` as:
-///   // Reads a vector from the source tensor
-///   %read = vector.transfer_read  %src
-///   // Transpose %read as specified in `outer_dims_perm` attribute
-///   %tr = vector.transpose %read
-///   // Reshape the data based on the target
-///   %sc = vector.shape_cast %tr
-///   // Write the result vector to the destination tensor.
-///   vector.transfer_write %sc into %dest
+/// Vectorize `linalg.unpack` into:
+///   * xfer_read -> vector.transpose -> vector.shape_cast -> xfer_write
+///
+/// The input-vector-sizes specify both the read and the write vector
+/// sizes and are passed as one array covering both operations, i.e.:
+///
+///  input-vector-sizes = [1, 1, 8, [8],  8, [8]]
+///                        \         /    \    /
+///                        read-sizes   write-sizes
+///
+/// (for brefity, in the diagram,
+///    * input-vector-sizes = `inputVectorSizes` + `inputScalableDims`
+/// )
+///
+/// If the vector sizes are not provided:
+///  * the vector sizes are determined by the operands,
+///  * the inBounds attribute is used instead of masking.
+///
+/// EXAMPLE (no vector sizes):
+/// ```
+///   %unpack = linalg.unpack  %src
+///    inner_dims_pos = [0, 1]
+///    inner_tiles = [8, 8]
+///    into %dest : tensor<1x1x8x8xf32> -> tensor<8x8xf32>
+/// ```
+/// is vectorized as:
+/// ```
+///   vector.transfer_write %sc into %dest : vector<8x8xf32>, tensor<8x8xf32>
+/// ```
----------------
hanhanW wrote:

You deleted few ops in https://github.com/llvm/llvm-project/pull/149293/commits/b073854c4e681f2a8c89528ae4108f8bf7c6193f. I think we should add the xfer_read, transpose and shape_cast back.

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