Leave out the cuDNN libraries MLX never loads

crates/stemd-server/src/cuda.rs

@@ -175,7 +175,7 @@ fn unpack_libraries(archive: &Path, into: &Path) -> Result<usize> {
         let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
             continue;
         };
-        if !name.to_ascii_lowercase().ends_with(".dll") {
+        if !wanted(name) {
             continue;
         }
 
@@ -194,6 +194,21 @@ fn unpack_libraries(archive: &Path, into: &Path) -> Result<usize> {
     Ok(placed)
 }
 
+/// cuDNN libraries MLX never loads: it uses only the graph API, and no other
+/// cuDNN library links these.
+const UNUSED: &[&str] = &[
+    "cudnn_adv64_9.dll",
+    "cudnn_ops64_9.dll",
+    "cudnn_cnn64_9.dll",
+    "cudnn_ext64_9.dll",
+];
+
+/// Whether an archive entry is a DLL to install.
+fn wanted(name: &str) -> bool {
+    let name = name.to_ascii_lowercase();
+    name.ends_with(".dll") && !UNUSED.contains(&name.as_str())
+}
+
 /// Fail on a read-only install directory before spending a gigabyte finding out.
 ///
 /// Under `Program Files` this is the ordinary case rather than a strange one,
@@ -271,6 +286,19 @@ mod tests {
         }
     }
 
+    #[test]
+    fn the_unused_cudnn_libraries_are_skipped() {
+        assert!(wanted("cudnn64_9.dll"));
+        assert!(wanted("cudnn_graph64_9.dll"));
+        assert!(wanted("cublasLt64_13.dll"));
+        assert!(!wanted("cudnn_ops64_9.dll"));
+        assert!(!wanted("CUDNN_ADV64_9.DLL"));
+        assert!(!wanted("cudnn_cnn64_9.dll"));
+        assert!(!wanted("cudnn_ext64_9.dll"));
+        assert!(wanted("cudnn_engines_tensor_ir64_9.dll"));
+        assert!(!wanted("cudnn.h"));
+    }
+
     /// Distinct names, since they share one staging directory.
     #[test]
     fn no_two_components_stage_under_the_same_name() {

docs/building.md

@@ -161,7 +161,7 @@ at run time.
 | `jit-headers.sh` | writes `jit-headers.txt`, the toolkit headers MLX's run-time compiler reads; rerun after an MLX or CUDA update |
 
 Every package carries CUDA, cuDNN, BLAS and avahi-compat in `/opt/stemd/lib`,
-about 2 GB unpacked. The driver (`libcuda.so.1`), the display stack and the C++
+about 1.7 GB unpacked. The driver (`libcuda.so.1`), the display stack and the C++
 runtime come from the machine. NVIDIA's libraries are copied unmodified, as
 their licences require; the executables find them through `DT_RPATH`
 (`$ORIGIN/../lib`), which also covers what cuDNN opens at run time.

packaging/linux/stage.sh

@@ -82,9 +82,16 @@ soname() { readelf -d "$1" | sed -n 's/.*(SONAME).*\[\(.*\)\]/\1/p'; }
 # cuDNN opens its engine libraries itself and NVRTC opens its builtins, so no
 # NEEDED entry names them. They go into the walk so their own dependencies come
 # too. Copied under their soname, which is the name they are opened by.
+# cuDNN's adv, ops, cnn and ext libraries are left out: MLX uses only the
+# graph API, which never loads them, and no other cuDNN library links them.
+unused_cudnn() {
+    case "$1" in libcudnn_adv.so*|libcudnn_ops.so*|libcudnn_cnn.so*|libcudnn_ext.so*) return 0 ;; esac
+    return 1
+}
 queue=("$exe")
 for f in "$cudnn"/lib/libcudnn*.so.9 "$cuda"/lib64/libnvrtc-builtins.so; do
     [[ -f "$f" ]] || continue
+    unused_cudnn "$(basename "$f")" && continue
     name="$(soname "$f")"
     [[ -f "$prefix/lib/$name" ]] && continue
     cp -L "$f" "$prefix/lib/$name"
@@ -102,6 +109,7 @@ while [[ ${#queue[@]} -gt 0 ]]; do
     resolved="$(ldd "$target" | awk '$2 == "=>" {print $1, ($3 == "not" ? "" : $3)}')"
     for name in $(readelf -d "$target" | sed -n 's/.*(NEEDED).*\[\(.*\)\]/\1/p'); do
         host_owned "$name" && continue
+        unused_cudnn "$name" && die "$(basename "$target") links $name, which is left out as unused"
         [[ "$seen" == *" $name "* ]] && continue
         seen="$seen$name "
         path="$(awk -v n="$name" '$1 == n {print $2}' <<<"$resolved")"