Fixed the comments, changed dump function

Author: cehongwang

core/runtime/TRTEngine.cpp

@@ -283,14 +283,12 @@ void TRTEngine::enable_profiling() {
 
 void TRTEngine::set_profile_format(std::string format) {
   if (format == "trex") {
-    profile_format = TraceFormat::kTREX;
+    this->trt_engine_profiler->set_profile_format(TraceFormat::kTREX);
   } else if (format == "perfetto") {
-    profile_format = TraceFormat::kPERFETTO;
+    this->trt_engine_profiler->set_profile_format(TraceFormat::kPERFETTO);
   } else {
     TORCHTRT_THROW_ERROR("Invalid profile format: " + format);
   }
-
-  profile_format = profile_format;
 }
 
 std::string TRTEngine::get_engine_layer_info() {

core/runtime/TRTEngine.h

@@ -192,7 +192,6 @@ struct TRTEngine : torch::CustomClassHolder {
 #else
   bool profile_execution = false;
 #endif
-  TraceFormat profile_format = TraceFormat::kPERFETTO;
   std::string device_profile_path;
   std::string input_profile_path;
   std::string output_profile_path;

core/runtime/TRTEngineProfiler.cpp

@@ -32,7 +32,11 @@ TRTEngineProfiler::TRTEngineProfiler(const std::string& name, const std::vector<
   }
 }
 
-void dump_trace(const std::string& path, const TRTEngineProfiler& value, TraceFormat format) {
+void TRTEngineProfiler::set_profile_format(TraceFormat format) {
+  this->profile_format = format;
+}
+
+void dump_trace(const std::string& path, const TRTEngineProfiler& value) {
   std::stringstream out;
   out << "[" << std::endl;
   double ts = 0.0;
@@ -48,17 +52,17 @@ void dump_trace(const std::string& path, const TRTEngineProfiler& value, TraceFo
 
     out << "  {" << std::endl;
     out << "    \"name\": \"" << layer_name << "\"," << std::endl;
-    if (format == kPERFETTO) {
+    if (value.profile_format == TraceFormat::kPERFETTO) {
       out << "    \"ph\": \"X\"," << std::endl;
       out << "    \"ts\": " << running_time * 1000 << "," << std::endl;
       out << "    \"dur\": " << elem.time * 1000 << "," << std::endl;
       out << "    \"tid\": 1," << std::endl;
       out << "    \"pid\": \"" << value.name << " Engine Execution\"," << std::endl;
+      out << "    \"args\": {}" << std::endl;
     } else { // kTREX
       out << "    \"timeMs\": " << elem.time << "," << std::endl;
       out << "    \"averageMs\": " << elem.time / elem.count << "," << std::endl;
       out << "    \"percentage\": " << (elem.time * 100.0 / ts) << "," << std::endl;
-      out << "    \"args\": {}" << std::endl;
     }
     out << "  }," << std::endl;
     running_time += elem.time;

core/runtime/TRTEngineProfiler.h

@@ -19,18 +19,19 @@ struct TRTEngineProfiler : public nvinfer1::IProfiler {
     float time{0};
     int count{0};
   };
-
+  void set_profile_format(TraceFormat format);
   virtual void reportLayerTime(const char* layerName, float ms) noexcept;
   TRTEngineProfiler(
       const std::string& name,
       const std::vector<TRTEngineProfiler>& srcProfilers = std::vector<TRTEngineProfiler>());
   friend std::ostream& operator<<(std::ostream& out, const TRTEngineProfiler& value);
-  friend void dump_trace(const std::string& path, const TRTEngineProfiler& value, TraceFormat format);
+  friend void dump_trace(const std::string& path, const TRTEngineProfiler& value);
 
  private:
   std::string name;
   std::vector<std::string> layer_names;
   std::map<std::string, Record> profile;
+  TraceFormat profile_format = TraceFormat::kPERFETTO;
 };
 
 } // namespace runtime

core/runtime/execute_engine.cpp

@@ -339,10 +339,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
 
     if (compiled_engine->profile_execution) {
       LOG_INFO(std::endl << *compiled_engine->trt_engine_profiler);
-      dump_trace(
-          compiled_engine->trt_engine_profile_path,
-          *compiled_engine->trt_engine_profiler,
-          compiled_engine->profile_format);
+      dump_trace(compiled_engine->trt_engine_profile_path, *compiled_engine->trt_engine_profiler);
       compiled_engine->dump_engine_layer_info();
     }
 
@@ -443,10 +440,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
 
     if (compiled_engine->profile_execution) {
       LOG_INFO(std::endl << *compiled_engine->trt_engine_profiler);
-      dump_trace(
-          compiled_engine->trt_engine_profile_path,
-          *compiled_engine->trt_engine_profiler,
-          compiled_engine->profile_format);
+      dump_trace(compiled_engine->trt_engine_profile_path, *compiled_engine->trt_engine_profiler);
       compiled_engine->dump_engine_layer_info();
     }
 

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -335,7 +335,9 @@ def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]:
         return tuple(outputs)
 
     def enable_profiling(
-        self, profiling_results_dir: Optional[str] = None, profile_format: str = "trex"
+        self,
+        profiling_results_dir: Optional[str] = None,
+        profile_format: str = "perfetto",
     ) -> None:
         """Enable the profiler to collect latency information about the execution of the engine