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_LITE_TF_TFL_PASSES_H_ 17 #define TENSORFLOW_COMPILER_MLIR_LITE_TF_TFL_PASSES_H_ 18 19 #include "llvm/ADT/Optional.h" 20 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project 21 #include "mlir/IR/BuiltinOps.h" // from @llvm-project 22 #include "mlir/Pass/PassManager.h" // from @llvm-project 23 #include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h" 24 #include "tensorflow/lite/toco/model_flags.pb.h" 25 #include "tensorflow/lite/toco/toco_flags.pb.h" 26 27 namespace tensorflow { 28 29 // Add the TF to TFLite passes, specified in the pass_config, into a 30 // pass_manager. The session object will be provided when the TF MLIR is 31 // imported from saved model version one and utilized for capturing resource 32 // variables. If the `saved_model_dir` directory path is provided, then the 33 // `tf_saved_model.asset` ops will be freezed. 34 void AddTFToTFLConversionPasses(llvm::StringRef saved_model_dir, 35 const toco::TocoFlags& toco_flags, 36 const mlir::TFL::PassConfig& pass_config, 37 mlir::OpPassManager* pass_manager); 38 39 // This is the early part of the conversion in isolation. This enables a caller 40 // to inject more information in the middle of the conversion before resuming it 41 // (like freezing variables for example). 42 void AddPreVariableFreezingTFToTFLConversionPasses( 43 const mlir::TFL::PassConfig& pass_config, 44 mlir::OpPassManager* pass_manager); 45 46 // This is the later part of the conversion in isolation. This enables a caller 47 // to resume the conversion after injecting more information in the middle of 48 // it. 49 void AddPostVariableFreezingTFToTFLConversionPasses( 50 llvm::StringRef saved_model_dir, const toco::TocoFlags& toco_flags, 51 const mlir::TFL::PassConfig& pass_config, 52 mlir::OpPassManager* pass_manager); 53 54 // Simplified API for TF->TFLite conversion with default flags. 55 void AddTFToTFLConversionPasses(const mlir::TFL::PassConfig& pass_config, 56 mlir::OpPassManager* pass_manager); 57 58 // Add the Quantization passes, specified in the quant_specs, into a pass 59 // manager. 60 void AddQuantizationPasses(const mlir::quant::QuantizationSpecs& quant_specs, 61 mlir::OpPassManager& pass_manager); 62 63 // Add the DynamicRangeQuantization passes, specified in the quant_specs, into a 64 // pass manager. 65 void AddDynamicRangeQuantizationPasses( 66 const mlir::quant::QuantizationSpecs& quant_specs, 67 mlir::OpPassManager& pass_manager); 68 } // namespace tensorflow 69 70 #endif // TENSORFLOW_COMPILER_MLIR_LITE_TF_TFL_PASSES_H_ 71