[Mlir-commits] [mlir] [mlir][ArmNeon] Add pack/mmt4d path to f32 linalg.matmul integration test (PR #225063)
Federico Bruzzone
llvmlistbot at llvm.org
Sun Oct 4 07:59:39 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/225063
>From be40d49f89d7873a2060e702e9a5b3350aa614e2 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 21 Sep 2026 13:16:44 +0200
Subject: [PATCH 1/5] [mlir][ArmNeon] Add pack/mmt4d path to f32 linalg.matmul
test
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Linalg/CPU/ArmNeon/matmul-f32.mlir | 269 ++++++++++++++----
1 file changed, 220 insertions(+), 49 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
index 3a55d2a890eef..4bac3a8ec81bb 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
@@ -1,52 +1,153 @@
// REQUIRES: target={{(aarch64|arm64).*}}
-// RUN: mlir-opt %s \
-// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
-// RUN: -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline -cse -canonicalize -convert-vector-to-scf \
-// RUN: -convert-vector-to-llvm="enable-arm-neon" -test-lower-to-llvm -o %t
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline \
+// DEFINE: -convert-bufferization-to-memref -cse -canonicalize -convert-vector-to-scf \
+// DEFINE: -convert-vector-to-llvm="enable-arm-neon" -test-lower-to-llvm \
+// DEFINE: -o %t
-// RUN: %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+neon" \
-// RUN: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils | \
-// RUN: FileCheck %s
+// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+neon" \
+// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
+
+// RUN: rm -f %t && %{compile} && FileCheck %s --input-file=%t -check-prefix CHECK-IR && %{run} | FileCheck %s
func.func @main() {
- // Matrix dimensions
- %K = arith.constant 3 : index
- %M = arith.constant 5 : index
- %N = arith.constant 15 : index
- %c0_f32 = arith.constant 0.0 : f32
-
- // Allocate the matrices
- %A_alloc = bufferization.alloc_tensor(%M, %K) : tensor<?x?xf32>
- %B_alloc = bufferization.alloc_tensor(%K, %N) : tensor<?x?xf32>
- %C_alloc = bufferization.alloc_tensor(%M, %N) : tensor<?x?xf32>
-
- // Initialise the matrices
- %pi = arith.constant 3.14 : f32
- %A = linalg.fill ins(%pi : f32) outs(%A_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>
- %B = linalg.fill ins(%pi : f32) outs(%B_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>
- %C_in = linalg.fill ins(%c0_f32 : f32) outs(%C_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>
-
- // Matmul
- %C_out = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>) outs(%C_in: tensor<?x?xf32>) -> tensor<?x?xf32>
-
- // Print and verify the output
- // CHECK-LABEL: NEON: START OF TEST OUTPUT
- vector.print str "NEON: START OF TEST OUTPUT\n"
-
- // CHECK-NEXT: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [5, 15] strides = [15, 1] data =
- // CHECK-COUNT-5: [29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788, 29.5788]
- %xf = tensor.cast %C_out : tensor<?x?xf32> to tensor<*xf32>
- call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()
-
- // CHECK-NEXT: NEON: END OF TEST OUTPUT
- vector.print str "NEON: END OF TEST OUTPUT\n"
+ %A = arith.constant dense<[
+ [1.0, 3.0, 5.0],
+ [2.0, 4.0, 6.0],
+ [3.0, 5.0, 7.0],
+ [4.0, 6.0, 8.0],
+ [5.0, 7.0, 9.0]
+ ]> : tensor<5x3xf32>
+ %B = arith.constant dense<[
+ [1.0, 4.0, 7.0, 10.0, 13.0, 16.0, 19.0, 22.0, 25.0, 28.0, 31.0, 34.0, 37.0, 40.0, 43.0],
+ [2.0, 5.0, 8.0, 11.0, 14.0, 17.0, 20.0, 23.0, 26.0, 29.0, 32.0, 35.0, 38.0, 41.0, 44.0],
+ [3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0, 27.0, 30.0, 33.0, 36.0, 39.0, 42.0, 45.0]
+ ]> : tensor<3x15xf32>
+ %C = arith.constant dense<[
+ [1.0, 6.0, 11.0, 16.0, 21.0, 26.0, 31.0, 36.0, 41.0, 46.0, 51.0, 56.0, 61.0, 66.0, 71.0],
+ [2.0, 7.0, 12.0, 17.0, 22.0, 27.0, 32.0, 37.0, 42.0, 47.0, 52.0, 57.0, 62.0, 67.0, 72.0],
+ [3.0, 8.0, 13.0, 18.0, 23.0, 28.0, 33.0, 38.0, 43.0, 48.0, 53.0, 58.0, 63.0, 68.0, 73.0],
+ [4.0, 9.0, 14.0, 19.0, 24.0, 29.0, 34.0, 39.0, 44.0, 49.0, 54.0, 59.0, 64.0, 69.0, 74.0],
+ [5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0, 50.0, 55.0, 60.0, 65.0, 70.0, 75.0]
+ ]> : tensor<5x15xf32>
+
+ // VARIANT: Matrix multiplication via linalg.mmt4d
+ // CHECK: Unranked Memref
+ // CHECK: [23, 55, 87, 119, 151, 183, 215, 247, 279, 311, 343, 375, 407, 439, 471]
+ // CHECK: [30, 71, 112, 153, 194, 235, 276, 317, 358, 399, 440, 481, 522, 563, 604]
+ // CHECK: [37, 87, 137, 187, 237, 287, 337, 387, 437, 487, 537, 587, 637, 687, 737]
+ // CHECK: [44, 103, 162, 221, 280, 339, 398, 457, 516, 575, 634, 693, 752, 811, 870]
+ // CHECK: [51, 119, 187, 255, 323, 391, 459, 527, 595, 663, 731, 799, 867, 935, 1003]
+ %C_mmt4d = func.call @matmul_via_mmt4d(%A, %B, %C) : (tensor<5x3xf32>, tensor<3x15xf32>, tensor<5x15xf32>) -> tensor<5x15xf32>
+ %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<5x15xf32> to tensor<*xf32>
+ vector.print str "RESULT FROM linalg.mmt4d:\n"
+ call @printMemrefF32(%C_mmt4d_cast) : (tensor<*xf32>) -> ()
+
+ // VARIANT: Matrix multiplication via linalg.matmul (cross-check)
+ // CHECK: Unranked Memref
+ // CHECK: [23, 55, 87, 119, 151, 183, 215, 247, 279, 311, 343, 375, 407, 439, 471]
+ // CHECK: [30, 71, 112, 153, 194, 235, 276, 317, 358, 399, 440, 481, 522, 563, 604]
+ // CHECK: [37, 87, 137, 187, 237, 287, 337, 387, 437, 487, 537, 587, 637, 687, 737]
+ // CHECK: [44, 103, 162, 221, 280, 339, 398, 457, 516, 575, 634, 693, 752, 811, 870]
+ // CHECK: [51, 119, 187, 255, 323, 391, 459, 527, 595, 663, 731, 799, 867, 935, 1003]
+ %C_matmul = func.call @matmul(%A, %B, %C) : (tensor<5x3xf32>, tensor<3x15xf32>, tensor<5x15xf32>) -> tensor<5x15xf32>
+ %C_matmul_cast = tensor.cast %C_matmul : tensor<5x15xf32> to tensor<*xf32>
+ vector.print str "RESULT FROM linalg.matmul:\n"
+ call @printMemrefF32(%C_matmul_cast) : (tensor<*xf32>) -> ()
+ // VARIANT: Matrix multiplication via linalg.matmul on dynamic shapes
+ // CHECK: Unranked Memref
+ // CHECK: [23, 55, 87, 119, 151, 183, 215, 247, 279, 311, 343, 375, 407, 439, 471]
+ // CHECK: [30, 71, 112, 153, 194, 235, 276, 317, 358, 399, 440, 481, 522, 563, 604]
+ // CHECK: [37, 87, 137, 187, 237, 287, 337, 387, 437, 487, 537, 587, 637, 687, 737]
+ // CHECK: [44, 103, 162, 221, 280, 339, 398, 457, 516, 575, 634, 693, 752, 811, 870]
+ // CHECK: [51, 119, 187, 255, 323, 391, 459, 527, 595, 663, 731, 799, 867, 935, 1003]
+ %A_dyn = tensor.cast %A : tensor<5x3xf32> to tensor<?x?xf32>
+ %B_dyn = tensor.cast %B : tensor<3x15xf32> to tensor<?x?xf32>
+ %C_dyn = tensor.cast %C : tensor<5x15xf32> to tensor<?x?xf32>
+ %C_matmul_dyn = func.call @matmul_dynamic(%A_dyn, %B_dyn, %C_dyn) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+ %C_matmul_dyn_cast = tensor.cast %C_matmul_dyn : tensor<?x?xf32> to tensor<*xf32>
+ vector.print str "RESULT FROM linalg.matmul (dynamic):\n"
+ call @printMemrefF32(%C_matmul_dyn_cast) : (tensor<*xf32>) -> ()
return
}
-module attributes {transform.with_named_sequence} {
- // Tile and vectorize the matmul.
+func.func private @matmul(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %C: tensor<5x15xf32>) -> tensor<5x15xf32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<5x3xf32>, tensor<3x15xf32>)
+ outs(%C: tensor<5x15xf32>) -> tensor<5x15xf32>
+ return %C_matmul : tensor<5x15xf32>
+}
+
+func.func private @matmul_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %C_matmul : tensor<?x?xf32>
+}
+
+// LHS packed tile: M0=4, K0=1.
+func.func private @pack_lhs(%A: tensor<5x3xf32>) -> tensor<2x3x4x1xf32> {
+ %pad = arith.constant 0.0 : f32
+ %A_pack_empty = tensor.empty() : tensor<2x3x4x1xf32>
+ %A_pack = linalg.pack %A
+ padding_value(%pad : f32)
+ inner_dims_pos = [0, 1]
+ inner_tiles = [4, 1]
+ into %A_pack_empty : tensor<5x3xf32> -> tensor<2x3x4x1xf32>
+ return %A_pack : tensor<2x3x4x1xf32>
+}
+
+// RHS packed tile: N0=4, K0=1.
+func.func private @pack_rhs(%B: tensor<3x15xf32>) -> tensor<4x3x4x1xf32> {
+ %pad = arith.constant 0.0 : f32
+ %B_pack_empty = tensor.empty() : tensor<4x3x4x1xf32>
+ %B_pack = linalg.pack %B
+ padding_value(%pad : f32)
+ outer_dims_perm = [1, 0]
+ inner_dims_pos = [1, 0]
+ inner_tiles = [4, 1]
+ into %B_pack_empty : tensor<3x15xf32> -> tensor<4x3x4x1xf32>
+ return %B_pack : tensor<4x3x4x1xf32>
+}
+
+func.func private @pack_acc(%C: tensor<5x15xf32>) -> tensor<2x4x4x4xf32> {
+ %pad = arith.constant 0.0 : f32
+ %C_pack_empty = tensor.empty() : tensor<2x4x4x4xf32>
+ %C_pack = linalg.pack %C
+ padding_value(%pad : f32)
+ outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1]
+ inner_tiles = [4, 4]
+ into %C_pack_empty : tensor<5x15xf32> -> tensor<2x4x4x4xf32>
+ return %C_pack : tensor<2x4x4x4xf32>
+}
+
+func.func private @unpack_acc(%C_packed: tensor<2x4x4x4xf32>) -> tensor<5x15xf32> {
+ %C_out_empty = tensor.empty() : tensor<5x15xf32>
+ %C_out_unpack = linalg.unpack %C_packed
+ outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1]
+ inner_tiles = [4, 4]
+ into %C_out_empty : tensor<2x4x4x4xf32> -> tensor<5x15xf32>
+ return %C_out_unpack: tensor<5x15xf32>
+}
+
+// CHECK-IR-LABEL: llvm.func @matmul_via_mmt4d
+// CHECK-IR-COUNT-4: llvm.intr.fmuladd
+func.func private @matmul_via_mmt4d(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %C: tensor<5x15xf32>) -> tensor<5x15xf32> {
+ %A_pack = func.call @pack_lhs(%A): (tensor<5x3xf32>) -> tensor<2x3x4x1xf32>
+ %B_pack = func.call @pack_rhs(%B): (tensor<3x15xf32>) -> tensor<4x3x4x1xf32>
+ %C_pack = func.call @pack_acc(%C): (tensor<5x15xf32>) -> tensor<2x4x4x4xf32>
+
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<2x3x4x1xf32>, tensor<4x3x4x1xf32>) outs(%C_pack : tensor<2x4x4x4xf32>) -> tensor<2x4x4x4xf32>
+
+ %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<2x4x4x4xf32>) -> tensor<5x15xf32>
+ return %C_out_unpack : tensor<5x15xf32>
+}
+
+module @transforms attributes { transform.with_named_sequence } {
+ // Tile and vectorize the plain `linalg.matmul`.
transform.named_sequence @tile_and_vectorize_matmul(%func
: !transform.op<"func.func"> {transform.readonly}) {
@@ -74,18 +175,88 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
- // Apply `tile_and_vectorize_matmul` to every function in the module.
transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
- %funcs = transform.structured.match ops{["func.func"]} in %module
- : (!transform.any_op) -> !transform.op<"func.func">
-
- transform.foreach %funcs : !transform.op<"func.func"> {
- ^bb2(%func : !transform.op<"func.func">):
- transform.include @tile_and_vectorize_matmul failures(propagate)
- (%func) : (!transform.op<"func.func">) -> ()
- }
+ %matmul_func = transform.structured.match ops{["func.func"]} attributes{sym_name = "matmul"} in %module
+ : (!transform.any_op) -> !transform.op<"func.func">
+ transform.include @tile_and_vectorize_matmul failures(propagate)
+ (%matmul_func) : (!transform.op<"func.func">) -> ()
+
+ %matmul_dynamic_func = transform.structured.match ops{["func.func"]} attributes{sym_name = "matmul_dynamic"} in %module
+ : (!transform.any_op) -> !transform.op<"func.func">
+ transform.include @tile_and_vectorize_matmul failures(propagate)
+ (%matmul_dynamic_func) : (!transform.op<"func.func">) -> ()
+
+ %mmt4d =transform.collect_matching @match_mmt4d in %module : (!transform.any_op) -> (!transform.any_op)
+ %mmt4d_func = transform.get_parent_op %mmt4d <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
+
+ // Tile the outer dims (m, n, k) one at a time; the inner tile
+ // (m0, n0, k0) = (4, 4, 1) is kept whole.
+ %tiled_mmt4d, %_:3 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 1, 0, 0, 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+
+ transform.structured.vectorize %tiled_mmt4d create_named_contraction
+ vector_sizes [1, 1, 1, 4, 4, 1] : !transform.any_op
+
+ transform.apply_patterns to %mmt4d_func {
+ transform.apply_patterns.vector.reduction_to_contract
+ transform.apply_patterns.vector.transfer_permutation_patterns
+ } : !transform.op<"func.func">
+
+ %mmt4d_func_h = transform.structured.hoist_redundant_vector_transfers %mmt4d_func
+ : (!transform.op<"func.func">) -> !transform.op<"func.func">
+ %all_loops = transform.structured.match interface{LoopLikeInterface} in %mmt4d_func_h
+ : (!transform.op<"func.func">) -> !transform.any_op
+ transform.apply_licm to %all_loops : !transform.any_op
+ transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op
+
+ transform.apply_patterns to %mmt4d_func_h {
+ transform.apply_patterns.vector.reduction_to_contract
+ transform.apply_patterns.vector.cast_away_vector_leading_one_dim
+ transform.apply_patterns.canonicalization
+ } : !transform.op<"func.func">
+
+ // Lower to outer products: with k0 = 1, each `vector.contract` is a rank-1
+ // update, i.e., one `vector.fma` per row of the 4x4 accumulator tile.
+ transform.apply_patterns to %mmt4d_func_h {
+ transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
+ transform.apply_patterns.vector.lower_outerproduct
+ transform.apply_patterns.canonicalization
+ } : !transform.op<"func.func">
+
+ %pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op
+ %unpack = transform.structured.match ops{["linalg.unpack"]} in %module : (!transform.any_op) -> !transform.any_op
+
+ %tiled_pack_op_p, %loops_pack:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [4, 4]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+ %func_op_pack = transform.get_parent_op %tiled_pack_op_p <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
+ transform.apply_patterns to %func_op_pack {
+ transform.apply_patterns.linalg.decompose_pack_unpack
+ transform.apply_patterns.linalg.decompose_pad
+ } : !transform.op<"func.func">
+ transform.apply_patterns to %func_op_pack {
+ transform.apply_patterns.tensor.fold_tensor_subset_ops
+ transform.apply_patterns.canonicalization
+ } : !transform.op<"func.func">
+
+ %func_op_unpack = transform.get_parent_op %tiled_unpack_op_p <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
+ transform.apply_patterns to %func_op_unpack {
+ transform.apply_patterns.linalg.decompose_pack_unpack
+ } : !transform.op<"func.func">
+ transform.apply_patterns to %func_op_unpack {
+ transform.apply_patterns.tensor.fold_tensor_subset_ops
+ transform.apply_patterns.canonicalization
+ } : !transform.op<"func.func">
transform.yield
}
+
+ transform.named_sequence @match_mmt4d(
+ %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
+ transform.match.operation_name %entry ["linalg.mmt4d"] : !transform.any_op
+ transform.yield %entry : !transform.any_op
+ }
}
func.func private @printMemrefF32(%ptr : tensor<*xf32>)
>From cbe9bfb31e507f66d7d52c585c6cb49b531780e9 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Sun, 27 Sep 2026 09:59:13 +0200
Subject: [PATCH 2/5] Address Andrzej comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Linalg/CPU/ArmNeon/matmul-f32.mlir | 128 +++++++++++++++++-
1 file changed, 123 insertions(+), 5 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
index 4bac3a8ec81bb..84e2701b6f56e 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
@@ -12,6 +12,47 @@
// RUN: rm -f %t && %{compile} && FileCheck %s --input-file=%t -check-prefix CHECK-IR && %{run} | FileCheck %s
+//===----------------------------------------------------------------------===//
+/// HIGH-LEVEL OVERVIEW
+///
+/// End-to-end test for computing matrix-multiplication (f32, accumulating to
+/// f32) 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)
+///
+/// is equivalent to:
+///
+/// linalg.matmul(A, B, C)
+///
+/// (implemented in both @matmul, static shapes, and @matmul_dynamic, dynamic
+/// shapes).
+///
+/// NOTES ON IMPLEMENTATION
+/// 1. NEON has no f32 matrix-multiply instruction (`fmmla` is SVE-only,
+/// `bfmmla` takes bf16 operands). The inner tile is therefore
+/// (M0, N0, K0) = (4, 4, 1): every vector.contract is a rank-1 update,
+/// lowered to outer products, i.e., one `fmla` per row of the 4x4
+/// accumulator tile.
+///
+/// 2. Packing is still worthwhile without a dedicated instruction: it gives
+/// the tiles a statically-known shape (no masking in the mmt4d loop) and
+/// makes them contiguous in memory.
+///
+/// 3. The MMT4D and Pack/Unpack Ops are kept in separate functions to isolate
+/// the corresponding lowering and lowering configs.
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// @main
+//
+// The main entry point that computes matrix multiplication via
+// linalg.mmt4d and linalg.matmul (static and dynamic shapes). The output
+// should be independent of the Linalg Op used.
+//===----------------------------------------------------------------------===//
func.func @main() {
%A = arith.constant dense<[
[1.0, 3.0, 5.0],
@@ -74,19 +115,36 @@ func.func @main() {
return
}
+//===----------------------------------------------------------------------===//
+// @matmul
+//
+// Implements matrix-multiplication via linalg.matmul, static shapes.
+//===----------------------------------------------------------------------===//
func.func private @matmul(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %C: tensor<5x15xf32>) -> tensor<5x15xf32> {
%C_matmul = linalg.matmul ins(%A, %B: tensor<5x3xf32>, tensor<3x15xf32>)
outs(%C: tensor<5x15xf32>) -> tensor<5x15xf32>
return %C_matmul : tensor<5x15xf32>
}
+//===----------------------------------------------------------------------===//
+// @matmul_dynamic
+//
+// Implements matrix-multiplication via linalg.matmul, dynamic shapes. Unlike
+// the other variants, this is the one that exercises masked vectorization
+// (masks computed from runtime dimensions).
+//===----------------------------------------------------------------------===//
func.func private @matmul_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
%C_matmul = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
return %C_matmul : tensor<?x?xf32>
}
-// LHS packed tile: M0=4, K0=1.
+//===----------------------------------------------------------------------===//
+// @pack_lhs
+//
+// Implements packing for the A matrix (LHS) in matrix multiplication. The
+// inner tile size is fixed: 4 * 1.
+//===----------------------------------------------------------------------===//
func.func private @pack_lhs(%A: tensor<5x3xf32>) -> tensor<2x3x4x1xf32> {
%pad = arith.constant 0.0 : f32
%A_pack_empty = tensor.empty() : tensor<2x3x4x1xf32>
@@ -98,7 +156,12 @@ func.func private @pack_lhs(%A: tensor<5x3xf32>) -> tensor<2x3x4x1xf32> {
return %A_pack : tensor<2x3x4x1xf32>
}
-// RHS packed tile: N0=4, K0=1.
+//===----------------------------------------------------------------------===//
+// @pack_rhs
+//
+// Implements packing for the B matrix (RHS) in matrix multiplication. The
+// inner tile size is fixed: 4 * 1.
+//===----------------------------------------------------------------------===//
func.func private @pack_rhs(%B: tensor<3x15xf32>) -> tensor<4x3x4x1xf32> {
%pad = arith.constant 0.0 : f32
%B_pack_empty = tensor.empty() : tensor<4x3x4x1xf32>
@@ -111,6 +174,12 @@ func.func private @pack_rhs(%B: tensor<3x15xf32>) -> tensor<4x3x4x1xf32> {
return %B_pack : tensor<4x3x4x1xf32>
}
+//===----------------------------------------------------------------------===//
+// @pack_acc
+//
+// Implements packing for the C matrix (accumulator) in matrix multiplication.
+// The inner tile size is fixed: 4 * 4.
+//===----------------------------------------------------------------------===//
func.func private @pack_acc(%C: tensor<5x15xf32>) -> tensor<2x4x4x4xf32> {
%pad = arith.constant 0.0 : f32
%C_pack_empty = tensor.empty() : tensor<2x4x4x4xf32>
@@ -123,6 +192,12 @@ func.func private @pack_acc(%C: tensor<5x15xf32>) -> tensor<2x4x4x4xf32> {
return %C_pack : tensor<2x4x4x4xf32>
}
+//===----------------------------------------------------------------------===//
+// @unpack_acc
+//
+// Implements unpacking for the C matrix (accumulator) in matrix
+// multiplication. The inner tile size is fixed: 4 * 4.
+//===----------------------------------------------------------------------===//
func.func private @unpack_acc(%C_packed: tensor<2x4x4x4xf32>) -> tensor<5x15xf32> {
%C_out_empty = tensor.empty() : tensor<5x15xf32>
%C_out_unpack = linalg.unpack %C_packed
@@ -133,21 +208,37 @@ func.func private @unpack_acc(%C_packed: tensor<2x4x4x4xf32>) -> tensor<5x15xf32
return %C_out_unpack: tensor<5x15xf32>
}
+//===----------------------------------------------------------------------===//
+// @matmul_via_mmt4d
+//
+// Implements matrix-multiplication via linalg.mmt4d.
+//===----------------------------------------------------------------------===//
// CHECK-IR-LABEL: llvm.func @matmul_via_mmt4d
// CHECK-IR-COUNT-4: llvm.intr.fmuladd
func.func private @matmul_via_mmt4d(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %C: tensor<5x15xf32>) -> tensor<5x15xf32> {
+ // Pack input matrices
%A_pack = func.call @pack_lhs(%A): (tensor<5x3xf32>) -> tensor<2x3x4x1xf32>
%B_pack = func.call @pack_rhs(%B): (tensor<3x15xf32>) -> tensor<4x3x4x1xf32>
%C_pack = func.call @pack_acc(%C): (tensor<5x15xf32>) -> tensor<2x4x4x4xf32>
+ // MMT4D
%mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<2x3x4x1xf32>, tensor<4x3x4x1xf32>) outs(%C_pack : tensor<2x4x4x4xf32>) -> tensor<2x4x4x4xf32>
+ // Unpack the output
%C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<2x4x4x4xf32>) -> tensor<5x15xf32>
return %C_out_unpack : tensor<5x15xf32>
}
+//===----------------------------------------------------------------------===//
+// TD Sequence
+//===----------------------------------------------------------------------===//
module @transforms attributes { transform.with_named_sequence } {
- // Tile and vectorize the plain `linalg.matmul`.
+ //===--------------------------------------------------------------------===//
+ // @tile_and_vectorize_matmul
+ //
+ // Tile and vectorize a plain `linalg.matmul` (shared by @matmul and
+ // @matmul_dynamic).
+ //===--------------------------------------------------------------------===//
transform.named_sequence @tile_and_vectorize_matmul(%func
: !transform.op<"func.func"> {transform.readonly}) {
@@ -176,6 +267,9 @@ module @transforms attributes { transform.with_named_sequence } {
}
transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+ //==========================================================================
+ // HANDLE PLAIN MATMUL (static + dynamic)
+ //==========================================================================
%matmul_func = transform.structured.match ops{["func.func"]} attributes{sym_name = "matmul"} in %module
: (!transform.any_op) -> !transform.op<"func.func">
transform.include @tile_and_vectorize_matmul failures(propagate)
@@ -186,22 +280,29 @@ module @transforms attributes { transform.with_named_sequence } {
transform.include @tile_and_vectorize_matmul failures(propagate)
(%matmul_dynamic_func) : (!transform.op<"func.func">) -> ()
- %mmt4d =transform.collect_matching @match_mmt4d in %module : (!transform.any_op) -> (!transform.any_op)
+ //==========================================================================
+ // HANDLE MMT4D
+ //==========================================================================
+ %mmt4d = transform.collect_matching @match_mmt4d in %module : (!transform.any_op) -> (!transform.any_op)
%mmt4d_func = transform.get_parent_op %mmt4d <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
+ // Step 1: Tile
// Tile the outer dims (m, n, k) one at a time; the inner tile
// (m0, n0, k0) = (4, 4, 1) is kept whole.
%tiled_mmt4d, %_:3 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 1, 0, 0, 0]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ // Step 2: Vectorize
transform.structured.vectorize %tiled_mmt4d create_named_contraction
vector_sizes [1, 1, 1, 4, 4, 1] : !transform.any_op
+ // Step 3: Simplify
transform.apply_patterns to %mmt4d_func {
transform.apply_patterns.vector.reduction_to_contract
transform.apply_patterns.vector.transfer_permutation_patterns
} : !transform.op<"func.func">
+ // Hoisting and LICM - not strictly required
%mmt4d_func_h = transform.structured.hoist_redundant_vector_transfers %mmt4d_func
: (!transform.op<"func.func">) -> !transform.op<"func.func">
%all_loops = transform.structured.match interface{LoopLikeInterface} in %mmt4d_func_h
@@ -209,6 +310,7 @@ module @transforms attributes { transform.with_named_sequence } {
transform.apply_licm to %all_loops : !transform.any_op
transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op
+ // Simplification
transform.apply_patterns to %mmt4d_func_h {
transform.apply_patterns.vector.reduction_to_contract
transform.apply_patterns.vector.cast_away_vector_leading_one_dim
@@ -216,21 +318,30 @@ module @transforms attributes { transform.with_named_sequence } {
} : !transform.op<"func.func">
// Lower to outer products: with k0 = 1, each `vector.contract` is a rank-1
- // update, i.e., one `vector.fma` per row of the 4x4 accumulator tile.
+ // update, i.e., one `fmla` per row of the 4x4 accumulator tile.
transform.apply_patterns to %mmt4d_func_h {
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
transform.apply_patterns.vector.lower_outerproduct
transform.apply_patterns.canonicalization
} : !transform.op<"func.func">
+ //==========================================================================
+ // HANDLE PACK + UNPACK
+ //==========================================================================
%pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op
%unpack = transform.structured.match ops{["linalg.unpack"]} in %module : (!transform.any_op) -> !transform.any_op
+ // 1.1 Tile the linalg.pack Op so that we can decompose it into e.g. tensor.pad
+ // and other lower-level Ops (see step 2.1)
%tiled_pack_op_p, %loops_pack:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+ // 1.2 Tile the linalg.unpack Op so that we can decompose it into e.g. tensor.pad
+ // and other lower-level Ops (see step 2.2)
%tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [4, 4]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ // 2.1. Decompose tiled PackOp into lower-level Ops + simplify
%func_op_pack = transform.get_parent_op %tiled_pack_op_p <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op_pack {
transform.apply_patterns.linalg.decompose_pack_unpack
@@ -241,6 +352,7 @@ module @transforms attributes { transform.with_named_sequence } {
transform.apply_patterns.canonicalization
} : !transform.op<"func.func">
+ // 2.2. Decompose tiled UnpackOp into lower-level Ops + simplify
%func_op_unpack = transform.get_parent_op %tiled_unpack_op_p <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op_unpack {
transform.apply_patterns.linalg.decompose_pack_unpack
@@ -252,6 +364,9 @@ module @transforms attributes { transform.with_named_sequence } {
transform.yield
}
+ //==========================================================================
+ // TD MATCHERS (helper hooks)
+ //==========================================================================
transform.named_sequence @match_mmt4d(
%entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
transform.match.operation_name %entry ["linalg.mmt4d"] : !transform.any_op
@@ -259,4 +374,7 @@ module @transforms attributes { transform.with_named_sequence } {
}
}
+//===----------------------------------------------------------------------===//
+// Function signatures
+//===----------------------------------------------------------------------===//
func.func private @printMemrefF32(%ptr : tensor<*xf32>)
>From b18048bd645b30efc80555e59dd092513c6a3755 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 29 Sep 2026 00:15:10 +0200
Subject: [PATCH 3/5] From `matmul-f32` to `pack-unpack-mmt4d-f32`
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../CPU/ArmNeon/{matmul-f32.mlir => pack-unpack-mmt4d-f32.mlir} | 0
1 file changed, 0 insertions(+), 0 deletions(-)
rename mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/{matmul-f32.mlir => pack-unpack-mmt4d-f32.mlir} (100%)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
similarity index 100%
rename from mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
rename to mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
>From b80c1b05680b7106b3d85d64384c9df467097de6 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 30 Sep 2026 23:03:27 +0200
Subject: [PATCH 4/5] Address Andrzej comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Linalg/CPU/ArmNeon/matmul-f32.mlir | 113 +++++++++++++++++
.../CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir | 114 ++++--------------
2 files changed, 138 insertions(+), 89 deletions(-)
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
new file mode 100644
index 0000000000000..c2d4f52b7678f
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
@@ -0,0 +1,113 @@
+// REQUIRES: target={{(aarch64|arm64).*}}
+
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline \
+// DEFINE: -convert-bufferization-to-memref -cse -canonicalize -convert-vector-to-scf \
+// DEFINE: -convert-vector-to-llvm="enable-arm-neon" -test-lower-to-llvm \
+// DEFINE: -o %t
+
+// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+neon" \
+// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
+
+// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+/// HIGH-LEVEL OVERVIEW
+///
+/// End-to-end test for `linalg.matmul` (f32, accumulating to f32) on dynamic
+/// shapes, tiled and vectorized for NEON (implemented in @matmul). Unlike the
+/// static shapes in pack-unpack-mmt4d-f32.mlir, the tile sizes here don't
+/// evenly divide the matrix dimensions, so this is also the variant that
+/// exercises masked vectorization (masks computed from runtime dimensions).
+///
+/// NOTES ON IMPLEMENTATION
+/// 1. NEON has no scalable vectors: N0 = 4 matches a full 128-bit register.
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// @main
+//===----------------------------------------------------------------------===//
+func.func @main() {
+ %A = arith.constant dense<[
+ [1.0, 3.0, 5.0],
+ [2.0, 4.0, 6.0],
+ [3.0, 5.0, 7.0],
+ [4.0, 6.0, 8.0],
+ [5.0, 7.0, 9.0]
+ ]> : tensor<5x3xf32>
+ %B = arith.constant dense<[
+ [1.0, 4.0, 7.0, 10.0, 13.0, 16.0, 19.0, 22.0, 25.0, 28.0, 31.0, 34.0, 37.0, 40.0, 43.0],
+ [2.0, 5.0, 8.0, 11.0, 14.0, 17.0, 20.0, 23.0, 26.0, 29.0, 32.0, 35.0, 38.0, 41.0, 44.0],
+ [3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0, 27.0, 30.0, 33.0, 36.0, 39.0, 42.0, 45.0]
+ ]> : tensor<3x15xf32>
+ %C = arith.constant dense<[
+ [1.0, 6.0, 11.0, 16.0, 21.0, 26.0, 31.0, 36.0, 41.0, 46.0, 51.0, 56.0, 61.0, 66.0, 71.0],
+ [2.0, 7.0, 12.0, 17.0, 22.0, 27.0, 32.0, 37.0, 42.0, 47.0, 52.0, 57.0, 62.0, 67.0, 72.0],
+ [3.0, 8.0, 13.0, 18.0, 23.0, 28.0, 33.0, 38.0, 43.0, 48.0, 53.0, 58.0, 63.0, 68.0, 73.0],
+ [4.0, 9.0, 14.0, 19.0, 24.0, 29.0, 34.0, 39.0, 44.0, 49.0, 54.0, 59.0, 64.0, 69.0, 74.0],
+ [5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0, 50.0, 55.0, 60.0, 65.0, 70.0, 75.0]
+ ]> : tensor<5x15xf32>
+
+ // CHECK: Unranked Memref
+ // CHECK: [23, 55, 87, 119, 151, 183, 215, 247, 279, 311, 343, 375, 407, 439, 471]
+ // CHECK: [30, 71, 112, 153, 194, 235, 276, 317, 358, 399, 440, 481, 522, 563, 604]
+ // CHECK: [37, 87, 137, 187, 237, 287, 337, 387, 437, 487, 537, 587, 637, 687, 737]
+ // CHECK: [44, 103, 162, 221, 280, 339, 398, 457, 516, 575, 634, 693, 752, 811, 870]
+ // CHECK: [51, 119, 187, 255, 323, 391, 459, 527, 595, 663, 731, 799, 867, 935, 1003]
+ %A_dyn = tensor.cast %A : tensor<5x3xf32> to tensor<?x?xf32>
+ %B_dyn = tensor.cast %B : tensor<3x15xf32> to tensor<?x?xf32>
+ %C_dyn = tensor.cast %C : tensor<5x15xf32> to tensor<?x?xf32>
+ %C_matmul = func.call @matmul(%A_dyn, %B_dyn, %C_dyn) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+ %C_matmul_cast = tensor.cast %C_matmul : tensor<?x?xf32> to tensor<*xf32>
+ vector.print str "RESULT FROM linalg.matmul:\n"
+ call @printMemrefF32(%C_matmul_cast) : (tensor<*xf32>) -> ()
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @matmul
+//
+// Implements matrix-multiplication via linalg.matmul on dynamic shapes.
+//===----------------------------------------------------------------------===//
+func.func private @matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %C_matmul : tensor<?x?xf32>
+}
+
+//===----------------------------------------------------------------------===//
+// TD Sequence
+//===----------------------------------------------------------------------===//
+module @transforms attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+ %matmul = transform.structured.match ops{["linalg.matmul"]} in %module
+ : (!transform.any_op) -> !transform.any_op
+ %func = transform.get_parent_op %matmul <isolated_from_above> : (!transform.any_op) -> !transform.op<"func.func">
+
+ %tiled_matmul, %loops:3 = transform.structured.tile_using_for %matmul tile_sizes [2, 4, 1]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+
+ transform.structured.vectorize %tiled_matmul vector_sizes [2, 4, 1] : !transform.any_op
+
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.reduction_to_contract
+ transform.apply_patterns.vector.transfer_permutation_patterns
+ transform.apply_patterns.vector.lower_masked_transfers
+ transform.apply_patterns.vector.sink_ops
+ } : !transform.op<"func.func">
+
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
+ transform.apply_patterns.vector.lower_outerproduct
+ } : !transform.op<"func.func">
+
+ transform.yield
+ }
+}
+
+//===----------------------------------------------------------------------===//
+// Function signatures
+//===----------------------------------------------------------------------===//
+func.func private @printMemrefF32(%ptr : tensor<*xf32>)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
index 84e2701b6f56e..0edc7a4ec61f1 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
@@ -2,8 +2,7 @@
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
-// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline \
-// DEFINE: -convert-bufferization-to-memref -cse -canonicalize -convert-vector-to-scf \
+// DEFINE: -cse -canonicalize -convert-vector-to-scf \
// DEFINE: -convert-vector-to-llvm="enable-arm-neon" -test-lower-to-llvm \
// DEFINE: -o %t
@@ -28,21 +27,23 @@
///
/// linalg.matmul(A, B, C)
///
-/// (implemented in both @matmul, static shapes, and @matmul_dynamic, dynamic
-/// shapes).
+/// (implemented in @matmul).
///
/// NOTES ON IMPLEMENTATION
-/// 1. NEON has no f32 matrix-multiply instruction (`fmmla` is SVE-only,
-/// `bfmmla` takes bf16 operands). The inner tile is therefore
-/// (M0, N0, K0) = (4, 4, 1): every vector.contract is a rank-1 update,
-/// lowered to outer products, i.e., one `fmla` per row of the 4x4
-/// accumulator tile.
+/// 1. @matmul is only lowered via `-test-lower-to-llvm`, with no tiling or
+/// vectorization -- it's there to provide a reference check, not to
+/// demonstrate a lowering path.
///
-/// 2. Packing is still worthwhile without a dedicated instruction: it gives
+/// 2. NEON has no instruction that multiplies two f32 matrices in one go, so
+/// the inner tile is (M0, N0, K0) = (4, 4, 1): every vector.contract is a
+/// rank-1 update, lowered to outer products, i.e. one `fmla` per row of
+/// the 4x4 accumulator tile.
+///
+/// 3. Packing is still worthwhile without a dedicated instruction: it gives
/// the tiles a statically-known shape (no masking in the mmt4d loop) and
/// makes them contiguous in memory.
///
-/// 3. The MMT4D and Pack/Unpack Ops are kept in separate functions to isolate
+/// 4. The MMT4D and Pack/Unpack Ops are kept in separate functions to isolate
/// the corresponding lowering and lowering configs.
//===----------------------------------------------------------------------===//
@@ -50,8 +51,8 @@
// @main
//
// The main entry point that computes matrix multiplication via
-// linalg.mmt4d and linalg.matmul (static and dynamic shapes). The output
-// should be independent of the Linalg Op used.
+// linalg.mmt4d and linalg.matmul. The output should be independent of the
+// Linalg Op used.
//===----------------------------------------------------------------------===//
func.func @main() {
%A = arith.constant dense<[
@@ -86,7 +87,7 @@ func.func @main() {
vector.print str "RESULT FROM linalg.mmt4d:\n"
call @printMemrefF32(%C_mmt4d_cast) : (tensor<*xf32>) -> ()
- // VARIANT: Matrix multiplication via linalg.matmul (cross-check)
+ // VARIANT: Matrix multiplication via linalg.matmul (reference check)
// CHECK: Unranked Memref
// CHECK: [23, 55, 87, 119, 151, 183, 215, 247, 279, 311, 343, 375, 407, 439, 471]
// CHECK: [30, 71, 112, 153, 194, 235, 276, 317, 358, 399, 440, 481, 522, 563, 604]
@@ -98,27 +99,15 @@ func.func @main() {
vector.print str "RESULT FROM linalg.matmul:\n"
call @printMemrefF32(%C_matmul_cast) : (tensor<*xf32>) -> ()
- // VARIANT: Matrix multiplication via linalg.matmul on dynamic shapes
- // CHECK: Unranked Memref
- // CHECK: [23, 55, 87, 119, 151, 183, 215, 247, 279, 311, 343, 375, 407, 439, 471]
- // CHECK: [30, 71, 112, 153, 194, 235, 276, 317, 358, 399, 440, 481, 522, 563, 604]
- // CHECK: [37, 87, 137, 187, 237, 287, 337, 387, 437, 487, 537, 587, 637, 687, 737]
- // CHECK: [44, 103, 162, 221, 280, 339, 398, 457, 516, 575, 634, 693, 752, 811, 870]
- // CHECK: [51, 119, 187, 255, 323, 391, 459, 527, 595, 663, 731, 799, 867, 935, 1003]
- %A_dyn = tensor.cast %A : tensor<5x3xf32> to tensor<?x?xf32>
- %B_dyn = tensor.cast %B : tensor<3x15xf32> to tensor<?x?xf32>
- %C_dyn = tensor.cast %C : tensor<5x15xf32> to tensor<?x?xf32>
- %C_matmul_dyn = func.call @matmul_dynamic(%A_dyn, %B_dyn, %C_dyn) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
- %C_matmul_dyn_cast = tensor.cast %C_matmul_dyn : tensor<?x?xf32> to tensor<*xf32>
- vector.print str "RESULT FROM linalg.matmul (dynamic):\n"
- call @printMemrefF32(%C_matmul_dyn_cast) : (tensor<*xf32>) -> ()
return
}
//===----------------------------------------------------------------------===//
// @matmul
//
-// Implements matrix-multiplication via linalg.matmul, static shapes.
+// Implements matrix-multiplication via linalg.matmul. Lowered only via
+// `-test-lower-to-llvm` (no tiling, no vectorization): this is a reference
+// check for @matmul_via_mmt4d, not a lowering path under test.
//===----------------------------------------------------------------------===//
func.func private @matmul(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %C: tensor<5x15xf32>) -> tensor<5x15xf32> {
%C_matmul = linalg.matmul ins(%A, %B: tensor<5x3xf32>, tensor<3x15xf32>)
@@ -126,19 +115,6 @@ func.func private @matmul(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %C: tensor<
return %C_matmul : tensor<5x15xf32>
}
-//===----------------------------------------------------------------------===//
-// @matmul_dynamic
-//
-// Implements matrix-multiplication via linalg.matmul, dynamic shapes. Unlike
-// the other variants, this is the one that exercises masked vectorization
-// (masks computed from runtime dimensions).
-//===----------------------------------------------------------------------===//
-func.func private @matmul_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
- %C_matmul = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)
- outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>
- return %C_matmul : tensor<?x?xf32>
-}
-
//===----------------------------------------------------------------------===//
// @pack_lhs
//
@@ -233,53 +209,7 @@ func.func private @matmul_via_mmt4d(%A: tensor<5x3xf32>, %B: tensor<3x15xf32>, %
// TD Sequence
//===----------------------------------------------------------------------===//
module @transforms attributes { transform.with_named_sequence } {
- //===--------------------------------------------------------------------===//
- // @tile_and_vectorize_matmul
- //
- // Tile and vectorize a plain `linalg.matmul` (shared by @matmul and
- // @matmul_dynamic).
- //===--------------------------------------------------------------------===//
- transform.named_sequence @tile_and_vectorize_matmul(%func
- : !transform.op<"func.func"> {transform.readonly}) {
-
- %matmul = transform.structured.match ops{["linalg.matmul"]} in %func
- : (!transform.op<"func.func">) -> !transform.any_op
-
- // NEON has no scalable vectors: N = 4 matches a full 128-bit register.
- %tiled_matmul, %loops:3 = transform.structured.tile_using_for %matmul tile_sizes [2, 4, 1]
- : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
-
- transform.structured.vectorize %tiled_matmul vector_sizes [2, 4, 1] : !transform.any_op
-
- transform.apply_patterns to %func {
- transform.apply_patterns.vector.reduction_to_contract
- transform.apply_patterns.vector.transfer_permutation_patterns
- transform.apply_patterns.vector.lower_masked_transfers
- transform.apply_patterns.vector.sink_ops
- } : !transform.op<"func.func">
-
- transform.apply_patterns to %func {
- transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
- transform.apply_patterns.vector.lower_outerproduct
- } : !transform.op<"func.func">
-
- transform.yield
- }
-
- transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
- //==========================================================================
- // HANDLE PLAIN MATMUL (static + dynamic)
- //==========================================================================
- %matmul_func = transform.structured.match ops{["func.func"]} attributes{sym_name = "matmul"} in %module
- : (!transform.any_op) -> !transform.op<"func.func">
- transform.include @tile_and_vectorize_matmul failures(propagate)
- (%matmul_func) : (!transform.op<"func.func">) -> ()
-
- %matmul_dynamic_func = transform.structured.match ops{["func.func"]} attributes{sym_name = "matmul_dynamic"} in %module
- : (!transform.any_op) -> !transform.op<"func.func">
- transform.include @tile_and_vectorize_matmul failures(propagate)
- (%matmul_dynamic_func) : (!transform.op<"func.func">) -> ()
-
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consumed}) {
//==========================================================================
// HANDLE MMT4D
//==========================================================================
@@ -361,6 +291,12 @@ module @transforms attributes { transform.with_named_sequence } {
transform.apply_patterns.tensor.fold_tensor_subset_ops
transform.apply_patterns.canonicalization
} : !transform.op<"func.func">
+
+ //==========================================================================
+ // BUFFERIZATION
+ //==========================================================================
+ %bufferize = transform.bufferization.one_shot_bufferize %module
+ <bufferize_function_boundaries = true> : (!transform.any_op) -> !transform.any_op
transform.yield
}
>From a5fe58ea43afad86ce309812360426bd236b92db Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Sun, 4 Oct 2026 16:58:58 +0200
Subject: [PATCH 5/5] Address comments (pt1)
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Linalg/CPU/ArmNeon/matmul-f32.mlir | 3 ---
.../CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir | 19 ++++++-------------
2 files changed, 6 insertions(+), 16 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
index c2d4f52b7678f..ddf01760df37c 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
@@ -20,9 +20,6 @@
/// static shapes in pack-unpack-mmt4d-f32.mlir, the tile sizes here don't
/// evenly divide the matrix dimensions, so this is also the variant that
/// exercises masked vectorization (masks computed from runtime dimensions).
-///
-/// NOTES ON IMPLEMENTATION
-/// 1. NEON has no scalable vectors: N0 = 4 matches a full 128-bit register.
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
index 0edc7a4ec61f1..1cc6c29dbee5a 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/pack-unpack-mmt4d-f32.mlir
@@ -2,8 +2,7 @@
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
-// DEFINE: -cse -canonicalize -convert-vector-to-scf \
-// DEFINE: -convert-vector-to-llvm="enable-arm-neon" -test-lower-to-llvm \
+// DEFINE: -cse -canonicalize -test-lower-to-llvm \
// DEFINE: -o %t
// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+neon" \
@@ -30,20 +29,14 @@
/// (implemented in @matmul).
///
/// NOTES ON IMPLEMENTATION
-/// 1. @matmul is only lowered via `-test-lower-to-llvm`, with no tiling or
-/// vectorization -- it's there to provide a reference check, not to
+/// 1. linalg.matmul is only lowered via `-test-lower-to-llvm`, with no tiling or
+/// vectorization - it's there to provide a reference check, not to
/// demonstrate a lowering path.
///
-/// 2. NEON has no instruction that multiplies two f32 matrices in one go, so
-/// the inner tile is (M0, N0, K0) = (4, 4, 1): every vector.contract is a
-/// rank-1 update, lowered to outer products, i.e. one `fmla` per row of
-/// the 4x4 accumulator tile.
+/// 2. Packing gives the tiles a statically-known shape (no masking in the
+/// mmt4d loop) and makes them contiguous in memory.
///
-/// 3. Packing is still worthwhile without a dedicated instruction: it gives
-/// the tiles a statically-known shape (no masking in the mmt4d loop) and
-/// makes them contiguous in memory.
-///
-/// 4. The MMT4D and Pack/Unpack Ops are kept in separate functions to isolate
+/// 3. The MMT4D and Pack/Unpack Ops are kept in separate functions to isolate
/// the corresponding lowering and lowering configs.
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list