Fix GPU handle pool singleton aliasing with tag-based template separation

Problem:
- BLAS and SOLVER handle pools were sharing the same static singleton instance
- C++ template instantiation created only one pool when HandleType/StreamType were identical
- This caused resource contamination between hipBLAS and hipSOLVER operations
- Led to potential memory corruption and unexpected GPU library behavior

Solution:
- Added Tag template parameter to HandlePool class with default DefaultTag
- Introduced BlasTag and SolverTag for unique pool instantiations
- Updated all handle pool typedefs to use three-argument template
- Each tag type now gets its own static singleton pool instance

Author: phambinhfin

jaxlib/gpu/blas_handle_pool.cc

@@ -41,4 +41,4 @@ template <>
   return Handle(pool, handle, stream);
 }
 
-}  // namespace jax
+}  // namespace jax

jaxlib/gpu/blas_handle_pool.h

@@ -22,7 +22,7 @@ limitations under the License.
 
 namespace jax {
 
-using BlasHandlePool = HandlePool<gpublasHandle_t, gpuStream_t>;
+using BlasHandlePool = HandlePool<gpublasHandle_t, gpuStream_t, BlasTag>;
 
 template <>
 absl::StatusOr<BlasHandlePool::Handle> BlasHandlePool::Borrow(

jaxlib/gpu/handle_pool.h

@@ -25,9 +25,20 @@ limitations under the License.
 
 namespace jax {
 
+// Tag types for unique pool instantiations
+struct DefaultTag {};
+struct BlasTag {};
+struct SolverTag {};
+
 // To avoid creating cublas/cusolver contexts in the middle of execution, we
 // maintain a pool of them.
-template <typename HandleType, typename StreamType>
+
+// The Tag template parameter ensures unique pool instantiations for different
+// handle types (BLAS, SOLVER, etc.). Without this tag, C++ template 
+// instantiation would create a single shared static pool when HandleType and 
+// StreamType are the same, leading to resource contamination between different
+// GPU library contexts (e.g., hipBLAS and hipSOLVER sharing the same pool).
+template <typename HandleType, typename StreamType, typename Tag = DefaultTag>
 class HandlePool {
  public:
   HandlePool() = default;
@@ -66,11 +77,11 @@ class HandlePool {
     HandleType get() { return handle_; }
 
    private:
-    friend class HandlePool<HandleType, StreamType>;
-    Handle(HandlePool<HandleType, StreamType>* pool, HandleType handle,
+    friend class HandlePool<HandleType, StreamType, Tag>;
+    Handle(HandlePool<HandleType, StreamType, Tag>* pool, HandleType handle,
            StreamType stream)
         : pool_(pool), handle_(handle), stream_(stream) {}
-    HandlePool<HandleType, StreamType>* pool_ = nullptr;
+    HandlePool<HandleType, StreamType, Tag>* pool_ = nullptr;
     HandleType handle_ = nullptr;
     StreamType stream_ = nullptr;
   };
@@ -80,31 +91,32 @@ class HandlePool {
   static absl::StatusOr<Handle> Borrow(StreamType stream);
 
  private:
-  static HandlePool<HandleType, StreamType>* Instance();
+  static HandlePool<HandleType, StreamType, Tag>* Instance();
 
   void Return(HandleType handle, StreamType stream);
 
   absl::Mutex mu_;
   std::map<StreamType, std::vector<HandleType>> handles_ ABSL_GUARDED_BY(mu_);
 };
 
-template <typename HandleType, typename StreamType>
-/*static*/ HandlePool<HandleType, StreamType>*
-HandlePool<HandleType, StreamType>::Instance() {
-  static auto* pool = new HandlePool<HandleType, StreamType>;
+template <typename HandleType, typename StreamType, typename Tag>
+/*static*/ HandlePool<HandleType, StreamType, Tag>*
+HandlePool<HandleType, StreamType, Tag>::Instance() {
+  static auto* pool = new HandlePool<HandleType, StreamType, Tag>;
   return pool;
 }
 
-template <typename HandleType, typename StreamType>
-void HandlePool<HandleType, StreamType>::Return(HandleType handle,
-                                                StreamType stream) {
+template <typename HandleType, typename StreamType, typename Tag>
+void HandlePool<HandleType, StreamType, Tag>::Return(HandleType handle,
+                                                    StreamType stream) {
   absl::MutexLock lock(&mu_);
   handles_[stream].push_back(handle);
 }
 
 // template <typename HandleType, typename StreamType>
 // HandlePool<HandleType, StreamType>::Borrow(StreamType stream)
 
+
 }  // namespace jax
 
 #endif  // JAXLIB_GPU_HANDLE_POOL_H_

jaxlib/gpu/solver_handle_pool.cc

@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include "jaxlib/gpu/solver_handle_pool.h"
-
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "jaxlib/gpu/gpu_kernel_helpers.h"
@@ -40,7 +39,7 @@ template <>
     pool->handles_[stream].pop_back();
   }
   if (stream) {
-    JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnSetStream(handle, stream)));
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnSetStream(handle, stream)));
   }
   return Handle(pool, handle, stream);
 }

jaxlib/gpu/solver_handle_pool.h

@@ -26,14 +26,14 @@ limitations under the License.
 
 namespace jax {
 
-using SolverHandlePool = HandlePool<gpusolverDnHandle_t, gpuStream_t>;
+using SolverHandlePool = HandlePool<gpusolverDnHandle_t, gpuStream_t, SolverTag>;
 
 template <>
 absl::StatusOr<SolverHandlePool::Handle> SolverHandlePool::Borrow(
     gpuStream_t stream);
 
 #ifdef JAX_GPU_CUDA
-using SpSolverHandlePool = HandlePool<cusolverSpHandle_t, gpuStream_t>;
+using SpSolverHandlePool = HandlePool<cusolverSpHandle_t, gpuStream_t, SolverTag>;
 
 template <>
 absl::StatusOr<SpSolverHandlePool::Handle> SpSolverHandlePool::Borrow(