[LLVMCPU] Add SME lowering-strategy tests for f64 and unsupported i8 matmuls (#24656)
`getMatmulAArch64SMEVectorSizes` (within `KernelDispatch.cpp`) only had
test coverage for the f32 matmul case.
This PR add two tests:
- `f64` which uses `[4]x[8]` SME tiles instead of the `[8]x[8]` used for
f32.
https://github.com/iree-org/iree/blob/50d0addfb49254e0b7b746959b360c79437281cc/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp#L1465-L1471
- `i8` which SME does not support
https://github.com/iree-org/iree/blob/50d0addfb49254e0b7b746959b360c79437281cc/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp#L1488-L1502
_AI disclaimer_: I used `Claude:Sonnet5` to write the checks.
---------
Signed-off-by: Federico Bruzzone <federico.bruzzone.i@gmail.com>
Co-authored-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
index 12f8285..d39c820 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel
@@ -56,6 +56,7 @@
"scalable_tile_and_vectorize_matmul.mlir",
"select_aarch64_lowering_strategy.mlir",
"select_aarch64_sme_lowering_strategy.mlir",
+ "select_aarch64_ssve_lowering_strategy.mlir",
"select_aarch64_sve_lowering_strategy.mlir",
"select_aarch64_sve_lowering_strategy_peeling.mlir",
"select_riscv_lowering_strategy.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
index e8d26a0..abc526f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt
@@ -51,6 +51,7 @@
"scalable_tile_and_vectorize_matmul.mlir"
"select_aarch64_lowering_strategy.mlir"
"select_aarch64_sme_lowering_strategy.mlir"
+ "select_aarch64_ssve_lowering_strategy.mlir"
"select_aarch64_sve_lowering_strategy.mlir"
"select_aarch64_sve_lowering_strategy_peeling.mlir"
"select_riscv_lowering_strategy.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sme_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sme_lowering_strategy.mlir
index 9b20c0e..2d05365 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sme_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sme_lowering_strategy.mlir
@@ -1,4 +1,5 @@
-// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file %s | FileCheck %s
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file %s | FileCheck %s --check-prefixes=CHECK,WITH-SME
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' --iree-llvmcpu-enable-scalable-vectorization=true --iree-llvmcpu-disable-arm-sme-tiling --split-input-file %s | FileCheck %s --check-prefixes=CHECK,DISABLE-ARM-SME
#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
func.func @transpose_f32(%2: tensor<32x32xf32>) -> tensor<32x32xf32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
@@ -9,8 +10,9 @@
} -> tensor<32x32xf32>
return %4 : tensor<32x32xf32>
}
-// CHECK: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [4, 32], vector_common_parallel = {{\[}}[4], [4]]>
-// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [4, 32], vector_common_parallel = {{\[}}[4], [4]]>
+// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [32, 32], vector_common_parallel = [4, 4]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
// CHECK: func.func @transpose_f32(
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
@@ -27,8 +29,9 @@
} -> tensor<32x32xf32>
return %4 : tensor<32x32xf32>
}
-// CHECK: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [4, 32], vector_common_parallel = {{\[}}[4], [4]]>
-// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [4, 32], vector_common_parallel = {{\[}}[4], [4]]>
+// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [32, 32], vector_common_parallel = [4, 4]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
// CHECK: func.func @transpose_output_indexing_map_f32(
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
@@ -45,8 +48,9 @@
} -> tensor<32x32xf64>
return %4 : tensor<32x32xf64>
}
-// CHECK: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [4, 32], vector_common_parallel = {{\[}}[2], [2]]>
-// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [4, 32], vector_common_parallel = {{\[}}[2], [2]]>
+// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [32, 32], vector_common_parallel = [2, 2]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
// CHECK: func.func @transpose_f64(
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
@@ -87,3 +91,52 @@
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+// -----
+
+#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
+func.func @matmul_tensors(%7: tensor<?x?xf32>, %8: tensor<?x?xf32>, %9: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %10 = linalg.matmul ins(%7, %8 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%9 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %10 : tensor<?x?xf32>
+}
+// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = [8, [8], 0], vector_reduction = [0, 0, 4]>
+// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = {{\[}}[8], [8], 0], vector_reduction = [0, 0, 1]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// CHECK: func.func @matmul_tensors(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+// -----
+
+// f64 matmul uses different SME tile sizes than f32: [4]x[8] instead of [8]x[8]
+// (each SME tile only holds half as many f64 elements per row as f32 ones).
+#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
+func.func @matmul_tensors_f64(%7: tensor<?x?xf64>, %8: tensor<?x?xf64>, %9: tensor<?x?xf64>) -> tensor<?x?xf64> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %10 = linalg.matmul ins(%7, %8 : tensor<?x?xf64>, tensor<?x?xf64>) outs(%9 : tensor<?x?xf64>) -> tensor<?x?xf64>
+ return %10 : tensor<?x?xf64>
+}
+// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = [8, [4], 0], vector_reduction = [0, 0, 2]>
+// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = {{\[}}[4], [8], 0], vector_reduction = [0, 0, 1]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// CHECK: func.func @matmul_tensors_f64(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering_config = #[[CONFIG]]
+
+// -----
+
+// SME tiling only supports f32 and f64 matmuls, so an i8 matmul should fall
+// back to the regular SVE tiling heuristic regardless of whether SME tiling
+// is enabled or disabled.
+#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
+func.func @matmul_tensors_i8i8_i32_with_sme(%7: tensor<?x?xi8>, %8: tensor<?x?xi8>, %9: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %10 = linalg.matmul ins(%7, %8 : tensor<?x?xi8>, tensor<?x?xi8>) outs(%9 : tensor<?x?xi32>) -> tensor<?x?xi32>
+ return %10 : tensor<?x?xi32>
+}
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = [8, [8], 0], vector_reduction = [0, 0, 4]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// CHECK: func.func @matmul_tensors_i8i8_i32_with_sme(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering_config = #[[CONFIG]]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_ssve_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_ssve_lowering_strategy.mlir
new file mode 100644
index 0000000..840d702
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_ssve_lowering_strategy.mlir
@@ -0,0 +1,23 @@
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file %s | FileCheck %s --check-prefixes=CHECK,WITH-SME
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' --iree-llvmcpu-enable-scalable-vectorization=true --iree-llvmcpu-disable-arm-sme-tiling --split-input-file %s | FileCheck %s --check-prefixes=CHECK,DISABLE-ARM-SME
+
+// A target with +sme but not +sve (i.e. no non-streaming SVE) still picks
+// scalable SME tile sizes when SME tiling is enabled: matmul lowering does
+// not itself require +sve. Note this goes through a different
+// pre-processing strategy than the +sve,+sme case: peeling, with an extra
+// `cache_parallel` tiling level.
+#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
+func.func @matmul_tensors_sme_no_sve(%7: tensor<?x?xf32>, %8: tensor<?x?xf32>, %9: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %10 = linalg.matmul ins(%7, %8 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%9 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %10 : tensor<?x?xf32>
+}
+// SME tiling disabled: falls back to non-scalable NEON tile sizes, since the
+// SVE fallback heuristic requires +sve (see TODO above).
+// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<cache_parallel = [64, 64, 0], distribution = [64, 64, 0], vector_common_parallel = [8, 8, 0], vector_reduction = [0, 0, 4]>
+// SME tiling enabled: still picks scalable [8]x[8] SME tile sizes.
+// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<cache_parallel = [64, 64, 0], distribution = [64, 64, 0], vector_common_parallel = {{\[}}[8], [8], 0], vector_reduction = [0, 0, 1]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>, {enable_loop_peeling}>
+// CHECK: func.func @matmul_tensors_sme_no_sve(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering_config = #[[CONFIG]]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
index b5d9cf8..7f79c40 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
@@ -1,7 +1,5 @@
// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' \
-// RUN: --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file %s | FileCheck %s --check-prefixes=CHECK,WITH-SME
-// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' \
-// RUN: --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file --iree-llvmcpu-disable-arm-sme-tiling %s | FileCheck %s --check-prefixes=CHECK,DISABLE-ARM-SME
+// RUN: --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file %s | FileCheck %s
#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
func.func @matmul_tensors(%7: tensor<?x?xf32>, %8: tensor<?x?xf32>, %9: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
@@ -75,24 +73,18 @@
// -----
-#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
-func.func @matmul_tensors(%7: tensor<?x?xf32>, %8: tensor<?x?xf32>, %9: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
- %10 = linalg.matmul ins(%7, %8 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%9 : tensor<?x?xf32>) -> tensor<?x?xf32>
- return %10 : tensor<?x?xf32>
+// f64 matmul on an SVE-only target.
+#executable_target_embedded_elf_arm_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>
+func.func @matmul_tensors_f64(%7: tensor<?x?xf64>, %8: tensor<?x?xf64>, %9: tensor<?x?xf64>) -> tensor<?x?xf64> attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %10 = linalg.matmul ins(%7, %8 : tensor<?x?xf64>, tensor<?x?xf64>) outs(%9 : tensor<?x?xf64>) -> tensor<?x?xf64>
+ return %10 : tensor<?x?xf64>
}
-// DISABLE-ARM-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = [8, [8], 0], vector_reduction = [0, 0, 4]>
-// DISABLE-ARM-SME-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
-// DISABLE-ARM-SME: func.func @matmul_tensors(
-// DISABLE-ARM-SME-SAME: translation_info = #[[TRANSLATION]]
-// DISABLE-ARM-SME: linalg.matmul
-// DISABLE-ARM-SME-SAME: lowering_config = #[[CONFIG]]
-
-// WITH-SME-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = {{\[}}[8], [8], 0], vector_reduction = [0, 0, 1]>
-// WITH-SME-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
-// WITH-SME: func.func @matmul_tensors(
-// WITH-SME-SAME: translation_info = #[[TRANSLATION]]
-// WITH-SME: linalg.matmul
-// WITH-SME-SAME: lowering_config = #[[CONFIG]]
+// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 0], vector_common_parallel = [8, [4], 0], vector_reduction = [0, 0, 2]>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>>
+// CHECK: func.func @matmul_tensors_f64(
+// CHECK-SAME: translation_info = #[[TRANSLATION]]
+// CHECK: linalg.matmul
+// CHECK-SAME: lowering_config = #[[CONFIG]]
// -----
