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

@@ -24,9 +24,9 @@ limitations under the License.
 namespace jax {
 
 template <>
-/*static*/ absl::StatusOr<BlasHandlePool::Handle> BlasHandlePool::Borrow(
-    gpuStream_t stream) {
-  BlasHandlePool* pool = Instance();
+/*static*/ absl::StatusOr<HandlePool<gpublasHandle_t, gpuStream_t, BlasTag>::Handle>
+HandlePool<gpublasHandle_t, gpuStream_t, BlasTag>::Borrow(gpuStream_t stream) {
+  auto* pool = Instance(HandleKind::BLAS);
   absl::MutexLock lock(&pool->mu_);
   gpublasHandle_t handle;
   if (pool->handles_[stream].empty()) {
@@ -36,7 +36,11 @@ template <>
     pool->handles_[stream].pop_back();
   }
   if (stream) {
-    JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpublasSetStream(handle, stream)));
+    if (pool->kind_ == HandleKind::BLAS) {
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpublasSetStream(handle, stream)));
+    } else {
+      return absl::InternalError("BlasHandlePool kind is not BLAS");
+    }
   }
   return Handle(pool, handle, stream);
 }

jaxlib/gpu/blas_handle_pool.h

@@ -1,3 +1,6 @@
+
+// Tag type for BLAS pool
+struct BlasTag {};
 /* Copyright 2024 The JAX Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,7 +25,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

@@ -18,17 +18,38 @@ limitations under the License.
 
 #include <map>
 #include <vector>
+#include <iostream>
+#include <typeinfo>
 
 #include "absl/base/thread_annotations.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 
 namespace jax {
+// Enum to distinguish between pool types
+enum class HandleKind {
+  UNKNOWN,
+  BLAS,
+  SOLVER
+};
+
+// 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 {
+  HandleKind kind() const { return kind_; }
+  void set_kind(HandleKind kind) { kind_ = kind; }
  public:
   HandlePool() = default;
 
@@ -66,11 +87,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 +101,40 @@ class HandlePool {
   static absl::StatusOr<Handle> Borrow(StreamType stream);
 
  private:
-  static HandlePool<HandleType, StreamType>* Instance();
+  static HandlePool<HandleType, StreamType, Tag>* Instance(HandleKind kind = HandleKind::UNKNOWN);
 
   void Return(HandleType handle, StreamType stream);
 
   absl::Mutex mu_;
   std::map<StreamType, std::vector<HandleType>> handles_ ABSL_GUARDED_BY(mu_);
+  HandleKind kind_ = HandleKind::UNKNOWN;
 };
 
-template <typename HandleType, typename StreamType>
-/*static*/ HandlePool<HandleType, StreamType>*
-HandlePool<HandleType, StreamType>::Instance() {
-  static auto* pool = new HandlePool<HandleType, StreamType>;
-  return pool;
+template <typename HandleType, typename StreamType, typename Tag>
+/*static*/ HandlePool<HandleType, StreamType, Tag>*
+HandlePool<HandleType, StreamType, Tag>::Instance(HandleKind kind) {
+  static std::map<HandleKind, HandlePool<HandleType, StreamType, Tag>*> pools;
+  auto it = pools.find(kind);
+  if (it == pools.end()) {
+    auto* pool = new HandlePool<HandleType, StreamType, Tag>;
+    pool->set_kind(kind);
+    pools[kind] = pool;
+    return pool;
+  }
+  return it->second;
 }
 
-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"
@@ -30,7 +29,7 @@ namespace jax {
 template <>
 /*static*/ absl::StatusOr<SolverHandlePool::Handle> SolverHandlePool::Borrow(
     gpuStream_t stream) {
-  SolverHandlePool* pool = Instance();
+  SolverHandlePool* pool = Instance(HandleKind::SOLVER);
   absl::MutexLock lock(&pool->mu_);
   gpusolverDnHandle_t handle;
   if (pool->handles_[stream].empty()) {
@@ -40,7 +39,11 @@ template <>
     pool->handles_[stream].pop_back();
   }
   if (stream) {
-    JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnSetStream(handle, stream)));
+    if (pool->kind_ == HandleKind::SOLVER) {
+      JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpusolverDnSetStream(handle, stream)));
+    } else {
+      return absl::InternalError("SolverHandlePool kind is not SOLVER");
+    }
   }
   return Handle(pool, handle, stream);
 }

jaxlib/gpu/solver_handle_pool.h

@@ -26,7 +26,7 @@ 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(