[Mlir-commits] [mlir] [mlir][ArmSME] Skip tile allocation for functions with no SME tile values (PR #227153)
Federico Bruzzone
llvmlistbot at llvm.org
Tue Sep 29 01:46:51 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/227153
>From 280d51891359fcdd3b4ada17dfefb11e88a28437 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 29 Sep 2026 01:04:07 +0200
Subject: [PATCH 1/2] [mlir][ArmSME] Skip tile allocation for functions with no
SME tile values
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmSME/Transforms/TileAllocation.cpp | 39 +++++++++++++++++--
.../Dialect/ArmSME/basic-tile-allocation.mlir | 24 ++++++++++++
.../ArmSME/tile-allocation-liveness.mlir | 18 +++++++++
3 files changed, 78 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index b9b2cecf8fd00..0a73ac32e2b47 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -800,12 +800,45 @@ struct TestTileAllocationPass
};
} // namespace
+/// Returns true if `function` contains any SME tile values, either as block
+/// arguments (e.g., function arguments, or arguments to blocks merging values
+/// from different paths) or as results of operations.
+static bool functionHasSMETileValues(FunctionOpInterface function) {
+ auto hasSMETileType = [](TypeRange types) {
+ return llvm::any_of(
+ types, [](Type type) { return isValidSMETileVectorType(type); });
+ };
+
+ bool hasSMETileValues = false;
+ function.walk([&](Block *block) -> WalkResult {
+ if (hasSMETileType(block->getArgumentTypes())) {
+ hasSMETileValues = true;
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ });
+ if (hasSMETileValues)
+ return true;
+ function.walk([&](Operation *op) -> WalkResult {
+ if (hasSMETileType(op->getResultTypes())) {
+ hasSMETileValues = true;
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ });
+ return hasSMETileValues;
+}
+
LogicalResult mlir::arm_sme::allocateSMETiles(FunctionOpInterface function,
bool dumpRanges) {
- if (function.empty()) {
- // TODO: Also return early if the function contains no ArmSME ops?
+ if (function.empty())
+ return success();
+
+ // Bail out early if the function has no SME tile values, this avoids
+ // running the (non-trivial) preprocessing and liveness analysis below
+ // for functions that have nothing to allocate.
+ if (!functionHasSMETileValues(function))
return success();
- }
LiveRange::Allocator liveRangeAllocator;
IRRewriter rewriter(function.getContext());
diff --git a/mlir/test/Dialect/ArmSME/basic-tile-allocation.mlir b/mlir/test/Dialect/ArmSME/basic-tile-allocation.mlir
index 8b46998d56b04..fbf5754f24417 100644
--- a/mlir/test/Dialect/ArmSME/basic-tile-allocation.mlir
+++ b/mlir/test/Dialect/ArmSME/basic-tile-allocation.mlir
@@ -404,3 +404,27 @@ func.func @za_q() {
"test.some_use"(%next_tile) : (vector<[1]x[1]xi128>) -> ()
return
}
+
+// -----
+
+// Functions with no ArmSME tile ops should be left completely unmodified.
+// This also means the preprocessing steps must not run.
+
+// CHECK-LABEL: func.func @no_sme_ops(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[A:.*]]: i32, %[[B:.*]]: i32) -> i32 {
+// CHECK-NEXT: cf.cond_br %[[COND]], ^bb1, ^bb2
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: cf.br ^bb3(%[[A]] : i32)
+// CHECK-NEXT: ^bb2:
+// CHECK-NEXT: cf.br ^bb3(%[[B]] : i32)
+// CHECK-NEXT: ^bb3(%[[R:.*]]: i32):
+// CHECK-NEXT: return %[[R]] : i32
+func.func @no_sme_ops(%cond: i1, %a: i32, %b: i32) -> i32 {
+ cf.cond_br %cond, ^bb1, ^bb2
+^bb1:
+ cf.br ^bb3(%a : i32)
+^bb2:
+ cf.br ^bb3(%b : i32)
+^bb3(%r: i32):
+ return %r : i32
+}
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
index 0b739feaf019d..f17e5d0aa3f0f 100644
--- a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
@@ -608,3 +608,21 @@ func.func @reactivate_inactive_live_range(%cond: i1) {
^bb3:
return
}
+
+// -----
+
+// A function argument that is a valid SME tile type must be processed
+// for tile allocation, even though it is not produced by an ArmSME op.
+
+// CHECK-LIVE-RANGE-LABEL: @tile_value_from_function_argument
+// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
+// CHECK-LIVE-RANGE: ^bb0:
+// CHECK-LIVE-RANGE-NEXT: E test.some_use
+
+// CHECK-LABEL: @tile_value_from_function_argument(
+// CHECK-SAME: %[[TILE:.*]]: vector<[4]x[4]xf32>
+func.func @tile_value_from_function_argument(%tile: vector<[4]x[4]xf32>) {
+ // CHECK: "test.some_use"(%[[TILE]])
+ "test.some_use"(%tile) : (vector<[4]x[4]xf32>) -> ()
+ return
+}
\ No newline at end of file
>From 8bd0c5d31cca6b657ba21cb9755a438f01c62184 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Tue, 29 Sep 2026 10:46:17 +0200
Subject: [PATCH 2/2] Address comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmSME/Transforms/TileAllocation.cpp | 34 +++++++------------
1 file changed, 13 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 0a73ac32e2b47..c9312462f9831 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -800,33 +800,25 @@ struct TestTileAllocationPass
};
} // namespace
-/// Returns true if `function` contains any SME tile values, either as block
-/// arguments (e.g., function arguments, or arguments to blocks merging values
-/// from different paths) or as results of operations.
+/// Returns true if `function` contains any SME tile values, either as
+/// function arguments or as results of operations.
+///
+/// NOTE: Block arguments of non-entry blocks do not need to be checked
+/// separately, as they are always fed by branch operands that trace back to
+/// either a function argument or an operation result.
static bool functionHasSMETileValues(FunctionOpInterface function) {
auto hasSMETileType = [](TypeRange types) {
return llvm::any_of(
types, [](Type type) { return isValidSMETileVectorType(type); });
};
-
- bool hasSMETileValues = false;
- function.walk([&](Block *block) -> WalkResult {
- if (hasSMETileType(block->getArgumentTypes())) {
- hasSMETileValues = true;
- return WalkResult::interrupt();
- }
- return WalkResult::advance();
- });
- if (hasSMETileValues)
+ if (hasSMETileType(function.getArgumentTypes()))
return true;
- function.walk([&](Operation *op) -> WalkResult {
- if (hasSMETileType(op->getResultTypes())) {
- hasSMETileValues = true;
- return WalkResult::interrupt();
- }
- return WalkResult::advance();
- });
- return hasSMETileValues;
+ return function
+ .walk([&](Operation *op) {
+ return hasSMETileType(op->getResultTypes()) ? WalkResult::interrupt()
+ : WalkResult::advance();
+ })
+ .wasInterrupted();
}
LogicalResult mlir::arm_sme::allocateSMETiles(FunctionOpInterface function,
More information about the Mlir-commits
mailing list