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

from __future__ import annotations

import functools
from collections.abc import Callable
from typing import Any

import numpy
from cuda.core._layout import _StridedLayout
from cuda.core._memory import Buffer
from cuda.core._stream import Stream

from ._dlpack import *


class StridedMemoryView:
    """A class holding metadata of a strided dense array/tensor.

    A :obj:`StridedMemoryView` instance can be created in three ways:

      1. Using the :obj:`args_viewable_as_strided_memory` decorator (recommended)
      2. Explicit construction relying on DLPack or CUDA Array Interface, see below.
      3. From :obj:`~_memory.Buffer` and shape and size tuples (see
         :meth:`from_buffer` classmethod)

    ``StridedMemoryView(obj, stream_ptr)`` can be used to create a view from
    objects supporting either DLPack (up to v1.0) or CUDA Array Interface
    (CAI) v3. When wrapping an arbitrary object it will try the DLPack protocol
    first, then the CAI protocol. A :obj:`BufferError` is raised if neither is
    supported.

    Since either way would take a consumer stream, for DLPack it is passed to
    ``obj.__dlpack__()`` as-is (except for :obj:`None`, see below); for CAI, a
    stream order will be established between the consumer stream and the
    producer stream (from ``obj.__cuda_array_interface__()["stream"]``), as if
    ``cudaStreamWaitEvent`` is called by this method.

    To opt-out of the stream ordering operation in either DLPack or CAI,
    please pass ``stream_ptr=-1``. Note that this deviates (on purpose)
    from the semantics of ``obj.__dlpack__(stream=None, ...)`` since ``cuda.core``
    does not encourage using the (legacy) default/null stream, but is
    consistent with the CAI's semantics. For DLPack, ``stream=-1`` will be
    internally passed to ``obj.__dlpack__()`` instead.

    Parameters
    ----------
    obj : Any
        Any objects that supports either DLPack (up to v1.0) or CUDA Array
        Interface (v3).
    stream_ptr: int
        The pointer address (as Python `int`) to the **consumer** stream.
        Stream ordering will be properly established unless ``-1`` is passed.


    Attributes
    -----------
    ptr : int
        Pointer to the tensor buffer (as a Python `int`).
    device_id : int
        The device ID for where the tensor is located. It is -1 for CPU tensors
        (meaning those only accessible from the host).
    is_device_accessible : bool
        Whether the tensor data can be accessed on the GPU.
    readonly: bool
        Whether the tensor data can be modified in place.
    exporting_obj : Any
        A reference to the original tensor object that is being viewed.
        If the view is created with :meth:`from_buffer`,
        it will be the Buffer instance passed to the method.

    """

    def __init__(self, obj: object=None, stream_ptr: int | None=None) -> None:
        ...

    @classmethod
    def from_dlpack(cls, obj: object, stream_ptr: int | None=None) -> StridedMemoryView:
        """Create a view from an object supporting the `DLPack <https://dmlc.github.io/dlpack/latest/>`_ protocol.

        Parameters
        ----------
        obj : object
            An object implementing the `DLPack <https://dmlc.github.io/dlpack/latest/>`_ protocol
            (via ``__dlpack__``).
        stream_ptr : int, optional
            Stream pointer for synchronization. If ``None``, no synchronization is performed.
        """

    @classmethod
    def from_cuda_array_interface(cls, obj: object, stream_ptr: int | None=None) -> StridedMemoryView:
        """Create a view from an object supporting the `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_ protocol.

        Parameters
        ----------
        obj : object
            An object implementing the `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_ protocol.
        stream_ptr : int, optional
            Stream pointer for synchronization. If ``None``, no synchronization is performed.
        """

    @classmethod
    def from_array_interface(cls, obj: object) -> StridedMemoryView:
        """Create a view from an object supporting the `__array_interface__ <https://numpy.org/doc/stable/reference/arrays.interface.html>`_ protocol.

        Parameters
        ----------
        obj : object
            An object implementing the `__array_interface__ <https://numpy.org/doc/stable/reference/arrays.interface.html>`_ protocol (e.g., a numpy array).
        """

    @classmethod
    def from_any_interface(cls, obj: object, stream_ptr: int | None=None) -> StridedMemoryView:
        """Create a view by automatically selecting the best available protocol.

        Tries `DLPack <https://dmlc.github.io/dlpack/latest/>`_ first, then falls back to
        `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_.
        ``torch.Tensor`` objects are transparently handled via a fast AOTI path
        regardless of which protocol is selected.

        Parameters
        ----------
        obj : object
            An object implementing `DLPack <https://dmlc.github.io/dlpack/latest/>`_ or
            `__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_.
        stream_ptr : int, optional
            Stream pointer for synchronization. If ``None``, no synchronization is performed.
        """

    @classmethod
    def from_buffer(cls, buffer: Buffer, shape: tuple[int, ...], strides: tuple[int, ...] | None=None, *, itemsize: int | None=None, dtype: numpy.dtype | None=None, is_readonly: bool=False) -> StridedMemoryView:
        """
        Creates a :obj:`StridedMemoryView` instance from a :obj:`~_memory.Buffer` and shape and strides tuples.
        The Buffer can be either allocation coming from a :obj:`MemoryResource` or an external allocation
        wrapped in a :obj:`~_memory.Buffer` object with ``Buffer.from_handle(ptr, size, owner=...)``.

        .. caution::
            When creating a :obj:`StridedMemoryView` from a :obj:`~_memory.Buffer`,
            no synchronization is performed. It is the user's responsibility to ensure
            the data in ``buffer`` is properly synchronized when consuming the view.

        Parameters
        ----------
        buffer : :obj:`~_memory.Buffer`
            The buffer to create the view from.
        shape : :obj:`tuple`
            The layout describing the shape, strides and itemsize of the elements in
            the buffer.
        strides : :obj:`tuple`
            The layout describing the shape, strides and itemsize of the elements in
            the buffer.
        dtype : :obj:`numpy.dtype`
            Optional dtype.
            If specified, the dtype's itemsize must match the layout's itemsize.
        is_readonly : bool, optional
            Whether the mark the view as readonly.
        """

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

    def view(self, layout: _StridedLayout | None=None, dtype: numpy.dtype | None=None) -> StridedMemoryView:
        """
        Creates a new view with adjusted layout and dtype.
        Same as calling :meth:`from_buffer` with the current buffer.
        """

    def as_tensor_map(self, box_dim: tuple[int, ...] | None=None, *, options: object=None, element_strides: tuple[int, ...] | None=None, data_type: object=None, interleave: object=None, swizzle: object=None, l2_promotion: object=None, oob_fill: object=None) -> object:
        """Create a tiled :obj:`TensorMapDescriptor` from this view.

        This is the public entry point for creating tiled tensor map
        descriptors in ``cuda.core``. Pass either ``box_dim`` and the
        individual keyword arguments directly, or provide bundled tiled
        options via ``options=``.
        """

    def copy_from(self, other: StridedMemoryView, stream: Stream, allocator: object=None, blocking: bool | None=None) -> None:
        """
        Copies the data from the other view into this view.

        The copy can be performed between following memory spaces:
        host-to-device, device-to-host, device-to-device (on the same device).

        Parameters
        ----------
        other : StridedMemoryView
            The view to copy data from.
        stream : Stream | None, optional
            The stream to schedule the copy on.
        allocator : MemoryResource | None, optional
            If temporary buffers are needed, the specified memory resources
            will be used to allocate the memory. If not specified, default
            resources will be used.
        blocking : bool | None, optional
            Whether the call should block until the copy is complete.
                * ``True``: the ``stream`` is synchronized with the host at the end of the call,
                  blocking until the copy is complete.
                * ``False``: if possible, the call returns immediately once the copy is scheduled.
                  However, in some cases of host-to-device or device-to-host copies, the call may
                  still synchronize with the host if necessary.
                * ``None`` (default):
                    * for device-to-device, it defaults to ``False`` (non-blocking),
                    * for host-to-device or device-to-host, it defaults to ``True`` (blocking).
        """

    def copy_to(self, other: StridedMemoryView, stream: Stream | None=None, allocator: object=None, blocking: bool | None=None) -> None:
        """
        Copies the data from this view into the ``other`` view.

        For details, see :meth:`copy_from`.
        """

    def __dlpack__(self, *, stream: int | None=None, max_version: tuple[int, int] | None=None, dl_device: tuple[int, int] | None=None, copy: bool | None=None) -> object:
        ...

    def __dlpack_device__(self) -> tuple[int, int]:
        ...

    @property
    def _layout(self) -> _StridedLayout:
        """
        The layout of the tensor. For StridedMemoryView created from DLPack or CAI,
        the layout is inferred from the tensor object's metadata.
        """

    @property
    def size(self) -> int:
        ...

    @property
    def shape(self) -> tuple[int, ...]:
        """
        Shape of the tensor.
        """

    @property
    def strides(self) -> tuple[int, ...] | None:
        """
        Strides of the tensor (in **counts**, not bytes).
        """

    @property
    def dtype(self) -> numpy.dtype | None:
        """
        Data type of the tensor.

        Supports standard NumPy dtypes as well as narrow data types (e.g., ``bfloat16``)
        when the optional `ml_dtypes <https://github.com/jax-ml/ml_dtypes>`_ package is
        installed. If ``ml_dtypes`` is not available and such a tensor is encountered,
        a :obj:`NotImplementedError` will be raised.
        """

    def __repr__(self) -> str:
        ...

class _StridedMemoryViewProxy:

    def view(self, stream_ptr=None) -> StridedMemoryView:
        ...

    def __init__(self, obj: object) -> None:
        ...
_SMV_DLPACK_EXCHANGE_API_CAPSULE = ...

def view_as_cai(obj, stream_ptr, view=None) -> StridedMemoryView:
    ...

def view_as_array_interface(obj, view=None) -> StridedMemoryView:
    ...

@functools.lru_cache
def _typestr2dtype(typestr: str) -> numpy.dtype:
    ...

@functools.lru_cache
def _typestr2itemsize(typestr: str) -> int:
    ...

def args_viewable_as_strided_memory(arg_indices: tuple[int, ...]) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """
    Decorator to create proxy objects to :obj:`StridedMemoryView` for the
    specified positional arguments.

    This allows array/tensor attributes to be accessed inside the function
    implementation, while keeping the function body array-library-agnostic (if
    desired).

    Inside the decorated function, the specified arguments become instances
    of an (undocumented) proxy type, regardless of its original source. A
    :obj:`StridedMemoryView` instance can be obtained by passing the (consumer)
    stream pointer (as a Python `int`) to the proxies's ``view()`` method. For
    example:

    .. code-block:: python

        @args_viewable_as_strided_memory((1,))
        def my_func(arg0, arg1, arg2, stream: Stream):
            # arg1 can be any object supporting DLPack or CUDA Array Interface
            view = arg1.view(stream.handle)
            assert isinstance(view, StridedMemoryView)
            ...

    Parameters
    ----------
    arg_indices : tuple
        The indices of the target positional arguments.
    """