[Metal Direct Pass] Loop Nest Generation

[nsmithtt]

Since we have a nested linalg.generic inside of the ttir.generic, we get a loop nest for free:

• --convert-linalg-to-affine-loops
• These passes can additionally manipulate the loop nest: --affine-loop-coalescing, --expand-strided-metadata

Taking the eltwise example, and running the above linalg pass yields:

%5 = "ttir.generic"(%1, %3, %4) <{
  grid = #tt.grid<1x1>,
  indexing_maps = [#map, #map, #map],
  iterator_types = [#parallel, #parallel], ...}> ({
^bb0(%arg2: memref<2x4x!tt.tile<32x32, f32>, #l1_>,
     %arg3: memref<2x4x!tt.tile<32x32, f32>, #l1_>,
     %arg4: memref<2x4x!tt.tile<32x32, f32>, #l1_>):
  affine.for %arg5 = 0 to 2 {
    affine.for %arg6 = 0 to 4 {
      %8 = affine.load %arg2[%arg5, %arg6] : memref<2x4x!tt.tile<32x32, f32>, #l1_>
      %9 = affine.load %arg3[%arg5, %arg6] : memref<2x4x!tt.tile<32x32, f32>, #l1_>
      %10 = "ttir.tile_maximum"(%8, %9) : (!tt.tile<32x32, f32>, !tt.tile<32x32, f32>) -> !tt.tile<32x32, f32>
      affine.store %10, %arg4[%arg5, %arg6] : memref<2x4x!tt.tile<32x32, f32>, #l1_>
    }
  }
  "ttir.yield"() : () -> ()
})