1 /* Copyright 2019 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 16 #ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_TF_MLIR_TRANSLATE_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_TF_MLIR_TRANSLATE_H_ 18 19 #include <string> 20 #include <unordered_set> 21 22 #include "absl/base/macros.h" 23 #include "absl/strings/string_view.h" 24 #include "absl/types/span.h" 25 #include "mlir/IR/BuiltinOps.h" // from @llvm-project 26 #include "mlir/IR/MLIRContext.h" // from @llvm-project 27 #include "tensorflow/cc/saved_model/loader.h" 28 #include "tensorflow/compiler/mlir/tensorflow/translate/mlir_import_options.h" 29 #include "tensorflow/stream_executor/lib/statusor.h" 30 31 namespace tensorflow { 32 33 using stream_executor::port::Status; 34 using stream_executor::port::StatusOr; 35 36 // TODO(antiagainst): Directly manipulating files in library functions is not 37 // a good idea. We should pass in a string/stream here. 38 39 // Converts a TensorFlow GraphDef contained in `input` param into a MLIR module. 40 // Creates MLIR entities into the given MLIR `context`. 41 StatusOr<mlir::OwningModuleRef> GraphdefToMlirTranslateFunction( 42 llvm::StringRef input, absl::string_view debug_info_file, 43 const std::vector<std::string>& input_arrays, 44 const std::vector<std::string>& input_dtypes, 45 const std::vector<llvm::Optional<std::vector<int>>>& input_shapes, 46 const std::vector<std::string>& output_arrays, 47 const std::vector<std::string>& control_output_arrays, 48 bool prune_unused_nodes, bool convert_legacy_fed_inputs, 49 bool graph_as_function, bool upgrade_legacy, 50 // TODO(jpienaar): Remove this. 51 bool enable_shape_inference, mlir::MLIRContext* context); 52 53 ABSL_DEPRECATED( 54 "Please use the other overload of this function which accepts structured " 55 "inputs instead of strings") 56 // Converts a TensorFlow GraphDef contained in `input` param into a MLIR module. 57 // Creates MLIR entities into the given MLIR `context`. 58 StatusOr<mlir::OwningModuleRef> GraphdefToMlirTranslateFunction( 59 llvm::StringRef input, absl::string_view debug_info_file, 60 absl::string_view input_arrays, absl::string_view input_dtypes, 61 absl::string_view input_shapes, absl::string_view output_arrays, 62 absl::string_view control_output_arrays, bool prune_unused_nodes, 63 bool convert_legacy_fed_inputs, bool graph_as_function, bool upgrade_legacy, 64 // TODO(jpienaar): Remove this. 65 bool enable_shape_inference, mlir::MLIRContext* context); 66 67 // Similar as the above function, but replaces all constant tensors 68 // with randomly generated splat values. 69 StatusOr<mlir::OwningModuleRef> GraphdefToSplattedMlirTranslateFunction( 70 llvm::StringRef input, absl::string_view debug_info_file, 71 const std::vector<std::string>& input_arrays, 72 const std::vector<std::string>& input_dtypes, 73 const std::vector<std::vector<int>>& input_shapes, 74 const std::vector<std::string>& output_arrays, 75 const std::vector<std::string>& control_output_arrays, 76 bool prune_unused_nodes, bool convert_legacy_fed_inputs, 77 bool graph_as_function, bool upgrade_legacy, bool enable_shape_inference, 78 mlir::MLIRContext* context); 79 80 ABSL_DEPRECATED( 81 "Please use the other overload of this function which accepts structured " 82 "inputs instead of strings") 83 // Similar as the above function, but replaces all constant tensors 84 // with randomly generated splat values. 85 StatusOr<mlir::OwningModuleRef> GraphdefToSplattedMlirTranslateFunction( 86 llvm::StringRef input, absl::string_view debug_info_file, 87 absl::string_view input_arrays, absl::string_view input_dtypes, 88 absl::string_view input_shapes, absl::string_view output_arrays, 89 absl::string_view control_output_arrays, bool prune_unused_nodes, 90 bool convert_legacy_fed_inputs, bool graph_as_function, bool upgrade_legacy, 91 // TODO(jpienaar): Remove this. 92 bool enable_shape_inference, mlir::MLIRContext* context); 93 94 // Converts a TensorFlow SavedModel stored in the directory with the given 95 // `saved_model_dir` into a MLIR module. Creates MLIR entities into the 96 // given MLIR `context`. 97 StatusOr<mlir::OwningModuleRef> SavedModelObjectGraphToMlirImport( 98 absl::string_view saved_model_dir, 99 const std::unordered_set<std::string>& tags, 100 absl::Span<std::string> exported_names, mlir::MLIRContext* context); 101 102 // Converts a TensorFlow V1 SavedModel stored in the directory with the given 103 // `saved_model_dir` into a MLIR module. Creates MLIR entities into the 104 // given MLIR `context`. 105 // 'saved_model_bundle' if not null, will be initialized with the model bundle. 106 StatusOr<mlir::OwningModuleRef> SavedModelSignatureDefsToMlirImport( 107 absl::string_view saved_model_dir, 108 const std::unordered_set<std::string>& tags, 109 absl::Span<std::string> exported_names, mlir::MLIRContext* context, 110 MLIRImportOptions options, bool lift_variables = true, 111 std::unique_ptr<tensorflow::SavedModelBundle>* saved_model_bundle = 112 nullptr); 113 114 // Converts a TensorFlow V1 SavedModel stored in the directory with the given 115 // `saved_model_dir` into a MLIR module. Creates MLIR entities into the 116 // given MLIR `context`. This does not create session internally so it is faster 117 // and does not perform any graph transformation. 118 StatusOr<mlir::OwningModuleRef> SavedModelSignatureDefsToMlirImportLite( 119 absl::string_view saved_model_dir, 120 const std::unordered_set<std::string>& tags, 121 absl::Span<std::string> exported_names, mlir::MLIRContext* context, 122 MLIRImportOptions options); 123 124 } // namespace tensorflow 125 126 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_TF_MLIR_TRANSLATE_H_ 127