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