diff --git a/Cargo.lock b/Cargo.lock index d560e9c..a59a76d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1113,8 +1113,8 @@ dependencies = [ [[package]] name = "sppark" -version = "0.1.15" -source = "git+https://github.com/argumentcomputer/sppark?rev=114feb69ebfe029e2caa5e6c2d372c695e5d4da8#114feb69ebfe029e2caa5e6c2d372c695e5d4da8" +version = "0.1.16" +source = "git+https://github.com/argumentcomputer/sppark?rev=9326ac014ec5415ebc4eab24c1246276714aee57#9326ac014ec5415ebc4eab24c1246276714aee57" dependencies = [ "cc", "which", diff --git a/Cargo.toml b/Cargo.toml index 3ae8b24..d923195 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,7 +37,7 @@ p3-util = { git = "https://github.com/Plonky3/Plonky3", rev = "3152b14a89067c837 # streams and a runtime compatible with the Lean executable's C++ linkage. [dependencies.sppark] git = "https://github.com/argumentcomputer/sppark" -rev = "114feb69ebfe029e2caa5e6c2d372c695e5d4da8" +rev = "9326ac014ec5415ebc4eab24c1246276714aee57" optional = true features = ["cuda"] diff --git a/build.rs b/build.rs index 03d9f10..ffc7b99 100644 --- a/build.rs +++ b/build.rs @@ -1,8 +1,13 @@ //! CUDA compilation is opt-in; CPU builds need no CUDA toolchain. +//! +//! Through `links = "multi_stark_cuda"`, dependents' build scripts receive +//! the toolchain this backend was compiled with as `DEP_MULTI_STARK_CUDA_*`: +//! `NVCC`, `ARCHS` (comma-separated) and `LIB_DIRS` (a search path for the +//! CUDA runtime), so their kernels target the same compiler and GPUs. use std::env; use std::ffi::{OsStr, OsString}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::process::Command; const CUDA_EXPORTS: &str = r"{ @@ -47,6 +52,11 @@ fn main() { let nvcc = nvcc_path(); let architectures = cuda_architectures(&nvcc); let compiler = || cuda_command(&nvcc, &architectures); + println!("cargo:nvcc={}", Path::new(&nvcc).display()); + println!("cargo:archs={}", architectures.join(",")); + let lib_dirs = env::join_paths(cuda_library_directories(&nvcc)) + .expect("CUDA library directories contain a path separator"); + println!("cargo:lib_dirs={}", Path::new(&lib_dirs).display()); let objects = [ (PathBuf::from("cuda/sppark_ntt.cu"), "sppark_ntt.o"), @@ -122,6 +132,23 @@ fn nvcc_path() -> OsString { .unwrap_or_else(|| "nvcc".into()) } +/// The CUDA toolkit's library directories: under `CUDA_HOME` or +/// `CUDA_PATH`, else the toolkit an absolute `nvcc` belongs to, else +/// `/usr/local/cuda`. +fn cuda_library_directories(nvcc: &OsStr) -> Vec { + let root = env::var_os("CUDA_HOME") + .or_else(|| env::var_os("CUDA_PATH")) + .map(PathBuf::from) + .or_else(|| { + let nvcc = Path::new(nvcc); + nvcc.is_absolute() + .then(|| nvcc.parent()?.parent().map(Path::to_path_buf)) + .flatten() + }) + .unwrap_or_else(|| PathBuf::from("/usr/local/cuda")); + vec![root.join("lib64"), root.join("targets/x86_64-linux/lib")] +} + fn cuda_architectures(nvcc: &OsStr) -> Vec { let output = command_output(nvcc, &["--list-gpu-code"]).unwrap_or_default(); let supported = String::from_utf8_lossy(&output); diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 02305fb..ec114ea 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -4224,3 +4224,19 @@ extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, } return static_cast(status); } + +// Return what the default pool retains (see `retain_default_pool`) on every +// visible device to the driver, once each device's queued work completes, so +// another process or a later stage can allocate it. Live allocations stay. +extern "C" int multi_stark_cuda_release_memory() { + int count = 0; + cudaError_t status = cudaGetDeviceCount(&count); + for (int device_id = 0; status == cudaSuccess && device_id < count; ++device_id) { + cudaMemPool_t pool = nullptr; + status = cudaSetDevice(device_id); + if (status == cudaSuccess) status = cudaDeviceSynchronize(); + if (status == cudaSuccess) status = cudaDeviceGetDefaultMemPool(&pool, device_id); + if (status == cudaSuccess) status = cudaMemPoolTrimTo(pool, 0); + } + return static_cast(status); +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index 8143ae2..b220be0 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -2240,6 +2240,15 @@ pub(crate) fn device_memory_info(device_id: i32) -> (usize, usize) { (free_bytes, total_bytes) } +/// Returns the device memory this backend's allocation pool retains between +/// proofs to the driver, on every visible device, after queued work completes. +/// Memory held by live buffers is unaffected. Call it once proving is done and +/// before another process or a different backend needs the devices. +pub fn release_device_memory() { + let status = unsafe { multi_stark_cuda_release_memory() }; + check_cuda(status, "CUDA memory release"); +} + /// Device bytes the prover keeps free for its later stages: a quarter of the /// device unless `MULTI_STARK_CUDA_MIN_FREE_BYTES` says otherwise. pub(crate) fn minimum_free_bytes(total_bytes: usize) -> usize { @@ -2923,6 +2932,7 @@ impl Drop for CudaMixedMerkleTree { } unsafe extern "C" { + fn multi_stark_cuda_release_memory() -> i32; fn multi_stark_cuda_memory_info( device_id: i32, free_bytes: *mut usize, diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs index 71e975d..a5ea164 100644 --- a/src/cuda/pcs.rs +++ b/src/cuda/pcs.rs @@ -190,9 +190,13 @@ where Dft: CudaPcsDft, Radix2DitParallel: TwoAdicSubgroupDft, { - let cpu = Radix2DitParallel::::default(); jobs.into_par_iter() .map(|(index, (domain, evaluations))| { + // At the pinned Plonky3, twiddle initialization schedules Rayon work + // while holding a spin lock, so sharing one cache across concurrent + // jobs can deadlock. Plonky3 4148206d computes twiddles outside the + // lock; share one DFT again once the Plonky3 pin includes it. + let cpu = Radix2DitParallel::::default(); let shift = F::GENERATOR / domain.shift(); if let Some(lde) = dft.try_coset_lde_batch_host(&evaluations, log_blowup, shift) { return (index, lde); @@ -2338,6 +2342,39 @@ mod tests { use crate::types::ExtVal; use p3_field::Field; + #[test] + #[ignore = "requires a CUDA device"] + fn concurrent_host_lde_jobs_match_serial_transforms() { + let dft = crate::cuda::CudaDft::new(0).with_auxiliary_devices(&[]); + let jobs = (0..32) + .map(|index| { + let log_height = 12 + index % 3; + let width = 1 + index % 4; + let domain = TwoAdicMultiplicativeCoset::new( + Goldilocks::from_usize(1 + index % 5), + log_height, + ) + .unwrap(); + let values = (0..(1 << log_height) * width) + .map(|i| Goldilocks::from_usize(i + index)) + .collect(); + (index, (domain, RowMajorMatrix::new(values, width))) + }) + .collect_vec(); + let expected = jobs + .iter() + .map(|(index, (domain, matrix))| { + let cpu = Radix2DitParallel::::default(); + let output = cpu + .coset_lde_batch(matrix.clone(), 2, Goldilocks::GENERATOR / domain.shift()) + .bit_reverse_rows() + .to_row_major_matrix(); + (*index, output) + }) + .collect_vec(); + assert_eq!(host_coset_lde_jobs(&dft, jobs, 2), expected); + } + #[test] #[ignore = "requires two CUDA devices"] fn auxiliary_host_ldes_preserve_matrix_order_cosets_and_canonical_words() {