Merge branch 'master' into dist_part

tensoradapter/include/tensoradapter.h

@@ -18,7 +18,7 @@ namespace tensoradapter {
 extern "C" {
 
 /*!
- * \brief Allocate an empty tensor
+ * \brief Allocate an empty tensor.
  *
  * \param shape The shape
  * \param dtype The data type
@@ -28,6 +28,24 @@ extern "C" {
 DLManagedTensor* TAempty(
     std::vector<int64_t> shape, DLDataType dtype, DLContext ctx);
 
+#ifdef DGL_USE_CUDA
+/*!
+ * \brief Allocate a piece of GPU memory via
+ * PyTorch's THCCachingAllocator.
+ *
+ * \param nbytes The size to be allocated.
+ * \return Pointer to the allocated memory.
+ */
+void* RawAlloc(size_t nbytes);
+
+/*!
+ * \brief Free the GPU memory.
+ *
+ * \param ptr Pointer to the memory to be freed.
+ */
+void RawDelete(void* ptr);
+#endif  // DGL_USE_CUDA
+
 }
 
 };  // namespace tensoradapter

tensoradapter/pytorch/CMakeLists.txt

@@ -17,6 +17,10 @@ list(GET TORCH_PREFIX_VER 0 TORCH_PREFIX)
 list(GET TORCH_PREFIX_VER 1 TORCH_VER)
 message(STATUS "Configuring for PyTorch ${TORCH_VER}")
 
+if(USE_CUDA)
+  add_definitions(-DDGL_USE_CUDA)
+endif()
+
 set(Torch_DIR "${TORCH_PREFIX}/Torch")
 message(STATUS "Setting directory to ${Torch_DIR}")
 find_package(Torch REQUIRED)

tensoradapter/pytorch/build.bat

@@ -11,7 +11,7 @@ IF x%1x == xx GOTO single
 
 FOR %%X IN (%*) DO (
 	DEL /S /Q *
-	"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
+	"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
 	msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
 	COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
 )
@@ -21,7 +21,7 @@ GOTO end
 :single
 
 DEL /S /Q *
-"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% .. -G "Visual Studio 16 2019" || EXIT /B 1
+"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% .. -G "Visual Studio 16 2019" || EXIT /B 1
 msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
 COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
 

tensoradapter/pytorch/build.sh

@@ -13,7 +13,7 @@ else
 	CPSOURCE=*.so
 fi
 
-CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
+CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST -DUSE_CUDA=$USE_CUDA"
 
 if [ $# -eq 0 ]; then
 	$CMAKE_COMMAND $CMAKE_FLAGS ..

tensoradapter/pytorch/torch.cpp

@@ -7,6 +7,9 @@
 #include <tensoradapter_exports.h>
 #include <torch/torch.h>
 #include <ATen/DLConvertor.h>
+#ifdef DGL_USE_CUDA
+#include <c10/cuda/CUDACachingAllocator.h>
+#endif  // DGL_USE_CUDA
 #include <vector>
 #include <iostream>
 
@@ -47,6 +50,16 @@ TA_EXPORTS DLManagedTensor* TAempty(
   return at::toDLPack(tensor);
 }
 
+#ifdef DGL_USE_CUDA
+TA_EXPORTS void* RawAlloc(size_t nbytes) {
+  return c10::cuda::CUDACachingAllocator::raw_alloc(nbytes);
+}
+
+TA_EXPORTS void RawDelete(void* ptr) {
+  c10::cuda::CUDACachingAllocator::raw_delete(ptr);
+}
+#endif  // DGL_USE_CUDA
+
 };
 
 };  // namespace tensoradapter

tests/compute/test_csrmm.py

@@ -5,6 +5,10 @@ import dgl
 from test_utils import parametrize_idtype
 import backend as F
 
+if F.backend_name == 'pytorch':
+    import torch
+    torch.backends.cuda.matmul.allow_tf32 = False
+
 def _random_simple_graph(idtype, dtype, ctx, M, N, max_nnz, srctype, dsttype, etype):
     src = np.random.randint(0, M, (max_nnz,))
     dst = np.random.randint(0, N, (max_nnz,))

tests/compute/test_edge_softmax_hetero.py

@@ -2,6 +2,7 @@ import dgl
 from dgl.ops import edge_softmax
 import dgl.function as fn
 from collections import Counter
+import math
 import numpy as np
 import scipy.sparse as ssp
 import itertools
@@ -17,8 +18,6 @@ rfuncs = {'sum': fn.sum, 'max': fn.max, 'min': fn.min, 'mean': fn.mean}
 fill_value = {'sum': 0, 'max': float("-inf")}
 feat_size = 2
 
-@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
-
 def create_test_heterograph(idtype):
     # test heterograph from the docstring, plus a user -- wishes -- game relation
     # 3 users, 2 games, 2 developers
@@ -37,8 +36,26 @@ def create_test_heterograph(idtype):
     assert g.idtype == idtype
     assert g.device == F.ctx()
     return g
+    
+@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
+def test_edge_softmax_unidirectional():
+    g = dgl.heterograph({
+        ('A', 'AB', 'B'): ([1,2,3,1,2,3,1,2,3],[0,0,0,1,1,1,2,2,2]),
+        ('B', 'BB', 'B'): ([0,1,2,0,1,2,0,1,2], [0,0,0,1,1,1,2,2,2])})
+    g = g.to(F.ctx())
+    g.edges['AB'].data['x'] = F.ones(9) * 2
+    g.edges['BB'].data['x'] = F.ones(9)
+    result = dgl.ops.edge_softmax(g, {'AB': g.edges['AB'].data['x'], 'BB': g.edges['BB'].data['x']})
 
+    ab = result['A', 'AB', 'B']
+    bb = result['B', 'BB', 'B']
+    e2 = F.zeros_like(ab) + math.exp(2) / ((math.exp(2) + math.exp(1)) * 3)
+    e1 = F.zeros_like(bb) + math.exp(1) / ((math.exp(2) + math.exp(1)) * 3)
+    assert F.allclose(ab, e2)
+    assert F.allclose(bb, e1)
 
+
+@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
 @pytest.mark.parametrize('g', get_cases(['clique']))
 @pytest.mark.parametrize('norm_by', ['src', 'dst'])
 # @pytest.mark.parametrize('shp', edge_softmax_shapes)
@@ -109,5 +126,4 @@ def test_edge_softmax(g, norm_by, idtype):
         assert F.allclose(grad_edata_hm, grad_edata_ht)
 
 if __name__ == '__main__':
-    test_edge_softmax()
-
+    test_edge_softmax_unidirectional()

tests/compute/test_sampling.py

@@ -24,59 +24,74 @@ def check_random_walk(g, metapath, traces, ntypes, prob=None, trace_eids=None):
                 u, v = g.find_edges(trace_eids[i, j], etype=metapath[j])
                 assert (u == traces[i, j]) and (v == traces[i, j + 1])
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="Random walk with non-uniform prob is not supported in GPU.")
-def test_non_uniform_random_walk():
+@pytest.mark.parametrize('use_uva', [True, False])
+def test_non_uniform_random_walk(use_uva):
+    if use_uva:
+        if F.ctx() == F.cpu():
+            pytest.skip('UVA biased random walk requires a GPU.')
+        if dgl.backend.backend_name != 'pytorch':
+            pytest.skip('UVA biased random walk is only supported with PyTorch.')
     g2 = dgl.heterograph({
             ('user', 'follow', 'user'): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0])
-        }).to(F.ctx())
+        })
     g4 = dgl.heterograph({
             ('user', 'follow', 'user'): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]),
             ('user', 'view', 'item'): ([0, 0, 1, 2, 3, 3], [0, 1, 1, 2, 2, 1]),
             ('item', 'viewed-by', 'user'): ([0, 1, 1, 2, 2, 1], [0, 0, 1, 2, 3, 3])
-        }).to(F.ctx())
+        })
 
-    g2.edata['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
-    g2.edata['p2'] = F.tensor([[3], [0], [3], [3], [3]], dtype=F.float32)
-    g4.edges['follow'].data['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
-    g4.edges['viewed-by'].data['p'] = F.tensor([1, 1, 1, 1, 1, 1], dtype=F.float32)
+    g2.edata['p'] = F.copy_to(F.tensor([3, 0, 3, 3, 3], dtype=F.float32), F.cpu())
+    g2.edata['p2'] = F.copy_to(F.tensor([[3], [0], [3], [3], [3]], dtype=F.float32), F.cpu())
+    g4.edges['follow'].data['p'] = F.copy_to(F.tensor([3, 0, 3, 3, 3], dtype=F.float32), F.cpu())
+    g4.edges['viewed-by'].data['p'] = F.copy_to(F.tensor([1, 1, 1, 1, 1, 1], dtype=F.float32), F.cpu())
 
-    traces, eids, ntypes = dgl.sampling.random_walk(
-        g2, [0, 1, 2, 3, 0, 1, 2, 3], length=4, prob='p', return_eids=True)
-    check_random_walk(g2, ['follow'] * 4, traces, ntypes, 'p', trace_eids=eids)
+    if use_uva:
+        for g in (g2, g4):
+            g.create_formats_()
+            g.pin_memory_()
+    elif F._default_context_str == 'gpu':
+        g2 = g2.to(F.ctx())
+        g4 = g4.to(F.ctx())
 
     try:
-        traces, ntypes = dgl.sampling.random_walk(
-            g2, [0, 1, 2, 3, 0, 1, 2, 3], length=4, prob='p2')
-        fail = False
-    except dgl.DGLError:
-        fail = True
-    assert fail
+        traces, eids, ntypes = dgl.sampling.random_walk(
+            g2, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g2.idtype),
+            length=4, prob='p', return_eids=True)
+        check_random_walk(g2, ['follow'] * 4, traces, ntypes, 'p', trace_eids=eids)
 
-    metapath = ['follow', 'view', 'viewed-by'] * 2
-    traces, eids, ntypes = dgl.sampling.random_walk(
-        g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath, prob='p', return_eids=True)
-    check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
-    traces, eids, ntypes = dgl.sampling.random_walk(
-        g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath, prob='p', restart_prob=0., return_eids=True)
-    check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
-    traces, eids, ntypes = dgl.sampling.random_walk(
-        g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath, prob='p',
-        restart_prob=F.zeros((6,), F.float32, F.cpu()), return_eids=True)
-    check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
-    traces, eids, ntypes = dgl.sampling.random_walk(
-        g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath + ['follow'], prob='p',
-        restart_prob=F.tensor([0, 0, 0, 0, 0, 0, 1], F.float32), return_eids=True)
-    check_random_walk(g4, metapath, traces[:, :7], ntypes[:7], 'p', trace_eids=eids)
-    assert (F.asnumpy(traces[:, 7]) == -1).all()
-
-def _use_uva():
-    if F._default_context_str == 'cpu':
-        return [False]
-    else:
-        return [True, False]
+        with pytest.raises(dgl.DGLError):
+            traces, ntypes = dgl.sampling.random_walk(
+                g2, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g2.idtype),
+                length=4, prob='p2')
 
-@pytest.mark.parametrize('use_uva', _use_uva())
+        metapath = ['follow', 'view', 'viewed-by'] * 2
+        traces, eids, ntypes = dgl.sampling.random_walk(
+            g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
+            metapath=metapath, prob='p', return_eids=True)
+        check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
+        traces, eids, ntypes = dgl.sampling.random_walk(
+            g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
+            metapath=metapath, prob='p', restart_prob=0., return_eids=True)
+        check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
+        traces, eids, ntypes = dgl.sampling.random_walk(
+            g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
+            metapath=metapath, prob='p',
+            restart_prob=F.zeros((6,), F.float32, F.ctx()), return_eids=True)
+        check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
+        traces, eids, ntypes = dgl.sampling.random_walk(
+            g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
+            metapath=metapath + ['follow'], prob='p',
+            restart_prob=F.tensor([0, 0, 0, 0, 0, 0, 1], F.float32), return_eids=True)
+        check_random_walk(g4, metapath, traces[:, :7], ntypes[:7], 'p', trace_eids=eids)
+        assert (F.asnumpy(traces[:, 7]) == -1).all()
+    finally:
+        for g in (g2, g4):
+            g.unpin_memory_()
+
+@pytest.mark.parametrize('use_uva', [True, False])
 def test_uniform_random_walk(use_uva):
+    if use_uva and F.ctx() == F.cpu():
+        pytest.skip('UVA random walk requires a GPU.')
     g1 = dgl.heterograph({
             ('user', 'follow', 'user'): ([0, 1, 2], [1, 2, 0])
         })
@@ -179,8 +194,10 @@ def test_pack_traces():
     assert F.array_equal(result[2], F.tensor([2, 7], dtype=F.int64))
     assert F.array_equal(result[3], F.tensor([0, 2], dtype=F.int64))
 
-@pytest.mark.parametrize('use_uva', _use_uva())
+@pytest.mark.parametrize('use_uva', [True, False])
 def test_pinsage_sampling(use_uva):
+    if use_uva and F.ctx() == F.cpu():
+        pytest.skip('UVA sampling requires a GPU.')
     def _test_sampler(g, sampler, ntype):
         seeds = F.copy_to(F.tensor([0, 2], dtype=g.idtype), F.ctx())
         neighbor_g = sampler(seeds)
@@ -626,12 +643,10 @@ def test_sample_neighbors_noprob():
     _test_sample_neighbors(False, None)
     #_test_sample_neighbors(True)
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors with probability is not implemented")
 def test_sample_neighbors_prob():
     _test_sample_neighbors(False, 'prob')
     #_test_sample_neighbors(True)
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
 def test_sample_neighbors_outedge():
     _test_sample_neighbors_outedge(False)
     #_test_sample_neighbors_outedge(True)
@@ -646,9 +661,8 @@ def test_sample_neighbors_topk_outedge():
     _test_sample_neighbors_topk_outedge(False)
     #_test_sample_neighbors_topk_outedge(True)
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
 def test_sample_neighbors_with_0deg():
-    g = dgl.graph(([], []), num_nodes=5)
+    g = dgl.graph(([], []), num_nodes=5).to(F.ctx())
     sg = dgl.sampling.sample_neighbors(g, F.tensor([1, 2], dtype=F.int64), 2, edge_dir='in', replace=False)
     assert sg.number_of_edges() == 0
     sg = dgl.sampling.sample_neighbors(g, F.tensor([1, 2], dtype=F.int64), 2, edge_dir='in', replace=True)
@@ -885,7 +899,6 @@ def test_sample_neighbors_etype_sorted_homogeneous(format_, direction):
     assert fail
 
 @pytest.mark.parametrize('dtype', ['int32', 'int64'])
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
 def test_sample_neighbors_exclude_edges_heteroG(dtype):
     d_i_d_u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300, size=100, dtype=dtype)))
     d_i_d_v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=d_i_d_u_nodes.shape, dtype=dtype))
@@ -898,7 +911,7 @@ def test_sample_neighbors_exclude_edges_heteroG(dtype):
         ('drug', 'interacts', 'drug'): (d_i_d_u_nodes, d_i_d_v_nodes),
         ('drug', 'interacts', 'gene'): (d_i_g_u_nodes, d_i_g_v_nodes),
         ('drug', 'treats', 'disease'): (d_t_d_u_nodes, d_t_d_v_nodes)
-    })
+    }).to(F.ctx())
 
     (U, V, EID) = (0, 1, 2)
 
@@ -951,11 +964,10 @@ def test_sample_neighbors_exclude_edges_heteroG(dtype):
                                                      etype=('drug','treats','disease'))))
 
 @pytest.mark.parametrize('dtype', ['int32', 'int64'])
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
 def test_sample_neighbors_exclude_edges_homoG(dtype):
     u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300,size=100, dtype=dtype)))
     v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=u_nodes.shape, dtype=dtype))
-    g = dgl.graph((u_nodes, v_nodes))
+    g = dgl.graph((u_nodes, v_nodes)).to(F.ctx())
 
     (U, V, EID) = (0, 1, 2)
 

tests/compute/test_sparse.py

@@ -293,6 +293,8 @@ def test_segment_reduce(reducer):
 @pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256])
 @pytest.mark.parametrize('dtype,tol', [(torch.float16,1e-2),(torch.float32,3e-3),(torch.float64,1e-4)])
 def test_segment_mm(idtype, feat_size, dtype, tol):
+    if F._default_context_str == 'cpu' and dtype == torch.float16:
+        pytest.skip("fp16 support for CPU linalg functions has been removed in PyTorch.")
     dev = F.ctx()
     # input
     a = torch.tensor(np.random.rand(100, feat_size)).to(dev).to(dtype)

tests/compute/test_transform.py

@@ -1625,8 +1625,11 @@ def test_remove_nodes(idtype):
 @parametrize_idtype
 def test_add_selfloop(idtype):
     # homogeneous graph
+
+    # test for fill_data is float
     g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
     g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
+    g.edata['he1'] = F.copy_to(F.tensor([[0., 1.], [2., 3.], [4., 5.]]), ctx=F.ctx())
     g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
     g = dgl.add_self_loop(g)
     assert g.number_of_nodes() == 3
@@ -1634,7 +1637,39 @@ def test_add_selfloop(idtype):
     u, v = g.edges(form='uv', order='eid')
     assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
     assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
-    assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 0, 0, 0], dtype=idtype))
+    assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 1, 1, 1], dtype=idtype))
+    assert F.array_equal(g.edata['he1'], F.tensor([[0., 1.], [2., 3.], [4., 5.],
+                                                   [1., 1.], [1., 1.], [1., 1.]]))
+
+    # test for fill_data is int
+    g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
+    g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
+    g.edata['he1'] = F.copy_to(F.tensor([[0, 1], [2, 3], [4, 5]], dtype=idtype), ctx=F.ctx())
+    g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
+    g = dgl.add_self_loop(g, fill_data=1)
+    assert g.number_of_nodes() == 3
+    assert g.number_of_edges() == 6
+    u, v = g.edges(form='uv', order='eid')
+    assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 1, 1, 1], dtype=idtype))
+    assert F.array_equal(g.edata['he1'], F.tensor([[0, 1], [2, 3], [4, 5],
+                                                   [1, 1], [1, 1], [1, 1]], dtype=idtype))
+
+    # test for fill_data is str
+    g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
+    g.edata['he'] = F.copy_to(F.tensor([1., 2., 3.]), ctx=F.ctx())
+    g.edata['he1'] = F.copy_to(F.tensor([[0., 1.], [2., 3.], [4., 5.]]), ctx=F.ctx())
+    g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
+    g = dgl.add_self_loop(g, fill_data='sum')
+    assert g.number_of_nodes() == 3
+    assert g.number_of_edges() == 6
+    u, v = g.edges(form='uv', order='eid')
+    assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(g.edata['he'], F.tensor([1., 2., 3., 3., 2., 1.]))
+    assert F.array_equal(g.edata['he1'], F.tensor([[0., 1.], [2., 3.], [4., 5.],
+                                                   [4., 5.], [2., 3.], [0., 1.]]))
 
     # bipartite graph
     g = dgl.heterograph(
@@ -1647,7 +1682,9 @@ def test_add_selfloop(idtype):
         raise_error = True
     assert raise_error
 
+    # test for fill_data is float
     g = create_test_heterograph5(idtype)
+    g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0., 1.], [1., 2.]]), ctx=F.ctx())
     g = dgl.add_self_loop(g, etype='follows')
     assert g.number_of_nodes('user') == 3
     assert g.number_of_nodes('game') == 2
@@ -1656,9 +1693,52 @@ def test_add_selfloop(idtype):
     u, v = g.edges(form='uv', order='eid', etype='follows')
     assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
     assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
-    assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 0, 0, 0], dtype=idtype))
+    assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 1, 1, 1], dtype=idtype))
+    assert F.array_equal(g.edges['follows'].data['h1'], F.tensor([[0., 1.], [1., 2.], [1., 1.],
+                                                                  [1., 1.], [1., 1.]]))
+    assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 2], dtype=idtype))
+
+    # test for fill_data is int
+    g = create_test_heterograph5(idtype)
+    g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0, 1], [1, 2]], dtype=idtype), ctx=F.ctx())
+    g = dgl.add_self_loop(g, fill_data=1, etype='follows')
+    assert g.number_of_nodes('user') == 3
+    assert g.number_of_nodes('game') == 2
+    assert g.number_of_edges('follows') == 5
+    assert g.number_of_edges('plays') == 2
+    u, v = g.edges(form='uv', order='eid', etype='follows')
+    assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 1, 1, 1], dtype=idtype))
+    assert F.array_equal(g.edges['follows'].data['h1'], F.tensor([[0, 1], [1, 2], [1, 1],
+                                                                  [1, 1], [1, 1]], dtype=idtype))
     assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 2], dtype=idtype))
 
+    # test for fill_data is str
+    g = dgl.heterograph({
+        ('user', 'follows', 'user'): (F.tensor([1, 2], dtype=idtype),
+                                      F.tensor([0, 1], dtype=idtype)),
+        ('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
+                                    F.tensor([0, 1], dtype=idtype))},
+        idtype=idtype, device=F.ctx())
+    g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
+    g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
+    g.edges['follows'].data['h'] = F.copy_to(F.tensor([1., 2.]), ctx=F.ctx())
+    g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0., 1.], [1., 2.]]), ctx=F.ctx())
+    g.edges['plays'].data['h'] = F.copy_to(F.tensor([1., 2.]), ctx=F.ctx())
+    g = dgl.add_self_loop(g, fill_data='mean', etype='follows')
+    assert g.number_of_nodes('user') == 3
+    assert g.number_of_nodes('game') == 2
+    assert g.number_of_edges('follows') == 5
+    assert g.number_of_edges('plays') == 2
+    u, v = g.edges(form='uv', order='eid', etype='follows')
+    assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
+    assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1., 2., 1., 2., 0.]))
+    assert F.array_equal(g.edges['follows'].data['h1'], F.tensor([[0., 1.], [1., 2.], [0., 1.],
+                                                                  [1., 2.], [0., 0.]]))
+    assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1., 2.]))
+
     raise_error = False
     try:
         g = dgl.add_self_loop(g, etype='plays')

tests/cpp/socket_communicator_test.cc

@@ -154,6 +154,14 @@ void start_server(int id) {
   receiver.Finalize();
 }
 
+TEST(SocketCommunicatorTest, TCPSocketBind) {
+  dgl::network::TCPSocket socket;
+  testing::internal::CaptureStderr();
+  EXPECT_EQ(socket.Bind("127.0.0", 50001), false);
+  const std::string stderr = testing::internal::GetCapturedStderr();
+  EXPECT_NE(stderr.find("Invalid IP: 127.0.0"), std::string::npos);
+}
+
 #else
 
 #include <windows.h>

tests/cugraph/test_basics.py

+import backend as F
+import dgl
+import numpy as np
+from dgl import DGLGraph
+import unittest
+import pytest
+import cugraph
+
+def test_dummy():
+    cg = cugraph.Graph()
+    assert cg is not None
+
+def test_to_cugraph_conversion():
+    g = dgl.graph((F.tensor([0, 1, 2, 3]), F.tensor([1, 0, 3, 2]))).to('cuda')
+    cugraph_g = g.to_cugraph()
+
+    assert cugraph_g.number_of_nodes()==g.number_of_nodes()
+    assert cugraph_g.number_of_edges()==g.number_of_edges()
+
+    assert cugraph_g.has_edge(0, 1)
+    assert cugraph_g.has_edge(1, 0)
+    assert cugraph_g.has_edge(3, 2)
+
+def test_from_cugraph_conversion():
+    # cudf is a dependency of cugraph
+    import cudf
+
+    # directed graph conversion test
+    cugraph_g = cugraph.Graph(directed=True)
+    df = cudf.DataFrame({"source":[0, 1, 2, 3],
+                     "destination":[1, 2, 3, 2]})
+
+    cugraph_g.from_cudf_edgelist(df)
+
+    g = dgl.from_cugraph(cugraph_g)
+
+    assert g.device.type == 'cuda'
+    assert g.number_of_nodes() == cugraph_g.number_of_nodes()
+    assert g.number_of_edges() == cugraph_g.number_of_edges()
+
+    # assert reverse edges are not present
+    assert g.has_edges_between(0, 1)
+    assert not g.has_edges_between(1, 0)
+    assert g.has_edges_between(1, 2)
+    assert not g.has_edges_between(2, 1)
+    assert g.has_edges_between(2, 3)
+
+    # undirected graph conversion test
+    cugraph_g = cugraph.Graph(directed=False)
+    df = cudf.DataFrame({"source":[0, 1, 2, 3],
+                     "destination":[1, 2, 3, 2]})
+
+    cugraph_g.from_cudf_edgelist(df)
+
+    g = dgl.from_cugraph(cugraph_g)
+
+    assert g.device.type == 'cuda'
+    assert g.number_of_nodes() == cugraph_g.number_of_nodes()
+    # assert reverse edges are present
+    assert g.has_edges_between(0, 1)
+    assert g.has_edges_between(1, 0)
+    assert g.has_edges_between(1, 2)
+    assert g.has_edges_between(2, 1)
+    assert g.has_edges_between(2, 3)

tests/dist/python/run_dist_objects.py

+import dgl
+import os
+import numpy as np
+import dgl.backend as F
+from dgl.distributed import load_partition_book
+
+mode = os.environ.get('DIST_DGL_TEST_MODE', "")
+graph_name = os.environ.get('DIST_DGL_TEST_GRAPH_NAME', 'random_test_graph')
+num_part = int(os.environ.get('DIST_DGL_TEST_NUM_PART'))
+num_servers_per_machine = int(os.environ.get('DIST_DGL_TEST_NUM_SERVER'))
+num_client_per_machine = int(os.environ.get('DIST_DGL_TEST_NUM_CLIENT'))
+shared_workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE')
+graph_path = os.environ.get('DIST_DGL_TEST_GRAPH_PATH')
+part_id = int(os.environ.get('DIST_DGL_TEST_PART_ID'))
+net_type = os.environ.get('DIST_DGL_TEST_NET_TYPE')
+ip_config = os.environ.get('DIST_DGL_TEST_IP_CONFIG', 'ip_config.txt')
+
+os.environ['DGL_DIST_MODE'] = 'distributed'
+
+def zeros_init(shape, dtype):
+    return F.zeros(shape, dtype=dtype, ctx=F.cpu())
+
+def run_server(graph_name, server_id, server_count, num_clients, shared_mem, keep_alive=False):
+    # server_count = num_servers_per_machine
+    g = dgl.distributed.DistGraphServer(server_id, ip_config,
+                        server_count, num_clients,
+                        graph_path + '/{}.json'.format(graph_name),
+                        disable_shared_mem=not shared_mem,
+                        graph_format=['csc', 'coo'], keep_alive=keep_alive,
+                        net_type=net_type)
+    print('start server', server_id)
+    g.start()
+
+##########################################
+############### DistTensor ###############
+##########################################
+
+def dist_tensor_test_sanity(data_shape, name=None):
+    local_rank = dgl.distributed.get_rank() % num_client_per_machine
+    dist_ten = dgl.distributed.DistTensor(data_shape,
+                                          F.int32,
+                                          init_func=zeros_init,
+                                          name=name)
+    # arbitrary value
+    stride = 3
+    pos = (part_id // 2) * num_client_per_machine + local_rank
+    if part_id % 2 == 0:
+        dist_ten[pos*stride:(pos+1)*stride] = F.ones((stride, 2), dtype=F.int32, ctx=F.cpu()) * (pos+1)
+
+    dgl.distributed.client_barrier()
+    assert F.allclose(dist_ten[pos*stride:(pos+1)*stride],
+                    F.ones((stride, 2), dtype=F.int32, ctx=F.cpu()) * (pos+1))
+
+
+def dist_tensor_test_destroy_recreate(data_shape, name):
+    dist_ten = dgl.distributed.DistTensor(data_shape, F.float32, name, init_func=zeros_init)
+    del dist_ten
+
+    dgl.distributed.client_barrier()
+
+    new_shape = (data_shape[0], 4)
+    dist_ten = dgl.distributed.DistTensor(new_shape, F.float32, name, init_func=zeros_init)
+
+def dist_tensor_test_persistent(data_shape):
+    dist_ten_name = 'persistent_dist_tensor'
+    dist_ten = dgl.distributed.DistTensor(data_shape, F.float32, dist_ten_name, init_func=zeros_init,
+                                          persistent=True)
+    del dist_ten
+    try:
+        dist_ten = dgl.distributed.DistTensor(data_shape, F.float32, dist_ten_name)
+        raise Exception('')
+    except:
+        pass
+
+
+def test_dist_tensor(g):
+    first_type = g.ntypes[0]
+    data_shape = (g.number_of_nodes(first_type), 2)
+    dist_tensor_test_sanity(data_shape)
+    dist_tensor_test_sanity(data_shape, name="DistTensorSanity")
+    dist_tensor_test_destroy_recreate(data_shape, name="DistTensorRecreate")
+    dist_tensor_test_persistent(data_shape)
+
+
+##########################################
+############# DistEmbedding ##############
+##########################################
+
+def dist_embedding_check_sanity(num_nodes, optimizer, name=None):
+    local_rank =  dgl.distributed.get_rank() % num_client_per_machine
+
+    emb = dgl.distributed.DistEmbedding(num_nodes, 1, name=name, init_func=zeros_init)
+    lr = 0.001
+    optim = optimizer(params=[emb], lr=lr)
+
+    stride = 3
+
+    pos = (part_id // 2) * num_client_per_machine + local_rank
+    idx = F.arange(pos*stride, (pos+1)*stride)
+
+    if part_id % 2 == 0:
+        with F.record_grad():
+            value = emb(idx)
+            optim.zero_grad()
+            loss = F.sum(value + 1, 0)
+        loss.backward()
+        optim.step()
+
+    dgl.distributed.client_barrier()
+    value = emb(idx)
+    F.allclose(value, F.ones((len(idx), 1), dtype=F.int32, ctx=F.cpu()) * -lr)
+
+    not_update_idx = F.arange(((num_part + 1) / 2) * num_client_per_machine * stride, num_nodes)
+    value = emb(not_update_idx)
+    assert np.all(F.asnumpy(value) == np.zeros((len(not_update_idx), 1)))
+
+
+def dist_embedding_check_existing(num_nodes):
+    dist_emb_name = "UniqueEmb"
+    emb = dgl.distributed.DistEmbedding(num_nodes, 1, name=dist_emb_name, init_func=zeros_init)
+    try:
+        emb1 = dgl.distributed.DistEmbedding(num_nodes, 2, name=dist_emb_name, init_func=zeros_init)
+        raise Exception('')
+    except:
+        pass
+
+def test_dist_embedding(g):
+    num_nodes = g.number_of_nodes(g.ntypes[0])
+    dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdagrad)
+    dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdagrad, name='SomeEmbedding')
+    dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdam, name='SomeEmbedding')
+
+    dist_embedding_check_existing(num_nodes)
+
+if mode == "server":
+    shared_mem = bool(int(os.environ.get('DIST_DGL_TEST_SHARED_MEM')))
+    server_id = int(os.environ.get('DIST_DGL_TEST_SERVER_ID'))
+    run_server(graph_name, server_id, server_count=num_servers_per_machine,
+               num_clients=num_part*num_client_per_machine, shared_mem=shared_mem, keep_alive=False)
+elif mode == "client":
+    os.environ['DGL_NUM_SERVER'] = str(num_servers_per_machine)
+    dgl.distributed.initialize(ip_config, net_type=net_type)
+
+    gpb, graph_name, _, _ = load_partition_book(graph_path + '/{}.json'.format(graph_name), part_id, None)
+    g = dgl.distributed.DistGraph(graph_name, gpb=gpb)
+
+    target_func_map = {"DistTensor": test_dist_tensor,
+                       "DistEmbedding": test_dist_embedding,
+                       }
+
+    target = os.environ.get("DIST_DGL_TEST_OBJECT_TYPE", "")
+    if target not in target_func_map:
+        for test_func in target_func_map.values():
+            test_func(g)
+    else:
+        target_func_map[target](g)
+
+else:
+    print("DIST_DGL_TEST_MODE has to be either server or client")
+    exit(1)
+

tests/dist/test_dist_objects.py

+import os
+import unittest
+import pytest
+import multiprocessing as mp
+import subprocess
+import utils
+import dgl
+import numpy as np
+import dgl.backend as F
+from dgl.distributed import partition_graph
+
+graph_name = os.environ.get('DIST_DGL_TEST_GRAPH_NAME', 'random_test_graph')
+target = os.environ.get('DIST_DGL_TEST_OBJECT_TYPE', '')
+shared_workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE')
+
+def create_graph(num_part, dist_graph_path, hetero):
+    if not hetero:
+        g = dgl.rand_graph(10000, 42000)
+        g.ndata['feat'] = F.unsqueeze(F.arange(0, g.number_of_nodes()), 1)
+        g.edata['feat'] = F.unsqueeze(F.arange(0, g.number_of_edges()), 1)
+        partition_graph(g, graph_name, num_part, dist_graph_path)
+    else:
+        from scipy import sparse as spsp
+        num_nodes = {'n1': 10000, 'n2': 10010, 'n3': 10020}
+        etypes = [('n1', 'r1', 'n2'),
+                ('n1', 'r2', 'n3'),
+                ('n2', 'r3', 'n3')]
+        edges = {}
+        for etype in etypes:
+            src_ntype, _, dst_ntype = etype
+            arr = spsp.random(num_nodes[src_ntype], num_nodes[dst_ntype], density=0.001, format='coo',
+                            random_state=100)
+            edges[etype] = (arr.row, arr.col)
+        g = dgl.heterograph(edges, num_nodes)
+        g.nodes['n1'].data['feat'] = F.unsqueeze(F.arange(0, g.number_of_nodes('n1')), 1)
+        g.edges['r1'].data['feat'] = F.unsqueeze(F.arange(0, g.number_of_edges('r1')), 1)
+        partition_graph(g, graph_name, num_part, dist_graph_path)
+
+
+@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
+@pytest.mark.parametrize("net_type", ['tensorpipe', 'socket'])
+@pytest.mark.parametrize("num_servers", [1, 4])
+@pytest.mark.parametrize("num_clients", [1, 4])
+@pytest.mark.parametrize("hetero", [False, True])
+@pytest.mark.parametrize("shared_mem", [False, True])
+def test_dist_objects(net_type, num_servers, num_clients, hetero, shared_mem):
+    if not shared_mem and num_servers > 1:
+        pytest.skip(f"Backup servers are not supported when shared memory is disabled")
+    ip_config = os.environ.get('DIST_DGL_TEST_IP_CONFIG', 'ip_config.txt')
+    workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE', '/shared_workspace/dgl_dist_tensor_test/')
+
+    ips = utils.get_ips(ip_config)
+    num_part = len(ips)
+
+    test_bin = os.path.join(os.environ.get(
+        'DIST_DGL_TEST_PY_BIN_DIR', '.'), 'run_dist_objects.py')
+
+    dist_graph_path = os.path.join(workspace, 'hetero_dist_graph' if hetero else 'dist_graph')
+    if not os.path.isdir(dist_graph_path):
+        create_graph(num_part, dist_graph_path, hetero)
+
+    base_envs = f"DIST_DGL_TEST_WORKSPACE={workspace} " \
+                f"DIST_DGL_TEST_NUM_PART={num_part} " \
+                f"DIST_DGL_TEST_NUM_SERVER={num_servers} " \
+                f"DIST_DGL_TEST_NUM_CLIENT={num_clients} " \
+                f"DIST_DGL_TEST_NET_TYPE={net_type} " \
+                f"DIST_DGL_TEST_GRAPH_PATH={dist_graph_path} " \
+                f"DIST_DGL_TEST_IP_CONFIG={ip_config} "
+
+    procs = []
+    # Start server
+    server_id = 0
+    for part_id, ip in enumerate(ips):
+        for _ in range(num_servers):
+            cmd_envs = base_envs + \
+                       f"DIST_DGL_TEST_SERVER_ID={server_id} " \
+                       f"DIST_DGL_TEST_PART_ID={part_id} " \
+                       f"DIST_DGL_TEST_SHARED_MEM={str(int(shared_mem))} " \
+                       f"DIST_DGL_TEST_MODE=server "
+            procs.append(utils.execute_remote(
+                f"{cmd_envs} python3 {test_bin}",
+                ip))
+            server_id += 1
+    # Start client processes
+    for part_id, ip in enumerate(ips):
+        for _ in range(num_clients):
+            cmd_envs = base_envs + \
+                       f"DIST_DGL_TEST_PART_ID={part_id} " \
+                       f"DIST_DGL_TEST_OBJECT_TYPE={target} " \
+                       f"DIST_DGL_TEST_MODE=client "
+            procs.append(utils.execute_remote(
+                f"{cmd_envs} python3 {test_bin}",
+                 ip))
+
+    for p in procs:
+        p.join()
+        assert p.exitcode == 0
+

tests/distributed/test_dist_graph_store.py

@@ -35,6 +35,15 @@ def run_server(graph_name, server_id, server_count, num_clients, shared_mem, kee
                         disable_shared_mem=not shared_mem,
                         graph_format=['csc', 'coo'], keep_alive=keep_alive)
     print('start server', server_id)
+    # verify dtype of underlying graph
+    cg = g.client_g
+    for k, dtype in dgl.distributed.dist_graph.FIELD_DICT.items():
+        if k in cg.ndata:
+            assert F.dtype(
+                cg.ndata[k]) == dtype, "Data type of {} in ndata should be {}.".format(k, dtype)
+        if k in cg.edata:
+            assert F.dtype(
+                cg.edata[k]) == dtype, "Data type of {} in edata should be {}.".format(k, dtype)
     g.start()
 
 def emb_init(shape, dtype):
@@ -219,6 +228,11 @@ def check_dist_graph(g, num_clients, num_nodes, num_edges):
     feats = F.squeeze(feats1, 1)
     assert np.all(F.asnumpy(feats == eids))
 
+    # Test edge_subgraph
+    sg = g.edge_subgraph(eids)
+    assert sg.num_edges() == len(eids)
+    assert F.array_equal(sg.edata[dgl.EID], eids)
+
     # Test init node data
     new_shape = (g.number_of_nodes(), 2)
     test1 = dgl.distributed.DistTensor(new_shape, F.int32)
@@ -485,6 +499,14 @@ def check_dist_graph_hetero(g, num_clients, num_nodes, num_edges):
     feats = F.squeeze(feats1, 1)
     assert np.all(F.asnumpy(feats == eids))
 
+    # Test edge_subgraph
+    sg = g.edge_subgraph({'r1': eids})
+    assert sg.num_edges() == len(eids)
+    assert F.array_equal(sg.edata[dgl.EID], eids)
+    sg = g.edge_subgraph({('n1', 'r1', 'n2'): eids})
+    assert sg.num_edges() == len(eids)
+    assert F.array_equal(sg.edata[dgl.EID], eids)
+
     # Test init node data
     new_shape = (g.number_of_nodes('n1'), 2)
     g.nodes['n1'].data['test1'] = dgl.distributed.DistTensor(new_shape, F.int32)

tests/distributed/test_distributed_sampling.py

@@ -84,7 +84,7 @@ def start_get_degrees_client(rank, tmpdir, disable_shared_mem, nids=None):
     gpb = None
     if disable_shared_mem:
         _, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_get_degrees.json', rank)
-    dgl.distributed.initialize("rpc_ip_config.txt", 1)
+    dgl.distributed.initialize("rpc_ip_config.txt")
     dist_graph = DistGraph("test_get_degrees", gpb=gpb)
     try:
         in_deg = dist_graph.in_degrees(nids)
@@ -160,9 +160,9 @@ def check_rpc_find_edges_shuffle(tmpdir, num_server):
 def create_random_hetero(dense=False, empty=False):
     num_nodes = {'n1': 210, 'n2': 200, 'n3': 220} if dense else \
         {'n1': 1010, 'n2': 1000, 'n3': 1020}
-    etypes = [('n1', 'r1', 'n2'),
-              ('n1', 'r2', 'n3'),
-              ('n2', 'r3', 'n3')]
+    etypes = [('n1', 'r12', 'n2'),
+              ('n1', 'r13', 'n3'),
+              ('n2', 'r23', 'n3')]
     edges = {}
     random.seed(42)
     for etype in etypes:
@@ -195,9 +195,18 @@ def check_rpc_hetero_find_edges_shuffle(tmpdir, num_server):
         time.sleep(1)
         pserver_list.append(p)
 
-    eids = F.tensor(np.random.randint(g.number_of_edges('r1'), size=100))
-    u, v = g.find_edges(orig_eid['r1'][eids], etype='r1')
-    du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids, etype='r1')
+    eids = F.tensor(np.random.randint(g.num_edges('r12'), size=100))
+    expect_except = False
+    try:
+        _, _ = g.find_edges(orig_eid['r12'][eids], etype=('n1', 'r12'))
+    except:
+        expect_except = True
+    assert expect_except
+    u, v = g.find_edges(orig_eid['r12'][eids], etype='r12')
+    u1, v1 = g.find_edges(orig_eid['r12'][eids], etype=('n1', 'r12', 'n2'))
+    assert F.array_equal(u, u1)
+    assert F.array_equal(v, v1)
+    du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids, etype='r12')
     du = orig_nid['n1'][du]
     dv = orig_nid['n2'][dv]
     assert F.array_equal(u, du)
@@ -335,7 +344,8 @@ def start_hetero_sample_client(rank, tmpdir, disable_shared_mem, nodes):
     return block, gpb
 
 def start_hetero_etype_sample_client(rank, tmpdir, disable_shared_mem, fanout=3,
-                                     nodes={'n3': [0, 10, 99, 66, 124, 208]}):
+                                     nodes={'n3': [0, 10, 99, 66, 124, 208]},
+                                     etype_sorted=False):
     gpb = None
     if disable_shared_mem:
         _, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json', rank)
@@ -358,7 +368,7 @@ def start_hetero_etype_sample_client(rank, tmpdir, disable_shared_mem, fanout=3,
     if gpb is None:
         gpb = dist_graph.get_partition_book()
     try:
-        sampled_graph = sample_etype_neighbors(dist_graph, nodes, dgl.ETYPE, fanout)
+        sampled_graph = sample_etype_neighbors(dist_graph, nodes, dgl.ETYPE, fanout, etype_sorted=etype_sorted)
         block = dgl.to_block(sampled_graph, nodes)
         block.edata[dgl.EID] = sampled_graph.edata[dgl.EID]
     except Exception as e:
@@ -461,7 +471,7 @@ def check_rpc_hetero_sampling_empty_shuffle(tmpdir, num_server):
     assert block.number_of_edges() == 0
     assert len(block.etypes) == len(g.etypes)
 
-def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server):
+def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server, etype_sorted=False):
     generate_ip_config("rpc_ip_config.txt", num_server, num_server)
 
     g = create_random_hetero(dense=True)
@@ -474,21 +484,22 @@ def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server):
     pserver_list = []
     ctx = mp.get_context('spawn')
     for i in range(num_server):
-        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling'))
+        p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling', ['csc', 'coo']))
         p.start()
         time.sleep(1)
         pserver_list.append(p)
 
     fanout = 3
     block, gpb = start_hetero_etype_sample_client(0, tmpdir, num_server > 1, fanout,
-                                                  nodes={'n3': [0, 10, 99, 66, 124, 208]})
+                                                  nodes={'n3': [0, 10, 99, 66, 124, 208]},
+                                                  etype_sorted=etype_sorted)
     print("Done sampling")
     for p in pserver_list:
         p.join()
 
-    src, dst = block.edges(etype=('n1', 'r2', 'n3'))
+    src, dst = block.edges(etype=('n1', 'r13', 'n3'))
     assert len(src) == 18
-    src, dst = block.edges(etype=('n2', 'r3', 'n3'))
+    src, dst = block.edges(etype=('n2', 'r23', 'n3'))
     assert len(src) == 18
 
     orig_nid_map = {ntype: F.zeros((g.number_of_nodes(ntype),), dtype=F.int64) for ntype in g.ntypes}
@@ -832,6 +843,7 @@ def test_rpc_sampling_shuffle(num_server):
         check_rpc_hetero_sampling_shuffle(Path(tmpdirname), num_server)
         check_rpc_hetero_sampling_empty_shuffle(Path(tmpdirname), num_server)
         check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server)
+        check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server, etype_sorted=True)
         check_rpc_hetero_etype_sampling_empty_shuffle(Path(tmpdirname), num_server)
         check_rpc_bipartite_sampling_empty(Path(tmpdirname), num_server)
         check_rpc_bipartite_sampling_shuffle(Path(tmpdirname), num_server)

tests/pytorch/test_dataloader.py

@@ -5,11 +5,9 @@ import dgl.ops as OPS
 import backend as F
 import unittest
 import torch
+import torch.distributed as dist
 from functools import partial
-from torch.utils.data import DataLoader
-from collections import defaultdict
 from collections.abc import Iterator, Mapping
-from itertools import product
 from test_utils import parametrize_idtype
 import pytest
 
@@ -40,7 +38,7 @@ def test_cluster_gcn(num_workers):
 def test_shadow(num_workers):
     g = dgl.data.CoraFullDataset()[0]
     sampler = dgl.dataloading.ShaDowKHopSampler([5, 10, 15])
-    dataloader = dgl.dataloading.NodeDataLoader(
+    dataloader = dgl.dataloading.DataLoader(
         g, torch.arange(g.num_nodes()), sampler,
         batch_size=5, shuffle=True, drop_last=False, num_workers=num_workers)
     for i, (input_nodes, output_nodes, subgraph) in enumerate(dataloader):
@@ -70,43 +68,73 @@ def test_saint(num_workers, mode):
     for sg in dataloader:
         pass
 
-@pytest.mark.parametrize('num_workers', [0, 4])
-def test_neighbor_nonuniform(num_workers):
-    g = dgl.graph(([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1]))
+@parametrize_idtype
+@pytest.mark.parametrize('mode', ['cpu', 'uva_cuda_indices', 'uva_cpu_indices', 'pure_gpu'])
+@pytest.mark.parametrize('use_ddp', [False, True])
+def test_neighbor_nonuniform(idtype, mode, use_ddp):
+    if mode != 'cpu' and F.ctx() == F.cpu():
+        pytest.skip('UVA and GPU sampling require a GPU.')
+    if use_ddp:
+        dist.init_process_group('gloo' if F.ctx() == F.cpu() else 'nccl',
+            'tcp://127.0.0.1:12347', world_size=1, rank=0)
+    g = dgl.graph(([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1])).astype(idtype)
     g.edata['p'] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
+    if mode in ('cpu', 'uva_cpu_indices'):
+        indices = F.copy_to(F.tensor([0, 1], idtype), F.cpu())
+    else:
+        indices = F.copy_to(F.tensor([0, 1], idtype), F.cuda())
+    if mode == 'pure_gpu':
+        g = g.to(F.cuda())
+    use_uva = mode.startswith('uva')
+
     sampler = dgl.dataloading.MultiLayerNeighborSampler([2], prob='p')
-    dataloader = dgl.dataloading.NodeDataLoader(g, [0, 1], sampler, batch_size=1, device=F.ctx())
-    for input_nodes, output_nodes, blocks in dataloader:
-        seed = output_nodes.item()
-        neighbors = set(input_nodes[1:].cpu().numpy())
-        if seed == 1:
-            assert neighbors == {5, 6}
-        elif seed == 0:
-            assert neighbors == {1, 2}
+    for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
+        dataloader = dgl.dataloading.NodeDataLoader(
+            g, indices, sampler,
+            batch_size=1, device=F.ctx(),
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp)
+        for input_nodes, output_nodes, blocks in dataloader:
+            seed = output_nodes.item()
+            neighbors = set(input_nodes[1:].cpu().numpy())
+            if seed == 1:
+                assert neighbors == {5, 6}
+            elif seed == 0:
+                assert neighbors == {1, 2}
 
     g = dgl.heterograph({
         ('B', 'BA', 'A'): ([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1]),
         ('C', 'CA', 'A'): ([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1]),
-        })
+        }).astype(idtype)
     g.edges['BA'].data['p'] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
     g.edges['CA'].data['p'] = torch.FloatTensor([0, 0, 1, 1, 0, 0, 1, 1])
-    sampler = dgl.dataloading.MultiLayerNeighborSampler([2], prob='p')
-    dataloader = dgl.dataloading.NodeDataLoader(
-        g, {'A': [0, 1]}, sampler, batch_size=1, device=F.ctx())
-    for input_nodes, output_nodes, blocks in dataloader:
-        seed = output_nodes['A'].item()
-        # Seed and neighbors are of different node types so slicing is not necessary here.
-        neighbors = set(input_nodes['B'].cpu().numpy())
-        if seed == 1:
-            assert neighbors == {5, 6}
-        elif seed == 0:
-            assert neighbors == {1, 2}
-
-        neighbors = set(input_nodes['C'].cpu().numpy())
-        if seed == 1:
-            assert neighbors == {7, 8}
-        elif seed == 0:
-            assert neighbors == {3, 4}
+    if mode == 'pure_gpu':
+        g = g.to(F.cuda())
+    for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
+        dataloader = dgl.dataloading.NodeDataLoader(
+            g, {'A': indices}, sampler,
+            batch_size=1, device=F.ctx(),
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp)
+        for input_nodes, output_nodes, blocks in dataloader:
+            seed = output_nodes['A'].item()
+            # Seed and neighbors are of different node types so slicing is not necessary here.
+            neighbors = set(input_nodes['B'].cpu().numpy())
+            if seed == 1:
+                assert neighbors == {5, 6}
+            elif seed == 0:
+                assert neighbors == {1, 2}
+
+            neighbors = set(input_nodes['C'].cpu().numpy())
+            if seed == 1:
+                assert neighbors == {7, 8}
+            elif seed == 0:
+                assert neighbors == {3, 4}
+
+    if use_ddp:
+        dist.destroy_process_group()
 
 def _check_dtype(data, dtype, attr_name):
     if isinstance(data, dict):
@@ -130,38 +158,37 @@ def _check_device(data):
 
 @parametrize_idtype
 @pytest.mark.parametrize('sampler_name', ['full', 'neighbor', 'neighbor2'])
-@pytest.mark.parametrize('pin_graph', [None, 'cuda_indices', 'cpu_indices'])
-def test_node_dataloader(idtype, sampler_name, pin_graph):
+@pytest.mark.parametrize('mode', ['cpu', 'uva_cuda_indices', 'uva_cpu_indices', 'pure_gpu'])
+@pytest.mark.parametrize('use_ddp', [False, True])
+def test_node_dataloader(idtype, sampler_name, mode, use_ddp):
+    if mode != 'cpu' and F.ctx() == F.cpu():
+        pytest.skip('UVA and GPU sampling require a GPU.')
+    if use_ddp:
+        dist.init_process_group('gloo' if F.ctx() == F.cpu() else 'nccl',
+            'tcp://127.0.0.1:12347', world_size=1, rank=0)
     g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
     g1.ndata['feat'] = F.copy_to(F.randn((5, 8)), F.cpu())
     g1.ndata['label'] = F.copy_to(F.randn((g1.num_nodes(),)), F.cpu())
-    indices = F.arange(0, g1.num_nodes(), idtype)
-    if F.ctx() != F.cpu():
-        if pin_graph:
-            g1.create_formats_()
-            g1.pin_memory_()
-            if pin_graph == 'cpu_indices':
-                indices = F.arange(0, g1.num_nodes(), idtype, F.cpu())
-            elif pin_graph == 'cuda_indices':
-                if F._default_context_str == 'gpu':
-                    indices = F.arange(0, g1.num_nodes(), idtype, F.cuda())
-                else:
-                    return  # skip
-        else:
-            g1 = g1.to('cuda')
+    if mode in ('cpu', 'uva_cpu_indices'):
+        indices = F.copy_to(F.arange(0, g1.num_nodes(), idtype), F.cpu())
+    else:
+        indices = F.copy_to(F.arange(0, g1.num_nodes(), idtype), F.cuda())
+    if mode == 'pure_gpu':
+        g1 = g1.to(F.cuda())
 
-    use_uva = pin_graph is not None and F.ctx() != F.cpu()
+    use_uva = mode.startswith('uva')
 
-    for num_workers in [0, 1, 2]:
-        sampler = {
-            'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
-            'neighbor': dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
-            'neighbor2': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
-        dataloader = dgl.dataloading.NodeDataLoader(
+    sampler = {
+        'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
+        'neighbor': dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
+        'neighbor2': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
+    for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
+        dataloader = dgl.dataloading.DataLoader(
             g1, indices, sampler, device=F.ctx(),
             batch_size=g1.num_nodes(),
-            num_workers=(num_workers if (pin_graph and F.ctx() == F.cpu()) else 0),
-            use_uva=use_uva)
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp)
         for input_nodes, output_nodes, blocks in dataloader:
             _check_device(input_nodes)
             _check_device(output_nodes)
@@ -169,8 +196,6 @@ def test_node_dataloader(idtype, sampler_name, pin_graph):
             _check_dtype(input_nodes, idtype, 'dtype')
             _check_dtype(output_nodes, idtype, 'dtype')
             _check_dtype(blocks, idtype, 'idtype')
-    if g1.is_pinned():
-        g1.unpin_memory_()
 
     g2 = dgl.heterograph({
          ('user', 'follow', 'user'): ([0, 0, 0, 1, 1, 1, 2], [1, 2, 3, 0, 2, 3, 0]),
@@ -180,73 +205,80 @@ def test_node_dataloader(idtype, sampler_name, pin_graph):
     }).astype(idtype)
     for ntype in g2.ntypes:
         g2.nodes[ntype].data['feat'] = F.copy_to(F.randn((g2.num_nodes(ntype), 8)), F.cpu())
-    indices = {nty: F.arange(0, g2.num_nodes(nty)) for nty in g2.ntypes}
-    if F.ctx() != F.cpu():
-        if pin_graph:
-            g2.create_formats_()
-            g2.pin_memory_()
-            if pin_graph == 'cpu_indices':
-                indices = {nty: F.arange(0, g2.num_nodes(nty), idtype, F.cpu()) for nty in g2.ntypes}
-            elif pin_graph == 'cuda_indices':
-                if F._default_context_str == 'gpu':
-                    indices = {nty: F.arange(0, g2.num_nodes(), idtype, F.cuda()) for nty in g2.ntypes}
-                else:
-                    return  # skip
-        else:
-            g2 = g2.to('cuda')
+    if mode in ('cpu', 'uva_cpu_indices'):
+        indices = {nty: F.copy_to(g2.nodes(nty), F.cpu()) for nty in g2.ntypes}
+    else:
+        indices = {nty: F.copy_to(g2.nodes(nty), F.cuda()) for nty in g2.ntypes}
+    if mode == 'pure_gpu':
+        g2 = g2.to(F.cuda())
 
     batch_size = max(g2.num_nodes(nty) for nty in g2.ntypes)
     sampler = {
         'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
         'neighbor': dgl.dataloading.MultiLayerNeighborSampler([{etype: 3 for etype in g2.etypes}] * 2),
         'neighbor2': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
+    for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
+        dataloader = dgl.dataloading.DataLoader(
+            g2, indices, sampler,
+            device=F.ctx(), batch_size=batch_size,
+            num_workers=num_workers,
+            use_uva=use_uva,
+            use_ddp=use_ddp)
+        assert isinstance(iter(dataloader), Iterator)
+        for input_nodes, output_nodes, blocks in dataloader:
+            _check_device(input_nodes)
+            _check_device(output_nodes)
+            _check_device(blocks)
+            _check_dtype(input_nodes, idtype, 'dtype')
+            _check_dtype(output_nodes, idtype, 'dtype')
+            _check_dtype(blocks, idtype, 'idtype')
 
-    dataloader = dgl.dataloading.NodeDataLoader(
-        g2, {nty: g2.nodes(nty) for nty in g2.ntypes},
-        sampler, device=F.ctx(), batch_size=batch_size,
-        num_workers=(num_workers if (pin_graph and F.ctx() == F.cpu()) else 0),
-        use_uva=use_uva)
-    assert isinstance(iter(dataloader), Iterator)
-    for input_nodes, output_nodes, blocks in dataloader:
-        _check_device(input_nodes)
-        _check_device(output_nodes)
-        _check_device(blocks)
-        _check_dtype(input_nodes, idtype, 'dtype')
-        _check_dtype(output_nodes, idtype, 'dtype')
-        _check_dtype(blocks, idtype, 'idtype')
-
-    if g2.is_pinned():
-        g2.unpin_memory_()
+    if use_ddp:
+        dist.destroy_process_group()
 
+@parametrize_idtype
 @pytest.mark.parametrize('sampler_name', ['full', 'neighbor'])
 @pytest.mark.parametrize('neg_sampler', [
     dgl.dataloading.negative_sampler.Uniform(2),
     dgl.dataloading.negative_sampler.GlobalUniform(15, False, 3),
     dgl.dataloading.negative_sampler.GlobalUniform(15, True, 3)])
-@pytest.mark.parametrize('pin_graph', [False, True])
-def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
-    g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4]))
-    if F.ctx() != F.cpu() and pin_graph:
-        g1.create_formats_()
-        g1.pin_memory_()
+@pytest.mark.parametrize('mode', ['cpu', 'uva', 'pure_gpu'])
+@pytest.mark.parametrize('use_ddp', [False, True])
+def test_edge_dataloader(idtype, sampler_name, neg_sampler, mode, use_ddp):
+    if mode != 'cpu' and F.ctx() == F.cpu():
+        pytest.skip('UVA and GPU sampling require a GPU.')
+    if mode == 'uva' and isinstance(neg_sampler, dgl.dataloading.negative_sampler.GlobalUniform):
+        pytest.skip("GlobalUniform don't support UVA yet.")
+    if use_ddp:
+        dist.init_process_group('gloo' if F.ctx() == F.cpu() else 'nccl',
+            'tcp://127.0.0.1:12347', world_size=1, rank=0)
+    g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
     g1.ndata['feat'] = F.copy_to(F.randn((5, 8)), F.cpu())
+    if mode == 'pure_gpu':
+        g1 = g1.to(F.cuda())
 
     sampler = {
         'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
         'neighbor': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
 
     # no negative sampler
-    dataloader = dgl.dataloading.EdgeDataLoader(
-        g1, g1.edges(form='eid'), sampler, device=F.ctx(), batch_size=g1.num_edges())
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
+    dataloader = dgl.dataloading.DataLoader(
+        g1, g1.edges(form='eid'), edge_sampler,
+        device=F.ctx(), batch_size=g1.num_edges(),
+        use_uva=(mode == 'uva'), use_ddp=use_ddp)
     for input_nodes, pos_pair_graph, blocks in dataloader:
         _check_device(input_nodes)
         _check_device(pos_pair_graph)
         _check_device(blocks)
 
     # negative sampler
-    dataloader = dgl.dataloading.EdgeDataLoader(
-        g1, g1.edges(form='eid'), sampler, device=F.ctx(),
-        negative_sampler=neg_sampler, batch_size=g1.num_edges())
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(
+        sampler, negative_sampler=neg_sampler)
+    dataloader = dgl.dataloading.DataLoader(
+        g1, g1.edges(form='eid'), edge_sampler,
+        device=F.ctx(), batch_size=g1.num_edges(),
+        use_uva=(mode == 'uva'), use_ddp=use_ddp)
     for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
         _check_device(input_nodes)
         _check_device(pos_pair_graph)
@@ -258,9 +290,12 @@ def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
          ('user', 'followed-by', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]),
          ('user', 'play', 'game'): ([0, 1, 1, 3, 5], [0, 1, 2, 0, 2]),
          ('game', 'played-by', 'user'): ([0, 1, 2, 0, 2], [0, 1, 1, 3, 5])
-    })
+    }).astype(idtype)
     for ntype in g2.ntypes:
         g2.nodes[ntype].data['feat'] = F.copy_to(F.randn((g2.num_nodes(ntype), 8)), F.cpu())
+    if mode == 'pure_gpu':
+        g2 = g2.to(F.cuda())
+
     batch_size = max(g2.num_edges(ety) for ety in g2.canonical_etypes)
     sampler = {
         'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
@@ -268,19 +303,23 @@ def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
         }[sampler_name]
 
     # no negative sampler
-    dataloader = dgl.dataloading.EdgeDataLoader(
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
+    dataloader = dgl.dataloading.DataLoader(
         g2, {ety: g2.edges(form='eid', etype=ety) for ety in g2.canonical_etypes},
-        sampler, device=F.ctx(), batch_size=batch_size)
+        edge_sampler, device=F.ctx(), batch_size=batch_size,
+        use_uva=(mode == 'uva'), use_ddp=use_ddp)
     for input_nodes, pos_pair_graph, blocks in dataloader:
         _check_device(input_nodes)
         _check_device(pos_pair_graph)
         _check_device(blocks)
 
     # negative sampler
-    dataloader = dgl.dataloading.EdgeDataLoader(
+    edge_sampler = dgl.dataloading.as_edge_prediction_sampler(
+        sampler, negative_sampler=neg_sampler)
+    dataloader = dgl.dataloading.DataLoader(
         g2, {ety: g2.edges(form='eid', etype=ety) for ety in g2.canonical_etypes},
-        sampler, device=F.ctx(), negative_sampler=neg_sampler,
-        batch_size=batch_size)
+        edge_sampler, device=F.ctx(),batch_size=batch_size,
+        use_uva=(mode == 'uva'), use_ddp=use_ddp)
 
     assert isinstance(iter(dataloader), Iterator)
     for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
@@ -289,8 +328,8 @@ def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
         _check_device(neg_pair_graph)
         _check_device(blocks)
 
-    if g1.is_pinned():
-        g1.unpin_memory_()
+    if use_ddp:
+        dist.destroy_process_group()
 
 def _create_homogeneous():
     s = torch.randint(0, 200, (1000,), device=F.ctx())

tests/pytorch/test_nn.py

@@ -412,6 +412,60 @@ def test_rgcn(idtype, O):
         h_new = rgc_bdd(g, h, r, norm)
         assert h_new.shape == (100, O)
 
+@parametrize_idtype
+@pytest.mark.parametrize('O', [1, 10, 40])
+def test_rgcn_default_nbasis(idtype, O):
+    ctx = F.ctx()
+    etype = []
+    g = dgl.from_scipy(sp.sparse.random(100, 100, density=0.1))
+    g = g.astype(idtype).to(F.ctx())
+    # 5 etypes
+    R = 5
+    for i in range(g.number_of_edges()):
+        etype.append(i % 5)
+    I = 10
+
+    h = th.randn((100, I)).to(ctx)
+    r = th.tensor(etype).to(ctx)
+    norm = th.rand((g.number_of_edges(), 1)).to(ctx)
+    sorted_r, idx = th.sort(r)
+    sorted_g = dgl.reorder_graph(g, edge_permute_algo='custom', permute_config={'edges_perm' : idx.to(idtype)})
+    sorted_norm = norm[idx]
+
+    rgc = nn.RelGraphConv(I, O, R).to(ctx)
+    th.save(rgc, tmp_buffer)  # test pickle
+    rgc_basis = nn.RelGraphConv(I, O, R, "basis").to(ctx)
+    th.save(rgc_basis, tmp_buffer)  # test pickle
+    if O % R == 0:
+        rgc_bdd = nn.RelGraphConv(I, O, R, "bdd").to(ctx)
+        th.save(rgc_bdd, tmp_buffer)  # test pickle
+
+    # basic usage
+    h_new = rgc(g, h, r)
+    assert h_new.shape == (100, O)
+    h_new_basis = rgc_basis(g, h, r)
+    assert h_new_basis.shape == (100, O)
+    if O % R == 0:
+        h_new_bdd = rgc_bdd(g, h, r)
+        assert h_new_bdd.shape == (100, O)
+
+    # sorted input
+    h_new_sorted = rgc(sorted_g, h, sorted_r, presorted=True)
+    assert th.allclose(h_new, h_new_sorted, atol=1e-4, rtol=1e-4)
+    h_new_basis_sorted = rgc_basis(sorted_g, h, sorted_r, presorted=True)
+    assert th.allclose(h_new_basis, h_new_basis_sorted, atol=1e-4, rtol=1e-4)
+    if O % R == 0:
+        h_new_bdd_sorted = rgc_bdd(sorted_g, h, sorted_r, presorted=True)
+        assert th.allclose(h_new_bdd, h_new_bdd_sorted, atol=1e-4, rtol=1e-4)
+
+    # norm input
+    h_new = rgc(g, h, r, norm)
+    assert h_new.shape == (100, O)
+    h_new = rgc_basis(g, h, r, norm)
+    assert h_new.shape == (100, O)
+    if O % R == 0:
+        h_new = rgc_bdd(g, h, r, norm)
+        assert h_new.shape == (100, O)
 
 @parametrize_idtype
 @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
@@ -647,10 +701,11 @@ def test_appnp_conv_e_weight(g, idtype):
 
 @parametrize_idtype
 @pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
-def test_gcn2conv_e_weight(g, idtype):
+@pytest.mark.parametrize("bias", [True, False])
+def test_gcn2conv_e_weight(g, idtype, bias):
     ctx = F.ctx()
     g = g.astype(idtype).to(ctx)
-    gcn2conv = nn.GCN2Conv(5, layer=2, alpha=0.5,
+    gcn2conv = nn.GCN2Conv(5, layer=2, alpha=0.5, bias=bias,
                            project_initial_features=True)
     feat = F.randn((g.number_of_nodes(), 5))
     eweight = F.ones((g.num_edges(), ))
@@ -1207,6 +1262,11 @@ def test_hetero_embedding(out_dim):
     assert embeds['user'].shape == (2, out_dim)
     assert embeds[('user', 'follows', 'user')].shape == (3, out_dim)
 
+    layer.reset_parameters()
+    embeds = layer.weight
+    assert embeds['user'].shape == (2, out_dim)
+    assert embeds[('user', 'follows', 'user')].shape == (3, out_dim)
+
     embeds = layer({
         'user': F.tensor([0], dtype=F.int64),
         ('user', 'follows', 'user'): F.tensor([0, 2], dtype=F.int64)
@@ -1561,4 +1621,4 @@ def test_dgn_conv(in_size, out_size, aggregators, scalers, delta,
     aggregators_non_eig = [aggr for aggr in aggregators if not aggr.startswith('dir')]
     model = nn.DGNConv(in_size, out_size, aggregators_non_eig, scalers, delta, dropout,
         num_towers, edge_feat_size, residual).to(dev)
-    model(g, h, edge_feat=e)
+    model(g, h, edge_feat=e)
\ No newline at end of file

tests/pytorch/test_unified_tensor.py

@@ -18,23 +18,24 @@ def start_unified_tensor_worker(dev_id, input, seq_idx, rand_idx, output_seq, ou
 def test_unified_tensor():
     test_row_size = 65536
     test_col_size = 128
-
     rand_test_size = 8192
 
+    device = th.device('cuda:0')
+
     input = th.rand((test_row_size, test_col_size))
-    input_unified = dgl.contrib.UnifiedTensor(input, device=th.device('cuda'))
+    input_unified = dgl.contrib.UnifiedTensor(input, device=device)
 
     seq_idx = th.arange(0, test_row_size)
+    # CPU indexing
     assert th.all(th.eq(input[seq_idx], input_unified[seq_idx]))
-
-    seq_idx = seq_idx.to(th.device('cuda'))
-    assert th.all(th.eq(input[seq_idx].to(th.device('cuda')), input_unified[seq_idx]))
+    # GPU indexing
+    assert th.all(th.eq(input[seq_idx].to(device), input_unified[seq_idx.to(device)]))
 
     rand_idx = th.randint(0, test_row_size, (rand_test_size,))
+    # CPU indexing
     assert th.all(th.eq(input[rand_idx], input_unified[rand_idx]))
-
-    rand_idx = rand_idx.to(th.device('cuda'))
-    assert th.all(th.eq(input[rand_idx].to(th.device('cuda')), input_unified[rand_idx]))
+    # GPU indexing
+    assert th.all(th.eq(input[rand_idx].to(device), input_unified[rand_idx.to(device)]))
 
 @unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
 @unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test')

tests/scripts/build_dgl.sh

@@ -3,7 +3,7 @@ set -e
 . /opt/conda/etc/profile.d/conda.sh
 
 if [ $# -ne 1 ]; then
-    echo "Device argument required, can be cpu or gpu"
+    echo "Device argument required, can be cpu, gpu or cugraph"
     exit -1
 fi
 
@@ -11,7 +11,9 @@ CMAKE_VARS="-DBUILD_CPP_TEST=ON -DUSE_OPENMP=ON -DBUILD_TORCH=ON"
 # This is a semicolon-separated list of Python interpreters containing PyTorch.
 # The value here is for CI.  Replace it with your own or comment this whole
 # statement for default Python interpreter.
-CMAKE_VARS="$CMAKE_VARS -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
+if [ "$1" != "cugraph" ]; then
+    CMAKE_VARS="$CMAKE_VARS -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
+fi
 
 #This is implemented to detect underlying architecture and enable arch specific optimization.
 arch=`uname -m`
@@ -19,7 +21,7 @@ if [[ $arch == *"x86"* ]]; then
   CMAKE_VARS="-DUSE_AVX=ON $CMAKE_VARS"
 fi
 
-if [ "$1" == "gpu" ]; then
+if [[ $1 != "cpu" ]]; then
     CMAKE_VARS="-DUSE_CUDA=ON -DUSE_NCCL=ON -DUSE_FP16=ON $CMAKE_VARS"
 fi
 
@@ -36,14 +38,23 @@ make -j
 popd
 
 pushd python
-for backend in pytorch mxnet tensorflow
-do
-conda activate "${backend}-ci"
-rm -rf build *.egg-info dist
-pip uninstall -y dgl
-# test install
-python3 setup.py install
-# test inplace build (for cython)
-python3 setup.py build_ext --inplace
-done
+if [[ $1 == "cugraph" ]]; then
+    rm -rf build *.egg-info dist
+    pip uninstall -y dgl
+    # test install
+    python3 setup.py install
+    # test inplace build (for cython)
+    python3 setup.py build_ext --inplace
+else
+    for backend in pytorch mxnet tensorflow
+    do
+    conda activate "${backend}-ci"
+    rm -rf build *.egg-info dist
+    pip uninstall -y dgl
+    # test install
+    python3 setup.py install
+    # test inplace build (for cython)
+    python3 setup.py build_ext --inplace
+    done
+fi
 popd