DLPack and CUDA Stream Limitations
==================================

The Python wrappers accept CUDA arrays through DLPack and return DLPack output
capsules. The guidance on raw capsules, input stream ordering, consumer stream
handoff, and output-buffer lifetime applies to ``VideoSuperRes``,
``VideoFrameGeneration``, and ``TrueHDR``.

The RGB8 output-conversion limitation applies only to effects configured for
RGB8 output. ``TrueHDR`` returns RGB10A2 output and does not use this conversion
path.

Raw Output Capsule
------------------

Each inference method returns an output dataclass whose ``image`` field is a raw
DLPack capsule. PyTorch accepts this format directly with
``torch.from_dlpack(result.image)``. CuPy also accepts it through its legacy
capsule compatibility path.

JAX 0.7.2 and later require an object that implements ``__dlpack__()`` and
``__dlpack_device__()`` instead of a raw capsule. An adapter wrapper is required
until the wheel exposes the full Python DLPack protocol. Additionally, copying
during DLPack import requires buffer semantics that the JAX PJRT C API backend
did not support. As a workaround, wrap the raw capsule, import the DLPack buffer
without copying, then create and synchronize a copy after the DLPack import.
Compatibility with JAX versions earlier than 0.7.2 and other JAX backend
configurations was not evaluated.

After completing the SDK producer work as described in `CUDA Stream Ordering`_,
use the following Python DLPack protocol adapter.
``result`` can be returned by ``VideoSuperRes.run()``,
``VideoFrameGeneration.run()``, ``VideoFrameGeneration.run_at_timestep()``, or
``TrueHDR.run()``. Set ``device_id`` to the CUDA device used by the effect.

.. code-block:: python

   import jax
   import jax.numpy as jnp

   class JaxDLPackAdapter:
       def __init__(self, capsule, device_id=0):
           self._capsule = capsule
           self._device_id = device_id

       def __dlpack__(self, stream=None, **kwargs):
           return self._capsule

       def __dlpack_device__(self):
           return (2, self._device_id)  # kDLCUDA

   adapter = JaxDLPackAdapter(result.image, device_id=0)
   output_view = jax.dlpack.from_dlpack(adapter, copy=False)
   output = jnp.array(output_view, copy=True)
   output.block_until_ready()

CUDA Stream Ordering
--------------------

``stream_ptr`` selects the CUDA stream used by the SDK effect and input
transfer, but it is not a complete producer-to-consumer synchronization
contract today.

The bindings acquire each input by calling the array producer's
``__dlpack__()`` without passing the SDK stream.

When input writes and ``run(..., stream_ptr=...)`` use the same stream, CUDA
orders the writes before the SDK input transfer, so no additional input
synchronization is required. When they use different streams, synchronize the
producer stream before calling ``run()``.

For output, ordering includes inference, RGB8 output conversion when
applicable, and the framework import or copy. If all these operations use the
same stream, CUDA orders them and device-wide synchronization is not required.
If the consumer imports or copies the output on a different stream, synchronize
the SDK producer stream before importing or copying the output.

The high-level wrappers currently call low-level ``get_output()`` without
forwarding ``stream_ptr``. Consequently, RGB8 inference may use the requested
SDK stream while output conversion uses stream 0. With a nonzero ``stream_ptr``,
do not assume that inference and conversion share a stream. Synchronizing only
the configured SDK stream may not cover the conversion work. In particular,
``non_blocking=True`` is unsafe because conversion may begin before inference
finishes.

For the current high-level RGB8 path, use blocking mode. Device-wide
synchronization before import is the conservative fallback because it waits for
both the SDK stream and the RGB8 conversion stream. The following example
assumes that ``sr`` is a loaded ``VideoSuperRes`` instance and ``frame`` is a
compatible CUDA input:

.. code-block:: python

   stream = torch.cuda.current_stream()
   result = sr.run(
       frame,
       stream_ptr=stream.cuda_stream,
       non_blocking=False,
   )

   # Wait for SDK work and RGB8 output conversion before import.
   torch.cuda.synchronize()
   output = torch.from_dlpack(result.image).clone()

   # Wait for the clone before reusing or closing the effect or accessing CPU data.
   stream.synchronize()

A synchronization only after ``clone()`` cannot order a clone that has already
raced SDK producer work. If inference, conversion, and the copy are all ordered
on the same stream, synchronizing that stream is sufficient and no pre-import
device-wide synchronization is needed. Device-wide synchronization may reduce
throughput.

Output Buffer Lifetime
----------------------

The DLPack output references SDK-owned memory. It remains valid only while the
effect and its current output buffer remain alive. The buffer may be reused or
released by another ``run()``, an output-size or encoding change, or
``close()``.

After SDK producer work has completed, import the capsule and immediately
clone/copy the result. If the copy and the next ``run()`` use the same stream,
CUDA orders the copy before reuse of the SDK output buffer. If they use
different streams, synchronize the copy stream before reusing the buffer.

Always wait for the copy before ``close()`` or an output-size or encoding change
because those host operations may release or replace the SDK buffer without
CUDA stream ordering. After the copy has completed, the copied output is
managed by the framework and can outlive the effect.

Do not retain a zero-copy tensor imported from ``result.image`` across another
``run()``, an output-size or encoding change, or ``close()``.

For CuPy with the current high-level RGB8 path, call
``cupy.cuda.runtime.deviceSynchronize()`` before importing the output when the
inference, conversion, and copy streams may differ. After the copy, synchronize
the CuPy copy stream before ``close()`` or an output-size or encoding change. If
the copy and next ``run()`` use the same stream, CUDA orders the copy before
buffer reuse.