tvm_ffi.use_raw_stream

Contents

tvm_ffi.use_raw_stream#

tvm_ffi.use_raw_stream(device, stream)[source]#

Create an FFI stream context with the given device and stream handle.

Parameters:

  • device (Device) – The device to which the stream belongs.

  • stream (int | c_void_p) – The stream handle (for example, a CUDA cudaStream_t as an integer, or 0).

Return type:

StreamContext

Returns:

context – The FFI stream context.

Examples

The example below uses a CPU device and a dummy stream handle. On CUDA, pass a real cudaStream_t integer.

import tvm_ffi

dev = tvm_ffi.device("cpu:0")
with tvm_ffi.use_raw_stream(dev, 0):
    # Within the context, the current stream for this device is set
    assert tvm_ffi.get_raw_stream(dev) == 0

See also

tvm_ffi.use_torch_stream()

Use a Torch stream or CUDA graph as the source of truth.

tvm_ffi.get_raw_stream()

Query the current FFI stream for a device.

Contents