1 /* Copyright 2018 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_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ 17 #define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ 18 19 #include <string> 20 21 #include "tensorflow/compiler/tf2tensorrt/convert/utils.h" 22 #include "tensorflow/core/framework/graph.pb.h" 23 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h" 24 #include "tensorflow/core/platform/logging.h" 25 26 #if GOOGLE_CUDA && GOOGLE_TENSORRT 27 28 #if !IS_TRT_VERSION_GE(7, 0, 0, 0) 29 #error From version 2.6, we only support NVIDIA TensorRT version 7 or newer. 30 #error Please update your environment and relaunch the compilation. 31 #endif 32 33 namespace tensorflow { 34 namespace tensorrt { 35 namespace convert { 36 37 class TRTOptimizationPass : public grappler::CustomGraphOptimizer { 38 public: 39 TRTOptimizationPass(const string& name = "TRTOptimizationPass") name_(name)40 : name_(name), 41 trt_logger_name_("DefaultLogger"), 42 minimum_segment_size_(3), 43 precision_mode_(TrtPrecisionMode::FP32), 44 maximum_batch_size_(-1), 45 is_dynamic_op_(false), 46 max_cached_batches_(1), 47 max_workspace_size_bytes_(256LL << 20), 48 use_calibration_(true), 49 use_implicit_batch_(true), 50 profile_strategy_(ProfileStrategy::kRange), 51 allow_build_at_runtime_(true) { 52 VLOG(1) << "Constructing " << name_; 53 } 54 name()55 string name() const override { return name_; }; 56 UsesFunctionLibrary()57 bool UsesFunctionLibrary() const override { return true; } 58 59 Status Init( 60 const RewriterConfig_CustomGraphOptimizer* config = nullptr) override; 61 62 Status Optimize(grappler::Cluster* cluster, 63 const grappler::GrapplerItem& item, 64 GraphDef* optimized_graph) override; 65 66 void PrintDebugInfo(grappler::Cluster* cluster, 67 const grappler::GrapplerItem& item); 68 69 private: 70 const string name_; 71 string trt_logger_name_; 72 int minimum_segment_size_; 73 TrtPrecisionMode precision_mode_; 74 int maximum_batch_size_; 75 bool is_dynamic_op_; 76 std::vector<int> batches_; 77 int max_cached_batches_; 78 int64_t max_workspace_size_bytes_; 79 bool use_calibration_; 80 bool use_implicit_batch_; 81 ProfileStrategy profile_strategy_; 82 bool allow_build_at_runtime_; 83 }; 84 85 } // namespace convert 86 } // namespace tensorrt 87 } // namespace tensorflow 88 89 #endif // GOOGLE_CUDA && GOOGLE_TENSORRT 90 #endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_ 91