[mlir][ArmSME] Simplify permutation map handling

[c-rhodes]

In -convert-vector-to-arm-sme the permutation_map is explicitly checked for transpose when converting xfer ops, but for 2-D vector types the only non-identity permutation map is transpose so this can be simplified.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Cullen Rhodes (c-rhodes)

Changes

In -convert-vector-to-arm-sme the permutation_map is explicitly checked for transpose when converting xfer ops, but for 2-D vector types the only non-identity permutation map is transpose so this can be simplified.

---

Full diff: https://github.com/llvm/llvm-project/pull/93515.diff

1 Files Affected:

• (modified) mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp (+14-18)

diff --git a/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp b/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
index 87923477766d1..c2f1584e43bac 100644
--- a/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
+++ b/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
@@ -65,20 +65,18 @@ struct TransferReadToArmSMELowering
       return rewriter.notifyMatchFailure(transferReadOp,
                                          "not inbounds transfer read");
 
-    arm_sme::TileSliceLayout layout;
-
-    AffineExpr d0, d1;
-    bindDims(transferReadOp.getContext(), d0, d1);
     AffineMap map = transferReadOp.getPermutationMap();
-    if (map.isIdentity())
-      layout = arm_sme::TileSliceLayout::Horizontal;
-    else if (map == AffineMap::get(map.getNumDims(), 0, {d1, d0},
-                                   transferReadOp.getContext()))
-      layout = arm_sme::TileSliceLayout::Vertical;
-    else
+    if (!map.isPermutation())
       return rewriter.notifyMatchFailure(transferReadOp,
                                          "unsupported permutation map");
 
+    // Note: For 2D vector types the only non-identity permutation is a simple
+    // transpose [1, 0].
+    bool transposed = !map.isIdentity();
+    arm_sme::TileSliceLayout layout =
+        transposed ? arm_sme::TileSliceLayout::Vertical
+                   : arm_sme::TileSliceLayout::Horizontal;
+
     // Padding isn't optional for transfer_read, but is only used in the case
     // of out-of-bounds accesses (not supported here) and/or masking. Mask is
     // optional, if it's not present don't pass padding.
@@ -137,19 +135,17 @@ struct TransferWriteToArmSMELowering
       return rewriter.notifyMatchFailure(writeOp,
                                          "not inbounds transfer write");
 
-    AffineExpr d0, d1;
-    bindDims(writeOp.getContext(), d0, d1);
     AffineMap map = writeOp.getPermutationMap();
-    bool isTranspose = (map == AffineMap::get(map.getNumDims(), 0, {d1, d0},
-                                              writeOp.getContext()));
-
-    if (!map.isIdentity() && !isTranspose)
+    if (!map.isPermutation())
       return rewriter.notifyMatchFailure(writeOp,
                                          "unsupported permutation map");
 
+    // Note: For 2D vector types the only non-identity permutation is a simple
+    // transpose [1, 0].
+    bool transposed = !map.isIdentity();
     arm_sme::TileSliceLayout layout =
-        isTranspose ? arm_sme::TileSliceLayout::Vertical
-                    : arm_sme::TileSliceLayout::Horizontal;
+        transposed ? arm_sme::TileSliceLayout::Vertical
+                   : arm_sme::TileSliceLayout::Horizontal;
 
     rewriter.replaceOpWithNewOp<arm_sme::TileStoreOp>(
         writeOp, writeOp.getVector(), writeOp.getSource(), writeOp.getIndices(),

[banach-space]

LGTM, thanks!