[Mlir-commits] [mlir] 145af0d - [mlir][ArmSME] Make tile allocation order deterministic (#229096)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 9 00:08:57 PDT 2026
Author: Federico Bruzzone
Date: 2026-10-09T09:08:51+02:00
New Revision: 145af0deae54b95c5e70cb158f060fdc5a39c94f
URL: https://github.com/llvm/llvm-project/commit/145af0deae54b95c5e70cb158f060fdc5a39c94f
DIFF: https://github.com/llvm/llvm-project/commit/145af0deae54b95c5e70cb158f060fdc5a39c94f.diff
LOG: [mlir][ArmSME] Make tile allocation order deterministic (#229096)
`gatherTileLiveRanges` used a `DenseMap<Value, LiveRange>`, so tile-ID
tie-breaking depended on pointer hash order: _non-deterministic_ in
nature. This PR switches it to `MapVector` for insertion-order
iteration.
Also restores the _move required_ diagnostic on `overlapping_branches`
(dropped in #227153).
Additionally, it adds the forward declaration of `MapVector` in
[`mlir/Support/LLVM.h`](https://github.com/llvm/llvm-project/blob/7bd5b596c3d0afff84caebe3be6f0e61190e937b/mlir/include/mlir/Support/LLVM.h)
to avoid the prefix `llvm::`.
---------
Signed-off-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Added:
Modified:
mlir/include/mlir/Support/LLVM.h
mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
Removed:
mlir/test/Dialect/ArmSME/tile-allocation-move-required.mlir
################################################################################
diff --git a/mlir/include/mlir/Support/LLVM.h b/mlir/include/mlir/Support/LLVM.h
index 03509d4af2a234..f9b7ae03179a2c 100644
--- a/mlir/include/mlir/Support/LLVM.h
+++ b/mlir/include/mlir/Support/LLVM.h
@@ -56,6 +56,9 @@ template <typename... PT>
class PointerUnion;
template <typename T, typename Vector, typename Set, unsigned N>
class SetVector;
+template <typename KeyT, typename ValueT, typename MapType, typename VectorType,
+ unsigned N>
+class MapVector;
template <typename T, unsigned N>
class SmallPtrSet;
template <typename T>
@@ -123,6 +126,11 @@ using DenseSet = llvm::DenseSet<ValueT, ValueInfoT>;
template <typename T, typename Vector = llvm::SmallVector<T, 0>,
typename Set = DenseSet<T>, unsigned N = 0>
using SetVector = llvm::SetVector<T, Vector, Set, N>;
+template <typename KeyT, typename ValueT,
+ typename MapType = DenseMap<KeyT, unsigned>,
+ typename VectorType = llvm::SmallVector<std::pair<KeyT, ValueT>, 0>,
+ unsigned N = 0>
+using MapVector = llvm::MapVector<KeyT, ValueT, MapType, VectorType, N>;
template <typename AllocatorTy = llvm::MallocAllocator>
using StringSet = llvm::StringSet<AllocatorTy>;
using llvm::MutableArrayRef;
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 8381a63c9973df..1786e3a69d5c97 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -370,13 +370,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).
+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;
+ 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 +447,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(MapVector<Value, LiveRange> &initialLiveRanges) {
+ MapVector<Value, LiveRange *> liveRanges;
for (auto &[value, liveRange] : initialLiveRanges) {
liveRanges.insert({value, &liveRange});
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index cd71db7ec783aa..b1775cf4ab67ae 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -32,8 +32,6 @@ using namespace mlir;
using namespace mlir::linalg;
using namespace mlir::scf;
-using llvm::MapVector;
-
#define DEBUG_TYPE "linalg-promotion"
/// Alloc a new buffer of `size` * `width` i8; where `width` is given by the
diff --git a/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir b/mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir
index cddcf96feba8ea..a9a0f8fa9b3026 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
diff erent 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
diff erent 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
diff erent 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 56b6bda081ab4b..00000000000000
--- 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
diff erent 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