Why does JAX bind a handle to each stream when calling BLAS?

[south-ocean]

Hello! I noticed that JAX binds a handle to each stream when calling BLAS, while in the Torch framework, a handle is shared across each device. If a handle is bound to each stream, does that increase the overhead of handle creation and management? Would it be better for JAX to modify its approach to share a handle for each device instead? Thank you!

[dfm]

Good question! I don't expect that this will have any major performance implications but I'd be interested if you found a case where you think it would. I'm not totally clear from the CUDA solver docs about the implications of using the same handle with multiple streams, but it sounds like it might be ok to reuse a handle as long as we call *SetStream right before every time we use it. I probably wouldn't list this as a high priority change until we see a case where this is performance critical. What do you think?

[south-ocean]

Yeah, we have now found a case:
because we have done some memory use on the BLAS handle in exchange for time performance, which causes the handle to be heavier, And we have done tests on torch and jax respectively, and found that there is a performance improvement on torch, because torch is bound to a handle for each device, and the performance improvement in jax is offset by the handle created many times. We observed that the same example was that torch would only create one handle, while jax would create multiple handles. So the way the JAX uses a per-stream binding handle has resulted in some handle-based optimizations not working. What do you think of the case?
I don't know much about cublas, I've seen the code related to hipblas, combined with the code in torch, there is always a performance loss in creating multiple handles and management, is it better to share a device for a handle than to bind a handle to a stream? It can also make Jax more adaptable in this aspect.

[dfm]

I must say that I don't find the comparison to Pytorch very compelling in the abstract like this since there are so many factors that can affect performance comparisons like this. But, if you can share an explicit, self-contained, JAX program that you expect to have significantly improved performance after such a change, I'd be very happy to dig into what's happening with such an example!