[Mlir-commits] [mlir] [mlir][ArmSME] Lower transfer_write + transpose to vertical store (PR #71181)

[Cullen Rhodes]

================
@@ -153,12 +171,35 @@ struct TransferWriteToArmSMELowering
     if (!arm_sme::isValidSMETileVectorType(vType))
       return failure();
 
+    assert(writeOp.getTransferRank() == 2 &&
+           "expected a permutation_map with result dims of the same rank as "
+           "the vector type");
+
     if (!llvm::isa<MemRefType>(writeOp.getSource().getType()))
       return failure();
 
+    // Out-of-bounds dims are not supported.
+    if (writeOp.hasOutOfBoundsDim())
+      return rewriter.notifyMatchFailure(writeOp,
+                                         "not inbounds transfer write");
+
+    arm_sme::TileSliceLayout layout;
+
+    AffineExpr d0, d1;
+    bindDims(writeOp.getContext(), d0, d1);
+    AffineMap map = writeOp.getPermutationMap();
+    if (map.isIdentity())
+      layout = arm_sme::TileSliceLayout::Horizontal;
+    else if (map == AffineMap::get(map.getNumDims(), 0, {d1, d0},
+                                   writeOp.getContext()))
+      layout = arm_sme::TileSliceLayout::Vertical;
+    else
+      return rewriter.notifyMatchFailure(writeOp,
+                                         "unsupported permutation map");
+
----------------
c-rhodes wrote:

that is nicer 👍 

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