lib/python3.12/site-packages/cudnn/wrapper.py

"""Wrapper for cuDNN frontend to improve user experience.

This wrapper provides a more user-friendly interface for cuDNN frontend.
It allows users to create a graph, add operations to the graph, and then
compile the graph to a cuDNN plan. This wrapper is designed to avoid
boilerplate code.

Key Features:
    - Automatic graph validation and compilation
    - Simplified tensor management
    - Support for both named and positional tensor inputs
    - Automatic workspace management
    - PyTorch integration with DLPack support

Example:
    >>> x = torch.randn(8, 56, 56, 64, device=device, dtype=torch.float16).permute(0, 3, 1, 2)
    >>> w = torch.randn(32, 3, 3, 64, device=device, dtype=torch.float16).permute(0, 3, 1, 2)
    >>> with Graph() as graph:
    ...     y = graph.conv_fprop(
    ...         image=x, weight=w,
    ...         padding=[1,1], stride=[1,1], dilation=[1,1],
    ...         compute_data_type=data_type.FLOAT,
    ...         name="conv2d",
    ...     )
    ...     y.set_output(True).set_data_type(data_type.HALF)
    ...     # Graph is automatically validated and compiled on exit
    >>> graph.set_io_tuples(["conv2d::image", "conv2d::weight"], ["conv2d::Y"])
    >>> # Execute the graph
    >>> output = graph(x, w)
"""

from collections import OrderedDict
import atexit
import itertools
import inspect
import logging
from typing import Any, Dict, List, Optional, Tuple, Union

import cudnn
import cudnn.datatypes
from cudnn import data_type, heur_mode

try:
    import torch
except ImportError:
    torch = None

__all__ = ["Graph", "data_type", "heur_mode", "cudnn"]

# typedefs for readability
CudnnHandle = int
_default_cudnn_handle = None

# Configure logging
logger = logging.getLogger(__name__)


def _graph_tensor(graph: cudnn.pygraph, tensor: "torch.Tensor") -> cudnn.tensor:
    """Create a tensor in the graph object.

    Args:
        graph: The cuDNN graph object to create the tensor in
        tensor: The dlpack tensor to create a graph tensor from

    Returns:
        A cuDNN tensor object representing the input tensor in the graph

    Note:
        If the input tensor has requires_grad=True, it will be detached
        before creating the graph tensor to avoid gradient tracking issues.
    """
    if hasattr(tensor, "requires_grad") and tensor.requires_grad:
        # PyTorch tensor with requires_grad=True need to be detached first
        return graph.tensor_like(tensor.detach())
    else:
        return graph.tensor_like(tensor)


def _find_tensor(
    tensor: Union[str, cudnn.tensor, "torch.Tensor"],
    tensor_map: Dict[str, cudnn.tensor],
    dlpack_map: Dict[int, cudnn.tensor],
) -> str:
    """Find the mapping name for a tensor used in a graph.

    This function searches for a tensor in the tensor map and returns its
    corresponding name. The tensor can be specified in multiple ways:
    - As a string (either the assigned tensor name or the node::input_name)
    - As a cuDNN tensor object
    - As a DLPack-compatible tensor (e.g., PyTorch tensor) that was used in creating the graph

    Args:
        tensor: The tensor to find, can be a string name, cuDNN tensor, or DLPack tensor
        tensor_map: Dictionary mapping tensor names to cuDNN tensor objects
        dlpack_map: Dictionary mapping DLPack tensor IDs to cuDNN tensor objects

    Returns:
        The key in tensor_map that the provided tensor is mapped to

    Raises:
        ValueError: If the tensor cannot be found in the tensor map
    """
    if isinstance(tensor, str):
        # look up by canonical name, then assigned name
        if tensor in tensor_map:
            return tensor  # this name is "node::input_name"
        for tensor_name, tensor_value in tensor_map.items():
            if tensor_value.get_name() == tensor:
                return tensor_name  # name is the assigned name of the tensor
    elif isinstance(tensor, int):
        # look up by tensor uid
        for tensor_name, tensor_value in tensor_map.items():
            if tensor_value.get_uid() == tensor:
                return tensor_name
    elif isinstance(tensor, cudnn.tensor):
        for tensor_name, tensor_value in tensor_map.items():
            if tensor is tensor_value:
                return tensor_name
    elif hasattr(tensor, "__dlpack__") and isinstance(dlpack_map, dict) and id(tensor) in dlpack_map:
        tensor = dlpack_map[id(tensor)]
        for tensor_name, tensor_value in tensor_map.items():
            if tensor_value == tensor:
                return tensor_name
    raise ValueError("Input not found in tensor map")


def _extract_tensor(name: str, tensor: cudnn.tensor, arg_dict: dict) -> Optional["torch.Tensor"]:
    """Extract a dlpack tensor from the arg_dict that matches the provided name or cudnn tensor

    Args:
        name: The name of the tensor to extract
        tensor: The cudnn tensor to extract
        arg_dict: The dictionary of arguments to extract the tensor from

    Returns:
        A dlpack tensor
    """
    if name in arg_dict:
        return arg_dict[name]  # match by canonical name
    if tensor in arg_dict:
        return arg_dict[tensor]  # match by cudnn tensor object
    try:
        return arg_dict[tensor.get_name()]  # match by assigned name
    except KeyError:
        pass
    try:
        return arg_dict[tensor.get_uid()]  # match by tensor uid
    except KeyError:
        return None  # not found


def _tensor_like(cudnn_tensor: cudnn.tensor, tensor_type: str = "pyt") -> "torch.Tensor":
    """Create a tensor like the provided cudnn tensor

    Args:
        cudnn_tensor: The cuDNN tensor to create a dlpack tensor from
        tensor_type: The type of tensor to create, currently only "pyt" is supported

    Returns:
        A dlpack tensor allocated that is like the provided cuDNN tensor
    """
    if tensor_type != "pyt":
        raise NotImplementedError("Only PyTorch tensor is supported for now")
    if not cudnn.datatypes.is_torch_available():
        raise RuntimeError("PyTorch is not available")
    dtype = cudnn.datatypes._cudnn_to_torch_data_type(cudnn_tensor.get_data_type())
    if dtype is None:
        raise TypeError(f"cuDNN uses an unsupported data type in PyTorch: {cudnn_tensor.get_data_type()}")
    tensor = torch.empty_strided(cudnn_tensor.get_dim(), cudnn_tensor.get_stride(), device="cuda", dtype=dtype)
    return tensor


def get_default_handle(stream: Optional["torch.cuda.Stream"] = None) -> CudnnHandle:
    """Get the default cuDNN handle and set to torch's current stream"""
    global _default_cudnn_handle
    if torch is None:
        raise RuntimeError("PyTorch is not available")
    if _default_cudnn_handle is None:
        _default_cudnn_handle = cudnn.create_handle()
    if stream is None:
        stream = torch.cuda.current_stream().cuda_stream
    cudnn.set_stream(handle=_default_cudnn_handle, stream=stream)
    return _default_cudnn_handle


def destroy_default_handle():
    if _default_cudnn_handle is not None:
        cudnn.destroy_handle(_default_cudnn_handle)


atexit.register(destroy_default_handle)


class Graph:
    """Wrapper object for cuDNN computation graph

    This class simplifies the process of creating, compiling, and executing
    cuDNN computation graphs. It handles common boilerplate code and provides
    a more Pythonic interface to the cuDNN frontend API.

    Key features:
    - Automatic graph validation and compilation
    - Simplified tensor management with PyTorch integration
    - Support for both named and positional tensor inputs
    - Automatic workspace management

    Note:
        The graph is automatically validated and compiled when exiting the
        context manager. Any errors in graph construction will be raised
        at that point.
    """

    __handle: Optional[CudnnHandle] = None  # holding the cudnn handle pointer

    def __init__(
        self,
        *,
        handle: Optional[CudnnHandle] = None,
        inputs: Optional[List[Union[str, "torch.Tensor", cudnn.tensor]]] = None,
        outputs: Optional[List[Union[str, "torch.Tensor", cudnn.tensor]]] = None,
        heuristics: Optional[List[heur_mode]] = None,
        workspace_alloc: bool = True,
        **kwargs,
    ) -> None:
        if torch is None:
            raise RuntimeError("PyTorch is not available")
        if inputs and not isinstance(inputs, (list, tuple)):
            raise ValueError("inputs must be a list or tuple")
        if outputs and not isinstance(outputs, (list, tuple)):
            raise ValueError("outputs must be a list or tuple")
        if heuristics and not isinstance(heuristics, (list, tuple)):
            raise ValueError("heuristics must be a list or tuple")
        if cudnn.backend_version() < 91200:
            raise RuntimeError("cuDNN version 9.12.0 or higher is required")
        self.__kwargs = kwargs
        self.__graph = None  # to hold the cudnn.pygraph object
        self.__tensor_map = {}  # obj id of dlpack tensor -> cudnn tensor
        self.__tensor_in = OrderedDict()  # canonical node::argname -> cudnn tensors used as the input
        self.__tensor_out = OrderedDict()  # canonical node::outname -> cudnn tensors produced by the node
        self.__tensor_unknown = []  # list of cuDNN tensors created by user directly
        self.__node_count = {}  # function name of graph node -> number of times used
        self.__node_names = set()  # set of assigned names of graph nodes, to check name collision
        self.__input_tuples = None  # tuple of input tensors, if set by set_io_tuples
        self.__output_tuples = None  # tuple of output tensors, if set by set_io_tuples
        self.__inputs = inputs or []  # hold the list of inputs, to be used by set_io_tuples() implicitly
        self.__outputs = outputs or []  # hold the list of outputs, to be used by set_io_tuples() implicitly
        self.__heuristics = heuristics or [heur_mode.A, heur_mode.FALLBACK]
        if not workspace_alloc:
            self.__workspace = False
        if handle:
            self.__handle = handle
        # silently replace the PyTorch dtype into cuDNN dtype
        for key in ["io_data_type", "intermediate_data_type", "compute_data_type"]:
            if key in kwargs:
                kwargs[key] = cudnn.datatypes._torch_to_cudnn_data_type(kwargs[key]) or kwargs[key]

    def __del__(self):
        pass

    def __enter__(self):
        if self.__graph is not None:
            raise RuntimeError("Graph already created")
        self.__graph = cudnn.pygraph(
            # Pass handle only if self.__handle is an existing handle
            **({"handle": self.__handle} if self.__handle not in ["auto", None] else {}),
            **self.__kwargs,
        )
        return self

    def __exit__(self, exc_type, exc_value, tb):
        """Exit the context manager, validating and compiling the graph.

        This method performs the following steps:
        1. Validates the graph structure
        2. Builds the operation graph
        3. Creates execution plans
        4. Checks hardware support
        5. Builds the final plans
        6. Allocates workspace memory

        Raises:
            ValidationError: If graph validation fails
            GraphStateError: If graph operations are performed in invalid order
            CudnnError: For other cuDNN-related errors
        """
        # if there is an exception, clean up and propagate
        if exc_type is not None:
            logger.error("Exception during graph construction: %s", exc_value)
            self.__graph = None
            raise
        # prepare the graph and build plans: Each should return None or raise exception
        self.__graph.validate()
        self.__graph.build_operation_graph()
        self.__graph.create_execution_plans(self.__heuristics)
        # TODO: let user select_behavior_notes() and select_numeric_notes() here
        self.__graph.check_support()
        self.__graph.build_plans()
        # Set up workspace if not forbidden by user, then set up I/O tensor orders
        if not hasattr(self, "__workspace"):
            self.__workspace = torch.empty(
                self.__graph.get_workspace_size(),
                device="cuda",
                dtype=torch.uint8,
            )
        if self.__inputs or self.__outputs:
            self.set_io_tuples(self.__inputs, self.__outputs)
        del self.__inputs, self.__outputs

        logger.debug("Inputs: %s", self.__tensor_in)
        logger.debug("Outputs: %s", self.__tensor_out)
        logger.debug("Node count: %s", self.__node_count)
        return self.__graph

    def __getattr__(self, name: str) -> Any:
        """Intercept method calls to the graph object, usually during graph
        construction, and handle them appropriately.
        """
        attr = getattr(self.__graph, name)
        # calling tensor_like is unnecessary, just pass through
        pass_through = [
            "get_workspace_size",
            "get_workspace_size_plan_at_index",
            "serialize",
            "deserialize",
            "query_tensor_attributes_of_uid",
        ]
        if name in pass_through:
            return attr
        # some methods are blocked and should not be called via wrapper
        # TODO should allow user select execution plan
        blocked_methods = [
            "build",
            "build_operation_graph",
            "build_plan_at_index",
            "build_plans",
            "check_support",
            "create_execution_plan",
            "create_execution_plans",
            "deselect_behavior_notes",
            "deselect_engines",
            "deselect_numeric_notes",
            "deselect_workspace_greater_than",
            "execute",
            "execute_plan_at_index",
            "get_behavior_notes",
            "get_behavior_notes_for_plan_at_index",
            "get_engine_count",
            "get_execution_plan_count",
            "get_knobs_for_engine",
            "get_plan_name_at_index",
            "key",
            "populate_cuda_graph",
            "query_tensor_attributes_of_uid",
            "select_behavior_notes",
            "select_numeric_notes",
            "update_cuda_graph",
            "validate",
        ]

        if name in blocked_methods:
            raise RuntimeError(f"Calling {name} via wrapper is not allowed")
        # non-methods: pass through. Probably not used but be safe
        if not inspect.ismethod(attr):
            return attr

        # tensor creation methods: capture the output
        def tensor_capture(*args, **kwargs):
            output = attr(*args, **kwargs)
            self.__tensor_unknown.append(output)
            return output

        if name in ["tensor", "tensor_like"]:
            return tensor_capture

        # other methods: wrap the method to intercept the arguments and return values
        def wrapper(*args, **kwargs):
            args = list(args)  # shallow copy, to allow in-place modification
            # determine the name of the graph node, the node may carry name attribute
            if name not in self.__node_count:
                self.__node_count[name] = 0
            self.__node_count[name] += 1
            if "name" in kwargs:
                node_name = kwargs["name"]
            else:
                node_name = f"{name}.{self.__node_count[name]-1}"
                kwargs["name"] = node_name
            if node_name in self.__node_names:
                raise ValueError(f"Node name {node_name} already used")
            self.__node_names.add(node_name)
            # process positional arguments for dlpack tensors
            for i, obj in enumerate(args):
                if hasattr(obj, "__dlpack__"):
                    obj_id = id(obj)
                    if obj_id not in self.__tensor_map:
                        self.__tensor_map[obj_id] = _graph_tensor(self.__graph, obj)
                    obj = args[i] = self.__tensor_map[obj_id]
                if isinstance(obj, cudnn.tensor):
                    self.__tensor_in[f"{node_name}::{i}"] = obj
            # process keyword arguments for dlpack tensors
            for key, obj in kwargs.items():
                if hasattr(obj, "__dlpack__"):
                    obj_id = id(obj)
                    if obj_id not in self.__tensor_map:
                        self.__tensor_map[obj_id] = _graph_tensor(self.__graph, obj)
                    obj = kwargs[key] = self.__tensor_map[obj_id]
                if isinstance(obj, cudnn.tensor):
                    self.__tensor_in[f"{node_name}::{key}"] = obj
            # capturing node output
            output = attr(*args, **kwargs)
            if isinstance(output, cudnn.tensor):
                output_list = [output]
            elif isinstance(output, (list, tuple)):
                output_list = output
            for i, obj in enumerate(output_list):
                if isinstance(obj, cudnn.tensor):
                    if hasattr(obj, "get_name") and obj.get_name():
                        tensor_name = obj.get_name()
                    else:
                        tensor_name = f"{node_name}::{i}"
                    self.__tensor_out[tensor_name] = obj
            return output

        return wrapper

    def __call__(self, *args, **kwargs):
        """Execute the graph with tensor dict"""
        if self.__graph is None:
            raise RuntimeError("Graph not created")
        if not self.__graph.get_execution_plan_count():
            raise RuntimeError("You should not invoke the graph before the context exits")
        if len(args) == 1 and isinstance(args[0], dict):
            return self.__call_with_tensor_dict(args[0], **kwargs)
        else:
            if len(args) > 0 and not self.__input_tuples:
                raise ValueError("You should not invoke the graph with positional arguments before running set_io_tuples()")
            if len(args) != len(self.__input_tuples):
                raise ValueError(f"Number of arguments ({len(args)}) does not match number of inputs ({len(self.__input_tuples)})")
            return self.__call_with_positional_args(*args, **kwargs)

    def __call_with_positional_args(self, *args, **kwargs) -> Union["torch.Tensor", Tuple["torch.Tensor", ...]]:
        """Execute the graph with positional arguments.

        Args:
            *args: Positional arguments to pass to the graph
            **kwargs: Additional keyword arguments to pass to the graph execution

        Returns:
            A single tensor or a tuple of tensors

        Note:
            This method is called by __call__() when the graph is executed with positional arguments.
            It is not intended to be called directly by the user. The `args` should be a list of dlpack tensors
            that matches the input order of `self.__input_tuples`.
        """
        # prepare the variant pack:
        # all non-virtual tensors in __tensor_in and __tensor_out should be filled
        variant_pack = {}
        for cudnn_tensor, user_tensor in zip(self.__input_tuples, args):
            variant_pack[cudnn_tensor.get_uid()] = user_tensor
        output_tuple = [_tensor_like(cudnn_tensor, "pyt") for cudnn_tensor in self.__output_tuples]
        for cudnn_tensor, user_tensor in zip(self.__output_tuples, output_tuple):
            variant_pack[cudnn_tensor.get_uid()] = user_tensor
        # execute the graph
        kwargs = dict(kwargs)  # shallow copy

        if "handle" not in kwargs:
            if self.__handle == "auto":
                kwargs["handle"] = get_default_handle()
            elif self.__handle is not None:
                kwargs["handle"] = self.__handle
            else:
                raise RuntimeError("Need to specify cudnn handle to execute graph")
        if "workspace" not in kwargs:
            if self.__workspace is False:
                raise RuntimeError("Need to specify workspace to execute graph")
            else:
                kwargs["workspace"] = self.__workspace
        self.__graph.execute(variant_pack, **kwargs)
        # return the output as a single tensor or a tuple
        if len(output_tuple) == 1:
            return output_tuple[0]
        else:
            return output_tuple

    def __call_with_tensor_dict(
        self,
        tensor_dict: Dict[str, "torch.Tensor"],
        **kwargs,
    ) -> Dict[str, "torch.Tensor"]:
        """Execute the graph with a dictionary of tensors.

        Args:
            tensor_dict: Dictionary of tensor names to tensors
            **kwargs: Additional keyword arguments to pass to the graph execution

        Returns:
            Dictionary of tensor names to tensors

        Raises:
            RuntimeError: If a non-virtual tensor in the graph is not found in
            `tensor_dict`, or the tensor in `tensor_dict` is not a dlpack tensor
        """
        # Notes
        """
        from arg tensor_dict -> variant_pack
        also from self.__tensor_in + self.__tensor_out -> variant_pack
        both: check if all non-virtual tensors are filled
        """
        # prepare the variant pack:
        # all non-virtual tensors in __tensor_in and __tensor_out should be filled
        variant_pack = {}
        missing_tensors = {}
        for name, tensor in itertools.chain(self.__tensor_in.items(), self.__tensor_out.items()):
            if tensor.get_uid() in variant_pack or tensor.get_is_virtual():
                continue  # already filled or not needed
            user_tensor = _extract_tensor(name, tensor, tensor_dict)
            if user_tensor is None:
                missing_tensors[tensor] = name  # overwriting existing entries
                continue
            if not hasattr(user_tensor, "__dlpack__"):
                raise RuntimeError(f"Tensor {name} is not provided as a dlpack tensor")
            variant_pack[tensor.get_uid()] = user_tensor
        # check if all non-virtual tensors are filled
        missing_inputs = []
        missing_outputs = []
        for tensor, name in missing_tensors.items():
            if tensor.get_uid() in variant_pack:
                continue  # already filled
            if name in self.__tensor_out:
                # output tensor not specified, should be created automatically
                variant_pack[tensor.get_uid()] = tensor_dict[name] = _tensor_like(tensor, "pyt")
                missing_outputs.append(name)
            else:
                # input tensor not specified, flag it as missing
                missing_inputs.append(name)
        if missing_inputs:
            raise RuntimeError(f"Non-virtual input tensors not found in variant pack: {missing_inputs}")
        if missing_outputs:
            logger.debug("Added output tensors: %s", missing_outputs)
        # execute the graph
        kwargs = dict(kwargs)  # shallow copy
        if "handle" not in kwargs:
            if self.__handle == "auto":
                kwargs["handle"] = get_default_handle()
            elif self.__handle is not None:
                kwargs["handle"] = self.__handle
            else:
                raise RuntimeError("Need to specify cudnn handle to execute graph")
        if "workspace" not in kwargs:
            if self.__workspace is False:
                raise RuntimeError("Need to specify workspace to execute graph")
            else:
                kwargs["workspace"] = self.__workspace
        self.__graph.execute(variant_pack, **kwargs)
        return tensor_dict  # by this time, the output tensors are updated

    def set_io_tuples(
        self,
        inputs: List[Union[str, "torch.Tensor", cudnn.tensor]],
        outputs: List[Union[str, "torch.Tensor", cudnn.tensor]],
    ) -> None:
        """Set order of input and output tensors to allow graph to be executed with positional arguments.

        Args:
            inputs: List of input tensors or names
            outputs: List of output tensors or names

        Raises:
            ValueError: If inputs or outputs are not lists or tuples
        """
        if not isinstance(inputs, (list, tuple)):
            raise ValueError("inputs must be a list or tuple")
        if not isinstance(outputs, (list, tuple)):
            raise ValueError("outputs must be a list or tuple")
        if not self.__graph.get_execution_plan_count():
            # raise RuntimeError("You should not invoke set_io_tuples() before the context exits")
            self.__inputs = inputs
            self.__outputs = outputs
            return

        # self.__tensor_out and self.__tensor_in are dict of str -> cudnn tensor
        # all non-virtual tensors should be either input or output

        # Convert "inputs" to a list of names that can be looked up in __tensor_in
        input_tensors = []
        tensors_found = set()
        for i, name in enumerate(inputs):
            try:
                if name in self.__tensor_unknown:
                    input_tensors.append(name)  # user-created cuDNN tensor object
                    continue
                name = _find_tensor(name, self.__tensor_in, self.__tensor_map)
                tensor = self.__tensor_in[name]
                if id(tensor) in tensors_found:
                    raise ValueError(f"Input at index {i} ({name}) is a duplicate")
                tensors_found.add(id(tensor))
                input_tensors.append(tensor)
            except ValueError:
                raise ValueError(f"Input at index {i} ({name}) not found in tensor map") from None
        # Convert "outputs" to a list of names that can be looked up in __tensor_out
        output_tensors = []
        for i, name in enumerate(outputs):
            try:
                if name in self.__tensor_unknown:
                    output_tensors.append(name)  # user-created cuDNN tensor object
                    continue
                name = _find_tensor(name, self.__tensor_out, self.__tensor_map)
                tensor = self.__tensor_out[name]
                if id(tensor) in tensors_found:
                    raise ValueError(f"Output at index {i} ({name}) is a duplicate")
                tensors_found.add(id(tensor))
                output_tensors.append(tensor)
            except ValueError:
                raise ValueError(f"Output at index {i} ({name}) not found in tensor map") from None
        # Verify that all input tensors are non-virtual
        for i, tensor in enumerate(input_tensors):
            if tensor.get_is_virtual():
                raise ValueError(f"Input at index {i} is a virtual tensor")
        # Verify that all non-virtual tensors are covered by input or output
        for name, tensor in self.__tensor_out.items():
            if not tensor.get_is_virtual() and tensor not in output_tensors:
                raise ValueError(f"Node output {name} is a non-virtual tensor but not specified as output")
        for name, tensor in self.__tensor_in.items():
            if not tensor.get_is_virtual() and id(tensor) not in tensors_found:
                raise ValueError(f"Node input {name} is a non-virtual tensor but not specified as input or output")
        # Set the input and output names
        self.__input_tuples = tuple(input_tensors)
        self.__output_tuples = tuple(output_tensors)