zzos/tensorflow
1
0
Fork 0
mirror of https://github.com/tensorflow/tensorflow.git synced 2026-09-28 05:13:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Reverts 4470f2f1dd

PiperOrigin-RevId: 987140670
This commit is contained in:
A. Unique TensorFlower
2026-09-25 01:02:54 -07:00
committed by TensorFlower Gardener
parent c96cc1156c
commit 2a3525476e
2 changed files with 0 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
Vendored
-22
View File
@@ -955,28 +955,6 @@ 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
Vendored
-1
View File
@@ -1921,7 +1921,6 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> {
}
}];
let hasVerifier = 1;
let hasFolder = 1;
}
def TPU_LogOp : TPU_Op<"log"> {