[Mlir-commits] [mlir] [mlir][sme] Add e2e test for lowering mmt4d to sme (PR #208226)
Federico Bruzzone
llvmlistbot at llvm.org
Sun Jul 12 02:23:39 PDT 2026
================
@@ -0,0 +1,362 @@
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE: -canonicalize -test-lower-to-arm-sme -convert-vector-to-llvm="enable-arm-sve" \
+// DEFINE: -test-lower-to-llvm
+// DEFINE: %{entry_point} = main
+// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sme"\
+// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib
+
+// RUN: %{compile} | %{run} | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+/// HIGH-LEVEL OVERVIEW
+///
+/// End-to-end test for computing matrix-multiplication using linalg.mmt4d. In
+/// particular, demonstrates how the following MLIR sequence (implemented in
+/// @matmul_via_mmt4d):
+///
+/// A_pack = linalg.pack A
+/// B_pack = linalg.pack B
+/// C_pack = linalg.pack C
+/// out_pack = linalg.mmt4d(A_pack, B_pack, C_pack)
+/// out = linalg.unpack out_pack
+///
+/// is equivalent to:
+///
+/// linalg.matmul(A, B, C)
+///
+/// (implemented in @matmul_via_matmul).
+///
+/// NOTES ON IMPLEMENTATION
+/// 1. The MMT4D example uses _scalable_ tile sizes for data tiling.
+/// * The matrix-multiplication dimensions that are scalable: M and N.
+///
+/// 2. The lowering of linalg.mmt4d leverages scalable vectorisation.
+/// * The matrix-multiplication dimension that's scalable: M, N (to match data
+/// tiling configuration).
+///
+/// 3. Neither `linalg.pack` nor `linalg.unpack` are vectorised ATM.
+///
+/// 4. The MMT4D and Pack/Unpack Ops are kept in seperate functions to isolate
+/// the corresponding lowering and lowering configs.
+///
+/// TODO: Ideally, we should consider fusion opportunities by moving
+/// pack/unpack/mmt4d ops into one function:
+/// * https://github.com/llvm/llvm-project/issues/159770
+/// TODO: Vectorize linalg.pack + linalg.unpack:
+/// * https://github.com/llvm/llvm-project/issues/159751
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// @main
+//
+// The main entry point that computes matrix multiplication via linalg.mmt4d
+// and linalg.matmul. Note, the output should be independent of the underlying
+// Linalg Op used, as well as SVE vector length.
+//===----------------------------------------------------------------------===//
+func.func @main() {
+ // Allocate and initialise the inputs
+ %A_empty = tensor.empty() : tensor<7x16xf32>
+ %B_empty = tensor.empty() : tensor<16x13xf32>
+
+ %c3 = arith.constant 3.0 : f32
+ %c4 = arith.constant 4.0 : f32
+ %A = linalg.fill ins(%c3 : f32) outs(%A_empty : tensor<7x16xf32>) -> tensor<7x16xf32>
+ %B = linalg.fill ins(%c4 : f32) outs(%B_empty : tensor<16x13xf32>) -> tensor<16x13xf32>
+ %C = arith.constant dense<[
+ [ 1., 8., 15., 22., 29., 36., 43., 50., 57., 64., 71., 78., 85.],
+ [ 2., 9., 16., 23., 30., 37., 44., 51., 58., 65., 72., 79., 86.],
+ [ 3., 10., 17., 24., 31., 38., 45., 52., 59., 66., 73., 80., 87.],
+ [ 4., 11., 18., 25., 32., 39., 46., 53., 60., 67., 74., 81., 88.],
+ [ 5., 12., 19., 26., 33., 40., 47., 54., 61., 68., 75., 82., 89.],
+ [ 6., 13., 20., 27., 34., 41., 48., 55., 62., 69., 76., 83., 90.],
+ [ 7., 14., 21., 28., 35., 42., 49., 56., 63., 70., 77., 84., 91.]
+ ]> : tensor<7x13xf32>
+
+ // VARIANT: Matrix multiplication via linalg.mmt4d
+ // CHECK: Unranked Memref
+ // CHECK: [193, 200, 207, 214, 221, 228, 235, 242, 249, 256, 263, 270, 277]
+ // CHECK: [194, 201, 208, 215, 222, 229, 236, 243, 250, 257, 264, 271, 278]
+ // CHECK: [195, 202, 209, 216, 223, 230, 237, 244, 251, 258, 265, 272, 279]
+ // CHECK: [196, 203, 210, 217, 224, 231, 238, 245, 252, 259, 266, 273, 280]
+ // CHECK: [197, 204, 211, 218, 225, 232, 239, 246, 253, 260, 267, 274, 281]
+ // CHECK: [198, 205, 212, 219, 226, 233, 240, 247, 254, 261, 268, 275, 282]
+ // CHECK: [199, 206, 213, 220, 227, 234, 241, 248, 255, 262, 269, 276, 283]
+ %C_mmt4d = func.call @matmul_via_mmt4d(%A, %B, %C) : (tensor<7x16xf32>, tensor<16x13xf32>, tensor<7x13xf32>) -> tensor<7x13xf32>
+ %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<7x13xf32> to tensor<*xf32>
+ vector.print str "--------------------------\n"
+ vector.print str "RESULT FROM linalg.mmt4d:\n"
+ vector.print str "--------------------------\n"
+ call @printMemrefF32(%C_mmt4d_cast) : (tensor<*xf32>) -> ()
+
+ // VARIANT: Matrix multiplication via linalg.matmul
+ // CHECK: Unranked Memref
+ // CHECK: [193, 200, 207, 214, 221, 228, 235, 242, 249, 256, 263, 270, 277]
+ // CHECK: [194, 201, 208, 215, 222, 229, 236, 243, 250, 257, 264, 271, 278]
+ // CHECK: [195, 202, 209, 216, 223, 230, 237, 244, 251, 258, 265, 272, 279]
+ // CHECK: [196, 203, 210, 217, 224, 231, 238, 245, 252, 259, 266, 273, 280]
+ // CHECK: [197, 204, 211, 218, 225, 232, 239, 246, 253, 260, 267, 274, 281]
+ // CHECK: [198, 205, 212, 219, 226, 233, 240, 247, 254, 261, 268, 275, 282]
+ // CHECK: [199, 206, 213, 220, 227, 234, 241, 248, 255, 262, 269, 276, 283]
+ %C_matmul = func.call @matmul(%A, %B, %C) : (tensor<7x16xf32>, tensor<16x13xf32>, tensor<7x13xf32>) -> tensor<7x13xf32>
+ %C_matmul_cast = tensor.cast %C_matmul : tensor<7x13xf32> to tensor<*xf32>
+ vector.print str "\n--------------------------\n"
+ vector.print str "RESULT FROM linalg.matmul:\n"
+ vector.print str "--------------------------\n"
+ call @printMemrefF32(%C_matmul_cast) : (tensor<*xf32>) -> ()
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @matmul_via_matmul
+//
+// Implements matrix-multiplication via linalg.matmul
+//===----------------------------------------------------------------------===//
+func.func private @matmul(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tensor<7x13xf32>) -> tensor<7x13xf32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<7x16xf32>, tensor<16x13xf32>)
+ outs(%C: tensor<7x13xf32>) -> tensor<7x13xf32>
+
+ return %C_matmul : tensor<7x13xf32>
+}
+
+//===----------------------------------------------------------------------===//
+// @matmul_via_mmt4d
+//
+// Implements matrix-multiplication via linalg.mmt4d
+//===----------------------------------------------------------------------===//
+func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x?x1xf32> {
----------------
FedericoBruzzone wrote:
A slight stylistic inconsistency: `pack_lhs` now fixes the outer dimension M as static at 1 (`tensor<1x16x?x1xf32>`), whereas `pack_acc` keeps it fully dynamic (`tensor<?x?x?x?xf32>`), even though the calculation is conceptually the same.
Is there a reason why we can't align with something like [ArmSVE/pack-unpack-mmt4d](https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/ArmSVE/pack-unpack-mmt4d.mlir)?
https://github.com/llvm/llvm-project/pull/208226
More information about the Mlir-commits
mailing list