Convert torch.cond to TensorRT IIfConditional

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)

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:

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,

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:

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],
+    )