[Mlir-commits] [mlir] [mlir][ArmNeon] Add linalg.matmul e2e tests (PR #212809)
Federico Bruzzone
llvmlistbot at llvm.org
Mon Aug 3 08:04:22 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/2] [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 e9aa9581abde7de03e311223447be1e8bfc1636e 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/2] 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..91ba625a2e191 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 (standard, non-transposed layout).
+ %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">
More information about the Mlir-commits
mailing list