[Mosaic TPU] Add some fold rules for the transpose op.

Reverts 2a3525476e

PiperOrigin-RevId: 987600180

third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc

@@ -955,6 +955,28 @@ LogicalResult TransposeOp::verify() {
   return success();
 }
 
+OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
+  if (llvm::is_sorted(getPermutation())) {
+    return getVector();
+  }
+  if (auto cst = dyn_cast_if_present<SplatElementsAttr>(adaptor.getVector())) {
+    return cst.reshape(getType());
+  }
+  if (auto prev_op = getVector().getDefiningOp<TransposeOp>()) {
+    ArrayRef<int64_t> prev_perm = prev_op.getPermutation();
+    ArrayRef<int64_t> curr_perm = getPermutation();
+    CHECK_EQ(prev_perm.size(), curr_perm.size());
+    SmallVector<int64_t> composed_perm(curr_perm.size());
+    for (size_t i = 0; i < curr_perm.size(); ++i) {
+      composed_perm[i] = prev_perm[curr_perm[i]];
+    }
+    if (llvm::is_sorted(composed_perm)) {
+      return prev_op.getVector();
+    }
+  }
+  return nullptr;
+}
+
 LogicalResult MemRefBitcastOp::verify() {
   auto src_ty = getInput().getType();
   auto tgt_ty = getType();

third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td

@@ -1921,6 +1921,7 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> {
     }
   }];
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def TPU_LogOp : TPU_Op<"log"> {