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

from __future__ import annotations

from collections.abc import Sequence

from cuda.core._device import Device
from cuda.core._graph import GraphBuilder
from cuda.core._host import Host
from cuda.core._memory._buffer import Buffer
from cuda.core._stream import Stream
from cuda.core._utils.cuda_utils import driver


def discard_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer]) -> None:
    """Discard a batch of managed-memory ranges.

    Requires CUDA 13+. For a single buffer, use
    :meth:`ManagedBuffer.discard` instead.

    Parameters
    ----------
    stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
        Stream for the asynchronous discard. First positional, required
        (mirrors :func:`launch`).
    buffers : Sequence[:class:`Buffer`]
        Two or more managed allocations to discard. Resident pages are
        released without prefetching new contents; subsequent access is
        satisfied by lazy migration.

    Raises
    ------
    NotImplementedError
        On a CUDA 12 build of ``cuda.core``.
    """

def _do_single_discard_py(buf: Buffer, stream: Stream | GraphBuilder | None) -> None:
    """Internal: single-buffer discard for ManagedBuffer.discard()."""

def _advise_one(buf: Buffer, advice: driver.CUmem_advise, location: Device | Host | None) -> None:
    """Internal: apply managed-memory advice to a single buffer.

    Used by :class:`ManagedBuffer` property setters. Not part of the
    public API.
    """

def prefetch_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer], locations: Device | Host | Sequence[Device | Host]) -> None:
    """Prefetch a batch of managed-memory ranges to target locations.

    Requires CUDA 13+. For a single buffer, use
    :meth:`ManagedBuffer.prefetch` instead.

    Parameters
    ----------
    stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
        Stream for the asynchronous prefetch. First positional, required
        (mirrors :func:`launch`).
    buffers : Sequence[:class:`Buffer`]
        Two or more managed allocations to operate on.
    locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...]
        Target location(s). A single location applies to all buffers; a
        sequence must match ``len(buffers)``.

    Notes
    -----
    On a CUDA 12 build, falls back to a Python-level loop calling
    ``cuMemPrefetchAsync`` per buffer (no batched driver entry point on
    CUDA 12). CUDA 13 builds use ``cuMemPrefetchBatchAsync`` directly.
    """

def _do_single_prefetch_py(buf: Buffer, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None:
    """Internal: single-buffer prefetch for ManagedBuffer.prefetch().

    Uses cuMemPrefetchAsync (works on CUDA 12 and 13).
    """

def _read_preferred_location_v2(buf: Buffer) -> Device | Host | None:
    ...

def discard_prefetch_batch(stream: Stream | GraphBuilder, buffers: Sequence[Buffer], locations: Device | Host | Sequence[Device | Host]) -> None:
    """Discard a batch of managed-memory ranges and prefetch them to target locations.

    Requires CUDA 13+. For a single buffer, use
    :meth:`ManagedBuffer.discard_prefetch` instead.

    Parameters
    ----------
    stream : :class:`~_stream.Stream` | :class:`~graph.GraphBuilder`
        Stream for the asynchronous operation. First positional, required
        (mirrors :func:`launch`).
    buffers : Sequence[:class:`Buffer`]
        Two or more managed allocations to discard and re-prefetch.
    locations : :class:`~cuda.core.Device` | :class:`~cuda.core.Host` | Sequence[...]
        Target location(s). A single location applies to all buffers;
        a sequence must match ``len(buffers)``.

    Raises
    ------
    NotImplementedError
        On a CUDA 12 build of ``cuda.core``.
    """

def _do_single_discard_prefetch_py(buf: Buffer, location: Device | Host | None, stream: Stream | GraphBuilder | None) -> None:
    """Internal: single-buffer discard+prefetch for
    ManagedBuffer.discard_prefetch()."""