[LLVMCPU] Fix scalable vectorization fallback on SME-only (no `+sve`) targets (#24693)
**Stacked** on #24656. Follow-up to #24661 / #24660.
This PR aims to removes the TODO introduced in #24656.
#24661 fixed the ArmSME lowering pipeline to not require classic `+sve`
when SME tiling is used. This PR fixes the same underlying problem for
the *fallback* path, which is what's used when the matmul falls back to
it (e.g., `--iree-llvmcpu-disable-arm-sme-tiling`, or op types SME
doesn't support): compiling a scalable-vectorized matmul on a
`+sme`-only target with SME tiling disabled crashed LLVM instruction
selection with `Cannot select: vscale`.
The fix is scoped to a single function:
`getMatmulVectorSizesUsingFillRegisterFileHeuristic` only checked
`hasAnySVEFeature`, so on a `+sme`-only target it always fell back to
non-scalable NEON tiles. It now also picks scalable tiles when `+sme` is
present *and* the user has explicitly passed
`--iree-llvmcpu-force-arm-streaming`.
Net effect:
1. Non-matmul ops: always Peeling.
4. `linalg.matmul`, `+sme-only`, `disable-arm-sme-tiling=true`, and
`force-arm-streaming=true`: still reaches SSVE (that's the actual fix),
but now goes through `Peeling + masked tail` instead of masking-only (as
in #24661).
Unrelated to #24689
---------
Signed-off-by: Federico Bruzzone <federico.bruzzone.i@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
index 83c8444..59e5407 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -1372,7 +1372,9 @@
auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(op);
scalableSizeFlags.resize(3, false);
if (isScalableVectorizationEnabled() &&
- hasAnySVEFeature(targetAttr.getConfiguration())) {
+ (hasAnySVEFeature(targetAttr.getConfiguration()) ||
+ (hasSMEFeature(targetAttr.getConfiguration()) &&
+ isArmStreamingForced()))) {
scalableSizeFlags[1] = true;
}
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index 452f2b8..91fa67f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -11,6 +11,7 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "iree/compiler/Codegen/LLVMCPU/Passes.h"
+#include "iree/compiler/Codegen/Utils/CPUUtils.h"
#include "iree/compiler/Codegen/Utils/CodegenOptions.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h"
#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
@@ -60,19 +61,6 @@
"LLVMCPUMmt4dVectorLowering pass."),
llvm::cl::init(false), llvm::cl::Hidden);
-// By default, IREE does not enable the Armv9-A streaming SVE mode in the
-// presence of scalable vectors (even when using `+sme`), as currently there's
-// no cost model of when it could be beneficial. This flag will effectively make
-// IREE/LLVM switch from SVE to SSVE in dispatch regions with supported
-// scalable vector operations.
-static llvm::cl::opt<bool> clForceArmStreaming(
- "iree-llvmcpu-force-arm-streaming",
- llvm::cl::desc(
- "Enables Armv9-A streaming SVE mode for any dispatch region that "
- "contains supported scalable vector operations (i.e., use SSVE rather "
- "than SVE). Requires the +sme feature flag."),
- llvm::cl::init(false), llvm::cl::Hidden);
-
static llvm::cl::opt<bool> clPatchFuncOps(
"iree-llvmcpu-debug-patch-func-ops",
llvm::cl::desc(
@@ -548,7 +536,7 @@
modulePassManager.addPass(mlir::arm_sme::createVectorLegalizationPass());
FunctionLikeNest(modulePassManager)
.addPredicatedPass(
- clForceArmStreaming,
+ isArmStreamingForced(),
[] {
// 1. Enable Armv9-A streaming mode without ZA (i.e., SSVE) for
// dispatch regions that contain scalable vectors when forced via
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
index 840d702..cbbc1be 100644
--- 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
@@ -1,21 +1,24 @@
// 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
+// RUN: iree-opt --pass-pipeline='builtin.module(iree-llvmcpu-select-lowering-strategy)' --iree-llvmcpu-enable-scalable-vectorization=true --iree-llvmcpu-disable-arm-sme-tiling --iree-llvmcpu-force-arm-streaming --split-input-file %s | FileCheck %s --check-prefixes=CHECK,FORCE-STREAMING
-// 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.
+// A target with +sme but not +sve. Checks the lowering strategy selected for
+// a matmul when SME tiling is disabled, with and without
+// --iree-llvmcpu-force-arm-streaming.
#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).
+// SME tiling disabled, streaming not forced: falls back to non-scalable
+// NEON tile sizes (the default).
// 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.
+// SME tiling enabled: picks scalable M/N tile sizes regardless of
+// whether streaming is forced.
// 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]>
+// SME tiling disabled, streaming forced explicitly: falls back to the
+// regular SVE-style tiling heuristic, scalable on the N dimension only.
+// FORCE-STREAMING-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]>
// 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]]
diff --git a/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.cpp
index 9b6d3ef..f5cb417 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.cpp
@@ -27,6 +27,19 @@
"target (e.g., +sve, +sve2 and/or +sme feature flags)"),
llvm::cl::init(false));
+// By default, IREE does not enable the Armv9-A streaming SVE mode in the
+// presence of scalable vectors (even when using `+sme`), as currently there's
+// no cost model of when it could be beneficial. This flag will effectively make
+// IREE/LLVM switch from SVE to SSVE in dispatch regions with supported
+// scalable vector operations.
+static llvm::cl::opt<bool> clForceArmStreaming(
+ "iree-llvmcpu-force-arm-streaming",
+ llvm::cl::desc(
+ "Enables Armv9-A streaming SVE mode for any dispatch region that "
+ "contains supported scalable vector operations (i.e., use SSVE rather "
+ "than SVE). Requires the +sme feature flag."),
+ llvm::cl::init(false), llvm::cl::Hidden);
+
static llvm::cl::opt<int> clVscaleFromUser(
"iree-experimental-vscale-value",
llvm::cl::desc(
@@ -112,6 +125,8 @@
bool isScalableVectorizationEnabled() { return clEnableScalableVectorization; }
+bool isArmStreamingForced() { return clForceArmStreaming; }
+
unsigned getUserVscaleValue() {
assert(clVscaleFromUser >= 1 && "Currently, vscale needs to be specified by "
"the user for host-side code!");
diff --git a/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.h b/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.h
index 9ccc64c..cc8bba2 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/CPUUtils.h
@@ -39,6 +39,11 @@
/// Returns if scalable vectorization is enabled or not.
bool isScalableVectorizationEnabled();
+/// Returns whether Armv9-A streaming SVE mode is forced for dispatch regions
+/// containing scalable vector operations, via
+/// `--iree-llvmcpu-force-arm-streaming`.
+bool isArmStreamingForced();
+
/// Returns the runtime value of vscale specified by the user. This is not meant
/// to be used for codegen and is meant to circumvent the current limitation on
/// host-side querying of this value at runtime. This is temporary until #21317
