Skip to content

Commit f07b6d6

Browse files
phambinhfinzahiqbal
authored andcommitted
Fix GPU handle pool singleton aliasing (#525)
Addresses handle pool singleton sharing issue between different GPU operation types in ROCm/HIP backend. (cherry picked from commit 4b5b368)
1 parent c641df3 commit f07b6d6

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

jaxlib/gpu/vendor.h

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,10 @@ typedef hipDoubleComplex gpuDoubleComplex;
450450
typedef hipComplex gpublasComplex;
451451
typedef hipDoubleComplex gpublasDoubleComplex;
452452

453-
typedef hipsolverHandle_t gpusolverDnHandle_t;
453+
// Create unique opaque pointer types for proper singleton separation - BLAS and SOLVER only
454+
typedef struct hipblasHandle_* gpublasHandle_t;
455+
typedef struct hipsolverHandle_* gpusolverDnHandle_t;
456+
454457
typedef hipblasFillMode_t gpublasFillMode_t;
455458
typedef hipsolverFillMode_t gpusolverFillMode_t;
456459
typedef struct hipblasHandle_* gpublasHandle_t;
@@ -502,6 +505,7 @@ typedef hipsparseDnVecDescr_t gpusparseDnVecDescr_t;
502505
// Wrapper functions for BLAS handles to ensure unique types
503506
#define gpublasCreate(handle) hipblasCreate(reinterpret_cast<hipblasHandle_t*>(handle))
504507
#define gpublasSetStream hipblasSetStream
508+
505509
#define gpublasSgeqrfBatched hipblasSgeqrfBatched
506510
#define gpublasDgeqrfBatched hipblasDgeqrfBatched
507511
#define gpublasCgeqrfBatched hipblasCgeqrfBatched
@@ -554,6 +558,7 @@ typedef hipsparseDnVecDescr_t gpusparseDnVecDescr_t;
554558
// Wrapper functions for SOLVER handles to ensure unique types
555559
#define gpusolverDnCreate(handle) hipsolverCreate(reinterpret_cast<hipsolverHandle_t*>(handle))
556560
#define gpusolverDnSetStream hipsolverSetStream
561+
557562
#define gpusolverDnCreateSyevjInfo hipsolverCreateSyevjInfo
558563
#define gpusolverDnDestroySyevjInfo hipsolverDestroySyevjInfo
559564
#define gpusolverDnSgeqrf hipsolverSgeqrf

0 commit comments

Comments
 (0)