Missing Transpose(Unary(...)) optimizations

[avik-pal]

module {
  func.func @main(%arg0: tensor<1x3x2x2xf32>) -> (tensor<1x3x2x2xf32>, tensor<1x3x2x2xf32>) {
    %cst = stablehlo.constant dense<0.317560464> : tensor<1x3x2x2xf32>
    %cst_0 = stablehlo.constant dense<1.23323452> : tensor<1x3x2x2xf32>
    %0 = stablehlo.multiply %cst_0, %arg0 : tensor<1x3x2x2xf32>
    %1 = stablehlo.add %cst, %0 : tensor<1x3x2x2xf32>
    %2 = stablehlo.transpose %1, dims = [3, 2, 1, 0] : (tensor<1x3x2x2xf32>) -> tensor<2x2x3x1xf32>
    %3 = stablehlo.cosine %2 : tensor<2x2x3x1xf32>
    %4 = stablehlo.sine %2 : tensor<2x2x3x1xf32>
    %5 = stablehlo.transpose %4, dims = [3, 2, 1, 0] : (tensor<2x2x3x1xf32>) -> tensor<1x3x2x2xf32>
    %6 = stablehlo.transpose %3, dims = [3, 2, 1, 0] : (tensor<2x2x3x1xf32>) -> tensor<1x3x2x2xf32>
    return %5, %6 : tensor<1x3x2x2xf32>, tensor<1x3x2x2xf32>
  }
}

[avik-pal]

This is probably a bit more involved, we need to check if uses of %2 are in unary operations where eliminating the transpose is a valid op.

    %2 = stablehlo.transpose %1, dims = [3, 2, 1, 0] : (tensor<1x3x2x2xf32>) -> tensor<2x2x3x1xf32>
    %3 = stablehlo.cosine %2 : tensor<2x2x3x1xf32>
    %4 = stablehlo.sine %2 : tensor<2x2x3x1xf32>

[avik-pal]

We need 2 transforms here:

• Transpose -> Unary Op -> Transpose can be simplified to Unary Op
• Transpose -> Unary Op can be simplified to Unary Op -> Transpose. Only apply this if the transpose is being used with other unary ops