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]: %timeit dev.set_current()
374 ns ± 1.17 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

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.