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

"""Tensor bridge: extract PyTorch tensor metadata via the AOTI stable C ABI.

PyTorch is NOT required at build time.  At runtime the AOTI symbols are
resolved from ``torch._C`` (which is loaded with ``RTLD_GLOBAL``).

The ``pyobj_to_aten_handle`` trick exploits the internal layout of
``THPVariable`` (PyTorch's Python tensor wrapper).

In PyTorch 2.10+ ``cdata`` is ``at::Tensor`` directly::

    struct THPVariable {
        PyObject_HEAD
        at::Tensor cdata;   // <-- &cdata is usable as AtenTensorHandle
        ...
    };

In PyTorch 2.3–2.9 ``cdata`` was ``c10::MaybeOwned<at::Tensor>``,
whose first member is ``bool isBorrowed_`` (padded to 8 bytes),
followed by the ``at::Tensor`` union member::

    struct THPVariable {
        PyObject_HEAD
        c10::MaybeOwned<at::Tensor> cdata;
        // MaybeOwned layout: { bool isBorrowed_ (8 bytes); at::Tensor own_; }
        ...
    };

In both cases the address of the ``at::Tensor`` inside ``cdata`` is
accepted by the AOTI stable C ABI functions as an ``AtenTensorHandle``.
The extra 8-byte skip for the ``isBorrowed_`` member is determined
at runtime from the PyTorch version (see ``_get_cdata_extra_offset``).

Offsetting past ``PyObject_HEAD`` gives us the handle
without any Python attribute access or method calls (~14 ns for all
7 metadata queries).

Credit: Emilio Castillo (ecastillo@nvidia.com) – original tensor-bridge POC.

.. note::

   This module must NOT be imported at ``cuda.core`` load time.  It is
   loaded lazily (by ``_memoryview.pyx``) only when the user actually
   passes a ``torch.Tensor``.  The caller must ensure that
   ``torch._C`` has been re-opened with ``RTLD_GLOBAL`` *before*
   importing this module so that the AOTI symbols are visible.
"""
from __future__ import annotations

import numpy
from cuda.core._memoryview import StridedMemoryView

AOTITorchError = int

def sync_torch_stream(device_index: int, consumer_s: int) -> int:
    """Establish stream ordering between PyTorch's current CUDA stream
    and the given consumer stream.

    Records an event on PyTorch's current stream (the producer) and makes
    the consumer stream wait on it.  This is a no-op if both streams are
    the same.
    """

def resolve_aoti_dtype(dtype_code: int) -> numpy.dtype:
    """Python-callable wrapper around _get_aoti_dtype (for lazy resolution)."""

def view_as_torch_tensor(obj: object, stream_ptr: object, view: StridedMemoryView | None=None) -> StridedMemoryView:
    """Create/populate a :class:`StridedMemoryView` from a ``torch.Tensor``.

    This is a fast path that avoids DLPack/CAI protocol overhead by
    reading tensor metadata directly through the AOTI stable C ABI.

    Parameters
    ----------
    obj : torch.Tensor
        The source tensor.
    stream_ptr : int or None
        Consumer stream pointer.  When not ``-1``, stream ordering is
        established between PyTorch's current CUDA stream (the producer)
        and the consumer stream, matching the DLPack contract.
    view : StridedMemoryView, optional
        If provided, populate this existing view in-place.  Otherwise a
        new instance is created.
    """