diff --git a/py/torch_tensorrt/dynamo/conversion/_SubgraphInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_SubgraphInterpreter.py new file mode 100644 index 00000000000..debd1a4509d --- /dev/null +++ b/py/torch_tensorrt/dynamo/conversion/_SubgraphInterpreter.py @@ -0,0 +1,122 @@ +import logging +from typing import Any, Optional, Sequence, Tuple + +import torch +from torch.fx.experimental.proxy_tensor import unset_fake_temporarily +from torch.utils._python_dispatch import _disable_current_modes +from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext +from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( + DYNAMO_CONVERTERS as CONVERTERS, +) +from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( + CallingConvention, +) +from torch_tensorrt.dynamo.conversion._TRTInterpreter import ( + UnsupportedOperatorException, +) +from torch_tensorrt.dynamo.conversion.converter_utils import get_node_name, to_torch + +_LOGGER = logging.getLogger(__name__) + + +class TRTSubgraphInterpreter(torch.fx.Interpreter): # type: ignore[misc] + """Convert an FX GraphModule into an existing TensorRT network. + + Unlike ``TRTInterpreter``, this does not create a builder, network, or + engine I/O bindings. Placeholders are bound to caller-provided values + (typically ``IIfConditionalInputLayer`` outputs) via ``Interpreter.run``. + """ + + def __init__( + self, + module: torch.fx.GraphModule, + ctx: ConversionContext, + name_prefix: str, + ) -> None: + super().__init__(module) + self.ctx = ctx + self.name_prefix = name_prefix + self._cur_node: Optional[torch.fx.Node] = None + self._cur_node_name: Optional[str] = None + + def run_node(self, n: torch.fx.Node) -> Any: + prev = self.ctx.current_node + self._cur_node = n + self._cur_node_name = f"{self.name_prefix}/{get_node_name(n)}" + self.ctx.current_node = n + try: + if _LOGGER.isEnabledFor(logging.DEBUG): + _LOGGER.debug( + "Converting cond-subgraph node %s (kind: %s)", + self._cur_node_name, + n.target, + ) + return super().run_node(n) + finally: + self.ctx.current_node = prev + + def get_attr(self, target: str, args: Any, kwargs: Any) -> Any: + del args, kwargs + with _disable_current_modes(), unset_fake_temporarily(): + attr = self.fetch_attr(target) + if isinstance(attr, torch.nn.Module): + return attr + if isinstance(attr, torch.nn.Parameter): + attr = attr.data + return to_torch(attr) + + def call_function(self, target: Any, args: Any, kwargs: Any) -> Any: + converter_packet = CONVERTERS.get(self._cur_node) + if converter_packet is None: + raise UnsupportedOperatorException( + f"Conversion of function {torch.typename(target)} not currently supported " + f"inside torch.cond subgraph '{self.name_prefix}'" + ) + + converter, calling_convention, converter_info = converter_packet + if converter_info.get("requires_output_allocator", False): + self.ctx.requires_output_allocator = True + _LOGGER.debug("%s requires output allocator", target) + if converter_info.get("requires_native_multidevice", False): + self.ctx.requires_native_multidevice = True + _LOGGER.debug("%s requires native multi-device support", target) + + if calling_convention is CallingConvention.LEGACY: + return converter(self.ctx.net, target, args, kwargs, self._cur_node_name) + return converter(self.ctx, target, args, kwargs, self._cur_node_name) + + def call_method(self, target: str, args: Any, kwargs: Any) -> Any: + converter_packet = CONVERTERS.get(self._cur_node) + if converter_packet is None: + raise UnsupportedOperatorException( + f"Conversion of method {target} not currently supported " + f"inside torch.cond subgraph '{self.name_prefix}'" + ) + converter, calling_convention, _ = converter_packet + if calling_convention is CallingConvention.LEGACY: + return converter(self.ctx.net, target, args, kwargs, self._cur_node_name) + return converter(self.ctx, target, args, kwargs, self._cur_node_name) + + def call_module(self, target: str, args: Any, kwargs: Any) -> Any: + del args, kwargs + raise UnsupportedOperatorException( + f"call_module '{target}' is not supported inside torch.cond subgraphs" + ) + + +def convert_subgraph( + ctx: ConversionContext, + gm: torch.fx.GraphModule, + operands: Sequence[Any], + name_prefix: str, +) -> Tuple[Any, ...]: + """Convert ``gm`` with ``operands`` bound to its placeholders. + + Returns the subgraph outputs as a tuple, matching torch.cond's convention + that branch graphs return a tuple even for a single tensor. + """ + interp = TRTSubgraphInterpreter(gm, ctx, name_prefix) + outputs = interp.run(*operands) + if not isinstance(outputs, (list, tuple)): + return (outputs,) + return tuple(outputs) diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py index 448d96a0ce4..eed5bb5d6a2 100644 --- a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py +++ b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py @@ -15,6 +15,7 @@ ) import numpy as np +import tensorrt as trt import torch import torch.fx from torch.fx.experimental.proxy_tensor import unset_fake_temporarily @@ -52,8 +53,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - _LOGGER: logging.Logger = logging.getLogger(__name__) TRT_INTERPRETER_CALL_PRE_OBSERVER: Observer[Callable[[torch.fx.GraphModule], None]] = ( @@ -542,7 +541,7 @@ def run_node(self, n: torch.fx.Node) -> torch.fx.Node: trt_node: torch.fx.Node = super().run_node(n) - if n.op == "get_attr": + if n.op == "get_attr" and isinstance(trt_node, torch.Tensor): self.const_mapping[str(n)] = (tuple(trt_node.shape), str(trt_node.dtype)) _LOGGER.info( @@ -713,9 +712,13 @@ def call_function(self, target: str, args: Any, kwargs: Any) -> Any: else: return converter(self.ctx, target, args, kwargs, self._cur_node_name) - def get_attr(self, target: str, args: Any, kwargs: Any) -> torch.Tensor: + def get_attr(self, target: str, args: Any, kwargs: Any) -> Any: with _disable_current_modes(), unset_fake_temporarily(): frozen_attr = self.fetch_attr(target) + # Cond (and other higher-order ops) store branch graphs as module + # attributes. Those must be passed through to the converter. + if isinstance(frozen_attr, torch.nn.Module): + return frozen_attr if isinstance(frozen_attr, torch.nn.Parameter): constant_tensor = frozen_attr.data else: diff --git a/py/torch_tensorrt/dynamo/conversion/__init__.py b/py/torch_tensorrt/dynamo/conversion/__init__.py index f41540a0cfb..1e1b204d13f 100644 --- a/py/torch_tensorrt/dynamo/conversion/__init__.py +++ b/py/torch_tensorrt/dynamo/conversion/__init__.py @@ -1,6 +1,7 @@ from . import ( aten_ops_converters, custom_ops_converters, + higher_order_ops_converters, ops_evaluators, plugins, prims_ops_converters, diff --git a/py/torch_tensorrt/dynamo/conversion/converter_utils.py b/py/torch_tensorrt/dynamo/conversion/converter_utils.py index 7a8b23ed4cc..721f4cf5e3e 100644 --- a/py/torch_tensorrt/dynamo/conversion/converter_utils.py +++ b/py/torch_tensorrt/dynamo/conversion/converter_utils.py @@ -94,8 +94,12 @@ def format_tensor_metadata(metadata: Union[Any, Sequence[Any]]) -> str: for arg in node.args: if isinstance(arg, torch.fx.Node): if arg.op == "get_attr": - shape, dtype = constant_mapping[str(arg)] - arg_repr = f"{shape}@{dtype}" + mapped = constant_mapping.get(str(arg)) + arg_repr = ( + f"{mapped[0]}@{mapped[1]}" + if mapped is not None + else f"attr:{arg.target}" + ) elif arg.meta.get("tensor_meta") is not None: arg_repr = format_tensor_metadata(arg.meta["tensor_meta"]) elif arg.meta.get("val") is not None: @@ -114,8 +118,10 @@ def format_tensor_metadata(metadata: Union[Any, Sequence[Any]]) -> str: # Format output tensors and arguments metadata_string += " | Outputs: (" if node.op == "get_attr": - shape, dtype = constant_mapping[str(node)] - node_repr = f"{shape}@{dtype}" + mapped = constant_mapping.get(str(node)) + node_repr = ( + f"{mapped[0]}@{mapped[1]}" if mapped is not None else f"attr:{node.target}" + ) elif node.meta.get("tensor_meta") is not None: node_repr = format_tensor_metadata(node.meta["tensor_meta"]) elif node.meta.get("val") is not None: diff --git a/py/torch_tensorrt/dynamo/conversion/higher_order_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/higher_order_ops_converters.py new file mode 100644 index 00000000000..557cbc923db --- /dev/null +++ b/py/torch_tensorrt/dynamo/conversion/higher_order_ops_converters.py @@ -0,0 +1,113 @@ +# mypy: disallow-untyped-decorators=False + +import logging +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import torch +from tensorrt import ITensor as TRTTensor +from torch.fx.node import Argument, Node, Target +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo._SourceIR import SourceIR +from torch_tensorrt.dynamo.conversion import impl +from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext +from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( + DYNAMO_CONVERTERS, + dynamo_tensorrt_converter, +) + +_LOGGER = logging.getLogger(__name__) + + +def _fetch_attr(mod: torch.nn.Module, target: str) -> Any: + cur: Any = mod + for atom in target.split("."): + cur = getattr(cur, atom) + return cur + + +def _branch_modules(node: Node) -> Optional[List[torch.fx.GraphModule]]: + """Return the true/false GraphModules captured on a higher_order.cond node.""" + gm = node.graph.owning_module + if gm is None or len(node.args) < 3: + return None + branches: List[torch.fx.GraphModule] = [] + for arg in node.args[1:3]: + if not isinstance(arg, Node) or arg.op != "get_attr": + return None + try: + attr = _fetch_attr(gm, str(arg.target)) + except AttributeError: + return None + if not isinstance(attr, torch.fx.GraphModule): + return None + branches.append(attr) + return branches + + +def _subgraph_is_supported(gm: torch.fx.GraphModule) -> bool: + """True if every computational node in ``gm`` (and nested cond branches) has a converter.""" + for node in gm.graph.nodes: + if node.op in ("placeholder", "output"): + continue + if node.op == "get_attr": + try: + attr = _fetch_attr(gm, str(node.target)) + except AttributeError: + return False + if isinstance(attr, torch.fx.GraphModule) and not _subgraph_is_supported( + attr + ): + return False + continue + if node.op == "call_function": + if node not in DYNAMO_CONVERTERS: + _LOGGER.debug( + "torch.cond subgraph %s has unsupported op %s", + gm._get_name(), + node.target, + ) + return False + continue + _LOGGER.debug( + "torch.cond subgraph %s has unsupported node.op %s", + gm._get_name(), + node.op, + ) + return False + return True + + +def cond_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Support cond only when both branch graphs are fully TRT-convertible.""" + del settings + branches = _branch_modules(node) + if not branches: + return False + return all(_subgraph_is_supported(branch) for branch in branches) + + +@dynamo_tensorrt_converter( + torch.ops.higher_order.cond, + capability_validator=cond_capability_validator, + supports_dynamic_shapes=True, +) +def higher_order_ops_cond( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + del kwargs + return impl.condition.cond( + ctx, + target, + SourceIR.UNKNOWN, + name, + pred=args[0], + true_fn=args[1], + false_fn=args[2], + operands=args[3], + ) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py index b7739c3b3fd..a9bc8bb5220 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Any, Optional, Tuple, Union import numpy as np import tensorrt as trt @@ -16,6 +16,7 @@ set_layer_name, ) from torch_tensorrt.dynamo.conversion.impl.elementwise import ne +from torch_tensorrt.dynamo.conversion.impl.shuffle import reshape as reshape_tensor def where( @@ -76,3 +77,72 @@ def select( select_layer = ctx.net.add_select(condition, input, other) set_layer_name(select_layer, target, name + "_select", source_ir) return select_layer.get_output(0) + + +def _as_sequence(value: Any) -> list[Any]: + if isinstance(value, (list, tuple)): + return list(value) + return [value] + + +def cond( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + pred: Union[TRTTensor, torch.Tensor, bool], + true_fn: torch.fx.GraphModule, + false_fn: torch.fx.GraphModule, + operands: Any, +) -> Tuple[TRTTensor, ...]: + """Convert torch.ops.higher_order.cond to a TensorRT IIfConditional. + + Both branches consume the same ``IIfConditionalInputLayer`` tensors. Layers + created while converting ``true_fn`` / ``false_fn`` are associated with the + corresponding branch by TensorRT via the path from those inputs to + ``add_output``. + """ + if not isinstance(true_fn, torch.fx.GraphModule) or not isinstance( + false_fn, torch.fx.GraphModule + ): + raise RuntimeError( + f"{name}: torch.cond branches must be GraphModules, got " + f"{type(true_fn)} and {type(false_fn)}" + ) + + if not isinstance(pred, TRTTensor): + pred = get_trt_tensor(ctx, pred, f"{name}_pred", dtype=torch.bool, min_rank=0) + if pred.dtype != trt.bool: + pred = cast_trt_tensor( + ctx, pred, torch.bool, f"{name}_pred_bool", target, source_ir + ) + # TensorRT requires a 0-D boolean predicate. + if len(pred.shape) != 0: + pred = reshape_tensor(ctx, target, source_ir, f"{name}_pred_scalar", pred, []) + + conditional = ctx.net.add_if_conditional() + conditional.name = name + conditional.set_condition(pred) + + wrapped_operands = [] + for i, operand in enumerate(_as_sequence(operands)): + if not isinstance(operand, TRTTensor): + operand = get_trt_tensor(ctx, operand, f"{name}_operand_{i}") + wrapped_operands.append(conditional.add_input(operand).get_output(0)) + + from torch_tensorrt.dynamo.conversion._SubgraphInterpreter import convert_subgraph + + true_outs = convert_subgraph(ctx, true_fn, wrapped_operands, f"{name}_true") + false_outs = convert_subgraph(ctx, false_fn, wrapped_operands, f"{name}_false") + if len(true_outs) != len(false_outs): + raise RuntimeError( + f"{name}: cond branches return different numbers of tensors " + f"({len(true_outs)} vs {len(false_outs)})" + ) + + outputs: list[TRTTensor] = [] + for i, (t_out, f_out) in enumerate(zip(true_outs, false_outs)): + layer = conditional.add_output(t_out, f_out) + set_layer_name(layer, target, f"{name}_out_{i}", source_ir) + outputs.append(layer.get_output(0)) + return tuple(outputs) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index cc4ff89ec8b..c642d5f3223 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -10,6 +10,7 @@ mark_constant_fold_exclusions, ) from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + iter_cond_subgraphs, trace_intermediate_node_outputs, ) @@ -171,6 +172,9 @@ def post_lowering( if fake_mode is not None: fake_tensor_updater.incremental_update(fake_mode) + for sub in iter_cond_subgraphs(gm): + post_lowering(sub, settings) + return gm @@ -192,6 +196,8 @@ def pre_export_lowering( ) gm = ep.graph_module gm = ATEN_PRE_LOWERING_PASSES(gm, settings) + for sub in iter_cond_subgraphs(gm): + ATEN_PRE_LOWERING_PASSES(sub, settings) return ep diff --git a/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py b/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py index 2a51bbac6fd..830b02c5505 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py @@ -14,6 +14,35 @@ def clean_up_graph_after_modifications( return gm +def iter_cond_subgraphs(gm: torch.fx.GraphModule) -> List[torch.fx.GraphModule]: + """Return GraphModules used as torch.cond true/false branches. + + Nested cond subgraphs are not flattened here; callers that recurse (e.g. + post_lowering) will visit them when processing each returned module. + """ + cond_op = getattr(getattr(torch.ops, "higher_order", None), "cond", None) + if cond_op is None: + return [] + subgraphs: List[torch.fx.GraphModule] = [] + for node in gm.graph.nodes: + if node.op != "call_function" or node.target is not cond_op: + continue + if len(node.args) < 3: + continue + for arg in node.args[1:3]: + if not isinstance(arg, torch.fx.Node) or arg.op != "get_attr": + continue + attr: Any = gm + try: + for atom in str(arg.target).split("."): + attr = getattr(attr, atom) + except AttributeError: + continue + if isinstance(attr, torch.fx.GraphModule): + subgraphs.append(attr) + return subgraphs + + def get_tensor_placeholders( gm: torch.fx.GraphModule, ) -> List[torch.fx.Node]: @@ -32,16 +61,16 @@ def get_tensor_placeholders( return placeholders -def find_complex_nodes(gm: torch.fx.GraphModule): - complex_nodes = [] - complexNodes = {} +def find_complex_nodes(gm: torch.fx.GraphModule) -> List[torch.fx.Node]: + complex_nodes: List[torch.fx.Node] = [] + complexNodes: Dict[str, bool] = {} for node in gm.graph.nodes: if is_node_complex(node, complexNodes): complex_nodes.append(node) return complex_nodes -def is_node_complex(node: torch.fx.Node, complexNodes): +def is_node_complex(node: Any, complexNodes: Dict[str, bool]) -> bool: if not isinstance(node, torch.fx.Node): return False if node.name in complexNodes: diff --git a/tests/py/dynamo/conversion/test_cond_aten.py b/tests/py/dynamo/conversion/test_cond_aten.py new file mode 100644 index 00000000000..b056793e0ef --- /dev/null +++ b/tests/py/dynamo/conversion/test_cond_aten.py @@ -0,0 +1,170 @@ +import torch +import torch.nn as nn +from parameterized import parameterized +from torch.testing._internal.common_utils import run_tests + +from .harness import DispatchTestCase + + +class TestCondConverter(DispatchTestCase): + @parameterized.expand( + [ + ("pred_true", True), + ("pred_false", False), + ] + ) + def test_cond_add_sub(self, _, pred): + class CondAddSub(nn.Module): + def forward(self, x, predicate): + return torch.cond( + predicate, + lambda value: value + 1, + lambda value: value - 1, + (x,), + ) + + self.run_test( + CondAddSub(), + [torch.randn(1, 4), torch.tensor(pred)], + use_dynamo_tracer=True, + enable_passes=True, + ) + + @parameterized.expand( + [ + ("pred_true", True), + ("pred_false", False), + ] + ) + def test_cond_one_element_pred(self, _, pred): + class CondAddSub(nn.Module): + def forward(self, x, predicate): + return torch.cond( + predicate, + lambda value: value + 1, + lambda value: value - 1, + (x,), + ) + + self.run_test( + CondAddSub(), + [torch.randn(2, 3), torch.tensor([pred])], + use_dynamo_tracer=True, + enable_passes=True, + ) + + @parameterized.expand( + [ + ("pred_true", True), + ("pred_false", False), + ] + ) + def test_cond_multi_output(self, _, pred): + class CondMulti(nn.Module): + def forward(self, x, y, predicate): + def true_fn(a, b): + return a + 1, b * 2 + + def false_fn(a, b): + return a - 1, b / 2 + + return torch.cond(predicate, true_fn, false_fn, (x, y)) + + self.run_test( + CondMulti(), + [torch.randn(2, 2), torch.randn(2, 2), torch.tensor(pred)], + use_dynamo_tracer=True, + enable_passes=True, + ) + + @parameterized.expand( + [ + ("pred_true", True), + ("pred_false", False), + ] + ) + def test_cond_identity_branch(self, _, pred): + class CondIdentity(nn.Module): + def forward(self, x, predicate): + return torch.cond( + predicate, + # torch.cond forbids returning an operand alias; clone is the + # documented workaround and still exercises a pass-through branch. + lambda value: value.clone(), + lambda value: value + 1, + (x,), + ) + + self.run_test( + CondIdentity(), + [torch.randn(3, 3), torch.tensor(pred)], + use_dynamo_tracer=True, + enable_passes=True, + ) + + @parameterized.expand( + [ + ("pred_true", True), + ("pred_false", False), + ] + ) + def test_cond_linear_outside(self, _, pred): + class ConditionalModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + + def forward(self, x, predicate): + x = torch.relu(self.linear(x)) + return torch.cond( + predicate, + lambda value: value + 1, + lambda value: value - 1, + (x,), + ) + + self.run_test( + ConditionalModel(), + [torch.ones(1, 4), torch.tensor(pred)], + use_dynamo_tracer=True, + enable_passes=True, + ) + + @parameterized.expand( + [ + ("true_true", True, True), + ("true_false", True, False), + ("false_true", False, True), + ("false_false", False, False), + ] + ) + def test_cond_nested(self, _, pred_outer, pred_inner): + class NestedCond(nn.Module): + def forward(self, x, p1, p2): + def true_fn(v, inner_pred): + return torch.cond( + inner_pred, + lambda a: a + 1, + lambda a: a + 2, + (v,), + ) + + def false_fn(v, inner_pred): + return v - 1 + + return torch.cond(p1, true_fn, false_fn, (x, p2)) + + self.run_test( + NestedCond(), + [ + torch.randn(2, 2), + torch.tensor(pred_outer), + torch.tensor(pred_inner), + ], + use_dynamo_tracer=True, + enable_passes=True, + ) + + +if __name__ == "__main__": + run_tests() diff --git a/tests/py/dynamo/models/test_cond.py b/tests/py/dynamo/models/test_cond.py new file mode 100644 index 00000000000..ee5ce8c5dfb --- /dev/null +++ b/tests/py/dynamo/models/test_cond.py @@ -0,0 +1,124 @@ +import unittest + +import pytest +import torch +import torch.nn as nn +import torch_tensorrt +from torch.testing._internal.common_utils import TestCase +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.conversion import DYNAMO_CONVERTERS as CONVERTERS +from torch_tensorrt.dynamo.lowering import ( + get_decompositions, + post_lowering, + pre_export_lowering, +) + + +def _cond_node(gm: torch.fx.GraphModule) -> torch.fx.Node: + return next( + n + for n in gm.graph.nodes + if n.op == "call_function" and n.target is torch.ops.higher_order.cond + ) + + +class _AddSubCond(nn.Module): + def forward(self, x, predicate): + return torch.cond( + predicate, + lambda value: value + 1, + lambda value: value - 1, + (x,), + ) + + +class _LinearCond(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + + def forward(self, x, predicate): + x = torch.relu(self.linear(x)) + return torch.cond( + predicate, + lambda value: value + 1, + lambda value: value - 1, + (x,), + ) + + +class _UnsupportedBranchCond(nn.Module): + def forward(self, x, predicate): + return torch.cond( + predicate, + lambda value: torch.lgamma(value), + lambda value: value - 1, + (x,), + ) + + +def _export_lowered(mod: nn.Module, inputs: tuple) -> torch.fx.GraphModule: + settings = CompilationSettings() + exported = torch.export.export(mod.eval(), inputs) + exported = pre_export_lowering(exported, settings) + exported = exported.run_decompositions(get_decompositions()) + return post_lowering(exported.module(), settings) + + +@pytest.mark.unit +class TestCondCompilation(TestCase): + def test_cond_is_supported_when_branches_are_convertible(self): + x = torch.randn(1, 4) + pred = torch.tensor(True) + gm = _export_lowered(_AddSubCond(), (x, pred)) + self.assertTrue(_cond_node(gm) in CONVERTERS) + + def test_cond_falls_back_when_branch_has_unsupported_op(self): + x = torch.randn(1, 4).abs() + 0.1 + pred = torch.tensor(True) + gm = _export_lowered(_UnsupportedBranchCond(), (x, pred)) + self.assertFalse(_cond_node(gm) in CONVERTERS) + + @unittest.skipIf(not torch.cuda.is_available(), "CUDA required") + def test_cond_require_full_compilation(self): + model = _AddSubCond().eval().cuda() + x = torch.randn(1, 4, device="cuda") + + trt_mod = torch_tensorrt.compile( + model, + ir="dynamo", + inputs=[x, torch.tensor(True, device="cuda")], + min_block_size=1, + require_full_compilation=True, + pass_through_build_failures=True, + cache_built_engines=False, + reuse_cached_engines=False, + ) + + for pred in (True, False): + predicate = torch.tensor(pred, device="cuda") + eager = model(x, predicate) + compiled = trt_mod(x, predicate) + torch.testing.assert_close(compiled, eager, rtol=1e-4, atol=1e-4) + + @unittest.skipIf(not torch.cuda.is_available(), "CUDA required") + def test_cond_linear_outside_require_full_compilation(self): + model = _LinearCond().eval().cuda() + x = torch.ones(1, 4, device="cuda") + + trt_mod = torch_tensorrt.compile( + model, + ir="dynamo", + inputs=[x, torch.tensor(True, device="cuda")], + min_block_size=1, + require_full_compilation=True, + pass_through_build_failures=True, + cache_built_engines=False, + reuse_cached_engines=False, + ) + + for pred in (True, False): + predicate = torch.tensor(pred, device="cuda") + eager = model(x, predicate) + compiled = trt_mod(x, predicate) + torch.testing.assert_close(compiled, eager, rtol=1e-4, atol=1e-4)