Wiring and making GPU accel. optional with an argument

Author: vrbouza

src/cli.rs

@@ -197,8 +197,8 @@ pub enum Commands {
     /// Calculate pairwise distances using sketches
     Dist {
         /// The .skm file used as the reference
-        #[arg(required = true)]
-        ref_db: String,
+        #[arg(required_unless_present = "list_gpus")]
+        ref_db: Option<String>,
 
         /// The .skm file used as the query (omit for ref v ref)
         #[arg(group = "query")]
@@ -240,6 +240,17 @@ pub enum Commands {
         /// minimum completeness for a sample to be corrected but the completeness correction
         #[arg(long, default_value_t = 0.64)]
         completeness_cutoff: f64,
+
+        /// Calculate dense Jaccard/ANI distances on the GPU (needs a build with
+        /// `--features wgpu`): alone, wgpu chooses, preferring a discrete GPU; with N, GPU N
+        /// of --list-gpus. kNN and core/accessory distances stay on the CPU. Give it after
+        /// the databases, or as --gpu=N, as it takes an optional value
+        #[arg(long, num_args = 0..=1, value_name = "N")]
+        gpu: Option<Option<usize>>,
+
+        /// List the GPUs --gpu N can choose from, and exit
+        #[arg(long)]
+        list_gpus: bool,
     },
 
     /// Building and querying with inverted indices (.ski)

src/distances/mod.rs

@@ -232,22 +232,6 @@ fn self_dists_all_stream_generic<const BB: usize, W: IoWrite + Send>(
     completeness_cutoff: f64,
     threads: usize,
 ) -> Result<(), Error> {
-    #[cfg(feature = "wgpu")]
-    {
-        if matches!(dist_type, DistType::Jaccard(..)) {
-            return dists_all_stream_wgpu::<BB, W>(
-                writer,
-                sketches,
-                None,
-                &dist_type,
-                completeness_vec,
-                None,
-                completeness_cutoff,
-            );
-        }
-        log::warn!("The wgpu backend does not support core/accessory distances yet; using the CPU");
-    }
-
     let ani = matches!(dist_type, DistType::Jaccard(_, _, true));
     let k_vals = match dist_type {
         DistType::Jaccard(k_idx, k_val, _) => Some((k_idx, k_val)),
@@ -333,22 +317,25 @@ fn self_dists_all_stream_generic<const BB: usize, W: IoWrite + Send>(
 
 
 /// GPU (wgpu) variant of [`self_dists_all_stream`] (`query_sketches` = None) and
-/// [`cross_dists_all_stream`], for Jaccard distances and ANI. Same line format; pairs come
-/// out tile by tile and chunk by chunk, in no global order, as in the CUDA code.
+/// [`cross_dists_all_stream`], for Jaccard distances and ANI, on GPU `gpu_device` of
+/// [`available_gpus`] (None: wgpu's choice). Same line format; pairs come out tile by tile
+/// and chunk by chunk, in no global order, as in the CUDA code.
 #[cfg(feature = "wgpu")]
-fn dists_all_stream_wgpu<const BB: usize, W: IoWrite + Send>(
+pub fn dists_all_stream_gpu<W: IoWrite + Send>(
     writer: &mut W,
     ref_sketches: &MultiSketch,
     query_sketches: Option<&MultiSketch>,
-    dist_type: &DistType,
+    dist_type: DistType,
     ref_completeness_vec: Option<&Vec<f64>>,
     query_completeness_vec: Option<&Vec<f64>>,
     completeness_cutoff: f64,
+    gpu_device: Option<usize>,
 ) -> Result<(), Error> {
     let pool_size = rayon::current_num_threads();
     if pool_size < 2 {
         bail!("GPU dense streaming requires a Rayon producer and writer; configure the pool with set_threads_with_writer()");
     }
+    let bin_bits = if ref_sketches.is_legacy_format() { LEGACY_BIN_BITS } else { BIN_BITS };
     let ref_names = <DistanceMatrix as Distances>::sketch_names(ref_sketches);
     let query_names = query_sketches.map(<DistanceMatrix as Distances>::sketch_names);
 
@@ -372,11 +359,12 @@ fn dists_all_stream_wgpu<const BB: usize, W: IoWrite + Send>(
         let gpu_result = pollster::block_on(get_distances_wgpu(
             ref_sketches,
             query_sketches,
-            dist_type,
+            &dist_type,
             ref_completeness_vec,
             query_completeness_vec,
             completeness_cutoff,
-            BB,
+            bin_bits,
+            gpu_device,
             |tile, dists| match &query_names {
                 // Self mode: both sides from the reference set, the lower index first
                 None => tile.consume(dists, &ref_names, &ref_names, false, &text_tx),
@@ -391,6 +379,33 @@ fn dists_all_stream_wgpu<const BB: usize, W: IoWrite + Send>(
     })
 }
 
+/// Stand-in for builds without the `wgpu` feature, so callers need no cfg of their own
+#[cfg(not(feature = "wgpu"))]
+pub fn dists_all_stream_gpu<W: IoWrite + Send>(
+    _writer: &mut W,
+    _ref_sketches: &MultiSketch,
+    _query_sketches: Option<&MultiSketch>,
+    _dist_type: DistType,
+    _ref_completeness_vec: Option<&Vec<f64>>,
+    _query_completeness_vec: Option<&Vec<f64>>,
+    _completeness_cutoff: f64,
+    _gpu_device: Option<usize>,
+) -> Result<(), Error> {
+    bail!("This sketchlib was built without GPU support: rebuild it with --features wgpu")
+}
+
+/// The GPUs [`dists_all_stream_gpu`] can use, one line each, `index: name (type, backend)`
+#[cfg(feature = "wgpu")]
+pub fn available_gpus() -> Result<Vec<String>, Error> {
+    Ok(crate::gpu::dists_wgpu::gpu_list())
+}
+
+/// Stand-in for builds without the `wgpu` feature
+#[cfg(not(feature = "wgpu"))]
+pub fn available_gpus() -> Result<Vec<String>, Error> {
+    bail!("This sketchlib was built without GPU support: rebuild it with --features wgpu")
+}
+
 /// Self query mode (dense, all distances)
 pub fn self_dists_knn<'a>(
     sketches: &'a MultiSketch,
@@ -709,22 +724,6 @@ fn cross_dists_all_stream_generic<const BB: usize, W: IoWrite + Send>(
     completeness_cutoff: f64,
     threads: usize,
 ) -> Result<(), Error> {
-    #[cfg(feature = "wgpu")]
-    {
-        if matches!(dist_type, DistType::Jaccard(..)) {
-            return dists_all_stream_wgpu::<BB, W>(
-                writer,
-                ref_sketches,
-                Some(query_sketches),
-                &dist_type,
-                ref_completeness_vec,
-                query_completeness_vec,
-                completeness_cutoff,
-            );
-        }
-        log::warn!("The wgpu backend does not support core/accessory distances yet; using the CPU");
-    }
-
     let ani = matches!(dist_type, DistType::Jaccard(_, _, true));
     let k_vals = match dist_type {
         DistType::Jaccard(k_idx, k_val, _) => Some((k_idx, k_val)),

src/gpu/dists_wgpu.rs

@@ -187,6 +187,7 @@ pub async fn get_distances_wgpu(
     query_completeness_vec: Option<&Vec<f64>>, // only read with query_sketches
     completeness_cutoff: f64,
     bin_bits: usize, // BB chosen by the CPU dispatch: BIN_BITS, or LEGACY_BIN_BITS
+    gpu_device: Option<usize>, // index in gpu_list(); None: wgpu's choice
     mut on_tile: impl FnMut(&DistTile, &[f32]) -> anyhow::Result<()>,
 ) -> anyhow::Result<()> {
     let &DistType::Jaccard(k_idx, _, ani) = dist_type else {
@@ -211,22 +212,12 @@ pub async fn get_distances_wgpu(
         "Reference and query sketches differ in k-mer lengths, sketch size, hash type or format"
     );
 
-    // TODO: this could be a parameter? Also, we should expose which GPU to select by index, apart than from PowerPreference, which might be quite useful for browser but a bit less in highly optimised environments.
-    let power_pref = wgpu::PowerPreference::None;
-    let backends = wgpu::Backends::all();
-
-    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
-        backends,
-        // This fills the remaining items with the "..THING" and the thing itself.
-        ..wgpu::InstanceDescriptor::new_without_display_handle()
-    });
-
-    let adapter = instance
-        .request_adapter(&wgpu::RequestAdapterOptions {
-            power_preference: power_pref, ..Default::default()
-        })
-        .await
-        .context("No GPU adapter found")?;
+    let instance = create_instance();
+    let adapter = select_adapter(&instance, gpu_device).await?;
+    let info = adapter.get_info();
+    if info.device_type == wgpu::DeviceType::Cpu {
+        log::warn!("{} is a software adapter, not a GPU: expect it to be slow", info.name);
+    }
 
     // Ask for the hardware limits rather than the WebGPU defaults, so the tile ranges and
     // workgroup sizes reflect what this device can actually do
@@ -239,7 +230,7 @@ pub async fn get_distances_wgpu(
         .await
         .context("Failed to get GPU device")?;
 
-    log::info!("Using this GPU: {}", adapter.get_info().name);
+    log::info!("Using this GPU: {} ({:?}, {:?})", info.name, info.device_type, info.backend);
 
     // =============================================
     // SETTING UP DATA!
@@ -670,6 +661,58 @@ pub async fn get_distances_wgpu(
     Ok(())
 }
 
+// Native backends only: with OpenGL too, a GPU would be listed twice
+const GPU_BACKENDS: wgpu::Backends = wgpu::Backends::PRIMARY;
+
+// The wgpu instance behind every GPU entry point
+fn create_instance() -> wgpu::Instance {
+    wgpu::Instance::new(wgpu::InstanceDescriptor {
+        backends: GPU_BACKENDS,
+        // This fills the remaining items with the "..THING" and the thing itself.
+        ..wgpu::InstanceDescriptor::new_without_display_handle()
+    })
+}
+
+// The adapter --gpu N names, by its index in enumerate_adapters() (as --list-gpus shows
+// it); with --gpu alone, the adapter wgpu prefers for performance, a discrete GPU if any
+async fn select_adapter(
+    instance: &wgpu::Instance,
+    gpu_device: Option<usize>,
+) -> anyhow::Result<wgpu::Adapter> {
+    match gpu_device {
+        Some(index) => {
+            let mut adapters = instance.enumerate_adapters(GPU_BACKENDS).await;
+            ensure!(
+                index < adapters.len(),
+                "There is no GPU {index}: `sketchlib dist --list-gpus` lists the {} available",
+                adapters.len()
+            );
+            Ok(adapters.swap_remove(index))
+        }
+        None => instance
+            .request_adapter(&wgpu::RequestAdapterOptions {
+                power_preference: wgpu::PowerPreference::HighPerformance,
+                ..Default::default()
+            })
+            .await
+            .context("No GPU adapter found"),
+    }
+}
+
+/// The GPUs --gpu N chooses from, one line each, `index: name (type, backend)`, in the
+/// order it numbers them
+pub fn gpu_list() -> Vec<String> {
+    let adapters = pollster::block_on(create_instance().enumerate_adapters(GPU_BACKENDS));
+    adapters
+        .iter()
+        .enumerate()
+        .map(|(index, adapter)| {
+            let info = adapter.get_info();
+            format!("{index}: {} ({:?}, {:?})", info.name, info.device_type, info.backend)
+        })
+        .collect()
+}
+
 // Workgroups needed for all-vs-all: each query row of n - q - 1 references gets its
 // own workgroups, as in CUDA's getBlockSize()
 fn self_dispatch_count(n: u32, workgroup_size: u32) -> u64 {

src/lib.rs

@@ -154,6 +154,7 @@ use std::time::Instant;
 extern crate arrayref;
 extern crate num_cpus;
 use anyhow::{bail, Error};
+use crate::distances::distance_matrix::DistType;
 #[cfg(not(target_arch = "wasm32"))]
 use indicatif::ParallelProgressIterator;
 #[cfg(not(target_arch = "wasm32"))]
@@ -316,7 +317,19 @@ pub fn main() -> Result<(), Error> {
             ref_completeness_file,
             query_completeness_file,
             completeness_cutoff,
+            gpu,
+            list_gpus,
         } => {
+            if *list_gpus {
+                for gpu_line in available_gpus()? {
+                    println!("{gpu_line}");
+                }
+                return Ok(());
+            }
+            // clap requires it without --list-gpus
+            let Some(ref_db) = ref_db.as_deref() else {
+                bail!("dist needs a reference database");
+            };
             if knn.is_none() {
                 // Dense path streams output via an internal writer thread running
                 // concurrently with compute (mirrors Sketch/Append/Inverted Build/Query).
@@ -353,6 +366,17 @@ pub fn main() -> Result<(), Error> {
                 panic!("Error setting k size: {e}");
             });
 
+            // The GPU covers dense Jaccard and ANI distances; anything else stays on the CPU.
+            // --gpu alone leaves the device to wgpu (gpu_device = None)
+            if gpu.is_some() && !cfg!(feature = "wgpu") {
+                bail!("This sketchlib was built without GPU support: rebuild it with --features wgpu");
+            }
+            let gpu_device = gpu.flatten();
+            let use_gpu = gpu.is_some() && knn.is_none() && matches!(dist_type, DistType::Jaccard(..));
+            if gpu.is_some() && !use_gpu {
+                log::warn!("--gpu only covers dense Jaccard and ANI distances (-k without --knn) for now; using the CPU");
+            }
+
             // Read queries if supplied. Note no subsetting here
             let queries = if let Some(query_db_name) = query_db {
                 log::info!("Loading query sketches from {query_db_name}.skm/.skd");
@@ -379,16 +403,29 @@ pub fn main() -> Result<(), Error> {
                             // Self mode (dense)
                             log::info!("Calculating all ref vs ref distances");
                             log::info!("Streaming out in long matrix form");
-                            self_dists_all_stream(
-                                &mut output_file,
-                                &references,
-                                n,
-                                dist_type,
-                                args.quiet,
-                                ref_completeness_vec.as_ref(),
-                                *completeness_cutoff,
-                                *threads,
-                            )?;
+                            if use_gpu {
+                                dists_all_stream_gpu(
+                                    &mut output_file,
+                                    &references,
+                                    None,
+                                    dist_type,
+                                    ref_completeness_vec.as_ref(),
+                                    None,
+                                    *completeness_cutoff,
+                                    gpu_device,
+                                )?;
+                            } else {
+                                self_dists_all_stream(
+                                    &mut output_file,
+                                    &references,
+                                    n,
+                                    dist_type,
+                                    args.quiet,
+                                    ref_completeness_vec.as_ref(),
+                                    *completeness_cutoff,
+                                    *threads,
+                                )?;
+                            }
                         }
                         Some(mut nn) => {
                             // Self mode (sparse): a genome cannot be its own neighbour
@@ -452,19 +489,32 @@ pub fn main() -> Result<(), Error> {
                             // Cross-query mode (dense, all pairs)
                             log::info!("Calculating all ref vs query distances");
                             log::info!("Streaming out in long matrix form");
-                            cross_dists_all_stream(
-                                &mut output_file,
-                                &references,
-                                &query_db,
-                                n,
-                                n_query,
-                                dist_type,
-                                args.quiet,
-                                ref_completeness_vec.as_ref(),
-                                query_completeness_vec.as_ref(),
-                                *completeness_cutoff,
-                                *threads,
-                            )?;
+                            if use_gpu {
+                                dists_all_stream_gpu(
+                                    &mut output_file,
+                                    &references,
+                                    Some(&query_db),
+                                    dist_type,
+                                    ref_completeness_vec.as_ref(),
+                                    query_completeness_vec.as_ref(),
+                                    *completeness_cutoff,
+                                    gpu_device,
+                                )?;
+                            } else {
+                                cross_dists_all_stream(
+                                    &mut output_file,
+                                    &references,
+                                    &query_db,
+                                    n,
+                                    n_query,
+                                    dist_type,
+                                    args.quiet,
+                                    ref_completeness_vec.as_ref(),
+                                    query_completeness_vec.as_ref(),
+                                    *completeness_cutoff,
+                                    *threads,
+                                )?;
+                            }
                         }
                     }
                 }