[Mlir-commits] [mlir] [mlir][ArmSME] Support vertical layout in load and store ops (PR #66758)

[Andrzej Warzyński]

================
@@ -204,29 +204,51 @@ struct LoadTileSliceToArmSMELowering
     auto allActiveMask = rewriter.create<vector::SplatOp>(loc, predTy, one);
 
     auto tileI32 = castTileIDToI32(tile, loc, rewriter);
-    // Create 'arm_sme.intr.ld1*.horiz' intrinsic to load ZA tile slice.
+    arm_sme::TileSliceLayout layout = loadTileSliceOp.getLayout();
+
+    // Create 'arm_sme.intr.ld1*.(horiz|vert)' intrinsic to load ZA tile slice.
     switch (tileElementWidth) {
     default:
       llvm_unreachable("unexpected element type!");
     case 8:
-      rewriter.create<arm_sme::aarch64_sme_ld1b_horiz>(loc, allActiveMask, ptr,
-                                                       tileI32, tileSliceI32);
+      if (layout == arm_sme::TileSliceLayout::Horizontal)
+        rewriter.create<arm_sme::aarch64_sme_ld1b_horiz>(
+            loc, allActiveMask, ptr, tileI32, tileSliceI32);
+      else
+        rewriter.create<arm_sme::aarch64_sme_ld1b_vert>(loc, allActiveMask, ptr,
+                                                        tileI32, tileSliceI32);
----------------
banach-space wrote:

Wondering how to get rid of these `if` statements. You could try this:

```suggestion
      bool isHorizontal = (layout == arm_sme::TileSliceLayout::Horizonta);
      isHorizontal ? 
        rewriter.create<arm_sme::aarch64_sme_ld1b_horiz>(
            loc, allActiveMask, ptr, tileI32, tileSliceI32) :
        rewriter.create<arm_sme::aarch64_sme_ld1b_vert>(
            loc, allActiveMask, ptr, tileI32, tileSliceI32);
```

or this:

```cpp
if (layout == arm_sme::TileSliceLayout::Horizontal) {
    switch (tileElementWidth) {
    default:
       llvm_unreachable("unexpected element type!");
    case 8:
    // (...)
    break;
} else {
    switch (tileElementWidth) {
    default:
       llvm_unreachable("unexpected element type!");
    case 8:
    // (...)
    break;
}
}
```

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