[Performance] Redirect AllocWorkspace to PyTorch's allocator if available

[yaox12]

Description

Related issues: #3933 #3957.

Redirect AllocWorkspace/FreeWorkspace to PyTorch's allocator via raw_alloc and raw_delete.

I run examples/pytorch/graphsage/node_classification.py and get the GPU memory footprints as:

	nvidia-smi	max_allocated	allocated	max_reserved	reserved
new allocator + pure_gpu	10629	9944	5542	9971	9971
old allocator + pure_gpu	10645	7930	5243	7958	7958
new allocator + uva	2531	550	266	1480	1480
old allocator + uva	2591	550	265	1241	1241

*The four columns on the right are reported by torch.cuda.max_allocated/allocated/max_reserved/reserved.

The total GPU memory footprints are close. Advantages are:

1. Users can release the reserved GPU memory via PyTorch's APIs if they'd like to.
2. When PyTorch reports OOM, users won't see a big discrepancy between the memory used by PyTorch and the GPUs capacity.

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

[dgl-bot]

Commit ID: c07819c

Build ID: 4

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[yaox12]

More backgrounds on PyTorch's CUDA allocator:

1. Call stack of relevant objects: CudaCachingAllocator device_allocator → THCCachingAllocator caching_allocator → std::vector<std::unique_ptr<DeviceCachingAllocator>> device_allocator

2. There’re mainly two ways in PyTorch to allocate CUDA memory:

   1. c10::cuda::CUDACachingAllocator::get()->allocate() [code]. It uses CudaCachingAllocator device_allocator and has a life cycle management (will be freed automatically). Inside, it calls THCCachingAllocator caching_allocator for the actual allocation. [code]
   2. c10::cuda::CUDACachingAllocator::raw_alloc()/raw_delete() [code]. It uses THCCachingAllocator caching_allocator directly and requires being freed via raw_delete manually.

cc @nv-dlasalle

[dgl-bot]

Commit ID: d210032

Build ID: 5

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: e2ec146

Build ID: 6

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[shintarok111]

I encountered the same issue as in #3933. Has the problem been resolved?