[Mlir-commits] [mlir] [mlir][sme] Add e2e test for lowering mmt4d to sme (PR #208226)
Stephen Long
llvmlistbot at llvm.org
Tue Jul 21 06:45:35 PDT 2026
https://github.com/steplong updated https://github.com/llvm/llvm-project/pull/208226
>From 27327d2c28378c73a64c24621a90e176c17cc7eb Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Tue, 7 Jul 2026 09:38:23 -0700
Subject: [PATCH 01/23] Copy pack-unpack-mmt4d.mlir from ArmSVE
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 403 ++++++++++++++++++
1 file changed, 403 insertions(+)
create mode 100644 mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
new file mode 100644
index 0000000000000..6192ed345debf
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -0,0 +1,403 @@
+// DEFINE: %{compile} = mlir-opt %s \
+// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
+// DEFINE: -cse -canonicalize -test-lower-to-llvm
+// DEFINE: %{entry_point} = main
+// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve"\
+// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
+
+// RUN: %{compile} | %{run} | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+/// HIGH-LEVEL OVERVIEW
+///
+/// End-to-end test for computing matrix-multiplication using linalg.mmt4d. In
+/// particular, demonstrates how the following MLIR sequence (implemented in
+/// @matmul_via_mmt4d):
+///
+/// A_pack = linalg.pack A
+/// B_pack = linalg.pack B
+/// C_pack = linalg.pack C
+/// out_pack = linalg.mmt4d(A_pack, B_pack, C_pack)
+///
+/// is equivalent to:
+///
+/// linalg.matmul(A, B, C)
+///
+/// (implemented in @matmul_via_matmul).
+///
+/// NOTES ON IMPLEMENTATION
+/// 1. The MMT4D example uses _scalable_ tile sizes for data tiling.
+/// * The matrix-multiplication dimension that's scalable: N.
+///
+/// 2. The lowering of linalg.mmt4d leverages scalable vectorisation.
+/// * The matrix-multiplication dimension that's scalable: N (to match data
+/// tiling configuration).
+///
+/// 3. Neither `linalg.pack` nor `linalg.unpack` are vectorised ATM.
+///
+/// 4. The MMT4D and Pack/Unpack Ops are kept in seperate functions to isolate
+/// the corresponding lowering and lowering configs.
+///
+/// TODO: Ideally, we should consider fusion opportunities by moving
+/// pack/unapack/mmt4d Ops into one function:
+/// * https://github.com/llvm/llvm-project/issues/159770
+/// TODO: Vectorize linalg.pack + linalg.unpack:
+/// * https://github.com/llvm/llvm-project/issues/159751
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// @main
+//
+// The main entry point that computes matrix multiplication via linalg.mmt4d
+// and linalg.matmul. Note, the output should be independent of the underlying
+// Linalg Op used, as well as SVE vector length.
+//===----------------------------------------------------------------------===//
+func.func @main() {
+ // Allocate and initialise the inputs
+ %A_empty = tensor.empty() : tensor<7x16xi32>
+ %B_empty = tensor.empty() : tensor<16x13xi32>
+
+ %c3 = arith.constant 3 : i32
+ %c4 = arith.constant 4 : i32
+ %A = linalg.fill ins(%c3 : i32) outs(%A_empty : tensor<7x16xi32>) -> tensor<7x16xi32>
+ %B = linalg.fill ins(%c4 : i32) outs(%B_empty : tensor<16x13xi32>) -> tensor<16x13xi32>
+ %C = arith.constant dense<[
+ [ 1, 8, 15, 22, 29, 36, 43, 50, 57, 64, 71, 78, 85],
+ [ 2, 9, 16, 23, 30, 37, 44, 51, 58, 65, 72, 79, 86],
+ [ 3, 10, 17, 24, 31, 38, 45, 52, 59, 66, 73, 80, 87],
+ [ 4, 11, 18, 25, 32, 39, 46, 53, 60, 67, 74, 81, 88],
+ [ 5, 12, 19, 26, 33, 40, 47, 54, 61, 68, 75, 82, 89],
+ [ 6, 13, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90],
+ [ 7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91]
+ ]> : tensor<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<7x16xi32>, tensor<16x13xi32>, tensor<7x13xi32>) -> tensor<7x13xi32>
+ %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<7x13xi32> to tensor<*xi32>
+ vector.print str "--------------------------\n"
+ vector.print str "RESULT FROM linalg.mmt4d:\n"
+ vector.print str "--------------------------\n"
+ call @printMemrefI32(%C_mmt4d_cast) : (tensor<*xi32>) -> ()
+
+ // VARIANT: Matrix multiplication via linalg.matmul
+ // CHECK: Unranked Memref
+ // CHECK: [193, 200, 207, 214, 221, 228, 235, 242, 249, 256, 263, 270, 277]
+ // CHECK: [194, 201, 208, 215, 222, 229, 236, 243, 250, 257, 264, 271, 278]
+ // CHECK: [195, 202, 209, 216, 223, 230, 237, 244, 251, 258, 265, 272, 279]
+ // CHECK: [196, 203, 210, 217, 224, 231, 238, 245, 252, 259, 266, 273, 280]
+ // CHECK: [197, 204, 211, 218, 225, 232, 239, 246, 253, 260, 267, 274, 281]
+ // CHECK: [198, 205, 212, 219, 226, 233, 240, 247, 254, 261, 268, 275, 282]
+ // CHECK: [199, 206, 213, 220, 227, 234, 241, 248, 255, 262, 269, 276, 283]
+ %C_matmul = func.call @matmul(%A, %B, %C) : (tensor<7x16xi32>, tensor<16x13xi32>, tensor<7x13xi32>) -> tensor<7x13xi32>
+ %C_matmul_cast = tensor.cast %C_matmul : tensor<7x13xi32> to tensor<*xi32>
+ vector.print str "\n--------------------------\n"
+ vector.print str "RESULT FROM linalg.matmul:\n"
+ vector.print str "--------------------------\n"
+ call @printMemrefI32(%C_matmul_cast) : (tensor<*xi32>) -> ()
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @matmul_via_matmul
+//
+// Implements matrix-multiplication via linalg.matmul
+//===----------------------------------------------------------------------===//
+func.func private @matmul(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<7x16xi32>, tensor<16x13xi32>)
+ outs(%C: tensor<7x13xi32>) -> tensor<7x13xi32>
+
+ return %C_matmul : tensor<7x13xi32>
+}
+
+//===----------------------------------------------------------------------===//
+// @matmul_via_mmt4d
+//
+// Implements matrix-multiplication via linalg.mmt4d
+//===----------------------------------------------------------------------===//
+func.func private @pack_lhs(%A: tensor<7x16xi32>) -> tensor<1x16x8x1xi32> {
+ %pad = arith.constant 0 : i32
+
+ %A_pack_empty = tensor.empty() : tensor<1x16x8x1xi32>
+ %A_pack = linalg.pack %A
+ padding_value(%pad : i32)
+ inner_dims_pos = [0, 1]
+ inner_tiles = [8, 1]
+ into %A_pack_empty : tensor<7x16xi32> -> tensor<1x16x8x1xi32>
+
+ return %A_pack : tensor<1x16x8x1xi32>
+}
+
+//===----------------------------------------------------------------------===//
+// @pack_rhs
+//
+// Implements packing for the B matrix (RHS) in matrix multiplication. The
+// inner tile size is "scalable": 8 * vscale.
+//===----------------------------------------------------------------------===//
+func.func private @pack_rhs(%B: tensor<16x13xi32>) -> tensor<?x16x?x1xi32> {
+ %pad = arith.constant 0 : i32
+
+ // Compute the outer tile size.
+ %vs = vector.vscale
+ %c8 = arith.constant 8 : index
+ %vs_c8 = arith.muli %vs, %c8 : index
+ %c13 = arith.constant 13 : index
+ %outer_tile_size = arith.ceildivui %c13, %vs_c8 : index
+
+ %B_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<?x16x?x1xi32>
+ %B_pack = linalg.pack %B
+ padding_value(%pad : i32)
+ outer_dims_perm = [1, 0]
+ inner_dims_pos = [1, 0]
+ inner_tiles = [%vs_c8, 1]
+ into %B_pack_empty : tensor<16x13xi32> -> tensor<?x16x?x1xi32>
+
+ return %B_pack : tensor<?x16x?x1xi32>
+}
+
+//===----------------------------------------------------------------------===//
+// @pack_acc
+//
+// Implements packing for the C matrix (accumulator) in matrix multiplication.
+// The inner tile size is "scalable": 8 * vscale
+//===----------------------------------------------------------------------===//
+func.func private @pack_acc(%C: tensor<7x13xi32>) -> tensor<1x?x8x?xi32> {
+ %pad = arith.constant 0 : i32
+
+ // Compute the outer tile size.
+ %c13 = arith.constant 13 : index
+ %vs = vector.vscale
+ %c8 = arith.constant 8 : index
+ %vs_c8 = arith.muli %vs, %c8 : index
+ %outer_tile_size = arith.ceildivui %c13, %vs_c8 : index
+
+ %C_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<1x?x8x?xi32>
+ %C_pack = linalg.pack %C
+ padding_value(%pad : i32)
+ outer_dims_perm = [0, 1]
+ inner_dims_pos = [0, 1]
+ inner_tiles = [8, %vs_c8] into %C_pack_empty : tensor<7x13xi32> -> tensor<1x?x8x?xi32>
+
+ return %C_pack : tensor<1x?x8x?xi32>
+}
+
+//===----------------------------------------------------------------------===//
+// @unpack_acc
+//
+// Implements unpacking for the C matrix (accumulator) in matrix
+// multiplication. The inner tile size is "scalable": 8 * vscale
+//===----------------------------------------------------------------------===//
+func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xi32>) -> tensor<7x13xi32> {
+ %vs = vector.vscale
+ %c8 = arith.constant 8 : index
+ %vs_c8 = arith.muli %vs, %c8 : index
+
+ %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 = [8, %vs_c8]
+ into %C_out_empty : tensor<1x?x8x?xi32> -> tensor<7x13xi32>
+
+ return %C_out_unpack: tensor<7x13xi32>
+}
+
+//===----------------------------------------------------------------------===//
+// Helper methods for printing
+//===----------------------------------------------------------------------===//
+func.func private @print_pack_A(%A_pack : tensor<1x16x8x1xi32>) -> () {
+ %A_pack_cast = tensor.cast %A_pack : tensor<1x16x8x1xi32> to tensor<*xi32>
+ call @printMemrefI32(%A_pack_cast) : (tensor<*xi32>) -> ()
+
+ return
+}
+
+func.func private @print_pack_B(%B_pack : tensor<?x16x?x1xi32>) -> () {
+ %B_pack_cast = tensor.cast %B_pack : tensor<?x16x?x1xi32> to tensor<*xi32>
+ call @printMemrefI32(%B_pack_cast) : (tensor<*xi32>) -> ()
+
+ return
+}
+
+func.func private @print_pack_C(%C_pack : tensor<1x?x8x?xi32>) -> () {
+ %C_pack_cast = tensor.cast %C_pack : tensor<1x?x8x?xi32> to tensor<*xi32>
+ call @printMemrefI32(%C_pack_cast) : (tensor<*xi32>) -> ()
+
+ return
+}
+
+//===----------------------------------------------------------------------===//
+// @matmul_via_mmt4d
+//
+// Implements matrix-multiplication via linalg.mmt4d
+//===----------------------------------------------------------------------===//
+func.func private @matmul_via_mmt4d(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {
+ // Pack input matrices
+ %A_pack = func.call @pack_lhs(%A): (tensor<7x16xi32>) -> tensor<1x16x8x1xi32>
+ %B_pack = func.call @pack_rhs(%B): (tensor<16x13xi32>) -> tensor<?x16x?x1xi32>
+ %C_pack = func.call @pack_acc(%C): (tensor<7x13xi32>) -> tensor<1x?x8x?xi32>
+
+ // Print the packed matrices (this is the only _visible_ part that changes
+ // when adjusting the SVE vector size).
+ func.call @print_pack_A(%A_pack) : (tensor<1x16x8x1xi32>) -> ()
+ func.call @print_pack_B(%B_pack) : (tensor<?x16x?x1xi32>) -> ()
+ func.call @print_pack_C(%C_pack) : (tensor<1x?x8x?xi32>) -> ()
+
+ // MMT4D
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x8x1xi32>, tensor<?x16x?x1xi32>) outs(%C_pack : tensor<1x?x8x?xi32>) -> tensor<1x?x8x?xi32>
+
+ // Unpack the output
+ %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x8x?xi32>) -> tensor<7x13xi32>
+
+ return %C_out_unpack : tensor<7x13xi32>
+}
+
+//===----------------------------------------------------------------------===//
+// TD Sequence
+//===----------------------------------------------------------------------===//
+module @transforms attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consumed}) {
+ //==========================================================================
+ // 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 parallel dims (note, the N dim is scalable!)
+ %tiled_mmt4d_parallel, %_:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, 8, [8], 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ // Tile reduction dims
+ %tiled_mmt4d, %_1:2 = transform.structured.tile_using_for %tiled_mmt4d_parallel tile_sizes [0, 0, 1, 0, 0, 1]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+
+ // Step 2: Vectorize linalg.mmt4d (note, the N dim is scalable!)
+ // TODO: Lower directly to named contractions: https://github.com/llvm/llvm-project/issues/159749
+ transform.structured.vectorize %tiled_mmt4d
+ vector_sizes [1, 1, 1, 8, [8], 1] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op
+
+ // Step 3: Simplify
+ // vector.multi_reduction --> vector.contract
+ // Generates a 6-dim vector.contract with the dim matching the original MMT4D Op
+ // and with the following split into parallel and reduction dims:
+ // * parallel, parallel, reduction, parallel, parallel, reduction
+ transform.apply_patterns to %mmt4d_func {
+ transform.apply_patterns.vector.reduction_to_contract
+ // Reduce the rank of xfer ops. This transforms vector.contract to be
+ // more matmul-like and to enable the lowering to outer product Ops.
+ 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
+ : (!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
+
+ // 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
+ 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)
+ %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [8, 1]
+ : (!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
+ 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">
+
+ // 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
+ } : !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">
+
+ //==========================================================================
+ // BUFFERIZATION
+ //==========================================================================
+ %bufferize = transform.bufferization.one_shot_bufferize %module
+ {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
+
+ //==========================================================================
+ // SIMPLIFY THE CONTRACT 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">
+
+ // Drop trailing unit dims (the correspondong pattern works only
+ // post-bufferization)
+ 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">
+
+ //==========================================================================
+ // LOWER CONTRACT TO FMA
+ //==========================================================================
+ transform.apply_patterns to %contract_func {
+ transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
+ transform.apply_patterns.vector.lower_outerproduct
+ } : !transform.op<"func.func">
+
+ 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
+ transform.yield %entry : !transform.any_op
+ }
+
+ 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
+ }
+}
+
+//===----------------------------------------------------------------------===//
+// Function signatures
+//===----------------------------------------------------------------------===//
+func.func private @printMemrefI32(%ptr : tensor<*xi32>)
>From 888bff0512276f88c6c4d70768587e19790f1b58 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Tue, 7 Jul 2026 09:40:44 -0700
Subject: [PATCH 02/23] Convert from I32 to F32
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 140 +++++++++---------
1 file changed, 70 insertions(+), 70 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 6192ed345debf..1b15c964d7bd3 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -1,9 +1,9 @@
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
-// DEFINE: -cse -canonicalize -test-lower-to-llvm
+// DEFINE: -canonicalize -test-lower-to-arm-sme -test-lower-to-llvm
// DEFINE: %{entry_point} = main
-// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve"\
-// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
+// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve,+sme"\
+// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib
// RUN: %{compile} | %{run} | FileCheck %s
@@ -30,7 +30,7 @@
/// * The matrix-multiplication dimension that's scalable: N.
///
/// 2. The lowering of linalg.mmt4d leverages scalable vectorisation.
-/// * The matrix-multiplication dimension that's scalable: N (to match data
+/// * The matrix-multiplication dimension that's scalable: M, N (to match data
/// tiling configuration).
///
/// 3. Neither `linalg.pack` nor `linalg.unpack` are vectorised ATM.
@@ -54,22 +54,22 @@
//===----------------------------------------------------------------------===//
func.func @main() {
// Allocate and initialise the inputs
- %A_empty = tensor.empty() : tensor<7x16xi32>
- %B_empty = tensor.empty() : tensor<16x13xi32>
+ %A_empty = tensor.empty() : tensor<7x16xf32>
+ %B_empty = tensor.empty() : tensor<16x13xf32>
- %c3 = arith.constant 3 : i32
- %c4 = arith.constant 4 : i32
- %A = linalg.fill ins(%c3 : i32) outs(%A_empty : tensor<7x16xi32>) -> tensor<7x16xi32>
- %B = linalg.fill ins(%c4 : i32) outs(%B_empty : tensor<16x13xi32>) -> tensor<16x13xi32>
+ %c3 = arith.constant 3.0 : f32
+ %c4 = arith.constant 4.0 : f32
+ %A = linalg.fill ins(%c3 : f32) outs(%A_empty : tensor<7x16xf32>) -> tensor<7x16xf32>
+ %B = linalg.fill ins(%c4 : f32) outs(%B_empty : tensor<16x13xf32>) -> tensor<16x13xf32>
%C = arith.constant dense<[
- [ 1, 8, 15, 22, 29, 36, 43, 50, 57, 64, 71, 78, 85],
- [ 2, 9, 16, 23, 30, 37, 44, 51, 58, 65, 72, 79, 86],
- [ 3, 10, 17, 24, 31, 38, 45, 52, 59, 66, 73, 80, 87],
- [ 4, 11, 18, 25, 32, 39, 46, 53, 60, 67, 74, 81, 88],
- [ 5, 12, 19, 26, 33, 40, 47, 54, 61, 68, 75, 82, 89],
- [ 6, 13, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90],
- [ 7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91]
- ]> : tensor<7x13xi32>
+ [ 1., 8., 15., 22., 29., 36., 43., 50., 57., 64., 71., 78., 85.],
+ [ 2., 9., 16., 23., 30., 37., 44., 51., 58., 65., 72., 79., 86.],
+ [ 3., 10., 17., 24., 31., 38., 45., 52., 59., 66., 73., 80., 87.],
+ [ 4., 11., 18., 25., 32., 39., 46., 53., 60., 67., 74., 81., 88.],
+ [ 5., 12., 19., 26., 33., 40., 47., 54., 61., 68., 75., 82., 89.],
+ [ 6., 13., 20., 27., 34., 41., 48., 55., 62., 69., 76., 83., 90.],
+ [ 7., 14., 21., 28., 35., 42., 49., 56., 63., 70., 77., 84., 91.]
+ ]> : tensor<7x13xf32>
// VARIANT: Matrix multiplication via linalg.mmt4d
// CHECK: Unranked Memref
@@ -80,12 +80,12 @@ func.func @main() {
// 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<7x16xi32>, tensor<16x13xi32>, tensor<7x13xi32>) -> tensor<7x13xi32>
- %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<7x13xi32> to tensor<*xi32>
+ %C_mmt4d = func.call @matmul_via_mmt4d(%A, %B, %C) : (tensor<7x16xf32>, tensor<16x13xf32>, tensor<7x13xf32>) -> tensor<7x13xf32>
+ %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<7x13xf32> to tensor<*xf32>
vector.print str "--------------------------\n"
vector.print str "RESULT FROM linalg.mmt4d:\n"
vector.print str "--------------------------\n"
- call @printMemrefI32(%C_mmt4d_cast) : (tensor<*xi32>) -> ()
+ call @printMemrefF32(%C_mmt4d_cast) : (tensor<*xf32>) -> ()
// VARIANT: Matrix multiplication via linalg.matmul
// CHECK: Unranked Memref
@@ -96,12 +96,12 @@ func.func @main() {
// 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<7x16xi32>, tensor<16x13xi32>, tensor<7x13xi32>) -> tensor<7x13xi32>
- %C_matmul_cast = tensor.cast %C_matmul : tensor<7x13xi32> to tensor<*xi32>
+ %C_matmul = func.call @matmul(%A, %B, %C) : (tensor<7x16xf32>, tensor<16x13xf32>, tensor<7x13xf32>) -> tensor<7x13xf32>
+ %C_matmul_cast = tensor.cast %C_matmul : tensor<7x13xf32> to tensor<*xf32>
vector.print str "\n--------------------------\n"
vector.print str "RESULT FROM linalg.matmul:\n"
vector.print str "--------------------------\n"
- call @printMemrefI32(%C_matmul_cast) : (tensor<*xi32>) -> ()
+ call @printMemrefF32(%C_matmul_cast) : (tensor<*xf32>) -> ()
return
}
@@ -111,11 +111,11 @@ func.func @main() {
//
// Implements matrix-multiplication via linalg.matmul
//===----------------------------------------------------------------------===//
-func.func private @matmul(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {
- %C_matmul = linalg.matmul ins(%A, %B: tensor<7x16xi32>, tensor<16x13xi32>)
- outs(%C: tensor<7x13xi32>) -> tensor<7x13xi32>
+func.func private @matmul(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tensor<7x13xf32>) -> tensor<7x13xf32> {
+ %C_matmul = linalg.matmul ins(%A, %B: tensor<7x16xf32>, tensor<16x13xf32>)
+ outs(%C: tensor<7x13xf32>) -> tensor<7x13xf32>
- return %C_matmul : tensor<7x13xi32>
+ return %C_matmul : tensor<7x13xf32>
}
//===----------------------------------------------------------------------===//
@@ -123,17 +123,17 @@ func.func private @matmul(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tenso
//
// Implements matrix-multiplication via linalg.mmt4d
//===----------------------------------------------------------------------===//
-func.func private @pack_lhs(%A: tensor<7x16xi32>) -> tensor<1x16x8x1xi32> {
- %pad = arith.constant 0 : i32
+func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x8x1xf32> {
+ %pad = arith.constant 0.0 : f32
- %A_pack_empty = tensor.empty() : tensor<1x16x8x1xi32>
+ %A_pack_empty = tensor.empty() : tensor<1x16x8x1xf32>
%A_pack = linalg.pack %A
- padding_value(%pad : i32)
+ padding_value(%pad : f32)
inner_dims_pos = [0, 1]
inner_tiles = [8, 1]
- into %A_pack_empty : tensor<7x16xi32> -> tensor<1x16x8x1xi32>
+ into %A_pack_empty : tensor<7x16xf32> -> tensor<1x16x8x1xf32>
- return %A_pack : tensor<1x16x8x1xi32>
+ return %A_pack : tensor<1x16x8x1xf32>
}
//===----------------------------------------------------------------------===//
@@ -142,8 +142,8 @@ func.func private @pack_lhs(%A: tensor<7x16xi32>) -> tensor<1x16x8x1xi32> {
// Implements packing for the B matrix (RHS) in matrix multiplication. The
// inner tile size is "scalable": 8 * vscale.
//===----------------------------------------------------------------------===//
-func.func private @pack_rhs(%B: tensor<16x13xi32>) -> tensor<?x16x?x1xi32> {
- %pad = arith.constant 0 : i32
+func.func private @pack_rhs(%B: tensor<16x13xf32>) -> tensor<?x16x?x1xf32> {
+ %pad = arith.constant 0.0 : f32
// Compute the outer tile size.
%vs = vector.vscale
@@ -152,15 +152,15 @@ func.func private @pack_rhs(%B: tensor<16x13xi32>) -> tensor<?x16x?x1xi32> {
%c13 = arith.constant 13 : index
%outer_tile_size = arith.ceildivui %c13, %vs_c8 : index
- %B_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<?x16x?x1xi32>
+ %B_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<?x16x?x1xf32>
%B_pack = linalg.pack %B
- padding_value(%pad : i32)
+ padding_value(%pad : f32)
outer_dims_perm = [1, 0]
inner_dims_pos = [1, 0]
inner_tiles = [%vs_c8, 1]
- into %B_pack_empty : tensor<16x13xi32> -> tensor<?x16x?x1xi32>
+ into %B_pack_empty : tensor<16x13xf32> -> tensor<?x16x?x1xf32>
- return %B_pack : tensor<?x16x?x1xi32>
+ return %B_pack : tensor<?x16x?x1xf32>
}
//===----------------------------------------------------------------------===//
@@ -169,8 +169,8 @@ func.func private @pack_rhs(%B: tensor<16x13xi32>) -> tensor<?x16x?x1xi32> {
// Implements packing for the C matrix (accumulator) in matrix multiplication.
// The inner tile size is "scalable": 8 * vscale
//===----------------------------------------------------------------------===//
-func.func private @pack_acc(%C: tensor<7x13xi32>) -> tensor<1x?x8x?xi32> {
- %pad = arith.constant 0 : i32
+func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<1x?x8x?xf32> {
+ %pad = arith.constant 0.0 : f32
// Compute the outer tile size.
%c13 = arith.constant 13 : index
@@ -179,14 +179,14 @@ func.func private @pack_acc(%C: tensor<7x13xi32>) -> tensor<1x?x8x?xi32> {
%vs_c8 = arith.muli %vs, %c8 : index
%outer_tile_size = arith.ceildivui %c13, %vs_c8 : index
- %C_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<1x?x8x?xi32>
+ %C_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<1x?x8x?xf32>
%C_pack = linalg.pack %C
- padding_value(%pad : i32)
+ padding_value(%pad : f32)
outer_dims_perm = [0, 1]
inner_dims_pos = [0, 1]
- inner_tiles = [8, %vs_c8] into %C_pack_empty : tensor<7x13xi32> -> tensor<1x?x8x?xi32>
+ inner_tiles = [8, %vs_c8] into %C_pack_empty : tensor<7x13xf32> -> tensor<1x?x8x?xf32>
- return %C_pack : tensor<1x?x8x?xi32>
+ return %C_pack : tensor<1x?x8x?xf32>
}
//===----------------------------------------------------------------------===//
@@ -195,41 +195,41 @@ func.func private @pack_acc(%C: tensor<7x13xi32>) -> tensor<1x?x8x?xi32> {
// Implements unpacking for the C matrix (accumulator) in matrix
// multiplication. The inner tile size is "scalable": 8 * vscale
//===----------------------------------------------------------------------===//
-func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xi32>) -> tensor<7x13xi32> {
+func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xf32>) -> tensor<7x13xf32> {
%vs = vector.vscale
%c8 = arith.constant 8 : index
%vs_c8 = arith.muli %vs, %c8 : index
- %C_out_empty = tensor.empty() : tensor<7x13xi32>
+ %C_out_empty = tensor.empty() : tensor<7x13xf32>
%C_out_unpack = linalg.unpack %C_packed
outer_dims_perm = [0, 1]
inner_dims_pos = [0, 1]
inner_tiles = [8, %vs_c8]
- into %C_out_empty : tensor<1x?x8x?xi32> -> tensor<7x13xi32>
+ into %C_out_empty : tensor<1x?x8x?xf32> -> tensor<7x13xf32>
- return %C_out_unpack: tensor<7x13xi32>
+ return %C_out_unpack: tensor<7x13xf32>
}
//===----------------------------------------------------------------------===//
// Helper methods for printing
//===----------------------------------------------------------------------===//
-func.func private @print_pack_A(%A_pack : tensor<1x16x8x1xi32>) -> () {
- %A_pack_cast = tensor.cast %A_pack : tensor<1x16x8x1xi32> to tensor<*xi32>
- call @printMemrefI32(%A_pack_cast) : (tensor<*xi32>) -> ()
+func.func private @print_pack_A(%A_pack : tensor<1x16x8x1xf32>) -> () {
+ %A_pack_cast = tensor.cast %A_pack : tensor<1x16x8x1xf32> to tensor<*xf32>
+ call @printMemrefF32(%A_pack_cast) : (tensor<*xf32>) -> ()
return
}
-func.func private @print_pack_B(%B_pack : tensor<?x16x?x1xi32>) -> () {
- %B_pack_cast = tensor.cast %B_pack : tensor<?x16x?x1xi32> to tensor<*xi32>
- call @printMemrefI32(%B_pack_cast) : (tensor<*xi32>) -> ()
+func.func private @print_pack_B(%B_pack : tensor<?x16x?x1xf32>) -> () {
+ %B_pack_cast = tensor.cast %B_pack : tensor<?x16x?x1xf32> to tensor<*xf32>
+ call @printMemrefF32(%B_pack_cast) : (tensor<*xf32>) -> ()
return
}
-func.func private @print_pack_C(%C_pack : tensor<1x?x8x?xi32>) -> () {
- %C_pack_cast = tensor.cast %C_pack : tensor<1x?x8x?xi32> to tensor<*xi32>
- call @printMemrefI32(%C_pack_cast) : (tensor<*xi32>) -> ()
+func.func private @print_pack_C(%C_pack : tensor<1x?x8x?xf32>) -> () {
+ %C_pack_cast = tensor.cast %C_pack : tensor<1x?x8x?xf32> to tensor<*xf32>
+ call @printMemrefF32(%C_pack_cast) : (tensor<*xf32>) -> ()
return
}
@@ -239,25 +239,25 @@ func.func private @print_pack_C(%C_pack : tensor<1x?x8x?xi32>) -> () {
//
// Implements matrix-multiplication via linalg.mmt4d
//===----------------------------------------------------------------------===//
-func.func private @matmul_via_mmt4d(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {
+func.func private @matmul_via_mmt4d(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tensor<7x13xf32>) -> tensor<7x13xf32> {
// Pack input matrices
- %A_pack = func.call @pack_lhs(%A): (tensor<7x16xi32>) -> tensor<1x16x8x1xi32>
- %B_pack = func.call @pack_rhs(%B): (tensor<16x13xi32>) -> tensor<?x16x?x1xi32>
- %C_pack = func.call @pack_acc(%C): (tensor<7x13xi32>) -> tensor<1x?x8x?xi32>
+ %A_pack = func.call @pack_lhs(%A): (tensor<7x16xf32>) -> tensor<1x16x8x1xf32>
+ %B_pack = func.call @pack_rhs(%B): (tensor<16x13xf32>) -> tensor<?x16x?x1xf32>
+ %C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<1x?x8x?xf32>
// Print the packed matrices (this is the only _visible_ part that changes
// when adjusting the SVE vector size).
- func.call @print_pack_A(%A_pack) : (tensor<1x16x8x1xi32>) -> ()
- func.call @print_pack_B(%B_pack) : (tensor<?x16x?x1xi32>) -> ()
- func.call @print_pack_C(%C_pack) : (tensor<1x?x8x?xi32>) -> ()
+ func.call @print_pack_A(%A_pack) : (tensor<1x16x8x1xf32>) -> ()
+ func.call @print_pack_B(%B_pack) : (tensor<?x16x?x1xf32>) -> ()
+ func.call @print_pack_C(%C_pack) : (tensor<1x?x8x?xf32>) -> ()
// MMT4D
- %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x8x1xi32>, tensor<?x16x?x1xi32>) outs(%C_pack : tensor<1x?x8x?xi32>) -> tensor<1x?x8x?xi32>
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x8x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<1x?x8x?xf32>) -> tensor<1x?x8x?xf32>
// Unpack the output
- %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x8x?xi32>) -> tensor<7x13xi32>
+ %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x8x?xf32>) -> tensor<7x13xf32>
- return %C_out_unpack : tensor<7x13xi32>
+ return %C_out_unpack : tensor<7x13xf32>
}
//===----------------------------------------------------------------------===//
@@ -400,4 +400,4 @@ module @transforms attributes { transform.with_named_sequence } {
//===----------------------------------------------------------------------===//
// Function signatures
//===----------------------------------------------------------------------===//
-func.func private @printMemrefI32(%ptr : tensor<*xi32>)
+func.func private @printMemrefF32(%ptr : tensor<*xf32>)
>From c7be7ca7cbc0c1687ac60556df78e7779fa80acc Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Tue, 7 Jul 2026 12:25:38 -0700
Subject: [PATCH 03/23] Add lowering sequence that hoists
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 98 ++++++++-----------
1 file changed, 39 insertions(+), 59 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 1b15c964d7bd3..f4a3650ed4ac9 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -272,44 +272,44 @@ module @transforms attributes { transform.with_named_sequence } {
%mmt4d_func = transform.get_parent_op %mmt4d {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">
// Step 1: Tile
- // Tile parallel dims (note, the N dim is scalable!)
- %tiled_mmt4d_parallel, %_:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, 8, [8], 0]
+ // Tile parallel dims (note, the M, N dim is scalable!)
+ %tiled_mmt4d_parallel, %_:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, [8], [8], 0]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
// Tile reduction dims
- %tiled_mmt4d, %_1:2 = transform.structured.tile_using_for %tiled_mmt4d_parallel tile_sizes [0, 0, 1, 0, 0, 1]
- : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ %tiled_mmt4d, %loop_k = transform.structured.tile_using_for %tiled_mmt4d_parallel tile_sizes [0, 0, 1, 0, 0, 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)
- // Step 2: Vectorize linalg.mmt4d (note, the N dim is scalable!)
+ // Step 2: Vectorize linalg.mmt4d (note, the M, N dim is scalable!)
// TODO: Lower directly to named contractions: https://github.com/llvm/llvm-project/issues/159749
transform.structured.vectorize %tiled_mmt4d
- vector_sizes [1, 1, 1, 8, [8], 1] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op
-
- // Step 3: Simplify
- // vector.multi_reduction --> vector.contract
- // Generates a 6-dim vector.contract with the dim matching the original MMT4D Op
- // and with the following split into parallel and reduction dims:
- // * parallel, parallel, reduction, parallel, parallel, reduction
- transform.apply_patterns to %mmt4d_func {
+ vector_sizes [1, 1, 1, [8], [8], 1] : !transform.any_op
+
+ // Step 3: Lower vector.mask %mask { vector.transfer_* } to vector.transfer_* %mask
+ transform.apply_patterns to %loop_k {
+ transform.apply_patterns.vector.lower_masked_transfers
+ } : !transform.op<"scf.for">
+
+ // Step 4: Hoist the C accumulator load/store out of the k-loop while still
+ // in tensor form, so transfer_write has a result value the loop can yield.
+ transform.apply_licm to %loop_k : !transform.op<"scf.for">
+ transform.loop.hoist_loop_invariant_subsets %loop_k : !transform.op<"scf.for">
+
+ // Lower to outerproduct
+ %func_pre = transform.structured.match ops{["func.func"]} in %module
+ : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func_pre {
transform.apply_patterns.vector.reduction_to_contract
- // Reduce the rank of xfer ops. This transforms vector.contract to be
- // more matmul-like and to enable the lowering to outer product Ops.
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
- : (!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
-
- // 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
- transform.apply_patterns.canonicalization
- } : !transform.op<"func.func">
+ transform.apply_patterns.canonicalization
+ } : !transform.any_op
+ transform.apply_patterns to %func_pre {
+ transform.apply_patterns.vector.reduction_to_contract
+ transform.apply_patterns.vector.cast_away_vector_leading_one_dim
+ transform.apply_patterns.vector.lower_contraction
+ lowering_strategy = "outerproduct"
+ transform.apply_patterns.vector.drop_unit_dims_with_shape_cast
+ transform.apply_patterns.canonicalization
+ } {apply_cse} : !transform.any_op
//==========================================================================
// HANDLE PACK + UNPACK
@@ -353,30 +353,16 @@ module @transforms attributes { transform.with_named_sequence } {
//==========================================================================
// BUFFERIZATION
//==========================================================================
- %bufferize = transform.bufferization.one_shot_bufferize %module
+ %bufferize = transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module
{bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
+ %func = transform.structured.match ops{["func.func"]} in %bufferize
+ : (!transform.any_op) -> !transform.any_op
- //==========================================================================
- // SIMPLIFY THE CONTRACT 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">
-
- // Drop trailing unit dims (the correspondong pattern works only
- // post-bufferization)
- 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">
-
- //==========================================================================
- // LOWER CONTRACT TO FMA
- //==========================================================================
- transform.apply_patterns to %contract_func {
- transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"
- transform.apply_patterns.vector.lower_outerproduct
- } : !transform.op<"func.func">
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.rank_reducing_subview_patterns
+ transform.apply_patterns.vector.drop_unit_dims_with_shape_cast
+ transform.apply_patterns.canonicalization
+ } {apply_cse} : !transform.any_op
transform.yield
}
@@ -389,12 +375,6 @@ module @transforms attributes { transform.with_named_sequence } {
transform.match.operation_name %entry ["linalg.mmt4d"] : !transform.any_op
transform.yield %entry : !transform.any_op
}
-
- 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 aef815fbba2f7b352b27ac579fc33cedc1c1d97d Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Tue, 7 Jul 2026 14:32:54 -0700
Subject: [PATCH 04/23] Removed print pack functions because they introduced
memcpy calls that required bringing in __arm_sc_memcpy
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 33 ++-----------------
1 file changed, 2 insertions(+), 31 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index f4a3650ed4ac9..745c076c265f8 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -1,6 +1,7 @@
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \
-// DEFINE: -canonicalize -test-lower-to-arm-sme -test-lower-to-llvm
+// DEFINE: -canonicalize -test-lower-to-arm-sme -convert-vector-to-llvm="enable-arm-sve" \
+// DEFINE: -test-lower-to-llvm
// DEFINE: %{entry_point} = main
// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve,+sme"\
// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib
@@ -210,30 +211,6 @@ func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xf32>) -> tensor<7x13xf32
return %C_out_unpack: tensor<7x13xf32>
}
-//===----------------------------------------------------------------------===//
-// Helper methods for printing
-//===----------------------------------------------------------------------===//
-func.func private @print_pack_A(%A_pack : tensor<1x16x8x1xf32>) -> () {
- %A_pack_cast = tensor.cast %A_pack : tensor<1x16x8x1xf32> to tensor<*xf32>
- call @printMemrefF32(%A_pack_cast) : (tensor<*xf32>) -> ()
-
- return
-}
-
-func.func private @print_pack_B(%B_pack : tensor<?x16x?x1xf32>) -> () {
- %B_pack_cast = tensor.cast %B_pack : tensor<?x16x?x1xf32> to tensor<*xf32>
- call @printMemrefF32(%B_pack_cast) : (tensor<*xf32>) -> ()
-
- return
-}
-
-func.func private @print_pack_C(%C_pack : tensor<1x?x8x?xf32>) -> () {
- %C_pack_cast = tensor.cast %C_pack : tensor<1x?x8x?xf32> to tensor<*xf32>
- call @printMemrefF32(%C_pack_cast) : (tensor<*xf32>) -> ()
-
- return
-}
-
//===----------------------------------------------------------------------===//
// @matmul_via_mmt4d
//
@@ -245,12 +222,6 @@ func.func private @matmul_via_mmt4d(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>,
%B_pack = func.call @pack_rhs(%B): (tensor<16x13xf32>) -> tensor<?x16x?x1xf32>
%C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<1x?x8x?xf32>
- // Print the packed matrices (this is the only _visible_ part that changes
- // when adjusting the SVE vector size).
- func.call @print_pack_A(%A_pack) : (tensor<1x16x8x1xf32>) -> ()
- func.call @print_pack_B(%B_pack) : (tensor<?x16x?x1xf32>) -> ()
- func.call @print_pack_C(%C_pack) : (tensor<1x?x8x?xf32>) -> ()
-
// MMT4D
%mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x8x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<1x?x8x?xf32>) -> tensor<1x?x8x?xf32>
>From 9d77231066583b8085c8060a77009cb3f9b74bed Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Thu, 9 Jul 2026 10:13:31 -0400
Subject: [PATCH 05/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
Co-authored-by: Ege Beysel <beyselege at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 745c076c265f8..76cded79bcce7 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -28,7 +28,7 @@
///
/// NOTES ON IMPLEMENTATION
/// 1. The MMT4D example uses _scalable_ tile sizes for data tiling.
-/// * The matrix-multiplication dimension that's scalable: N.
+/// * The matrix-multiplication dimensions that are scalable: M and N.
///
/// 2. The lowering of linalg.mmt4d leverages scalable vectorisation.
/// * The matrix-multiplication dimension that's scalable: M, N (to match data
>From a15409a499e9b76eb6a9bcfe013bd546607e4f94 Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Thu, 9 Jul 2026 10:13:47 -0400
Subject: [PATCH 06/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
Co-authored-by: Ege Beysel <beyselege at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 76cded79bcce7..b71995f9b49a7 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -40,7 +40,7 @@
/// the corresponding lowering and lowering configs.
///
/// TODO: Ideally, we should consider fusion opportunities by moving
-/// pack/unapack/mmt4d Ops into one function:
+/// pack/unpack/mmt4d ops into one function:
/// * https://github.com/llvm/llvm-project/issues/159770
/// TODO: Vectorize linalg.pack + linalg.unpack:
/// * https://github.com/llvm/llvm-project/issues/159751
>From 8e4bca2ca57ce5c8ced18691a95f3468ceb23445 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Thu, 9 Jul 2026 07:18:49 -0700
Subject: [PATCH 07/23] Remove +sve flag
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index b71995f9b49a7..893710fd800d2 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -3,7 +3,7 @@
// DEFINE: -canonicalize -test-lower-to-arm-sme -convert-vector-to-llvm="enable-arm-sve" \
// DEFINE: -test-lower-to-llvm
// DEFINE: %{entry_point} = main
-// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve,+sme"\
+// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sme"\
// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib
// RUN: %{compile} | %{run} | FileCheck %s
>From 457f16cf5289905ee0b49c287480cd1f7ba3bd55 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Thu, 9 Jul 2026 07:22:33 -0700
Subject: [PATCH 08/23] Add linalg.unpack comment
---
.../Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 1 +
1 file changed, 1 insertion(+)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 893710fd800d2..d6a6bedc75a10 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -19,6 +19,7 @@
/// B_pack = linalg.pack B
/// C_pack = linalg.pack C
/// out_pack = linalg.mmt4d(A_pack, B_pack, C_pack)
+/// out = linalg.unpack out_pack
///
/// is equivalent to:
///
>From d608b7d4f6291ff3f435d488f82f89da912f25bf Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Thu, 9 Jul 2026 07:35:29 -0700
Subject: [PATCH 09/23] Added create_named_contractions. Removing 2nd
reduction_to_contract causes test to fail
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index d6a6bedc75a10..32984a04f6a99 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -252,9 +252,8 @@ module @transforms attributes { transform.with_named_sequence } {
: (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)
// Step 2: Vectorize linalg.mmt4d (note, the M, N dim is scalable!)
- // TODO: Lower directly to named contractions: https://github.com/llvm/llvm-project/issues/159749
- transform.structured.vectorize %tiled_mmt4d
- vector_sizes [1, 1, 1, [8], [8], 1] : !transform.any_op
+ transform.structured.vectorize %tiled_mmt4d vector_sizes [1, 1, 1, [8], [8], 1] {create_named_contraction}
+ : !transform.any_op
// Step 3: Lower vector.mask %mask { vector.transfer_* } to vector.transfer_* %mask
transform.apply_patterns to %loop_k {
@@ -270,7 +269,6 @@ module @transforms attributes { transform.with_named_sequence } {
%func_pre = transform.structured.match ops{["func.func"]} in %module
: (!transform.any_op) -> !transform.any_op
transform.apply_patterns to %func_pre {
- transform.apply_patterns.vector.reduction_to_contract
transform.apply_patterns.vector.transfer_permutation_patterns
transform.apply_patterns.canonicalization
} : !transform.any_op
>From c1290ec4e38aa1c53dd4285ace01665dce5ac65c Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Thu, 9 Jul 2026 07:47:50 -0700
Subject: [PATCH 10/23] Fix some comments
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 32984a04f6a99..06dae3ee89042 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -244,14 +244,14 @@ module @transforms attributes { transform.with_named_sequence } {
%mmt4d_func = transform.get_parent_op %mmt4d {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">
// Step 1: Tile
- // Tile parallel dims (note, the M, N dim is scalable!)
+ // Tile parallel dims (note, the M, N dims are scalable!)
%tiled_mmt4d_parallel, %_:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, [8], [8], 0]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
// Tile reduction dims
%tiled_mmt4d, %loop_k = transform.structured.tile_using_for %tiled_mmt4d_parallel tile_sizes [0, 0, 1, 0, 0, 0]
: (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)
- // Step 2: Vectorize linalg.mmt4d (note, the M, N dim is scalable!)
+ // Step 2: Vectorize linalg.mmt4d (note, the M, N dims are scalable!)
transform.structured.vectorize %tiled_mmt4d vector_sizes [1, 1, 1, [8], [8], 1] {create_named_contraction}
: !transform.any_op
>From 5e79874fa5aceaa9a854ce19646a871213ae0af0 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Thu, 9 Jul 2026 09:35:52 -0700
Subject: [PATCH 11/23] Make A_pack inner_tiles [%vscale_c8, 1]
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 18 +++++++++++-------
1 file changed, 11 insertions(+), 7 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 06dae3ee89042..7b8b3b4c74ff1 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -125,17 +125,21 @@ func.func private @matmul(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tenso
//
// Implements matrix-multiplication via linalg.mmt4d
//===----------------------------------------------------------------------===//
-func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x8x1xf32> {
+func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x?x1xf32> {
%pad = arith.constant 0.0 : f32
- %A_pack_empty = tensor.empty() : tensor<1x16x8x1xf32>
+ %vs = vector.vscale
+ %c8 = arith.constant 8 : index
+ %vs_c8 = arith.muli %vs, %c8 : index
+
+ %A_pack_empty = tensor.empty(%vs_c8) : tensor<1x16x?x1xf32>
%A_pack = linalg.pack %A
padding_value(%pad : f32)
inner_dims_pos = [0, 1]
- inner_tiles = [8, 1]
- into %A_pack_empty : tensor<7x16xf32> -> tensor<1x16x8x1xf32>
+ inner_tiles = [%vs_c8, 1]
+ into %A_pack_empty : tensor<7x16xf32> -> tensor<1x16x?x1xf32>
- return %A_pack : tensor<1x16x8x1xf32>
+ return %A_pack : tensor<1x16x?x1xf32>
}
//===----------------------------------------------------------------------===//
@@ -219,12 +223,12 @@ func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xf32>) -> tensor<7x13xf32
//===----------------------------------------------------------------------===//
func.func private @matmul_via_mmt4d(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tensor<7x13xf32>) -> tensor<7x13xf32> {
// Pack input matrices
- %A_pack = func.call @pack_lhs(%A): (tensor<7x16xf32>) -> tensor<1x16x8x1xf32>
+ %A_pack = func.call @pack_lhs(%A): (tensor<7x16xf32>) -> tensor<1x16x?x1xf32>
%B_pack = func.call @pack_rhs(%B): (tensor<16x13xf32>) -> tensor<?x16x?x1xf32>
%C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<1x?x8x?xf32>
// MMT4D
- %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x8x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<1x?x8x?xf32>) -> tensor<1x?x8x?xf32>
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x?x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<1x?x8x?xf32>) -> tensor<1x?x8x?xf32>
// Unpack the output
%C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x8x?xf32>) -> tensor<7x13xf32>
>From ff4a55c252b87045781e1a375221091b43c4780f Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Thu, 9 Jul 2026 10:12:37 -0700
Subject: [PATCH 12/23] Make C_pack and C_out_unpack's inner tiles [%vs_c8,
%vs_c8]
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 33 +++++++++++--------
1 file changed, 19 insertions(+), 14 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 7b8b3b4c74ff1..f0348c1679723 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -175,24 +175,26 @@ func.func private @pack_rhs(%B: tensor<16x13xf32>) -> tensor<?x16x?x1xf32> {
// Implements packing for the C matrix (accumulator) in matrix multiplication.
// The inner tile size is "scalable": 8 * vscale
//===----------------------------------------------------------------------===//
-func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<1x?x8x?xf32> {
+func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<?x?x?x?xf32> {
%pad = arith.constant 0.0 : f32
- // Compute the outer tile size.
+ // Compute the outer tile sizes.
+ %c7 = arith.constant 7 : index
%c13 = arith.constant 13 : index
%vs = vector.vscale
%c8 = arith.constant 8 : index
%vs_c8 = arith.muli %vs, %c8 : index
- %outer_tile_size = arith.ceildivui %c13, %vs_c8 : index
+ %outer_tile_size_0 = arith.ceildivui %c7, %vs_c8 : index
+ %outer_tile_size_1 = arith.ceildivui %c13, %vs_c8 : index
- %C_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<1x?x8x?xf32>
+ %C_pack_empty = tensor.empty(%outer_tile_size_0, %outer_tile_size_1, %vs_c8, %vs_c8) : tensor<?x?x?x?xf32>
%C_pack = linalg.pack %C
padding_value(%pad : f32)
outer_dims_perm = [0, 1]
inner_dims_pos = [0, 1]
- inner_tiles = [8, %vs_c8] into %C_pack_empty : tensor<7x13xf32> -> tensor<1x?x8x?xf32>
+ inner_tiles = [%vs_c8, %vs_c8] into %C_pack_empty : tensor<7x13xf32> -> tensor<?x?x?x?xf32>
- return %C_pack : tensor<1x?x8x?xf32>
+ return %C_pack : tensor<?x?x?x?xf32>
}
//===----------------------------------------------------------------------===//
@@ -201,7 +203,7 @@ func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<1x?x8x?xf32> {
// Implements unpacking for the C matrix (accumulator) in matrix
// multiplication. The inner tile size is "scalable": 8 * vscale
//===----------------------------------------------------------------------===//
-func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xf32>) -> tensor<7x13xf32> {
+func.func private @unpack_acc(%C_packed: tensor<?x?x?x?xf32>) -> tensor<7x13xf32> {
%vs = vector.vscale
%c8 = arith.constant 8 : index
%vs_c8 = arith.muli %vs, %c8 : index
@@ -210,8 +212,8 @@ func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xf32>) -> tensor<7x13xf32
%C_out_unpack = linalg.unpack %C_packed
outer_dims_perm = [0, 1]
inner_dims_pos = [0, 1]
- inner_tiles = [8, %vs_c8]
- into %C_out_empty : tensor<1x?x8x?xf32> -> tensor<7x13xf32>
+ inner_tiles = [%vs_c8, %vs_c8]
+ into %C_out_empty : tensor<?x?x?x?xf32> -> tensor<7x13xf32>
return %C_out_unpack: tensor<7x13xf32>
}
@@ -225,13 +227,13 @@ func.func private @matmul_via_mmt4d(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>,
// Pack input matrices
%A_pack = func.call @pack_lhs(%A): (tensor<7x16xf32>) -> tensor<1x16x?x1xf32>
%B_pack = func.call @pack_rhs(%B): (tensor<16x13xf32>) -> tensor<?x16x?x1xf32>
- %C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<1x?x8x?xf32>
+ %C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<?x?x?x?xf32>
// MMT4D
- %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x?x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<1x?x8x?xf32>) -> tensor<1x?x8x?xf32>
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x?x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
// Unpack the output
- %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x8x?xf32>) -> tensor<7x13xf32>
+ %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<?x?x?x?xf32>) -> tensor<7x13xf32>
return %C_out_unpack : tensor<7x13xf32>
}
@@ -297,8 +299,11 @@ module @transforms attributes { transform.with_named_sequence } {
: (!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)
- %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [8, 1]
+ // and other lower-level Ops (see step 2). Tile by [8*vscale, 8*vscale]
+ // (scalable, equal to the inner tile sizes) so the tiled outer dims are
+ // statically 1 and DecomposeOuterUnitDimsUnPackOp can fire.
+ %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [[8], [8]]
+ inner_tile_alignments = [Equal, Equal]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
// 2.1. Decompose tiled PackOp into lower-level Ops + simplify
>From d3cf3bb016b8edc571a7e56209c722f9efb619b0 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Mon, 13 Jul 2026 06:38:23 -0700
Subject: [PATCH 13/23] Remove dead %mmt4d_func
---
.../Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index f0348c1679723..db15eff92cad3 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -247,7 +247,6 @@ module @transforms attributes { transform.with_named_sequence } {
// 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 parallel dims (note, the M, N dims are scalable!)
>From bb0524298418dfdfe92a0a1b9a87a8735caa9b42 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Mon, 13 Jul 2026 06:41:16 -0700
Subject: [PATCH 14/23] Fix comment for pack_lhs
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index db15eff92cad3..b599b95cd19fa 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -121,9 +121,10 @@ func.func private @matmul(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tenso
}
//===----------------------------------------------------------------------===//
-// @matmul_via_mmt4d
+// @pack_lhs
//
-// Implements matrix-multiplication via linalg.mmt4d
+// Implements packing for the A matrix (LHS) in matrix multiplication. The
+// inner tile size is "scalable": 8 * vscale.
//===----------------------------------------------------------------------===//
func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x?x1xf32> {
%pad = arith.constant 0.0 : f32
>From 6bbc2112fef26cf55990ec419c54e97955ed7bc4 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Mon, 13 Jul 2026 08:54:55 -0700
Subject: [PATCH 15/23] Add flag assume_dynamic_dims_match_vec_sizes when
vectorizing
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index b599b95cd19fa..0db52af86c17e 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -258,7 +258,7 @@ module @transforms attributes { transform.with_named_sequence } {
: (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)
// Step 2: Vectorize linalg.mmt4d (note, the M, N dims are scalable!)
- transform.structured.vectorize %tiled_mmt4d vector_sizes [1, 1, 1, [8], [8], 1] {create_named_contraction}
+ transform.structured.vectorize %tiled_mmt4d vector_sizes [1, 1, 1, [8], [8], 1] {assume_dynamic_dims_match_vec_sizes, create_named_contraction}
: !transform.any_op
// Step 3: Lower vector.mask %mask { vector.transfer_* } to vector.transfer_* %mask
>From bee3857e57916b02fe0693ddcc429b6985001734 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Mon, 13 Jul 2026 08:58:12 -0700
Subject: [PATCH 16/23] Remove TODO about fusion
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 3 ---
1 file changed, 3 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 0db52af86c17e..a520816bce2f0 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -40,9 +40,6 @@
/// 4. The MMT4D and Pack/Unpack Ops are kept in seperate functions to isolate
/// the corresponding lowering and lowering configs.
///
-/// TODO: Ideally, we should consider fusion opportunities by moving
-/// pack/unpack/mmt4d ops into one function:
-/// * https://github.com/llvm/llvm-project/issues/159770
/// TODO: Vectorize linalg.pack + linalg.unpack:
/// * https://github.com/llvm/llvm-project/issues/159751
//===----------------------------------------------------------------------===//
>From 19b681bc5ccaf1ebb04aafa23089126556cf634d Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Tue, 14 Jul 2026 16:19:44 -0400
Subject: [PATCH 17/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index a520816bce2f0..a516148c80a7e 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -32,7 +32,7 @@
/// * The matrix-multiplication dimensions that are scalable: M and N.
///
/// 2. The lowering of linalg.mmt4d leverages scalable vectorisation.
-/// * The matrix-multiplication dimension that's scalable: M, N (to match data
+/// * The matrix-multiplication dimensions that are scalable: M, N (to match data
/// tiling configuration).
///
/// 3. Neither `linalg.pack` nor `linalg.unpack` are vectorised ATM.
>From 9a5f92cacfa68e7a5bc30e9022a1bd51086b34f1 Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Tue, 14 Jul 2026 16:20:26 -0400
Subject: [PATCH 18/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index a516148c80a7e..7923ff09aa9d9 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -171,7 +171,7 @@ func.func private @pack_rhs(%B: tensor<16x13xf32>) -> tensor<?x16x?x1xf32> {
// @pack_acc
//
// Implements packing for the C matrix (accumulator) in matrix multiplication.
-// The inner tile size is "scalable": 8 * vscale
+// The inner tile sizes are "scalable": 8 * vscale, 8 * vscale
//===----------------------------------------------------------------------===//
func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<?x?x?x?xf32> {
%pad = arith.constant 0.0 : f32
>From 5a3b8546c149ec007a5aa673b477d0937b1a0ffe Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Tue, 14 Jul 2026 16:20:44 -0400
Subject: [PATCH 19/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 7923ff09aa9d9..78bcd5cf63ce5 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -199,7 +199,7 @@ func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<?x?x?x?xf32> {
// @unpack_acc
//
// Implements unpacking for the C matrix (accumulator) in matrix
-// multiplication. The inner tile size is "scalable": 8 * vscale
+// multiplication. The inner tile sizes are "scalable": 8 * vscale, 8 * vscale
//===----------------------------------------------------------------------===//
func.func private @unpack_acc(%C_packed: tensor<?x?x?x?xf32>) -> tensor<7x13xf32> {
%vs = vector.vscale
>From 70243043c22de9587bbf2f9e506f9ded83b71169 Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Tue, 14 Jul 2026 16:21:01 -0400
Subject: [PATCH 20/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 78bcd5cf63ce5..34abbdf31ab47 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -121,7 +121,7 @@ func.func private @matmul(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>, %C: tenso
// @pack_lhs
//
// Implements packing for the A matrix (LHS) in matrix multiplication. The
-// inner tile size is "scalable": 8 * vscale.
+// inner tile size for dim M is "scalable": 8 * vscale.
//===----------------------------------------------------------------------===//
func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x?x1xf32> {
%pad = arith.constant 0.0 : f32
>From 45937bcf98c6ab252f023657ad2a2eadafbca90e Mon Sep 17 00:00:00 2001
From: Stephen Long <63318318+steplong at users.noreply.github.com>
Date: Tue, 14 Jul 2026 16:21:19 -0400
Subject: [PATCH 21/23] Update
mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 34abbdf31ab47..27322db143b30 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -144,7 +144,7 @@ func.func private @pack_lhs(%A: tensor<7x16xf32>) -> tensor<1x16x?x1xf32> {
// @pack_rhs
//
// Implements packing for the B matrix (RHS) in matrix multiplication. The
-// inner tile size is "scalable": 8 * vscale.
+// inner tile size for dim N is "scalable": 8 * vscale.
//===----------------------------------------------------------------------===//
func.func private @pack_rhs(%B: tensor<16x13xf32>) -> tensor<?x16x?x1xf32> {
%pad = arith.constant 0.0 : f32
>From 6466c104f84c5cf67df73361dce2ab1afc055cf8 Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Fri, 17 Jul 2026 12:37:19 -0700
Subject: [PATCH 22/23] Add comment about create_named_contraction
---
.../Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 3 +++
1 file changed, 3 insertions(+)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index 27322db143b30..ed8c12992b3a2 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -276,6 +276,9 @@ module @transforms attributes { transform.with_named_sequence } {
transform.apply_patterns.canonicalization
} : !transform.any_op
transform.apply_patterns to %func_pre {
+ // TODO: `create_named_contraction` above lowers directly to a named contraction. If we want to remove
+ // `transform.apply_patterns.vector.reduction_to_contract` below, `CombineContractBroadcastMask`
+ // will have to be moved to a dedicated TD Op.
transform.apply_patterns.vector.reduction_to_contract
transform.apply_patterns.vector.cast_away_vector_leading_one_dim
transform.apply_patterns.vector.lower_contraction
>From 55a72ad0312bab0f6f6caafad91fb19c1f38c3ef Mon Sep 17 00:00:00 2001
From: Stephen Long <steplong at quicinc.com>
Date: Fri, 17 Jul 2026 13:55:58 -0700
Subject: [PATCH 23/23] Fix up dims of pack_lhs, pack_rhs, pack_acc, unpack_acc
---
.../Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir | 24 +++++++++----------
1 file changed, 11 insertions(+), 13 deletions(-)
diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
index ed8c12992b3a2..73e286f960008 100644
--- a/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
+++ b/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/pack-unpack-mmt4d.mlir
@@ -173,26 +173,24 @@ func.func private @pack_rhs(%B: tensor<16x13xf32>) -> tensor<?x16x?x1xf32> {
// Implements packing for the C matrix (accumulator) in matrix multiplication.
// The inner tile sizes are "scalable": 8 * vscale, 8 * vscale
//===----------------------------------------------------------------------===//
-func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<?x?x?x?xf32> {
+func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<1x?x?x?xf32> {
%pad = arith.constant 0.0 : f32
- // Compute the outer tile sizes.
- %c7 = arith.constant 7 : index
+ // Compute the outer tile size.
%c13 = arith.constant 13 : index
%vs = vector.vscale
%c8 = arith.constant 8 : index
%vs_c8 = arith.muli %vs, %c8 : index
- %outer_tile_size_0 = arith.ceildivui %c7, %vs_c8 : index
- %outer_tile_size_1 = arith.ceildivui %c13, %vs_c8 : index
+ %outer_tile_size = arith.ceildivui %c13, %vs_c8 : index
- %C_pack_empty = tensor.empty(%outer_tile_size_0, %outer_tile_size_1, %vs_c8, %vs_c8) : tensor<?x?x?x?xf32>
+ %C_pack_empty = tensor.empty(%outer_tile_size, %vs_c8, %vs_c8) : tensor<1x?x?x?xf32>
%C_pack = linalg.pack %C
padding_value(%pad : f32)
outer_dims_perm = [0, 1]
inner_dims_pos = [0, 1]
- inner_tiles = [%vs_c8, %vs_c8] into %C_pack_empty : tensor<7x13xf32> -> tensor<?x?x?x?xf32>
+ inner_tiles = [%vs_c8, %vs_c8] into %C_pack_empty : tensor<7x13xf32> -> tensor<1x?x?x?xf32>
- return %C_pack : tensor<?x?x?x?xf32>
+ return %C_pack : tensor<1x?x?x?xf32>
}
//===----------------------------------------------------------------------===//
@@ -201,7 +199,7 @@ func.func private @pack_acc(%C: tensor<7x13xf32>) -> tensor<?x?x?x?xf32> {
// Implements unpacking for the C matrix (accumulator) in matrix
// multiplication. The inner tile sizes are "scalable": 8 * vscale, 8 * vscale
//===----------------------------------------------------------------------===//
-func.func private @unpack_acc(%C_packed: tensor<?x?x?x?xf32>) -> tensor<7x13xf32> {
+func.func private @unpack_acc(%C_packed: tensor<1x?x?x?xf32>) -> tensor<7x13xf32> {
%vs = vector.vscale
%c8 = arith.constant 8 : index
%vs_c8 = arith.muli %vs, %c8 : index
@@ -211,7 +209,7 @@ func.func private @unpack_acc(%C_packed: tensor<?x?x?x?xf32>) -> tensor<7x13xf32
outer_dims_perm = [0, 1]
inner_dims_pos = [0, 1]
inner_tiles = [%vs_c8, %vs_c8]
- into %C_out_empty : tensor<?x?x?x?xf32> -> tensor<7x13xf32>
+ into %C_out_empty : tensor<1x?x?x?xf32> -> tensor<7x13xf32>
return %C_out_unpack: tensor<7x13xf32>
}
@@ -225,13 +223,13 @@ func.func private @matmul_via_mmt4d(%A: tensor<7x16xf32>, %B: tensor<16x13xf32>,
// Pack input matrices
%A_pack = func.call @pack_lhs(%A): (tensor<7x16xf32>) -> tensor<1x16x?x1xf32>
%B_pack = func.call @pack_rhs(%B): (tensor<16x13xf32>) -> tensor<?x16x?x1xf32>
- %C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<?x?x?x?xf32>
+ %C_pack = func.call @pack_acc(%C): (tensor<7x13xf32>) -> tensor<1x?x?x?xf32>
// MMT4D
- %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x?x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
+ %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x?x1xf32>, tensor<?x16x?x1xf32>) outs(%C_pack : tensor<1x?x?x?xf32>) -> tensor<1x?x?x?xf32>
// Unpack the output
- %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<?x?x?x?xf32>) -> tensor<7x13xf32>
+ %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x?x?xf32>) -> tensor<7x13xf32>
return %C_out_unpack : tensor<7x13xf32>
}
More information about the Mlir-commits
mailing list