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 #ifndef TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTION_OPTIONS_H_ 16 #define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTION_OPTIONS_H_ 17 18 #include "absl/types/optional.h" 19 #include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h" 20 #include "tensorflow/core/protobuf/config.pb.h" 21 #include "tensorflow/core/public/session_options.h" 22 #include "tensorflow/core/tfrt/runtime/runtime.h" 23 24 namespace tensorflow { 25 namespace tfrt_stub { 26 27 // General options for graph execution. 28 struct GraphExecutionOptions { GraphExecutionOptionsGraphExecutionOptions29 explicit GraphExecutionOptions(const tensorflow::tfrt_stub::Runtime* rt) 30 : runtime(rt) { 31 DCHECK(runtime); 32 } 33 34 // If true, when creating an optimized subgraph, Placer and Grappler will 35 // also run on the functions. 36 bool run_placer_grappler_on_functions = false; 37 38 // If true, the function optimizer in the grappler will be enabled, and 39 // optimizations like function inlining will be applied. 40 bool enable_grappler_function_optimizer = false; 41 42 // Whether to enable TFRT GPU. 43 bool enable_tfrt_gpu = false; 44 45 // Runtime configuration. Refer to tensorflow::tfrt_stub::Runtime class for 46 // more details. It must not be nullptr; 47 const tensorflow::tfrt_stub::Runtime* runtime = nullptr; 48 49 // Model metadata used for monitoring and tracing. 50 tensorflow::SessionMetadata model_metadata; 51 52 tensorflow::TfrtCompileOptions compile_options; 53 }; 54 55 // Per-request options for graph execution. 56 struct GraphExecutionRunOptions { 57 absl::optional<std::chrono::system_clock::time_point> deadline; 58 59 // Priority of the request. Larger number means higher priority. 60 int priority = 0; 61 62 // If true, the input specs will be checked before running, and an error 63 // will be raised upon mismatch. 64 bool validate_input_specs = false; 65 66 // TODO(b/239749833) Remove after b/239749833 is fixed. 67 // If true, the input specs will be checked before running, and an error 68 // will be logged upon mismatch. 69 bool validate_input_specs_dry_run = false; 70 71 // The thread pool used for this run. If it is nullptr, a default one set 72 // in the tensorflow::tfrt_stub::Runtime will be used. 73 tensorflow::tfrt_stub::WorkQueueInterface* work_queue = nullptr; 74 75 // If true, the cost of the op will be measured at the execution time. 76 bool enable_cost_measurement = false; 77 }; 78 79 // Creates the default `SessionOptions` from a `GraphExecutionOptions`. 80 // The created `SessionOptions` contains the Grappler configs. 81 tensorflow::SessionOptions CreateDefaultSessionOptions( 82 const GraphExecutionOptions& options); 83 84 // Updates TPU target to fallback if bridge uncompatible, otherwise TPU runtime. 85 void UpdateTpuTargetByBridgeCompatibility( 86 tensorflow::tfrt_stub::GraphExecutionOptions& options, 87 const tensorflow::GraphDef& graph_def); 88 89 } // namespace tfrt_stub 90 } // namespace tensorflow 91 92 #endif // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTION_OPTIONS_H_ 93