[Mlir-commits] [mlir] [mlir][ArmSME] Make tile allocation order deterministic (PR #229096)
Federico Bruzzone
llvmlistbot at llvm.org
Mon Oct 5 08:47:53 PDT 2026
https://github.com/FedericoBruzzone updated https://github.com/llvm/llvm-project/pull/229096
>From 7bd5b596c3d0afff84caebe3be6f0e61190e937b Mon Sep 17 00:00:00 2001
From: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Date: Mon, 5 Oct 2026 16:14:38 +0200
Subject: [PATCH] [mlir][ArmSME] Make tile allocation order deterministic
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
---
.../ArmSME/Transforms/TileAllocation.cpp | 13 +++++----
.../ArmSME/tile-allocation-liveness.mlir | 27 ++++++++++---------
.../ArmSME/tile-allocation-move-required.mlir | 24 -----------------
3 files changed, 22 insertions(+), 42 deletions(-)
delete mode 100644 mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 8381a63c9973d..7aca1b042f5fe 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -53,6 +53,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "llvm/ADT/IntervalMap.h"
+#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -370,13 +371,15 @@ generateOperationNumbering(FunctionOpInterface function) {
return operationToIndexMap;
}
-/// Gather live ranges for SME tiles from the MLIR liveness analysis.
-DenseMap<Value, LiveRange>
+/// Gather live ranges for SME tiles from the MLIR liveness analysis. Uses a
+/// `MapVector` (rather than a `DenseMap`) so that iteration order is
+/// deterministic (insertion order).
+llvm::MapVector<Value, LiveRange>
gatherTileLiveRanges(DenseMap<Operation *, unsigned> const &operationToIndexMap,
LiveRange::Allocator &liveRangeAllocator,
Liveness &liveness, FunctionOpInterface function) {
assert(!operationToIndexMap.empty() && "expected operation numbering");
- DenseMap<Value, LiveRange> liveRanges;
+ llvm::MapVector<Value, LiveRange> liveRanges;
/// Defines or updates a live range for an SME tile value. Live-ins may update
/// an existing live range (rather than define a new one). Note: If
/// `liveAtBlockEntry` is true then `firstUseOrDef` is the first operation in
@@ -445,8 +448,8 @@ static void forEachPredecessorTileValue(BlockArgument blockArg,
/// Coalesce live ranges where it would prevent unnecessary tile moves.
SmallVector<LiveRange *>
-coalesceTileLiveRanges(DenseMap<Value, LiveRange> &initialLiveRanges) {
- DenseMap<Value, LiveRange *> liveRanges;
+coalesceTileLiveRanges(llvm::MapVector<Value, LiveRange> &initialLiveRanges) {
+ llvm::MapVector<Value, LiveRange *> liveRanges;
for (auto &[value, liveRange] : initialLiveRanges) {
liveRanges.insert({value, &liveRange});
}
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
index cddcf96feba8e..a9a0f8fa9b302 100644
--- a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
+++ b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
@@ -241,21 +241,22 @@ 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. 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>
+// This means that %vecA and %vecB are both allocated to different tiles (which is not legal).
+
+func.func @overlapping_branches(%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>
+ // expected-note at below {{tile operand is:}}
+ %vecB = arm_sme.load_tile_slice %src[%c0], %mask, %tile, %c1 : memref<?x?xi8>, vector<[16]xi1>, vector<[16]x[16]xi8>
+ // expected-error at below {{op tile operand allocated to different SME virtial tile (move required)}}
+ %ret = scf.if %cond -> vector<[16]x[16]xi8> {
+ scf.yield %vecA : vector<[16]x[16]xi8>
} else {
- scf.yield %vecB : vector<[4]x[4]xf32>
+ scf.yield %vecB : vector<[16]x[16]xi8>
}
- "test.some_use"(%tile) : (vector<[4]x[4]xf32>) -> ()
+ "test.some_use"(%ret) : (vector<[16]x[16]xi8>) -> ()
return
}
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
deleted file mode 100644
index 56b6bda081ab4..0000000000000
--- a/mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
+++ /dev/null
@@ -1,24 +0,0 @@
-// 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
-}
More information about the Mlir-commits
mailing list