o ԰i,j @sxdZddlmZddlZddlZddlZddlZddlmZm Z m Z m Z m Z m Z ddlZddlZddlmZmZzddlZWn eyIdZYnwgdZeZdaeeZdejdd d ejfd d Zde eejd fd e eejfde eejfd efddZdedejde d e d fddZ!d"dejded d fddZ"d#de dd efddZ#ddZ$e%e$Gd d!d!Z&dS)$aWrapper 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) ) OrderedDictN)AnyDictListOptionalTupleUnion) data_type heur_mode)Graphr r cudnngraphtensor torch.TensorreturncCs(t|dr|jr||S||S)aCreate 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. requires_grad)hasattrr tensor_likedetach)r rrA/home/ubuntu/.local/lib/python3.10/site-packages/cudnn/wrapper.py _graph_tensor:s r tensor_map dlpack_mapcCst|tr"||vr |S|D]\}}||kr|Sqt dt|tr>|D]\}}||kr9|Sq+t dt|tjrY|D] \}}||urT|SqHt dt|drt|t rt ||vr|t |}|D] \}}||kr|Sqst d)a7Find 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 __dlpack__zInput not found in tensor map) isinstancestritemsget_nameintget_uidr rrdictid ValueError)rrr tensor_name tensor_valuerrr _find_tensorOs6      r&namearg_dictcCsf||vr||S||vr||Sz||WSty Ynwz||WSty2YdSw)a3Extract 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 N)rKeyErrorr )r'rr(rrr_extract_tensors   r*pyt cudnn_tensor tensor_typecCsz|dkrtdtjstdtj|}|dur&td|tj | d|d}t || | }|S)a0Create 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 r+z(Only PyTorch tensor is supported for nowPyTorch is not availableNz0cuDNN uses an unsupported data type in PyTorch: cudadevicedtype) NotImplementedErrorr datatypesis_torch_available RuntimeError_cudnn_to_torch_data_type get_data_type TypeErrortorchemptyget_dim as_strided get_stride)r,r-r2rrrr _tensor_likes  r?streamztorch.cuda.StreamcCsFtdurtdtdurta|durtjj}tjt|dtS)z>Get the default cuDNN handle and set to torch's current streamNr.)handler@) r:r6_default_cudnn_handler create_handler/current_stream cuda_stream set_stream)r@rrrget_default_handles rGcCstdur ttdSdSN)rBr destroy_handlerrrrdestroy_default_handlesrJc @s*eZdZUdZdZeeed<dddddddeedeee e de j fd eee e de j fd eee d ed df d dZddZddZddZde d efddZddZd e dedffddZdee dfd ee dffddZdee e de j fd ee e de j fd dfd d!ZdS)"r aWrapper 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. N_Graph__handleT)rAinputsoutputs heuristicsworkspace_allocrArLrrMrNrOrcKs tdurtd|rt|ttfstd|r"t|ttfs"td|r/t|ttfs/tdtdkr9td||_d|_ i|_ t |_ t |_ g|_i|_t|_d|_d|_|p]g|_|pbg|_|pktjtjg|_|srd|_|rw||_dD]}||vrtj||p||||<qydS) Nr.inputs must be a list or tupleoutputs must be a list or tuplez"heuristics must be a list or tuplei@dz*cuDNN version 9.12.0 or higher is requiredF) io_data_typeintermediate_data_typecompute_data_type)r:r6rlisttupler#r backend_version_Graph__kwargs _Graph__graph_Graph__tensor_mapr_Graph__tensor_in_Graph__tensor_out_Graph__tensor_unknown_Graph__node_countset_Graph__node_names_Graph__input_tuples_Graph__output_tuples_Graph__inputs_Graph__outputsr AFALLBACK_Graph__heuristics_Graph__workspacerKr4_torch_to_cudnn_data_type)selfrArLrMrNrOkwargskeyrrr__init__s@    zGraph.__init__cCsdSrHrrjrrr__del__sz Graph.__del__cCsB|jdur tdtjdi|jdvrd|jini|j|_|S)NzGraph already created)autoNrAr)rYr6r pygraphrKrXrnrrr __enter__s zGraph.__enter__cCs|durtd|d|_|j|j|j|j|j|jt |ds;t j |j dt j d|_|jsA|jrI||j|j|`|`td|jtd|jtd|j|jS) a%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 Nz'Exception during graph construction: %s __workspacer/r0z Inputs: %sz Outputs: %szNode count: %s)loggererrorrYvalidatebuild_operation_graphcreate_execution_plansrg check_support build_plansrr:r;get_workspace_sizeuint8rhrcrd set_io_tuplesdebugr[r\r^)rjexc_type exc_valuetbrrr__exit__s,       zGraph.__exit__r'cs|tjgd}|vrSgd}|vr tddts'Sfdd}dvr4|Sfdd }|S) N)r{ get_workspace_size_plan_at_index serialize deserializequery_tensor_attributes_of_uid)buildrwbuild_plan_at_indexrzrycreate_execution_planrxdeselect_behavior_notesdeselect_enginesdeselect_numeric_notesdeselect_workspace_greater_thanexecuteexecute_plan_at_indexget_behavior_notes$get_behavior_notes_for_plan_at_indexget_engine_countget_execution_plan_countget_knobs_for_engineget_plan_name_at_indexrlpopulate_cuda_graphrselect_behavior_notesselect_numeric_notesupdate_cuda_graphrvzCalling z via wrapper is not allowedcs|i|}j||SrH)r]append)argsrkoutput)attrrjrrtensor_capturems z)Graph.__getattr__..tensor_capture)rrc st|}jvrdj<jd7<d|vr |d}ndjd}||d<|jvr=td|dj|t|D]4\}}t|drkt|}|jvrbt j |j|<j|}||<t |t j r{|j|d|<qG|D]4\}}t|drt|}|jvrt j |j|<j|}||<t |t j r|j|d|<q|i|}t |t j r|g}n t |ttfr|}t|D]$\}}t |t j rt|d r|r|} n|d|} |j| <q|S) Nrr'.z Node name z already usedrz::r)rUr^r`r#add enumeraterr"rZrrYrr rr[rrVrr\) rrk node_nameiobjobj_idrlr output_listr$rr'rjrrwrappervsT               z"Graph.__getattr__..wrapper)getattrrYr6inspectismethod)rjr' pass_throughblocked_methodsrrrrr __getattr__;s  /zGraph.__getattr__cOs|jdur td|jstdt|dkr*t|dtr*|j|dfi|St|dkr7|js7tdt|t|jkrPtdt|dt|jd |j |i|S) z"Execute the graph with tensor dictNzGraph not createdz8You should not invoke the graph before the context exitsrrzXYou should not invoke the graph with positional arguments before running set_io_tuples()zNumber of arguments (z#) does not match number of inputs ()) rYr6rlenrr!_Graph__call_with_tensor_dictrar#!_Graph__call_with_positional_args)rjrrkrrr__call__s   zGraph.__call__)r.cOsi}t|j|D] \}}|||<qdd|jD}t|j|D] \}}|||<q!t|}d|vrN|jdkr?t|d<n|jdurJ|j|d<ntdd|vr`|jdur[td |j|d<|j j |fi|t |d krt|d S|S) a?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`. cSsg|]}t|dqSr+)r?).0r,rrr sz5Graph.__call_with_positional_args..rArpN-Need to specify cudnn handle to execute graph workspaceF*Need to specify workspace to execute graphrr) ziprar rbr!rKrGr6rhrYrr)rjrrk variant_packr, user_tensor output_tuplerrr__call_with_positional_argss*       z!Graph.__call_with_positional_args tensor_dictc Ks i}i}t|j|jD]1\}}||vs|r qt|||}|dur/|||<qt|dssF        0