[Sampling] New sampling pipeline plus asynchronous prefetching (#3665)

* initial update

* more

* more

* multi-gpu example

* cluster gcn, finalize homogeneous

* more explanation

* fix

* bunch of fixes

* fix

* RGAT example and more fixes

* shadow-gnn sampler and some changes in unit test

* fix

* wth

* more fixes

* remove shadow+node/edge dataloader tests for possible ux changes

* lints

* add legacy dataloading import just in case

* fix

* update pylint for f-strings

* fix

* lint

* lint

* lint again

* cherry-picking commit fa9f494

* oops

* fix

* add sample_neighbors in dist_graph

* fix

* lint

* fix

* fix

* fix

* fix tutorial

* fix

* fix

* fix

* fix warning

* remove debug

* add get_foo_storage apis

* lint

python/dgl/sampling/__init__.py

@@ -10,3 +10,4 @@ from .pinsage import *
 from .neighbor import *
 from .node2vec_randomwalk import *
 from .negative import *
+from . import utils

python/dgl/sampling/negative.py

@@ -3,6 +3,8 @@
 from numpy.polynomial import polynomial
 from .._ffi.function import _init_api
 from .. import backend as F
+from .. import utils
+from ..heterograph import DGLHeteroGraph
 
 __all__ = [
     'global_uniform_negative_sampling']
@@ -99,5 +101,7 @@ def global_uniform_negative_sampling(
     src, dst = _CAPI_DGLGlobalUniformNegativeSampling(
         g._graph, etype_id, num_samples, 3, exclude_self_loops, replace, redundancy)
     return F.from_dgl_nd(src), F.from_dgl_nd(dst)
+DGLHeteroGraph.global_uniform_negative_sampling = utils.alias_func(
+    global_uniform_negative_sampling)
 
 _init_api('dgl.sampling.negative', __name__)

python/dgl/storages/__init__.py

+"""Feature storage classes for DataLoading"""
+from .. import backend as F
+
+from .base import *
+from .numpy import *
+if F.get_preferred_backend() == 'pytorch':
+    from .pytorch_tensor import *
+else:
+    from .tensor import *

python/dgl/utils/__init__.py

@@ -4,3 +4,4 @@ from .data import *
 from .checks import *
 from .shared_mem import *
 from .filter import *
+from .exception import *

tests/compute/test_subgraph.py

@@ -594,3 +594,6 @@ def test_khop_out_subgraph(idtype):
     assert edge_set == {(0, 1)}
     assert F.array_equal(F.astype(inv['user'], idtype), F.tensor([0], idtype))
     assert F.array_equal(F.astype(inv['game'], idtype), F.tensor([0], idtype))
+
+if __name__ == '__main__':
+    test_khop_out_subgraph(F.int64)

tests/compute/test_transform.py

@@ -2308,3 +2308,4 @@ def test_module_add_edge(idtype):
 
 if __name__ == '__main__':
     test_partition_with_halo()
+    test_module_heat_kernel(F.int32)