[TensorRT] Optimize performance via only place TRT ops to gpu device. (#860)

Signed-off-by: Tao Peng <[email protected]>

Author: shanshanpt

tensorflow/compiler/tf2tensorrt/tool/tf2trt.py

@@ -63,7 +63,7 @@ def GetConversionParams(run_params):
       rewriter_config_template=None,
       max_workspace_size_bytes=1 << 25, 
       precision_mode=run_params['precision_mode'],
-      minimum_segment_size=2,
+      minimum_segment_size=run_params['minimum_segment_size'],
       is_dynamic_op=run_params['dynamic_engine'],
       maximum_cached_engines=1,
       use_calibration=run_params['use_calibration'],
@@ -96,6 +96,7 @@ def ConvertGraph(run_params, saved_model_dir, trt_saved_model_dir):
     'use_calibration': False,
     'max_batch_size': 1024,
     'convert_online': False,
+    'minimum_segment_size': 4,
     'use_ev': True,
   }
 

tensorflow/core/common_runtime/placer.cc

@@ -28,6 +28,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/util/env_var.h"
 #include "tensorflow/core/util/dump_graph.h"
 #include "tensorflow/core/util/port.h"
 
@@ -230,6 +231,29 @@ Status Placer::Run() {
     colocation_graph.AssignGpuStreamIdx(node);    
   }
 
+  bool place_trtop_on_gpu_only = false;
+  TF_CHECK_OK(tensorflow::ReadBoolFromEnvVar(
+      "PLACE_TRT_OP_ON_GPU_ONLY", false, &place_trtop_on_gpu_only));
+  // Keep TRTEngineOp On GPU Only
+  if (place_trtop_on_gpu_only) {
+    std::string cpu_name, gpu_name;
+    for (auto d : devices_->devices()) {
+      if (d->name().find("device:CPU:") != std::string::npos) {
+        cpu_name = d->name();
+      } else if (d->name().find("device:GPU:") != std::string::npos) {
+        gpu_name = d->name();
+      }
+    }
+
+    for (Node* n : graph_->op_nodes()) {
+      if (n->type_string() == "TRTEngineOp") {
+        n->set_assigned_device_name(gpu_name);
+      } else {
+        n->set_assigned_device_name(cpu_name);
+      }
+    }
+  }
+
   if (VLOG_IS_ON(3)) {
     DumpGraphToFile("placer_output", *graph_, nullptr);
   }