dmlc--dgl
nv-dlasalle
020f02498c
* Add uva by default to embedding * More updates * Update optimizer * Add new uva functions * Expose new pinned memory function * Add unit tests * Update formatting * Fix unit test * Handle auto UVA case when training is on CPU * Allow per-embedding decisions for whether to use UVA * Address spares_optim.py comments * Remove unused templates * Update unit test * Use dgl allocate memory for pinning * allow automatically unpin * workaround for d2h copy with a different dtype * fix linting * update error message * update copyright Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
70 行
2.2 KiB
Python
70 行
2.2 KiB
Python
"""Utility functions related to pinned memory tensors."""
|
|
|
|
from ..base import DGLError
|
|
from .. import backend as F
|
|
from .._ffi.function import _init_api
|
|
|
|
def pin_memory_inplace(tensor):
|
|
"""Register the tensor into pinned memory in-place (i.e. without copying).
|
|
Users are required to save the returned dgl.ndarray object to avoid being unpinned.
|
|
|
|
Parameters
|
|
----------
|
|
tensor : Tensor
|
|
The tensor to be pinned.
|
|
|
|
Returns
|
|
-------
|
|
dgl.ndarray
|
|
The dgl.ndarray object that holds the pinning status and shares the same
|
|
underlying data with the tensor.
|
|
"""
|
|
if F.backend_name in ['mxnet', 'tensorflow']:
|
|
raise DGLError("The {} backend does not support pinning " \
|
|
"tensors in-place.".format(F.backend_name))
|
|
|
|
# needs to be writable to allow in-place modification
|
|
try:
|
|
nd_array = F.zerocopy_to_dgl_ndarray_for_write(tensor)
|
|
nd_array.pin_memory_()
|
|
return nd_array
|
|
except Exception as e:
|
|
raise DGLError("Failed to pin memory in-place due to: {}".format(e))
|
|
|
|
def gather_pinned_tensor_rows(tensor, rows):
|
|
"""Directly gather rows from a CPU tensor given an indices array on CUDA devices,
|
|
and returns the result on the same CUDA device without copying.
|
|
|
|
Parameters
|
|
----------
|
|
tensor : Tensor
|
|
The tensor. Must be in pinned memory.
|
|
rows : Tensor
|
|
The rows to gather. Must be a CUDA tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The result with the same device as :attr:`rows`.
|
|
"""
|
|
return F.from_dgl_nd(_CAPI_DGLIndexSelectCPUFromGPU(F.to_dgl_nd(tensor), F.to_dgl_nd(rows)))
|
|
|
|
def scatter_pinned_tensor_rows(dest, rows, source):
|
|
"""Directly scatter rows from a GPU tensor given an indices array on CUDA devices,
|
|
to a pinned tensor on the CPU.
|
|
|
|
Parameters
|
|
----------
|
|
dest : Tensor
|
|
The tensor on the CPU to scatter rows to. Must be in pinned memory.
|
|
rows : Tensor
|
|
The rows to scatter. Must be a CUDA tensor with unique entries.
|
|
source : Tensor
|
|
The tensor on the GPU to scatter rows from.
|
|
"""
|
|
_CAPI_DGLIndexScatterGPUToCPU(F.to_dgl_nd(dest), F.to_dgl_nd(rows),
|
|
F.to_dgl_nd(source))
|
|
|
|
|
|
_init_api("dgl.ndarray.uvm", __name__)
|