o
    i                     @  s   d Z ddlmZ ddlmZ ddlmZ ddlmZ ddl	m
Z
 ddlmZ ddlmZ dd	lmZ dd
lmZmZ eddd\ZZdgZG dd dZdS )zX
Decorators and context managers for NVIDIA Tools Extension to profile MONAI components
    )annotations)defaultdict)wraps)Any)Function)Module)	Optimizer)Dataset)ensure_tupleoptional_importztorch._C._nvtxz:NVTX is not installed. Are you sure you have a CUDA build?)
descriptorRangec                   @  sb   e Zd ZU dZeeZded< 				dd ddZd!ddZ	dd Z
d"ddZdd Zdd ZdS )#r   a  
    A decorator and context manager for NVIDIA Tools Extension (NVTX) Range for profiling.
    When used as a decorator it encloses a specific method of the object with an NVTX Range.
    When used as a context manager, it encloses the runtime context (created by with statement) with an NVTX Range.

    Args:
        name: the name to be associated to the range
        methods: (only when used as decorator) the name of a method (or a list of the name of the methods)
            to be wrapped by NVTX range.
            If None (default), the method(s) will be inferred based on the object's type for various MONAI components,
            such as Networks, Losses, Functions, Transforms, and Datasets.
            Otherwise, it look up predefined methods: "forward", "__call__", "__next__", "__getitem__"
        append_method_name: if append the name of the methods to be decorated to the range's name
            If None (default), it appends the method's name only if we are annotating more than one method.
        recursive: if set to True, it will recursively annotate every individual module in a list
            or in a chain of modules (chained using Compose). Default to False.

    dictname_counterNFname
str | Nonemethodsstr | tuple[str, ...] | Noneappend_method_namebool | None	recursiveboolreturnNonec                 C  s   || _ || _|| _|| _d S N)r   r   r   r   )selfr   r   r   r    r   R/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/utils/nvtx.py__init__7   s   
zRange.__init__objr   c                 C  s8  | j du r.t|ttfrt|dd |D S ddlm} t||r+tdd|j|_d| _ | j	d u rgt|j
}d|v s@d	|v rH| d
|j	 }| j|  d7  < | j| dkrd| d
| j|  | _	n|| _	| jd u rs| || _nt| j| _| jd u rt| jdkrd| _nd| _| jD ]
}| ||| j q|S )NTc                 s  s    | ]
}t d d|V  qdS )Tr   N)r   ).0tr   r   r   	<genexpr>F   s    z!Range.__call__.<locals>.<genexpr>r   )Composer    FZCuCIMTorchVision_   )r   
isinstancelisttupletypeZmonai.transforms.composer$   r   
transformsr   __name__r   r   _get_methodr
   r   len_decorate_method)r   r   r$   r   methodr   r   r   __call__C   s2   






zRange.__call__c                   s   |r
| j  d| n| j |drt|}n|}t|| t  fdd}|drAG dd d|}t||| ||_d S t||| d S )N.__c                    s$   t   | i |}t   |S r   )_nvtx
rangePushArangePop)argskwargsoutputZ
_temp_funcr   r   r   range_wrapper}   s   
z-Range._decorate_method.<locals>.range_wrapperc                   @  s   e Zd ZdS )z7Range._decorate_method.<locals>.NVTXRangeDecoratedClassN)r-   
__module____qualname__r   r   r   r   NVTXRangeDecoratedClass   s    r?   )r   
startswithr+   getattrr   setattr	__class__)r   r   r1   r   ownerr<   r?   r   r;   r   r0   o   s   




zRange._decorate_methodr*   c                 C  s   t |trdg}t
|S t |trdg}t
|S t |tr%ddg}t
|S t |tr1dg}t
|S g d}g }|D ]}t||rE|| q9t|dk rXtdt	| d| t
|S )	Nforwardstepbackward__getitem__)rE   r2   __next__rH   r'   z*The method to be wrapped for this object [zm] is not recognized.The name of the method should be provided or the object should have one of these methods:)r(   r   r   r   r	   hasattrappendr/   
ValueErrorr+   r
   )r   r   method_listZdefault_methodsr1   r   r   r   r.      s2   





zRange._get_methodc                 C  s>   | j d u r| jd  d7  < d| jd  | _ t| j  d S )Ncontextr'   Zcontext_)r   r   r5   r6   )r   r   r   r   	__enter__   s   
zRange.__enter__c                 C  s   t   d S r   )r5   r7   )r   r+   value	tracebackr   r   r   __exit__   s   zRange.__exit__)NNNF)
r   r   r   r   r   r   r   r   r   r   )r   r   r   r   )r   r   r   r*   )r-   r=   r>   __doc__r   intr   __annotations__r   r2   r0   r.   rO   rR   r   r   r   r   r   !   s   
 
,
!N)rS   
__future__r   collectionsr   	functoolsr   typingr   Ztorch.autogradr   torch.nnr   torch.optimr   torch.utils.datar	   monai.utilsr
   r   r5   r&   __all__r   r   r   r   r   <module>   s   
