utils.py

from __future__ import annotations

import logging
from dataclasses import fields, replace
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union

import numpy as np
import tensorrt as trt
import torch
from torch._subclasses.fake_tensor import FakeTensor
from torch_tensorrt._Device import Device
from torch_tensorrt._enums import dtype
from torch_tensorrt._Input import Input
from torch_tensorrt.dynamo import _defaults
from torch_tensorrt.dynamo._defaults import default_device
from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
from torch_tensorrt.dynamo._settings import CompilationSettings

from packaging import version

from .types import TRTDataType

logger = logging.getLogger(__name__)

COSINE_THRESHOLD = 0.99
DYNAMIC_DIM = -1
RTOL = 5e-3
ATOL = 5e-3


class Frameworks(Enum):
    NUMPY = "numpy"
    TORCH = "torch"
    TRT = "trt"


DataTypeEquivalence: Dict[
    TRTDataType, Dict[Frameworks, Union[TRTDataType, np.dtype, torch.dtype]]
] = {
    trt.int8: {
        Frameworks.NUMPY: np.int8,
        Frameworks.TORCH: torch.int8,
        Frameworks.TRT: trt.int8,
    },
    trt.int32: {
        Frameworks.NUMPY: np.int32,
        Frameworks.TORCH: torch.int32,
        Frameworks.TRT: trt.int32,
    },
    trt.int64: {
        Frameworks.NUMPY: np.int64,
        Frameworks.TORCH: torch.int64,
        Frameworks.TRT: trt.int64,
    },
    trt.float16: {
        Frameworks.NUMPY: np.float16,
        Frameworks.TORCH: torch.float16,
        Frameworks.TRT: trt.float16,
    },
    trt.float32: {
        Frameworks.NUMPY: np.float32,
        Frameworks.TORCH: torch.float32,
        Frameworks.TRT: trt.float32,
    },
    trt.bool: {
        Frameworks.NUMPY: bool,
        Frameworks.TORCH: torch.bool,
        Frameworks.TRT: trt.bool,
    },
}

if trt.__version__ >= "7.0":
    DataTypeEquivalence[trt.bool] = {
        Frameworks.NUMPY: np.bool_,
        Frameworks.TORCH: torch.bool,
        Frameworks.TRT: trt.bool,
    }


def use_python_runtime_parser(use_python_runtime: Optional[bool] = None) -> bool:
    """Parses a user-provided input argument regarding Python runtime

    Automatically handles cases where the user has not specified a runtime (None)

    Returns True if the Python runtime should be used, False if the C++ runtime should be used
    """
    using_python_runtime = use_python_runtime
    reason = ""

    # Runtime was manually specified by the user
    if using_python_runtime is not None:
        reason = "as requested by user"
    # Runtime was not manually specified by the user, automatically detect runtime
    else:
        try:
            from torch_tensorrt.dynamo.runtime import TorchTensorRTModule  # noqa: F401

            using_python_runtime = False
            reason = "since C++ dependency was detected as present"
        except ImportError:
            using_python_runtime = True
            reason = "since import failed, C++ dependency not installed"

    logger.info(
        f"Using {'Python-only' if using_python_runtime else 'Default'} Torch-TRT Runtime ({reason})"
    )

    return using_python_runtime


def cosine_similarity(gt_tensor: torch.Tensor, pred_tensor: torch.Tensor) -> float:
    gt_tensor = gt_tensor.flatten().to(torch.float32)
    pred_tensor = pred_tensor.flatten().to(torch.float32)
    if torch.sum(gt_tensor) == 0.0 or torch.sum(pred_tensor) == 0.0:
        if torch.allclose(gt_tensor, pred_tensor, atol=1e-4, rtol=1e-4, equal_nan=True):
            return 1.0
    res_t = torch.nn.functional.cosine_similarity(
        gt_tensor, pred_tensor, dim=0, eps=1e-6
    )
    res: float = res_t.cpu().detach().item()

    return res


def input_is_dynamic(inputs: Sequence[Union[Input, torch.Tensor]]) -> bool:
    """
    Return true if the provided inputs are `torch_tensorrt.Input` objects and have dynamic shapes.
    """
    return not any(isinstance(input, torch.Tensor) for input in inputs) and any(
        input.shape_mode == Input._ShapeMode.DYNAMIC for input in inputs
    )


def get_torch_tensor(
    input: Input,
    device: torch.device,
    mode: str = "",
) -> Union[int, torch.Tensor]:
    if input.is_shape_tensor:
        # TODO: All the shape tensors we've encountered so far are plain integers.
        # Validate this assumption on more models.
        return input.shape["opt_shape"][0]

    if len(mode) > 0:
        return input.example_tensor(mode).to(device)
    else:
        return input.torch_tensor.to(device)


def get_torch_inputs(
    inputs: Sequence[Input] | Dict[str, Any],
    device: Union[Device, torch.device, str],
    mode: str = "",
) -> Sequence[Union[int, torch.Tensor]] | Dict[str, Union[int, torch.Tensor]]:
    """
    Return the torch_tensor from the Input object. If mode is set, this implies
    user is using dynamic shaped inputs and return the corresponding input based
    on the mode requested.
    """
    device = to_torch_device(device)

    if isinstance(inputs, dict):
        result_dict: Dict[str, Union[int, torch.Tensor]] = {}
        for k, v in inputs.items():
            if isinstance(v, (list, tuple, dict)):
                result_dict[k] = get_torch_inputs(v, device)
            elif isinstance(v, Input):
                result_dict[k] = get_torch_tensor(v, device, mode)
        return result_dict
    else:
        result_list: List[Union[int, torch.Tensor]] = []
        for input in inputs:
            if isinstance(input, Input):
                result_list.append(get_torch_tensor(input, device, mode))
            elif isinstance(input, torch.Tensor):
                result_list.append(input.to(device))
            else:
                raise AssertionError(f"Input type {type(input)} is not a valid type")
        return result_list


def get_model_device(module: torch.fx.GraphModule) -> torch.device:
    """
    Returns the device on which the module parameters exist.
    """
    device = None
    for parameter in list(module.parameters()):
        if isinstance(parameter, (torch.nn.parameter.Parameter, torch.Tensor)):
            return parameter.device

    for buffer in list(module.buffers()):
        if isinstance(buffer, (torch.Tensor)):
            return buffer.device

    if device is None:
        device = to_torch_device(default_device())

    return device


def set_log_level(parent_logger: Any, level: Any) -> None:
    """
    Sets the log level to the user provided level.
    This is used to set debug logging at a global level
    at entry points of tracing, dynamo and torch_compile compilation.
    """
    if parent_logger:
        parent_logger.setLevel(level)


def prepare_inputs(
    inputs: Input | torch.Tensor | Sequence[Any] | Dict[Any, Any],
    disable_memory_format_check: bool = False,
) -> Any:
    if isinstance(inputs, Input):
        return inputs

    elif isinstance(inputs, torch.Tensor):
        return Input.from_tensor(
            inputs, disable_memory_format_check=disable_memory_format_check
        )

    elif isinstance(inputs, (list, tuple)):
        torchtrt_input_list = []
        for input_obj in inputs:
            torchtrt_input = prepare_inputs(
                input_obj, disable_memory_format_check=disable_memory_format_check
            )
            torchtrt_input_list.append(torchtrt_input)

        return (
            torchtrt_input_list
            if isinstance(inputs, list)
            else tuple(torchtrt_input_list)
        )

    elif isinstance(inputs, dict):
        torchtrt_inputs_dict: Dict[Any, Any] = dict()

        for key, input_obj in inputs.items():
            torchtrt_input = prepare_inputs(
                input_obj, disable_memory_format_check=disable_memory_format_check
            )
            torchtrt_inputs_dict[key] = torchtrt_input

        return torchtrt_inputs_dict

    else:
        raise ValueError(
            f"Invalid input type {type(inputs)} encountered in the dynamo_compile input parsing. "
            + "Allowed input types: {torch_tensorrt.Input, torch.Tensor, list, tuple, dict}"
        )


def parse_complex_tensor_structs(
    inputs: Input | torch.Tensor | Sequence[Any] | Dict[Any, Any],
    attribute_to_extract: str,
    apply_fn: Callable[[Any], Any] = lambda x: x,
) -> Any:
    """Parses complex structures of Tensors and returns a mirrored structure
    Extracts key attributes of each singular element, while reconstructing the struct
    Optionally applies a function to each attribute before returning
    """
    if isinstance(inputs, (torch.Tensor, Input)):
        return apply_fn(getattr(inputs, attribute_to_extract, None))
    elif isinstance(inputs, (int, float, bool)):
        # inputs is a python scalar value
        inputs_torch = torch.tensor(inputs)
        return apply_fn(getattr(inputs_torch, attribute_to_extract, None))

    elif isinstance(inputs, (list, tuple)):
        torchtrt_input_list = []
        for input_obj in inputs:
            torchtrt_input = parse_complex_tensor_structs(
                input_obj, attribute_to_extract, apply_fn
            )
            torchtrt_input_list.append(torchtrt_input)

        return (
            torchtrt_input_list
            if isinstance(inputs, list)
            else tuple(torchtrt_input_list)
        )

    elif isinstance(inputs, dict):
        torchtrt_inputs_dict: Dict[Any, Any] = dict()

        for key, input_obj in inputs.items():
            torchtrt_input = parse_complex_tensor_structs(
                input_obj, attribute_to_extract, apply_fn
            )
            torchtrt_inputs_dict[key] = torchtrt_input

        return torchtrt_inputs_dict

    else:
        raise ValueError(
            f"Invalid input type {type(inputs)} encountered during Dynamo input parsing. "
            + "Allowed input types: {torch_tensorrt.Input, torch.Tensor, list, tuple, dict}"
        )


def contains_sym_int(tensor: torch.Tensor) -> bool:
    """
    Returns true if the given tensor has symbolic shape.
    """
    return any(isinstance(dim, torch.SymInt) for dim in tensor)


def extract_var_range_info(symbolic_integer: torch.SymInt) -> Dict[str, int]:
    """
    This function returns the min, max, opt values of a symbolic integer.
    """
    node = symbolic_integer.node
    expr = node.expr
    shape_env = node.shape_env
    # An expr can be a independent SymInt node (eg: s0 or s1) or a composition of them eg: (48*s0 or s0*s1).
    # In the case of expr which has symbolic computation, bound_sympy evaluates them.
    # https://pytorch.org/docs/stable/generated/torch.fx.experimental.symbolic_shapes.ShapeEnv.html#torch.fx.experimental.symbolic_shapes.ShapeEnv.bound_sympy
    # expr.xreplace replaces the symbolic variables with their current values and computes the expression.
    var_range = shape_env.var_to_range.get(expr, None) or shape_env.bound_sympy(expr)
    var_val = shape_env.var_to_val.get(expr, None) or expr.xreplace(
        shape_env.var_to_val
    )
    assert var_range, var_val
    min_val, max_val, opt_val = int(var_range.lower), int(var_range.upper), int(var_val)
    # Torchdynamo 0/1 specialization outlier
    min_val = 1 if min_val == 2 else min_val
    min_max_opt = {}
    min_max_opt["min"] = min_val
    min_max_opt["max"] = max_val
    min_max_opt["opt"] = opt_val

    return min_max_opt


def unwrap_tensor_shape(
    tensor: Union[torch.Tensor, FakeTensor, torch.SymInt]
) -> Sequence[Union[int, Tuple[int, int]]]:
    """
    This is a helper function used to print/return the shape of the tensor.
    For regular torch.tensor's, it returns the static shape.
    For symbolic tensors, eg:(1, s0, 4), this function returns [1, [min, max], 4]. The min
    and max correspond to the lower and upper values of s0 symbolic dimension.
    """
    tensor_shape: List[Union[int, Tuple[int, int]]] = []
    # for dimension in tensor.shape:
    if isinstance(tensor, int):
        tensor_shape.append(tensor)
    elif isinstance(tensor, torch.SymInt):
        min_max_opt = extract_var_range_info(tensor)
        tensor_shape.append((min_max_opt["min"], min_max_opt["max"]))
    elif isinstance(tensor, (torch.Tensor, FakeTensor)):
        for dimension in tensor.shape:
            tensor_shape.extend(unwrap_tensor_shape(dimension))

    return tuple(tensor_shape)


def unwrap_tensor_dtype(tensor: Union[torch.Tensor, FakeTensor, torch.SymInt]) -> Any:
    """
    Returns the dtype of torch.tensor or FakeTensor. For symbolic integers, we return int64
    """
    if isinstance(tensor, (torch.Tensor, FakeTensor)):
        return tensor.dtype
    elif isinstance(tensor, torch.SymInt):
        return torch.int64
    else:
        raise ValueError(f"Found invalid tensor type {type(tensor)}")


def get_graph_io_attrs(
    io_nodes: Sequence[torch.fx.Node], attr_type: str
) -> Sequence[Any]:
    """
    Returns a list of attributes (shapes or dtypes) of the I/O nodes
    """
    assert attr_type in ["shape", "dtype"]
    attr_fn = unwrap_tensor_shape if attr_type == "shape" else unwrap_tensor_dtype
    graph_io_attrs = []
    for node in io_nodes:
        if "val" in node.meta:
            metadata = node.meta["val"]
            if isinstance(metadata, (tuple, list)):
                for tensor in metadata:
                    graph_io_attrs.append(attr_fn(tensor))
            else:
                graph_io_attrs.append(attr_fn(metadata))

    return graph_io_attrs


def parse_graph_io(module: torch.fx.GraphModule, dryrun_tracker: Any) -> None:
    """
    Parse the graph I/O shape/dtype info for the whole graph and store in the dryrun tracker
    """
    # Parse inputs of the graph
    input_nodes = [node for node in module.graph.nodes if node.op == "placeholder"]
    input_shapes = get_graph_io_attrs(input_nodes, "shape")
    input_dtypes = get_graph_io_attrs(input_nodes, "dtype")
    dryrun_tracker.input_shapes = input_shapes
    dryrun_tracker.input_dtypes = input_dtypes

    # Parse outputs of the graph
    mark_output_nodes = [node for node in module.graph.nodes if node.op == "output"]
    output_nodes = []
    for node in mark_output_nodes:
        output_nodes.extend(node.all_input_nodes)
    output_shapes = get_graph_io_attrs(output_nodes, "shape")
    output_dtypes = get_graph_io_attrs(output_nodes, "dtype")
    dryrun_tracker.output_shapes = output_shapes
    dryrun_tracker.output_dtypes = output_dtypes


def to_torch_device(device: Optional[Union[Device, torch.device, str]]) -> torch.device:
    """Cast a device-type to torch.device

    Returns the corresponding torch.device
    """
    if isinstance(device, Device):
        return device.to(torch.device)

    elif isinstance(device, torch.device):
        return device

    elif device is None:
        return torch.device(torch.cuda.current_device())

    else:
        return torch.device(device)


def to_torch_tensorrt_device(
    device: Optional[Union[Device, torch.device, str]]
) -> Device:
    """Cast a device-type to torch_tensorrt.Device

    Returns the corresponding torch_tensorrt.Device
    """
    return Device._from(device)


def parse_dynamo_kwargs(
    kwargs: Any,
) -> Tuple[CompilationSettings, Optional[BaseEngineCache]]:
    """Parses the kwargs field of a Dynamo backend

    Args:
        kwargs: Keyword arguments dictionary provided to the backend
    Returns:
        CompilationSettings object with relevant kwargs
    """
    # Initialize an empty CompilationSettings object
    settings = CompilationSettings()

    # If the user specifies keyword args, overwrite those fields in settings
    # Validate all specified kwargs to ensure they are true fields of the dataclass
    #
    # Note: kwargs provided by torch.compile are wrapped in the "options" key
    if kwargs:
        if "options" in kwargs and len(kwargs) == 1:
            kwargs = kwargs["options"]

        valid_attrs = {attr.name for attr in fields(settings)}
        valid_kwargs = {k: v for k, v in kwargs.items() if k in valid_attrs}
        settings = replace(settings, **valid_kwargs)

    # TODO: Remove once Dynamo precisions refactoring is complete
    if "enabled_precisions" in kwargs:
        enabled_precisions = {dtype._from(e) for e in kwargs["enabled_precisions"]}

        if len(enabled_precisions) == 0:
            logger.info(
                f"No precision specified, defaulting to {_defaults.ENABLED_PRECISION}"
            )
            enabled_precisions = _defaults.ENABLED_PRECISIONS

        settings.enabled_precisions = enabled_precisions

    # Parse input runtime specification
    settings.use_python_runtime = use_python_runtime_parser(settings.use_python_runtime)

    # Ensure device is a torch_tensorrt Device
    settings.device = to_torch_tensorrt_device(settings.device)

    # Check and update device settings
    if "device" not in kwargs:
        logger.info(
            f"Device not specified, using Torch default current device - cuda:{settings.device.gpu_id}. "
            "If this is incorrect, please specify an input device, via the device keyword."
        )

    # Ignore and warn about require_full_compilation flag
    if settings.require_full_compilation:
        logger.warning(
            "Detected require_full_compilation=True for a torch.compile run. "
            "This option has no effect in torch.compile."
        )
        settings.require_full_compilation = False

    # If cache_built_engines and reuse_cached_engines are True but custom_engine_cache is not provided,
    # then create a default disk engine cache

    engine_cache = None
    if kwargs.get("cache_built_engines") or kwargs.get("reuse_cached_engines"):
        if kwargs.get("custom_engine_cache") is not None:
            engine_cache = kwargs.get("custom_engine_cache")
        else:
            from torch_tensorrt.dynamo._engine_cache import DiskEngineCache

            engine_cache_dir = kwargs.get(
                "engine_cache_dir", _defaults.ENGINE_CACHE_DIR
            )
            engine_cache_size = kwargs.get(
                "engine_cache_size", _defaults.ENGINE_CACHE_SIZE
            )
            engine_cache = DiskEngineCache(engine_cache_dir, engine_cache_size)

    if kwargs.get("torch_executed_ops"):
        settings.torch_executed_ops = kwargs.get("torch_executed_ops")

    logger.info("Compilation Settings: %s\n", settings)

    return settings, engine_cache


def req_torch_version(min_torch_version: str = "2.dev") -> Callable[..., Any]:
    """
    Create a decorator which verifies the Torch version installed
    against a specified version range

    Args:
        min_torch_version (str): The minimum required Torch version
        for the decorated function to work properly

    Returns:
        A decorator which raises a descriptive error message if
        an unsupported Torch version is used
    """

    def nested_decorator(f: Callable[..., Any]) -> Callable[..., Any]:
        def function_wrapper(*args: Any, **kwargs: Any) -> Any:
            # Parse minimum and current Torch versions
            min_version = version.parse(min_torch_version)
            current_version = version.parse(torch.__version__)

            if current_version < min_version:
                raise AssertionError(
                    f"Expected Torch version {min_torch_version} or greater, "
                    + f"when calling {f}. Detected version {torch.__version__}"
                )
            else:
                return f(*args, **kwargs)

        return function_wrapper

    return nested_decorator


def check_module_output(
    new_module: torch.fx.GraphModule,
    refitted_module: torch.fx.GraphModule,
    arg_inputs: Any,
    kwarg_inputs: Any = None,
) -> bool:
    old_outputs, new_outputs = refitted_module(*arg_inputs), new_module(
        *arg_inputs, **kwarg_inputs
    )
    if type(old_outputs) != type(new_outputs):
        logger.warning("The output types are different. Output check is skipped.")
        return True
    return check_output_equal(old_outputs, new_outputs)


def check_output_equal(
    output1: Any,
    output2: Any,
    rtol: float = RTOL,
    atol: float = ATOL,
) -> bool:

    if type(output1) != type(output2):
        logger.warning(
            "The output types are different. Check_output_equal will always return false."
        )
        return False

    if isinstance(output1, torch.Tensor):
        if output1.shape != output2.shape:
            return False
        return torch.allclose(output1, output2, rtol, atol)  # type: ignore

    elif isinstance(output1, (tuple, list)):
        if len(output1) != len(output2):
            return False
        for a, b in zip(output1, output2):
            if not check_output_equal(a, b):
                return False
            return True

    elif isinstance(output1, dict):
        if output1.keys() != output2.keys():
            return False
        for a, b in zip(output1.values(), output2.values()):
            if not check_output_equal(a, b):
                return False
        return True

    logger.warning(
        "The output type is not supported to be checked. Check_output_equal will always return false."
    )
    return False


def get_flat_args_with_check(
    exported_program: torch.export.ExportedProgram,
    args: list[Any],
    kwargs: dict[str, Any],
) -> tuple[Any, Any]:
    """Flatten args, kwargs using pytree, then, check specs.

    Args:
        args: List[Any] original args passed to __call__
        kwargs: Dict[str, Any] original kwargs passed to __call

    Returns:
        A tuple of (flat_args, received_spec)
        flat_args is flattend args / kwargs
        received_spec is the pytree spec produced while flattening the
        tuple (args, kwargs)
    """
    import torch.utils._pytree as pytree
    from torch.export._tree_utils import reorder_kwargs

    in_spec = exported_program.call_spec.in_spec
    if in_spec is not None:
        kwargs = reorder_kwargs(kwargs, in_spec)
    flat_args_with_path, received_spec = pytree.tree_flatten_with_path((args, kwargs))
    flat_args = tuple(x[1] for x in flat_args_with_path)
    return flat_args, received_spec