[Mlir-commits] [mlir] [mlir][ArmSME] Support filling liveness 'holes' in the tile allocator (PR #98350)

[Andrzej Warzyński]

================
@@ -430,3 +430,133 @@ func.func @cond_branch_with_backedge(%slice: vector<[4]xf32>) {
   // Live here: %finalTileA, %finalTileB, %finalTileC, %finalTileD
   return
 }
+
+// -----
+
+//  CHECK-LIVE-RANGE-LABEL: @fill_holes_in_tile_liveness
+//        CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
+//        CHECK-LIVE-RANGE: ^bb0:
+//        CHECK-LIVE-RANGE: S  arm_sme.get_tile
+//        CHECK-LIVE-RANGE: E  cf.cond_br
+//        CHECK-LIVE-RANGE: ^bb1:
+//        CHECK-LIVE-RANGE:  S arm_sme.get_tile
+//        CHECK-LIVE-RANGE:  | test.dummy
+//        CHECK-LIVE-RANGE:  E test.some_use
+//        CHECK-LIVE-RANGE:    cf.br
+//        CHECK-LIVE-RANGE: ^bb2:
+//        CHECK-LIVE-RANGE: |  test.dummy
+//        CHECK-LIVE-RANGE: |  test.dummy
+//        CHECK-LIVE-RANGE: |  test.dummy
+//        CHECK-LIVE-RANGE: E  test.some_use
+//        CHECK-LIVE-RANGE:    cf.br
+
+// Here there's a 'hole' in the liveness of %tileA (in bb1) where another value
+// can reuse the tile ID (0) assigned to %tileA.
+
+// CHECK-LABEL: @fill_holes_in_tile_liveness
+func.func @fill_holes_in_tile_liveness(%cond: i1) {
+  // CHECK: arm_sme.get_tile {tile_id = 0 : i32}
+  %tileA = arm_sme.get_tile : vector<[4]x[4]xf32>
+  cf.cond_br %cond, ^bb2, ^bb1
+^bb1:
+  // CHECK: arm_sme.get_tile {tile_id = 0 : i32}
+  %tileB = arm_sme.get_tile : vector<[4]x[4]xf32>
+  "test.dummy"(): () -> ()
+  "test.some_use"(%tileB) : (vector<[4]x[4]xf32>) -> ()
+  cf.br ^bb3
+^bb2:
+  "test.dummy"(): () -> ()
+  "test.dummy"(): () -> ()
+  "test.dummy"(): () -> ()
+  "test.some_use"(%tileA) : (vector<[4]x[4]xf32>) -> ()
+  cf.br ^bb3
+^bb3:
+  return
+}
----------------
banach-space wrote:

It would be helpful to annotate this a bit more. Otherwise, it's tricky to see the "key" things being tested/demonstrated/exercised. It's clear to me _today_, but let's make it future-proof :)


```suggestion
// Note the hole in the live range of %tileA - this is an opportunity to re-use the same virtual tile for %tileB

//  CHECK-LIVE-RANGE-LABEL: @fill_holes_in_tile_liveness
//        CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
//        CHECK-LIVE-RANGE: ^bb0:
//        CHECK-LIVE-RANGE: S  arm_sme.get_tile
//        CHECK-LIVE-RANGE: E  cf.cond_br
//        CHECK-LIVE-RANGE: ^bb1:
//        CHECK-LIVE-RANGE:  S arm_sme.get_tile
//        CHECK-LIVE-RANGE:  | test.dummy
//        CHECK-LIVE-RANGE:  E test.some_use
//        CHECK-LIVE-RANGE:    cf.br
//        CHECK-LIVE-RANGE: ^bb2:
//        CHECK-LIVE-RANGE: |  test.dummy
//        CHECK-LIVE-RANGE: |  test.dummy
//        CHECK-LIVE-RANGE: |  test.dummy
//        CHECK-LIVE-RANGE: E  test.some_use
//        CHECK-LIVE-RANGE:    cf.br

// As noted above, there's a 'hole' in the liveness of %tileA (in bb1) where another value
// can reuse the tile ID assigned to %tileA. That's what's being verified here (i.e. that tile_id for `%tileA` and `%tileB` are identical.

// CHECK-LABEL: @fill_holes_in_tile_liveness
func.func @fill_holes_in_tile_liveness(%cond: i1) {
  // CHECK: arm_sme.get_tile {tile_id = 0 : i32}
  %tileA = arm_sme.get_tile : vector<[4]x[4]xf32>
  cf.cond_br %cond, ^bb2, ^bb1
^bb1:
  // CHECK: arm_sme.get_tile {tile_id = 0 : i32}
  %tileB = arm_sme.get_tile : vector<[4]x[4]xf32>
  "test.dummy"(): () -> ()
  "test.some_use"(%tileB) : (vector<[4]x[4]xf32>) -> ()
  cf.br ^bb3
^bb2:
  "test.dummy"(): () -> ()
  "test.dummy"(): () -> ()
  "test.dummy"(): () -> ()
  "test.some_use"(%tileA) : (vector<[4]x[4]xf32>) -> ()
  cf.br ^bb3
^bb3:
  return
}
```

Also, could you use REGEX to capture "0" (i.e. the value of the tile ID)? That would highlight that the value ought to be identical on L458 and L462.

https://github.com/llvm/llvm-project/pull/98350