tvm_ffi.use_torch_stream

tvm_ffi.use_torch_stream#

tvm_ffi.use_torch_stream(context=None)[source]#

Create an FFI stream context with a Torch stream or graph.

cuda graph or current stream if None provided.

Parameters:

context (Any = None) – The wrapped torch stream or cuda graph.

Returns:

context – The ffi stream context wrapping torch stream context.

Return type:

tvm_ffi.TorchStreamContext

Examples

s = torch.cuda.Stream()
with tvm_ffi.use_torch_stream(torch.cuda.stream(s)):
    ...

g = torch.cuda.CUDAGraph()
with tvm_ffi.use_torch_stream(torch.cuda.graph(g)):
    ...

Note

When working with raw cudaStream_t handle, using tvm_ffi.use_raw_stream() instead.

Contents