[llvm-branch-commits] [mlir] [MLIR] Legalize certain `vector.transfer_read` ops of scalable vectors (PR #143146)

[Momchil Velikov via llvm-branch-commits]

================
@@ -298,16 +298,139 @@ struct LegalizeSVEMaskLoadConversion : public OpRewritePattern<memref::LoadOp> {
   }
 };
 
+/// Transforms a `transfer_read` operation so it reads vector of a type that
+/// can be mapped to an LLVM type. This is done by collapsing trailing
+/// dimensions so we obtain a vector type with a single scalable dimension in
+/// the rightmost position.
+///
+/// Example:
+/// ```
+/// %v = vector.transfer_read %M[%i, %j, %c0, %c0], %c0_i8
+///   {in_bounds = [false, true, true, true]}
+///   : memref<?x?x2x8xi8>, vector<2x[4]x2x8xi8>
+/// ```
+/// is rewritten to
+/// ```
+/// %collapse_shape = memref.collapse_shape %M [[0], [1, 2, 3]]
+///   : memref<?x?x2x8xi8> into memref<?x?xi8>
+/// %0 = vector.transfer_read  %collapse_shape[%i, %j], %c0_i8
+///   {in_bounds = [false, true]}
+///   : memref<?x?xi8>, vector<2x[64]xi8>
+/// %1 = vector.shape_cast %0 : vector<2x[64]xi8> to vector<2x[4]x2x8xi8>
+/// ```
+struct LegalizeTransferRead : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
+                                PatternRewriter &rewriter) const override {
+
+    // Do not try to transform masked reads. For example, if we have a transfer
+    // to a `vector<[4]x4xi8>` we could have a mask like
+    //    1 1 1 0
+    //    1 1 1 0
+    //    1 1 1 0
+    //    0 0 0 0
+    // Flattening this mask would look like
+    //    1 1 1 0 1 1 1 0 1 1 1 0 0 0 0 0
+    // and we have not yet figured out an efficient way to build such a mask,
+    // neither from the mask operand, nor from the original `vector.create_mask`
+    // operation (if visible at all).
+    if (readOp.isMasked() || readOp.getMask())
+      return rewriter.notifyMatchFailure(readOp,
+                                         "masked transfers not-supported");
+
+    if (!readOp.getPermutationMap().isMinorIdentity())
+      return rewriter.notifyMatchFailure(readOp, "non-identity permutation");
+
+    // We handle transfers of vectors with rank >= 2 and a single scalable
+    // dimension.
----------------
momchil-velikov wrote:

Comment added.

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