wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/utils/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""Internal utilities."""
from .checks import *
from .data import *
from .exception import *
from .filter import *
from .internal import *
from .pin_memory import *
from .shared_mem import *
try:
from packaging import version
except ImportError:
# If packaging isn't installed, try and use the vendored copy in setuptools
from setuptools.extern.packaging import version
python/dgl/utils/checks.py
+295
View File
@@ -0,0 +1,295 @@
"""Checking and logging utilities."""
# pylint: disable=invalid-name
from __future__ import absolute_import, division
from collections.abc import Mapping
from .. import backend as F
from .._ffi.function import _init_api
from ..base import DGLError
def prepare_tensor(g, data, name):
"""Convert the data to ID tensor and check its ID type and context.
If the data is already in tensor type, raise error if its ID type
and context does not match the graph's.
Otherwise, convert it to tensor type of the graph's ID type and
ctx and return.
Parameters
----------
g : DGLGraph
Graph.
data : int, iterable of int, tensor
Data.
name : str
Name of the data.
Returns
-------
Tensor
Data in tensor object.
"""
if F.is_tensor(data):
if F.dtype(data) != g.idtype:
raise DGLError(
f'Expect argument "{name}" to have data type {g.idtype}. '
f"But got {F.dtype(data)}."
)
if F.context(data) != g.device and not g.is_pinned():
raise DGLError(
f'Expect argument "{name}" to have device {g.device}. '
f"But got {F.context(data)}."
)
ret = data
else:
data = F.tensor(data)
if not (
F.ndim(data) > 0 and F.shape(data)[0] == 0
) and F.dtype( # empty tensor
data
) not in (
F.int32,
F.int64,
):
raise DGLError(
'Expect argument "{}" to have data type int32 or int64,'
" but got {}.".format(name, F.dtype(data))
)
ret = F.copy_to(F.astype(data, g.idtype), g.device)
if F.ndim(ret) == 0:
ret = F.unsqueeze(ret, 0)
if F.ndim(ret) > 1:
raise DGLError(
'Expect a 1-D tensor for argument "{}". But got {}.'.format(
name, ret
)
)
return ret
def prepare_tensor_dict(g, data, name):
"""Convert a dictionary of data to a dictionary of ID tensors.
Calls ``prepare_tensor`` on each key-value pair.
Parameters
----------
g : DGLGraph
Graph.
data : dict[str, (int, iterable of int, tensor)]
Data dict.
name : str
Name of the data.
Returns
-------
dict[str, tensor]
"""
return {
key: prepare_tensor(g, val, '{}["{}"]'.format(name, key))
for key, val in data.items()
}
def prepare_tensor_or_dict(g, data, name):
"""Convert data to either a tensor or a dictionary depending on input type.
Parameters
----------
g : DGLGraph
Graph.
data : dict[str, (int, iterable of int, tensor)]
Data dict.
name : str
Name of the data.
Returns
-------
tensor or dict[str, tensor]
"""
return (
prepare_tensor_dict(g, data, name)
if isinstance(data, Mapping)
else prepare_tensor(g, data, name)
)
def parse_edges_arg_to_eid(g, edges, etid, argname="edges"):
"""Parse the :attr:`edges` argument and return an edge ID tensor.
The resulting edge ID tensor has the same ID type and device of :attr:`g`.
Parameters
----------
g : DGLGraph
Graph
edges : pair of Tensor, Tensor, iterable[int]
Argument for specifying edges.
etid : int
Edge type ID.
argname : str, optional
Argument name.
Returns
-------
Tensor
Edge ID tensor
"""
if isinstance(edges, tuple):
u, v = edges
u = prepare_tensor(g, u, "{}[0]".format(argname))
v = prepare_tensor(g, v, "{}[1]".format(argname))
eid = g.edge_ids(u, v, etype=g.canonical_etypes[etid])
else:
eid = prepare_tensor(g, edges, argname)
return eid
def check_all_same_idtype(glist, name):
"""Check all the graphs have the same idtype."""
if len(glist) == 0:
return
idtype = glist[0].idtype
for i, g in enumerate(glist):
if g.idtype != idtype:
raise DGLError(
"Expect {}[{}] to have {} type ID, but got {}.".format(
name, i, idtype, g.idtype
)
)
def check_device(data, device):
"""Check if data is on the target device.
Parameters
----------
data : Tensor or dict[str, Tensor]
device: Backend device.
Returns
-------
Bool: True if the data is on the target device.
"""
if isinstance(data, dict):
for v in data.values():
if v.device != device:
return False
elif data.device != device:
return False
return True
def check_all_same_device(glist, name):
"""Check all the graphs have the same device."""
if len(glist) == 0:
return
device = glist[0].device
for i, g in enumerate(glist):
if g.device != device:
raise DGLError(
"Expect {}[{}] to be on device {}, but got {}.".format(
name, i, device, g.device
)
)
def check_all_same_schema(schemas, name):
"""Check the list of schemas are the same."""
if len(schemas) == 0:
return
for i, schema in enumerate(schemas):
if schema != schemas[0]:
raise DGLError(
"Expect all graphs to have the same schema on {}, "
"but graph {} got\n\t{}\nwhich is different from\n\t{}.".format(
name, i, schema, schemas[0]
)
)
def check_all_same_schema_for_keys(schemas, keys, name):
"""Check the list of schemas are the same on the given keys."""
if len(schemas) == 0:
return
head = None
keys = set(keys)
for i, schema in enumerate(schemas):
if not keys.issubset(schema.keys()):
raise DGLError(
"Expect all graphs to have keys {} on {}, "
"but graph {} got keys {}.".format(keys, name, i, schema.keys())
)
if head is None:
head = {k: schema[k] for k in keys}
else:
target = {k: schema[k] for k in keys}
if target != head:
raise DGLError(
"Expect all graphs to have the same schema for keys {} on {}, "
"but graph {} got \n\t{}\n which is different from\n\t{}.".format(
keys, name, i, target, head
)
)
def check_valid_idtype(idtype):
"""Check whether the value of the idtype argument is valid (int32/int64)
Parameters
----------
idtype : data type
The framework object of a data type.
"""
if idtype not in [None, F.int32, F.int64]:
raise DGLError(
"Expect idtype to be a framework object of int32/int64, "
"got {}".format(idtype)
)
def is_sorted_srcdst(src, dst, num_src=None, num_dst=None):
"""Checks whether an edge list is in ascending src-major order (e.g., first
sorted by ``src`` and then by ``dst``).
Parameters
----------
src : IdArray
The tensor of source nodes for each edge.
dst : IdArray
The tensor of destination nodes for each edge.
num_src : int, optional
The number of source nodes.
num_dst : int, optional
The number of destination nodes.
Returns
-------
bool, bool
Whether ``src`` is in ascending order, and whether ``dst`` is
in ascending order with respect to ``src``.
"""
# for some versions of MXNET and TensorFlow, num_src and num_dst get
# incorrectly marked as floats, so force them as integers here
if num_src is None:
num_src = int(F.as_scalar(F.max(src, dim=0) + 1))
if num_dst is None:
num_dst = int(F.as_scalar(F.max(dst, dim=0) + 1))
src = F.zerocopy_to_dgl_ndarray(src)
dst = F.zerocopy_to_dgl_ndarray(dst)
sorted_status = _CAPI_DGLCOOIsSorted(src, dst, num_src, num_dst)
row_sorted = sorted_status > 0
col_sorted = sorted_status > 1
return row_sorted, col_sorted
_init_api("dgl.utils.checks")
python/dgl/utils/data.py
+409
View File
@@ -0,0 +1,409 @@
"""Data utilities."""
from collections import namedtuple
import networkx as nx
import scipy as sp
from .. import backend as F
from ..base import DGLError
from . import checks
def elist2tensor(elist, idtype):
"""Function to convert an edge list to edge tensors.
Parameters
----------
elist : iterable of int pairs
List of (src, dst) node ID pairs.
idtype : int32, int64, optional
Integer ID type. Must be int32 or int64.
Returns
-------
(Tensor, Tensor)
Edge tensors.
"""
if len(elist) == 0:
u, v = [], []
else:
u, v = zip(*elist)
u = list(u)
v = list(v)
return F.tensor(u, idtype), F.tensor(v, idtype)
def scipy2tensor(spmat, idtype):
"""Function to convert a scipy matrix to a sparse adjacency matrix tuple.
Note that the data array of the scipy matrix is discarded.
Parameters
----------
spmat : scipy.sparse.spmatrix
SciPy sparse matrix.
idtype : int32, int64, optional
Integer ID type. Must be int32 or int64.
Returns
-------
(str, tuple[Tensor])
A tuple containing the format as well as the list of tensors representing
the sparse matrix.
"""
if spmat.format in ["csr", "csc"]:
indptr = F.tensor(spmat.indptr, idtype)
indices = F.tensor(spmat.indices, idtype)
data = F.tensor([], idtype)
return SparseAdjTuple(spmat.format, (indptr, indices, data))
else:
spmat = spmat.tocoo()
row = F.tensor(spmat.row, idtype)
col = F.tensor(spmat.col, idtype)
return SparseAdjTuple("coo", (row, col))
def networkx2tensor(nx_graph, idtype, edge_id_attr_name=None):
"""Function to convert a networkx graph to edge tensors.
Parameters
----------
nx_graph : nx.Graph
NetworkX graph.
idtype : int32, int64, optional
Integer ID type. Must be int32 or int64.
edge_id_attr_name : str, optional
Key name for edge ids in the NetworkX graph. If not found, we
will consider the graph not to have pre-specified edge ids. (Default: None)
Returns
-------
(Tensor, Tensor)
Edge tensors.
"""
if not nx_graph.is_directed():
nx_graph = nx_graph.to_directed()
# Relabel nodes using consecutive integers
nx_graph = nx.convert_node_labels_to_integers(nx_graph, ordering="sorted")
has_edge_id = edge_id_attr_name is not None
if has_edge_id:
num_edges = nx_graph.number_of_edges()
src = [0] * num_edges
dst = [0] * num_edges
for u, v, attr in nx_graph.edges(data=True):
eid = int(attr[edge_id_attr_name])
if eid < 0 or eid >= nx_graph.number_of_edges():
raise DGLError(
"Expect edge IDs to be a non-negative integer smaller than {:d}, "
"got {:d}".format(num_edges, eid)
)
src[eid] = u
dst[eid] = v
else:
src = []
dst = []
for e in nx_graph.edges:
src.append(e[0])
dst.append(e[1])
src = F.tensor(src, idtype)
dst = F.tensor(dst, idtype)
return src, dst
SparseAdjTuple = namedtuple("SparseAdjTuple", ["format", "arrays"])
def graphdata2tensors(
data, idtype=None, bipartite=False, infer_node_count=True, **kwargs
):
"""Function to convert various types of data to edge tensors and infer
the number of nodes.
Parameters
----------
data : graph data
Various kinds of graph data. Possible data types are:
- ``(row, col)``
- ``('coo', (row, col))``
- ``('csr', (indptr, indices, edge_ids))``
- ``('csc', (indptr, indices, edge_ids))``
- SciPy sparse matrix
- NetworkX graph
idtype : int32, int64, optional
Integer ID type. If None, try infer from the data and if fail use
int64.
bipartite : bool, optional
Whether infer number of nodes of a bipartite graph --
num_src and num_dst can be different.
infer_node_count : bool, optional
Whether infer number of nodes at all. If False, num_src and num_dst
are returned as None.
kwargs
- edge_id_attr_name : The name (str) of the edge attribute that stores the edge
IDs in the NetworkX graph.
- top_map : The dictionary mapping the original IDs of the source nodes to the
new ones.
- bottom_map : The dictionary mapping the original IDs of the destination nodes
to the new ones.
Returns
-------
data : SparseAdjTuple
A tuple with the sparse matrix format and the adjacency matrix tensors.
num_src : int
Number of source nodes.
num_dst : int
Number of destination nodes.
"""
# Convert tuple to SparseAdjTuple
if isinstance(data, tuple):
if not isinstance(data[0], str):
# (row, col) format, convert to ('coo', (row, col))
data = ("coo", data)
data = SparseAdjTuple(*data)
if idtype is None and not (
isinstance(data, SparseAdjTuple) and F.is_tensor(data.arrays[0])
):
# preferred default idtype is int64
# if data is tensor and idtype is None, infer the idtype from tensor
idtype = F.int64
checks.check_valid_idtype(idtype)
if isinstance(data, SparseAdjTuple) and (
not all(F.is_tensor(a) for a in data.arrays)
):
# (Iterable, Iterable) type data, convert it to (Tensor, Tensor)
if len(data.arrays[0]) == 0:
# force idtype for empty list
data = SparseAdjTuple(
data.format, tuple(F.tensor(a, idtype) for a in data.arrays)
)
else:
# convert the iterable to tensor and keep its native data type so we can check
# its validity later
data = SparseAdjTuple(
data.format, tuple(F.tensor(a) for a in data.arrays)
)
num_src, num_dst = None, None
if isinstance(data, SparseAdjTuple):
if idtype is not None:
data = SparseAdjTuple(
data.format, tuple(F.astype(a, idtype) for a in data.arrays)
)
if infer_node_count:
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
elif isinstance(data, list):
src, dst = elist2tensor(data, idtype)
data = SparseAdjTuple("coo", (src, dst))
if infer_node_count:
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
elif isinstance(data, sp.sparse.spmatrix):
# We can get scipy matrix's number of rows and columns easily.
if infer_node_count:
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
data = scipy2tensor(data, idtype)
elif isinstance(data, nx.Graph):
# We can get networkx graph's number of sources and destinations easily.
if infer_node_count:
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
edge_id_attr_name = kwargs.get("edge_id_attr_name", None)
if bipartite:
top_map = kwargs.get("top_map")
bottom_map = kwargs.get("bottom_map")
src, dst = networkxbipartite2tensors(
data,
idtype,
top_map=top_map,
bottom_map=bottom_map,
edge_id_attr_name=edge_id_attr_name,
)
else:
src, dst = networkx2tensor(
data, idtype, edge_id_attr_name=edge_id_attr_name
)
data = SparseAdjTuple("coo", (src, dst))
else:
raise DGLError("Unsupported graph data type:", type(data))
return data, num_src, num_dst
def networkxbipartite2tensors(
nx_graph, idtype, top_map, bottom_map, edge_id_attr_name=None
):
"""Function to convert a networkx bipartite to edge tensors.
Parameters
----------
nx_graph : nx.Graph
NetworkX graph. It must follow the bipartite graph convention of networkx.
Each node has an attribute ``bipartite`` with values 0 and 1 indicating
which set it belongs to.
top_map : dict
The dictionary mapping the original node labels to the node IDs for the source type.
bottom_map : dict
The dictionary mapping the original node labels to the node IDs for the destination type.
idtype : int32, int64, optional
Integer ID type. Must be int32 or int64.
edge_id_attr_name : str, optional
Key name for edge ids in the NetworkX graph. If not found, we
will consider the graph not to have pre-specified edge ids. (Default: None)
Returns
-------
(Tensor, Tensor)
Edge tensors.
"""
has_edge_id = edge_id_attr_name is not None
if has_edge_id:
num_edges = nx_graph.number_of_edges()
src = [0] * num_edges
dst = [0] * num_edges
for u, v, attr in nx_graph.edges(data=True):
if u not in top_map:
raise DGLError(
"Expect the node {} to have attribute bipartite=0 "
"with edge {}".format(u, (u, v))
)
if v not in bottom_map:
raise DGLError(
"Expect the node {} to have attribute bipartite=1 "
"with edge {}".format(v, (u, v))
)
eid = int(attr[edge_id_attr_name])
if eid < 0 or eid >= nx_graph.number_of_edges():
raise DGLError(
"Expect edge IDs to be a non-negative integer smaller than {:d}, "
"got {:d}".format(num_edges, eid)
)
src[eid] = top_map[u]
dst[eid] = bottom_map[v]
else:
src = []
dst = []
for e in nx_graph.edges:
u, v = e[0], e[1]
if u not in top_map:
raise DGLError(
"Expect the node {} to have attribute bipartite=0 "
"with edge {}".format(u, (u, v))
)
if v not in bottom_map:
raise DGLError(
"Expect the node {} to have attribute bipartite=1 "
"with edge {}".format(v, (u, v))
)
src.append(top_map[u])
dst.append(bottom_map[v])
src = F.tensor(src, dtype=idtype)
dst = F.tensor(dst, dtype=idtype)
return src, dst
def infer_num_nodes(data, bipartite=False):
"""Function for inferring the number of nodes.
Parameters
----------
data : graph data
Supported types are:
* SparseTuple ``(sparse_fmt, arrays)`` where ``arrays`` can be either ``(src, dst)`` or
``(indptr, indices, data)``.
* SciPy matrix.
* NetworkX graph.
bipartite : bool, optional
Whether infer number of nodes of a bipartite graph --
num_src and num_dst can be different.
Returns
-------
num_src : int
Number of source nodes.
num_dst : int
Number of destination nodes.
or
None
If the inference failed.
"""
if isinstance(data, tuple) and len(data) == 2:
if not isinstance(data[0], str):
raise TypeError(
"Expected sparse format as a str, but got %s" % type(data[0])
)
if data[0] == "coo":
# ('coo', (src, dst)) format
u, v = data[1]
nsrc = F.as_scalar(F.max(u, dim=0)) + 1 if len(u) > 0 else 0
ndst = F.as_scalar(F.max(v, dim=0)) + 1 if len(v) > 0 else 0
elif data[0] == "csr":
# ('csr', (indptr, indices, eids)) format
indptr, indices, _ = data[1]
nsrc = F.shape(indptr)[0] - 1
ndst = (
F.as_scalar(F.max(indices, dim=0)) + 1
if len(indices) > 0
else 0
)
elif data[0] == "csc":
# ('csc', (indptr, indices, eids)) format
indptr, indices, _ = data[1]
ndst = F.shape(indptr)[0] - 1
nsrc = (
F.as_scalar(F.max(indices, dim=0)) + 1
if len(indices) > 0
else 0
)
else:
raise ValueError("unknown format %s" % data[0])
elif isinstance(data, sp.sparse.spmatrix):
nsrc, ndst = data.shape[0], data.shape[1]
elif isinstance(data, nx.Graph):
if data.number_of_nodes() == 0:
nsrc = ndst = 0
elif not bipartite:
nsrc = ndst = data.number_of_nodes()
else:
nsrc = len(
{n for n, d in data.nodes(data=True) if d["bipartite"] == 0}
)
ndst = data.number_of_nodes() - nsrc
else:
return None
if not bipartite:
nsrc = ndst = max(nsrc, ndst)
return nsrc, ndst
def to_device(data, device):
"""Transfer the tensor or dictionary of tensors to the given device.
Nothing will happen if the device of the original tensor is the same as target device.
Parameters
----------
data : Tensor or dict[str, Tensor]
The data.
device : device
The target device.
Returns
-------
Tensor or dict[str, Tensor]
The output data.
"""
if isinstance(data, dict):
return {k: F.copy_to(v, device) for k, v in data.items()}
else:
return F.copy_to(data, device)
python/dgl/utils/exception.py
+61
View File
@@ -0,0 +1,61 @@
"""Exception wrapper classes to properly display exceptions under multithreading or
multiprocessing.
"""
import sys
import traceback
# The following code is borrowed from PyTorch. Basically when a subprocess or thread
# throws an exception, you will need to wrap the exception with ExceptionWrapper class
# and put it in the queue you are normally retrieving from.
# NOTE [ Python Traceback Reference Cycle Problem ]
#
# When using sys.exc_info(), it is important to **not** store the exc_info[2],
# which is the traceback, because otherwise you will run into the traceback
# reference cycle problem, i.e., the traceback holding reference to the frame,
# and the frame (which holds reference to all the object in its temporary scope)
# holding reference the traceback.
class KeyErrorMessage(str):
r"""str subclass that returns itself in repr"""
def __repr__(self): # pylint: disable=invalid-repr-returned
return self
class ExceptionWrapper(object):
r"""Wraps an exception plus traceback to communicate across threads"""
def __init__(self, exc_info=None, where="in background"):
# It is important that we don't store exc_info, see
# NOTE [ Python Traceback Reference Cycle Problem ]
if exc_info is None:
exc_info = sys.exc_info()
self.exc_type = exc_info[0]
self.exc_msg = "".join(traceback.format_exception(*exc_info))
self.where = where
def reraise(self):
r"""Reraises the wrapped exception in the current thread"""
# Format a message such as: "Caught ValueError in DataLoader worker
# process 2. Original Traceback:", followed by the traceback.
msg = "Caught {} {}.\nOriginal {}".format(
self.exc_type.__name__, self.where, self.exc_msg
)
if self.exc_type == KeyError:
# KeyError calls repr() on its argument (usually a dict key). This
# makes stack traces unreadable. It will not be changed in Python
# (https://bugs.python.org/issue2651), so we work around it.
msg = KeyErrorMessage(msg)
elif getattr(self.exc_type, "message", None):
# Some exceptions have first argument as non-str but explicitly
# have message field
raise self.exc_type(message=msg)
try:
exception = self.exc_type(msg)
except TypeError:
# If the exception takes multiple arguments, don't try to
# instantiate since we don't know how to
raise RuntimeError(msg) from None
raise exception
python/dgl/utils/filter.py
+75
View File
@@ -0,0 +1,75 @@
"""Utilities for finding overlap or missing items in arrays."""
from .. import backend as F
from .._ffi.function import _init_api
class Filter(object):
"""Class used to either find the subset of IDs that are in this
filter, or the subset of IDs that are not in this filter
given a second set of IDs.
Examples
--------
>>> import torch as th
>>> from dgl.utils import Filter
>>> f = Filter(th.tensor([3,2,9], device=th.device('cuda')))
>>> f.find_included_indices(th.tensor([0,2,8,9], device=th.device('cuda')))
tensor([1,3])
>>> f.find_excluded_indices(th.tensor([0,2,8,9], device=th.device('cuda')))
tensor([0,2], device='cuda')
"""
def __init__(self, ids):
"""Create a new filter from a given set of IDs. This currently is only
implemented for the GPU.
Parameters
----------
ids : IdArray
The unique set of IDs to keep in the filter.
"""
self._filter = _CAPI_DGLFilterCreateFromSet(
F.zerocopy_to_dgl_ndarray(ids)
)
def find_included_indices(self, test):
"""Find the index of the IDs in `test` that are in this filter.
Parameters
----------
test : IdArray
The set of IDs to to test with.
Returns
-------
IdArray
The index of IDs in `test` that are also in this filter.
"""
return F.zerocopy_from_dgl_ndarray(
_CAPI_DGLFilterFindIncludedIndices(
self._filter, F.zerocopy_to_dgl_ndarray(test)
)
)
def find_excluded_indices(self, test):
"""Find the index of the IDs in `test` that are not in this filter.
Parameters
----------
test : IdArray
The set of IDs to to test with.
Returns
-------
IdArray
The index of IDs in `test` that are not in this filter.
"""
return F.zerocopy_from_dgl_ndarray(
_CAPI_DGLFilterFindExcludedIndices(
self._filter, F.zerocopy_to_dgl_ndarray(test)
)
)
_init_api("dgl.utils.filter")
python/dgl/utils/internal.py
+1178
View File
File diff suppressed because it is too large Load Diff
python/dgl/utils/pin_memory.py
+77
View File
@@ -0,0 +1,77 @@
"""Utility functions related to pinned memory tensors."""
from .. import backend as F
from .._ffi.function import _init_api
from ..base import DGLError
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__)
python/dgl/utils/shared_mem.py
+9
View File
@@ -0,0 +1,9 @@
"""Shared memory utilities.
For compatibility with older code that uses ``dgl.utils.shared_mem`` namespace; the
content has been moved to ``dgl.ndarray`` module.
"""
from ..ndarray import ( # pylint: disable=unused-import
create_shared_mem_array,
get_shared_mem_array,
)