From bf707bf17712ed0fa713e8f21f9502b75800951b Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 22 Jun 2026 20:42:43 +0800 Subject: [PATCH 1/8] perf: optimize cuVS IVF-PQ artifact builds --- backends/cuvs_26_02/Cargo.lock | 1 + backends/cuvs_26_02/Cargo.toml | 1 + .../lance_cuvs_backend_cu12/__init__.py | 10 + backends/cuvs_26_02/src/backend.rs | 661 ++++++++++-------- backends/cuvs_26_02/src/cuda.rs | 90 ++- 5 files changed, 455 insertions(+), 308 deletions(-) diff --git a/backends/cuvs_26_02/Cargo.lock b/backends/cuvs_26_02/Cargo.lock index b25228a..ffbb090 100644 --- a/backends/cuvs_26_02/Cargo.lock +++ b/backends/cuvs_26_02/Cargo.lock @@ -4657,6 +4657,7 @@ dependencies = [ "lance-index", "lance-io", "lance-linalg", + "libc", "log", "ndarray 0.16.1", "pyo3", diff --git a/backends/cuvs_26_02/Cargo.toml b/backends/cuvs_26_02/Cargo.toml index b3cd29f..ac44a54 100644 --- a/backends/cuvs_26_02/Cargo.toml +++ b/backends/cuvs_26_02/Cargo.toml @@ -25,6 +25,7 @@ cuvs = "=26.2.0" cuvs-sys = "=26.2.0" futures = "0.3" half = { version = "2.5", default-features = false, features = ["num-traits", "std"] } +libc = "0.2" # This crate depends on unreleased Lance vector-build APIs. Keep the Rust-side # Lance dependencies pinned until those APIs stabilize. lance = { git = "https://github.com/lance-format/lance.git", rev = "6112a34bfe38618f07c099217dc3d89fd39ca6bb" } diff --git a/backends/cuvs_26_02/python/lance_cuvs_backend_cu12/__init__.py b/backends/cuvs_26_02/python/lance_cuvs_backend_cu12/__init__.py index 9f7cf54..9f87b1f 100644 --- a/backends/cuvs_26_02/python/lance_cuvs_backend_cu12/__init__.py +++ b/backends/cuvs_26_02/python/lance_cuvs_backend_cu12/__init__.py @@ -4,6 +4,7 @@ import os import tempfile +from collections.abc import Mapping from ._native import ( IvfPqArtifactOutput, @@ -31,6 +32,7 @@ def train_ivf_pq( max_iters: int = 50, num_bits: int = 8, filter_nan: bool = True, + storage_options: Mapping[str, str] | None = None, ) -> IvfPqTrainingOutput: """Train an IVF_PQ model with cuVS. @@ -55,6 +57,8 @@ def train_ivf_pq( Number of bits per PQ code. cuVS currently supports only ``8`` here. filter_nan: Whether to drop null or non-finite vectors before training. + storage_options: + Optional Lance storage options used when opening ``dataset_uri``. Returns ------- @@ -72,6 +76,7 @@ def train_ivf_pq( max_iters=max_iters, num_bits=num_bits, filter_nan=filter_nan, + storage_options=dict(storage_options) if storage_options is not None else None, ) @@ -83,6 +88,7 @@ def build_ivf_pq_artifact( artifact_uri: str | os.PathLike[str] | None = None, batch_size: int = 1024 * 128, filter_nan: bool = True, + storage_options: Mapping[str, str] | None = None, ) -> IvfPqArtifactOutput: """Encode a dataset into a partition-local IVF_PQ artifact. @@ -102,6 +108,9 @@ def build_ivf_pq_artifact( Number of rows per transform batch. filter_nan: Whether to drop null or non-finite vectors during artifact build. + storage_options: + Optional Lance storage options used when opening ``dataset_uri`` and + writing ``artifact_uri``. Returns ------- @@ -119,6 +128,7 @@ def build_ivf_pq_artifact( training=training, batch_size=batch_size, filter_nan=filter_nan, + storage_options=dict(storage_options) if storage_options is not None else None, ) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index f9df4aa..167fc3a 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -2,40 +2,34 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use crate::cuda::{ - CudaEvent, CuvsIvfPqIndex, DeviceTensor, HostTensorView, PinnedHostBuffer, check_cuvs, - copy_tensor_to_host_f32_2d, copy_tensor_to_host_f32_3d, create_index_params, - destroy_index_params, ivf_centroids_from_host, make_tensor_view, matrix_from_vectors, - pq_codebook_from_host, + CudaEvent, CuvsIvfPqIndex, DeviceTensor, HostTensorView, MatrixBuffer, PinnedHostBuffer, + RegisteredHostBuffer, check_cuvs, copy_tensor_to_host_f32_2d, copy_tensor_to_host_f32_3d, + create_index_params, destroy_index_params, ivf_centroids_from_host, make_tensor_view, + matrix_from_vectors, pq_codebook_from_host, }; use arrow::compute::filter; use arrow_array::cast::AsArray; -use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, ListArray, RecordBatch, UInt8Array, UInt32Array, -}; -use arrow_buffer::{OffsetBuffer, ScalarBuffer}; +use arrow_array::{Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use cuvs::Resources; -use futures::{FutureExt, SinkExt, TryStreamExt, channel::mpsc, future::LocalBoxFuture}; +use futures::{FutureExt, SinkExt, StreamExt, TryStreamExt, channel::mpsc, future::LocalBoxFuture}; use lance::dataset::Dataset; +use lance::index::vector::PartitionArtifactBuilder; use lance::index::vector::utils::infer_vector_dim; use lance_arrow::FixedSizeListArrayExt; use lance_core::{Error, ROW_ID, Result}; -use lance_file::version::LanceFileVersion; -use lance_file::writer::{FileWriter, FileWriterOptions}; -use lance_index::vector::ivf::shuffler::IvfShuffler; use lance_index::vector::utils::is_finite; use lance_index::vector::{PART_ID_COLUMN, PQ_CODE_COLUMN}; -use lance_io::local::to_local_path; use lance_linalg::distance::DistanceType; use log::warn; use std::collections::HashMap; use std::sync::Arc; +use std::time::{Duration, Instant}; const PARTITION_ARTIFACT_METADATA_FILE_NAME: &str = "metadata.lance"; -const PARTITION_ARTIFACT_FILE_VERSION: &str = "2.2"; -const PARTITION_ARTIFACT_UNSORTED_FILE_NAME: &str = "unsorted.lance"; -const PARTITION_ARTIFACT_SHUFFLE_BATCHES_PER_FILE: usize = 16; const PIPELINE_SLOTS: usize = 2; +const DEFAULT_SCAN_FRAGMENT_READAHEAD: usize = 64; +const DEFAULT_SCAN_IO_BUFFER_SIZE: u64 = 16 * 1024 * 1024 * 1024; /// A trained cuVS IVF_PQ model that can be reused for artifact builds. /// @@ -183,6 +177,7 @@ impl VectorBuildBackend for CuvsVectorBuildBackend { async move { match params.kind { VectorIndexKind::IvfPq(build_params) => { + let train_start = Instant::now(); let trained = train_ivf_pq( dataset, ¶ms.column, @@ -195,6 +190,11 @@ impl VectorBuildBackend for CuvsVectorBuildBackend { params.filter_nan, ) .await?; + eprintln!( + "cuVS train_ivf_pq time: {:.3}s", + train_start.elapsed().as_secs_f64() + ); + let artifact_start = Instant::now(); let files = assign_ivf_pq_to_artifact( dataset, ¶ms.column, @@ -205,6 +205,11 @@ impl VectorBuildBackend for CuvsVectorBuildBackend { None, ) .await?; + eprintln!( + "cuVS assign_ivf_pq_to_artifact time: {:.3}s files={}", + artifact_start.elapsed().as_secs_f64(), + files.len() + ); Ok(VectorIndexBuildOutput::PartitionArtifact( PartitionArtifactBuildOutput { artifact_uri: params.artifact_uri, @@ -286,116 +291,6 @@ fn vector_column_to_fsl(batch: &RecordBatch, column: &str) -> Result Result { - let ivf_offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0i32, ivf_centroids.len() as i32])); - let pq_offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0i32, pq_codebook.len() as i32])); - let ivf_list = ListArray::new( - Arc::new(Field::new( - "_ivf_centroids_item", - ivf_centroids.data_type().clone(), - false, - )), - ivf_offsets, - Arc::new(ivf_centroids.clone()), - None, - ); - let pq_list = ListArray::new( - Arc::new(Field::new( - "_pq_codebook_item", - pq_codebook.data_type().clone(), - false, - )), - pq_offsets, - Arc::new(pq_codebook.clone()), - None, - ); - let schema = Arc::new(ArrowSchema::new(vec![ - Field::new("_ivf_centroids", ivf_list.data_type().clone(), false), - Field::new("_pq_codebook", pq_list.data_type().clone(), false), - ])); - Ok(RecordBatch::try_new( - schema, - vec![Arc::new(ivf_list), Arc::new(pq_list)], - )?) -} - -fn metadata_writer_options() -> Result { - Ok(FileWriterOptions { - format_version: Some( - PARTITION_ARTIFACT_FILE_VERSION - .parse::() - .map_err(|error| { - Error::invalid_input(format!( - "invalid partition artifact file version '{}': {}", - PARTITION_ARTIFACT_FILE_VERSION, error - )) - })?, - ), - ..Default::default() - }) -} - -async fn write_partition_artifact_metadata( - artifact_uri: &str, - trained: &TrainedIvfPqIndex, - storage_options: Option<&HashMap>, -) -> Result<()> { - let registry = Arc::new(lance_io::object_store::ObjectStoreRegistry::default()); - let params = if let Some(storage_options) = storage_options { - lance_io::object_store::ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - lance_io::object_store::StorageOptionsAccessor::with_static_options( - storage_options.clone(), - ), - )), - ..Default::default() - } - } else { - lance_io::object_store::ObjectStoreParams::default() - }; - let (object_store, root_dir) = - lance::io::ObjectStore::from_uri_and_params(registry, artifact_uri, ¶ms) - .await - .map_err(|error| Error::io(error.to_string()))?; - let path = root_dir.child(PARTITION_ARTIFACT_METADATA_FILE_NAME); - let batch = build_metadata_batch(&trained.ivf_centroids, &trained.pq_codebook)?; - let mut writer = FileWriter::try_new( - object_store.create(&path).await?, - lance_core::datatypes::Schema::try_from(batch.schema().as_ref())?, - metadata_writer_options()?, - )?; - writer.add_schema_metadata( - "lance:index_build:artifact_version".to_string(), - "1".to_string(), - ); - writer.add_schema_metadata( - "lance:index_build:distance_type".to_string(), - trained.metric_type.to_string(), - ); - writer.add_schema_metadata( - "lance:index_build:num_partitions".to_string(), - trained.num_partitions.to_string(), - ); - writer.add_schema_metadata( - "lance:index_build:num_sub_vectors".to_string(), - trained.num_sub_vectors.to_string(), - ); - writer.add_schema_metadata( - "lance:index_build:num_bits".to_string(), - trained.num_bits.to_string(), - ); - writer.add_schema_metadata( - "lance:index_build:dimension".to_string(), - trained.dimension.to_string(), - ); - writer.write_batch(&batch).await?; - writer.finish().await?; - Ok(()) -} - fn build_partition_batch( row_ids: Arc, partitions: &[u32], @@ -438,7 +333,6 @@ fn build_partition_batch( } struct TransformSlot { - input_host: PinnedHostBuffer, input_device: DeviceTensor, labels_host: PinnedHostBuffer, labels_device: DeviceTensor, @@ -448,10 +342,31 @@ struct TransformSlot { h2d_done: CudaEvent, transform_done: CudaEvent, output_ready: CudaEvent, + input_vectors: Option, + input_registration: Option, row_ids: Option>, rows: usize, } +struct PreparedTransformBatch { + row_ids: Arc, + vectors: FixedSizeListArray, + input_registration: Option, +} + +#[derive(Default)] +struct ArtifactPrepareStats { + input_batches: usize, + input_rows: usize, + scan_wait: Duration, + send: Duration, + vector: Duration, + filter: Duration, + matrix: Duration, + register: Duration, + registered_bytes: usize, +} + impl TransformSlot { fn try_new( resources: &Resources, @@ -460,7 +375,6 @@ impl TransformSlot { code_width: usize, ) -> Result { Ok(Self { - input_host: PinnedHostBuffer::try_new(max_rows * dimension)?, input_device: DeviceTensor::try_new(resources, &[max_rows, dimension])?, labels_host: PinnedHostBuffer::try_new(max_rows)?, labels_device: DeviceTensor::try_new(resources, &[max_rows])?, @@ -470,6 +384,8 @@ impl TransformSlot { h2d_done: CudaEvent::try_new()?, transform_done: CudaEvent::try_new()?, output_ready: CudaEvent::try_new()?, + input_vectors: None, + input_registration: None, row_ids: None, rows: 0, }) @@ -483,25 +399,41 @@ impl TransformSlot { &mut self, trained: &TrainedIvfPqIndex, stream: cuvs_sys::cudaStream_t, - row_ids: Arc, - matrix: &[f32], - rows: usize, - dimension: usize, + prepared: PreparedTransformBatch, ) -> Result<()> { let code_width = trained.pq_code_width(); - self.input_host.copy_from_slice(matrix)?; + let row_ids = prepared.row_ids; + let vectors = prepared.vectors; + let matrix = matrix_from_vectors(&vectors)?; + let (input_slice, rows, dimension, keep_input_vectors) = match &matrix { + MatrixBuffer::Borrowed { values, rows, cols } => (*values, *rows, *cols, true), + MatrixBuffer::Owned(array) => ( + array + .as_slice_memory_order() + .ok_or_else(|| Error::io("transform matrix is not contiguous"))?, + array.nrows(), + array.ncols(), + false, + ), + }; + self.input_device.set_shape(&[rows, dimension])?; self.labels_device.set_shape(&[rows])?; self.codes_device.set_shape(&[rows, code_width])?; self.rows = rows; self.row_ids = Some(row_ids); + self.input_registration = prepared.input_registration; self.h2d_start.record(stream)?; - self.input_device.copy_from_host_async( - &trained.resources, - self.input_host.prefix(rows * dimension)?, - )?; + self.input_device + .copy_from_host_async(&trained.resources, input_slice)?; self.h2d_done.record(stream)?; + if keep_input_vectors { + self.input_vectors = Some(vectors); + } else { + self.h2d_done.synchronize()?; + self.input_vectors = None; + } check_cuvs( unsafe { cuvs_sys::cuvsIvfPqTransform( @@ -531,6 +463,8 @@ impl TransformSlot { } self.output_ready.synchronize()?; + self.input_registration = None; + self.input_vectors = None; let row_ids = self .row_ids .take() @@ -546,31 +480,212 @@ impl TransformSlot { } } -async fn for_each_transformed_batch( - dataset: &Dataset, - column: &str, - trained: &TrainedIvfPqIndex, +#[derive(Default)] +struct ArtifactBuildStats { + input_batches: usize, + input_rows: usize, + output_batches: usize, + output_rows: usize, + scan_wait: Duration, + drain: Duration, + send: Duration, + prepare_send: Duration, + vector: Duration, + filter: Duration, + matrix: Duration, + launch: Duration, + register: Duration, + registered_bytes: usize, +} + +impl ArtifactBuildStats { + fn merge_prepare(&mut self, prepare: ArtifactPrepareStats) { + self.input_batches += prepare.input_batches; + self.input_rows += prepare.input_rows; + self.scan_wait += prepare.scan_wait; + self.prepare_send += prepare.send; + self.vector += prepare.vector; + self.filter += prepare.filter; + self.matrix += prepare.matrix; + self.register += prepare.register; + self.registered_bytes += prepare.registered_bytes; + } + + fn record_output(&mut self, batch: &RecordBatch) { + self.output_batches += 1; + self.output_rows += batch.num_rows(); + } + + fn log(&self) { + eprintln!( + "cuVS artifact stages: input_batches={} input_rows={} output_batches={} output_rows={} scan_wait_s={:.3} drain_s={:.3} send_s={:.3} prepare_send_s={:.3} vector_s={:.3} filter_s={:.3} matrix_s={:.3} launch_s={:.3}", + self.input_batches, + self.input_rows, + self.output_batches, + self.output_rows, + secs(self.scan_wait), + secs(self.drain), + secs(self.send), + secs(self.prepare_send), + secs(self.vector), + secs(self.filter), + secs(self.matrix), + secs(self.launch), + ); + eprintln!( + "cuVS artifact h2d registration: register_s={:.3} registered_gib={:.3}", + secs(self.register), + self.registered_bytes as f64 / (1024.0 * 1024.0 * 1024.0), + ); + } +} + +fn secs(duration: Duration) -> f64 { + duration.as_secs_f64() +} + +async fn prepare_transform_batches( + dataset: Dataset, + column: String, batch_size: usize, filter_nan: bool, - mut on_batch: F, -) -> Result<()> -where - F: FnMut(RecordBatch) -> Fut, - Fut: std::future::Future>, -{ - let code_width = trained.pq_code_width(); + mut prepared_tx: mpsc::Sender, +) -> Result { let mut scanner = dataset.scan(); - scanner.project(&[column])?; + scanner.project(&[&column])?; if dataset .schema() - .field(column) + .field(&column) .is_some_and(|field| field.nullable && filter_nan) { scanner.filter(&format!("{column} is not null"))?; } scanner.with_row_id(); scanner.batch_size(batch_size); + scanner.scan_in_order(false); + scanner.fragment_readahead(DEFAULT_SCAN_FRAGMENT_READAHEAD); + scanner.io_buffer_size(DEFAULT_SCAN_IO_BUFFER_SIZE); let mut stream = scanner.try_into_stream().await?; + let mut stats = ArtifactPrepareStats::default(); + + loop { + let scan_start = Instant::now(); + let Some(batch) = stream.try_next().await? else { + stats.scan_wait += scan_start.elapsed(); + break; + }; + stats.scan_wait += scan_start.elapsed(); + stats.input_batches += 1; + stats.input_rows += batch.num_rows(); + + let vector_start = Instant::now(); + let vectors = vector_column_to_fsl(&batch, &column)?; + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::invalid_input(format!("transform batch is missing {ROW_ID}")))?; + stats.vector += vector_start.elapsed(); + + let filter_start = Instant::now(); + let (filtered_row_ids, filtered_vectors) = if filter_nan { + let finite_mask = is_finite(&vectors); + let valid_rows = finite_mask.true_count(); + if valid_rows == 0 { + continue; + } + if valid_rows != vectors.len() { + warn!( + "{} vectors are ignored during partition assignment because they are null or non-finite", + vectors.len() - valid_rows + ); + } + + let filtered_row_ids = if valid_rows == row_ids.len() { + row_ids.clone() + } else { + filter(row_ids.as_ref(), &finite_mask)? + }; + let filtered_vectors = if valid_rows == vectors.len() { + vectors + } else { + let vector_column = batch.column_by_name(&column).ok_or_else(|| { + Error::invalid_input(format!( + "transform batch is missing vector column '{column}'" + )) + })?; + let field = batch + .schema() + .field_with_name(&column) + .map_err(|_| { + Error::invalid_input(format!( + "transform batch schema is missing field '{column}'" + )) + })? + .clone(); + let filtered_vectors = filter(vector_column.as_ref(), &finite_mask)?; + vector_column_to_fsl( + &RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![field])), + vec![filtered_vectors], + )?, + &column, + )? + }; + (filtered_row_ids, filtered_vectors) + } else { + (row_ids.clone(), vectors) + }; + stats.filter += filter_start.elapsed(); + + let matrix_start = Instant::now(); + let matrix = matrix_from_vectors(&filtered_vectors)?; + let input_registration = match &matrix { + MatrixBuffer::Borrowed { values, .. } => { + let register_start = Instant::now(); + let registration = match RegisteredHostBuffer::try_new(values) { + Ok(registration) => Some(registration), + Err(error) => { + warn!( + "failed to register host vector buffer for CUDA H2D; falling back to pageable memory: {error}" + ); + None + } + }; + stats.register += register_start.elapsed(); + stats.registered_bytes += registration + .as_ref() + .map(RegisteredHostBuffer::original_bytes) + .unwrap_or_default(); + registration + } + MatrixBuffer::Owned(_) => None, + }; + stats.matrix += matrix_start.elapsed(); + + let prepared = PreparedTransformBatch { + row_ids: filtered_row_ids, + vectors: filtered_vectors, + input_registration, + }; + let send_start = Instant::now(); + prepared_tx + .send(prepared) + .await + .map_err(|error| Error::io(format!("failed to forward prepared batch: {error}")))?; + stats.send += send_start.elapsed(); + } + + Ok(stats) +} + +async fn append_transformed_batches_to_artifact( + dataset: &Dataset, + column: &str, + trained: &TrainedIvfPqIndex, + batch_size: usize, + filter_nan: bool, + append_tx: &mut mpsc::Sender>, +) -> Result<()> { + let code_width = trained.pq_code_width(); let cuda_stream = trained .resources .get_cuda_stream() @@ -586,89 +701,96 @@ where }) .collect::>>()?; let mut next_slot = 0usize; - - loop { - let Some(batch) = stream.try_next().await? else { - break; - }; + let mut stats = ArtifactBuildStats::default(); + let (prepared_tx, mut prepared_rx) = mpsc::channel::(PIPELINE_SLOTS); + let prepare_task = tokio::spawn(prepare_transform_batches( + dataset.clone(), + column.to_string(), + batch_size, + filter_nan, + prepared_tx, + )); + + while let Some(prepared) = prepared_rx.next().await { let slot = &mut slots[next_slot]; - if let Some(transformed) = slot.drain_to_batch(code_width)? { - on_batch(transformed).await?; - } - - let vectors = vector_column_to_fsl(&batch, column)?; - let row_ids = batch - .column_by_name(ROW_ID) - .ok_or_else(|| Error::invalid_input(format!("transform batch is missing {ROW_ID}")))?; - let finite_mask = is_finite(&vectors); - let valid_rows = finite_mask.true_count(); - if valid_rows == 0 { - continue; - } - if valid_rows != vectors.len() { - warn!( - "{} vectors are ignored during partition assignment because they are null or non-finite", - vectors.len() - valid_rows - ); - } - - let filtered_row_ids = if valid_rows == row_ids.len() { - row_ids.clone() - } else { - filter(row_ids.as_ref(), &finite_mask)? - }; - let filtered_vectors = if valid_rows == vectors.len() { - vectors + let drain_start = Instant::now(); + let transformed = if let Some(transformed) = slot.drain_to_batch(code_width)? { + stats.drain += drain_start.elapsed(); + stats.record_output(&transformed); + Some(transformed) } else { - let vector_column = batch.column_by_name(column).ok_or_else(|| { - Error::invalid_input(format!( - "transform batch is missing vector column '{column}'" - )) - })?; - let field = batch - .schema() - .field_with_name(column) - .map_err(|_| { - Error::invalid_input(format!( - "transform batch schema is missing field '{column}'" - )) - })? - .clone(); - let filtered_vectors = filter(vector_column.as_ref(), &finite_mask)?; - vector_column_to_fsl( - &RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![field])), - vec![filtered_vectors], - )?, - column, - )? + stats.drain += drain_start.elapsed(); + None }; - let matrix = matrix_from_vectors(&filtered_vectors)?; - let matrix_view = matrix.view()?; - let input_slice = matrix_view - .as_slice_memory_order() - .ok_or_else(|| Error::io("transform matrix is not contiguous"))?; - - slot.launch( - trained, - cuda_stream, - filtered_row_ids, - input_slice, - matrix.rows(), - matrix_view.ncols(), - )?; + let launch_start = Instant::now(); + slot.launch(trained, cuda_stream, prepared)?; + stats.launch += launch_start.elapsed(); + + if let Some(transformed) = transformed { + let send_start = Instant::now(); + append_tx.send(Ok(transformed)).await.map_err(|error| { + Error::io(format!("failed to forward transformed batch: {error}")) + })?; + stats.send += send_start.elapsed(); + } next_slot = (next_slot + 1) % PIPELINE_SLOTS; } + let prepare_stats = prepare_task + .await + .map_err(|error| Error::io(format!("prepare transform task failed: {error}")))??; + stats.merge_prepare(prepare_stats); for slot in &mut slots { + let drain_start = Instant::now(); if let Some(transformed) = slot.drain_to_batch(code_width)? { - on_batch(transformed).await?; + stats.drain += drain_start.elapsed(); + stats.record_output(&transformed); + let send_start = Instant::now(); + append_tx.send(Ok(transformed)).await.map_err(|error| { + Error::io(format!("failed to forward transformed batch: {error}")) + })?; + stats.send += send_start.elapsed(); + } else { + stats.drain += drain_start.elapsed(); } } + stats.log(); Ok(()) } +async fn append_artifact_batches( + mut artifact: PartitionArtifactBuilder, + mut rx: mpsc::Receiver>, +) -> Result> { + let mut batches = 0usize; + let mut rows = 0usize; + let mut append_time = Duration::default(); + while let Some(batch) = rx.next().await { + let batch = batch?; + batches += 1; + rows += batch.num_rows(); + let append_start = Instant::now(); + artifact.append_batch(&batch).await?; + append_time += append_start.elapsed(); + } + + let finish_start = Instant::now(); + let files = artifact + .finish(PARTITION_ARTIFACT_METADATA_FILE_NAME, None) + .await?; + let finish_time = finish_start.elapsed(); + eprintln!( + "cuVS artifact append task: batches={} rows={} append_s={:.3} finish_s={:.3} files={}", + batches, + rows, + secs(append_time), + secs(finish_time), + files.len() + ); + Ok(files) +} + /// Train an IVF_PQ model with cuVS and return Arrow-native training outputs. /// /// This function performs only the backend-owned training step. The returned @@ -866,78 +988,39 @@ pub async fn assign_ivf_pq_to_artifact( filter_nan: bool, storage_options: Option<&HashMap>, ) -> Result> { - let registry = Arc::new(lance_io::object_store::ObjectStoreRegistry::default()); - let params = if let Some(storage_options) = storage_options { - lance_io::object_store::ObjectStoreParams { - storage_options_accessor: Some(Arc::new( - lance_io::object_store::StorageOptionsAccessor::with_static_options( - storage_options.clone(), - ), - )), - ..Default::default() - } - } else { - lance_io::object_store::ObjectStoreParams::default() - }; - let (object_store, root_dir) = - lance::io::ObjectStore::from_uri_and_params(registry, artifact_uri, ¶ms) - .await - .map_err(|error| Error::io(error.to_string()))?; - if !object_store.is_local() { - return Err(Error::not_supported( - "partition artifact builds currently require a local filesystem artifact_uri", - )); + let artifact = PartitionArtifactBuilder::try_new( + artifact_uri, + trained.num_partitions, + trained.pq_code_width(), + storage_options, + ) + .await?; + + let (mut append_tx, append_rx) = mpsc::channel::>(PIPELINE_SLOTS); + let append_task = tokio::spawn(append_artifact_batches(artifact, append_rx)); + + let append_start = Instant::now(); + let append_result = append_transformed_batches_to_artifact( + dataset, + column, + trained, + batch_size, + filter_nan, + &mut append_tx, + ) + .await; + drop(append_tx); + if let Err(error) = append_result { + append_task.abort(); + return Err(error); } - - std::fs::create_dir_all(to_local_path(&root_dir)) - .map_err(|error| Error::io(format!("failed to create artifact directory: {error}")))?; - - let (tx, rx) = mpsc::channel::>(PIPELINE_SLOTS); - let mut shuffler = IvfShuffler::try_new( - trained.num_partitions as u32, - Some(root_dir.clone()), - true, - None, - )?; - shuffler = shuffler.with_format_version(LanceFileVersion::V2_0); - - let shuffle_task = tokio::spawn(async move { - shuffler.write_unsorted_stream(rx).await?; - shuffler - .write_partitioned_shuffles(PARTITION_ARTIFACT_SHUFFLE_BATCHES_PER_FILE, PIPELINE_SLOTS) - .await - }); - - let produce_result = - for_each_transformed_batch(dataset, column, trained, batch_size, filter_nan, |batch| { - let mut tx = tx.clone(); - async move { - tx.send(Ok(batch)).await.map_err(|error| { - Error::io(format!("failed to forward transformed batch: {error}")) - }) - } - }) - .await; - drop(tx); - - produce_result?; - write_partition_artifact_metadata(artifact_uri, trained, storage_options).await?; - - let mut files = shuffle_task + eprintln!( + "cuVS artifact append_transformed_batches time: {:.3}s", + append_start.elapsed().as_secs_f64() + ); + let files = append_task .await - .map_err(|error| Error::io(format!("partition artifact shuffle task failed: {error}")))??; - - let unsorted_path = root_dir.child(PARTITION_ARTIFACT_UNSORTED_FILE_NAME); - let unsorted_local_path = to_local_path(&unsorted_path); - if std::path::Path::new(&unsorted_local_path).exists() { - std::fs::remove_file(&unsorted_local_path).map_err(|error| { - Error::io(format!( - "failed to remove temporary unsorted buffer: {error}" - )) - })?; - } - - files.push(PARTITION_ARTIFACT_METADATA_FILE_NAME.to_string()); + .map_err(|error| Error::io(format!("partition artifact append task failed: {error}")))??; Ok(files) } diff --git a/backends/cuvs_26_02/src/cuda.rs b/backends/cuvs_26_02/src/cuda.rs index 94d36f6..efa61a8 100644 --- a/backends/cuvs_26_02/src/cuda.rs +++ b/backends/cuvs_26_02/src/cuda.rs @@ -20,6 +20,8 @@ pub(crate) type CudaEventHandle = *mut c_void; unsafe extern "C" { fn cudaMallocHost(ptr: *mut *mut c_void, size: usize) -> cuvs_sys::cudaError_t; fn cudaFreeHost(ptr: *mut c_void) -> cuvs_sys::cudaError_t; + fn cudaHostRegister(ptr: *mut c_void, size: usize, flags: u32) -> cuvs_sys::cudaError_t; + fn cudaHostUnregister(ptr: *mut c_void) -> cuvs_sys::cudaError_t; fn cudaEventCreate(event: *mut CudaEventHandle) -> cuvs_sys::cudaError_t; fn cudaEventDestroy(event: CudaEventHandle) -> cuvs_sys::cudaError_t; fn cudaEventRecord( @@ -71,13 +73,6 @@ impl MatrixBuffer<'_> { Self::Owned(array) => Ok(array.view()), } } - - pub(crate) fn rows(&self) -> usize { - match self { - Self::Borrowed { rows, .. } => *rows, - Self::Owned(array) => array.nrows(), - } - } } pub(crate) struct HostTensorView { @@ -290,6 +285,75 @@ impl Drop for DeviceTensor { } } +pub(crate) struct RegisteredHostBuffer { + ptr: *mut c_void, + original_bytes: usize, +} + +// CUDA host registration owns a process-local address range; unregistering it +// from the consumer task is safe because CUDA runtime calls are thread-safe. +unsafe impl Send for RegisteredHostBuffer {} + +impl RegisteredHostBuffer { + pub(crate) fn try_new(slice: &[T]) -> Result { + let original_bytes = std::mem::size_of_val(slice); + if original_bytes == 0 { + return Ok(Self { + ptr: ptr::null_mut(), + original_bytes: 0, + }); + } + + let page_size = page_size()?; + let start = slice.as_ptr() as usize; + let end = start + .checked_add(original_bytes) + .ok_or_else(|| Error::io("registered host buffer size overflow"))?; + let aligned_start = start & !(page_size - 1); + let aligned_end = end + .checked_add(page_size - 1) + .ok_or_else(|| Error::io("registered host buffer alignment overflow"))? + & !(page_size - 1); + let bytes = aligned_end - aligned_start; + let ptr = aligned_start as *mut c_void; + + check_cuda( + unsafe { cudaHostRegister(ptr, bytes, 0) }, + "register host buffer", + )?; + Ok(Self { + ptr, + original_bytes, + }) + } + + pub(crate) fn original_bytes(&self) -> usize { + self.original_bytes + } +} + +impl Drop for RegisteredHostBuffer { + fn drop(&mut self) { + if !self.ptr.is_null() { + let _ = unsafe { cudaHostUnregister(self.ptr) }; + } + } +} + +fn page_size() -> Result { + let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) }; + if page_size <= 0 { + return Err(Error::io("failed to resolve system page size")); + } + let page_size = page_size as usize; + if !page_size.is_power_of_two() { + return Err(Error::io(format!( + "system page size {page_size} is not a power of two" + ))); + } + Ok(page_size) +} + pub(crate) struct PinnedHostBuffer { ptr: *mut T, len: usize, @@ -340,18 +404,6 @@ impl PinnedHostBuffer { } Ok(&mut self.as_mut_slice()[..len]) } - - pub(crate) fn copy_from_slice(&mut self, src: &[T]) -> Result<()> { - if src.len() > self.len { - return Err(Error::io(format!( - "pinned host buffer length {} is smaller than source length {}", - self.len, - src.len() - ))); - } - self.prefix_mut(src.len())?.copy_from_slice(src); - Ok(()) - } } impl Drop for PinnedHostBuffer { From a9c5b2012b3e99cdb48c3833fb8bb9676c5ef90e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 22 Jun 2026 21:11:21 +0800 Subject: [PATCH 2/8] perf: avoid random training reads for cuVS IVF-PQ --- backends/cuvs_26_02/src/backend.rs | 92 ++++++++++++++++++++++++------ 1 file changed, 74 insertions(+), 18 deletions(-) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index 167fc3a..2e9dedf 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -7,12 +7,14 @@ use crate::cuda::{ create_index_params, destroy_index_params, ivf_centroids_from_host, make_tensor_view, matrix_from_vectors, pq_codebook_from_host, }; -use arrow::compute::filter; +use arrow::compute::{concat_batches, filter}; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use cuvs::Resources; -use futures::{FutureExt, SinkExt, StreamExt, TryStreamExt, channel::mpsc, future::LocalBoxFuture}; +use futures::{ + FutureExt, SinkExt, StreamExt, TryStreamExt, channel::mpsc, future::LocalBoxFuture, stream, +}; use lance::dataset::Dataset; use lance::index::vector::PartitionArtifactBuilder; use lance::index::vector::utils::infer_vector_dim; @@ -23,6 +25,7 @@ use lance_index::vector::{PART_ID_COLUMN, PQ_CODE_COLUMN}; use lance_linalg::distance::DistanceType; use log::warn; use std::collections::HashMap; +use std::ops::Range; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -30,6 +33,8 @@ const PARTITION_ARTIFACT_METADATA_FILE_NAME: &str = "metadata.lance"; const PIPELINE_SLOTS: usize = 2; const DEFAULT_SCAN_FRAGMENT_READAHEAD: usize = 64; const DEFAULT_SCAN_IO_BUFFER_SIZE: u64 = 16 * 1024 * 1024 * 1024; +const TRAINING_SAMPLE_CHUNK_ROWS: usize = 8 * 1024; +const TRAINING_SAMPLE_BATCH_READAHEAD: usize = 64; /// A trained cuVS IVF_PQ model that can be reused for artifact builds. /// @@ -544,6 +549,62 @@ fn secs(duration: Duration) -> f64 { duration.as_secs_f64() } +fn training_sample_ranges(num_rows: usize, sample_rows: usize) -> Vec> { + let sample_rows = sample_rows.min(num_rows); + if sample_rows == 0 { + return Vec::new(); + } + if sample_rows == num_rows { + return vec![0..num_rows as u64]; + } + + let chunk_rows = TRAINING_SAMPLE_CHUNK_ROWS.min(sample_rows); + let num_chunks = sample_rows.div_ceil(chunk_rows); + let mut remaining = sample_rows; + let mut ranges = Vec::with_capacity(num_chunks); + for chunk_idx in 0..num_chunks { + let rows = chunk_rows.min(remaining); + remaining -= rows; + let max_start = num_rows - rows; + let start = if num_chunks == 1 { + max_start / 2 + } else { + ((chunk_idx as u128 * max_start as u128) / (num_chunks - 1) as u128) as usize + }; + ranges.push(start as u64..(start + rows) as u64); + } + ranges +} + +async fn sample_training_vectors( + dataset: &Dataset, + column: &str, + sample_rows: usize, +) -> Result { + let num_rows = dataset.count_rows(None).await?; + if num_rows == 0 { + return Err(Error::invalid_input( + "cuVS training requires at least one training vector", + )); + } + + let ranges = training_sample_ranges(num_rows, sample_rows); + let projection = Arc::new(dataset.schema().project(&[column])?); + let stream = dataset.take_scan( + Box::pin(stream::iter(ranges.into_iter().map(Ok))), + projection, + TRAINING_SAMPLE_BATCH_READAHEAD, + ); + let batches = stream.try_collect::>().await?; + let Some(schema) = batches.first().map(RecordBatch::schema) else { + return Err(Error::invalid_input( + "cuVS training sample did not return any vectors", + )); + }; + let batch = concat_batches(&schema, &batches)?; + Ok(vector_column_to_fsl(&batch, column)?) +} + async fn prepare_transform_batches( dataset: Dataset, column: String, @@ -850,25 +911,20 @@ pub async fn train_ivf_pq( ))); } - let num_rows = dataset.count_rows(None).await?; - if num_rows == 0 { - return Err(Error::invalid_input( - "cuVS training requires at least one training vector", - )); - } - let train_rows = num_rows - .min((num_partitions * sample_rate).max(256 * 256)) - .max(1); + let train_rows = (num_partitions * sample_rate).max(256 * 256).max(1); + let sample_start = Instant::now(); + let train_vectors = sample_training_vectors(dataset, column, train_rows).await?; + eprintln!( + "cuVS train sample time: {:.3}s rows={}", + sample_start.elapsed().as_secs_f64(), + train_vectors.len() + ); let train_vectors = if filter_nan { - let batch = dataset.scan().project(&[column])?.try_into_batch().await?; - let vectors = vector_column_to_fsl(&batch, column)?; - let mask = is_finite(&vectors); - let filtered = filter(&vectors, &mask)?.as_fixed_size_list().clone(); + let mask = is_finite(&train_vectors); + let filtered = filter(&train_vectors, &mask)?.as_fixed_size_list().clone(); filtered.slice(0, train_rows.min(filtered.len())) } else { - let projection = dataset.schema().project(&[column])?; - let batch = dataset.sample(train_rows, &projection, None).await?; - vector_column_to_fsl(&batch, column)? + train_vectors }; if train_vectors.is_empty() { return Err(Error::invalid_input( From 90d14cc36a81be008438443169f09549e6e763c4 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 23 Jun 2026 05:08:36 +0800 Subject: [PATCH 3/8] perf: pipeline cuVS artifact preparation --- backends/cuvs_26_02/src/backend.rs | 416 +++++++++++++++++++++-------- 1 file changed, 299 insertions(+), 117 deletions(-) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index 2e9dedf..61b0607 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -9,9 +9,11 @@ use crate::cuda::{ }; use arrow::compute::{concat_batches, filter}; use arrow_array::cast::AsArray; +use arrow_array::types::Float32Type; use arrow_array::{Array, ArrayRef, FixedSizeListArray, RecordBatch, UInt8Array, UInt32Array}; use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use cuvs::Resources; +use futures::lock::Mutex; use futures::{ FutureExt, SinkExt, StreamExt, TryStreamExt, channel::mpsc, future::LocalBoxFuture, stream, }; @@ -24,6 +26,7 @@ use lance_index::vector::utils::is_finite; use lance_index::vector::{PART_ID_COLUMN, PQ_CODE_COLUMN}; use lance_linalg::distance::DistanceType; use log::warn; +use ndarray::Array2; use std::collections::HashMap; use std::ops::Range; use std::sync::Arc; @@ -31,8 +34,9 @@ use std::time::{Duration, Instant}; const PARTITION_ARTIFACT_METADATA_FILE_NAME: &str = "metadata.lance"; const PIPELINE_SLOTS: usize = 2; -const DEFAULT_SCAN_FRAGMENT_READAHEAD: usize = 64; +const DEFAULT_SCAN_FRAGMENT_READAHEAD: usize = 0; const DEFAULT_SCAN_IO_BUFFER_SIZE: u64 = 16 * 1024 * 1024 * 1024; +const DEFAULT_PREPARE_WORKERS: usize = 1; const TRAINING_SAMPLE_CHUNK_ROWS: usize = 8 * 1024; const TRAINING_SAMPLE_BATCH_READAHEAD: usize = 64; @@ -348,23 +352,70 @@ struct TransformSlot { transform_done: CudaEvent, output_ready: CudaEvent, input_vectors: Option, + input_matrix: Option>, input_registration: Option, row_ids: Option>, rows: usize, } +enum PreparedMatrix { + F32Arrow { + vectors: FixedSizeListArray, + rows: usize, + dimension: usize, + }, + Owned(Array2), +} + +impl PreparedMatrix { + fn rows(&self) -> usize { + match self { + Self::F32Arrow { rows, .. } => *rows, + Self::Owned(array) => array.nrows(), + } + } + + fn dimension(&self) -> usize { + match self { + Self::F32Arrow { dimension, .. } => *dimension, + Self::Owned(array) => array.ncols(), + } + } + + fn input_slice(&self) -> Result<&[f32]> { + match self { + Self::F32Arrow { vectors, .. } => { + let values = vectors.values().as_primitive::(); + Ok(values.values().as_ref()) + } + Self::Owned(array) => array + .as_slice_memory_order() + .ok_or_else(|| Error::io("transform matrix is not contiguous")), + } + } +} + struct PreparedTransformBatch { row_ids: Arc, - vectors: FixedSizeListArray, + matrix: PreparedMatrix, input_registration: Option, } #[derive(Default)] -struct ArtifactPrepareStats { +struct ArtifactScannerStats { input_batches: usize, input_rows: usize, scan_wait: Duration, send: Duration, +} + +#[derive(Default)] +struct ArtifactPrepareStats { + workers: usize, + input_batches: usize, + input_rows: usize, + raw_wait: Duration, + send: Duration, vector: Duration, filter: Duration, matrix: Duration, @@ -390,6 +441,7 @@ impl TransformSlot { transform_done: CudaEvent::try_new()?, output_ready: CudaEvent::try_new()?, input_vectors: None, + input_matrix: None, input_registration: None, row_ids: None, rows: 0, @@ -408,19 +460,10 @@ impl TransformSlot { ) -> Result<()> { let code_width = trained.pq_code_width(); let row_ids = prepared.row_ids; - let vectors = prepared.vectors; - let matrix = matrix_from_vectors(&vectors)?; - let (input_slice, rows, dimension, keep_input_vectors) = match &matrix { - MatrixBuffer::Borrowed { values, rows, cols } => (*values, *rows, *cols, true), - MatrixBuffer::Owned(array) => ( - array - .as_slice_memory_order() - .ok_or_else(|| Error::io("transform matrix is not contiguous"))?, - array.nrows(), - array.ncols(), - false, - ), - }; + let matrix = prepared.matrix; + let rows = matrix.rows(); + let dimension = matrix.dimension(); + let input_slice = matrix.input_slice()?; self.input_device.set_shape(&[rows, dimension])?; self.labels_device.set_shape(&[rows])?; @@ -433,11 +476,15 @@ impl TransformSlot { self.input_device .copy_from_host_async(&trained.resources, input_slice)?; self.h2d_done.record(stream)?; - if keep_input_vectors { - self.input_vectors = Some(vectors); - } else { - self.h2d_done.synchronize()?; - self.input_vectors = None; + match matrix { + PreparedMatrix::F32Arrow { vectors, .. } => { + self.input_vectors = Some(vectors); + self.input_matrix = None; + } + PreparedMatrix::Owned(array) => { + self.input_vectors = None; + self.input_matrix = Some(array); + } } check_cuvs( unsafe { @@ -470,6 +517,7 @@ impl TransformSlot { self.output_ready.synchronize()?; self.input_registration = None; self.input_vectors = None; + self.input_matrix = None; let row_ids = self .row_ids .take() @@ -487,11 +535,17 @@ impl TransformSlot { #[derive(Default)] struct ArtifactBuildStats { + scanner_tasks: usize, + prepare_workers: usize, input_batches: usize, input_rows: usize, + prepared_batches: usize, + prepared_rows: usize, output_batches: usize, output_rows: usize, scan_wait: Duration, + raw_send: Duration, + raw_wait: Duration, drain: Duration, send: Duration, prepare_send: Duration, @@ -504,10 +558,19 @@ struct ArtifactBuildStats { } impl ArtifactBuildStats { + fn merge_scanner(&mut self, scanner: ArtifactScannerStats) { + self.scanner_tasks += 1; + self.input_batches += scanner.input_batches; + self.input_rows += scanner.input_rows; + self.scan_wait += scanner.scan_wait; + self.raw_send += scanner.send; + } + fn merge_prepare(&mut self, prepare: ArtifactPrepareStats) { - self.input_batches += prepare.input_batches; - self.input_rows += prepare.input_rows; - self.scan_wait += prepare.scan_wait; + self.prepare_workers += prepare.workers; + self.prepared_batches += prepare.input_batches; + self.prepared_rows += prepare.input_rows; + self.raw_wait += prepare.raw_wait; self.prepare_send += prepare.send; self.vector += prepare.vector; self.filter += prepare.filter; @@ -523,12 +586,18 @@ impl ArtifactBuildStats { fn log(&self) { eprintln!( - "cuVS artifact stages: input_batches={} input_rows={} output_batches={} output_rows={} scan_wait_s={:.3} drain_s={:.3} send_s={:.3} prepare_send_s={:.3} vector_s={:.3} filter_s={:.3} matrix_s={:.3} launch_s={:.3}", + "cuVS artifact stages: scanner_tasks={} prepare_workers={} input_batches={} input_rows={} prepared_batches={} prepared_rows={} output_batches={} output_rows={} scan_wait_s={:.3} raw_send_s={:.3} raw_wait_s={:.3} drain_s={:.3} send_s={:.3} prepare_send_s={:.3} vector_s={:.3} filter_s={:.3} matrix_s={:.3} launch_s={:.3}", + self.scanner_tasks, + self.prepare_workers, self.input_batches, self.input_rows, + self.prepared_batches, + self.prepared_rows, self.output_batches, self.output_rows, secs(self.scan_wait), + secs(self.raw_send), + secs(self.raw_wait), secs(self.drain), secs(self.send), secs(self.prepare_send), @@ -549,6 +618,21 @@ fn secs(duration: Duration) -> f64 { duration.as_secs_f64() } +fn prepare_workers_from_env() -> usize { + std::env::var("LANCE_CUVS_PREPARE_WORKERS") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|workers| *workers > 0) + .unwrap_or(DEFAULT_PREPARE_WORKERS) +} + +fn scan_fragment_readahead_from_env() -> usize { + std::env::var("LANCE_CUVS_SCAN_FRAGMENT_READAHEAD") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_SCAN_FRAGMENT_READAHEAD) +} + fn training_sample_ranges(num_rows: usize, sample_rows: usize) -> Vec> { let sample_rows = sample_rows.min(num_rows); if sample_rows == 0 { @@ -605,13 +689,13 @@ async fn sample_training_vectors( Ok(vector_column_to_fsl(&batch, column)?) } -async fn prepare_transform_batches( +async fn scan_transform_batches( dataset: Dataset, column: String, batch_size: usize, filter_nan: bool, - mut prepared_tx: mpsc::Sender, -) -> Result { + mut raw_tx: mpsc::Sender, +) -> Result { let mut scanner = dataset.scan(); scanner.project(&[&column])?; if dataset @@ -624,10 +708,10 @@ async fn prepare_transform_batches( scanner.with_row_id(); scanner.batch_size(batch_size); scanner.scan_in_order(false); - scanner.fragment_readahead(DEFAULT_SCAN_FRAGMENT_READAHEAD); + scanner.fragment_readahead(scan_fragment_readahead_from_env()); scanner.io_buffer_size(DEFAULT_SCAN_IO_BUFFER_SIZE); let mut stream = scanner.try_into_stream().await?; - let mut stats = ArtifactPrepareStats::default(); + let mut stats = ArtifactScannerStats::default(); loop { let scan_start = Instant::now(); @@ -639,93 +723,165 @@ async fn prepare_transform_batches( stats.input_batches += 1; stats.input_rows += batch.num_rows(); - let vector_start = Instant::now(); - let vectors = vector_column_to_fsl(&batch, &column)?; - let row_ids = batch - .column_by_name(ROW_ID) - .ok_or_else(|| Error::invalid_input(format!("transform batch is missing {ROW_ID}")))?; - stats.vector += vector_start.elapsed(); - - let filter_start = Instant::now(); - let (filtered_row_ids, filtered_vectors) = if filter_nan { - let finite_mask = is_finite(&vectors); - let valid_rows = finite_mask.true_count(); - if valid_rows == 0 { - continue; - } - if valid_rows != vectors.len() { - warn!( - "{} vectors are ignored during partition assignment because they are null or non-finite", - vectors.len() - valid_rows - ); - } + let send_start = Instant::now(); + raw_tx + .send(batch) + .await + .map_err(|error| Error::io(format!("failed to forward raw batch: {error}")))?; + stats.send += send_start.elapsed(); + } - let filtered_row_ids = if valid_rows == row_ids.len() { - row_ids.clone() - } else { - filter(row_ids.as_ref(), &finite_mask)? - }; - let filtered_vectors = if valid_rows == vectors.len() { - vectors - } else { - let vector_column = batch.column_by_name(&column).ok_or_else(|| { + Ok(stats) +} + +fn prepare_transform_batch( + batch: RecordBatch, + column: &str, + filter_nan: bool, + stats: &mut ArtifactPrepareStats, +) -> Result> { + stats.input_batches += 1; + stats.input_rows += batch.num_rows(); + + let vector_start = Instant::now(); + let vectors = vector_column_to_fsl(&batch, column)?; + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::invalid_input(format!("transform batch is missing {ROW_ID}")))?; + stats.vector += vector_start.elapsed(); + + let filter_start = Instant::now(); + let (filtered_row_ids, filtered_vectors) = if filter_nan { + let finite_mask = is_finite(&vectors); + let valid_rows = finite_mask.true_count(); + if valid_rows == 0 { + stats.filter += filter_start.elapsed(); + return Ok(None); + } + if valid_rows != vectors.len() { + warn!( + "{} vectors are ignored during partition assignment because they are null or non-finite", + vectors.len() - valid_rows + ); + } + + let filtered_row_ids = if valid_rows == row_ids.len() { + row_ids.clone() + } else { + filter(row_ids.as_ref(), &finite_mask)? + }; + let filtered_vectors = if valid_rows == vectors.len() { + vectors + } else { + let vector_column = batch.column_by_name(column).ok_or_else(|| { + Error::invalid_input(format!( + "transform batch is missing vector column '{column}'" + )) + })?; + let field = batch + .schema() + .field_with_name(column) + .map_err(|_| { Error::invalid_input(format!( - "transform batch is missing vector column '{column}'" + "transform batch schema is missing field '{column}'" )) - })?; - let field = batch - .schema() - .field_with_name(&column) - .map_err(|_| { - Error::invalid_input(format!( - "transform batch schema is missing field '{column}'" - )) - })? - .clone(); - let filtered_vectors = filter(vector_column.as_ref(), &finite_mask)?; - vector_column_to_fsl( - &RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![field])), - vec![filtered_vectors], - )?, - &column, - )? - }; - (filtered_row_ids, filtered_vectors) - } else { - (row_ids.clone(), vectors) + })? + .clone(); + let filtered_vectors = filter(vector_column.as_ref(), &finite_mask)?; + vector_column_to_fsl( + &RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![field])), + vec![filtered_vectors], + )?, + column, + )? }; - stats.filter += filter_start.elapsed(); - - let matrix_start = Instant::now(); - let matrix = matrix_from_vectors(&filtered_vectors)?; - let input_registration = match &matrix { - MatrixBuffer::Borrowed { values, .. } => { - let register_start = Instant::now(); - let registration = match RegisteredHostBuffer::try_new(values) { - Ok(registration) => Some(registration), - Err(error) => { - warn!( - "failed to register host vector buffer for CUDA H2D; falling back to pageable memory: {error}" - ); - None - } - }; - stats.register += register_start.elapsed(); - stats.registered_bytes += registration - .as_ref() - .map(RegisteredHostBuffer::original_bytes) - .unwrap_or_default(); - registration - } - MatrixBuffer::Owned(_) => None, + (filtered_row_ids, filtered_vectors) + } else { + (row_ids.clone(), vectors) + }; + stats.filter += filter_start.elapsed(); + + let matrix_start = Instant::now(); + let matrix = matrix_from_vectors(&filtered_vectors)?; + stats.matrix += matrix_start.elapsed(); + + let (prepared_matrix, input_registration) = match matrix { + MatrixBuffer::Borrowed { values, rows, cols } => { + let register_start = Instant::now(); + let registration = match RegisteredHostBuffer::try_new(values) { + Ok(registration) => Some(registration), + Err(error) => { + warn!( + "failed to register host vector buffer for CUDA H2D; falling back to pageable memory: {error}" + ); + None + } + }; + stats.register += register_start.elapsed(); + stats.registered_bytes += registration + .as_ref() + .map(RegisteredHostBuffer::original_bytes) + .unwrap_or_default(); + ( + PreparedMatrix::F32Arrow { + vectors: filtered_vectors, + rows, + dimension: cols, + }, + registration, + ) + } + MatrixBuffer::Owned(array) => (PreparedMatrix::Owned(array), None), + }; + + Ok(Some(PreparedTransformBatch { + row_ids: filtered_row_ids, + matrix: prepared_matrix, + input_registration, + })) +} + +async fn prepare_transform_batches( + column: String, + filter_nan: bool, + raw_rx: Arc>>, + mut prepared_tx: mpsc::Sender, +) -> Result { + let mut stats = ArtifactPrepareStats { + workers: 1, + ..Default::default() + }; + + loop { + let raw_wait_start = Instant::now(); + let batch = { + let mut raw_rx = raw_rx.lock().await; + raw_rx.next().await }; - stats.matrix += matrix_start.elapsed(); + stats.raw_wait += raw_wait_start.elapsed(); - let prepared = PreparedTransformBatch { - row_ids: filtered_row_ids, - vectors: filtered_vectors, - input_registration, + let Some(batch) = batch else { + break; + }; + let column = column.clone(); + let (prepared, batch_stats) = tokio::task::spawn_blocking(move || { + let mut batch_stats = ArtifactPrepareStats::default(); + let prepared = prepare_transform_batch(batch, &column, filter_nan, &mut batch_stats)?; + Ok::<_, Error>((prepared, batch_stats)) + }) + .await + .map_err(|error| Error::io(format!("prepare transform blocking task failed: {error}")))??; + stats.input_batches += batch_stats.input_batches; + stats.input_rows += batch_stats.input_rows; + stats.vector += batch_stats.vector; + stats.filter += batch_stats.filter; + stats.matrix += batch_stats.matrix; + stats.register += batch_stats.register; + stats.registered_bytes += batch_stats.registered_bytes; + + let Some(prepared) = prepared else { + continue; }; let send_start = Instant::now(); prepared_tx @@ -763,14 +919,34 @@ async fn append_transformed_batches_to_artifact( .collect::>>()?; let mut next_slot = 0usize; let mut stats = ArtifactBuildStats::default(); - let (prepared_tx, mut prepared_rx) = mpsc::channel::(PIPELINE_SLOTS); - let prepare_task = tokio::spawn(prepare_transform_batches( + let prepare_workers = prepare_workers_from_env(); + if prepare_workers > 1 { + eprintln!( + "cuVS artifact prepare: using {} workers behind a single scanner", + prepare_workers + ); + } + let (raw_tx, raw_rx) = mpsc::channel::(prepare_workers); + let raw_rx = Arc::new(Mutex::new(raw_rx)); + let scanner_task = tokio::spawn(scan_transform_batches( dataset.clone(), column.to_string(), batch_size, filter_nan, - prepared_tx, + raw_tx, )); + let (prepared_tx, mut prepared_rx) = mpsc::channel::(PIPELINE_SLOTS); + let prepare_tasks = (0..prepare_workers) + .map(|_| { + tokio::spawn(prepare_transform_batches( + column.to_string(), + filter_nan, + raw_rx.clone(), + prepared_tx.clone(), + )) + }) + .collect::>(); + drop(prepared_tx); while let Some(prepared) = prepared_rx.next().await { let slot = &mut slots[next_slot]; @@ -797,10 +973,16 @@ async fn append_transformed_batches_to_artifact( } next_slot = (next_slot + 1) % PIPELINE_SLOTS; } - let prepare_stats = prepare_task + for prepare_task in prepare_tasks { + let prepare_stats = prepare_task + .await + .map_err(|error| Error::io(format!("prepare transform task failed: {error}")))??; + stats.merge_prepare(prepare_stats); + } + let scanner_stats = scanner_task .await - .map_err(|error| Error::io(format!("prepare transform task failed: {error}")))??; - stats.merge_prepare(prepare_stats); + .map_err(|error| Error::io(format!("scanner transform task failed: {error}")))??; + stats.merge_scanner(scanner_stats); for slot in &mut slots { let drain_start = Instant::now(); From 4fd93010b05fff3b7dae8e6aaed0057c6120863e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 23 Jun 2026 11:48:14 +0800 Subject: [PATCH 4/8] perf: report cuVS artifact GPU event timings --- backends/cuvs_26_02/src/backend.rs | 57 ++++++++++++++++++++++++------ backends/cuvs_26_02/src/cuda.rs | 14 ++++++++ 2 files changed, 61 insertions(+), 10 deletions(-) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index 61b0607..d597ba0 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -401,6 +401,13 @@ struct PreparedTransformBatch { input_registration: Option, } +struct DrainedTransformBatch { + batch: RecordBatch, + h2d: Duration, + transform: Duration, + d2h: Duration, +} + #[derive(Default)] struct ArtifactScannerStats { input_batches: usize, @@ -509,12 +516,15 @@ impl TransformSlot { Ok(()) } - fn drain_to_batch(&mut self, code_width: usize) -> Result> { + fn drain_to_batch(&mut self, code_width: usize) -> Result> { if !self.has_pending_output() { return Ok(None); } self.output_ready.synchronize()?; + let h2d = self.h2d_done.elapsed_since(&self.h2d_start)?; + let transform = self.transform_done.elapsed_since(&self.h2d_done)?; + let d2h = self.output_ready.elapsed_since(&self.transform_done)?; self.input_registration = None; self.input_vectors = None; self.input_matrix = None; @@ -529,7 +539,12 @@ impl TransformSlot { code_width, )?; self.rows = 0; - Ok(Some(batch)) + Ok(Some(DrainedTransformBatch { + batch, + h2d, + transform, + d2h, + })) } } @@ -553,6 +568,9 @@ struct ArtifactBuildStats { filter: Duration, matrix: Duration, launch: Duration, + gpu_h2d: Duration, + gpu_transform: Duration, + gpu_d2h: Duration, register: Duration, registered_bytes: usize, } @@ -584,6 +602,13 @@ impl ArtifactBuildStats { self.output_rows += batch.num_rows(); } + fn record_drained(&mut self, drained: &DrainedTransformBatch) { + self.record_output(&drained.batch); + self.gpu_h2d += drained.h2d; + self.gpu_transform += drained.transform; + self.gpu_d2h += drained.d2h; + } + fn log(&self) { eprintln!( "cuVS artifact stages: scanner_tasks={} prepare_workers={} input_batches={} input_rows={} prepared_batches={} prepared_rows={} output_batches={} output_rows={} scan_wait_s={:.3} raw_send_s={:.3} raw_wait_s={:.3} drain_s={:.3} send_s={:.3} prepare_send_s={:.3} vector_s={:.3} filter_s={:.3} matrix_s={:.3} launch_s={:.3}", @@ -611,6 +636,12 @@ impl ArtifactBuildStats { secs(self.register), self.registered_bytes as f64 / (1024.0 * 1024.0 * 1024.0), ); + eprintln!( + "cuVS artifact gpu events: h2d_s={:.3} transform_s={:.3} d2h_s={:.3}", + secs(self.gpu_h2d), + secs(self.gpu_transform), + secs(self.gpu_d2h), + ); } } @@ -953,7 +984,7 @@ async fn append_transformed_batches_to_artifact( let drain_start = Instant::now(); let transformed = if let Some(transformed) = slot.drain_to_batch(code_width)? { stats.drain += drain_start.elapsed(); - stats.record_output(&transformed); + stats.record_drained(&transformed); Some(transformed) } else { stats.drain += drain_start.elapsed(); @@ -966,9 +997,12 @@ async fn append_transformed_batches_to_artifact( if let Some(transformed) = transformed { let send_start = Instant::now(); - append_tx.send(Ok(transformed)).await.map_err(|error| { - Error::io(format!("failed to forward transformed batch: {error}")) - })?; + append_tx + .send(Ok(transformed.batch)) + .await + .map_err(|error| { + Error::io(format!("failed to forward transformed batch: {error}")) + })?; stats.send += send_start.elapsed(); } next_slot = (next_slot + 1) % PIPELINE_SLOTS; @@ -988,11 +1022,14 @@ async fn append_transformed_batches_to_artifact( let drain_start = Instant::now(); if let Some(transformed) = slot.drain_to_batch(code_width)? { stats.drain += drain_start.elapsed(); - stats.record_output(&transformed); + stats.record_drained(&transformed); let send_start = Instant::now(); - append_tx.send(Ok(transformed)).await.map_err(|error| { - Error::io(format!("failed to forward transformed batch: {error}")) - })?; + append_tx + .send(Ok(transformed.batch)) + .await + .map_err(|error| { + Error::io(format!("failed to forward transformed batch: {error}")) + })?; stats.send += send_start.elapsed(); } else { stats.drain += drain_start.elapsed(); diff --git a/backends/cuvs_26_02/src/cuda.rs b/backends/cuvs_26_02/src/cuda.rs index efa61a8..d96c9e2 100644 --- a/backends/cuvs_26_02/src/cuda.rs +++ b/backends/cuvs_26_02/src/cuda.rs @@ -29,6 +29,11 @@ unsafe extern "C" { stream: cuvs_sys::cudaStream_t, ) -> cuvs_sys::cudaError_t; fn cudaEventSynchronize(event: CudaEventHandle) -> cuvs_sys::cudaError_t; + fn cudaEventElapsedTime( + ms: *mut f32, + start: CudaEventHandle, + end: CudaEventHandle, + ) -> cuvs_sys::cudaError_t; } pub(crate) struct CuvsIvfPqIndex { @@ -438,6 +443,15 @@ impl CudaEvent { "synchronize CUDA event", ) } + + pub(crate) fn elapsed_since(&self, start: &Self) -> Result { + let mut ms = 0.0f32; + check_cuda( + unsafe { cudaEventElapsedTime(&mut ms, start.raw, self.raw) }, + "measure CUDA event elapsed time", + )?; + Ok(std::time::Duration::from_secs_f64(ms as f64 / 1000.0)) + } } impl Drop for CudaEvent { From 27607401584a1e7bf4e5a97ec401b9a1d1ef58b1 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 23 Jun 2026 14:22:57 +0800 Subject: [PATCH 5/8] perf: instrument cuVS artifact scheduling --- backends/cuvs_26_02/src/backend.rs | 57 ++++++++++++++++++++++++++++-- 1 file changed, 54 insertions(+), 3 deletions(-) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index d597ba0..ae90c72 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -36,6 +36,7 @@ const PARTITION_ARTIFACT_METADATA_FILE_NAME: &str = "metadata.lance"; const PIPELINE_SLOTS: usize = 2; const DEFAULT_SCAN_FRAGMENT_READAHEAD: usize = 0; const DEFAULT_SCAN_IO_BUFFER_SIZE: u64 = 16 * 1024 * 1024 * 1024; +const DEFAULT_SCAN_BATCH_READAHEAD: usize = 32; const DEFAULT_PREPARE_WORKERS: usize = 1; const TRAINING_SAMPLE_CHUNK_ROWS: usize = 8 * 1024; const TRAINING_SAMPLE_BATCH_READAHEAD: usize = 64; @@ -408,6 +409,13 @@ struct DrainedTransformBatch { d2h: Duration, } +#[derive(Default)] +struct LaunchTimings { + h2d_enqueue: Duration, + transform_call: Duration, + d2h_enqueue: Duration, +} + #[derive(Default)] struct ArtifactScannerStats { input_batches: usize, @@ -464,7 +472,8 @@ impl TransformSlot { trained: &TrainedIvfPqIndex, stream: cuvs_sys::cudaStream_t, prepared: PreparedTransformBatch, - ) -> Result<()> { + ) -> Result { + let mut timings = LaunchTimings::default(); let code_width = trained.pq_code_width(); let row_ids = prepared.row_ids; let matrix = prepared.matrix; @@ -480,8 +489,10 @@ impl TransformSlot { self.input_registration = prepared.input_registration; self.h2d_start.record(stream)?; + let h2d_enqueue_start = Instant::now(); self.input_device .copy_from_host_async(&trained.resources, input_slice)?; + timings.h2d_enqueue += h2d_enqueue_start.elapsed(); self.h2d_done.record(stream)?; match matrix { PreparedMatrix::F32Arrow { vectors, .. } => { @@ -493,6 +504,7 @@ impl TransformSlot { self.input_matrix = Some(array); } } + let transform_call_start = Instant::now(); check_cuvs( unsafe { cuvs_sys::cuvsIvfPqTransform( @@ -505,15 +517,18 @@ impl TransformSlot { }, "transform vectors with IVF_PQ", )?; + timings.transform_call += transform_call_start.elapsed(); self.transform_done.record(stream)?; + let d2h_enqueue_start = Instant::now(); self.labels_device .copy_to_host_async(&trained.resources, self.labels_host.prefix_mut(rows)?)?; self.codes_device.copy_to_host_async( &trained.resources, self.codes_host.prefix_mut(rows * code_width)?, )?; + timings.d2h_enqueue += d2h_enqueue_start.elapsed(); self.output_ready.record(stream)?; - Ok(()) + Ok(timings) } fn drain_to_batch(&mut self, code_width: usize) -> Result> { @@ -571,6 +586,9 @@ struct ArtifactBuildStats { gpu_h2d: Duration, gpu_transform: Duration, gpu_d2h: Duration, + launch_h2d_enqueue: Duration, + launch_transform_call: Duration, + launch_d2h_enqueue: Duration, register: Duration, registered_bytes: usize, } @@ -609,6 +627,12 @@ impl ArtifactBuildStats { self.gpu_d2h += drained.d2h; } + fn record_launch_timings(&mut self, timings: LaunchTimings) { + self.launch_h2d_enqueue += timings.h2d_enqueue; + self.launch_transform_call += timings.transform_call; + self.launch_d2h_enqueue += timings.d2h_enqueue; + } + fn log(&self) { eprintln!( "cuVS artifact stages: scanner_tasks={} prepare_workers={} input_batches={} input_rows={} prepared_batches={} prepared_rows={} output_batches={} output_rows={} scan_wait_s={:.3} raw_send_s={:.3} raw_wait_s={:.3} drain_s={:.3} send_s={:.3} prepare_send_s={:.3} vector_s={:.3} filter_s={:.3} matrix_s={:.3} launch_s={:.3}", @@ -642,6 +666,13 @@ impl ArtifactBuildStats { secs(self.gpu_transform), secs(self.gpu_d2h), ); + eprintln!( + "cuVS artifact launch cpu: h2d_enqueue_s={:.3} transform_call_s={:.3} d2h_enqueue_s={:.3}", + secs(self.launch_h2d_enqueue), + secs(self.launch_transform_call), + secs(self.launch_d2h_enqueue), + ); + eprintln!("cuVS artifact max rss: {} KiB", max_rss_kib()); } } @@ -649,6 +680,16 @@ fn secs(duration: Duration) -> f64 { duration.as_secs_f64() } +fn max_rss_kib() -> i64 { + let mut usage = std::mem::MaybeUninit::::uninit(); + let status = unsafe { libc::getrusage(libc::RUSAGE_SELF, usage.as_mut_ptr()) }; + if status == 0 { + unsafe { usage.assume_init().ru_maxrss } + } else { + -1 + } +} + fn prepare_workers_from_env() -> usize { std::env::var("LANCE_CUVS_PREPARE_WORKERS") .ok() @@ -664,6 +705,14 @@ fn scan_fragment_readahead_from_env() -> usize { .unwrap_or(DEFAULT_SCAN_FRAGMENT_READAHEAD) } +fn scan_batch_readahead_from_env() -> usize { + std::env::var("LANCE_CUVS_SCAN_BATCH_READAHEAD") + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(DEFAULT_SCAN_BATCH_READAHEAD) +} + fn training_sample_ranges(num_rows: usize, sample_rows: usize) -> Vec> { let sample_rows = sample_rows.min(num_rows); if sample_rows == 0 { @@ -740,6 +789,7 @@ async fn scan_transform_batches( scanner.batch_size(batch_size); scanner.scan_in_order(false); scanner.fragment_readahead(scan_fragment_readahead_from_env()); + scanner.batch_readahead(scan_batch_readahead_from_env()); scanner.io_buffer_size(DEFAULT_SCAN_IO_BUFFER_SIZE); let mut stream = scanner.try_into_stream().await?; let mut stats = ArtifactScannerStats::default(); @@ -992,8 +1042,9 @@ async fn append_transformed_batches_to_artifact( }; let launch_start = Instant::now(); - slot.launch(trained, cuda_stream, prepared)?; + let launch_timings = slot.launch(trained, cuda_stream, prepared)?; stats.launch += launch_start.elapsed(); + stats.record_launch_timings(launch_timings); if let Some(transformed) = transformed { let send_start = Instant::now(); From 217a51699a19bf5f9a9f881797d425f7b488428c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 23 Jun 2026 15:44:48 +0800 Subject: [PATCH 6/8] perf: break down cuVS artifact drain time --- backends/cuvs_26_02/src/backend.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index ae90c72..cca0e8e 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -407,6 +407,8 @@ struct DrainedTransformBatch { h2d: Duration, transform: Duration, d2h: Duration, + sync: Duration, + build_batch: Duration, } #[derive(Default)] @@ -536,7 +538,9 @@ impl TransformSlot { return Ok(None); } + let sync_start = Instant::now(); self.output_ready.synchronize()?; + let sync = sync_start.elapsed(); let h2d = self.h2d_done.elapsed_since(&self.h2d_start)?; let transform = self.transform_done.elapsed_since(&self.h2d_done)?; let d2h = self.output_ready.elapsed_since(&self.transform_done)?; @@ -547,18 +551,22 @@ impl TransformSlot { .row_ids .take() .ok_or_else(|| Error::io("transform slot is missing row ids"))?; + let build_batch_start = Instant::now(); let batch = build_partition_batch( row_ids, self.labels_host.prefix(self.rows)?, self.codes_host.prefix(self.rows * code_width)?, code_width, )?; + let build_batch = build_batch_start.elapsed(); self.rows = 0; Ok(Some(DrainedTransformBatch { batch, h2d, transform, d2h, + sync, + build_batch, })) } } @@ -589,6 +597,8 @@ struct ArtifactBuildStats { launch_h2d_enqueue: Duration, launch_transform_call: Duration, launch_d2h_enqueue: Duration, + drain_sync: Duration, + drain_build_batch: Duration, register: Duration, registered_bytes: usize, } @@ -625,6 +635,8 @@ impl ArtifactBuildStats { self.gpu_h2d += drained.h2d; self.gpu_transform += drained.transform; self.gpu_d2h += drained.d2h; + self.drain_sync += drained.sync; + self.drain_build_batch += drained.build_batch; } fn record_launch_timings(&mut self, timings: LaunchTimings) { @@ -672,6 +684,11 @@ impl ArtifactBuildStats { secs(self.launch_transform_call), secs(self.launch_d2h_enqueue), ); + eprintln!( + "cuVS artifact drain cpu: sync_s={:.3} build_batch_s={:.3}", + secs(self.drain_sync), + secs(self.drain_build_batch), + ); eprintln!("cuVS artifact max rss: {} KiB", max_rss_kib()); } } From 4968beb8e87d45107c824703ade8d2a07e6c55be Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 23 Jun 2026 18:21:52 +0800 Subject: [PATCH 7/8] fix: copy strided cuVS tensors correctly --- backends/cuvs_26_02/src/cuda.rs | 66 ++++++++++++++++++++++++++------- 1 file changed, 52 insertions(+), 14 deletions(-) diff --git a/backends/cuvs_26_02/src/cuda.rs b/backends/cuvs_26_02/src/cuda.rs index d96c9e2..2b4d160 100644 --- a/backends/cuvs_26_02/src/cuda.rs +++ b/backends/cuvs_26_02/src/cuda.rs @@ -22,6 +22,16 @@ unsafe extern "C" { fn cudaFreeHost(ptr: *mut c_void) -> cuvs_sys::cudaError_t; fn cudaHostRegister(ptr: *mut c_void, size: usize, flags: u32) -> cuvs_sys::cudaError_t; fn cudaHostUnregister(ptr: *mut c_void) -> cuvs_sys::cudaError_t; + fn cudaMemcpy2DAsync( + dst: *mut c_void, + dpitch: usize, + src: *const c_void, + spitch: usize, + width: usize, + height: usize, + kind: cuvs_sys::cudaMemcpyKind, + stream: cuvs_sys::cudaStream_t, + ) -> cuvs_sys::cudaError_t; fn cudaEventCreate(event: *mut CudaEventHandle) -> cuvs_sys::cudaError_t; fn cudaEventDestroy(event: CudaEventHandle) -> cuvs_sys::cudaError_t; fn cudaEventRecord( @@ -584,20 +594,48 @@ pub(crate) fn copy_tensor_to_host_f32_2d( ))); } let mut array = Array2::::zeros((shape[0], shape[1])); - check_cuda( - unsafe { - cuvs_sys::cudaMemcpyAsync( - array.as_mut_ptr() as *mut _, - tensor.dl_tensor.data, - tensor_num_bytes(tensor), - cuvs_sys::cudaMemcpyKind_cudaMemcpyDefault, - resources - .get_cuda_stream() - .map_err(|e| Error::io(e.to_string()))?, - ) - }, - "copy tensor to host", - )?; + let stream = resources + .get_cuda_stream() + .map_err(|e| Error::io(e.to_string()))?; + if tensor.dl_tensor.strides.is_null() { + check_cuda( + unsafe { + cuvs_sys::cudaMemcpyAsync( + array.as_mut_ptr() as *mut _, + tensor.dl_tensor.data, + tensor_num_bytes(tensor), + cuvs_sys::cudaMemcpyKind_cudaMemcpyDefault, + stream, + ) + }, + "copy tensor to host", + )?; + } else { + let row_stride = unsafe { *tensor.dl_tensor.strides } as usize; + let col_stride = unsafe { *tensor.dl_tensor.strides.add(1) } as usize; + if col_stride != 1 { + return Err(Error::not_supported(format!( + "copying 2D tensors with non-unit column stride is not supported: strides=({}, {})", + row_stride, col_stride + ))); + } + let row_bytes = shape[1] * std::mem::size_of::(); + check_cuda( + unsafe { + cudaMemcpy2DAsync( + array.as_mut_ptr().cast::(), + row_bytes, + tensor.dl_tensor.data, + row_stride * std::mem::size_of::(), + row_bytes, + shape[0], + cuvs_sys::cudaMemcpyKind_cudaMemcpyDefault, + stream, + ) + }, + "copy strided tensor to host", + )?; + } resources .sync_stream() .map_err(|e| Error::io(e.to_string()))?; From 00865d9749fc7acb5d41a9c8edc137744cdbaa89 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 23 Jun 2026 20:26:16 +0800 Subject: [PATCH 8/8] perf: allow enabling cuVS RMM pool --- backends/cuvs_26_02/src/backend.rs | 5 +++-- backends/cuvs_26_02/src/cuda.rs | 32 ++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/backends/cuvs_26_02/src/backend.rs b/backends/cuvs_26_02/src/backend.rs index cca0e8e..dff1a5b 100644 --- a/backends/cuvs_26_02/src/backend.rs +++ b/backends/cuvs_26_02/src/backend.rs @@ -4,8 +4,8 @@ use crate::cuda::{ CudaEvent, CuvsIvfPqIndex, DeviceTensor, HostTensorView, MatrixBuffer, PinnedHostBuffer, RegisteredHostBuffer, check_cuvs, copy_tensor_to_host_f32_2d, copy_tensor_to_host_f32_3d, - create_index_params, destroy_index_params, ivf_centroids_from_host, make_tensor_view, - matrix_from_vectors, pq_codebook_from_host, + create_index_params, destroy_index_params, enable_rmm_pool_from_env, ivf_centroids_from_host, + make_tensor_view, matrix_from_vectors, pq_codebook_from_host, }; use arrow::compute::{concat_batches, filter}; use arrow_array::cast::AsArray; @@ -1220,6 +1220,7 @@ pub async fn train_ivf_pq( } let matrix = matrix_from_vectors(&train_vectors)?; + enable_rmm_pool_from_env()?; let resources = Resources::new().map_err(|error| Error::io(error.to_string()))?; let index = CuvsIvfPqIndex::try_new()?; let params = create_index_params( diff --git a/backends/cuvs_26_02/src/cuda.rs b/backends/cuvs_26_02/src/cuda.rs index 2b4d160..7169f88 100644 --- a/backends/cuvs_26_02/src/cuda.rs +++ b/backends/cuvs_26_02/src/cuda.rs @@ -499,6 +499,38 @@ pub(crate) fn check_cuda(status: cuvs_sys::cudaError_t, context: &str) -> Result } } +pub(crate) fn enable_rmm_pool_from_env() -> Result<()> { + let Some(config) = std::env::var("LANCE_CUVS_RMM_POOL").ok() else { + return Ok(()); + }; + let (initial, max) = match config.split_once(',') { + Some((initial, max)) => ( + initial.parse::().map_err(|error| { + Error::invalid_input(format!( + "invalid LANCE_CUVS_RMM_POOL initial percent '{initial}': {error}" + )) + })?, + max.parse::().map_err(|error| { + Error::invalid_input(format!( + "invalid LANCE_CUVS_RMM_POOL max percent '{max}': {error}" + )) + })?, + ), + None => { + let percent = config.parse::().map_err(|error| { + Error::invalid_input(format!( + "invalid LANCE_CUVS_RMM_POOL percent '{config}': {error}" + )) + })?; + (percent, percent) + } + }; + check_cuvs( + unsafe { cuvs_sys::cuvsRMMPoolMemoryResourceEnable(initial, max, false) }, + "enable RMM pool memory resource", + ) +} + pub(crate) fn cuvs_distance_type(metric_type: DistanceType) -> Result { match metric_type { DistanceType::L2 => Ok(cuvs_sys::cuvsDistanceType::L2Expanded),