dmlc--dgl
Xin Yao
960092be02
* add set_stream * add .record_stream for NDArray and HeteroGraph * refactor dgl stream Python APIs * test record_stream * add unit test for record stream * use pytorch's stream * fix lint * fix cpu build * address comments * address comments * add record stream tests for dgl.graph * record frames and update dataloder * add docstring * update frame * add backend check for record_stream * remove CUDAThreadEntry::stream * record stream for newly created formats * fix bug * fix cpp test * fix None c_void_p to c_handle
43 行
1.1 KiB
Python
43 行
1.1 KiB
Python
# pylint: disable=invalid-name, unused-import
|
|
"""Runtime stream APIs which are mainly for internal test use only.
|
|
For applications, please use PyTorch's stream management, of which DGL is aware.
|
|
"""
|
|
from __future__ import absolute_import
|
|
|
|
import ctypes
|
|
from .base import _LIB, check_call, _FFI_MODE
|
|
from .runtime_ctypes import DGLStreamHandle
|
|
|
|
IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError
|
|
|
|
def to_dgl_stream_handle(cuda_stream):
|
|
""" Convert torch.cuda.Stream to DGL stream handle
|
|
|
|
Parameters
|
|
----------
|
|
cuda_stream : torch.cuda.Stream.
|
|
|
|
Returns
|
|
-------
|
|
DGLStreamHandle
|
|
DGLStreamHandle of the input ``cuda_stream``.
|
|
"""
|
|
return ctypes.c_void_p(cuda_stream.cuda_stream)
|
|
|
|
def _dgl_get_stream(ctx):
|
|
"""Get the current CUDA stream of the given DGL context.
|
|
|
|
Parameters
|
|
----------
|
|
ctx : DGL context.
|
|
|
|
Returns
|
|
-------
|
|
DGLStreamHandle
|
|
DGLStreamHandle of the current CUDA stream.
|
|
"""
|
|
current_cuda_stream = DGLStreamHandle()
|
|
check_call(_LIB.DGLGetStream(
|
|
ctx.device_type, ctx.device_id, ctypes.byref(current_cuda_stream)))
|
|
return current_cuda_stream
|