Propagate `stablehlo.reshape` and `stablehlo.transpose` through unary operations

[rsuderman]

Some operations do not manipulate shape at all (specifically unary operations). In these cases we should push shape manipulations through the unary op. E.g.

func.func @main(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = stablehlo.reshape(%arg0) : (tensor<4x4xf32>) -> tensor<16xf32>
  %1 = stablehlo.cos(%0) : (tensor<16xf32>) -> tensor<16xf32>
  %2 = stablehlo.reshape(%1) : (tensor<16xf32>) -> tensor<4x4xf32>
  return %2 : tensor<4x4xf32>
}

Has shape manipulations that are unnecessary. These can be corrected by pushing the shape manipulation the elementwise operations:

func.func @main(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = stablehlo.cos(%arg0) : (tensor<4x4xf32>) -> tensor<4x4xf32>
  %1 = stablehlo.reshape(%0) : (tensor<4x4xf32>) -> tensor<16xf32>
  %2 = stablehlo.reshape(%1) : (tensor<16xf32>) -> tensor<4x4xf32>
  return %2 : tensor<4x4xf32>
}

And ultimately simplified to:

func.func @main(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = stablehlo.cos(%arg0) : (tensor<4x4xf32>) -> tensor<4x4xf32>
  return %0 : tensor<4x4xf32>
}

This is also true with stablehlo.transpose.

[benvanik]

(+1 this can really help clean up IR and makes further stages easier to read as less of this happens in the more verbose forms!)

[allieculp]

@NatashaKnk Marking as a P1 for you - please edit as needed.

[NatashaKnk]

First part covered in 14494 Is there anybody working on the second part? (i.e. merging the shape manipulation ops together) I can also pick it up in a bit but as we've discussed it might be a good first bug :)

[alexander-shaposhnikov]

I'd be happy to pick up the second part P.S. if one can flesh out a bit more details - this would be very helpful