from __future__ import annotations import logging from typing import 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 logger = logging.getLogger(__name__) def _extract_downstream_get_nodes( module_node: torch.fx.Node, output_indices: Set[int] ) -> Sequence[torch.fx.Node]: """Extracts downstream users of a node which get the item at a particular index Certain module-type nodes have multiple outputs (tuple of outputs). This function returns downstream nodes which call the _operator.getitem function, which extracts the element at a particular index in the tuple Args: module_node: FX module-type node to analyze output_index: Indices in the module node output to search for Returns: List of nodes which get the item at the specified index in the module node output """ get_nodes = [] # Iterate over all downstream users of the node object for user in module_node.users: # If the user is a "get" node accessing the specified index, store it if _get_qualified_name(user.target) == "_operator.getitem" and ( user.args[1] in output_indices ): get_nodes.append(user) return get_nodes 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}" 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 # 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}, ) # 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() # 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): cast_node_64bit = gm.graph.call_function( torch.ops.aten._to_copy.default, args=(module_node,), kwargs={"dtype": dtype_64bit}, ) # 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 ) # 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() def repair_double_inputs( parent_graph: torch.fx.GraphModule, submodule: torch.fx.GraphModule, submodule_inputs: Sequence[Input], device: torch.device, submodule_name: Optional[str] = None, ) -> Sequence[Input]: """Fixes all Long/Double type inputs 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: 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 """ submodule_torch_inputs = get_torch_inputs(submodule_inputs, device) num_submodule_inputs = len(submodule_inputs) repaired_outputs_once = False # For each input to the TRT subgraph, check if its type is long/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 submodule inputs in accordance with inserted casts dtype_32bit = torch.float32 submodule_torch_inputs = ( list(submodule_torch_inputs[:position]) + [ param.to(dtype_32bit), ] + list(submodule_torch_inputs[position + 1 :]) ) # Set the 32bit inputs and their types to the submodule Inputs for idx in range(len(submodule_inputs)): submodule_inputs[idx].torch_tensor = submodule_torch_inputs[idx] submodule_inputs[idx].dtype = dtype._from( submodule_torch_inputs[idx].dtype ) return submodule_inputs