`iPldZddlmZddlmZmZddlmZmZm Z m Z ddl Z ddl m cmZddlZddlZddlmZddlmZdd ZddZGddeZGddeZdS)z Base classes for cuDNN API wrappers. This module provides abstract base classes that define common interfaces for cuDNN API wrapper classes, including validation, compilation, and execution patterns. ) annotations)ABCabstractmethod)AnyListTupleOptionalN)_convert_to_cutlass_data_typeaintbreturnc||zdz |zSN)r r s b/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/cudnn/api_base.pyceil_divrs EAI! nboolc&|dko ||dz zdkS)zCheck if n is a power of 2.rrr)rs r is_power_of_2rs q5 'a1q5ka''rceZdZdZdZedBdZedCdDd Zedd d dEdZdFdZ dGdZ dHdZ dIdZ dIdZ dJdZdJd ZdKdLd#Z dKdMd%Z dKdMd&Z dKdNd+Z dOdPd2Z dQdRd6ZdSd9ZdSd:ZdSd;ZdTdUd?Z dVdWdAZdS)XAPIBaseaMAbstract base class for cuDNN API wrappers. This class defines the common interface that all API wrapper implementations should follow, including configuration validation, compilation, and execution. Provides common functionality: - Logging via self._logger - Support validation tracking via self._is_supported - Compiled kernel caching via self._compiled_kernel - Stream management helpers Subclasses should implement the abstract methods to provide API-specific validation logic and execution behavior. Example: >>> class MyKernelAPI(APIBase): ... def __init__(self, sample_input, sample_output, config): ... super().__init__() ... self.sample_input = sample_input ... self.sample_output = sample_output ... self.config = config ... self._kernel = MyKernel ... ... def check_support(self) -> bool: ... # Validate inputs and configuration ... assert self.sample_input.dtype == torch.float32 ... self._is_supported = True ... return True ... ... def compile(self, current_stream=None): ... current_stream = self._get_default_stream(current_stream) ... self._ensure_support_checked() ... # Create and compile kernel ... kernel = self._kernel(self.config) ... self._compiled_kernel = cute.compile(kernel, ...) ... ... def execute(self, input_tensor, output_tensor, ... current_stream=None, skip_compile=False): ... current_stream = self._get_default_stream(current_stream) ... if not skip_compile: ... self._compiled_kernel(input_tensor, output_tensor, current_stream) ... else: ... # Direct execution without cached compilation ... kernel = self._kernel(self.config) ... kernel(input_tensor, output_tensor, current_stream) cd|_d|_d|_d|_t j|jj|_dS)aInitialize the API base. Sets up: - self._is_supported: Flag indicating if configuration is validated - self._kernel: Kernel instance - self._compiled_kernel: Cache for compiled kernel - self._logger: Logger instance for this class FN) _is_supported_kernel_compiled_kernel_interpret_uint8_as_fp4x2logging getLogger __class____name___loggerselfs r__init__zAPIBase.__init__Qs># $).&()@AA rrrcdS)aCheck if the current configuration is supported by the kernel. This method should validate: - Input/output tensor shapes and strides - Data types compatibility - Hardware capabilities (compute capability, memory, etc.) - Configuration parameters (tile sizes, cluster shapes, etc.) Implementations should set self._is_supported = True if valid. :return: True if the configuration is supported :rtype: bool :raises AssertionError: If a configuration requirement is not met Example: >>> def check_support(self) -> bool: ... self._logger.debug("Checking support") ... assert self.input.dtype in {torch.float16, torch.float32} ... assert self.input.shape[0] % 16 == 0, "Shape must be 16-aligned" ... self._is_supported = True ... return True Nrr&s r check_supportzAPIBase.check_support`s 0 rNcurrent_streamOptional[cuda.CUstream]NonecdS)aZCompile the kernel with the current configuration. This method should: 1. Ensure support has been checked (use self._ensure_support_checked()) 2. Get default stream if needed (use self._get_default_stream()) 3. Create the underlying kernel implementation 4. Compile the kernel using cute.compile() 5. Cache the compiled kernel in self._compiled_kernel :param current_stream: CUDA stream for compilation (optional) :type current_stream: cuda.CUstream or None :raises AssertionError: If the configuration is not supported Example: >>> def compile(self, current_stream=None): ... current_stream = self._get_default_stream(current_stream) ... self._ensure_support_checked() ... ... kernel = self._kernel(self.config) ... self._compiled_kernel = cute.compile( ... kernel, ... self.sample_input, ... self.sample_output, ... current_stream ... ) Nr)r'r+s rcompilezAPIBase.compilezs 8 rF)r+ skip_compiler0rcdS)akExecute the kernel with the provided inputs. This method should handle two execution modes: 1. With compiled kernel (skip_compile=False): Use self._compiled_kernel 2. Without compiled kernel (skip_compile=True): Create and execute kernel directly (JIT) :param args: Positional arguments (typically input/output tensors) :param current_stream: CUDA stream for execution (optional) :type current_stream: cuda.CUstream or None :param skip_compile: If False, use cached compiled kernel; If True, create and execute kernel directly :type skip_compile: bool :param kwargs: Additional keyword arguments for execution :return: Execution result (if any) :raises AssertionError: If compiled kernel is not available when skip_compile=False Example: >>> def execute(self, input_tensor, output_tensor, ... current_stream=None, skip_compile=False): ... current_stream = self._get_default_stream(current_stream) ... ... if not skip_compile: ... assert self._compiled_kernel is not None, "Kernel not compiled" ... self._logger.debug("Executing with compiled kernel") ... self._compiled_kernel(input_tensor, output_tensor, current_stream) ... else: ... self._logger.debug("Executing without compiled kernel (JIT)") ... kernel = self._kernel(self.config) ... kernel(input_tensor, output_tensor, current_stream) Nr)r'r+r0argskwargss rexecutezAPIBase.executes L rc |j|ddi|S)aConvenience method to execute the kernel. This is a shorthand for calling execute() with skip_compile=True, which bypasses the cached compiled kernel and executes directly. This is useful for one-off executions or when you want to ensure fresh compilation. :param args: Positional arguments passed to execute() :param kwargs: Keyword arguments passed to execute() :return: Result from execute() Example: >>> api = MyKernelAPI(...) >>> api.check_support() >>> # Direct execution without pre-compilation >>> api(input_tensor, output_tensor) # Equivalent to execute(..., skip_compile=True) r0T)r4)r'r2r3s r__call__zAPIBase.__call__s $t|T?????rc|jsE|j|jjd|s JddSdS)aHelper to ensure check_support() was called before compilation. If check_support() has not been called yet (self._is_supported is False), this method will automatically call it. This prevents compilation with invalid configurations. :raises AssertionError: If check_support() returns False or raises Example: >>> def compile(self, current_stream=None): ... self._ensure_support_checked() # Automatic validation ... # ... rest of compilation z2: check_support not previously called, calling nowzUnsupported configurationN)rr%infor#r$r*r&s r_ensure_support_checkedzAPIBase._ensure_support_checkedsg! E L  !8lll m m m%%'' D D)D D D' E E D Drstream cuda.CUstreamc|E|j|jjdtjS|S)aGet default CUDA stream if none provided. This is a convenience helper to handle optional stream parameters. If a stream is provided, it is returned as-is. If None, the default CUDA stream is returned. :param stream: CUDA stream or None :type stream: cuda.CUstream or None :return: CUDA stream (either provided or default) :rtype: cuda.CUstream Example: >>> def compile(self, current_stream=None): ... current_stream = self._get_default_stream(current_stream) ... # Now current_stream is guaranteed to be a valid stream Nz/: No CUDA stream provided, using default stream)r%debugr#r$cutlasscudadefault_stream)r'r:s r_get_default_streamzAPIBase._get_default_streamsF" > L  $."9jjj k k k<..00 0 rtensorOptional[torch.Tensor]ndimr namestrc|b|j|krW|jd|d|d|jt ||jz D]}|d}|S)aPad a tensor by unsqueezing at dim -1 until it reaches ndim rank. - If tensor is None, returns None. - Unsqueezes at dim -1 until tensor.ndim == ndim. - Logs final reshape for traceability. :param tensor: The tensor to pad (or None) :param ndim: Target rank (pad trailing dims until reached) :param name: Logical tensor name for logging :return: The padded tensor (or None) NzPadding  to zD from )rDr%r8shaperange unsqueezer'rBrDrE_s r_pad_tensor_to_ndimzAPIBase._pad_tensor_to_ndims"  V[4%7%7 L  NNN4NN NN O O O4&+-.. . .))"-- rc 8||j|kr|jd|d|jd|dt |j|z D]}|d}|j|kr)|jd|d|jd|d|S) aUnpad a tensor by squeezing at dim -1 until it reaches ndim rank. - If tensor is None, returns None. - Squeezes at dim -1 until tensor.ndim == ndim. - Logs final reshape for traceability. :param tensor: The tensor to unpad (or None) :param ndim: Target rank (squeeze trailing dims until reached) :param name: Logical tensor name for logging :return: The unpadded tensor (or None) Nz Unpadding z from rHDrIz resulted in shape z , expected )rDr%r8rJrKsqueezecriticalrMs r_unpad_tensor_to_ndimzAPIBase._unpad_tensor_to_ndims"  V[4%7%7 L  P4PPv|PPPPP Q Q Q6;-.. , ,++{d"" %%&l4&l&lFL&l&lei&l&l&lmmm rtensor_or_dtypetorch.Tensor | torch.dtypec|dSt|tjr|jn|}|tjkp|jo|tjkS)aCheck if tensor or dtype is an FP4x2 packed datatype. :param tensor_or_dtype: The torch tensor or dtype to check :type tensor_or_dtype: torch.Tensor | torch.dtype :return: True if tensor/dtype is an FP4x2 packed type :rtype: bool NF) isinstancetorchTensordtypefloat4_e2m1fn_x2r uint8r'rUr[s r _is_fp4x2zAPIBase._is_fp4x2,sS  "5)3OU\)R)Rg%%Xg//mT5S5lX]afalXlmrc|dSt|tjr|jn|}|tjtjhvS)zCheck if tensor or dtype is an FP8 datatype. :param tensor_or_dtype: The torch tensor or dtype to check :type tensor_or_dtype: torch.Tensor | torch.dtype :return: True if tensor/dtype is an FP8 type :rtype: bool NF)rXrYrZr[ float8_e5m2 float8_e4m3fnr^s r_is_fp8zAPIBase._is_fp89sD  "5)3OU\)R)Rg%%Xg*E,?@@@r torch.Tensorc Rtdt|Dd}|k|jd|d|jd|dt d|d|jd|d|S)zReturn index of innermost contiguous dimension (stride == 1). :raises RuntimeError: If no dimension with stride 1 is found. c3,K|]\}}|dk |VdS)rNr.0iss r z4APIBase._get_innermost_stride_dim..Ks*GG$!QQAGGrNztensor z has shape: z stride u= – innermost contiguous (stride == 1) dimension not found. )next enumeratestrider%rSrJ RuntimeError)r'rBrEidxs r_get_innermost_stride_dimz!APIBase._get_innermost_stride_dimFs GG)FMMOO"<"<GGGNN ; L ! !Q$QQFLQQ&--//QQQ    ` ` `6< ` `QWQ^Q^Q`Q` ` ` `aa a rOptional[Tuple[int, ...]]c |dS||rn|||tfdt|jD}|jd|d|jd||S|jS)aKGet the logical shape of a tensor, handling FP4x2 packed datatypes. For FP4x2 datatypes, two values are packed per byte. The innermost contiguous dimension (with stride 1) contains packed values, so the logical shape for that dimension is 2x the physical shape. :param tensor: The tensor to get shape from (or None) :type tensor: torch.Tensor or None :param name: Logical tensor name for logging :type name: str :return: The logical shape tuple (or None if tensor is None) :rtype: Tuple[int, ...] or None NrEc38K|]\}}|kr|dzn|VdSNr)rirjdiminnermost_dim_indexs rrlz(APIBase._tensor_shape..js:mm61cQ*=%=%=#''3mmmmmmr FP4x2 tensor z: physical shape z -> logical shape )r_rrtuplernrJr%r=)r'rBrErJrzs @r _tensor_shapezAPIBase._tensor_shapeSs$ >4 >>& ! ! "&"@"@d"@"S"S mmmmU^_e_kUlUlmmmmmE L  mtmmflmmfkmm n n nL< rcn|dS||r|||tfdt|D}|jd|d|d||S|S)a[Get the logical stride of a tensor, handling FP4x2 packed datatypes. For FP4x2 datatypes, two values are packed per byte. The strides must be adjusted to reflect logical element spacing. All strides are multiplied by 2 since each physical element contains 2 logical elements. :param tensor: The tensor to get stride from (or None) :type tensor: torch.Tensor or None :param name: Logical tensor name for logging :type name: str :return: The logical stride tuple (or None if tensor is None) :rtype: Tuple[int, ...] or None Nruc38K|]\}}|kr|dzn|VdSrwr)rirjrkrzs rrlz)APIBase._tensor_stride..s:llAQ*=%=%=AEE1llllllrr{z: physical stride z -> logical stride )r_rrr|rnror%r=)r'rBrEstridesrzs @r_tensor_stridezAPIBase._tensor_strideps$ >4 >>& ! ! #"&"@"@d"@"S"S llllQZ[a[h[h[j[jQkQklllllG L  ttttv}}ttkrtt u u uN==?? "rtensor_or_shapetorch.Tensor | Tuple[int, ...]rJ'Tuple[int, ...] | List[Tuple[int, ...]]cz|dSt|tjr|||n|}t|tr||krt |d|d|nPt|t r||vrt |d|d|nt dt||S)aYCheck if the shape of a tensor matches the expected shape(s). :param tensor_or_shape: The tensor to get shape from or the shape to check :type tensor_or_shape: torch.Tensor | Tuple[int, ...] :param shape: expected shape or list of expected shapes :type shape: Tuple[int, ...] | List[Tuple[int, ...]] :param name: Logical tensor name for logging :type name: str :raises ValueError: If the shape of the tensor does not match the expected shape(s) :return: The logical shape of the tensor :rtype: Optional[Tuple[int, ...]] Nruz! tensor shape mismatch: expected , got z( tensor shape mismatch: expected one of z*Expected shape to be a tuple or list, got )rXrYrZr}r| ValueErrorlisttype)r'rrJrE tensor_shapes r_check_tensor_shapezAPIBase._check_tensor_shapes$  "4ISTcejeqIrIrHt))/)EEEyH eU # # Yu$$ D!f!f5!f!fXd!f!fggg% t $ $ Y5(( D!m!mRW!m!m_k!m!mnnn)W$u++WWXX Xrtensor_or_stridero1Optional[Tuple[int, ...] | List[Tuple[int, ...]]] stride_orderextra_error_msg1Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]]c|dSt|tjr|||n|}t dt t |dD}|t|tr*||kr#|d|d|}|r|d |z }t|nht|tr(||vr#|d |d|}|r|d |z }t|n+d t|}|r|d |z }t||t|tr*||kr#|d |d|}|r|d |z }t|nht|tr(||vr#|d |d|}|r|d |z }t|n+dt|}|r|d |z }t|||fS)aqCheck if the stride of a tensor matches the expected stride(s) or stride order(s). :param tensor_or_stride: The tensor to get stride from or the stride to check :type tensor_or_stride: torch.Tensor | Tuple[int, ...] :param stride: The expected stride(s) :type stride: Tuple[int, ...] | List[Tuple[int, ...]] :param stride_order: The expected stride order(s) :type stride_order: Tuple[int, ...] | List[Tuple[int, ...]] :param name: Logical tensor name for logging :type name: str :param extra_error_msg: Extra error message to add to the error :type extra_error_msg: str :raises ValueError: If the stride of the tensor does not match the expected stride order :return: The stride and stride order of the tensor :rtype: Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]] N)NNruc3 K|] \}}|V dSNrrhs rrlz/APIBase._check_tensor_stride..&#g#g$!QA#g#g#g#g#g#grc|dSrrxs rz.APIBase._check_tensor_stride.. abcdaerkeyz" tensor stride mismatch: expected r: z) tensor stride mismatch: expected one of z+Expected stride to be a tuple or list, got z( tensor stride order mismatch: expected z/ tensor stride order mismatch: expected one of z1Expected stride order to be a tuple or list, got ) rXrYrZrr|sortedrnrrr) r'rrorrEr tensor_stridetensor_stride_order error_msgs r_check_tensor_stridezAPIBase._check_tensor_strides0  #:LVWginiuLvLvM++,<4+HHH}M ##g#g&=9Q9QWeWe2f2f2f#g#g#ggg  &%(( , F**#' h h6 h hYf h hI&<!%;/%;%;; $Y/// + FD)) , ..#' o oRX o o`m o oI&<!%;/%;%;; $Y/// / Y$v,,XX "8!7o!7!77I +++  #,.. ,&,66#' z zQ] z zex z zI&<!%;/%;%;; $Y/// 7 L$// ,&l::#'!B!BXd!B!Bl!B!BI&<!%;/%;%;; $Y/// ; ePTUaPbPbdd "8!7o!7!77I +++111rr[torch.dtype | List[torch.dtype]Optional[torch.dtype]c|dSt|tjr|jn|}t|tjr*||kr#|d|d|}|r|d|z }t |n\t|t r(||vr#|d|d|}|r|d|z }t |nt dt ||S)a6Check if the dtype of a tensor or dtype matches the expected dtype(s). :param tensor_or_dtype: The tensor to get dtype from or the dtype to check :type tensor_or_dtype: torch.Tensor | torch.dtype :param dtype: The expected dtype(s) :type dtype: torch.dtype | List[torch.dtype] :param name: Logical tensor name for logging :type name: str :raises ValueError: If the dtype of the tensor does not match the expected dtype(s) :return: The dtype of the tensor :rtype: Optional[torch.dtype] Nz dtype mismatch: expected rrz! dtype mismatch: expected one of z0Expected dtype to be a torch.dtype or list, got )rXrYrZr[rrr)r'rUr[rEr tensor_dtypers r _check_dtypezAPIBase._check_dtypes*&  "40:?EL0Y0Yn,,_n eU[ ) ) _u$$#ZZuZZLZZ "8!7o!7!77I +++ % t $ $ _5((#aaeaaS_aa "8!7o!7!77I +++ ) ]PTUZP[P[]]^^ ^r conditionrc(|rt|dS)a Raise a ValueError if the condition is true. :param condition: The condition to check :type condition: bool :param error_msg: The error message to raise :type error_msg: str :raises ValueError: If the condition is true N)rr'rrs r_value_error_ifzAPIBase._value_error_ifs$  (Y'' ' ( (rc(|rt|dS)aRaise a NotImplementedError if the condition is true. :param condition: The condition to check :type condition: bool :param error_msg: The error message to raise :type error_msg: str :raises NotImplementedError: If the condition is true N)NotImplementedErrorrs r_not_implemented_error_ifz!APIBase._not_implemented_error_if!s$  1%i00 0 1 1rc(|rt|dS)aRaise a RuntimeError if the condition is true. :param condition: The condition to check :type condition: bool :param error_msg: The error message to raise :type error_msg: str :raises RuntimeError: If the condition is true N)rprs r_runtime_error_ifzAPIBase._runtime_error_if-s$  *y)) ) * *r assumed_align cute.Pointerc|dStjt|j|j|tjj|S)a:Make a cute.Pointer for a tensor. :param tensor: The tensor to make a cute.Pointer for :type tensor: torch.Tensor :param assumed_align: The assumed alignment of the tensor :type assumed_align: int :return: A cute.Pointer for the tensor :rtype: cute.Pointer N)interpret_uint8_as_fp4x2r) cuteruntimemake_ptrr r[r data_ptr AddressSpacegmem)r'rBrs r_make_cute_pointerzAPIBase._make_cute_pointer9s\ >4|$$ )&,QUQo p p p OO     "' %   r5Tuple[cute.Pointer, Tuple[int, ...], Tuple[int, ...]]c |dS|||}|||}|||}tdt t |dD}|||fS)aMake a cute.Pointer, shape, and order for a tensor. :param tensor: The tensor to make a cute.Pointer, shape, and order for :type tensor: torch.Tensor :param assumed_align: The assumed alignment of the tensor :type assumed_align: int :param name: Logical tensor name for logging :type name: str :return: A cute.Pointer, shape, and stride order for the tensor :rtype: Tuple[cute.Pointer, Tuple[int, ...], Tuple[int, ...]] N)NNNrruc3 K|] \}}|V dSrrrhs rrlz7APIBase._make_cute_tensor_descriptor.._rrc|dSrrrs rrz6APIBase._make_cute_tensor_descriptor.._rrr)rr}rr|rrn)r'rBrrE tensor_ptrrrrs r_make_cute_tensor_descriptorz$APIBase._make_cute_tensor_descriptorLs >##,,V=,QQ ))&t)<< ++F+>> ##g#g&=9Q9QWeWe2f2f2f#g#g#ggg<)<<>> result = TupleDict(a=1, b=2, c=3) >>> x, y, z = result # Unpacks as (1, 2, 3) >>> result['a'] # Returns 1 >>> result[0] # Returns 1 (integer indexing) ctj|i|t||_dSr)superr(rkeys_keys)r'r2r3r#s rr(zTupleDict.__init__qs:$)&)))$))++&& rc*fdjDS)z;Iterate over values in insertion order for tuple unpacking.c3(K|] }|V dSrr)rikr's rrlz%TupleDict.__iter__..xs',,AQ,,,,,,r)rr&s`r__iter__zTupleDict.__iter__vs,,,,,,,,rcTt|trr|dks|t|jkr(t d|dt|jdt |j|St |S)z-Support both string keys and integer indices.rzindex z! out of range for TupleDict with z items)rXr lenr IndexErrorr __getitem__)r'rr#s rrzTupleDict.__getitem__zs c3   8Qww#TZ00 !g#!g!gPSTXT^P_P_!g!g!ghhh77&&tz#77 7ww""3'''r)r$rrrr(rr __classcell__)r#s@rrrcsj  ''''' ---(((((((((rr)r r r r rr )rr rr)r __future__rabcrrtypingrrrr r!cuda.bindings.driverbindingsdriverr?r>rY cutlass.cutercudnn.datatypesr rrrdictrrrrrsH#"""""########------------######### 999999(((( =====c===D((((((((((r