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

[Andrzej Warzyński]

================
@@ -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");
+
----------------
banach-space wrote:

Not a fan of nested "else if" :)
```suggestion
     bool isTranspose = (map == AffineMap::get(map.getNumDims(), 0, {d1, d0},
                                               transferReadOp.getContext()));
    if (!map.isIdentity() && !isTranspose)
        return rewriter.notifyMatchFailure(transferReadOp,
                                           "unsupported permutation map");
     arm_sme::TileSliceLayout layout =
         isTranspose ? arm_sme::TileSliceLayout::Vertical
                     : arm_sme::TileSliceLayout::Horizontal;
```

Apologies for the formatting - GitHub is not great for these things.

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