[Feature] Enable UVA sampling with CPU indices

[BarclayII]

Description

This PR enables specifying indices on CPU for UVA sampling to avoid duplicating the indices for every GPU.

It also adds a utility function dgl.multiprocessing.call_once_and_share() which calls a function in a single process and share the result to other processes. Requires using dgl.multiprocessing.spawn instead of torch.multiprocessing.spawn (the signatures are the same).

Fixes #3855 and #3893 .

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change,
  or have been fixed to be compatible with this change
• Related issue is referred in this PR

[dgl-bot]

To trigger regression tests:

• @dgl-bot run [instance-type] [which tests] [compare-with-branch];
  For example: @dgl-bot run g4dn.4xlarge all dmlc/master or @dgl-bot run c5.9xlarge kernel,api dmlc/master

[jermainewang]

I didn't follow the context very carefully but saw no change on user experience. I'm good with the PR. Will let @nv-dlasalle approve.

[yaox12]

Since @nv-dlasalle is out until next week, I'd like to approve this PR.

[kkranen]

Has anyone taken a look at how using this feature increases load on host memory? While this is an awesome solution for enabling larger graphs, it may cause problems from some memory BW-limited solutions.

[kkranen]

It seems that dgl.multiprocessing.spawn does not exist in main:
module 'dgl.multiprocessing' has no attribute 'spawn'
Should we be executing spawning by setting multiprocessing context?

[TristonC]

@BarclayII Could you help to answer Kyle's question?

[BarclayII]

Ah the docstring was wrong. It no longer requires dgl.multiprocessing.spawn. I can fix that.
Although for question from @kkranen, you need to spawn the process with torch.multiprocessing.spawn and set up a PyTorch distributed process group (torch.distributed.init_process_group).

[BarclayII]

Has anyone taken a look at how using this feature increases load on host memory? While this is an awesome solution for enabling larger graphs, it may cause problems from some memory BW-limited solutions.

I think the load increase should be fine since this PR only changes how we handle CPU indices, and we only require copying CPU indices (which has the same size of the minibatch itself) to GPU.