chore: import upstream snapshot with attribution

python/dgl/distributed/shared_mem_utils.py

@@ -0,0 +1,26 @@
+"""Define utility functions for shared memory."""
+
+from .. import backend as F, ndarray as nd
+from .._ffi.ndarray import empty_shared_mem
+
+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)