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/Transforms/Passes.h" // from @llvm-project 16 #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" 17 #include "tensorflow/compiler/mlir/tfrt/transforms/passes.h" 18 19 namespace tensorflow { 20 namespace tfrt_compiler { 21 namespace { 22 23 // Reorder tf.Assert ops or tf.If ops that contains only tf.Assert ops to the 24 // end of the function, in order to avoid unnecessary control dependencies 25 // between tf.Assert and other ops. 26 class ReorderTfAssertPass 27 : public mlir::PassWrapper<ReorderTfAssertPass, 28 mlir::OperationPass<mlir::ModuleOp>> { 29 public: getArgument() const30 llvm::StringRef getArgument() const final { return "tfrt-reorder-tf-assert"; } getDescription() const31 llvm::StringRef getDescription() const final { 32 return "Move tf.Assert to the end of the function to avoid unnecessary " 33 "control dependencies"; 34 } 35 runOnOperation()36 void runOnOperation() override { 37 auto module = getOperation(); 38 for (auto func_op : module.getOps<mlir::FuncOp>()) { 39 ProcessFunction(func_op); 40 } 41 } 42 ProcessFunction(mlir::FuncOp func_op)43 void ProcessFunction(mlir::FuncOp func_op) { 44 auto& block = func_op.front(); 45 46 llvm::SmallVector<mlir::Operation*, 2> assert_ops; 47 for (mlir::Operation& op : block) { 48 if (auto assert_op = llvm::dyn_cast<mlir::TF::AssertOp>(&op)) { 49 assert_ops.push_back(assert_op); 50 } 51 52 if (auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(&op)) { 53 if (IsAssertOnlyIfOp(if_op)) { 54 assert_ops.push_back(if_op); 55 } 56 } 57 } 58 59 auto& return_op = block.back(); 60 61 for (auto assert_op : assert_ops) { 62 assert_op->moveBefore(&return_op); 63 } 64 } 65 IsAssertOnlyIfOp(mlir::TF::IfOp op)66 bool IsAssertOnlyIfOp(mlir::TF::IfOp op) { 67 // If the results of the if op are used by some other ops, we cannot reorder 68 // it. 69 if (!op->use_empty()) return false; 70 71 // Only reorder if both branches are non-side-effecting or containing only 72 // Assert ops. 73 if (IsFunctionNonSideEffectingOrAssert(op.then_function()) && 74 IsFunctionNonSideEffectingOrAssert(op.else_function())) 75 return true; 76 77 return false; 78 } 79 IsFunctionNonSideEffectingOrAssert(mlir::FuncOp func_op)80 bool IsFunctionNonSideEffectingOrAssert(mlir::FuncOp func_op) { 81 auto& block = func_op.front(); 82 for (mlir::Operation& op : block) { 83 if (!llvm::isa<mlir::TF::AssertOp>(&op) && 84 !mlir::MemoryEffectOpInterface::hasNoEffect(&op)) 85 return false; 86 } 87 return true; 88 } 89 }; 90 91 } // namespace 92 93 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateReorderTfAssertPass()94CreateReorderTfAssertPass() { 95 return std::make_unique<ReorderTfAssertPass>(); 96 } 97 98 static mlir::PassRegistration<ReorderTfAssertPass> register_pass( 99 CreateReorderTfAssertPass); 100 101 } // namespace tfrt_compiler 102 } // namespace tensorflow 103