slice(transpose(...)) --> transpose(slice(...)) if transpose is only being used in slice

[avik-pal]

enables optimizations of routines like (if transpose is being used exclusively for slice operations)

module {
  func.func @main(%arg0: tensor<32x16xf32>, %arg1: tensor<32x8xf32>) -> (tensor<32x8xf32>, tensor<32x8xf32>) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<32x16xf32>) -> tensor<16x32xf32>
    %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<32x8xf32>) -> tensor<8x32xf32>
    %2 = stablehlo.slice %0 [0:8, 0:32] : (tensor<16x32xf32>) -> tensor<8x32xf32>
    %3 = stablehlo.slice %0 [8:16, 0:32] : (tensor<16x32xf32>) -> tensor<8x32xf32>
    %4 = stablehlo.exponential %3 : tensor<8x32xf32>
    %5 = stablehlo.multiply %1, %4 : tensor<8x32xf32>
    %6 = stablehlo.add %5, %2 : tensor<8x32xf32>
    %7 = stablehlo.transpose %6, dims = [1, 0] : (tensor<8x32xf32>) -> tensor<32x8xf32>
    %8 = stablehlo.transpose %3, dims = [1, 0] : (tensor<8x32xf32>) -> tensor<32x8xf32>
    return %7, %8 : tensor<32x8xf32>, tensor<32x8xf32>
  }
}

[avik-pal]

One caveat however would be if the new generated transposes aren't eliminated, then we might end up in a worse situation than we started at 😅

[wsmoses]

longer term I feel like this is something where we should do some sort of "register allocation for transpose" style thing. Or more simply look at the chain of transposes to see if we terminate on transposes and/or could otherwise get rid of (rather than looking on single op basis)