Thanks for looking at this!

While I understand that multi-GPU models are necessary for some workloads, it's not trivial to make sure that we come out ahead of a single GPU in terms of performance. And in the meantime changes like this can add overhead to both the code and execution that impacts the single-GPU case, which is what almost everyone will be using.

We therefore want to think carefully and do our own experiments before we take on this topic. We'd likely target data-parallel multi-GPU workloads rather than model parallelism, as that's the most common requirement people will have for the main usage of Thinc, which is models for spaCy.

Unfortunately that means this isn't really a good target for user contribution. It's a change that will have significant impact for the library, so we can only get to this when we're really ready to work on it.

I expect the performance impact to come primarily from python function call overhead when encountering context __enter__ and __exit__ calls. Those happen once for each ops primitive.

For CPU, the context itself does nothing. For GPU, both methods end up calling cudaSetDevice. It is fairly common to see codes that call this on every blas operation (e.g. page 2 of SLATE port to AMD and Intel platforms). From NVIDIA's docs, it is described as:

This function will do no synchronization with the previous or new device, and should be considered a very low overhead call.

The HIP docs go further,

This function does no synchronization with the previous or new device, and has very little runtime overhead. Applications can use hipSetDevice to quickly switch the default device before making a HIP runtime call which uses the default device.
The default device is stored in thread-local-storage for each thread. Thread-pool implementations may inherit the default device of the previous thread. A good practice is to always call hipSetDevice at the start of HIP coding sequency to establish a known standard device.

Note also that the change to CupyOps.asarray has no overhead, since both cupy.asarray and numpy.asarray state that no copy is performed if the array already matches the call's request.

Considering @honnibal's earlier reply, I'll go ahead and close this particular PR, but we'll add the idea to our backlog of things to explore further. Thanks again!