Fixed the comments

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -66,7 +66,6 @@ def cross_compile_for_windows(
         Set[Union[torch.dtype, dtype]], Tuple[Union[torch.dtype, dtype]]
     ] = _defaults.ENABLED_PRECISIONS,
     engine_capability: EngineCapability = _defaults.ENGINE_CAPABILITY,
-    debug: bool = False,
     num_avg_timing_iters: int = _defaults.NUM_AVG_TIMING_ITERS,
     workspace_size: int = _defaults.WORKSPACE_SIZE,
     dla_sram_size: int = _defaults.DLA_SRAM_SIZE,
@@ -187,7 +186,7 @@ def cross_compile_for_windows(
             f"Cross compile for windows is only supported on x86-64 Linux architecture, current platform: {platform.system()=}, {platform.architecture()[0]=}"
         )
 
-    if debug:
+    if kwargs.get("debug", False):
         warnings.warn(
             "`debug` is deprecated. Please use `torch_tensorrt.dynamo.Debugger` to configure debugging options.",
             DeprecationWarning,
@@ -404,7 +403,6 @@ def compile(
         Set[Union[torch.dtype, dtype]], Tuple[Union[torch.dtype, dtype]]
     ] = _defaults.ENABLED_PRECISIONS,
     engine_capability: EngineCapability = _defaults.ENGINE_CAPABILITY,
-    debug: bool = False,
     num_avg_timing_iters: int = _defaults.NUM_AVG_TIMING_ITERS,
     workspace_size: int = _defaults.WORKSPACE_SIZE,
     dla_sram_size: int = _defaults.DLA_SRAM_SIZE,
@@ -523,7 +521,7 @@ def compile(
         torch.fx.GraphModule: Compiled FX Module, when run it will execute via TensorRT
     """
 
-    if debug:
+    if kwargs.get("debug", False):
         warnings.warn(
             "`debug` is deprecated. Please use `torch_tensorrt.dynamo.Debugger` for debugging functionality",
             DeprecationWarning,
@@ -946,7 +944,8 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
                             trt_module.enable_profiling()
                     else:
                         path = os.path.join(
-                            _debugger_settings.logging_dir, "engine_visualization"
+                            _debugger_settings.logging_dir,
+                            "engine_visualization_profile",
                         )
                         os.makedirs(path, exist_ok=True)
                         trt_module.enable_profiling(
@@ -990,7 +989,6 @@ def convert_exported_program_to_serialized_trt_engine(
     enabled_precisions: (
         Set[torch.dtype | dtype] | Tuple[torch.dtype | dtype]
     ) = _defaults.ENABLED_PRECISIONS,
-    debug: bool = False,
     assume_dynamic_shape_support: bool = _defaults.ASSUME_DYNAMIC_SHAPE_SUPPORT,
     workspace_size: int = _defaults.WORKSPACE_SIZE,
     min_block_size: int = _defaults.MIN_BLOCK_SIZE,
@@ -1092,7 +1090,7 @@ def convert_exported_program_to_serialized_trt_engine(
     Returns:
         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
     """
-    if debug:
+    if kwargs.get("debug", False):
         warnings.warn(
             "`debug` is deprecated. Please use `torch_tensorrt.dynamo.Debugger` to configure debugging options.",
             DeprecationWarning,
@@ -1181,7 +1179,6 @@ def convert_exported_program_to_serialized_trt_engine(
     compilation_options = {
         "assume_dynamic_shape_support": assume_dynamic_shape_support,
         "enabled_precisions": enabled_precisions,
-        "debug": debug,
         "workspace_size": workspace_size,
         "min_block_size": min_block_size,
         "torch_executed_ops": torch_executed_ops,

py/torch_tensorrt/dynamo/_defaults.py

@@ -49,6 +49,7 @@
 L2_LIMIT_FOR_TILING = -1
 USE_DISTRIBUTED_MODE_TRACE = False
 OFFLOAD_MODULE_TO_CPU = False
+DEBUG_LOGGING_DIR = os.path.join(tempfile.gettempdir(), "torch_tensorrt/debug_logs")
 
 
 def default_device() -> Device:

py/torch_tensorrt/dynamo/debug/_Debugger.py

@@ -8,6 +8,7 @@
 from unittest import mock
 
 import torch
+from torch_tensorrt.dynamo._defaults import DEBUG_LOGGING_DIR
 from torch_tensorrt.dynamo.debug._DebuggerConfig import DebuggerConfig
 from torch_tensorrt.dynamo.debug._supports_debugger import (
     _DEBUG_ENABLED_CLS,
@@ -32,7 +33,7 @@ def __init__(
         save_engine_profile: bool = False,
         profile_format: str = "perfetto",
         engine_builder_monitor: bool = True,
-        logging_dir: str = tempfile.gettempdir(),
+        logging_dir: str = DEBUG_LOGGING_DIR,
         save_layer_info: bool = False,
     ):
         """Initialize a debugger for TensorRT conversion.

py/torch_tensorrt/dynamo/debug/_DebuggerConfig.py

@@ -1,12 +1,13 @@
-import tempfile
 from dataclasses import dataclass
 
+from torch_tensorrt.dynamo._defaults import DEBUG_LOGGING_DIR
+
 
 @dataclass
 class DebuggerConfig:
     log_level: str = "debug"
     save_engine_profile: bool = False
     engine_builder_monitor: bool = True
-    logging_dir: str = tempfile.gettempdir()
+    logging_dir: str = DEBUG_LOGGING_DIR
     profile_format: str = "perfetto"
     save_layer_info: bool = False

tools/debug/engine_visualization/README.md

@@ -7,3 +7,5 @@ pip install git+https://github.com/NVIDIA/TensorRT.git#subdirectory=tools/experi
 sudo apt --yes install graphviz
 ```
 
+## Usage
+The example usage can be found in `draw_engine_graph_example.py`. We use `torch_tensorrt.dynamo.debugger` to first output the engine profile info that required by TREX. Note that only when the compilation settings `use_python_runtime=False` can it produce TREX profiling. When it is saved to a folder, we call `draw_engine` on the same directory where the profile files are saved, which is in the subdirectory `engine_visualization_profile`.

tools/debug/engine_visualization/draw_engine_graph_example.py

@@ -5,6 +5,7 @@
 import torch
 import torch_tensorrt as torch_tensorrt
 import torchvision.models as models
+from torch_tensorrt.dynamo._defaults import DEBUG_LOGGING_DIR
 
 inputs = [torch.rand((1, 3, 224, 224)).to("cuda")]
 model = models.resnet18(pretrained=False).eval().to("cuda")
@@ -14,12 +15,13 @@
 # min_block_size = 0
 use_python_runtime = False
 torch_executed_ops = {}
-logging_dir = "/home/profile"
+
 with torch_tensorrt.dynamo.Debugger(
     "graphs",
-    logging_dir=logging_dir,
+    logging_dir=DEBUG_LOGGING_DIR,
     capture_fx_graph_after=["constant_fold"],
     save_engine_profile=True,
+    profile_format="trex",
 ):
     trt_gm = torch_tensorrt.dynamo.compile(
         exp_program,
@@ -32,5 +34,5 @@
 
     from draw_engine_graph import draw_engine
 
-    draw_engine(os.path.join(logging_dir, "engine_visualization"))
+    draw_engine(os.path.join(DEBUG_LOGGING_DIR, "engine_visualization_profile"))
 print()