# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause """ 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. """ from __future__ import annotations from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Any, List, Tuple, Optional import logging import cuda.bindings.driver as cuda import cutlass import torch import cutlass.cute as cute from cudnn.datatypes import _convert_to_cutlass_data_type def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b def is_power_of_2(n: int) -> bool: """Check if n is a power of 2.""" return n > 0 and (n & (n - 1)) == 0 @dataclass(frozen=True) class TensorDesc: """Metadata needed to validate/compile tensor signatures without storage.""" dtype: torch.dtype shape: Tuple[int, ...] stride: Tuple[int, ...] stride_order: Tuple[int, ...] ndim: int = field(init=False) def __post_init__(self): shape = tuple(self.shape) stride = tuple(self.stride) stride_order = tuple(self.stride_order) ndim = len(shape) if len(stride) != ndim: raise ValueError(f"Stride rank mismatch: expected {ndim}, got {len(stride)}") if len(stride_order) != ndim: raise ValueError(f"Stride order rank mismatch: expected {ndim}, got {len(stride_order)}") if tuple(sorted(stride_order)) != tuple(range(ndim)): raise ValueError(f"Stride order must be a permutation of [0, {ndim - 1}], got {stride_order}") object.__setattr__(self, "shape", shape) object.__setattr__(self, "stride", stride) object.__setattr__(self, "stride_order", stride_order) object.__setattr__(self, "ndim", ndim) @staticmethod def _normalize_dim(dim: int, ndim: int, *, allow_new_dim: bool = False) -> int: min_dim = -ndim - (1 if allow_new_dim else 0) max_dim = ndim if allow_new_dim else ndim - 1 if dim < min_dim or dim > max_dim: raise IndexError(f"Dimension out of range (expected to be in range of [{min_dim}, {max_dim}], but got {dim})") if dim < 0: dim += ndim + (1 if allow_new_dim else 0) return dim @staticmethod def _compute_stride_order(shape: Tuple[int, ...], stride: Tuple[int, ...]) -> Tuple[int, ...]: return tuple(i for i, _ in sorted(enumerate(stride), key=lambda x: (x[1], shape[x[0]]))) @staticmethod def _numel(shape: Tuple[int, ...]) -> int: numel = 1 for size in shape: numel *= size return numel @staticmethod def _compute_contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: if not shape: return () strides = [0] * len(shape) running = 1 for i in range(len(shape) - 1, -1, -1): strides[i] = running running *= max(shape[i], 1) return tuple(strides) @staticmethod def _compute_view_stride( old_shape: Tuple[int, ...], old_stride: Tuple[int, ...], new_shape: Tuple[int, ...], ) -> Optional[Tuple[int, ...]]: old_numel = TensorDesc._numel(old_shape) if old_numel == 0: return TensorDesc._compute_contiguous_stride(new_shape) new_stride = [0] * len(new_shape) view_dim = len(new_shape) - 1 tensor_numel = 1 view_numel = 1 for tensor_dim in range(len(old_shape) - 1, -1, -1): tensor_numel *= old_shape[tensor_dim] is_contiguous_chunk_end = tensor_dim == 0 or ( old_shape[tensor_dim - 1] != 1 and old_stride[tensor_dim - 1] != tensor_numel * old_stride[tensor_dim] ) if is_contiguous_chunk_end: while view_dim >= 0 and (view_numel < tensor_numel or new_shape[view_dim] == 1): new_stride[view_dim] = view_numel * old_stride[tensor_dim] view_numel *= new_shape[view_dim] view_dim -= 1 if view_numel != tensor_numel: return None if tensor_dim > 0: tensor_numel = 1 view_numel = 1 if view_dim != -1: return None return tuple(new_stride) def _with_layout(self, shape: Tuple[int, ...], stride: Tuple[int, ...]) -> "TensorDesc": return TensorDesc( dtype=self.dtype, shape=shape, stride=stride, stride_order=self._compute_stride_order(shape, stride), ) def __len__(self) -> int: if self.ndim == 0: raise TypeError("len() of a 0-d tensor") return self.shape[0] def size(self, dim: Optional[int] = None) -> torch.Size | int: if dim is None: return torch.Size(self.shape) dim = self._normalize_dim(int(dim), self.ndim) return self.shape[dim] def permute(self, *dims: int | Tuple[int, ...] | List[int]) -> "TensorDesc": if len(dims) == 1 and isinstance(dims[0], (tuple, list)): dims = tuple(dims[0]) dims = tuple(int(d) for d in dims) if len(dims) != self.ndim: raise RuntimeError(f"permute(): expected {self.ndim} dims, got {len(dims)}") dims = tuple(self._normalize_dim(d, self.ndim) for d in dims) if len(set(dims)) != self.ndim: raise RuntimeError(f"permute(): dims must be unique, got {dims}") new_shape = tuple(self.shape[d] for d in dims) new_stride = tuple(self.stride[d] for d in dims) return self._with_layout(new_shape, new_stride) def transpose(self, dim0: int, dim1: int) -> "TensorDesc": dim0 = self._normalize_dim(dim0, self.ndim) dim1 = self._normalize_dim(dim1, self.ndim) if dim0 == dim1: return self dims = list(range(self.ndim)) dims[dim0], dims[dim1] = dims[dim1], dims[dim0] return self.permute(dims) def squeeze(self, dim: Optional[int | Tuple[int, ...] | List[int]] = None) -> "TensorDesc": if dim is None: keep_dims = [i for i, size in enumerate(self.shape) if size != 1] elif isinstance(dim, (tuple, list)): squeeze_dims = tuple(self._normalize_dim(int(d), self.ndim) for d in dim) if len(set(squeeze_dims)) != len(squeeze_dims): raise RuntimeError(f"squeeze(): dims must be unique, got {squeeze_dims}") squeeze_dims = {d for d in squeeze_dims if self.shape[d] == 1} keep_dims = [i for i in range(self.ndim) if i not in squeeze_dims] else: squeeze_dim = self._normalize_dim(int(dim), self.ndim) if self.shape[squeeze_dim] != 1: return self keep_dims = [i for i in range(self.ndim) if i != squeeze_dim] new_shape = tuple(self.shape[i] for i in keep_dims) new_stride = tuple(self.stride[i] for i in keep_dims) if new_shape == self.shape and new_stride == self.stride: return self return self._with_layout(new_shape, new_stride) def unsqueeze(self, dim: int) -> "TensorDesc": dim = self._normalize_dim(dim, self.ndim, allow_new_dim=True) if dim >= self.ndim: inserted_stride = 1 else: inserted_stride = self.stride[dim] * self.shape[dim] new_shape = self.shape[:dim] + (1,) + self.shape[dim:] new_stride = self.stride[:dim] + (inserted_stride,) + self.stride[dim:] return self._with_layout(new_shape, new_stride) def contiguous(self) -> "TensorDesc": contiguous_stride = self._compute_contiguous_stride(self.shape) if contiguous_stride == self.stride: return self return self._with_layout(self.shape, contiguous_stride) def view(self, *shape: int | Tuple[int, ...] | List[int]) -> "TensorDesc": if len(shape) == 1 and isinstance(shape[0], (tuple, list)): shape = tuple(shape[0]) new_shape = tuple(int(s) for s in shape) old_numel = self._numel(self.shape) infer_dim = None known_numel = 1 for i, size in enumerate(new_shape): if size == -1: if infer_dim is not None: raise RuntimeError("only one dimension can be inferred") infer_dim = i continue if size < 0: raise RuntimeError(f"invalid shape dimension {size}") known_numel *= size if infer_dim is not None: if known_numel == 0 or old_numel % known_numel != 0: raise RuntimeError(f"shape '{new_shape}' is invalid for input of size {old_numel}") inferred_size = old_numel // known_numel new_shape = new_shape[:infer_dim] + (inferred_size,) + new_shape[infer_dim + 1 :] known_numel *= inferred_size if known_numel != old_numel: raise RuntimeError(f"shape '{new_shape}' is invalid for input of size {old_numel}") new_stride = self._compute_view_stride(self.shape, self.stride, new_shape) if new_stride is None: raise RuntimeError( "view size is not compatible with input tensor's size and stride " "(at least one dimension spans across two contiguous subspaces)" ) return self._with_layout(new_shape, new_stride) def as_strided( self, size: Tuple[int, ...] | List[int], stride: Tuple[int, ...] | List[int], storage_offset: int = 0, ) -> "TensorDesc": if storage_offset != 0: raise RuntimeError("TensorDesc.as_strided(): non-zero storage_offset is unsupported") if not isinstance(size, (tuple, list)) or not isinstance(stride, (tuple, list)): raise TypeError("TensorDesc.as_strided(): size and stride must be tuple/list") size = tuple(int(s) for s in size) stride = tuple(int(s) for s in stride) if len(size) != len(stride): raise RuntimeError(f"TensorDesc.as_strided(): mismatch in length of size ({len(size)}) and stride ({len(stride)})") if any(s < 0 for s in size): raise RuntimeError(f"TensorDesc.as_strided(): invalid size, got {size}") if any(s < 0 for s in stride): raise RuntimeError(f"TensorDesc.as_strided(): invalid stride, got {stride}") return self._with_layout(size, stride) class APIBase(ABC): """Abstract 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): ... 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): ... current_stream = self._get_default_stream(current_stream) ... self._compiled_kernel(input_tensor, output_tensor, current_stream) """ def __init__(self): """Initialize 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 """ self._is_supported = False self._kernel = None self._compiled_kernel = None self._interpret_uint8_as_fp4x2 = False self._logger = logging.getLogger(self.__class__.__name__) @abstractmethod def check_support(self) -> bool: """Check 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 """ pass @abstractmethod def compile(self) -> None: """Compile the kernel with the current configuration. This method should: 1. Ensure support has been checked (use self._ensure_support_checked()) 2. Create the underlying kernel implementation and fake cute tensors from the sample tensor descriptors 3. Compile the kernel using cute.compile() 4. Cache the compiled kernel in self._compiled_kernel :raises AssertionError: If the configuration is not supported Example: >>> def compile(self): ... self._ensure_support_checked() ... ... kernel = self._kernel(self.config) ... sample_input_cute = self._make_fake_cute_tensor_from_desc(self.sample_input) ... sample_output_cute = self._make_fake_cute_tensor_from_desc(self.sample_output) ... ... self._compiled_kernel = cute.compile( ... kernel, ... sample_input_cute, ... sample_output_cute ... ) """ pass @abstractmethod def execute( self, *args, current_stream: Optional[cuda.CUstream] = None, **kwargs, ) -> Any: """Execute the kernel with the provided inputs. This method should execute using the cached compiled kernel. :param args: Positional arguments (typically input/output tensors) :param current_stream: CUDA stream for execution (optional) :type current_stream: cuda.CUstream or None :param kwargs: Additional keyword arguments for execution :return: Execution result (if any) :raises AssertionError: If compiled kernel is not available Example: >>> def execute(self, input_tensor, output_tensor, current_stream=None): ... current_stream = self._get_default_stream(current_stream) ... 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) """ pass def __call__(self, *args, **kwargs) -> Any: """Convenience method to execute the kernel. This is a shorthand for compiling (if needed) and then executing. :param args: Positional arguments passed to execute() :param kwargs: Keyword arguments passed to execute() :return: Result from execute() Example: >>> api = MyKernelAPI(...) >>> api.check_support() >>> api.compile() >>> api(input_tensor, output_tensor) """ if self._compiled_kernel is None: self.compile() return self.execute(*args, **kwargs) def _ensure_support_checked(self) -> None: """Helper 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): ... self._ensure_support_checked() # Automatic validation ... # ... rest of compilation """ if not self._is_supported: self._logger.info(f"{self.__class__.__name__}: check_support not previously called, calling now") assert self.check_support(), "Unsupported configuration" def _get_default_stream(self, stream: Optional[cuda.CUstream]) -> cuda.CUstream: """Get 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 execute(self, input_tensor, output_tensor, current_stream=None): ... current_stream = self._get_default_stream(current_stream) ... # Now current_stream is guaranteed to be a valid stream """ if stream is None: self._logger.debug(f"{self.__class__.__name__}: No CUDA stream provided, using default stream") return cutlass.cuda.default_stream() return stream def _pad_tensor_to_ndim( self, tensor: Optional[torch.Tensor | TensorDesc], ndim: int, name: str, ) -> Optional[torch.Tensor | TensorDesc]: """Pad a tensor/descriptor by unsqueezing at dim -1 until it reaches ndim rank. - If tensor is None, returns None. - Unsqueezes at dim -1 until tensor/descriptor rank == ndim. - Logs final reshape for traceability. :param tensor: The tensor/descriptor to pad (or None) :param ndim: Target rank (pad trailing dims until reached) :param name: Logical tensor name for logging :return: The padded tensor/descriptor (or None) """ if tensor is None: return None if tensor.ndim < ndim: self._logger.info(f"Padding {name} to {ndim}D from {tensor.shape}") for _ in range(ndim - tensor.ndim): tensor = tensor.unsqueeze(-1) return tensor def _unpad_tensor_to_ndim( self, tensor: Optional[torch.Tensor | TensorDesc], ndim: int, name: str, ) -> Optional[torch.Tensor | TensorDesc]: """Unpad a tensor/descriptor by squeezing at dim -1 until it reaches ndim rank. - If tensor is None, returns None. - Squeezes at dim -1 until tensor/descriptor rank == ndim. - Logs final reshape for traceability. :param tensor: The tensor/descriptor to unpad (or None) :param ndim: Target rank (squeeze trailing dims until reached) :param name: Logical tensor name for logging :return: The unpadded tensor/descriptor (or None) """ if tensor is None: return None if tensor.ndim > ndim: self._logger.info(f"Unpadding {name} from {tensor.shape} to {ndim}D") for _ in range(tensor.ndim - ndim): if tensor.shape and tensor.shape[-1] == 1: tensor = tensor.squeeze(-1) else: break if tensor.ndim != ndim: self._logger.critical(f"Unpadding {name} resulted in shape {tensor.shape}, expected {ndim}D") return tensor def _is_fp4x2(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: """Check 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 """ if tensor_or_dtype is None: return False dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype return (dtype == torch.float4_e2m1fn_x2) or (self._interpret_uint8_as_fp4x2 and dtype == torch.uint8) def _is_fp8(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: """Check 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 """ if tensor_or_dtype is None: return False dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype return dtype in {torch.float8_e5m2, torch.float8_e4m3fn} def _get_innermost_stride_dim(self, tensor: torch.Tensor, name: str = "") -> int: """Return index of innermost contiguous dimension (stride == 1). :raises RuntimeError: If no dimension with stride 1 is found. """ idx = next((i for i, s in enumerate(tensor.stride()) if s == 1), None) if idx is None: self._logger.critical( f"tensor {name} has shape: {tensor.shape} stride {tensor.stride()} – innermost contiguous (stride == 1) dimension not found. " ) raise RuntimeError(f"tensor {name} has shape: {tensor.shape} stride {tensor.stride()} – innermost contiguous (stride == 1) dimension not found. ") return idx def _tensor_shape( self, tensor: Optional[torch.Tensor | TensorDesc], name: str = "", ) -> Optional[Tuple[int, ...]]: """Get 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 """ if tensor is None: return None if isinstance(tensor, TensorDesc): return tensor.shape if self._is_fp4x2(tensor): innermost_dim_index = self._get_innermost_stride_dim(tensor, name=name) shape = tuple(dim * 2 if i == innermost_dim_index else dim for i, dim in enumerate(tensor.shape)) self._logger.debug(f"FP4x2 tensor {name}: physical shape {tensor.shape} -> logical shape {shape}") return shape else: return tensor.shape def _tensor_stride( self, tensor: Optional[torch.Tensor | TensorDesc], name: str = "", ) -> Optional[Tuple[int, ...]]: """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 """ if tensor is None: return None if isinstance(tensor, TensorDesc): return tensor.stride if self._is_fp4x2(tensor): innermost_dim_index = self._get_innermost_stride_dim(tensor, name=name) strides = tuple(s * 2 if i != innermost_dim_index else s for i, s in enumerate(tensor.stride())) self._logger.debug(f"FP4x2 tensor {name}: physical stride {tensor.stride()} -> logical stride {strides}") return strides else: return tensor.stride() def _check_tensor_shape( self, tensor_or_shape: torch.Tensor | TensorDesc | Tuple[int, ...], shape: Tuple[int, ...] | List[Tuple[int, ...]], name: str = "", ) -> Optional[Tuple[int, ...]]: """Check 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 | TensorDesc | 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, ...]] """ if tensor_or_shape is None: return None tensor_shape = self._tensor_shape(tensor_or_shape, name=name) if isinstance(tensor_or_shape, (torch.Tensor, TensorDesc)) else tensor_or_shape if isinstance(shape, tuple): if tensor_shape != shape: raise ValueError(f"{name} tensor shape mismatch: expected {shape}, got {tensor_shape}") elif isinstance(shape, list): if tensor_shape not in shape: raise ValueError(f"{name} tensor shape mismatch: expected one of {shape}, got {tensor_shape}") else: raise ValueError(f"Expected shape to be a tuple or list, got {type(shape)}") return tensor_shape def _check_tensor_stride( self, tensor_or_stride: torch.Tensor | TensorDesc | Tuple[int, ...], stride: Optional[Tuple[int, ...] | List[Tuple[int, ...]]] = None, stride_order: Optional[Tuple[int, ...] | List[Tuple[int, ...]]] = None, name: str = "", extra_error_msg: str = "", ) -> Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]]: """Check 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 | TensorDesc | 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, ...]]] """ if tensor_or_stride is None: return None, None if isinstance(tensor_or_stride, TensorDesc): tensor_stride = tensor_or_stride.stride tensor_stride_order = tensor_or_stride.stride_order elif isinstance(tensor_or_stride, torch.Tensor): tensor_stride = self._tensor_stride(tensor_or_stride, name=name) tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: x[1])) else: tensor_stride = tensor_or_stride tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: x[1])) if stride is not None: if isinstance(stride, tuple): if tensor_stride != stride: error_msg = f"{name} tensor stride mismatch: expected {stride}, got {tensor_stride}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) elif isinstance(stride, list): if tensor_stride not in stride: error_msg = f"{name} tensor stride mismatch: expected one of {stride}, got {tensor_stride}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) else: error_msg = f"Expected stride to be a tuple or list, got {type(stride)}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) if stride_order is not None: if isinstance(stride_order, tuple): if tensor_stride_order != stride_order: error_msg = f"{name} tensor stride order mismatch: expected {stride_order}, got {tensor_stride_order}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) elif isinstance(stride_order, list): if tensor_stride_order not in stride_order: error_msg = f"{name} tensor stride order mismatch: expected one of {stride_order}, got {tensor_stride_order}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) else: error_msg = f"Expected stride order to be a tuple or list, got {type(stride_order)}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) return tensor_stride, tensor_stride_order def _check_dtype( self, tensor_or_dtype: torch.Tensor | TensorDesc | torch.dtype, dtype: torch.dtype | List[torch.dtype], name: str = "", extra_error_msg: str = "", ) -> Optional[torch.dtype]: """Check 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 | TensorDesc | 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] """ if tensor_or_dtype is None: return None tensor_dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype if isinstance(dtype, torch.dtype): if tensor_dtype != dtype: error_msg = f"{name} dtype mismatch: expected {dtype}, got {tensor_dtype}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) elif isinstance(dtype, list): if tensor_dtype not in dtype: error_msg = f"{name} dtype mismatch: expected one of {dtype}, got {tensor_dtype}" if extra_error_msg: error_msg += f": {extra_error_msg}" raise ValueError(error_msg) else: raise ValueError(f"Expected dtype to be a torch.dtype or list, got {type(dtype)}") return tensor_dtype def _value_error_if(self, condition: bool, error_msg: str) -> None: """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 """ if condition: raise ValueError(error_msg) def _not_implemented_error_if(self, condition: bool, error_msg: str) -> None: """Raise 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 """ if condition: raise NotImplementedError(error_msg) def _runtime_error_if(self, condition: bool, error_msg: str) -> None: """Raise 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 """ if condition: raise RuntimeError(error_msg) def _make_fake_cute_tensor_like( self, tensor: torch.Tensor, assumed_align: int = 16, name: str = "", ) -> cute.Pointer: """Make a fake tensor like the provided tensor. :param tensor: The tensor to make a fake tensor like :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 fake tensor like the provided tensor :rtype: cute.Pointer """ return self._make_fake_cute_tensor_from_desc( self._make_tensor_desc(tensor, name=name), assumed_align=assumed_align, ) def _make_tensor_desc(self, tensor: Optional[torch.Tensor], name: str = "") -> Optional[TensorDesc]: """Capture logical tensor metadata that is sufficient for validation/compile.""" if tensor is None: return None tensor_shape = self._tensor_shape(tensor, name=name) tensor_stride = self._tensor_stride(tensor, name=name) tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: (x[1], tensor_shape[x[0]]))) return TensorDesc( dtype=tensor.dtype, shape=tensor_shape, stride=tensor_stride, stride_order=tensor_stride_order, ) def _make_fake_cute_tensor_from_desc( self, tensor_desc: Optional[TensorDesc], assumed_align: int = 16, ) -> Optional[cute.Pointer]: """Build a fake cute tensor from a descriptor.""" if tensor_desc is None: return None return self._make_fake_cute_tensor( dtype=tensor_desc.dtype, shape=tensor_desc.shape, stride=tensor_desc.stride, assumed_align=assumed_align, ) def _make_fake_cute_tensor( self, dtype: torch.dtype, shape: Tuple[int, ...], stride: Tuple[int, ...], assumed_align: int = 16, ) -> cute.Pointer: """Make a fake tensor. :param dtype: The dtype of the tensor :type dtype: torch.dtype :param shape: The shape of the tensor :type shape: Tuple[int, ...] :param stride: The stride of the tensor :type stride: Tuple[int, ...] :param assumed_align: The assumed alignment of the tensor :type assumed_align: int :return: A fake tensor :rtype: cute.Pointer """ return cute.runtime.make_fake_tensor( dtype=_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2), shape=shape, stride=stride, assumed_align=assumed_align, ) class TupleDict(dict): """A dictionary that supports tuple unpacking. This class extends dict to allow unpacking like a tuple while still providing dictionary-style key access. The unpacking order is determined by the _keys attribute which preserves insertion order. Example: >>> 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) """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Store keys in order for tuple unpacking self._keys = list(self.keys()) def __iter__(self): """Iterate over values in insertion order for tuple unpacking.""" return (self[k] for k in self._keys) def __getitem__(self, key): """Support both string keys and integer indices.""" if isinstance(key, int): if key < 0 or key >= len(self._keys): raise IndexError(f"index {key} out of range for TupleDict with {len(self._keys)} items") return super().__getitem__(self._keys[key]) return super().__getitem__(key)