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

from __future__ import annotations

from dataclasses import dataclass

from cuda.core._memory._managed_buffer import ManagedBuffer
from cuda.core._memory._memory_pool import _MemPool
from cuda.core._stream import Stream
from cuda.core.graph import GraphBuilder
from cuda.core.typing import ManagedMemoryLocationType


@dataclass
class ManagedMemoryResourceOptions:
    """Customizable :obj:`~_memory.ManagedMemoryResource` options.

    Attributes
    ----------
    preferred_location : int | None, optional
        A location identifier (device ordinal or NUMA node ID) whose
        meaning depends on ``preferred_location_type``.
        (Default to ``None``)

    preferred_location_type : ManagedMemoryLocationType | str | None, optional
        Controls how ``preferred_location`` is interpreted.

        When set to ``None`` (the default), legacy behavior is used:
        ``preferred_location`` is interpreted as a device ordinal,
        ``-1`` for host, or ``None`` for no preference.

        When set explicitly, the type determines both the kind of
        preferred location and the valid values for
        ``preferred_location``:

        - ``"device"``: prefer a specific GPU. ``preferred_location``
          must be a device ordinal (``>= 0``).
        - ``"host"``: prefer host memory (OS-managed NUMA placement).
          ``preferred_location`` must be ``None``.
        - ``"host_numa"``: prefer a specific host NUMA node.
          ``preferred_location`` must be a NUMA node ID (``>= 0``),
          or ``None`` to derive the NUMA node from the current CUDA
          device's ``host_numa_id`` attribute (requires an active
          CUDA context).

        (Default to ``None``)
    """
    preferred_location: int | None = None
    preferred_location_type: ManagedMemoryLocationType | str | None = None

class ManagedMemoryResource(_MemPool):
    """
    A managed memory resource managing a stream-ordered memory pool.

    Managed memory is accessible from both the host and device, with automatic
    migration between them as needed.

    Parameters
    ----------
    options : ManagedMemoryResourceOptions
        Memory resource creation options.

        If set to `None`, the memory resource uses the driver's current
        stream-ordered memory pool. If no memory pool is set as current,
        the driver's default memory pool is used.

        If not set to `None`, a new memory pool is created, which is owned by
        the memory resource.

        When using an existing (current or default) memory pool, the returned
        managed memory resource does not own the pool (`is_handle_owned` is
        `False`), and closing the resource has no effect.

    Notes
    -----
    IPC (Inter-Process Communication) is not currently supported for managed
    memory pools.
    """

    def __init__(self, options: ManagedMemoryResourceOptions | dict[str, object] | None=None) -> None:
        ...

    def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> ManagedBuffer:
        """Allocate a managed-memory buffer of the requested size.

        Parameters
        ----------
        size : int
            The size of the buffer to allocate, in bytes.
        stream : :obj:`~_stream.Stream`
            Keyword-only. The stream on which to perform the allocation
            asynchronously. Must be passed explicitly; pass
            ``device.default_stream`` to use the default stream.

        Returns
        -------
        ManagedBuffer
            A :class:`ManagedBuffer` (a :class:`Buffer` subclass) that
            exposes the property-style advice API
            (``read_mostly``, ``preferred_location``, ``accessed_by``)
            and instance methods (``prefetch``, ``discard``,
            ``discard_prefetch``).
        """

    @property
    def device_id(self) -> int:
        """The preferred device ordinal, or -1 if the preferred location is not a device."""

    @property
    def preferred_location(self) -> tuple[ManagedMemoryLocationType, int | None] | None:
        """The preferred location for managed memory allocations.

        Returns ``None`` if no preferred location is set (driver decides),
        or a tuple ``(type, id)`` where *type* is one of ``"device"``,
        ``"host"``, or ``"host_numa"``, and *id* is the device ordinal,
        ``None`` (for ``"host"``), or the NUMA node ID, respectively.
        """

    @property
    def is_device_accessible(self) -> bool:
        """Return True. This memory resource provides device-accessible buffers."""

    @property
    def is_host_accessible(self) -> bool:
        """Return True. This memory resource provides host-accessible buffers."""

    @property
    def is_managed(self) -> bool:
        """Return True. This memory resource provides managed (unified) memory buffers."""
__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions']

def reset_concurrent_access_warning() -> None:
    """Reset the concurrent access warning flag for testing purposes."""