Skip to content

Commit d4a0f1d

Browse files
MacDueMacDue
committed
[mlir][ArmSME] Support filling liveness 'holes' in the tile allocator
Holes in a live range are points where the corresponding value does not need to be in a tile/register. If the tile allocator keeps track of these holes it can reuse tiles for more values (avoiding spills). Take this simple example: ```mlir func.func @example(%cond: i1) { %tileA = arm_sme.get_tile : vector<[4]x[4]xf32> cf.cond_br %cond, ^bb2, ^bb1 ^bb1: // If we end up here we never use %tileA again! "test.some_use"(%tileB) : (vector<[4]x[4]xf32>) -> () cf.br ^bb3 ^bb2: "test.some_use"(%tileA) : (vector<[4]x[4]xf32>) -> () cf.br ^bb3 ^bb3: return } ``` If you were to calculate the liveness of %tileA and %tileB. You'd see there is a hole in the liveness of %tileA in bb1: ``` %tileA %tileB ^bb0: Live ^bb1: Live ^bb2: Live ``` The tile allocator can make use of that hole and reuse the tile ID it assigned to %tileA for %tileB.
1 parent c2fe75f commit d4a0f1d

File tree

2 files changed

+218
-25
lines changed

2 files changed

+218
-25
lines changed

‎mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp

Lines changed: 88 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -153,10 +153,18 @@ class TileAllocator {
153153
return failure();
154154
}
155155

156+
/// Acquires a specific tile ID. Asserts the tile is initially free.
157+
void acquireTileId(ArmSMETileType tileType, unsigned tileId) {
158+
TileMask tileMask = getMasks(tileType)[tileId];
159+
assert((tilesInUse & tileMask) == TileMask::kNone &&
160+
"cannot acquire allocated tile!");
161+
tilesInUse |= tileMask;
162+
}
163+
156164
/// Releases a previously allocated tile ID.
157165
void releaseTileId(ArmSMETileType tileType, unsigned tileId) {
158166
TileMask tileMask = getMasks(tileType)[tileId];
159-
assert((tilesInUse & tileMask) != TileMask::kNone &&
167+
assert((tilesInUse & tileMask) == tileMask &&
160168
"cannot release unallocated tile!");
161169
tilesInUse ^= tileMask;
162170
}
@@ -289,6 +297,11 @@ struct LiveRange {
289297
.valid();
290298
}
291299

300+
/// Returns true if this range overlaps with `point`.
301+
bool overlaps(uint64_t point) const {
302+
return ranges->lookup(point) == kValidLiveRange;
303+
}
304+
292305
/// Unions this live range with `otherRange`, aborts if the ranges overlap.
293306
void unionWith(LiveRange const &otherRange) {
294307
for (auto it = otherRange.ranges->begin(); it != otherRange.ranges->end();
@@ -488,76 +501,126 @@ coalesceTileLiveRanges(DenseMap<Value, LiveRange> &initialLiveRanges) {
488501
return std::move(coalescedLiveRanges);
489502
}
490503

491-
/// Choose a live range to spill (via some heuristics). This picks either an
492-
/// active live range from `activeRanges` or the new live range `newRange`.
504+
/// Choose a live range to spill (via some heuristics). This picks either a live
505+
/// range from `activeRanges`, `inactiveRanges`, or the new live range
506+
/// `newRange`. Note: All live ranges in `activeRanges` and `inactiveRanges` are
507+
/// assumed to overlap with `newRange`.
493508
LiveRange *chooseSpillUsingHeuristics(ArrayRef<LiveRange *> activeRanges,
509+
ArrayRef<LiveRange *> inactiveRanges,
494510
LiveRange *newRange) {
511+
auto allOverlappingRanges =
512+
llvm::concat<LiveRange>(llvm::make_pointee_range(activeRanges),
513+
llvm::make_pointee_range(inactiveRanges));
514+
495515
// Heuristic: Spill trivially copyable operations (usually free).
496-
auto isTrivialSpill = [&](LiveRange *allocatedRange) {
497-
return isTileTypeGreaterOrEqual(allocatedRange->getTileType(),
516+
auto isTrivialSpill = [&](LiveRange &allocatedRange) {
517+
return isTileTypeGreaterOrEqual(allocatedRange.getTileType(),
498518
newRange->getTileType()) &&
499-
allocatedRange->values.size() == 1 &&
519+
allocatedRange.values.size() == 1 &&
500520
isTriviallyCloneableTileOp(
501-
allocatedRange->values[0]
502-
.getDefiningOp<ArmSMETileOpInterface>());
521+
allocatedRange.values[0].getDefiningOp<ArmSMETileOpInterface>());
503522
};
504-
if (isTrivialSpill(newRange))
523+
if (isTrivialSpill(*newRange))
505524
return newRange;
506-
auto trivialSpill = llvm::find_if(activeRanges, isTrivialSpill);
507-
if (trivialSpill != activeRanges.end())
508-
return *trivialSpill;
525+
auto trivialSpill = llvm::find_if(allOverlappingRanges, isTrivialSpill);
526+
if (trivialSpill != allOverlappingRanges.end())
527+
return &*trivialSpill;
509528

510529
// Heuristic: Spill the range that ends last (with a compatible tile type).
511-
auto isSmallerTileTypeOrEndsEarlier = [](LiveRange *a, LiveRange *b) {
512-
return !isTileTypeGreaterOrEqual(a->getTileType(), b->getTileType()) ||
513-
a->end() < b->end();
530+
auto isSmallerTileTypeOrEndsEarlier = [](LiveRange &a, LiveRange &b) {
531+
return !isTileTypeGreaterOrEqual(a.getTileType(), b.getTileType()) ||
532+
a.end() < b.end();
514533
};
515-
LiveRange *lastActiveLiveRange = *std::max_element(
516-
activeRanges.begin(), activeRanges.end(), isSmallerTileTypeOrEndsEarlier);
517-
if (!isSmallerTileTypeOrEndsEarlier(lastActiveLiveRange, newRange))
518-
return lastActiveLiveRange;
534+
LiveRange &lastActiveLiveRange = *std::max_element(
535+
allOverlappingRanges.begin(), allOverlappingRanges.end(),
536+
isSmallerTileTypeOrEndsEarlier);
537+
if (!isSmallerTileTypeOrEndsEarlier(lastActiveLiveRange, *newRange))
538+
return &lastActiveLiveRange;
519539
return newRange;
520540
}
521541

522542
/// Greedily allocate tile IDs to live ranges. Spill using simple heuristics.
523-
/// Note: This does not attempt to fill holes in active live ranges.
524543
void allocateTilesToLiveRanges(
525544
ArrayRef<LiveRange *> liveRangesSortedByStartPoint) {
526545
TileAllocator tileAllocator;
527546
SetVector<LiveRange *> activeRanges;
547+
SetVector<LiveRange *> inactiveRanges;
528548
for (LiveRange *nextRange : liveRangesSortedByStartPoint) {
529-
// Release tile IDs from live ranges that have ended.
530549
activeRanges.remove_if([&](LiveRange *activeRange) {
550+
// Check for live ranges that have expired.
531551
if (activeRange->end() <= nextRange->start()) {
532552
tileAllocator.releaseTileId(activeRange->getTileType(),
533553
*activeRange->tileId);
534554
return true;
535555
}
556+
// Check for live ranges that have become inactive.
557+
if (!activeRange->overlaps(nextRange->start())) {
558+
tileAllocator.releaseTileId(activeRange->getTileType(),
559+
*activeRange->tileId);
560+
inactiveRanges.insert(activeRange);
561+
return true;
562+
}
563+
return false;
564+
});
565+
inactiveRanges.remove_if([&](LiveRange *inactiveRange) {
566+
// Check for live ranges that have expired.
567+
if (inactiveRange->end() <= nextRange->start()) {
568+
return true;
569+
}
570+
// Check for live ranges that have become active.
571+
if (inactiveRange->overlaps(nextRange->start())) {
572+
tileAllocator.acquireTileId(inactiveRange->getTileType(),
573+
*inactiveRange->tileId);
574+
activeRanges.insert(inactiveRange);
575+
return true;
576+
}
536577
return false;
537578
});
538579

580+
// Collect inactive live ranges that overlap with the current new live
581+
// range. We need to acquire the tile IDs of overlapping inactive ranges to
582+
// prevent two (overlapping) live ranges from getting the same tile ID.
583+
SmallVector<LiveRange *> overlappingInactiveRanges;
584+
for (LiveRange *inactiveRange : inactiveRanges) {
585+
if (inactiveRange->overlaps(*nextRange)) {
586+
tileAllocator.acquireTileId(inactiveRange->getTileType(),
587+
*inactiveRange->tileId);
588+
overlappingInactiveRanges.push_back(inactiveRange);
589+
}
590+
}
591+
539592
// Allocate a tile ID to `nextRange`.
540593
auto rangeTileType = nextRange->getTileType();
541594
auto tileId = tileAllocator.allocateTileId(rangeTileType);
542595
if (succeeded(tileId)) {
543596
nextRange->tileId = *tileId;
544597
} else {
545-
LiveRange *rangeToSpill =
546-
chooseSpillUsingHeuristics(activeRanges.getArrayRef(), nextRange);
598+
LiveRange *rangeToSpill = chooseSpillUsingHeuristics(
599+
activeRanges.getArrayRef(), overlappingInactiveRanges, nextRange);
547600
if (rangeToSpill != nextRange) {
548-
// Spill an active live range (so release its tile ID first).
601+
// Spill an (in)active live range (so release its tile ID first).
549602
tileAllocator.releaseTileId(rangeToSpill->getTileType(),
550603
*rangeToSpill->tileId);
551-
activeRanges.remove(rangeToSpill);
552604
// This will always succeed after a spill (of an active live range).
553605
nextRange->tileId = *tileAllocator.allocateTileId(rangeTileType);
606+
// Remove the live range from the active/inactive sets.
607+
if (!activeRanges.remove(rangeToSpill)) {
608+
bool removed = inactiveRanges.remove(rangeToSpill);
609+
assert(removed && "expected a range to be removed!");
610+
}
554611
}
555612
rangeToSpill->tileId = tileAllocator.allocateInMemoryTileId();
556613
}
557614

558615
// Insert the live range into the active ranges.
559616
if (nextRange->tileId < kInMemoryTileIdBase)
560617
activeRanges.insert(nextRange);
618+
619+
// Release tiles reserved for inactive live ranges.
620+
for (LiveRange *range : overlappingInactiveRanges) {
621+
if (*range->tileId < kInMemoryTileIdBase)
622+
tileAllocator.releaseTileId(range->getTileType(), *range->tileId);
623+
}
561624
}
562625
}
563626

‎mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -430,3 +430,133 @@ func.func @cond_branch_with_backedge(%slice: vector<[4]xf32>) {
430430
// Live here: %finalTileA, %finalTileB, %finalTileC, %finalTileD
431431
return
432432
}
433+
434+
// -----
435+
436+
// CHECK-LIVE-RANGE-LABEL: @fill_holes_in_tile_liveness
437+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
438+
// CHECK-LIVE-RANGE: ^bb0:
439+
// CHECK-LIVE-RANGE: S arm_sme.get_tile
440+
// CHECK-LIVE-RANGE: E cf.cond_br
441+
// CHECK-LIVE-RANGE: ^bb1:
442+
// CHECK-LIVE-RANGE: S arm_sme.get_tile
443+
// CHECK-LIVE-RANGE: | test.dummy
444+
// CHECK-LIVE-RANGE: E test.some_use
445+
// CHECK-LIVE-RANGE: cf.br
446+
// CHECK-LIVE-RANGE: ^bb2:
447+
// CHECK-LIVE-RANGE: | test.dummy
448+
// CHECK-LIVE-RANGE: | test.dummy
449+
// CHECK-LIVE-RANGE: | test.dummy
450+
// CHECK-LIVE-RANGE: E test.some_use
451+
// CHECK-LIVE-RANGE: cf.br
452+
453+
// Here there's a 'hole' in the liveness of %tileA (in bb1) where another value
454+
// can reuse the tile ID (0) assigned to %tileA.
455+
456+
// CHECK-LABEL: @fill_holes_in_tile_liveness
457+
func.func @fill_holes_in_tile_liveness(%cond: i1) {
458+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
459+
%tileA = arm_sme.get_tile : vector<[4]x[4]xf32>
460+
cf.cond_br %cond, ^bb2, ^bb1
461+
^bb1:
462+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
463+
%tileB = arm_sme.get_tile : vector<[4]x[4]xf32>
464+
"test.dummy"(): () -> ()
465+
"test.some_use"(%tileB) : (vector<[4]x[4]xf32>) -> ()
466+
cf.br ^bb3
467+
^bb2:
468+
"test.dummy"(): () -> ()
469+
"test.dummy"(): () -> ()
470+
"test.dummy"(): () -> ()
471+
"test.some_use"(%tileA) : (vector<[4]x[4]xf32>) -> ()
472+
cf.br ^bb3
473+
^bb3:
474+
return
475+
}
476+
477+
// -----
478+
479+
// CHECK-LIVE-RANGE-LABEL: @holes_in_tile_liveness_inactive_overlaps
480+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
481+
// CHECK-LIVE-RANGE: ^bb0:
482+
// CHECK-LIVE-RANGE: S arm_sme.get_tile
483+
// CHECK-LIVE-RANGE: E cf.cond_br
484+
// CHECK-LIVE-RANGE: ^bb1:
485+
// CHECK-LIVE-RANGE: S arm_sme.get_tile
486+
// CHECK-LIVE-RANGE: | test.dummy
487+
// CHECK-LIVE-RANGE: | test.some_use
488+
// CHECK-LIVE-RANGE: | arm_sme.copy_tile
489+
// CHECK-LIVE-RANGE: E cf.br
490+
// CHECK-LIVE-RANGE: ^bb2:
491+
// CHECK-LIVE-RANGE: | test.dummy
492+
// CHECK-LIVE-RANGE: | test.dummy
493+
// CHECK-LIVE-RANGE: | test.dummy
494+
// CHECK-LIVE-RANGE: |S arm_sme.get_tile
495+
// CHECK-LIVE-RANGE: E| test.some_use
496+
// CHECK-LIVE-RANGE: | arm_sme.copy_tile
497+
// CHECK-LIVE-RANGE: E cf.br
498+
// CHECK-LIVE-RANGE: ^bb3:
499+
// CHECK-LIVE-RANGE: E test.some_use
500+
// CHECK-LIVE-RANGE: func.return
501+
502+
// This tests an edge case in inactive live ranges. The first live range is
503+
// inactive at the start of ^bb1. If the tile allocator did not check if the
504+
// second live range overlapped the first it would wrongly re-use tile ID 0
505+
// (as the first live range is inactive so tile ID 0 is free). This would mean
506+
// in ^bb2 two overlapping live ranges would have the same tile ID (bad!).
507+
508+
// CHECK-LABEL: @holes_in_tile_liveness_inactive_overlaps
509+
func.func @holes_in_tile_liveness_inactive_overlaps(%cond: i1) {
510+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
511+
%tileA = arm_sme.get_tile : vector<[4]x[4]xf32>
512+
cf.cond_br %cond, ^bb2, ^bb1
513+
^bb1:
514+
// CHECK: arm_sme.get_tile {tile_id = 1 : i32}
515+
%tileB = arm_sme.get_tile : vector<[4]x[4]xf32>
516+
"test.dummy"(): () -> ()
517+
"test.some_use"(%tileB) : (vector<[4]x[4]xf32>) -> ()
518+
cf.br ^bb3(%tileB: vector<[4]x[4]xf32>)
519+
^bb2:
520+
"test.dummy"(): () -> ()
521+
"test.dummy"(): () -> ()
522+
"test.dummy"(): () -> ()
523+
// CHECK: arm_sme.get_tile {tile_id = 1 : i32}
524+
%tileC = arm_sme.get_tile : vector<[4]x[4]xf32>
525+
"test.some_use"(%tileA) : (vector<[4]x[4]xf32>) -> ()
526+
cf.br ^bb3(%tileC: vector<[4]x[4]xf32>)
527+
^bb3(%tile: vector<[4]x[4]xf32>):
528+
"test.some_use"(%tile) : (vector<[4]x[4]xf32>) -> ()
529+
return
530+
}
531+
532+
// -----
533+
534+
// This is the same as the previous example, but changes the tile types to
535+
// vector<[16]x[16]xi8>. This means in bb1 the allocator will need to spill the
536+
// first live range (which is inactive).
537+
538+
// Note: The live ranges are the same as the previous example (so are not checked).
539+
540+
// CHECK-LABEL: @spill_inactive_live_range
541+
func.func @spill_inactive_live_range(%cond: i1) {
542+
// CHECK: arm_sme.get_tile {tile_id = 16 : i32}
543+
%tileA = arm_sme.get_tile : vector<[16]x[16]xi8>
544+
cf.cond_br %cond, ^bb2, ^bb1
545+
^bb1:
546+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
547+
%tileB = arm_sme.get_tile : vector<[16]x[16]xi8>
548+
"test.dummy"(): () -> ()
549+
"test.some_use"(%tileB) : (vector<[16]x[16]xi8>) -> ()
550+
cf.br ^bb3(%tileB: vector<[16]x[16]xi8>)
551+
^bb2:
552+
"test.dummy"(): () -> ()
553+
"test.dummy"(): () -> ()
554+
"test.dummy"(): () -> ()
555+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
556+
%tileC = arm_sme.get_tile : vector<[16]x[16]xi8>
557+
"test.some_use"(%tileA) : (vector<[16]x[16]xi8>) -> ()
558+
cf.br ^bb3(%tileC: vector<[16]x[16]xi8>)
559+
^bb3(%tile: vector<[16]x[16]xi8>):
560+
"test.some_use"(%tile) : (vector<[16]x[16]xi8>) -> ()
561+
return
562+
}

0 commit comments

Comments
 (0)