fix(dynamo): make truncate_double metadata match engine precision

py/torch_tensorrt/dynamo/conversion/_conversion.py

@@ -221,7 +221,9 @@ def interpret_module_to_result(
     """
 
     symbolic_shape_expressions = extract_symbolic_shape_expressions(
-        module, inputs=inputs
+        module,
+        inputs=inputs,
+        truncate_double=settings.truncate_double,
     )
     if symbolic_shape_expressions is None:
         raise RuntimeError(

py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py

@@ -19,6 +19,7 @@
 def extract_symbolic_shape_expressions(
     module: torch.fx.GraphModule,
     inputs: Optional[Sequence[Input]] = None,
+    truncate_double: bool = False,
 ) -> Optional[Dict[str, List[Dict[str, Any]]]]:
     """
     Extract symbolic shape expressions from an FX graph.
@@ -32,6 +33,8 @@ def extract_symbolic_shape_expressions(
             name. Used as the dtype source of truth for scalar inputs, which
             have no dtype of their own in FX metadata. Falls back to a
             best-effort default when not provided.
+        truncate_double: Record float64 tensor bindings as float32, matching
+            the precision TensorRT builds when double truncation is enabled
 
     Returns:
         Dict with 'inputs' and 'outputs' keys, each containing a list of dicts with shape_exprs and dtype,
@@ -79,7 +82,11 @@ def extract_symbolic_shape_expressions(
             input_info.append(
                 {
                     "shape_exprs": shape_exprs,
-                    "dtype": input_val.dtype,
+                    "dtype": (
+                        torch.float32
+                        if truncate_double and input_val.dtype == torch.float64
+                        else input_val.dtype
+                    ),
                     "name": input_node.name,
                 }
             )
@@ -134,7 +141,11 @@ def extract_symbolic_shape_expressions(
             output_info.append(
                 {
                     "shape_exprs": shape_exprs,
-                    "dtype": out_val.dtype,
+                    "dtype": (
+                        torch.float32
+                        if truncate_double and out_val.dtype == torch.float64
+                        else out_val.dtype
+                    ),
                 }
             )
         elif isinstance(out_val, (torch.SymInt, torch.SymFloat, int, float, bool)):

py/torch_tensorrt/dynamo/conversion/truncate_double.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional, Sequence, Set
+from typing import Any, Dict, Optional, Sequence, Set
 
 import torch
 from torch.fx.node import _get_qualified_name
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._Input import Input
-from torch_tensorrt.dynamo.utils import get_torch_inputs
+from torch_tensorrt.dynamo.utils import get_output_metadata, get_torch_inputs
 
 logger = logging.getLogger(__name__)
 
@@ -40,124 +40,150 @@ def _extract_downstream_get_nodes(
     return get_nodes
 
 
+def _metadata_dtype(metadata: Dict[str, Any]) -> Optional[torch.dtype]:
+    """Return the dtype of tensor metadata, ignoring scalar outputs."""
+    value = metadata.get("val")
+    if isinstance(value, torch.Tensor):
+        return value.dtype
+
+    tensor_meta = metadata.get("tensor_meta")
+    return getattr(tensor_meta, "dtype", None)
+
+
+def _metadata_to_dtype(
+    metadata: Dict[str, Any], target_dtype: torch.dtype
+) -> Dict[str, Any]:
+    """Copy tensor metadata while changing its dtype."""
+    updated = metadata.copy()
+    value = updated.get("val")
+    if isinstance(value, torch.Tensor):
+        updated["val"] = value.to(target_dtype)
+
+    tensor_meta = updated.get("tensor_meta")
+    if tensor_meta is not None and hasattr(tensor_meta, "_replace"):
+        updated["tensor_meta"] = tensor_meta._replace(dtype=target_dtype)
+
+    return updated
+
+
+def _find_module_node(gm: torch.fx.GraphModule, submodule_name: str) -> torch.fx.Node:
+    for node in gm.graph.nodes:
+        if node.op == "call_module" and str(node.target) == submodule_name:
+            return node
+
+    raise AssertionError(
+        f"Sought module node {submodule_name}, could not find in graph:\n{gm.graph}"
+    )
+
+
 def _repair_64bit_input(
     gm: torch.fx.GraphModule,
     position: int,
     submodule_name: str,
-    submodule_outputs: Optional[torch.Tensor | Sequence[torch.Tensor]],
-    dtype: torch.dtype,
 ) -> None:
-    """Fixes a single Long/Double input to a TRT-accelerated subgraph
-
-    In-Place modifies the provided graph
-
-    Inserts a cast to the 32-bit equivalent type for TRT, then if necessary,
-    inserts an upcast back to the 64-bit type for subsequent Torch operations
-
-    Args:
-        gm: FX GraphModule enclosing the TRT subgraph
-        position: Index in the submodule inputs at which the long or double input is found
-        submodule_name: Name of TRT-accelerated subgraph module in FX graph
-        submodule_outputs: Output tensor(s) of TRT-accelerated subgraph (used for dtypes/structure)
-        dtype: Data type of tensor at position in submodule (double/long)
-    """
-    assert dtype in (
-        torch.float64,
-    ), f"dtype argument must be torch.float64, got {dtype}"
-
+    """Downcast a single double input at a TRT boundary to float32."""
     logger.info(
         f"Downcasting a 64-bit input at position {position} of submodule {submodule_name}"
     )
 
-    # Determine target data type in 32 and 64 bit forms
-    dtype_64bit = dtype
     dtype_32bit = torch.float32
+    module_node = _find_module_node(gm, submodule_name)
 
-    # Find the node representing the submodule in the graph
-    module_node = None
-
-    # Iterate over all nodes in the graph, seeking target module name match
-    for n in gm.graph.nodes:
-        if n.op == "call_module" and str(n.target) == submodule_name:
-            module_node = n
-            break
-
-    if module_node is None:
-        raise AssertionError(
-            f"Sought module node {submodule_name}, could not find in graph:\n{gm.graph}"
-        )
-
-    # Extract the 64-bit node of the input
     node_64bit = module_node.all_input_nodes[position]
-
-    # Prior to the module, insert a cast to the 32-bit equivalent node
     with gm.graph.inserting_before(module_node):
         node_32bit = gm.graph.call_function(
             torch.ops.aten._to_copy.default,
             args=(node_64bit,),
             kwargs={"dtype": dtype_32bit},
         )
+        node_32bit.meta = _metadata_to_dtype(node_64bit.meta, dtype_32bit)
 
-    # Replace 64-bit input to TRT module with new 32-bit cast node
     module_node.replace_input_with(node_64bit, node_32bit)
 
-    output_positions_64bit = set()
+    gm.graph.eliminate_dead_code()
+    gm.graph.lint()
+    gm.recompile()
 
-    # Determine if any outputs of the model are 64-bit type and store their indices
-    if submodule_outputs is not None:
-        outputs_list = (
-            [submodule_outputs]
-            if isinstance(submodule_outputs, torch.Tensor)
-            else submodule_outputs
-        )
 
-        for output_position, output in enumerate(outputs_list):
-            if output.dtype == dtype_64bit:
-                output_positions_64bit.add(output_position)
-
-    # Only enter this code block if there exists a 64-bit output
-    # This implies a cast is needed, since TRT cannot output 64-bit tensors
-    if output_positions_64bit:
-        # Determine whether the outputs of the module are tuple-type or not
-        is_collection_output = False
-        if isinstance(submodule_outputs, tuple):
-            is_collection_output = True
-
-        if not is_collection_output:
-            # If the output is a single tensor, insert a cast back to int64
-            with gm.graph.inserting_after(module_node):
+def _repair_64bit_outputs(
+    gm: torch.fx.GraphModule,
+    submodule_name: str,
+    submodule_output_metadata: Sequence[Dict[str, Any]],
+    is_collection_output: bool,
+) -> None:
+    """Correct output metadata and insert restoring casts for double outputs.
+
+    A partition can have a float64 output even when every runtime input is
+    float32 -- e.g. a float32 input added to a float64 weight/constant -- so
+    this must run independently of whether any input needed repair. The
+    output dtypes come from the partition's FX metadata; compilation must not
+    execute the partition merely to discover information already recorded
+    there.
+    """
+    dtype_64bit = torch.float64
+    dtype_32bit = torch.float32
+
+    output_positions_64bit: Set[int] = set()
+    original_output_metadata = list(submodule_output_metadata)
+    truncated_output_metadata = []
+    for output_position, metadata in enumerate(original_output_metadata):
+        if _metadata_dtype(metadata) == dtype_64bit:
+            output_positions_64bit.add(output_position)
+            truncated_output_metadata.append(_metadata_to_dtype(metadata, dtype_32bit))
+        else:
+            truncated_output_metadata.append(metadata.copy())
+
+    if not output_positions_64bit:
+        return
+
+    module_node = _find_module_node(gm, submodule_name)
+
+    # The call_module node describes the actual engine boundary. Preserve its
+    # container convention while correcting tensor dtypes to what TRT emits.
+    for key in ("val", "tensor_meta"):
+        values = [
+            metadata[key] for metadata in truncated_output_metadata if key in metadata
+        ]
+        if not values:
+            continue
+        current = module_node.meta.get(key)
+        if isinstance(current, tuple):
+            module_node.meta[key] = tuple(values)
+        elif isinstance(current, list) or len(values) > 1:
+            module_node.meta[key] = values
+        else:
+            module_node.meta[key] = values[0]
+
+    if not is_collection_output:
+        with gm.graph.inserting_after(module_node):
+            cast_node_64bit = gm.graph.call_function(
+                torch.ops.aten._to_copy.default,
+                args=(module_node,),
+                kwargs={"dtype": dtype_64bit},
+            )
+            cast_node_64bit.meta = original_output_metadata[0].copy()
+
+        module_node.replace_all_uses_with(
+            cast_node_64bit, delete_user_cb=lambda user: user != cast_node_64bit
+        )
+    else:
+        get_nodes = _extract_downstream_get_nodes(module_node, output_positions_64bit)
+        for get_node in get_nodes:
+            output_position = get_node.args[1]
+            get_node.meta = truncated_output_metadata[output_position].copy()
+            with gm.graph.inserting_after(get_node):
                 cast_node_64bit = gm.graph.call_function(
                     torch.ops.aten._to_copy.default,
-                    args=(module_node,),
+                    args=(get_node,),
                     kwargs={"dtype": dtype_64bit},
                 )
+                cast_node_64bit.meta = original_output_metadata[output_position].copy()
 
-            # Replace all uses of the TRT module (except the cast node) with the 64-bit equivalent
-            module_node.replace_all_uses_with(
-                cast_node_64bit, delete_user_cb=lambda user: (user != cast_node_64bit)
-            )
-
-        else:
-            # If the output is a tuple of tensors, extract downstream users for each 64-bit output
-            get_nodes = _extract_downstream_get_nodes(
-                module_node, output_positions_64bit
+            get_node.replace_all_uses_with(
+                cast_node_64bit,
+                delete_user_cb=lambda user: user != cast_node_64bit,
             )
 
-            # For each downstream user, append a cast node back to the 64-bit precision
-            for get_node in get_nodes:
-                with gm.graph.inserting_after(get_node):
-                    cast_node_64bit = gm.graph.call_function(
-                        torch.ops.aten._to_copy.default,
-                        args=(get_node,),
-                        kwargs={"dtype": torch.float64},
-                    )
-
-                get_node.replace_all_uses_with(
-                    cast_node_64bit,
-                    delete_user_cb=lambda user: (user != cast_node_64bit),
-                )
-
-    # Clean up graph and ensure invariants are preserved
     gm.graph.eliminate_dead_code()
     gm.graph.lint()
     gm.recompile()
@@ -170,46 +196,34 @@ def repair_double_inputs(
     device: torch.device,
     submodule_name: Optional[str] = None,
 ) -> Sequence[Input]:
-    """Fixes all Long/Double type inputs to a TRT-accelerated subgraph
+    """Repair float64 inputs and outputs at a TensorRT partition boundary.
 
     In-Place modifies the provided graph
 
-    Inserts a cast to the 32-bit equivalent type for TRT, then if necessary,
-    inserts an upcast back to the 64-bit type for subsequent Torch operations
+    Casts float64 runtime inputs to float32 for TensorRT and independently
+    restores float64 source outputs after the engine.
 
     Args:
         parent_graph: FX GraphModule enclosing the TRT subgraph
         submodule: Child submodule to repair inputs on
         submodule_inputs: Input tensor(s) of TRT-accelerated subgraph (used for dtypes/structure)
         submodule_name: Optionally specify the name of the submodule target in the parent graph
     Returns:
-        New submodule inputs, updated accordingly with long/double truncation
+        New submodule inputs, updated for float64 truncation
     """
     submodule_torch_inputs = get_torch_inputs(submodule_inputs, device)
     num_submodule_inputs = len(submodule_inputs)
-    repaired_outputs_once = False
+    name = submodule_name if submodule_name is not None else submodule._get_name()
+    output_node = next(node for node in submodule.graph.nodes if node.op == "output")
+    is_collection_output = isinstance(output_node.args[0], (tuple, list))
+    submodule_output_metadata = get_output_metadata(submodule)
 
-    # For each input to the TRT subgraph, check if its type is long/double
+    # For each input to the TRT subgraph, check if its type is double.
     for position in range(num_submodule_inputs):
         param = submodule_torch_inputs[position]
 
-        # If the data type of the input is long/double, insert necessary
-        # casts to replace the operation
         if isinstance(param, torch.Tensor) and param.dtype == torch.float64:
-            # Ensure outputs are only repaired once per submodule to avoid
-            # unnecessary ops showing up in the graph
-            if not repaired_outputs_once:
-                submodule_outputs = submodule(*submodule_torch_inputs)
-
-            _repair_64bit_input(
-                parent_graph,
-                position,
-                submodule_name if submodule_name is not None else submodule._get_name(),
-                None if repaired_outputs_once else submodule_outputs,
-                param.dtype,
-            )
-
-            repaired_outputs_once = True
+            _repair_64bit_input(parent_graph, position, name)
 
             # Repair submodule inputs in accordance with inserted casts
             dtype_32bit = torch.float32
@@ -228,4 +242,10 @@ def repair_double_inputs(
                     submodule_torch_inputs[idx].dtype
                 )
 
+    # A partition can have a float64 output (e.g. from a float64 weight) even
+    # when no runtime input was float64, so this runs unconditionally.
+    _repair_64bit_outputs(
+        parent_graph, name, submodule_output_metadata, is_collection_output
+    )
+
     return submodule_inputs