fix: Update logging scheme for converters

- Add debug statements for converter registrations in both FX and Dynamo
- Print operator counts along with registry support message

Author: gs-olive

py/torch_tensorrt/dynamo/__init__.py

@@ -3,9 +3,9 @@
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from ._settings import *
-    from .compile import compile
     from .aten_tracer import trace
     from .converter_registry import (
         DYNAMO_CONVERTERS,
         dynamo_tensorrt_converter,
     )
+    from .compile import compile

py/torch_tensorrt/dynamo/common_utils/converter_utils.py renamed to py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -1,3 +1,4 @@
+import logging
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, Optional, Sequence, Union
 from enum import Enum, auto
@@ -6,6 +7,9 @@
 from torch_tensorrt.fx.converter_registry import CONVERTERS
 
 
+logger = logging.getLogger(__name__)
+
+
 class ConverterPriority(Enum):
     """Enum to set a converter's priority in the registry"""
 
@@ -85,6 +89,10 @@ def register_converter(converter):
         else:
             DYNAMO_ATEN_CONVERTERS[key] = [converter_support]
 
+        logger.debug(
+            f"Converter for {key} added to Dynamo ATen Converter Registry with priority: {priority}"
+        )
+
         return converter
 
     def disable_converter(converter):

py/torch_tensorrt/dynamo/converter_registry.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Callable, Dict, List, Optional, Sequence, Set
+from typing import Dict, List, Optional, Sequence, Set
 
 import torch
 
@@ -55,10 +55,6 @@ def __init__(
         )
 
         self.min_block_size = min_block_size
-        logger.debug(
-            "Initialized Capability-Based Partitioner with available Converters:\n"
-            + f"{CONVERTERS.display_all_available_converters()}"
-        )
 
     def propose_partitions(self) -> List[Partition]:
         # Propose partitions using the default, then refine the results
@@ -114,8 +110,8 @@ def __init__(self, support_dict=None, torch_executed_ops=set()):
         super().__init__(support_dict)
 
         # Initialize sets of supported/unsupported operators
-        self.supported_operators = set()
-        self.unsupported_operators = set()
+        self.supported_operators = {}
+        self.unsupported_operators = {}
         self.torch_executed_ops = torch_executed_ops
 
     def is_node_supported(
@@ -130,12 +126,18 @@ def is_node_supported(
         if node in CONVERTERS and node_name not in self.torch_executed_ops:
             # If node is a proper, supported computational node, store the operator
             if not node.is_impure():
-                self.supported_operators.add(node_name)
+                if node_name not in self.supported_operators:
+                    self.supported_operators[node_name] = 1
+                else:
+                    self.supported_operators[node_name] += 1
 
             return True
         else:
             if not node.is_impure():
-                self.unsupported_operators.add(node_name)
+                if node_name not in self.unsupported_operators:
+                    self.unsupported_operators[node_name] = 1
+                else:
+                    self.unsupported_operators[node_name] += 1
 
             return False
 
@@ -147,15 +149,16 @@ def print_support_overview(self, num_trt_blocks: Optional[int] = None):
 
         # Reformat support messages for debugger to print node overview as a single string
         supported_nodes_str = "\nSupported Nodes:\n"
-        for node_name in self.supported_operators:
-            supported_nodes_str += f"- {node_name}\n"
+        for node_name, count in self.supported_operators.items():
+            supported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
 
         logger.debug(supported_nodes_str)
 
-        if len(self.unsupported_operators) != 0:
+        if self.unsupported_operators:
             unsupported_nodes_str = "\nUnsupported or Excluded Nodes:\n"
-            for node_name in self.unsupported_operators:
-                unsupported_nodes_str += f"- {node_name}\n"
+            for node_name, count in self.unsupported_operators.items():
+                unsupported_nodes_str += f"- {node_name} + Operator Count: {count}\n"
+
             logger.debug(unsupported_nodes_str)
         else:
             logger.debug("\nAll Nodes Supported\n")

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Any, Callable, Dict
 
 from torch.fx.node import Target
@@ -7,6 +8,9 @@
 NO_EXPLICIT_BATCH_DIM_SUPPORT = {}
 
 
+logger = logging.getLogger(__name__)
+
+
 def tensorrt_converter(
     key: Target,
     no_implicit_batch_dim: bool = False,
@@ -19,6 +23,13 @@ def register_converter(converter):
             NO_IMPLICIT_BATCH_DIM_SUPPORT[key] = converter
         if no_explicit_batch_dim:
             NO_EXPLICIT_BATCH_DIM_SUPPORT[key] = converter
+
+        logger.debug(
+            f"Converter for {key} added to FX Converter Registry "
+            + f"{'without' if no_explicit_batch_dim else 'with'} Explicit Batch Dim Support + "
+            + f"{'without' if no_implicit_batch_dim else 'with'} Implicit Batch Dim Support"
+        )
+
         return converter
 
     def disable_converter(converter):