1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project 16 #include "mlir/Transforms/Passes.h" // from @llvm-project 17 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" 18 #include "tensorflow/compiler/mlir/tfrt/transforms/passes.h" 19 20 namespace tensorflow { 21 namespace tfrt_compiler { 22 namespace { 23 24 // Reorder tf.Assert ops or tf.If ops that contains only tf.Assert ops to the 25 // end of the function, in order to avoid unnecessary control dependencies 26 // between tf.Assert and other ops. 27 class ReorderTfAssertPass 28 : public mlir::PassWrapper<ReorderTfAssertPass, 29 mlir::OperationPass<mlir::ModuleOp>> { 30 public: MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReorderTfAssertPass)31 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReorderTfAssertPass) 32 33 llvm::StringRef getArgument() const final { return "tfrt-reorder-tf-assert"; } getDescription() const34 llvm::StringRef getDescription() const final { 35 return "Move tf.Assert to the end of the function to avoid unnecessary " 36 "control dependencies"; 37 } 38 runOnOperation()39 void runOnOperation() override { 40 auto module = getOperation(); 41 for (auto func_op : module.getOps<mlir::func::FuncOp>()) { 42 ProcessFunction(func_op); 43 } 44 } 45 ProcessFunction(mlir::func::FuncOp func_op)46 void ProcessFunction(mlir::func::FuncOp func_op) { 47 auto& block = func_op.front(); 48 49 llvm::SmallVector<mlir::Operation*, 2> assert_ops; 50 for (mlir::Operation& op : block) { 51 if (auto assert_op = llvm::dyn_cast<mlir::TF::AssertOp>(&op)) { 52 assert_ops.push_back(assert_op); 53 } 54 55 if (auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(&op)) { 56 if (IsAssertOnlyIfOp(if_op)) { 57 assert_ops.push_back(if_op); 58 } 59 } 60 } 61 62 auto& return_op = block.back(); 63 64 for (auto assert_op : assert_ops) { 65 assert_op->moveBefore(&return_op); 66 } 67 } 68 IsAssertOnlyIfOp(mlir::TF::IfOp op)69 bool IsAssertOnlyIfOp(mlir::TF::IfOp op) { 70 // If the results of the if op are used by some other ops, we cannot reorder 71 // it. 72 if (!op->use_empty()) return false; 73 74 // Only reorder if both branches are non-side-effecting or containing only 75 // Assert ops. 76 if (IsFunctionNonSideEffectingOrAssert(op.then_function()) && 77 IsFunctionNonSideEffectingOrAssert(op.else_function())) 78 return true; 79 80 return false; 81 } 82 IsFunctionNonSideEffectingOrAssert(mlir::func::FuncOp func_op)83 bool IsFunctionNonSideEffectingOrAssert(mlir::func::FuncOp func_op) { 84 auto& block = func_op.front(); 85 for (mlir::Operation& op : block) { 86 if (!llvm::isa<mlir::TF::AssertOp>(&op) && 87 !mlir::MemoryEffectOpInterface::hasNoEffect(&op)) 88 return false; 89 } 90 return true; 91 } 92 }; 93 94 } // namespace 95 96 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateReorderTfAssertPass()97CreateReorderTfAssertPass() { 98 return std::make_unique<ReorderTfAssertPass>(); 99 } 100 101 static mlir::PassRegistration<ReorderTfAssertPass> register_pass( 102 CreateReorderTfAssertPass); 103 104 } // namespace tfrt_compiler 105 } // namespace tensorflow 106