Make Device.set_current() faster

[leofang]

Description

closes#739

Before this PR:

In [4]: %timeit dev.set_current()
1.4 μs ± 10.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

With this PR:

In [4]: %timeitdev.set_current()
374ns ± 1.17nsperloop (mean ± std. dev. of7runs, 1,000,000loopseach)

We want to minimize the number of calls to CUDA APIs, which add up to the overheads (several hundreds of nanoseconds per call). This is achieved by combining two changes:

1. We lazily retain a reference to the primary context of each device
   • The only scenario where this could fail is if someone calls cudaDeviceReset() somewhere, in which case the pointers to the primary contests would be invalidated and become dangling. However, the CUDA team strongly discourages using this API (it does not solve any real issue, e.g. restoring from sticky errors), and it is impossible to get things right especially in the Python land if it is called. Too many CUDA-related objects floating around.
2. We unconditionally set the needed primary context to current without extra checks

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[leofang]

/ok to test 92d24ce

[leofang]

/ok to test b9b3c87

[github-actions]

Doc Preview CI
Preview removed because the pull request was closed or merged.