Missing Transpose of Binary Op optimization

[avik-pal]

module {
  func.func @main(%arg0: tensor<16x4xf32>) -> tensor<16x4xf32> {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<16x4xf32>) -> tensor<4x16xf32>
    %1 = stablehlo.subtract %0, %0 : tensor<4x16xf32>
    %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<4x16xf32>) -> tensor<16x4xf32>
    return %2 : tensor<16x4xf32>
  }
}

missed it in my last PR. Probably not working since %0 has 2 uses, which can be specialized here

[wsmoses]

Perhaps the better check would be “only used in this op” rather than a specific use count

[avik-pal]

How do I run the transpose binary_op optimization before transpose -> reshape optimization? Else we get this which could have been just sqrt and divide

module {
  func.func @main(%arg0: tensor<4x1x1x1xf32>) -> tensor<4x1x1x1xf32> {
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<1x1x1x4xf32>
    %0 = stablehlo.sqrt %arg0 : tensor<4x1x1x1xf32>
    %1 = stablehlo.reshape %0 : (tensor<4x1x1x1xf32>) -> tensor<1x1x1x4xf32>
    %2 = stablehlo.divide %cst, %1 : tensor<1x1x1x4xf32>
    %3 = stablehlo.reshape %2 : (tensor<1x1x1x4xf32>) -> tensor<4x1x1x1xf32>
    return %3 : tensor<4x1x1x1xf32>
  }
}

[wsmoses]

Longer term there's some ongoing work to use e-graphs to automatically search for what the optimal ordering of rewrite rules is [and this is a useful example] (cc @aryavohra @smjleo @ftynse @albertcohen @jbachurski )

However here probably the solution is to generalize the rsqrt optimization to handle intermediate reshape/transposes (which is done similarly elsewhere)