# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_tensor_map.pyx

from __future__ import annotations

from dataclasses import dataclass

import numpy
from cuda.bindings import cydriver
from cuda.core._device import Device


class TensorMapDataType:
    """Data types for tensor map descriptors.

    These correspond to the ``CUtensorMapDataType`` driver enum values.
    """
    UINT8 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8
    UINT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16
    UINT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32
    INT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32
    UINT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64
    INT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64
    FLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16
    FLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32
    FLOAT64 = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64
    BFLOAT16 = cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16
    FLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ
    TFLOAT32 = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32
    TFLOAT32_FTZ = cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ

class TensorMapInterleave:
    """Interleave layout for tensor map descriptors.

    These correspond to the ``CUtensorMapInterleave`` driver enum values.
    """
    NONE = cydriver.CU_TENSOR_MAP_INTERLEAVE_NONE
    INTERLEAVE_16B = cydriver.CU_TENSOR_MAP_INTERLEAVE_16B
    INTERLEAVE_32B = cydriver.CU_TENSOR_MAP_INTERLEAVE_32B

class TensorMapSwizzle:
    """Swizzle mode for tensor map descriptors.

    These correspond to the ``CUtensorMapSwizzle`` driver enum values.
    """
    NONE = cydriver.CU_TENSOR_MAP_SWIZZLE_NONE
    SWIZZLE_32B = cydriver.CU_TENSOR_MAP_SWIZZLE_32B
    SWIZZLE_64B = cydriver.CU_TENSOR_MAP_SWIZZLE_64B
    SWIZZLE_128B = cydriver.CU_TENSOR_MAP_SWIZZLE_128B

class TensorMapL2Promotion:
    """L2 promotion mode for tensor map descriptors.

    These correspond to the ``CUtensorMapL2promotion`` driver enum values.
    """
    NONE = cydriver.CU_TENSOR_MAP_L2_PROMOTION_NONE
    L2_64B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_64B
    L2_128B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_128B
    L2_256B = cydriver.CU_TENSOR_MAP_L2_PROMOTION_L2_256B

class TensorMapOOBFill:
    """Out-of-bounds fill mode for tensor map descriptors.

    These correspond to the ``CUtensorMapFloatOOBfill`` driver enum values.
    """
    NONE = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
    NAN_REQUEST_ZERO_FMA = cydriver.CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA

class TensorMapIm2ColWideMode:
    """Im2col wide mode for tensor map descriptors.

    This enum is always defined for API stability, but the
    :meth:`TensorMapDescriptor._from_im2col_wide` factory requires a CUDA 13+
    build and will raise otherwise.
    """
    W = 0
    W128 = 1

@dataclass
class TensorMapDescriptorOptions:
    """Options for :meth:`cuda.core.StridedMemoryView.as_tensor_map`.

    Attributes
    ----------
    box_dim : tuple[int, ...]
        Tile size for each tensor dimension, expressed in elements.
    element_strides : tuple[int, ...], optional
        Per-dimension element traversal strides.
    data_type : object, optional
        Explicit dtype override. Prefer NumPy or ``ml_dtypes`` dtype objects;
        :class:`TensorMapDataType` remains accepted for compatibility.
    interleave : TensorMapInterleave, optional
        Interleave layout. Default ``NONE``.
    swizzle : TensorMapSwizzle, optional
        Swizzle mode. Default ``NONE``.
    l2_promotion : TensorMapL2Promotion, optional
        L2 promotion mode. Default ``NONE``.
    oob_fill : TensorMapOOBFill, optional
        Out-of-bounds fill mode. Default ``NONE``.
    """
    box_dim: tuple[int, ...]
    element_strides: tuple[int, ...] | None = None
    data_type: object = None
    interleave: TensorMapInterleave = TensorMapInterleave.NONE
    swizzle: TensorMapSwizzle = TensorMapSwizzle.NONE
    l2_promotion: TensorMapL2Promotion = TensorMapL2Promotion.NONE
    oob_fill: TensorMapOOBFill = TensorMapOOBFill.NONE

    def __post_init__(self) -> None:
        ...

class TensorMapDescriptor:
    """Describes a TMA (Tensor Memory Accelerator) tensor map for Hopper+ GPUs.

    A ``TensorMapDescriptor`` wraps the opaque 128-byte ``CUtensorMap`` struct
    used by the hardware TMA unit for efficient bulk data movement between
    global and shared memory.

    Public tiled descriptors are created via
    :meth:`cuda.core.StridedMemoryView.as_tensor_map`. Specialized
    ``_from_*`` helpers remain private while this API surface settles, and
    descriptors can be passed directly to :func:`~cuda.core.launch` as a
    kernel argument.
    """

    def __init__(self):
        ...

    @property
    def device(self) -> Device | None:
        """Return the :obj:`~cuda.core.Device` associated with this descriptor."""

    @classmethod
    def _from_tiled(cls, view, box_dim=None, *, options=None, element_strides=None, data_type=None, interleave=..., swizzle=..., l2_promotion=..., oob_fill=...):
        """Create a tiled TMA descriptor from a validated view.

        Parameters
        ----------
        view : StridedMemoryView
            A device-accessible view with a 16-byte-aligned pointer.
        box_dim : tuple of int, optional
            The size of each tile dimension (in elements). Must have the
            same rank as the tensor and each value must be in [1, 256].
            Specified in the same (row-major) order as the tensor shape.
            Required unless ``options`` is provided.
        options : TensorMapDescriptorOptions or mapping, optional
            Bundled tiled-descriptor options. When provided, do not also pass
            ``box_dim`` or the individual option kwargs.
        element_strides : tuple of int, optional
            Per-dimension element traversal strides. Default is all 1s.
            Specified in the same (row-major) order as the tensor shape.
        data_type : dtype-like or TensorMapDataType, optional
            Explicit dtype override. If ``None``, inferred from the tensor's
            dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
            accepted for compatibility.
        interleave : TensorMapInterleave
            Interleave layout. Default ``NONE``.
        swizzle : TensorMapSwizzle
            Swizzle mode. Default ``NONE``.
        l2_promotion : TensorMapL2Promotion
            L2 promotion mode. Default ``NONE``.
        oob_fill : TensorMapOOBFill
            Out-of-bounds fill mode. Default ``NONE``.

        Returns
        -------
        TensorMapDescriptor

        Raises
        ------
        ValueError
            If the tensor rank is outside [1, 5], the pointer is not
            16-byte aligned, or dimension/stride constraints are violated.
        """

    @classmethod
    def _from_im2col(cls, view, pixel_box_lower_corner, pixel_box_upper_corner, channels_per_pixel, pixels_per_column, *, element_strides=None, data_type=None, interleave=..., swizzle=..., l2_promotion=..., oob_fill=...):
        """Create an im2col TMA descriptor from a validated view.

        Im2col layout is used for convolution-style data access patterns.

        Parameters
        ----------
        view : StridedMemoryView
            A device-accessible view with a 16-byte-aligned pointer.
        pixel_box_lower_corner : tuple of int
            Lower corner of the pixel bounding box for each spatial
            dimension (rank - 2 elements). Specified in row-major order
            matching the tensor's spatial dimensions.
        pixel_box_upper_corner : tuple of int
            Upper corner of the pixel bounding box for each spatial
            dimension (rank - 2 elements). Specified in row-major order
            matching the tensor's spatial dimensions.
        channels_per_pixel : int
            Number of channels per pixel.
        pixels_per_column : int
            Number of pixels per column.
        element_strides : tuple of int, optional
            Per-dimension element traversal strides. Default is all 1s.
        data_type : dtype-like or TensorMapDataType, optional
            Explicit dtype override. If ``None``, inferred from the tensor's
            dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
            accepted for compatibility.
        interleave : TensorMapInterleave
            Interleave layout. Default ``NONE``.
        swizzle : TensorMapSwizzle
            Swizzle mode. Default ``NONE``.
        l2_promotion : TensorMapL2Promotion
            L2 promotion mode. Default ``NONE``.
        oob_fill : TensorMapOOBFill
            Out-of-bounds fill mode. Default ``NONE``.

        Returns
        -------
        TensorMapDescriptor

        Raises
        ------
        ValueError
            If the tensor rank is outside [3, 5], the pointer is not
            16-byte aligned, or other constraints are violated.
        """

    @classmethod
    def _from_im2col_wide(cls, view, pixel_box_lower_corner_width, pixel_box_upper_corner_width, channels_per_pixel, pixels_per_column, *, element_strides=None, data_type=None, interleave=..., mode=..., swizzle=..., l2_promotion=..., oob_fill=...):
        """Create an im2col-wide TMA descriptor from a validated view.

        Im2col-wide layout loads elements exclusively along the W (width)
        dimension. This variant is supported on compute capability 10.0+
        (Blackwell and later).

        Parameters
        ----------
        view : StridedMemoryView
            A device-accessible view with a 16-byte-aligned pointer.
        pixel_box_lower_corner_width : int
            Lower corner of the pixel bounding box along the W dimension.
        pixel_box_upper_corner_width : int
            Upper corner of the pixel bounding box along the W dimension.
        channels_per_pixel : int
            Number of channels per pixel.
        pixels_per_column : int
            Number of pixels per column.
        element_strides : tuple of int, optional
            Per-dimension element traversal strides. Default is all 1s.
        data_type : dtype-like or TensorMapDataType, optional
            Explicit dtype override. If ``None``, inferred from the tensor's
            dtype. Prefer NumPy or ``ml_dtypes`` dtype objects; the enum is
            accepted for compatibility.
        interleave : TensorMapInterleave
            Interleave layout. Default ``NONE``.
        mode : TensorMapIm2ColWideMode
            Im2col wide mode. Default ``W``.
        swizzle : TensorMapSwizzle
            Swizzle mode. Default ``SWIZZLE_128B``.
        l2_promotion : TensorMapL2Promotion
            L2 promotion mode. Default ``NONE``.
        oob_fill : TensorMapOOBFill
            Out-of-bounds fill mode. Default ``NONE``.

        Returns
        -------
        TensorMapDescriptor

        Raises
        ------
        ValueError
            If the tensor rank is outside [3, 5], the pointer is not
            16-byte aligned, or other constraints are violated.
        """

    def replace_address(self, tensor: object) -> None:
        """Replace the global memory address in this tensor map descriptor.

        This is useful when the tensor data has been reallocated but the
        shape, strides, and other parameters remain the same.

        Parameters
        ----------
        tensor : object
            Any object supporting DLPack or ``__cuda_array_interface__``,
            or a :obj:`~cuda.core.StridedMemoryView`. Must refer to
            device-accessible memory with a 16-byte-aligned pointer.
        """

    def __repr__(self) -> str:
        ...
__all__ = ['TensorMapDescriptor', 'TensorMapDescriptorOptions']
_TMA_DT_UINT8 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT8)
_TMA_DT_UINT16 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT16)
_TMA_DT_UINT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT32)
_TMA_DT_INT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT32)
_TMA_DT_UINT64 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_UINT64)
_TMA_DT_INT64 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_INT64)
_TMA_DT_FLOAT16 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT16)
_TMA_DT_FLOAT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32)
_TMA_DT_FLOAT64 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT64)
_TMA_DT_BFLOAT16 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16)
_TMA_DT_FLOAT32_FTZ = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ)
_TMA_DT_TFLOAT32 = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32)
_TMA_DT_TFLOAT32_FTZ = int(cydriver.CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ)
_NUMPY_DTYPE_TO_TMA = {numpy.dtype(numpy.uint8): _TMA_DT_UINT8, numpy.dtype(numpy.uint16): _TMA_DT_UINT16, numpy.dtype(numpy.uint32): _TMA_DT_UINT32, numpy.dtype(numpy.int32): _TMA_DT_INT32, numpy.dtype(numpy.uint64): _TMA_DT_UINT64, numpy.dtype(numpy.int64): _TMA_DT_INT64, numpy.dtype(numpy.float16): _TMA_DT_FLOAT16, numpy.dtype(numpy.float32): _TMA_DT_FLOAT32, numpy.dtype(numpy.float64): _TMA_DT_FLOAT64}
_TMA_DATA_TYPE_SIZE = {_TMA_DT_UINT8: 1, _TMA_DT_UINT16: 2, _TMA_DT_UINT32: 4, _TMA_DT_INT32: 4, _TMA_DT_UINT64: 8, _TMA_DT_INT64: 8, _TMA_DT_FLOAT16: 2, _TMA_DT_FLOAT32: 4, _TMA_DT_FLOAT64: 8, _TMA_DT_BFLOAT16: 2, _TMA_DT_FLOAT32_FTZ: 4, _TMA_DT_TFLOAT32: 4, _TMA_DT_TFLOAT32_FTZ: 4}

def _normalize_tensor_map_data_type(data_type):
    ...

def _normalize_tensor_map_sequence(name, values):
    ...

def _require_tensor_map_enum(name, value, enum_type):
    ...

def _coerce_tensor_map_descriptor_options(box_dim, options, *, element_strides, data_type, interleave, swizzle, l2_promotion, oob_fill):
    ...

def _resolve_data_type(view, data_type):
    """Resolve the TMA data type from an explicit value or the view's dtype."""

def _get_validated_view(tensor):
    """Obtain a device-accessible StridedMemoryView with a 16-byte-aligned pointer."""

def _require_view_device(view, expected_device_id, operation):
    """Ensure device-local tensors match the current CUDA device.

    DLPack reports host/managed CUDA memory as ``kDLCUDAHost`` /
    ``kDLCUDAManaged`` with ``device_id=0`` regardless of the current device,
    so only true ``kDLCUDA`` tensors are rejected by device-id mismatch.
    """

def _compute_byte_strides(shape, strides, elem_size):
    """Compute byte strides from element strides or C-contiguous fallback.

    Returns a tuple of byte strides in row-major order.
    """

def _validate_element_strides(element_strides, rank):
    """Validate or default element_strides to all-ones."""