cuInit before querying compute capability

olupton · olupton · commit c97d95577146 · 2024-04-04T15:27:57.000Z
diff --git a/jaxlib/cuda/versions_helpers.cc b/jaxlib/cuda/versions_helpers.cc
@@ -86,6 +86,7 @@ size_t CudnnGetVersion() {
 }
 int CudaComputeCapability(int device) {
   int major, minor;
+  JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpuInit(0)));
   JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpuDeviceGetAttribute(
       &major, GPU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device)));
   JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpuDeviceGetAttribute(
@@ -102,4 +103,4 @@ int CudaDeviceCount() {
 }
 
 
-}  // namespace jax::cuda
+}  // namespace jax::cuda

Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,7 @@ size_t CudnnGetVersion() {`
`86`	`86`	`}`
`87`	`87`	`int CudaComputeCapability(int device) {`
`88`	`88`	`int major, minor;`
	`89`	`+ JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpuInit(0)));`
`89`	`90`	`JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpuDeviceGetAttribute(`
`90`	`91`	`&major, GPU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device)));`
`91`	`92`	`JAX_THROW_IF_ERROR(JAX_AS_STATUS(gpuDeviceGetAttribute(`
`@@ -102,4 +103,4 @@ int CudaDeviceCount() {`
`102`	`103`	`}`
`103`	`104`
`104`	`105`
`105`		`-} // namespace jax::cuda`
	`106`	`+} // namespace jax::cuda`