[mlir][ArmSME] Lower loads/stores of (.Q) 128-bit tiles to intrinsics

[benmxwl-arm]

Summary

This follows from D155306.

Loads and stores of 128-bit tiles have been confirmed to work in the
load-store-128-bit-tile.mlir integration test. However, there is
currently a bug in QEMU (see: https://gitlab.com/qemu-project/qemu/-/issues/1833)
which means this test produces incorrect results (a patch for this issue
is available but not yet in any released version of QEMU). Until a
fixed version of QEMU is available the integration test is expected to fail.

[c-rhodes]

I've left a few minor comments but otherwise LGTM, thanks Ben!

mlir/test/Dialect/ArmSME/vector-ops-to-llvm.mlir

223–224

nit: vertically align

383–385

nit: vertically align

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

17–18

I wonder if XFAIL would make more sense since it will alert us when QEMU starts behaving correctly?

40

nit: naming

57

nit: MLIR comments use two slashes

63

nit: missing grammar

64–65

missing memref.dealloc (in the other tests I used memref.alloca instead, no strong preference either way)

70–71

It's not clear what the rationale for using i8s is, please could you clarify this in a comment?

[awarzynski]

LGTM, ta!

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

57

[nit] "128-bit tile" could be read as "the size of this tile is 128 bit". Perhaps something like `128-bit tile, e.g. ZA{n}.q, in bytes" instead?

[benmxwl-arm]

• Address various nits :)

[benmxwl-arm]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

57

The triple slashes are a practice copied from elsewhere in LLVM to distinguish between CHECK/RUN comments and normal comments :)

[c-rhodes]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

57

The triple slashes are a practice copied from elsewhere in LLVM to distinguish between CHECK/RUN comments and normal comments :)

Wasnt aware of that, I figured it was Doxygen hangover, thanks for clarifying!

[awarzynski]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

57

Same. I don't mind what we use, but lets try to be consistent ;-)

[dcaballe]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

11

Just my opinion here, feel free to ignore: I would recommend keeping the run command as simple as possible or, at least, as easy to copy-paste as possible. When a test fails and we want to quickly repro/debug locally, copy-pasting multiple lines and remove all the surrounding decoration is not very exciting :)

[benmxwl-arm]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

11

To run/debug locally I do:

./bin/llvm-lit <path-to-test> -va

This will echo the commands lit is running, after they've been expanded and made into single liners. I find this easier than trying to copy them from the tests themselves (which vary a lot) :)

[awarzynski]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

11

I also like this style and like Ben use -va when I need something that I can copy & paste to re-run. I find that much more convenient than copying things from a test file - all the variables (including filenames) get expanded for you auto-magically :)

mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp

matchAndRewrite(arm_sme::LoadTileSliceOp loadTileSliceOp,
     case 32:
       rewriter.create<arm_sme::aarch64_sme_ld1w_horiz>(loc, allActiveMask, ptr,
                                                        tileI32, tileSliceI32);
       break;
     case 64:
       rewriter.create<arm_sme::aarch64_sme_ld1d_horiz>(loc, allActiveMask, ptr,
                                                        tileI32, tileSliceI32);
       break;
+    case 128:
+      rewriter.create<arm_sme::aarch64_sme_ld1q_horiz>(loc, allActiveMask, ptr,
+                                                       tileI32, tileSliceI32);
+      break;
     }
 
     // The load intrinsics have no result, replace 'arm_sme.tile_load' with
     // 'arm_sme.cast_tile_to_vector' to preserve dataflow.
     rewriter.replaceOpWithNewOp<arm_sme::CastTileToVector>(loadTileSliceOp,
                                                            tileType, tile);
 
     return success();
matchAndRewrite(arm_sme::StoreTileSliceOp storeTileSliceOp,
     case 32:
       rewriter.replaceOpWithNewOp<arm_sme::aarch64_sme_st1w_horiz>(
           storeTileSliceOp, allActiveMask, ptr, tileI32, tileSliceI32);
       break;
     case 64:
       rewriter.replaceOpWithNewOp<arm_sme::aarch64_sme_st1d_horiz>(
           storeTileSliceOp, allActiveMask, ptr, tileI32, tileSliceI32);
       break;
+    case 128:
+      rewriter.replaceOpWithNewOp<arm_sme::aarch64_sme_st1q_horiz>(
+          storeTileSliceOp, allActiveMask, ptr, tileI32, tileSliceI32);
+      break;
     }
 
     return success();
   }
 };
 
 } // namespace
 
 void mlir::configureArmSMELegalizeForExportTarget(
     LLVMConversionTarget &target) {
   target.addLegalOp<
       scf::ForOp, scf::YieldOp, arm_sme::CastTileToVector,
       arm_sme::CastVectorToTile, arm_sme::aarch64_sme_zero,
       arm_sme::aarch64_sme_str, arm_sme::aarch64_sme_ld1b_horiz,
       arm_sme::aarch64_sme_ld1h_horiz, arm_sme::aarch64_sme_ld1w_horiz,
-      arm_sme::aarch64_sme_ld1d_horiz, arm_sme::aarch64_sme_st1b_horiz,
-      arm_sme::aarch64_sme_st1h_horiz, arm_sme::aarch64_sme_st1w_horiz,
-      arm_sme::aarch64_sme_st1d_horiz, arm_sme::aarch64_sme_za_enable,
+      arm_sme::aarch64_sme_ld1d_horiz, arm_sme::aarch64_sme_ld1q_horiz,
+      arm_sme::aarch64_sme_st1b_horiz, arm_sme::aarch64_sme_st1h_horiz,
+      arm_sme::aarch64_sme_st1w_horiz, arm_sme::aarch64_sme_st1d_horiz,
+      arm_sme::aarch64_sme_st1q_horiz, arm_sme::aarch64_sme_za_enable,
       arm_sme::aarch64_sme_za_disable>();
   target.addLegalOp<GetTileID>();
 
   // Mark 'func.func' ops as legal if either:
   //   1. no 'arm_za' function attribute is present.
   //   2. the 'arm_za' function attribute is present and the first op in the
   //      function is an 'arm_sme::aarch64_sme_za_enable' intrinsic.
   target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp funcOp) {

mlir/lib/Dialect/ArmSME/Utils/Utils.cpp

 static constexpr unsigned MinStreamingVectorLengthInBits = 128;
 
 unsigned mlir::arm_sme::getSMETileSliceMinNumElts(Type type) {
   assert(isValidSMETileElementType(type) && "invalid tile type!");
   return MinStreamingVectorLengthInBits / type.getIntOrFloatBitWidth();
 }
 
 bool mlir::arm_sme::isValidSMETileElementType(Type type) {
-  // TODO: add support for i128.
   return type.isInteger(8) || type.isInteger(16) || type.isInteger(32) ||
-         type.isInteger(64) || type.isF16() || type.isBF16() || type.isF32() ||
-         type.isF64();
+         type.isInteger(64) || type.isInteger(128) || type.isF16() ||
+         type.isBF16() || type.isF32() || type.isF64() || type.isF128();
 }
 
 bool mlir::arm_sme::isValidSMETileVectorType(VectorType vType) {
   if ((vType.getRank() != 2) && vType.allDimsScalable())
     return false;
 
-  // TODO: add support for i128.
   auto elemType = vType.getElementType();
   if (!isValidSMETileElementType(elemType))
     return false;
 
   unsigned minNumElts = arm_sme::getSMETileSliceMinNumElts(elemType);
   if (vType.getShape() != ArrayRef<int64_t>({minNumElts, minNumElts}))
     return false;
 
   return true;
 }

mlir/test/Dialect/ArmSME/vector-ops-to-llvm.mlir

 func.func @vector_load_f64(%arg0 : memref<?x?xf64>) -> vector<[2]x[2]xf64> {
   %c0 = arith.constant 0 : index
   %tile = vector.load %arg0[%c0, %c0] : memref<?x?xf64>, vector<[2]x[2]xf64>
   return %tile : vector<[2]x[2]xf64>
 }
 
 // -----
 
+// CHECK-LABEL: @vector_load_i128(
+// CHECK-SAME:                   %[[ARG0:.*]]: memref<?x?xi128>)
+// CHECK-DAG: %[[TILE_ID:.*]] = arm_sme.get_tile_id : i128
+// CHECK-DAG: %[[CAST_TILE_TO_VECTOR:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID]] : i128 to vector<[1]x[1]xi128>
+// CHECK:       %[[TILE_ID_I32:.*]] = arith.trunci %[[TILE_ID]] : i128 to i32
+// CHECK:       arm_sme.intr.ld1q.horiz
+func.func @vector_load_i128(%arg0 : memref<?x?xi128>) -> vector<[1]x[1]xi128> {
+  %c0 = arith.constant 0 : index
+  %tile = vector.load %arg0[%c0, %c0] : memref<?x?xi128>, vector<[1]x[1]xi128>
+  return %tile : vector<[1]x[1]xi128>
+}
+
+// -----
+
 // CHECK-LABEL: @vector_store_i8(
 // CHECK-SAME:                   %[[TILE:.*]]: vector<[16]x[16]xi8>,
 // CHECK-SAME:                   %[[ARG0:.*]]: memref<?x?xi8>)
 // CHECK-DAG:  %[[MEM_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref<?x?xi8> to !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG:  %[[MIN_SVL_B:.*]] = arith.constant 16 : index
 // CHECK-DAG:  %[[C0_I64:.*]] = builtin.unrealized_conversion_cast %[[C0]] : index to i64

 // CHECK:       %[[CAST_VECTOR_TO_TILE:.*]] = arm_sme.cast_vector_to_tile %[[TILE]] : vector<[2]x[2]xf64> to i64
 // CHECK:       %[[TILE_ID_I32:.*]] = arith.trunci %[[CAST_VECTOR_TO_TILE]] : i64 to i32
 // CHECK:       arm_sme.intr.st1d.horiz
 func.func @vector_store_f64(%tile : vector<[2]x[2]xf64>, %arg0 : memref<?x?xf64>) {
   %c0 = arith.constant 0 : index
   vector.store %tile, %arg0[%c0, %c0] : memref<?x?xf64>, vector<[2]x[2]xf64>
   return
 }
+
+// -----
+
+// CHECK-LABEL: @vector_store_i128(
+// CHECK-SAME:                   %[[TILE:.*]]: vector<[1]x[1]xi128>,
+// CHECK-SAME:                   %[[ARG0:.*]]: memref<?x?xi128>)
+// CHECK:       %[[CAST_VECTOR_TO_TILE:.*]] = arm_sme.cast_vector_to_tile %[[TILE]] : vector<[1]x[1]xi128> to i128
+// CHECK:       %[[TILE_ID_I32:.*]] = arith.trunci %[[CAST_VECTOR_TO_TILE]] : i128 to i32
+// CHECK:       arm_sme.intr.st1q.horiz
+func.func @vector_store_i128(%tile : vector<[1]x[1]xi128>, %arg0 : memref<?x?xi128>) {
+  %c0 = arith.constant 0 : index
+  vector.store %tile, %arg0[%c0, %c0] : memref<?x?xi128>, vector<[1]x[1]xi128>
+  return
+}

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/load-store-128-bit-tile.mlir

+// DEFINE: %{entry_point} = test_load_store_zaq0
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE:   -enable-arm-streaming="mode=locally enable-za" \
+// DEFINE:   -convert-vector-to-arm-sme -convert-arm-sme-to-scf \
+// DEFINE:   -convert-vector-to-llvm="enable-arm-sme" -cse -canonicalize \
+// DEFINE:   -allocate-arm-sme-tiles -test-lower-to-llvm
+// DEFINE: %{run} = %mcr_aarch64_cmd \
+// DEFINE:  -march=aarch64 -mattr=+sve,+sme \
+// DEFINE:  -e %{entry_point} -entry-point-result=void \
+// DEFINE:  -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
+
+// RUN: %{compile} | %{run} | FileCheck %s
+
+/// Note: The SME ST1Q/LD1Q instructions are currently broken in QEMU
+/// see: https://gitlab.com/qemu-project/qemu/-/issues/1833
+/// This test is disabled until a fixed version of QEMU is available.
+/// FIXME: Enable this test once QEMU is confirmed to be fixed.
+// UNSUPPORTED: {{.*}}
+
+func.func @print_i8s(%arg0: memref<?xi8>, %len: index) {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  scf.for %i = %c0 to %len step %c16 {
+    %v = vector.load %arg0[%i] : memref<?xi8>, vector<16xi8>
+    vector.print %v : vector<16xi8>
+  }
+  return
+}
+
+llvm.func @printCString(!llvm.ptr<i8>)
+
+func.func @printStr(%s: !llvm.ptr<array<17 x i8>>) {
+  %c0 = llvm.mlir.constant(0 : index) : i64
+  %str_bytes = llvm.getelementptr %s[%c0, %c0]
+    : (!llvm.ptr<array<17 x i8>>, i64, i64) -> !llvm.ptr<i8>
+  llvm.call @printCString(%str_bytes) : (!llvm.ptr<i8>) -> ()
+  return
+}
+
+func.func @vector_copy_i128(%arg0: memref<?x?xi128>, %arg1: memref<?x?xi128>) {
+  %c0 = arith.constant 0 : index
+  %tile = vector.load %arg0[%c0, %c0] : memref<?x?xi128>, vector<[1]x[1]xi128>
+  vector.store %tile, %arg1[%c0, %c0] : memref<?x?xi128>, vector<[1]x[1]xi128>
+  return
+}
+
+func.func @test_load_store_zaq0() {
+  %init_a_str = llvm.mlir.addressof @init_tile_a : !llvm.ptr<array<17 x i8>>
+  %init_b_str = llvm.mlir.addressof @init_tile_b : !llvm.ptr<array<17 x i8>>
+  %final_a_str = llvm.mlir.addressof @final_tile_a : !llvm.ptr<array<17 x i8>>
+  %final_b_str = llvm.mlir.addressof @final_tile_b : !llvm.ptr<array<17 x i8>>
+
+  %c0 = arith.constant 0 : index
+  %min_elts_q = arith.constant 1 : index
+  %bytes_per_128_bit = arith.constant 16 : index
+
+  /// Calculate the size of an 128-bit tile in bytes:
+  %vscale = vector.vscale
+  %svl_q = arith.muli %min_elts_q, %vscale : index
+  %zaq_size = arith.muli %svl_q, %svl_q : index
+  %zaq_size_bytes = arith.muli %zaq_size, %bytes_per_128_bit : index
+
+  /// Allocate memory two 128-bit tiles (A and B) and fill them a constant.
+  %tile_a_bytes = memref.alloc(%zaq_size_bytes) : memref<?xi8>
+  %tile_b_bytes = memref.alloc(%zaq_size_bytes) : memref<?xi8>
+  %fill_a_i8 = arith.constant 7 : i8
+  %fill_b_i8 = arith.constant 64 : i8
+  linalg.fill ins(%fill_a_i8 : i8) outs(%tile_a_bytes : memref<?xi8>)
+  linalg.fill ins(%fill_b_i8 : i8) outs(%tile_b_bytes : memref<?xi8>)
+
+  /// Get an 128-bit view of the memory for tiles A and B:
+  %tile_a = memref.view %tile_a_bytes[%c0][%svl_q, %svl_q] :
+    memref<?xi8> to memref<?x?xi128>
+  %tile_b = memref.view %tile_b_bytes[%c0][%svl_q, %svl_q] :
+    memref<?xi8> to memref<?x?xi128>
+
+  // CHECK-LABEL: INITIAL TILE A:
+  // CHECK: ( 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7 )
+  func.call @printStr(%init_a_str) : (!llvm.ptr<array<17 x i8>>) -> ()
+  func.call @print_i8s(%tile_a_bytes, %zaq_size_bytes) : (memref<?xi8>, index) -> ()
+  vector.print punctuation <newline>
+
+  // CHECK-LABEL: INITIAL TILE B:
+  // CHECK: ( 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64 )
+  func.call @printStr(%init_b_str) : (!llvm.ptr<array<17 x i8>>) -> ()
+  func.call @print_i8s(%tile_b_bytes, %zaq_size_bytes) : (memref<?xi8>, index) -> ()
+  vector.print punctuation <newline>
+
+  /// Load tile A and store it to tile B:
+  func.call @vector_copy_i128(%tile_a, %tile_b) : (memref<?x?xi128>, memref<?x?xi128>) -> ()
+
+  // CHECK-LABEL: FINAL TILE A:
+  // CHECK: ( 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7 )
+  func.call @printStr(%final_a_str) : (!llvm.ptr<array<17 x i8>>) -> ()
+  func.call @print_i8s(%tile_a_bytes, %zaq_size_bytes) : (memref<?xi8>, index) -> ()
+  vector.print punctuation <newline>
+
+  // CHECK-LABEL: FINAL TILE B:
+  // CHECK: ( 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7 )
+  func.call @printStr(%final_b_str) : (!llvm.ptr<array<17 x i8>>) -> ()
+  func.call @print_i8s(%tile_b_bytes, %zaq_size_bytes) : (memref<?xi8>, index) -> ()
+
+  return
+}
+
+llvm.mlir.global internal constant @init_tile_a ("INITIAL TILE A:\0A\00")
+llvm.mlir.global internal constant @init_tile_b ("INITIAL TILE B:\0A\00")
+llvm.mlir.global internal constant @final_tile_a("  FINAL TILE A:\0A\00")
+llvm.mlir.global internal constant @final_tile_b("  FINAL TILE B:\0A\00")