[Mlir-commits] [mlir] [mlir][ArmNeon] Add linalg.matmul e2e tests (PR #212809)
Federico Bruzzone
llvmlistbot at llvm.org
Wed Aug 5 01:12:08 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/212809
>From df3125b21378215639f0ac440383e6d1baab6e03 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 29 Jul 2026 17:53:52 +0200
Subject: [PATCH 1/4] [mlir][ArmNeon] Add linalg.matmul e2e tests
Add end-to-end integration tests for `linalg.matmul` on Arm NEON,
mirroring the existing ArmSVE/ArmSME `Linalg/CPU` tests (previously
no NEON coverage existed here).
- `matmul.mlir`: plain f32 case, generic vectorize + outerproduct lowering.
- `matmul-i8mm.mlir`: i8->i32 case exercising FEAT_I8MM (`smmla`) via
`apply_patterns.arm_neon.vector_contract_to_i8mm`, using a transposed-RHS
`indexing_maps` on `linalg.matmul` (mirrors ArmSME's transpose-A trick)
so vectorization produces the contract shape the pattern expects.
`linalg.mmt4d` NEON coverage and ArmSVE/ArmSME test-format unification
are follow-ups.
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Linalg/CPU/ArmNeon/matmul-i8mm.mlir | 139 ++++++++++++++++++
.../Dialect/Linalg/CPU/ArmNeon/matmul.mlir | 96 ++++++++++++
2 files changed, 235 insertions(+)
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir
new file mode 100644
index 0000000000000..fba6ab7220dc6
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir
@@ -0,0 +1,139 @@
+// REQUIRES: arm-emulator
+
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline -cse -canonicalize -convert-vector-to-scf \
+// DEFINE: -convert-vector-to-llvm="enable-arm-neon enable-arm-i8mm" -test-lower-to-llvm \
+// DEFINE: -o %t
+
+// DEFINE: %{run} = %mcr_aarch64_cmd %t -e main -entry-point-result=void --march=aarch64 --mattr="+neon,+i8mm" \
+// 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
+
+//===----------------------------------------------------------------------===//
+// Tiles, vectorizes and lowers a `linalg.matmul` down to Arm's FEAT_I8MM
+// `smmla` instruction via `transform.apply_patterns.arm_neon.vector_contract_to_i8mm`
+// (LowerContractionToNeonI8MMPattern).
+//
+// That pattern expects a `vector.contract` with LHS vector<MxKxi8>, RHS
+// vector<NxKxi8> (RHS read "N-major", i.e. logically transposed relative to
+// what a plain `linalg.matmul` produces), and ACC/OUT vector<MxNxi32> -- see
+// #packed_maps in Vector/CPU/ArmNeon/vector-contract-i8mm.mlir. To get there,
+// this test gives `linalg.matmul` an explicit `indexing_maps` attribute that
+// reads the second operand "N-major" (== `MatmulTransposeBOp`'s default
+// maps, see LinalgOps.cpp), feeding it the second input pre-transposed
+// (NxK instead of KxN). This mirrors how ArmSME/matmul-transpose-a.mlir
+// transposes the *LHS* instead, for SME's own hardware constraints.
+//===----------------------------------------------------------------------===//
+
+// CHECK-IR-LABEL: llvm.func @main
+// CHECK-IR-COUNT-4: arm_neon.intr.smmla
+func.func @main() {
+ // A: MxK = 4x8.
+ %A = arith.constant dense<[
+ [-35, -27, -36, -31, 23, -34, -8, -33],
+ [-20, 17, -32, -47, 37, 22, -7, -21],
+ [ -7, -35, 20, -4, 39, 46, -23, 40],
+ [ 40, 27, 37, 43, 38, -6, 37, 49]
+ ]> : tensor<4x8xi8>
+
+ // B, transposed: NxK = 4x8 (row n holds column n of the logical KxN RHS).
+ %Bt = arith.constant dense<[
+ [-17, -50, -1, 48, -13, 22, 39, 33],
+ [-35, -24, 37, -32, 33, 30, -11, -17],
+ [-28, 31, 3, -44, -15, -27, 22, 35],
+ [-23, 39, 48, 26, -23, 32, -39, -38]
+ ]> : tensor<4x8xi8>
+
+ // C: MxN = 4x4, non-zero to also exercise the "+ ACC" part of `smmla`.
+ %C = arith.constant dense<[
+ [-44, 20, 44, -46],
+ [ -8, 25, -34, 26],
+ [-20, -36, -3, 39],
+ [-48, -31, -25, -21]
+ ]> : tensor<4x4xi32>
+
+ %A_dyn = tensor.cast %A : tensor<4x8xi8> to tensor<?x?xi8>
+ %Bt_dyn = tensor.cast %Bt : tensor<4x8xi8> to tensor<?x?xi8>
+ %C_dyn = tensor.cast %C : tensor<4x4xi32> to tensor<?x?xi32>
+
+ %res = linalg.matmul
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>
+ ]
+ ins(%A_dyn, %Bt_dyn : tensor<?x?xi8>, tensor<?x?xi8>)
+ outs(%C_dyn : tensor<?x?xi32>) -> tensor<?x?xi32>
+
+ // 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 = [4, 4] strides = [4, 1] data =
+ // CHECK-NEXT: [-1999, 1941, 685, -2879]
+ // CHECK-NEXT: [-3705, 2952, 987, -685]
+ // CHECK-NEXT: [2565, 4157, -1589, -357]
+ // CHECK-NEXT: [2383, -2252, 32, -1365]
+ %xf = tensor.cast %res : tensor<?x?xi32> to tensor<*xi32>
+ call @printMemrefI32(%xf) : (tensor<*xi32>) -> ()
+
+ // CHECK-NEXT: NEON: END OF TEST OUTPUT
+ vector.print str "NEON: END OF TEST OUTPUT\n"
+
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ // Tile, vectorize, then lower the `vector.contract` straight to FEAT_I8MM ops.
+ transform.named_sequence @tile_and_vectorize_matmul(%func
+ : !transform.op<"func.func"> {transform.readonly}) {
+
+ // Step 0: Get a handle to the matmul op, if any.
+ %matmul = transform.structured.match ops{["linalg.matmul"]} in %func
+ : (!transform.op<"func.func">) -> !transform.any_op
+
+ // Step 1: Tile to the FEAT_I8MM tile shape (M=N=4, K=8). This is the
+ // whole problem size here, so tiling produces a single tile, no tail.
+ %tiled_matmul, %loops:3 = transform.structured.tile_using_for %matmul
+ tile_sizes [4, 4, 8]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+
+ // Step 2: Vectorize directly to a named `vector.contract`.
+ transform.structured.vectorize %tiled_matmul vector_sizes [4, 4, 8]
+ {create_named_contraction} : !transform.any_op
+
+ // Step 3: M, N, K are static and match the tile/vector sizes, so
+ // vectorization masks are trivially full tile; clean them up.
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.transfer_permutation_patterns
+ transform.apply_patterns.vector.lower_masked_transfers
+ transform.apply_patterns.vector.sink_ops
+ } : !transform.op<"func.func">
+
+ // Step 4: Lower `vector.contract` straight to FEAT_I8MM ops,
+ // instead of the generic outerproduct lowering, which would
+ // exercise `smmla`.
+ transform.apply_patterns to %func {
+ transform.apply_patterns.arm_neon.vector_contract_to_i8mm
+ } : !transform.op<"func.func">
+
+ 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"> {
+ ^bb0(%func : !transform.op<"func.func">):
+ transform.include @tile_and_vectorize_matmul failures(propagate)
+ (%func) : (!transform.op<"func.func">) -> ()
+ }
+ transform.yield
+ }
+}
+
+func.func private @printMemrefI32(%ptr : tensor<*xi32>)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir
new file mode 100644
index 0000000000000..82ba17e43c969
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir
@@ -0,0 +1,96 @@
+// REQUIRES: arm-emulator
+
+// 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
+
+// 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
+
+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"
+
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ // Tile and vectorize the matmul.
+ transform.named_sequence @tile_and_vectorize_matmul(%func
+ : !transform.op<"func.func"> {transform.readonly}) {
+
+ // Step 0: Get a handle to the matmul op, if any.
+ %matmul = transform.structured.match ops{["linalg.matmul"]} in %func
+ : (!transform.op<"func.func">) -> !transform.any_op
+
+ // Step 1: Tile. NEON has no scalable vectors, so sizes are static:
+ // N = 4 matches a full 128-bit NEON register of f32 elements.
+ %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)
+
+ // Step 2: Vectorize directly to a named `vector.contract`.
+ transform.structured.vectorize %tiled_matmul vector_sizes [2, 4, 1] : !transform.any_op
+
+ // Step 3: Lower `vector.multi_reduction` to `vector.contract` (+ some helpful patterns)
+ 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">
+
+ // Step 4: Lower `vector.contract` to `vector.fma`.
+ 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
+ }
+
+ // 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">) -> ()
+ }
+ transform.yield
+ }
+}
+
+func.func private @printMemrefF32(%ptr : tensor<*xf32>)
>From cac61b2e38854eed867f02050e6b74fd8e80d34d Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 3 Aug 2026 17:03:53 +0200
Subject: [PATCH 2/4] Address comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmNeon/{matmul.mlir => matmul-f32.mlir} | 9 +--
.../{matmul-i8mm.mlir => matmul-i8.mlir} | 68 +++++++------------
2 files changed, 26 insertions(+), 51 deletions(-)
rename mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/{matmul.mlir => matmul-f32.mlir} (89%)
rename mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/{matmul-i8mm.mlir => matmul-i8.mlir} (59%)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
similarity index 89%
rename from mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir
rename to mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
index 82ba17e43c969..3a55d2a890eef 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-f32.mlir
@@ -1,4 +1,4 @@
-// REQUIRES: arm-emulator
+// REQUIRES: target={{(aarch64|arm64).*}}
// RUN: mlir-opt %s \
// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \
@@ -50,19 +50,15 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @tile_and_vectorize_matmul(%func
: !transform.op<"func.func"> {transform.readonly}) {
- // Step 0: Get a handle to the matmul op, if any.
%matmul = transform.structured.match ops{["linalg.matmul"]} in %func
: (!transform.op<"func.func">) -> !transform.any_op
- // Step 1: Tile. NEON has no scalable vectors, so sizes are static:
- // N = 4 matches a full 128-bit NEON register of f32 elements.
+ // 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)
- // Step 2: Vectorize directly to a named `vector.contract`.
transform.structured.vectorize %tiled_matmul vector_sizes [2, 4, 1] : !transform.any_op
- // Step 3: Lower `vector.multi_reduction` to `vector.contract` (+ some helpful patterns)
transform.apply_patterns to %func {
transform.apply_patterns.vector.reduction_to_contract
transform.apply_patterns.vector.transfer_permutation_patterns
@@ -70,7 +66,6 @@ module attributes {transform.with_named_sequence} {
transform.apply_patterns.vector.sink_ops
} : !transform.op<"func.func">
- // Step 4: Lower `vector.contract` to `vector.fma`.
transform.apply_patterns to %func {
transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
transform.apply_patterns.vector.lower_outerproduct
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
similarity index 59%
rename from mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir
rename to mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
index fba6ab7220dc6..8905c6462fb19 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8mm.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
@@ -11,21 +11,9 @@
// RUN: rm -f %t && %{compile} && FileCheck %s --input-file=%t -check-prefix CHECK-IR && %{run} | FileCheck %s
-//===----------------------------------------------------------------------===//
-// Tiles, vectorizes and lowers a `linalg.matmul` down to Arm's FEAT_I8MM
-// `smmla` instruction via `transform.apply_patterns.arm_neon.vector_contract_to_i8mm`
-// (LowerContractionToNeonI8MMPattern).
-//
-// That pattern expects a `vector.contract` with LHS vector<MxKxi8>, RHS
-// vector<NxKxi8> (RHS read "N-major", i.e. logically transposed relative to
-// what a plain `linalg.matmul` produces), and ACC/OUT vector<MxNxi32> -- see
-// #packed_maps in Vector/CPU/ArmNeon/vector-contract-i8mm.mlir. To get there,
-// this test gives `linalg.matmul` an explicit `indexing_maps` attribute that
-// reads the second operand "N-major" (== `MatmulTransposeBOp`'s default
-// maps, see LinalgOps.cpp), feeding it the second input pre-transposed
-// (NxK instead of KxN). This mirrors how ArmSME/matmul-transpose-a.mlir
-// transposes the *LHS* instead, for SME's own hardware constraints.
-//===----------------------------------------------------------------------===//
+// Lowers a vanilla `linalg.matmul` down to Arm's FEAT_I8MM `smmla`.
+// `LowerContractionToNeonI8MMPattern` expects the RHS transposed (N-major);
+// `transform.structured.transpose_matmul <rhs>` gets us there from a plain matmul.
// CHECK-IR-LABEL: llvm.func @main
// CHECK-IR-COUNT-4: arm_neon.intr.smmla
@@ -38,13 +26,17 @@ func.func @main() {
[ 40, 27, 37, 43, 38, -6, 37, 49]
]> : tensor<4x8xi8>
- // B, transposed: NxK = 4x8 (row n holds column n of the logical KxN RHS).
- %Bt = arith.constant dense<[
- [-17, -50, -1, 48, -13, 22, 39, 33],
- [-35, -24, 37, -32, 33, 30, -11, -17],
- [-28, 31, 3, -44, -15, -27, 22, 35],
- [-23, 39, 48, 26, -23, 32, -39, -38]
- ]> : tensor<4x8xi8>
+ // B: KxN = 8x4.
+ %B = arith.constant dense<[
+ [-17, -35, -28, -23],
+ [-50, -24, 31, 39],
+ [ -1, 37, 3, 48],
+ [ 48, -32, -44, 26],
+ [-13, 33, -15, -23],
+ [ 22, 30, -27, 32],
+ [ 39, -11, 22, -39],
+ [ 33, -17, 35, -38]
+ ]> : tensor<8x4xi8>
// C: MxN = 4x4, non-zero to also exercise the "+ ACC" part of `smmla`.
%C = arith.constant dense<[
@@ -54,18 +46,9 @@ func.func @main() {
[-48, -31, -25, -21]
]> : tensor<4x4xi32>
- %A_dyn = tensor.cast %A : tensor<4x8xi8> to tensor<?x?xi8>
- %Bt_dyn = tensor.cast %Bt : tensor<4x8xi8> to tensor<?x?xi8>
- %C_dyn = tensor.cast %C : tensor<4x4xi32> to tensor<?x?xi32>
-
%res = linalg.matmul
- indexing_maps = [
- affine_map<(d0, d1, d2) -> (d0, d2)>,
- affine_map<(d0, d1, d2) -> (d1, d2)>,
- affine_map<(d0, d1, d2) -> (d0, d1)>
- ]
- ins(%A_dyn, %Bt_dyn : tensor<?x?xi8>, tensor<?x?xi8>)
- outs(%C_dyn : tensor<?x?xi32>) -> tensor<?x?xi32>
+ ins(%A, %B : tensor<4x8xi8>, tensor<8x4xi8>)
+ outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
// Print and verify the output
// CHECK-LABEL: NEON: START OF TEST OUTPUT
@@ -76,7 +59,7 @@ func.func @main() {
// CHECK-NEXT: [-3705, 2952, 987, -685]
// CHECK-NEXT: [2565, 4157, -1589, -357]
// CHECK-NEXT: [2383, -2252, 32, -1365]
- %xf = tensor.cast %res : tensor<?x?xi32> to tensor<*xi32>
+ %xf = tensor.cast %res : tensor<4x4xi32> to tensor<*xi32>
call @printMemrefI32(%xf) : (tensor<*xi32>) -> ()
// CHECK-NEXT: NEON: END OF TEST OUTPUT
@@ -90,31 +73,28 @@ module attributes {transform.with_named_sequence} {
transform.named_sequence @tile_and_vectorize_matmul(%func
: !transform.op<"func.func"> {transform.readonly}) {
- // Step 0: Get a handle to the matmul op, if any.
%matmul = transform.structured.match ops{["linalg.matmul"]} in %func
: (!transform.op<"func.func">) -> !transform.any_op
- // Step 1: Tile to the FEAT_I8MM tile shape (M=N=4, K=8). This is the
- // whole problem size here, so tiling produces a single tile, no tail.
- %tiled_matmul, %loops:3 = transform.structured.tile_using_for %matmul
+ %transposed_matmul = transform.structured.transpose_matmul %matmul <rhs>
+ : (!transform.any_op) -> (!transform.any_op)
+
+ // M=N=4, K=8: FEAT_I8MM's native tile shape.
+ %tiled_matmul, %loops:3 = transform.structured.tile_using_for %transposed_matmul
tile_sizes [4, 4, 8]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
- // Step 2: Vectorize directly to a named `vector.contract`.
transform.structured.vectorize %tiled_matmul vector_sizes [4, 4, 8]
{create_named_contraction} : !transform.any_op
- // Step 3: M, N, K are static and match the tile/vector sizes, so
- // vectorization masks are trivially full tile; clean them up.
transform.apply_patterns to %func {
transform.apply_patterns.vector.transfer_permutation_patterns
transform.apply_patterns.vector.lower_masked_transfers
transform.apply_patterns.vector.sink_ops
} : !transform.op<"func.func">
- // Step 4: Lower `vector.contract` straight to FEAT_I8MM ops,
- // instead of the generic outerproduct lowering, which would
- // exercise `smmla`.
+ // Lower straight to FEAT_I8MM ops instead of the generic outerproduct
+ // path, which would never emit `smmla`.
transform.apply_patterns to %func {
transform.apply_patterns.arm_neon.vector_contract_to_i8mm
} : !transform.op<"func.func">
>From cb0863d1f04e8ec0e89fb4166f2d6fe227bb0c6b Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 5 Aug 2026 10:00:18 +0200
Subject: [PATCH 3/4] Address comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir | 277 +++++++++++++-----
1 file changed, 197 insertions(+), 80 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
index 8905c6462fb19..94c73f639fd3b 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
@@ -2,7 +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 -cse -canonicalize -convert-vector-to-scf \
+// DEFINE: -cse -canonicalize -convert-vector-to-scf \
// DEFINE: -convert-vector-to-llvm="enable-arm-neon enable-arm-i8mm" -test-lower-to-llvm \
// DEFINE: -o %t
@@ -11,108 +11,225 @@
// RUN: rm -f %t && %{compile} && FileCheck %s --input-file=%t -check-prefix CHECK-IR && %{run} | FileCheck %s
-// Lowers a vanilla `linalg.matmul` down to Arm's FEAT_I8MM `smmla`.
-// `LowerContractionToNeonI8MMPattern` expects the RHS transposed (N-major);
-// `transform.structured.transpose_matmul <rhs>` gets us there from a plain matmul.
+// End-to-end test for `linalg.matmul` on i8 operands accumulating to i32,
+// lowered via `linalg.pack -> linalg.mmt4d -> linalg.unpack` down to Arm's
+// FEAT_I8MM `smmla`. Packing gives the inner tiles a statically-known shape,
+// so the vectorized `vector.contract` never needs masking, and
+// `linalg.mmt4d`'s RHS is already N-major (that's the "t" in "mmt4d"),
+// exactly what `LowerContractionToNeonI8MMPattern` expects -- no
+// transpose_matmul step needed here, unlike a plain `linalg.matmul`.
-// CHECK-IR-LABEL: llvm.func @main
-// CHECK-IR-COUNT-4: arm_neon.intr.smmla
func.func @main() {
- // A: MxK = 4x8.
- %A = arith.constant dense<[
- [-35, -27, -36, -31, 23, -34, -8, -33],
- [-20, 17, -32, -47, 37, 22, -7, -21],
- [ -7, -35, 20, -4, 39, 46, -23, 40],
- [ 40, 27, 37, 43, 38, -6, 37, 49]
- ]> : tensor<4x8xi8>
-
- // B: KxN = 8x4.
- %B = arith.constant dense<[
- [-17, -35, -28, -23],
- [-50, -24, 31, 39],
- [ -1, 37, 3, 48],
- [ 48, -32, -44, 26],
- [-13, 33, -15, -23],
- [ 22, 30, -27, 32],
- [ 39, -11, 22, -39],
- [ 33, -17, 35, -38]
- ]> : tensor<8x4xi8>
-
- // C: MxN = 4x4, non-zero to also exercise the "+ ACC" part of `smmla`.
+ %A_empty = tensor.empty() : tensor<7x16xi8>
+ %B_empty = tensor.empty() : tensor<16x13xi8>
+
+ %c3 = arith.constant 3 : i8
+ %c4 = arith.constant 4 : i8
+ %A = linalg.fill ins(%c3 : i8) outs(%A_empty : tensor<7x16xi8>) -> tensor<7x16xi8>
+ %B = linalg.fill ins(%c4 : i8) outs(%B_empty : tensor<16x13xi8>) -> tensor<16x13xi8>
%C = arith.constant dense<[
- [-44, 20, 44, -46],
- [ -8, 25, -34, 26],
- [-20, -36, -3, 39],
- [-48, -31, -25, -21]
- ]> : tensor<4x4xi32>
-
- %res = linalg.matmul
- ins(%A, %B : tensor<4x8xi8>, tensor<8x4xi8>)
- outs(%C : tensor<4x4xi32>) -> tensor<4x4xi32>
-
- // 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 = [4, 4] strides = [4, 1] data =
- // CHECK-NEXT: [-1999, 1941, 685, -2879]
- // CHECK-NEXT: [-3705, 2952, 987, -685]
- // CHECK-NEXT: [2565, 4157, -1589, -357]
- // CHECK-NEXT: [2383, -2252, 32, -1365]
- %xf = tensor.cast %res : tensor<4x4xi32> to tensor<*xi32>
- call @printMemrefI32(%xf) : (tensor<*xi32>) -> ()
-
- // CHECK-NEXT: NEON: END OF TEST OUTPUT
- vector.print str "NEON: END OF TEST OUTPUT\n"
+ [ 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<7x13xi32>
+
+ // 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<7x16xi8>, tensor<16x13xi8>, tensor<7x13xi32>) -> tensor<7x13xi32>
+ %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<7x13xi32> to tensor<*xi32>
+ vector.print str "RESULT FROM linalg.mmt4d:\n"
+ call @printMemrefI32(%C_mmt4d_cast) : (tensor<*xi32>) -> ()
+
+ // VARIANT: Matrix multiplication via linalg.matmul (cross-check)
+ // 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<7x16xi8>, tensor<16x13xi8>, tensor<7x13xi32>) -> tensor<7x13xi32>
+ %C_matmul_cast = tensor.cast %C_matmul : tensor<7x13xi32> to tensor<*xi32>
+ vector.print str "RESULT FROM linalg.matmul:\n"
+ call @printMemrefI32(%C_matmul_cast) : (tensor<*xi32>) -> ()
return
}
-module attributes {transform.with_named_sequence} {
- // Tile, vectorize, then lower the `vector.contract` straight to FEAT_I8MM ops.
- transform.named_sequence @tile_and_vectorize_matmul(%func
- : !transform.op<"func.func"> {transform.readonly}) {
+func.func private @matmul(%A: tensor<7x16xi8>, %B: tensor<16x13xi8>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<7x16xi8>, tensor<16x13xi8>)
+ outs(%C: tensor<7x13xi32>) -> tensor<7x13xi32>
+ return %C_matmul : tensor<7x13xi32>
+}
- %matmul = transform.structured.match ops{["linalg.matmul"]} in %func
- : (!transform.op<"func.func">) -> !transform.any_op
+// LHS packed tile: M0=4, K0=8 (K0 must be a multiple of 8 for FEAT_I8MM).
+func.func private @pack_lhs(%A: tensor<7x16xi8>) -> tensor<2x2x4x8xi8> {
+ %pad = arith.constant 0 : i8
+ %A_pack_empty = tensor.empty() : tensor<2x2x4x8xi8>
+ %A_pack = linalg.pack %A
+ padding_value(%pad : i8)
+ inner_dims_pos = [0, 1]
+ inner_tiles = [4, 8]
+ into %A_pack_empty : tensor<7x16xi8> -> tensor<2x2x4x8xi8>
+ return %A_pack : tensor<2x2x4x8xi8>
+}
+
+// RHS packed tile: N0=4, K0=8.
+func.func private @pack_rhs(%B: tensor<16x13xi8>) -> tensor<4x2x4x8xi8> {
+ %pad = arith.constant 0 : i8
+ %B_pack_empty = tensor.empty() : tensor<4x2x4x8xi8>
+ %B_pack = linalg.pack %B
+ padding_value(%pad : i8)
+ outer_dims_perm = [1, 0]
+ inner_dims_pos = [1, 0]
+ inner_tiles = [4, 8]
+ into %B_pack_empty : tensor<16x13xi8> -> tensor<4x2x4x8xi8>
+ return %B_pack : tensor<4x2x4x8xi8>
+}
+
+func.func private @pack_acc(%C: tensor<7x13xi32>) -> tensor<2x4x4x4xi32> {
+ %pad = arith.constant 0 : i32
+ %C_pack_empty = tensor.empty() : tensor<2x4x4x4xi32>
+ %C_pack = linalg.pack %C
+ padding_value(%pad : i32)
+ outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1]
+ inner_tiles = [4, 4]
+ into %C_pack_empty : tensor<7x13xi32> -> tensor<2x4x4x4xi32>
+ return %C_pack : tensor<2x4x4x4xi32>
+}
+
+func.func private @unpack_acc(%C_packed: tensor<2x4x4x4xi32>) -> tensor<7x13xi32> {
+ %C_out_empty = tensor.empty() : tensor<7x13xi32>
+ %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<2x4x4x4xi32> -> tensor<7x13xi32>
+ return %C_out_unpack: tensor<7x13xi32>
+}
+
+// CHECK-IR-LABEL: llvm.func @matmul_via_mmt4d
+// CHECK-IR-COUNT-4: arm_neon.intr.smmla
+func.func private @matmul_via_mmt4d(%A: tensor<7x16xi8>, %B: tensor<16x13xi8>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {
+ %A_pack = func.call @pack_lhs(%A): (tensor<7x16xi8>) -> tensor<2x2x4x8xi8>
+ %B_pack = func.call @pack_rhs(%B): (tensor<16x13xi8>) -> tensor<4x2x4x8xi8>
+ %C_pack = func.call @pack_acc(%C): (tensor<7x13xi32>) -> tensor<2x4x4x4xi32>
- %transposed_matmul = transform.structured.transpose_matmul %matmul <rhs>
- : (!transform.any_op) -> (!transform.any_op)
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<2x2x4x8xi8>, tensor<4x2x4x8xi8>) outs(%C_pack : tensor<2x4x4x4xi32>) -> tensor<2x4x4x4xi32>
- // M=N=4, K=8: FEAT_I8MM's native tile shape.
- %tiled_matmul, %loops:3 = transform.structured.tile_using_for %transposed_matmul
- tile_sizes [4, 4, 8]
- : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<2x4x4x4xi32>) -> tensor<7x13xi32>
+ return %C_out_unpack : tensor<7x13xi32>
+}
- transform.structured.vectorize %tiled_matmul vector_sizes [4, 4, 8]
+module @transforms attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consumed}) {
+ %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 parallel dims (m, n, k, m0, n0, k0): full inner tiles, one outer
+ // iteration at a time.
+ %tiled_mmt4d_parallel, %_:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, 4, 4, 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ // Tile reduction dims: k0=8 is the full inner extent (FEAT_I8MM handles
+ // the whole 8-wide reduction in one instruction, no further split).
+ %tiled_mmt4d, %_1:2 = transform.structured.tile_using_for %tiled_mmt4d_parallel tile_sizes [0, 0, 1, 0, 0, 8]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+ // Vectorize directly to a named `vector.contract` (compact 2-operand
+ // form) instead of the generic broadcast form, since
+ // LowerContractionToNeonI8MMPattern requires LHS/RHS rank <= 2.
+ transform.structured.vectorize %tiled_mmt4d vector_sizes [1, 1, 1, 4, 4, 8]
{create_named_contraction} : !transform.any_op
- transform.apply_patterns to %func {
+ transform.apply_patterns to %mmt4d_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">
- // Lower straight to FEAT_I8MM ops instead of the generic outerproduct
- // path, which would never emit `smmla`.
- transform.apply_patterns to %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">
+
+ %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">
+
+ %bufferize = transform.bufferization.one_shot_bufferize %module
+ {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
+
+ %contract = transform.collect_matching @match_contract in %bufferize : (!transform.any_op) -> (!transform.any_op)
+ %contract_func = transform.get_parent_op %contract {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">
+
+ transform.apply_patterns to %contract_func {
+ transform.apply_patterns.tensor.fold_tensor_subset_ops
+ transform.apply_patterns.vector.drop_inner_most_unit_dims_from_xfer_ops
+ transform.apply_patterns.canonicalization
+ } : !transform.op<"func.func">
+
+ // Target FEAT_I8MM directly -- by this point the data is packed and
+ // statically shaped, so no masking survives to block the pattern.
+ transform.apply_patterns to %contract_func {
transform.apply_patterns.arm_neon.vector_contract_to_i8mm
} : !transform.op<"func.func">
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.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
+ }
- transform.foreach %funcs : !transform.op<"func.func"> {
- ^bb0(%func : !transform.op<"func.func">):
- transform.include @tile_and_vectorize_matmul failures(propagate)
- (%func) : (!transform.op<"func.func">) -> ()
- }
- transform.yield
+ transform.named_sequence @match_contract(
+ %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
+ transform.match.operation_name %entry ["vector.contract"] : !transform.any_op
+ transform.yield %entry : !transform.any_op
}
}
>From 52d5cadb5c6302557754f86e967c0802e96ddda9 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 5 Aug 2026 10:11:08 +0200
Subject: [PATCH 4/4] Address comments pt2
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/test/CMakeLists.txt | 3 +++
.../Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir | 2 +-
mlir/test/Integration/lit.local.cfg | 6 +++++-
mlir/test/lit.site.cfg.py.in | 3 +++
4 files changed, 12 insertions(+), 2 deletions(-)
diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index e0c32cd4bd9a0..c8b56caa5b9ec 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -40,11 +40,13 @@ if (MLIR_INCLUDE_INTEGRATION_TESTS)
option(MLIR_RUN_CUDA_SM90_TESTS "Run CUDA H100 tests.")
option(MLIR_RUN_ARM_SVE_TESTS "Run Arm SVE tests.")
option(MLIR_RUN_ARM_SME_TESTS "Run Arm SME tests.")
+ option(MLIR_RUN_ARM_I8MM_TESTS "Run Arm I8MM tests.")
# Check whether an emulator is required - if yes then make sure that it's
# been set.
check_emulator(MLIR_RUN_ARM_SVE_TESTS "HWCAP_SVE" ARM_EMULATOR_EXECUTABLE)
check_emulator(MLIR_RUN_ARM_SME_TESTS "HWCAP2_SME" ARM_EMULATOR_EXECUTABLE)
+ check_emulator(MLIR_RUN_ARM_I8MM_TESTS "HWCAP2_I8MM" ARM_EMULATOR_EXECUTABLE)
# The native target may not be enabled when cross compiling, raise an error.
if(NOT MLIR_ENABLE_EXECUTION_ENGINE)
@@ -82,6 +84,7 @@ llvm_canonicalize_cmake_booleans(
MLIR_RUN_X86_TESTS
MLIR_RUN_ARM_SVE_TESTS
MLIR_RUN_ARM_SME_TESTS
+ MLIR_RUN_ARM_I8MM_TESTS
MLIR_RUN_CUDA_SM80_TESTS
MLIR_RUN_CUDA_SM80_LT_TESTS
MLIR_RUN_CUDA_SM90_TESTS
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
index 94c73f639fd3b..70798738045ab 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmNeon/matmul-i8.mlir
@@ -1,4 +1,4 @@
-// REQUIRES: arm-emulator
+// REQUIRES: mlir_arm_i8mm_tests
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
diff --git a/mlir/test/Integration/lit.local.cfg b/mlir/test/Integration/lit.local.cfg
index 5f16b1cc3cc43..fc287f4ac878e 100644
--- a/mlir/test/Integration/lit.local.cfg
+++ b/mlir/test/Integration/lit.local.cfg
@@ -10,7 +10,11 @@ def configure_aarch64_mcr_cmd():
# NOTE: If the SVE tests are disabled and the SME tests are enabled to run
# under emulation, the SVE specific RUN lines in the SparseTensor tests
# will run under emulation.
- if not (config.mlir_run_arm_sve_tests or config.mlir_run_arm_sme_tests):
+ if not (
+ config.mlir_run_arm_sve_tests
+ or config.mlir_run_arm_sme_tests
+ or config.mlir_run_arm_i8mm_tests
+ ):
config.substitutions.append(("%mcr_aarch64_cmd", mcr_cmd))
return
diff --git a/mlir/test/lit.site.cfg.py.in b/mlir/test/lit.site.cfg.py.in
index 30b353f117bb2..e53c88648abe7 100644
--- a/mlir/test/lit.site.cfg.py.in
+++ b/mlir/test/lit.site.cfg.py.in
@@ -55,6 +55,9 @@ config.mlir_run_arm_sve_tests = @MLIR_RUN_ARM_SVE_TESTS@
if config.mlir_run_arm_sve_tests:
config.available_features.add("mlir_arm_sve_tests")
config.mlir_run_arm_sme_tests = @MLIR_RUN_ARM_SME_TESTS@
+config.mlir_run_arm_i8mm_tests = @MLIR_RUN_ARM_I8MM_TESTS@
+if config.mlir_run_arm_i8mm_tests:
+ config.available_features.add("mlir_arm_i8mm_tests")
config.mlir_run_x86_tests = @MLIR_RUN_X86_TESTS@
config.mlir_run_riscv_vector_tests = "@MLIR_RUN_RISCV_VECTOR_TESTS@"
config.mlir_run_cuda_tensor_core_tests = @MLIR_RUN_CUDA_TENSOR_CORE_TESTS@
More information about the Mlir-commits
mailing list