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_SAVED_MODEL_SAVED_MODEL_IMPORT_INPUT_H_ 16 #define TENSORFLOW_CORE_TFRT_SAVED_MODEL_SAVED_MODEL_IMPORT_INPUT_H_ 17 18 #include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h" 19 #include "tensorflow/core/tfrt/fallback/fallback_state.h" 20 #include "tensorflow/core/tfrt/utils/tfrt_graph_execution_state.h" 21 22 namespace tensorflow { 23 namespace tfrt_stub { 24 25 // TfrtSavedModelMLIRImportInput implements SavedModelMLIRImportInput, so that 26 // it can perform customization (eg. Placer and Grappler) on the input graph to 27 // the MLIR importer. 28 class TfrtSavedModelMLIRImportInput : public SavedModelMLIRImportInput { 29 public: 30 static StatusOr<TfrtSavedModelMLIRImportInput> Create( 31 const FallbackState& fallback_state, const MetaGraphDef* meta_graph_def, 32 const GraphDebugInfo& debug_info, 33 bool run_placer_grappler_on_nested_functions = false, 34 bool enable_tfrt_gpu = false); 35 36 TfrtSavedModelMLIRImportInput( 37 const MetaGraphDef* meta_graph_def, const GraphDebugInfo& debug_info, 38 std::unique_ptr<TfrtGraphExecutionState> graph_execution_state); 39 40 StatusOr<const tensorflow::Graph*> GetSubGraph( 41 absl::string_view name, GraphImportConfig& graph_import_config) override; 42 43 // Return the time used by grappler. GetGrapplerDuration()44 absl::Duration GetGrapplerDuration() const { return grappler_duration_; } 45 46 // Return the time used by functionalization. GetFunctionalizationDuration()47 absl::Duration GetFunctionalizationDuration() const { 48 return functionalization_duration_; 49 } 50 51 private: 52 std::unique_ptr<TfrtGraphExecutionState> graph_execution_state_; 53 absl::flat_hash_map<std::string, std::unique_ptr<tensorflow::Graph>> 54 optimized_graphs_; 55 56 absl::Duration functionalization_duration_; 57 absl::Duration grappler_duration_; 58 }; 59 60 } // namespace tfrt_stub 61 } // namespace tensorflow 62 63 #endif // TENSORFLOW_CORE_TFRT_SAVED_MODEL_SAVED_MODEL_IMPORT_INPUT_H_ 64