dmlc--dgl
Da Zheng
50962c433f
* refactor graph store for new kvstore. * fix kvstore. * fix. * fix lint complains. * add docstring. * small fix in graph partition book. * fix. * remove tests * disable sampling. * Revert "disable sampling." This reverts commit 70451008f61ff1481d5dadbf8bd199470aee559d. * Revert "remove tests" This reverts commit 1394364243bdd73b669abde6193a74e2cda5521d.
23 行
736 B
Python
23 行
736 B
Python
"""Define utility functions for shared memory."""
|
|
|
|
from .. import backend as F
|
|
from .._ffi.ndarray import empty_shared_mem
|
|
from .. import ndarray as nd
|
|
|
|
DTYPE_DICT = F.data_type_dict
|
|
DTYPE_DICT = {DTYPE_DICT[key]:key for key in DTYPE_DICT}
|
|
|
|
def _get_ndata_path(graph_name, ndata_name):
|
|
return "/" + graph_name + "_node_" + ndata_name
|
|
|
|
def _get_edata_path(graph_name, edata_name):
|
|
return "/" + graph_name + "_edge_" + edata_name
|
|
|
|
def _to_shared_mem(arr, name):
|
|
dlpack = F.zerocopy_to_dlpack(arr)
|
|
dgl_tensor = nd.from_dlpack(dlpack)
|
|
new_arr = empty_shared_mem(name, True, F.shape(arr), DTYPE_DICT[F.dtype(arr)])
|
|
dgl_tensor.copyto(new_arr)
|
|
dlpack = new_arr.to_dlpack()
|
|
return F.zerocopy_from_dlpack(dlpack)
|