[Mlir-commits] [mlir] [mlir][ArmSME] Skip tile allocation for functions with no SME tile ops (PR #227153)
Federico Bruzzone
llvmlistbot at llvm.org
Mon Oct 5 03:28:50 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/6] [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/6] 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,
>From 67183921549e0fe8a3a8e7b543a53927651b7d8d Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 30 Sep 2026 08:29:08 +0200
Subject: [PATCH 3/6] Address comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmSME/Transforms/TileAllocation.cpp | 26 ++++-----------
.../ArmSME/tile-allocation-liveness.mlir | 32 ++++++-------------
2 files changed, 16 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index c9312462f9831..4cfadbace6b83 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -800,24 +800,12 @@ struct TestTileAllocationPass
};
} // namespace
-/// 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); });
- };
- if (hasSMETileType(function.getArgumentTypes()))
- return true;
+/// Returns true if `function` contains any ArmSME tile ops. If it does not,
+/// there is nothing to allocate: any value not produced/consumed by such an
+/// op will never be assigned a tile ID, so tile allocation is a no-op.
+static bool functionHasArmSMETileOps(FunctionOpInterface function) {
return function
- .walk([&](Operation *op) {
- return hasSMETileType(op->getResultTypes()) ? WalkResult::interrupt()
- : WalkResult::advance();
- })
+ .walk([&](ArmSMETileOpInterface) { return WalkResult::interrupt(); })
.wasInterrupted();
}
@@ -826,10 +814,10 @@ LogicalResult mlir::arm_sme::allocateSMETiles(FunctionOpInterface function,
if (function.empty())
return success();
- // Bail out early if the function has no SME tile values, this avoids
+ // Bail out early if the function has no ArmSME tile ops, this avoids
// running the (non-trivial) preprocessing and liveness analysis below
// for functions that have nothing to allocate.
- if (!functionHasSMETileValues(function))
+ if (!functionHasArmSMETileOps(function))
return success();
LiveRange::Allocator liveRangeAllocator;
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
index f17e5d0aa3f0f..cddcf96feba8e 100644
--- a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
@@ -241,11 +241,15 @@ func.func @non_overlapping_branches(%cond: i1) {
// -----
// Here %vecA and %vecB are not merged into the same live range (as they are unknown values).
-// This means that %vecA and %vecB are both allocated to different tiles (which is not legal).
-
-// expected-note at below {{tile operand is: <block argument> of type 'vector<[4]x[4]xf32>'}}
-func.func @overlapping_branches(%cond: i1, %vecA: vector<[4]x[4]xf32>, %vecB: vector<[4]x[4]xf32>) {
- // expected-error at below {{op tile operand allocated to different SME virtial tile (move required)}}
+// This means that %vecA and %vecB are both allocated to different tiles. As `arm_sme.get_tile`
+// is trivially cloneable this is resolved by cloning it, rather than an error.
+
+// CHECK-LABEL: @overlapping_branches
+// CHECK: arm_sme.get_tile {tile_id = [[ID:.*]] : i32} : vector<[4]x[4]xf32>
+// CHECK: arm_sme.get_tile {tile_id = [[ID]] : i32} : vector<[4]x[4]xf32>
+func.func @overlapping_branches(%cond: i1) {
+ %vecA = arm_sme.get_tile : vector<[4]x[4]xf32>
+ %vecB = arm_sme.get_tile : vector<[4]x[4]xf32>
%tile = scf.if %cond -> vector<[4]x[4]xf32> {
scf.yield %vecA : vector<[4]x[4]xf32>
} else {
@@ -607,22 +611,4 @@ func.func @reactivate_inactive_live_range(%cond: i1) {
cf.br ^bb3
^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 2f4cb516b6f9d5f459dc514f37f74eaa00b438a8 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Wed, 30 Sep 2026 19:43:03 +0200
Subject: [PATCH 4/6] Address Andrzej & Benjamin suggestions
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmSME/Transforms/TileAllocation.cpp | 18 +++++++-----------
1 file changed, 7 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 4cfadbace6b83..21a7ef5d13258 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -800,10 +800,8 @@ struct TestTileAllocationPass
};
} // namespace
-/// Returns true if `function` contains any ArmSME tile ops. If it does not,
-/// there is nothing to allocate: any value not produced/consumed by such an
-/// op will never be assigned a tile ID, so tile allocation is a no-op.
-static bool functionHasArmSMETileOps(FunctionOpInterface function) {
+/// Returns true if `function` contains any ArmSME tile ops.
+static bool hasArmSMETileOps(FunctionOpInterface function) {
return function
.walk([&](ArmSMETileOpInterface) { return WalkResult::interrupt(); })
.wasInterrupted();
@@ -811,13 +809,11 @@ static bool functionHasArmSMETileOps(FunctionOpInterface function) {
LogicalResult mlir::arm_sme::allocateSMETiles(FunctionOpInterface function,
bool dumpRanges) {
- if (function.empty())
- return success();
-
- // Bail out early if the function has no ArmSME tile ops, this avoids
- // running the (non-trivial) preprocessing and liveness analysis below
- // for functions that have nothing to allocate.
- if (!functionHasArmSMETileOps(function))
+ // Bail out early if the function has no ArmSME tile ops: there is nothing
+ // to allocate (a value not produced/consumed by such an op never gets a
+ // tile ID), so this avoids the (non-trivial) preprocessing and liveness
+ // analysis below for functions that have nothing to allocate.
+ if (!hasArmSMETileOps(function))
return success();
LiveRange::Allocator liveRangeAllocator;
>From de6a470ab8ddb71a91fe172d663b14f564264433 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 5 Oct 2026 11:37:51 +0200
Subject: [PATCH 5/6] Add tile-allocation-move-required.mlir
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmSME/tile-allocation-move-required.mlir | 24 +++++++++++++++++++
1 file changed, 24 insertions(+)
create mode 100644 mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
new file mode 100644
index 0000000000000..56b6bda081ab4
--- /dev/null
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
@@ -0,0 +1,24 @@
+// RUN: not mlir-opt %s -convert-scf-to-cf -test-arm-sme-tile-allocation 2>&1 | FileCheck %s
+
+// `arm_sme.load_tile_slice` has operands, so it can't be cloned for free: merging two
+// independently-allocated tiles here is a genuine conflict that needs a move.
+//
+// NOTE: Checked with plain FileCheck, not -verify-diagnostics: which of %vecA/%vecB
+// gets named in the note is order-dependent, so it isn't deterministic across runs.
+
+// CHECK: error: 'arm_sme.copy_tile' op tile operand allocated to different SME virtial tile (move required)
+// CHECK: note: tile operand is:
+func.func @overlapping_branches_with_real_tile_values(%cond: i1, %src: memref<?x?xi8>, %mask: vector<[16]xi1>) {
+ %tile = arm_sme.get_tile : vector<[16]x[16]xi8>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %vecA = arm_sme.load_tile_slice %src[%c0], %mask, %tile, %c0 : memref<?x?xi8>, vector<[16]xi1>, vector<[16]x[16]xi8>
+ %vecB = arm_sme.load_tile_slice %src[%c0], %mask, %tile, %c1 : memref<?x?xi8>, vector<[16]xi1>, vector<[16]x[16]xi8>
+ %ret = scf.if %cond -> vector<[16]x[16]xi8> {
+ scf.yield %vecA : vector<[16]x[16]xi8>
+ } else {
+ scf.yield %vecB : vector<[16]x[16]xi8>
+ }
+ "test.some_use"(%ret) : (vector<[16]x[16]xi8>) -> ()
+ return
+}
>From df41d63477964b67cd3c6b3342bfab52786cce65 Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 5 Oct 2026 12:27:26 +0200
Subject: [PATCH 6/6] Address comments
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 21a7ef5d13258..8381a63c9973d 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -812,7 +812,7 @@ LogicalResult mlir::arm_sme::allocateSMETiles(FunctionOpInterface function,
// Bail out early if the function has no ArmSME tile ops: there is nothing
// to allocate (a value not produced/consumed by such an op never gets a
// tile ID), so this avoids the (non-trivial) preprocessing and liveness
- // analysis below for functions that have nothing to allocate.
+ // analysis.
if (!hasArmSMETileOps(function))
return success();
More information about the Mlir-commits
mailing list