[RFC] Faster CPU sampling through fused sampling+compaction

## 🚀 Feature

This issue is for discussing fused sampling kernels that sample and compact the sampled graph in one step, and directly return a block in CSC(compressed sparse column) format. This completely bypasses the expensive ``to_block`` function. Experimental implementation is [here](https://github.com/hesham-mostafa/dgl/tree/sampling_optim)

## Motivation

In sampling-based GNN training on CPUs, we observe that graph sampling takes up a significant portion of the epoch time. Graph sampling often ends up taking more time than the actual GNN forward+backward time. The current sampling pipeline in DGL is illustrated in the figure below:


In the current sampling pipeline, creating a graph sample that can be passed to a GNN is a 2-step process:
1. ``dgl.sampling.sample_neighbors`` or similar functions takes the seed nodes and sample some neighbors of the seed nodes. They return a copy of the original graph in COO(coordinate) format that only contains the sampled edges.
2. ``dgl.to_block`` compacts the sampled graph to create a bipartite graph. ``dgl.to_block`` reindexes the nodes so that the indices of the source nodes and the target nodes are consecutive and zero-based. It is common to include the target nodes at the beginning of the target nodes (``include_dst_in_src = True``).

``dgl.to_block`` takes up the majority of the sampling time as it involves expensive hashing operations that are needed to compact the graph. Compaction is the process of taking the original node indices and renumbering them so that they are contiguous and start from zero. We argue that several of the steps needed to compact the sampled graph can be very cheaply implemented if they are integrated with the actual graph sampling steps. The graph sampling step would take longer, but the extra time is much less than the time needed by the ``dgl.to_block`` function.

## Alternatives

There is currently no alternatives in DGL to the 2-step sampling process outlined above.

## Pitch

Implement C++ functions that sample and compact the sampled graph on the fly, and directly returns a compacted bipartite graph as a CSR matrix. This graph can be directly passed to a GNN, bypassing the ``dgl.to_block`` function.

## Additional context

The actual sampling of the graph neighbors is mostly done using functions such as ``CSRRowWisePick`` defined [here](https://github.com/dmlc/dgl/blob/master/src/array/cpu/rowwise_pick.h). These functions return COO matrices that contain the uncompacted topology of the sampled graph using the original node indices. We propose to define a parallel set of functions that sample the graph, and compact it on the fly. These functions will return compacted graphs in the form of CSR matrices that can be directly passed to GNNs. This bypasses the ``dgl.to_block`` function. The reasons we want to fuse compaction to sampling are:
1. Compacting the target/seed nodes is trivial during sampling. We assume the seed nodes are unique, and the ith seed node is simply mapped to target node i.
2. We have immediate access to the number of source nodes for each target node (the fan-in), as we get this directly from the ``NumPicksFn`` before we start sampling. We can thus construct the ``indptr`` component of our CSR matrix practically for free.

We still have to compact the source nodes. In our current implementation, for compacting the source nodes, we use a naive O(N) hashing where N is the number of nodes in the graph. It is very fast but takes up more memory than more elaborate techniques. A sample implementation of a fused ``CSRRowWisePick`` is [here](https://github.com/hesham-mostafa/dgl/blob/sampling_optim/src/array/cpu/rowwise_pick.h#L535).

## API and usage scenarios

We propose to add a new flag to ``dgl.sampling.sample_neighbors``. If this flag is set, ``sample_neighbors`` will dispatch to a C++ function that does combined sampling and compaction. The flag will be ``False`` by default.

The current ``dgl.dataloading.NeighborSampler`` runs the 2-step process of sampling the graph, followed by compaction using ``dgl.to_block``. We propose to add a new sampler ``dgl.dataloading.NeighborSamplerFused`` that will call the ``dgl.sampling.sample_neighbors`` so that it dispatches to the combined sampling and compaction kernels. ``dgl.dataloading.NeighborSamplerFused`` thus does not call ``dgl.to_block``.

## Experimental results

We implemented the fused sampling+compaction [here](https://github.com/hesham-mostafa/dgl/tree/sampling_optim). We measured performance on a 36-core Icelake processor with 256GB of RAM. We downloaded the stochastic training of GNNs tutorial code from [here](https://docs.dgl.ai/tutorials/large/L1_large_node_classification.html). We modified it by adding these lines:

```python

fused = True #Enable fused sampling+compaction
if fused:
sampler = dgl.dataloading.NeighborSamplerFused([15, 10, 5])
else:
sampler = dgl.dataloading.NeighborSampler([15, 10, 5])

```

We tested on ``ogbn-arxiv``, ``ogbn-products``, and ``ogbn-papers100M``. We used a 3-layer GraphSage model with hidden layer size of 128, batch size of 1024, and sampling fanouts of [15,10,5]. We tested with 0,2, and 4 DataLoader workers and enabled the dataloader's CPU affinity for dataloader workers larger than 0. The 'sampling time' is the time spent waiting on the dataloader. This will only accurately reflect the time needed to sample and compact the graph when the number of workers is zero. The 'gnn time' is the time spent in the GNN's forward+backward pass. We noticed that the 'gnn time' when using the fused sampling+compaction kernels is typically lower than when using DGL's vanilla sampling. This could be due to the fact that our kernels return a graph in CSC format directly, while ``dgl.to_block`` returns in COO format that will need to be converted to CSC in the forward pass.




The results indicate that the fused sampling+compaction kernel runs significantly faster than the vanilla sampling approach that uses ``dgl.to_block`` (in some cases up to ~2.5X faster). As we increase dataloader workers, the parallelized sampling gradually hides the sampling latency, but the fused sampling+compaction kernels still lead to significantly faster epoch times.