feat: upgrade ceno-gpu

Cargo.toml

@@ -82,6 +82,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 uint = "0.8"
 
 ceno_gpu = { path = "utils/cuda_hal", package = "cuda_hal" }
+cudarc = { version = "0.17.3", features = ["driver", "cuda-version-from-build-system"] }
 
 [profile.dev]
 lto = "thin"

ceno_zkvm/Cargo.toml

@@ -15,7 +15,7 @@ ceno-examples = { path = "../examples-builder" }
 ceno_emul = { path = "../ceno_emul" }
 ceno_gpu = { workspace = true, optional = true }
 ceno_host = { path = "../ceno_host" }
-cudarc = { version = "0.13.0", features = ["driver", "cuda-version-from-build-system"], optional = true }
+cudarc = { workspace = true, optional = true }
 either.workspace = true
 ff_ext.workspace = true
 gkr_iop = { path = "../gkr_iop" }

gkr_iop/Cargo.toml

@@ -12,7 +12,7 @@ version.workspace = true
 [dependencies]
 bincode.workspace = true
 ceno_gpu = { workspace = true, optional = true }
-cudarc = { version = "0.13.0", features = ["driver", "cuda-version-from-build-system"], optional = true }
+cudarc = { workspace = true, optional = true }
 either.workspace = true
 ff_ext.workspace = true
 itertools.workspace = true

gkr_iop/src/gpu/mod.rs

@@ -27,44 +27,32 @@ pub mod gpu_prover {
         common::{
             basefold::utils::convert_ceno_to_gpu_basefold_commitment,
             buffer::BufferImpl,
+            get_ceno_gpu_device_id,
             mle::{
                 build_mle_as_ceno, ordered_sparse32_selector_gpu, rotation_next_base_mle_gpu,
                 rotation_selector_gpu,
             },
         },
     };
-    use cudarc::driver::{CudaDevice, DriverError};
+
     use once_cell::sync::Lazy;
     use std::sync::{Arc, Mutex, MutexGuard};
 
     pub type BB31Base = p3::babybear::BabyBear;
     pub type BB31Ext = ff_ext::BabyBearExt4;
 
-    pub static CUDA_DEVICE: Lazy<Result<Arc<CudaDevice>, DriverError>> =
-        Lazy::new(|| CudaDevice::new(0));
-
     #[allow(clippy::type_complexity)]
     pub static CUDA_HAL: Lazy<
         Result<Arc<Mutex<CudaHalBB31>>, Box<dyn std::error::Error + Send + Sync>>,
     > = Lazy::new(|| {
-        let device = CUDA_DEVICE
-            .as_ref()
-            .map_err(|e| format!("Device init failed: {:?}", e))?;
-        device.bind_to_thread()?;
-
-        CudaHalBB31::new()
+        // can be overridden by env variable `CENO_GPU_DEVICE_ID`
+        let device_id: usize = get_ceno_gpu_device_id(0);
+        CudaHalBB31::new(device_id)
             .map(|hal| Arc::new(Mutex::new(hal)))
             .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
     });
 
     pub fn get_cuda_hal() -> Result<MutexGuard<'static, CudaHalBB31>, String> {
-        let device = CUDA_DEVICE
-            .as_ref()
-            .map_err(|e| format!("Device not available: {:?}", e))?;
-        device
-            .bind_to_thread()
-            .map_err(|e| format!("Failed to bind device to thread: {:?}", e))?;
-
         let hal_arc = CUDA_HAL
             .as_ref()
             .map_err(|e| format!("HAL not available: {:?}", e))?;