[Feature] Import PyTorch's pin_memory() method for DGL graph structure

[chang-l]

Description

This PR imports pytorch's CachingHostAllocator [ref] and its associated functions (like CachingHostAllocator_recordEvent and CachingHostAllocator_emptyCache) to DGL's tensoradaptor. It enhances/enriches DGL's native memory pinning strategy (cudaHostRegister), by using backend's pinned memory pool (viacudaHostAlloc) to avoid overhead from cudaHostRegister and frequent memory allocations/deallocations.

While cudaHostRegister (DGLGraph.pin_memory_) is more suitable to pin large data (like entire feature/graph) ONCE for future frequent GPU accesses (no memory allocation involved but with the compromising of speed), pytorch's pinned memory pool (DGLGraph(Index).pin_memory) is better for objects (small memory blocks) that could be created/destroyed often (like DGLGraph/feature for batches).

Based on this PR, it would be easier to implement pin_memory() method for DGLGraph object, and let PyTorch's dataloader pick it up when pin_memory=True to enable async H2D copying of (sub)graph structure (or feature if needed). Complete dataloader support should be included in a separate PR.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• I've leverage the tools to beautify the python and c++ code.
• The PR is complete and small, read the Google eng practice (CL equals to PR) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
• 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
• If the PR is for a new model/paper, I've updated the example index here.

Changes

C++ (DGL core):

• Imported getCachingHostAllocator, CachingHostAllocator_recordEvent and CachingHostAllocator_emptyCache to completely support pytorch's cachingHostAllocator.
• Added/propagated pin_memory* methods to all necessary classes(ndarry/csr/coo/unitgraph/heterograph).
• Added RecordedCopyDataFromTo to fuse data copy (H2D) and following CachingHostAllocator_recordEvent to manage the (pinned)resource consistently withbackend PyTorch.
• Empty the caching pool if calling cudaHostRegister to release resources/prepare enough page-locked memory needed by cudaHostRegister

Python:

• Added outplace pin_memory() for heterograph_index
• Unit test

To-dos

(maybe need separate PRs)

Python:

• Add outplace pin_memory() for heterograph or other equivalents to integrate with dataloader.

[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

[chang-l]

@jermainewang Pls help add other reviewers. Thanks.

[frozenbugs]

Hi Chang, sorry being picky on the code quality. Overall the PR looks very good to me, just some small tweak to make it shiny.

[chang-l]

No worries @frozenbugs I am always happy to iterate and improve in every aspect.

Thank you for taking the time and review this PR. I will address them all tomorrow.

[yaox12]

LGTM except for some comments issues.

[dgl-bot]

Commit ID: 2c6f7b087423f22420c7544446d7d8bfe8b37fb0

Build ID: 19

Status: ⚪️ CI test cancelled due to overrun.

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 4616369cecd36909b298b466cd31e1a38be2cc6e

Build ID: 20

Status: ❌ CI test failed in Stage [Tensorflow CPU Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: cf08a1c

Build ID: 21

Status: ⚪️ CI test cancelled due to overrun.

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 32dcb35

Build ID: 22

Status: ❌ CI test failed in Stage [Lint Check].

Report path: link

Full logs path: link

[chang-l]

Sorry for the delay. Ready for the next round @frozenbugs @BarclayII @yaox12

[dgl-bot]

Commit ID: ebcfa02

Build ID: 23

Status: ❌ CI test failed in Stage [Lint Check].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 9908b4b55d09ef87b36781b11112b70f026e2605

Build ID: 24

Status: ⚪️ CI test cancelled due to overrun.

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 1407e1912b7b06fb15dabdfe143b3bc2a6d0bb42

Build ID: 25

Status: ❌ CI test failed in Stage [Torch CPU (Win64) Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: ced69b9

Build ID: 26

Status: ✅ CI test succeeded.

Report path: link

Full logs path: link

[frozenbugs]

LGTM, overall this is a giant PR which we should try to avoid if we can.

[dgl-bot]

Commit ID: b6090eb

Build ID: 27

Status: ✅ CI test succeeded.

Report path: link

Full logs path: link

[BarclayII]

I'm good. Thanks a lot!