[StableHLO] Add transpose simplification

PiperOrigin-RevId: 820804015

Author: GleasonK

stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_simplification.mlir

@@ -1810,6 +1810,15 @@ func.func @transpose_is_not_reshape(%arg0: tensor<1x4x5x2xf32>) -> tensor<2x4x1x
   return %0 : tensor<2x4x1x5xf32>
 }
 
+// CHECK-LABEL: @transpose_of_transpose
+func.func @transpose_of_transpose(%arg0 : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> {
+  %0 = stablehlo.transpose %arg0, dims = [3,2,1,0] : (tensor<1x2x3x4xf32>) -> tensor<4x3x2x1xf32>
+  %1 = stablehlo.transpose %0, dims = [3,2,1,0] : (tensor<4x3x2x1xf32>) -> tensor<1x2x3x4xf32>
+  // CHECK-NOT: stablehlo.transpose
+  // CHECK: return %arg0
+  return %1 : tensor<1x2x3x4xf32>
+}
+
 // -----
 
 ////////

stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp

@@ -1309,6 +1309,17 @@ struct SortOpSetDimension : public SimplifyOpRewritePattern<SortOp> {
 // TransposeOp
 /////////////////////////////////
 
+DenseI64ArrayAttr getMergedTransposePermutation(OpBuilder& b,
+                                                ArrayRef<int64_t> childPerm,
+                                                ArrayRef<int64_t> parentPerm) {
+  SmallVector<int64_t> mergedPerm;
+  mergedPerm.reserve(parentPerm.size());
+  for (int64_t parentIdx : parentPerm) {
+    mergedPerm.push_back(childPerm[parentIdx]);
+  }
+  return b.getDenseI64ArrayAttr(mergedPerm);
+}
+
 // Pattern: transpose(X, [no_mem_layout_change...]) -> reshape(X)
 struct TransposeIsReshape final : SimplifyOpRewritePattern<TransposeOp> {
   using SimplifyOpRewritePattern::SimplifyOpRewritePattern;

stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.td

@@ -120,6 +120,8 @@ def InvertBroadcastDims : NativeCodeCall<"getInvertedBroadcastDimensions($_build
 
 def MergeBroadcastDims : NativeCodeCall<"getMergedBroadcastDimensions($_builder, $0, $1)">;
 
+def MergePermutations : NativeCodeCall<"getMergedTransposePermutation($_builder, $0, $1)">;
+
 def StableHLO_ConvertOpWithShape : NativeCodeCall<
     "$_builder.create<stablehlo::ConvertOp>($_loc, $0.getType(), $1)">;
 
@@ -539,6 +541,12 @@ def TransposeOp_RemoveNoop
   : Pat<(StableHLO_TransposeOp $lhs, IotaDims:$dims),
         (replaceWithValue $lhs)>;
 
+// Pattern: transpose(transpose(X)) -> transpose(X)
+def TransposeOp_TransposeOfTranspose
+  : Pat<(StableHLO_TransposeOp
+          (StableHLO_TransposeOp $child, $child_dims), $dims),
+        (StableHLO_TransposeOp $child, (MergePermutations $child_dims, $dims))>;
+
 ////////
 // GetTupleElementOp