o Ơi@sddlmZddlZddlmZmZddlmZmZm Z m Z ddl Z ddl ZddlZddlmZmZddlmZGdddeZGd d d eZGd d d eZGd ddeZGdddeZdS)) annotationsN)Enumauto)AnyOptionalTypeUnion)ENABLED_FEATURESneeds_torch_tensorrt_runtime)is_tensorrt_version_supportedc@s"eZdZdZeZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZeZeZe Ze ZeZeZeZeZe Ze Ze Ze Ze Ze Ze Z e Z!e Z"eZ#e$dddZ%e& dd d d Z'e& dd!ddZ( dd"ddZ)d#ddZ*d$ddZ+d%ddZ,e Z-eZ.dS)&dtypezfEnum to describe data types to Torch-TensorRT, has compatibility with torch, tensorrt and numpy dtypestrreturnboolcCs.t|tjrdSt|trt|tjrdSdS)NTF) isinstancenpr type issubclassgeneric)r rI/home/ubuntu/.local/lib/python3.10/site-packages/torch_tensorrt/_enums.py _is_np_objvs   zdtype._is_np_objF7Union[torch.dtype, trt.DataType, np.dtype, dtype, type] use_defaultcCspt|tjrs|tjkrtjS|tjkrtjS|tjkrtjS|tj kr&tj S|tj kr.tj S|tj kr6tjS|tjkr>tjS|tjkrFtjS|tjkrNtjS|tjkrVtjS|tjkr^tjS|rltd|dtjStd|t|tjr|tjjkrtjS|tjjkrtjS|tjj krtj S|tjj!krtj S|tjj"krtjS|tjj#krtjS|tjj$krtjS|tjj%krtjS|tjj&krtjSt'dr|tjj(krtj)Std|t*|r@|t+jkrtjS|t+jkrtjS|t+j krtj S|t+j,krtjS|t+j-krtjS|t+j.krtjS|t+jkr tjS|t+j/kr)tjS|r8td|dtjStdt0|t|trH|St1j2rddl3m4}t||jr||jjkrctjS||jj krmtj S||jjkrwtjS||jjkrtjS||jjkrtjS||jj5krtjS||jjkrtjS||jj6krtj6Std |td |d ) aCreate a Torch-TensorRT dtype from another library's dtype system. Takes a dtype enum from one of numpy, torch, and tensorrt and create a ``torch_tensorrt.dtype``. If the source dtype system is not supported or the type is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.dtype.try_from()`` Arguments: t (Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype)): Data type enum from another library use_default (bool): In some cases a catch all type (such as ``torch_tensorrt.dtype.f32``) is sufficient, so instead of throwing an exception, return default value. Returns: dtype: Equivalent ``torch_tensorrt.dtype`` to ``t`` Raises: TypeError: Unsupported data type or unknown source Examples: .. code:: py # Succeeds float_dtype = torch_tensorrt.dtype._from(torch.float) # Returns torch_tensorrt.dtype.f32 # Throws exception float_dtype = torch_tensorrt.dtype._from(torch.complex128) zQGiven dtype that does not have direct mapping to Torch-TensorRT supported types (z+), defaulting to torch_tensorrt.dtype.floatzyProvided an unsupported data type as a data type for translation (support: bool, int, long, half, float, bfloat16), got: 10.8.0zsProvided an unsupported data type as a data type for translation (support: bool, int, half, float, bfloat16), got: zpProvided an unsupported data type as an input data type (support: bool, int, long, half, float, bfloat16), got: r_ChProvided an unsupported data type as an input data type (support: bool, int32, long, half, float), got: z|tjkrt(jS|tjkrt(jS|tjkrt(j)S|tjkrt(j*S|tjkrt(j+S|tj krt(j S|tjkr"t(j,S|tjkr+t(j-S|tjkr4t(j.S|r:t(j,Std|tkrE|St/j0rdd l1m2}||jkr|tjkr_|jj S|tjkri|jjS|tjkrs|jj)S|tjkr}|jjS|tjkr|jjS|tjkr|jjS|tjkr|jjS|tj3kr|jj3Std |td |) aConvert dtype into the equivalent type in [torch, numpy, tensorrt] Converts ``self`` into one of numpy, torch, and tensorrt equivalent dtypes. If ``self`` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.dtype.try_to()`` Arguments: t (Union(Type(torch.dtype), Type(tensorrt.DataType), Type(numpy.dtype), Type(dtype))): Data type enum from another library to convert to use_default (bool): In some cases a catch all type (such as ``torch.float``) is sufficient, so instead of throwing an exception, return default value. Returns: Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype): dtype equivalent ``torch_tensorrt.dtype`` from library enum ``t`` Raises: TypeError: Unsupported data type or unknown target Examples: .. code:: py # Succeeds float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float # Failure float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Throws exception z8Given dtype that does not have direct mapping to torch (z), defaulting to torch.floatzUnsupported torch dtype (had: rrzUnsupported tensorrt dtypezUnsupported numpy dtyperrrz;Provided unsupported destination type for dtype conversion )4rr r!r r#r"r'intr%r$r)r(r+r*r-r,r/r.r1rLr2rr4r3r5r6r7r8r9r:r;r=r<r>r?r@rArBr rCrr&rErFrGr0rHr rJrKrrMselfr rrrrrtoGs $                                           zdtype.to;Optional[Union[torch.dtype, trt.DataType, np.dtype, dtype]]c CsTz |||}|WSttfy)}ztjd|dddWYd}~dSd}~ww)aConvert dtype into the equivalent type in [torch, numpy, tensorrt] Converts ``self`` into one of numpy, torch, and tensorrt equivalent dtypes. If ``self`` is not supported in the target library, then returns ``None``. Arguments: t (Union(Type(torch.dtype), Type(tensorrt.DataType), Type(numpy.dtype), Type(dtype))): Data type enum from another library to convert to use_default (bool): In some cases a catch all type (such as ``torch.float``) is sufficient, so instead of throwing an exception, return default value. Returns: Optional(Union(torch.dtype, tensorrt.DataType, numpy.dtype, dtype)): dtype equivalent ``torch_tensorrt.dtype`` from library enum ``t`` Examples: .. code:: py # Succeeds float_dtype = torch_tensorrt.dtype.f32.to(torch.dtype) # Returns torch.float # Failure float_dtype = torch_tensorrt.dtype.bf16.to(numpy.dtype) # Returns None z/torch_tensorrt.dtype conversion to target type  failedTrTNr_rVr7r5rWr^r rrXrYrrrtry_tos  z dtype.try_toothercCt|}t|j|jkSN)r rOrvaluer^reother_rrr__eq__ z dtype.__eq__r\cC t|jSrghashrhr^rrr__hash__  zdtype.__hash__N)r rrrF)r rrrrr )r rPrrrrQ)r r[rrrrP)r r[rrrr`)rerPrrrr\)/__name__ __module__ __qualname____doc__rrMr!r#r'r%r-r/r1r2r4r)r+r r"r&r$rEfloat8fp8float4rDr,fp16rFr.fp32rGrLfp64r0r3 staticmethodr classmethodrOrZr_rdrkrqrr\rrrrr s| +   ' r c@seZdZdZeZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZ eZe ZeZedddZeddd Zdd d ZdddZdddZdddZdS) memory_formatf;Union[torch.memory_format, trt.TensorFormat, memory_format]rcCst|tjr%|tjkrtjS|tjkrtjS|tjkrtjStdtt|t j r|t j j kr4tj S|t j jkr=tjS|t j jkrFtjS|t j jkrOtjS|t j jkrXtjS|t j jkratjS|t j jkrjtjS|t j jkrstjS|t j jkr|tjS|t j jkrtjS|t j j krtj!S|t j j"krtj#S|t j j$krtj%Stdtt|tr|St&j'rddl(m)}t||j r||j jkrtjS||j jkrtjSt*dtd)aCreate a Torch-TensorRT memory format enum from another library memory format enum. Takes a memory format enum from one of torch, and tensorrt and create a ``torch_tensorrt.memory_format``. If the source is not supported or the memory format is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.memory_format.try_from()`` Arguments: f (Union(torch.memory_format, tensorrt.TensorFormat, memory_format)): Memory format enum from another library Returns: memory_format: Equivalent ``torch_tensorrt.memory_format`` to ``f`` Raises: TypeError: Unsupported memory format or unknown source Examples: .. code:: py torchtrt_linear = torch_tensorrt.memory_format._from(torch.contiguous) z7Provided an unsupported memory format for tensor, got: z7Provided an unsupported tensor format for tensor, got: rrZProvided an unsupported tensor format (support: NCHW/contiguous_format, NHWC/channel_last)z=Provided unsupported source type for memory_format conversion)+rrrcontiguous_format contiguous channels_lastchannels_last_3dr7r r8r9 TensorFormatLINEARlinearCHW2chw2HWC8hwc8CHW4chw4CHW16chw16CHW32chw32DHWC8dhwc8CDHW32cdhw32HWChwc DLA_LINEAR dla_linearDLA_HWC4dla_hwc4HWC16hwc16DHWCdhwcr rJrKrrV)rNrrrrrrOsj                       zmemory_format._fromOptional[memory_format]c CRzt|}|WSttfy(}ztjd|dddWYd}~dSd}~ww)aCreate a Torch-TensorRT memory format enum from another library memory format enum. Takes a memory format enum from one of torch, and tensorrt and create a ``torch_tensorrt.memory_format``. If the source is not supported or the memory format is not supported in Torch-TensorRT, then ``None`` will be returned. Arguments: f (Union(torch.memory_format, tensorrt.TensorFormat, memory_format)): Memory format enum from another library Returns: Optional(memory_format): Equivalent ``torch_tensorrt.memory_format`` to ``f`` Examples: .. code:: py torchtrt_linear = torch_tensorrt.memory_format.try_from(torch.contiguous) rSz' to torch_tensorrt.memory_format failedTrTN)rrOrVr7r5rW)rNrrXrYrrrrZs  zmemory_format.try_fromr MUnion[Type[torch.memory_format], Type[trt.TensorFormat], Type[memory_format]]cCs|tjkr!|tjkr tjS|tjkrtjS|tjkrtjStd|tjkr|tj kr/tjj S|tj kr8tjj S|tj krAtjjS|tjkrJtjjS|tjkrStjjS|tjkr\tjjS|tjkretjjS|tjkrntjjS|tjkrwtjjS|tjkrtjjS|tjkrtjjS|tjkrtjj S|tj!krtjj"Std|tkr|St#j$rddl%m&}||jkr|tjkr|jjS|tjkr|jjSt'dtd)aConvert ``memory_format`` into the equivalent type in torch or tensorrt Converts ``self`` into one of torch or tensorrt equivalent memory format. If ``self`` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.memory_format.try_to()`` Arguments: t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))): Memory format type enum from another library to convert to Returns: Union(torch.memory_format, tensorrt.TensorFormat, memory_format): Memory format equivalent ``torch_tensorrt.memory_format`` in enum ``t`` Raises: TypeError: Unknown target type or unsupported memory format Examples: .. code:: py # Succeeds tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous zUnsupported torch dtypez"Unsupported tensorrt memory formatrrrzBProvided unsupported destination type for memory format conversion)(rrrrrrr7r8rrrrrrrrrrrrrrrrrrrrrrrrrrrr rJrKrrVr^r rrrrr_sf                      zmemory_format.toEOptional[Union[torch.memory_format, trt.TensorFormat, memory_format]]c CRz||}|WSttfy(}ztjd|dddWYd}~dSd}~ww)aConvert ``memory_format`` into the equivalent type in torch or tensorrt Converts ``self`` into one of torch or tensorrt equivalent memory format. If ``self`` is not supported in the target library, then ``None`` will be returned Arguments: t (Union(Type(torch.memory_format), Type(tensorrt.TensorFormat), Type(memory_format))): Memory format type enum from another library to convert to Returns: Optional(Union(torch.memory_format, tensorrt.TensorFormat, memory_format)): Memory format equivalent ``torch_tensorrt.memory_format`` in enum ``t`` Examples: .. code:: py # Succeeds tf = torch_tensorrt.memory_format.linear.to(torch.dtype) # Returns torch.contiguous z7torch_tensorrt.memory_format conversion to target type raTrTNrbr^r rXrYrrrrdn  zmemory_format.try_torercCst|}|j|jkSrg)rrOrhrirrrrks  zmemory_format.__eq__r\cCrmrgrnrprrrrqrrzmemory_format.__hash__N)rrrr)rrrr)r rrr)r rrr)rerrrrt)rurvrwrxrrrrrrrrrrrrrrrrrrrOrZr_rdrkrqrrrrrsN          Z  ! [ #rc@steZdZdZeZ eZ eZ edddZ eddd Z ddddZ ddddZ d ddZ d!ddZdS)" DeviceTypez#Type of device TensorRT will targetd!Union[trt.DeviceType, DeviceType]rcCst|tjr|tjjkrtjS|tjjkrtjStdt|tr#|StjrHddlm }t||jrH||jjkr;tjS||jjkrDtjStdt d)abCreate a Torch-TensorRT device type enum from a TensorRT device type enum. Takes a device type enum from tensorrt and create a ``torch_tensorrt.DeviceType``. If the source is not supported or the device type is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.DeviceType.try_from()`` Arguments: d (Union(tensorrt.DeviceType, DeviceType)): Device type enum from another library Returns: DeviceType: Equivalent ``torch_tensorrt.DeviceType`` to ``d`` Raises: TypeError: Unknown source type or unsupported device type Examples: .. code:: py torchtrt_dla = torch_tensorrt.DeviceType._from(tensorrt.DeviceType.DLA) 6Provided an unsupported device type (support: GPU/DLA)rrz:Provided unsupported source type for DeviceType conversion) rr8rGPUDLArVr rJrKrr7)rNrrrrrrOs*        zDeviceType._fromOptional[DeviceType]c Cr)a Create a Torch-TensorRT device type enum from a TensorRT device type enum. Takes a device type enum from tensorrt and create a ``torch_tensorrt.DeviceType``. If the source is not supported or the device type is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.DeviceType.try_from()`` Arguments: d (Union(tensorrt.DeviceType, DeviceType)): Device type enum from another library Returns: DeviceType: Equivalent ``torch_tensorrt.DeviceType`` to ``d`` Examples: .. code:: py torchtrt_dla = torch_tensorrt.DeviceType._from(tensorrt.DeviceType.DLA) rSz$ to torch_tensorrt.DeviceType failedTrTN)rrOrVr7r5rW)rNrrXrYrrrrZs  zDeviceType.try_fromFr -Union[Type[trt.DeviceType], Type[DeviceType]]rrcCs|tjkr!|tjkrtjjS|tjkrtjjS|rtjjStd|tkr'|StjrKddlm}||jkrK|tjkr>|jjS|tjkrG|jjStdt d)aConvert ``DeviceType`` into the equivalent type in tensorrt Converts ``self`` into one of torch or tensorrt equivalent device type. If ``self`` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.DeviceType.try_to()`` Arguments: t (Union(Type(tensorrt.DeviceType), Type(DeviceType))): Device type enum from another library to convert to Returns: Union(tensorrt.DeviceType, DeviceType): Device type equivalent ``torch_tensorrt.DeviceType`` in enum ``t`` Raises: TypeError: Unknown target type or unsupported device type Examples: .. code:: py # Succeeds trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA rrrz@Provided unsupported destination type for device type conversion) r8rrrrVr rJrKrr7r]rrrr_ s2       z DeviceType.to+Optional[Union[trt.DeviceType, DeviceType]]c CsVz |j||d}|WSttfy*}ztjd|dddWYd}~dSd}~ww)aConvert ``DeviceType`` into the equivalent type in tensorrt Converts ``self`` into one of torch or tensorrt equivalent memory format. If ``self`` is not supported in the target library, then ``None`` will be returned. Arguments: t (Union(Type(tensorrt.DeviceType), Type(DeviceType))): Device type enum from another library to convert to Returns: Optional(Union(tensorrt.DeviceType, DeviceType)): Device type equivalent ``torch_tensorrt.DeviceType`` in enum ``t`` Examples: .. code:: py # Succeeds trt_dla = torch_tensorrt.DeviceType.DLA.to(tensorrt.DeviceType) # Returns tensorrt.DeviceType.DLA rRz4torch_tensorrt.DeviceType conversion to target type raTrTNrbrcrrrrdHs zDeviceType.try_torecCrfrg)rrOrrhrirrrrkirlzDeviceType.__eq__r\cCrmrgrnrprrrrqmrrzDeviceType.__hash__N)rrrr)rrrrrs)r rrrrr)r rrrrr)rerrrrt)rurvrwrxrUNKNOWNrrrrOrZr_rdrkrqrrrrrs$ 5 # A !rc@sleZdZdZeZ eZ eZ edddZ eddd Z dd d Z dddZ dddZ dddZdS)EngineCapabilityzr EngineCapability determines the restrictions of a network during build time and what runtime it targets. c-Union[trt.EngineCapability, EngineCapability]rcCst|tjr%|tjjkrtjS|tjjkrtjS|tjjkr!tjStdt|tr,|StjrZddl m }t||jrZ||jjkrDtjS||jjkrMtjS||jjkrVtjStdt d)aCreate a Torch-TensorRT Engine capability enum from a TensorRT Engine capability enum. Takes a device type enum from tensorrt and create a ``torch_tensorrt.EngineCapability``. If the source is not supported or the engine capability is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.EngineCapability.try_from()`` Arguments: c (Union(tensorrt.EngineCapability, EngineCapability)): Engine capability enum from another library Returns: EngineCapability: Equivalent ``torch_tensorrt.EngineCapability`` to ``c`` Raises: TypeError: Unknown source type or unsupported engine capability Examples: .. code:: py torchtrt_ec = torch_tensorrt.EngineCapability._from(tensorrt.EngineCapability.SAFETY) )Provided an unsupported engine capabilityrrz@Provided unsupported source type for EngineCapability conversion) rr8rSTANDARDSAFETYDLA_STANDALONErVr rJrKrr7)rNrrrrrrOs.          zEngineCapability._fromOptional[EngineCapability]c CsRzt|}|WSttfy(}ztjd|dddWYd}~dSd}~ww)abCreate a Torch-TensorRT engine capability enum from a TensorRT engine capability enum. Takes a device type enum from tensorrt and create a ``torch_tensorrt.EngineCapability``. If the source is not supported or the engine capability level is not supported in Torch-TensorRT, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.EngineCapability.try_from()`` Arguments: c (Union(tensorrt.EngineCapability, EngineCapability)): Engine capability enum from another library Returns: EngineCapability: Equivalent ``torch_tensorrt.EngineCapability`` to ``c`` Examples: .. code:: py torchtrt_safety_ec = torch_tensorrt.EngineCapability._from(tensorrt.EngineCapability.SAEFTY) rSz) to torch_tensorrt.EngineCapablity failedTrTN)rrOrVr7r5rW)rrXrYrrrrZrzEngineCapability.try_fromr 9Union[Type[trt.EngineCapability], Type[EngineCapability]]cCs|tjkr$|tjkrtjjS|tjkrtjjS|tjkr tjjStd|tkr*|StjrWddlm }||jkrW|tjkrA|jjS|tjkrJ|jjS|tjkrS|jjStdt d)aConvert ``EngineCapability`` into the equivalent type in tensorrt Converts ``self`` into one of torch or tensorrt equivalent engine capability. If ``self`` is not supported in the target library, then an exception will be raised. As such it is not recommended to use this method directly. Alternatively use ``torch_tensorrt.EngineCapability.try_to()`` Arguments: t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))): Engine capability enum from another library to convert to Returns: Union(tensorrt.EngineCapability, EngineCapability): Engine capability equivalent ``torch_tensorrt.EngineCapability`` in enum ``t`` Raises: TypeError: Unknown target type or unsupported engine capability Examples: .. code:: py # Succeeds torchtrt_dla_ec = torch_tensorrt.EngineCapability.DLA_STANDALONE.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA rrrzKProvided unsupported destination type for engine capability type conversion) r8rrrrrVr rJrKrr7rrrrr_s.         zEngineCapability.to7Optional[Union[trt.EngineCapability, EngineCapability]]c Cr)a"Convert ``EngineCapability`` into the equivalent type in tensorrt Converts ``self`` into one of torch or tensorrt equivalent engine capability. If ``self`` is not supported in the target library, then ``None`` will be returned. Arguments: t (Union(Type(tensorrt.EngineCapability), Type(EngineCapability))): Engine capability enum from another library to convert to Returns: Optional(Union(tensorrt.EngineCapability, EngineCapability)): Engine capability equivalent ``torch_tensorrt.EngineCapability`` in enum ``t`` Examples: .. code:: py # Succeeds trt_dla_ec = torch_tensorrt.EngineCapability.DLA.to(tensorrt.EngineCapability) # Returns tensorrt.EngineCapability.DLA_STANDALONE z9torch_tensorrt.EngineCapablity conversion to target type raTrTNrbrrrrrd!s  zEngineCapability.try_torercCrfrg)rrOrrhrirrrrk@rlzEngineCapability.__eq__r\cCrmrgrnrprrrrqDrrzEngineCapability.__hash__N)rrrr)rrrr)r rrr)r rrr)rerrrrt)rurvrwrxrrrrrrOrZr_rdrkrqrrrrrqs  9  " 9 rc@sTeZdZdZeZ eZ eZ eZe d ddZ d ddZ e d dd Z d S) PlatformzZ Specifies a target OS and CPU architecture that a Torch-TensorRT program targets rcCsddl}|dr(|drtjS|dr%tjStj S|dr=|dr=tjStj S)z Returns an enum for the current platform Torch-TensorRT is running on Returns: Platform: Current platform rNlinuxaarch64x86_64windowsamd64) platformsystemlower startswithmachiner LINUX_AARCH64 LINUX_X86_64 WIN_X86_64r)rNrrrrcurrent_platformdszPlatform.current_platformrIcCrmrg)rInamerprrr__str__|rrzPlatform.__str__cCsZtjj}|tjkrtjj}|S|tjkr tjj}|S|tj kr+tjj }|Srg) ropstensorrt_platform_unknownrr_platform_linux_x86_64r_platform_linux_aarch64r_platform_win_x86_64)r^valrrr_to_serialized_rt_platforms     z#Platform._to_serialized_rt_platformN)rr)rrI)rurvrwrxrrrrrrrrr rrrrrrHs  r) __future__rr5enumrrtypingrrrrnumpyrrr8rtorch_tensorrt._featuresr r torch_tensorrt._utilsr r rrrrrrrrs.    WX