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 "tensorflow/core/tfrt/graph_executor/graph_execution_options.h" 16 17 #include "tensorflow/core/protobuf/rewriter_config.pb.h" 18 // TODO(b/200579737): using FunctionRegistry is simpler than the OSS trick. 19 #include "tensorflow/core/tfrt/utils/bridge_graph_analysis.h" 20 21 namespace tensorflow { 22 namespace tfrt_stub { 23 CreateDefaultSessionOptions(const GraphExecutionOptions & options)24tensorflow::SessionOptions CreateDefaultSessionOptions( 25 const GraphExecutionOptions& options) { 26 tensorflow::SessionOptions session_options; 27 auto& config = session_options.config; 28 29 config.mutable_graph_options() 30 ->mutable_rewrite_options() 31 ->set_disable_meta_optimizer(!options.compile_options.enable_grappler); 32 33 // The following configs are constant. 34 35 // Setting use_tfrt to true avoids grappler logic that lowers to v1 control 36 // flow. Note that other function inlining (e.g. on StatefulPartitionedCall) 37 // is still enabled. 38 config.mutable_experimental()->set_use_tfrt(true); 39 if (options.enable_grappler_function_optimizer) { 40 config.mutable_graph_options() 41 ->mutable_rewrite_options() 42 ->set_function_optimization(tensorflow::RewriterConfig::ON); 43 } else { 44 config.mutable_graph_options() 45 ->mutable_rewrite_options() 46 ->set_function_optimization(tensorflow::RewriterConfig::OFF); 47 } 48 // Do not skip grappler optimization even for small graphs. 49 config.mutable_graph_options() 50 ->mutable_rewrite_options() 51 ->set_min_graph_nodes(-1); 52 53 return session_options; 54 } 55 UpdateTpuTargetByBridgeCompatibility(tensorflow::tfrt_stub::GraphExecutionOptions & options,const tensorflow::GraphDef & graph_def)56void UpdateTpuTargetByBridgeCompatibility( 57 tensorflow::tfrt_stub::GraphExecutionOptions& options, 58 const tensorflow::GraphDef& graph_def) { 59 if (options.compile_options.tpu_target == 60 tensorflow::TfrtTpuInfraTarget::kBridgeFallback) { 61 auto s = tfrt::CheckTpuMlirBridgeCompatibility(graph_def); 62 if (!s.ok()) { 63 LOG(INFO) 64 << "TFRT detected Bridge unsupported feature, using TF fallback"; 65 options.compile_options.tpu_target = 66 tensorflow::TfrtTpuInfraTarget::kTfFallback; 67 } else { 68 options.compile_options.tpu_target = 69 tensorflow::TfrtTpuInfraTarget::kTpurt; 70 } 71 } 72 LOG(INFO) << "TFRT uses TPU target " << options.compile_options.tpu_target; 73 } 74 75 } // namespace tfrt_stub 76 } // namespace tensorflow 77