[LLVMCPU] Fix ArmSME requiring classic SVE, breaking SME-only targets (e.g., Apple Silicon) (#24661)
`enableAArch64SME` required `hasAnySVEFeature` in addition to
`hasSMEFeature`, even though SME runs in Streaming SVE mode and does not
architecturally need classic (non-streaming) SVE. On SME-only targets
(e.g. Apple Silicon, which implements `FEAT_SME`/`FEAT_SME2` but not
`FEAT_SVE`), this meant:
- **`+sme` alone**: `enableAArch64SME` was `false`, so the ArmSME
lowering pipeline never ran. Vectorization silently fell back to scalar
code, but a leftover `vscale` was never inside a function marked
"streaming", so it had no legal instruction-selection pattern:
`iree-compile` crashed with `LLVM ERROR: Cannot select: vscale`.
- **`+sve,+sme` together** (the only combination the existing SME e2e
tests use): compiles fine, but a separate, Darwin-specific mechanism in
`AArch64FrameLowering`/`AArch64PrologueEpilogue` is gated on `hasSVE()`.
With `+sve` set, LLVM emits `cntd` directly in the function prologue
*before* `smstart`, which is illegal on hardware without non-streaming
SVE and crashes at runtime with `SIGILL`.
Neither combination worked; does that mean that, prior to this PR, it
wasn't possible to use IREE's SME support on this class of hardware? If
so, it implies that, until now, IREE couldn't leverage SME on Apple M4
Pro processors.
But I feel like I might be missing something...
Fixes #24660.
See: https://github.com/llvm/llvm-project/issues/204853
@banach-space @egebeysel @AGindinson @maxbartel
---------
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/LLVMCPULowerExecutableTarget.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
index 5ca501e..edb73bd 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerExecutableTarget.cpp
@@ -144,7 +144,8 @@
pipelineOpts.lowerToAVX2 = hasAVX2Feature(targetConfig);
pipelineOpts.enableVectorMasking =
isX86(targetConfig) || isRISCV(targetConfig) ||
- (isAArch64(targetConfig) && hasAnySVEFeature(targetConfig));
+ (isAArch64(targetConfig) &&
+ (hasAnySVEFeature(targetConfig) || hasSMEFeature(targetConfig)));
// TODO(#16956): The decomposition of attention op leads to complex control
// flow, which leads to non-trivial stack allocation. Enforcing masking for
// targets that do not support native vector masking enables the
@@ -160,9 +161,11 @@
if (isAArch64(targetConfig) && hasAttentionOp(funcOp)) {
pipelineOpts.enableVectorMasking = true;
}
- pipelineOpts.enableAArch64SME = isAArch64(targetConfig) &&
- hasAnySVEFeature(targetConfig) &&
- hasSMEFeature(targetConfig);
+
+ // SME does not require classic (non-streaming) SVE: some targets
+ // (e.g., Apple Silicon) implement SME without it.
+ pipelineOpts.enableAArch64SME =
+ isAArch64(targetConfig) && hasSMEFeature(targetConfig);
pipelineOpts.enableAArch64I8mm =
isAArch64(targetConfig) && hasI8mmFeature(targetConfig);
pipelineOpts.enablePeeling = isOptEnabled(funcOp, getEnableLoopPeelingStr());
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_vector_masking_tests.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_vector_masking_tests.mlir
index e084135..5062610 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_vector_masking_tests.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_vector_masking_tests.mlir
@@ -234,3 +234,53 @@
// CHECK: scf.for
// CHECK-COUNT-2: vector.maskedload
// CHECK: vector.maskedstore
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<constants = 2, bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+#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"}>
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// This is @mask_dynamic_generic_add above, verbatim, with cpu_features
+// "+sve" swapped for "+sme": see the comment below for why it must still
+// mask.
+func.func @mask_dynamic_generic_add_sme() attributes {hal.executable.target = #executable_target_embedded_elf_arm_64_} {
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : i32
+ %1 = hal.interface.constant.load layout(#pipeline_layout) ordinal(1) : i32
+ %2 = arith.index_cast %0 : i32 to index
+ %3 = arith.index_cast %1 : i32 to index
+ %dim0 = iree_tensor_ext.dispatch.workload.ordinal %2, 0 : index
+ %dim1 = iree_tensor_ext.dispatch.workload.ordinal %3, 1 : index
+ %4 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32>>{%dim0, %dim1}
+ %5 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32>>{%dim0, %dim1}
+ %6 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<?x?xf32>>{%dim0, %dim1}
+ %7 = iree_tensor_ext.dispatch.tensor.load %4, offsets = [0, 0], sizes = [%dim0, %dim1], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32>>{%dim0, %dim1} -> tensor<?x?xf32>
+ %8 = iree_tensor_ext.dispatch.tensor.load %5, offsets = [0, 0], sizes = [%dim0, %dim1], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32>>{%dim0, %dim1} -> tensor<?x?xf32>
+ %9 = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>
+ %10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %11 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%7, %8 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%10 : tensor<?x?xf32>) {
+ ^bb0(%in: f32, %in_0: f32, %out: f32):
+ %12 = arith.addf %in, %in_0 : f32
+ linalg.yield %12 : f32
+ } -> tensor<?x?xf32>
+ iree_tensor_ext.dispatch.tensor.store %11, %6, offsets = [0, 0], sizes = [%dim0, %dim1], strides = [1, 1] : tensor<?x?xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<?x?xf32>>{%dim0, %dim1}
+ return
+}
+
+// SME does not require classic (non-streaming) SVE: masking must
+// still be enabled for SME-only targets (e.g. Apple Silicon), same as for SVE.
+// The vectorizer picks the peeling strategy here (SME's own 2D-tiling
+// heuristics are matmul-specific and don't apply to this elementwise op), so
+// masking shows up on the peeled tail loop, same shape as the SVE case above.
+
+// CHECK-LABEL: func.func @mask_dynamic_generic_add_sme
+// Main loop
+// CHECK: scf.for
+// Peeled loop:
+// CHECK: scf.for
+// CHECK-COUNT-2: vector.maskedload
+// CHECK: vector.maskedstore
diff --git a/runtime/src/iree/builtins/device/arch/arm_64/sme_abi.c b/runtime/src/iree/builtins/device/arch/arm_64/sme_abi.c
index 86c4a86..3acc8ab 100644
--- a/runtime/src/iree/builtins/device/arch/arm_64/sme_abi.c
+++ b/runtime/src/iree/builtins/device/arch/arm_64/sme_abi.c
@@ -43,3 +43,20 @@
IREE_DEVICE_EXPORT void __arm_za_disable(void) {
// No-op as not needed yet.
}
+
+// Unlike the stubs above, this one must return a real value: it's called
+// from the prologue, *before* `smstart`, to save the pre-streaming VG for
+// unwind info on targets lacking classic SVE, where `cntd` is illegal outside
+// Streaming SVE mode (e.g., Apple Silicon). LLVM only generates calls to this
+// function on platforms that lack SVE, i.e. for `!hasSVE && hasSME`
+// (otherwise it inlines `cntd` directly, which is always legal when SVE is
+// present). A normal host link gets this for free from compiler-rt's
+// sme-abi.S. IREE's embedded linking is freestanding and has none, hence the
+// local copy.
+// See https://github.com/llvm/llvm-project/issues/204853.
+//
+// ASSUMPTION: In IREE we only use `arm_locally_streaming` functions. That
+// means we can assume `__arm_get_current_vg` will be called outside of
+// streaming mode. So by assuming `!hasSVE && hasSME` is safe just returning
+// 0.
+IREE_DEVICE_EXPORT int64_t __arm_get_current_vg(void) { return 0; }
diff --git a/tests/e2e/matmul/BUILD.bazel b/tests/e2e/matmul/BUILD.bazel
index 872298a..cf42715 100644
--- a/tests/e2e/matmul/BUILD.bazel
+++ b/tests/e2e/matmul/BUILD.bazel
@@ -96,7 +96,7 @@
target_backends_and_drivers = [
("llvm-cpu", "local-task"),
],
- target_cpu_features_variants = ["arm_64:sme:+sve,+sme"],
+ target_cpu_features_variants = ["arm_64:sme:+sme"],
test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
test_type = "matmul",
) for dtype in [
diff --git a/tests/e2e/matmul/CMakeLists.txt b/tests/e2e/matmul/CMakeLists.txt
index e31ef11..40f516d 100644
--- a/tests/e2e/matmul/CMakeLists.txt
+++ b/tests/e2e/matmul/CMakeLists.txt
@@ -35,7 +35,7 @@
LABELS
"requires-arm-sme"
TARGET_CPU_FEATURES_VARIANTS
- "arm_64:sme:+sve,+sme"
+ "arm_64:sme:+sme"
)
iree_generated_e2e_runner_test(
@@ -62,7 +62,7 @@
LABELS
"requires-arm-sme"
TARGET_CPU_FEATURES_VARIANTS
- "arm_64:sme:+sve,+sme"
+ "arm_64:sme:+sme"
)
iree_generated_e2e_runner_test(
@@ -89,7 +89,7 @@
LABELS
"requires-arm-sme"
TARGET_CPU_FEATURES_VARIANTS
- "arm_64:sme:+sve,+sme"
+ "arm_64:sme:+sme"
)
iree_generated_e2e_runner_test(
@@ -115,7 +115,7 @@
LABELS
"requires-arm-sme"
TARGET_CPU_FEATURES_VARIANTS
- "arm_64:sme:+sve,+sme"
+ "arm_64:sme:+sme"
)
iree_generated_e2e_runner_test(
