[Mlir-commits] [mlir] [mlir][ArmSME] Add support for vector.transfer_read with transpose (PR #67527)

[Diego Caballero]

================
@@ -59,6 +59,67 @@ getSMETileAndCastToVector(PatternRewriter &rewriter, Location loc,
 
 namespace {
 
+/// Conversion pattern for vector.transfer_read op with transpose permutation
+/// map to vertical arm_sme.tile_load (in-flight transpose).
+///
+///   vector.transfer_read ...  permutation_map: (d0, d1) -> (d1, d0)
+///
+/// is converted to:
+///
+///   arm_sme.tile_load ... <vertical>
+struct TransferReadPermutationToArmSMELowering
+    : public OpRewritePattern<vector::TransferReadOp> {
+  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
+                                PatternRewriter &rewriter) const final {
+    // The permutation map must have two results.
+    if (transferReadOp.getTransferRank() != 2)
+      return rewriter.notifyMatchFailure(transferReadOp,
+                                         "not a 2 result permutation map");
+
+    auto vectorType = transferReadOp.getVectorType();
+    if (!arm_sme::isValidSMETileVectorType(vectorType))
+      return rewriter.notifyMatchFailure(transferReadOp,
+                                         "not a valid vector type for SME");
+
+    if (!llvm::isa<MemRefType>(transferReadOp.getSource().getType()))
+      return rewriter.notifyMatchFailure(transferReadOp, "not a memref source");
+
+    if (transferReadOp.getMask())
+      // TODO: support masking.
+      return rewriter.notifyMatchFailure(transferReadOp,
+                                         "masking not yet supported");
+
+    // Out-of-bounds dims are not supported.
+    if (transferReadOp.hasOutOfBoundsDim())
+      return rewriter.notifyMatchFailure(transferReadOp,
+                                         "not inbounds transfer read");
+
+    AffineMap map = transferReadOp.getPermutationMap();
+
+    // Permutation map doesn't perform permutation, can be lowered to
+    // vector.load by TransferReadToVectorLoadLowering and then
+    // arm_sme.tile_load by VectorLoadToArmSMELowering.
+    if (map.isIdentity())
----------------
dcaballe wrote:

I would move this check earlier in the function as it would be a very common early exit

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