fix: Release GPU memory and bump sppark

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"]
 

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<PathBuf> {
+    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<String> {
     let output = command_output(nvcc, &["--list-gpu-code"]).unwrap_or_default();
     let supported = String::from_utf8_lossy(&output);

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<int>(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<int>(status);
+}

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,

src/cuda/pcs.rs

@@ -190,9 +190,13 @@ where
     Dft: CudaPcsDft<F>,
     Radix2DitParallel<F>: TwoAdicSubgroupDft<F>,
 {
-    let cpu = Radix2DitParallel::<F>::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::<F>::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::<Goldilocks>::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() {