[Mlir-commits] [mlir] [mlir][ArmSME] Simplify permutation map handling (PR #93515)

[Cullen Rhodes]

https://github.com/c-rhodes created https://github.com/llvm/llvm-project/pull/93515

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.

>From b75793a7a7165f9ef3dd343d176ec7a966f0e954 Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Thu, 23 May 2024 14:01:32 +0000
Subject: [PATCH] [mlir][ArmSME] Simplify permutation map handling

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.
---
 .../VectorToArmSME/VectorToArmSME.cpp         | 32 ++++++++-----------
 1 file changed, 14 insertions(+), 18 deletions(-)

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(),