[Performance] Enable using pinned memory for transfers in SparseAdam optimizer.

[nv-dlasalle]

Description

In python 1.8, when non_blocking=True in a GPU->CPU copy, the output tensor is allocated in pinned memory (the cause of #2760). However, this now results in a performance regression between 0.6 and 0.7 when using torch>=1.8.

This PR re-enables using pinned memory for those transfers, while synchronizing afterwards to ensure correctness. This cuts time in the optimizer dramatically (on my system nearly 2x--but it will vary from system to system):

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 my best 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

Changes

Used cuda events to allow specifying non_blocking=True when transferring from the gradient computation device (GPU) to the state storage device (CPU).

This also re-arranges some operations in order to better in order to reduce the amount of time the GPU needs to wait on the CPU to finish slicing.

Remove variables used for setting non_blocking=False, as this is the default.

[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

[classicsong]

LGTM