Skip to content

Commit 55fe7eb

Browse files
MacDueMacDue
committed
Fixups
1 parent d4a0f1d commit 55fe7eb

File tree

2 files changed

+103
-46
lines changed

2 files changed

+103
-46
lines changed

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

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,7 @@ struct LiveRange {
297297
.valid();
298298
}
299299

300-
/// Returns true if this range overlaps with `point`.
300+
/// Returns true if this range is active at `point` in the program.
301301
bool overlaps(uint64_t point) const {
302302
return ranges->lookup(point) == kValidLiveRange;
303303
}
@@ -531,29 +531,37 @@ LiveRange *chooseSpillUsingHeuristics(ArrayRef<LiveRange *> activeRanges,
531531
return !isTileTypeGreaterOrEqual(a.getTileType(), b.getTileType()) ||
532532
a.end() < b.end();
533533
};
534-
LiveRange &lastActiveLiveRange = *std::max_element(
534+
LiveRange &latestEndingLiveRange = *std::max_element(
535535
allOverlappingRanges.begin(), allOverlappingRanges.end(),
536536
isSmallerTileTypeOrEndsEarlier);
537-
if (!isSmallerTileTypeOrEndsEarlier(lastActiveLiveRange, *newRange))
538-
return &lastActiveLiveRange;
537+
if (!isSmallerTileTypeOrEndsEarlier(latestEndingLiveRange, *newRange))
538+
return &latestEndingLiveRange;
539539
return newRange;
540540
}
541541

542542
/// Greedily allocate tile IDs to live ranges. Spill using simple heuristics.
543543
void allocateTilesToLiveRanges(
544544
ArrayRef<LiveRange *> liveRangesSortedByStartPoint) {
545545
TileAllocator tileAllocator;
546+
// `activeRanges` = Live ranges that need to be in a tile at the current point
547+
// in the program.
546548
SetVector<LiveRange *> activeRanges;
549+
// `inactiveRanges` = Live ranges that _do not_ need to be in a tile
550+
// at the current point in the program but could become active again later.
551+
// An inactive section of a live range can be seen as a 'hole' in the live
552+
// range, where it is possible to re-use the live range's tile ID _before_ has
553+
// it has ended. This allows reusing tiles more (so avoids spills).
547554
SetVector<LiveRange *> inactiveRanges;
548555
for (LiveRange *nextRange : liveRangesSortedByStartPoint) {
556+
// Update the `activeRanges` at `newRange->start()`.
549557
activeRanges.remove_if([&](LiveRange *activeRange) {
550-
// Check for live ranges that have expired.
558+
// 1. Check for live ranges that have expired.
551559
if (activeRange->end() <= nextRange->start()) {
552560
tileAllocator.releaseTileId(activeRange->getTileType(),
553561
*activeRange->tileId);
554562
return true;
555563
}
556-
// Check for live ranges that have become inactive.
564+
// 2. Check for live ranges that have become inactive.
557565
if (!activeRange->overlaps(nextRange->start())) {
558566
tileAllocator.releaseTileId(activeRange->getTileType(),
559567
*activeRange->tileId);
@@ -562,12 +570,13 @@ void allocateTilesToLiveRanges(
562570
}
563571
return false;
564572
});
573+
// Update the `inactiveRanges` at `newRange->start()`.
565574
inactiveRanges.remove_if([&](LiveRange *inactiveRange) {
566-
// Check for live ranges that have expired.
575+
// 1. Check for live ranges that have expired.
567576
if (inactiveRange->end() <= nextRange->start()) {
568577
return true;
569578
}
570-
// Check for live ranges that have become active.
579+
// 2. Check for live ranges that have become active.
571580
if (inactiveRange->overlaps(nextRange->start())) {
572581
tileAllocator.acquireTileId(inactiveRange->getTileType(),
573582
*inactiveRange->tileId);
@@ -577,9 +586,9 @@ void allocateTilesToLiveRanges(
577586
return false;
578587
});
579588

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.
589+
// Collect inactive live ranges that overlap with the new live range. We
590+
// need to reserve the tile IDs of overlapping inactive ranges to prevent
591+
// two (overlapping) live ranges from getting the same tile ID.
583592
SmallVector<LiveRange *> overlappingInactiveRanges;
584593
for (LiveRange *inactiveRange : inactiveRanges) {
585594
if (inactiveRange->overlaps(*nextRange)) {

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

Lines changed: 83 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -436,30 +436,32 @@ func.func @cond_branch_with_backedge(%slice: vector<[4]xf32>) {
436436
// CHECK-LIVE-RANGE-LABEL: @fill_holes_in_tile_liveness
437437
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
438438
// 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
439+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
440+
// CHECK-LIVE-RANGE-NEXT: E cf.cond_br
441+
// CHECK-LIVE-RANGE-NEXT: ^bb1:
442+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
443+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
444+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
445+
// CHECK-LIVE-RANGE-NEXT: cf.br
446+
// CHECK-LIVE-RANGE-NEXT: ^bb2:
447+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
448+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
449+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
450+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
451+
// CHECK-LIVE-RANGE-NEXT: cf.br
452452

453453
// 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.
454+
// can reuse the tile ID assigned to %tileA. The liveness for %tileB is
455+
// entirely within the 'hole' in %tileA's live range, so %tileB should get the
456+
// same tile ID as %tileA.
455457

456458
// CHECK-LABEL: @fill_holes_in_tile_liveness
457459
func.func @fill_holes_in_tile_liveness(%cond: i1) {
458-
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
460+
// CHECK: arm_sme.get_tile {tile_id = [[TILE_ID_A:.*]] : i32}
459461
%tileA = arm_sme.get_tile : vector<[4]x[4]xf32>
460462
cf.cond_br %cond, ^bb2, ^bb1
461463
^bb1:
462-
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
464+
// CHECK: arm_sme.get_tile {tile_id = [[TILE_ID_A]] : i32}
463465
%tileB = arm_sme.get_tile : vector<[4]x[4]xf32>
464466
"test.dummy"(): () -> ()
465467
"test.some_use"(%tileB) : (vector<[4]x[4]xf32>) -> ()
@@ -479,25 +481,25 @@ func.func @fill_holes_in_tile_liveness(%cond: i1) {
479481
// CHECK-LIVE-RANGE-LABEL: @holes_in_tile_liveness_inactive_overlaps
480482
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
481483
// 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
484+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
485+
// CHECK-LIVE-RANGE-NEXT: E cf.cond_br
486+
// CHECK-LIVE-RANGE-NEXT: ^bb1:
487+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
488+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
489+
// CHECK-LIVE-RANGE-NEXT: | test.some_use
490+
// CHECK-LIVE-RANGE-NEXT: | arm_sme.copy_tile
491+
// CHECK-LIVE-RANGE-NEXT: E cf.br
492+
// CHECK-LIVE-RANGE-NEXT: ^bb2:
493+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
494+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
495+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
496+
// CHECK-LIVE-RANGE-NEXT: |S arm_sme.get_tile
497+
// CHECK-LIVE-RANGE-NEXT: E| test.some_use
498+
// CHECK-LIVE-RANGE-NEXT: | arm_sme.copy_tile
499+
// CHECK-LIVE-RANGE-NEXT: E cf.br
500+
// CHECK-LIVE-RANGE-NEXT: ^bb3:
501+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
502+
// CHECK-LIVE-RANGE-NEXT: func.return
501503

502504
// This tests an edge case in inactive live ranges. The first live range is
503505
// inactive at the start of ^bb1. If the tile allocator did not check if the
@@ -560,3 +562,49 @@ func.func @spill_inactive_live_range(%cond: i1) {
560562
"test.some_use"(%tile) : (vector<[16]x[16]xi8>) -> ()
561563
return
562564
}
565+
566+
// -----
567+
568+
// CHECK-LIVE-RANGE-LABEL: @reactivate_inactive_live_range
569+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
570+
// CHECK-LIVE-RANGE: ^bb0:
571+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
572+
// CHECK-LIVE-RANGE-NEXT: E cf.cond_br
573+
// CHECK-LIVE-RANGE-NEXT: ^bb1:
574+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
575+
// CHECK-LIVE-RANGE-NEXT: | test.dummy
576+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
577+
// CHECK-LIVE-RANGE-NEXT: cf.br
578+
// CHECK-LIVE-RANGE-NEXT: ^bb2:
579+
// CHECK-LIVE-RANGE-NEXT: | S arm_sme.get_tile
580+
// CHECK-LIVE-RANGE-NEXT: | | test.dummy
581+
// CHECK-LIVE-RANGE-NEXT: | | test.dummy
582+
// CHECK-LIVE-RANGE-NEXT: | E test.some_use
583+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
584+
// CHECK-LIVE-RANGE-NEXT: cf.br
585+
586+
// Here the live range for %tileA becomes inactive in bb1 (so %tileB gets tile
587+
// ID 0 too). Then in bb2 the live range for tileA is reactivated as it overlaps
588+
// with the start of %tileC's live range (which means %tileC gets tile ID 1).
589+
590+
func.func @reactivate_inactive_live_range(%cond: i1) {
591+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
592+
%tileA = arm_sme.get_tile : vector<[4]x[4]xf32>
593+
cf.cond_br %cond, ^bb2, ^bb1
594+
^bb1:
595+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
596+
%tileB = arm_sme.get_tile : vector<[16]x[16]xi8>
597+
"test.dummy"(): () -> ()
598+
"test.some_use"(%tileB) : (vector<[16]x[16]xi8>) -> ()
599+
cf.br ^bb3
600+
^bb2:
601+
// CHECK: arm_sme.get_tile {tile_id = 1 : i32}
602+
%tileC = arm_sme.get_tile : vector<[4]x[4]xf32>
603+
"test.dummy"(): () -> ()
604+
"test.dummy"(): () -> ()
605+
"test.some_use"(%tileC) : (vector<[4]x[4]xf32>) -> ()
606+
"test.some_use"(%tileA) : (vector<[4]x[4]xf32>) -> ()
607+
cf.br ^bb3
608+
^bb3:
609+
return
610+
}

0 commit comments

Comments
 (0)