Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion py/torch_tensorrt/dynamo/_compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -1390,7 +1390,20 @@ def preserve_module_specs(
cpu_memory_budget=settings.cpu_memory_budget,
)

dryrun_tracker.unsupported_ops = supported_ops.unsupported_operators
dryrun_tracker.unsupported_ops = supported_ops.fallback_operators

if supported_ops.fallback_operators:
named = "; ".join(
f"{node_name} + Operator Count: {count} "
f"(Reasons: {', '.join(sorted(supported_ops.fallback_reasons[node_name]))})"
for node_name, count in sorted(supported_ops.fallback_operators.items())
)
logger.info(
"%d operator(s) will run in PyTorch: %s. "
"Compile with dryrun=True for the full report.",
len(supported_ops.fallback_operators),
named,
)

# The global partitioner leaves non-TRT nodes as-is
if not settings.use_fast_partitioner:
Expand Down
33 changes: 32 additions & 1 deletion py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import logging
from typing import Collection, Dict, List, Optional, Tuple
from typing import Collection, Dict, List, Optional, Set, Tuple

import torch
import torch.fx.passes.operator_support as ops
Expand Down Expand Up @@ -39,9 +39,22 @@ def __init__(self, torch_executed_ops: Collection[Target] = set()) -> None:
# Initialize sets of supported/unsupported operators
self.supported_operators: Dict[str, int] = {}
self.unsupported_operators: Dict[str, int] = {}
# Keep impure refusals out of the counters used to decide full support.
self.fallback_operators: Dict[str, int] = {}
self.fallback_reasons: Dict[str, Set[str]] = {}
self.torch_executed_ops = torch_executed_ops
self._non_target_device_cache: Dict[torch.fx.Node, bool] = {}

def _record_fallback(
self, node: torch.fx.Node, node_name: str, reason: str
) -> None:
# Structural nodes must not make a fully supported graph report fallback.
if node.op in CALLABLE_NODE_OPS:
self.fallback_operators[node_name] = (
self.fallback_operators.get(node_name, 0) + 1
)
self.fallback_reasons.setdefault(node_name, set()).add(reason)

def is_node_supported(
self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node
) -> bool:
Expand All @@ -65,6 +78,7 @@ def is_node_supported(
"non-target device region",
node_name,
)
self._record_fallback(node, node_name, "explicit non-target device region")
return False

if TorchTensorRTOperatorSupport._exceeds_max_tensor_rank(node):
Expand All @@ -73,6 +87,7 @@ def is_node_supported(
self.unsupported_operators[node_name] = (
self.unsupported_operators.get(node_name, 0) + 1
)
self._record_fallback(node, node_name, "tensor rank exceeds TensorRT limit")
return False

if TorchTensorRTOperatorSupport._has_complex_dtype(node):
Expand All @@ -81,6 +96,7 @@ def is_node_supported(
self.unsupported_operators[node_name] = (
self.unsupported_operators.get(node_name, 0) + 1
)
self._record_fallback(node, node_name, "complex tensor dtype")
return False

if (
Expand All @@ -94,6 +110,11 @@ def is_node_supported(
self.unsupported_operators[node_name] = (
self.unsupported_operators.get(node_name, 0) + 1
)
self._record_fallback(
node,
node_name,
"data-dependent output shape (fallback_data_dependent_ops=True)",
)
return False

if (
Expand All @@ -116,6 +137,16 @@ def is_node_supported(
else:
self.unsupported_operators[node_name] += 1

self._record_fallback(
node,
node_name,
(
"excluded by torch_executed_ops"
if node_name in self.torch_executed_ops
or node.target in self.torch_executed_ops
else "no validated TensorRT converter"
),
)
return False

def print_support_overview(self, num_trt_blocks: Optional[int] = None) -> None:
Expand Down
34 changes: 33 additions & 1 deletion py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
import logging
from typing import Collection, Dict, List, Mapping, Optional, Sequence, Tuple
from typing import Collection, Dict, List, Mapping, Optional, Sequence, Set, Tuple

import tensorrt as trt
import torch
from torch.fx.graph_module import GraphModule
from torch.fx.node import Target
from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
from torch.fx.passes.operator_support import OperatorSupport, SupportDict
from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
from torch.utils._pytree import tree_flatten
from torch_tensorrt.dynamo._defaults import (
MIN_BLOCK_SIZE,
Expand Down Expand Up @@ -145,6 +146,9 @@ def __init__(
# Initialize sets of supported/unsupported operators
self.supported_operators: Dict[str, int] = {}
self.unsupported_operators: Dict[str, int] = {}
# Keep impure refusals out of the counters used to decide full support.
self.fallback_operators: Dict[str, int] = {}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be nice to include some metadata about reasons for fallback?

self.fallback_reasons: Dict[str, Set[str]] = {}
self.torch_executed_ops: Collection[Target] = torch_executed_ops
self._non_target_device_cache: Dict[torch.fx.Node, bool] = {}

Expand Down Expand Up @@ -254,6 +258,16 @@ def _requires_output_allocator(node: torch.fx.Node) -> bool:
"requires_output_allocator", False
)

def _record_fallback(
self, node: torch.fx.Node, node_name: str, reason: str
) -> None:
# Structural nodes must not make a fully supported graph report fallback.
if node.op in CALLABLE_NODE_OPS:
self.fallback_operators[node_name] = (
self.fallback_operators.get(node_name, 0) + 1
)
self.fallback_reasons.setdefault(node_name, set()).add(reason)

def is_node_supported(
self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node
) -> bool:
Expand All @@ -274,6 +288,7 @@ def is_node_supported(
"non-target device region",
node_name,
)
self._record_fallback(node, node_name, "explicit non-target device region")
return False

if self._exceeds_max_tensor_rank(node):
Expand All @@ -282,6 +297,7 @@ def is_node_supported(
self.unsupported_operators[node_name] = (
self.unsupported_operators.get(node_name, 0) + 1
)
self._record_fallback(node, node_name, "tensor rank exceeds TensorRT limit")
return False

if self._has_complex_dtype(node):
Expand All @@ -291,6 +307,7 @@ def is_node_supported(
self.unsupported_operators[node_name] = (
self.unsupported_operators.get(node_name, 0) + 1
)
self._record_fallback(node, node_name, "complex tensor dtype")
return False

if (
Expand All @@ -304,6 +321,11 @@ def is_node_supported(
self.unsupported_operators[node_name] = (
self.unsupported_operators.get(node_name, 0) + 1
)
self._record_fallback(
node,
node_name,
"data-dependent output shape (fallback_data_dependent_ops=True)",
)
return False

if (
Expand All @@ -326,6 +348,16 @@ def is_node_supported(
else:
self.unsupported_operators[node_name] += 1

self._record_fallback(
node,
node_name,
(
"excluded by torch_executed_ops"
if node_name in self.torch_executed_ops
or node.target in self.torch_executed_ops
else "no validated TensorRT converter"
),
)
return False

def print_support_overview(
Expand Down
166 changes: 166 additions & 0 deletions tests/py/dynamo/partitioning/test_000_fallback_reasons.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,166 @@
from unittest.mock import patch

import tensorrt as trt
import torch
from parameterized import parameterized
from torch.testing._internal.common_utils import TestCase, run_tests
from torch_tensorrt.dynamo._settings import CompilationSettings
from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
DYNAMO_CONVERTERS,
ConverterRegistry,
ConverterSupport,
)
from torch_tensorrt.dynamo.partitioning._adjacency_partitioner import OpSupportTester
from torch_tensorrt.dynamo.partitioning._global_partitioner import (
TorchTensorRTOperatorSupport,
)

SUPPORT_CLASSES = [
("fast", OpSupportTester),
("global", TorchTensorRTOperatorSupport),
]


class TestFallbackReasons(TestCase):
def setUp(self):
super().setUp()
for name, value in (
("compilation_settings", CompilationSettings()),
("disallowed_targets", set()),
):
patcher = patch.object(DYNAMO_CONVERTERS, name, value)
patcher.start()
self.addCleanup(patcher.stop)

@staticmethod
def _node(target, input_value, output_value):
graph = torch.fx.Graph()
x = graph.placeholder("x")
x.meta["val"] = input_value
node = graph.call_function(target, (x,))
node.meta["val"] = output_value
graph.output(node)
return node

@parameterized.expand(SUPPORT_CLASSES)
def test_refusal_reasons(self, _, support_class):
x = torch.empty(2, device="cuda")
high_rank = torch.empty((1,) * (trt.Dims.MAX_DIMS + 1), device="cuda")
complex_value = torch.empty(2, dtype=torch.complex64, device="cuda")
cases = [
(
torch.ops.aten._to_copy.default,
x,
x.cpu(),
"explicit non-target device region",
),
(
torch.ops.aten.clone.default,
high_rank,
high_rank,
"tensor rank exceeds TensorRT limit",
),
(
torch.ops.aten.clone.default,
complex_value,
complex_value,
"complex tensor dtype",
),
(
torch.ops.aten.nonzero.default,
x,
torch.empty(1, 1, dtype=torch.int64, device="cuda"),
"data-dependent output shape (fallback_data_dependent_ops=True)",
),
(torch.ops.aten.rand_like.default, x, x, "no validated TensorRT converter"),
]
DYNAMO_CONVERTERS.compilation_settings.fallback_data_dependent_ops = True
for target, input_value, output_value, reason in cases:
with self.subTest(reason=reason):
support = support_class()
node = self._node(target, input_value, output_value)
name = ConverterRegistry.qualified_name_or_str(target)
if target == torch.ops.aten.nonzero.default:
self.assertTrue(
DYNAMO_CONVERTERS[node][2]["requires_output_allocator"]
)
self.assertFalse(support.is_node_supported({}, node))
self.assertEqual(support.fallback_operators, {name: 1})
self.assertEqual(support.fallback_reasons, {name: {reason}})
if target == torch.ops.aten.rand_like.default:
self.assertTrue(node.is_impure())
self.assertEqual(support.unsupported_operators, {})

@parameterized.expand(SUPPORT_CLASSES)
def test_requested_fallback(self, _, support_class):
x = torch.empty(2, device="cuda")
target = torch.ops.aten.relu.default
node = self._node(target, x, x)
name = ConverterRegistry.qualified_name_or_str(target)
self.assertIn(node, DYNAMO_CONVERTERS)
for excluded in (name, target):
with self.subTest(excluded=excluded):
support = support_class(torch_executed_ops={excluded})
self.assertFalse(support.is_node_supported({}, node))
self.assertEqual(support.fallback_operators, {name: 1})
self.assertEqual(
support.fallback_reasons,
{name: {"excluded by torch_executed_ops"}},
)

@parameterized.expand(SUPPORT_CLASSES)
def test_rejected_converter_is_not_reported_as_missing(self, _, support_class):
x = torch.empty(2, device="cuda")
target = torch.ops.aten.clone.default
node = self._node(target, x, x)
converters = {
target: [
ConverterSupport(
converter_implementation=lambda *args: None,
capability_validator=lambda node, settings: False,
)
]
}
with patch.object(DYNAMO_CONVERTERS, "registries", [converters]):
self.assertIsNotNone(DYNAMO_CONVERTERS.get_unvalidated(target))
support = support_class()
self.assertFalse(support.is_node_supported({}, node))
name = ConverterRegistry.qualified_name_or_str(target)
self.assertEqual(
support.fallback_reasons,
{name: {"no validated TensorRT converter"}},
)

@parameterized.expand(SUPPORT_CLASSES)
def test_same_operator_keeps_multiple_reasons(self, _, support_class):
support = support_class()
target = torch.ops.aten.clone.default
for value in (
torch.empty((1,) * (trt.Dims.MAX_DIMS + 1), device="cuda"),
torch.empty(2, dtype=torch.complex64, device="cuda"),
):
self.assertFalse(
support.is_node_supported({}, self._node(target, value, value))
)
name = ConverterRegistry.qualified_name_or_str(target)
self.assertEqual(support.fallback_operators, {name: 2})
self.assertEqual(
support.fallback_reasons,
{name: {"tensor rank exceeds TensorRT limit", "complex tensor dtype"}},
)

@parameterized.expand(SUPPORT_CLASSES)
def test_supported_and_structural_nodes_have_no_fallback(self, _, support_class):
x = torch.empty(2, device="cuda")
node = self._node(torch.ops.aten.relu.default, x, x)
support = support_class()
self.assertTrue(support.is_node_supported({}, node))
for structural in node.graph.nodes:
if structural.op in ("placeholder", "output"):
support.is_node_supported({}, structural)
self.assertEqual(support.fallback_operators, {})
self.assertEqual(support.fallback_reasons, {})


if __name__ == "__main__":
run_tests()
Loading
Loading