Skip to main content

cuvs/
prepared.rs

1use std::{marker::PhantomData, mem};
2
3use cudf::{
4    TypeId,
5    core::Table,
6    frame::dataframe::{ValidatedFixedSizeListExport, ValidatedUniformListExport},
7    interop::contract::{
8        ActiveColumnReadLease, BufferBinding, ColumnLayout, CompletionAuthority, LeaseCompletion, PhysicalValidity,
9        Placement, PreparedColumnExport, PreparedExportProfile, PreparedTableExport, ReadStartAuthority,
10    },
11};
12
13use crate::{
14    device::{DeviceElement, DeviceTensor, MatrixLayout, OwnedMatrix, ValidatedDeviceTensor},
15    error::{Result, from_cudf, from_interop, invalid_input},
16    resources::CuvsResources,
17};
18
19/// An activated prepared-export read plus its lifetime-bound tensor view.
20#[must_use = "prepared tensor reads must be completed explicitly after cuVS work has been enqueued"]
21pub struct PreparedTensorLease<'export, 'producer, 'consumer, T: DeviceElement> {
22    tensor: DeviceTensor<'export, T>,
23    lease: Option<ActiveColumnReadLease<'export, 'producer, 'consumer>>,
24    _consumer: PhantomData<&'consumer CuvsResources>,
25}
26
27impl<'export, T: DeviceElement> PreparedTensorLease<'export, '_, '_, T> {
28    pub fn tensor(&self) -> &DeviceTensor<'export, T> {
29        &self.tensor
30    }
31
32    /// Complete producer lifetime protection after all cuVS reads have been
33    /// enqueued onto the runtime stream.
34    pub fn complete(mut self) -> Result<LeaseCompletion> {
35        let operation = "PreparedTensorLease::complete";
36        let lease =
37            self.lease.take().ok_or_else(|| invalid_input(operation, "prepared tensor lease was already completed"))?;
38        lease.complete().map_err(|error| from_interop(operation, error))
39    }
40}
41
42impl<T: DeviceElement> std::fmt::Debug for PreparedTensorLease<'_, '_, '_, T> {
43    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        formatter
45            .debug_struct("PreparedTensorLease")
46            .field("tensor", &self.tensor)
47            .field("completed", &self.lease.is_none())
48            .finish_non_exhaustive()
49    }
50}
51
52/// A prepared list read with floating finite validation or integer structural validation.
53#[must_use = "prepared tensor reads must be completed explicitly after cuVS work has been enqueued"]
54pub struct PreparedValidatedTensorLease<'export, 'producer, 'consumer, T: DeviceElement> {
55    tensor: ValidatedDeviceTensor<'export, T>,
56    lease: Option<ActiveColumnReadLease<'export, 'producer, 'consumer>>,
57    _consumer: PhantomData<&'consumer CuvsResources>,
58}
59
60impl<'export, T: DeviceElement> PreparedValidatedTensorLease<'export, '_, '_, T> {
61    pub fn tensor(&self) -> &ValidatedDeviceTensor<'export, T> {
62        &self.tensor
63    }
64
65    pub fn complete(mut self) -> Result<LeaseCompletion> {
66        let operation = "PreparedValidatedTensorLease::complete";
67        let lease =
68            self.lease.take().ok_or_else(|| invalid_input(operation, "prepared tensor lease was already completed"))?;
69        lease.complete().map_err(|error| from_interop(operation, error))
70    }
71}
72
73impl<T: DeviceElement> std::fmt::Debug for PreparedValidatedTensorLease<'_, '_, '_, T> {
74    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        formatter
76            .debug_struct("PreparedValidatedTensorLease")
77            .field("tensor", &self.tensor)
78            .field("completed", &self.lease.is_none())
79            .finish_non_exhaustive()
80    }
81}
82
83impl CuvsResources {
84    /// Activates a fixed-width prepared column as a tensor on this runtime.
85    ///
86    /// # Safety
87    ///
88    /// For an export without a retained storage owner, the caller must finish
89    /// all queued tensor reads before the producer borrow ends, including when
90    /// the lease is forgotten. Call `complete` after the last submission and
91    /// keep storage alive through any subsequent error cleanup. A destructor
92    /// alone cannot guarantee this for borrowed asynchronous input. Exports
93    /// with retained owners keep those owners in the active lease.
94    ///
95    /// ```compile_fail,E0133
96    /// use cudf::interop::contract::PreparedColumnExport;
97    /// use cuvs::CuvsResources;
98    /// fn activate(resources: &CuvsResources, export: &PreparedColumnExport<'_>) {
99    ///     let _ = resources.activate_prepared_tensor::<f32>(export, &[3], None);
100    /// }
101    /// ```
102    pub unsafe fn activate_prepared_tensor<'export, 'producer, 'consumer, T: DeviceElement>(
103        &'consumer self,
104        export: &'export PreparedColumnExport<'producer>,
105        shape: &[usize],
106        strides: Option<&[usize]>,
107    ) -> Result<PreparedTensorLease<'export, 'producer, 'consumer, T>> {
108        let operation = "CuvsResources::activate_prepared_tensor";
109        self.validate_exact_binding(operation)?;
110        require_device_export(export, operation)?;
111        let column = export.column();
112        if column.data_type().id() != T::INTEROP_TYPE {
113            return Err(invalid_input(
114                operation,
115                format!(
116                    "prepared column dtype {:?} does not match tensor element {}",
117                    column.data_type().id(),
118                    std::any::type_name::<T>()
119                ),
120            ));
121        }
122        if !matches!(column.validity(), PhysicalValidity::NoValidityBuffer) {
123            return Err(invalid_input(operation, "cuVS tensor inputs must not carry a validity mask"));
124        }
125        let values = match column.layout() {
126            ColumnLayout::FixedWidth { values } => *values,
127            _ => return Err(invalid_input(operation, "cuVS prepared tensors require a fixed-width column layout")),
128        };
129        let binding = export
130            .bindings()
131            .iter()
132            .find(|binding| binding.id() == values)
133            .ok_or_else(|| invalid_input(operation, "prepared tensor values binding is missing"))?;
134        let binding_device = require_binding_device(binding, self.device().ordinal() as i32, operation)?;
135        let logical_elements = shape.iter().try_fold(1usize, |count, dimension| {
136            count
137                .checked_mul(*dimension)
138                .ok_or_else(|| invalid_input(operation, "prepared tensor shape overflowed usize"))
139        })?;
140        if logical_elements != column.len() {
141            return Err(invalid_input(
142                operation,
143                format!("prepared tensor shape contains {logical_elements} elements, column contains {}", column.len()),
144            ));
145        }
146        let byte_offset = column
147            .offset()
148            .checked_mul(mem::size_of::<T>())
149            .ok_or_else(|| invalid_input(operation, "prepared tensor byte offset overflowed usize"))?;
150
151        let runtime_binding = self.runtime().execution_binding().map_err(|error| from_interop(operation, error))?;
152        let stream = runtime_binding.cuda_stream();
153        let lease = export
154            .try_activate_read_lease(
155                ReadStartAuthority::ConsumerStream { stream },
156                CompletionAuthority::SameDeviceStream { stream },
157            )
158            .map_err(|error| from_interop(operation, error))?;
159        let tensor = DeviceTensor::from_parts(
160            binding.bytes().base_ptr() as *mut std::ffi::c_void,
161            binding_device,
162            shape,
163            strides,
164            binding.bytes().byte_len(),
165            byte_offset,
166            self.binding_id(),
167            operation,
168        )?;
169        Ok(PreparedTensorLease { tensor, lease: Some(lease), _consumer: PhantomData })
170    }
171
172    /// Activate a structurally validated numeric `FixedSizeList`
173    /// child as a row-major tensor under the retained root export's single
174    /// read lease.
175    ///
176    /// Tensor dimensions are verified against the cuDF witness proving uniform
177    /// list-offset deltas across active rows.
178    pub fn activate_validated_fixed_size_list_tensor<'witness, 'producer, 'consumer, T: DeviceElement>(
179        &'consumer self,
180        validated: &'witness ValidatedFixedSizeListExport<'producer>,
181    ) -> Result<PreparedValidatedTensorLease<'witness, 'witness, 'consumer, T>> {
182        let operation = "CuvsResources::activate_validated_fixed_size_list_tensor";
183        self.validate_exact_binding(operation)?;
184        require_element::<T>(validated.element(), operation)?;
185        let export = validated.export();
186        require_device_export(export, operation)?;
187        if validated.binding_id() != self.binding_id() {
188            return Err(invalid_input(
189                operation,
190                "prepared tensor and cuVS resources must use the same exact runtime binding",
191            )
192            .with_context("tensor_binding_id", validated.binding_id())
193            .with_context("resource_binding_id", self.binding_id()));
194        }
195        let binding = validated.child_values_binding();
196        let binding_device = require_binding_device(binding, self.device().ordinal() as i32, operation)?;
197
198        let runtime_binding = self.runtime().execution_binding().map_err(|error| from_interop(operation, error))?;
199        let stream = runtime_binding.cuda_stream();
200        let lease = export
201            .try_activate_read_lease(
202                ReadStartAuthority::ConsumerStream { stream },
203                CompletionAuthority::SameDeviceStream { stream },
204            )
205            .map_err(|error| from_interop(operation, error))?;
206        let tensor = DeviceTensor::from_parts(
207            binding.bytes().base_ptr() as *mut std::ffi::c_void,
208            binding_device,
209            &[validated.rows(), validated.dimension()],
210            None,
211            binding.bytes().byte_len(),
212            validated.child_byte_offset(),
213            validated.binding_id(),
214            operation,
215        )
216        .map(ValidatedDeviceTensor::new)?;
217        Ok(PreparedValidatedTensorLease { tensor, lease: Some(lease), _consumer: PhantomData })
218    }
219
220    /// Activate a structurally validated runtime-uniform
221    /// numeric `List` child as a row-major tensor under the retained root
222    /// export's single read lease.
223    ///
224    /// Tensor dimensions are derived directly from the cuDF witness of active
225    /// list-offset widths.
226    pub fn activate_validated_uniform_list_tensor<'witness, 'producer, 'consumer, T: DeviceElement>(
227        &'consumer self,
228        validated: &'witness ValidatedUniformListExport<'producer>,
229    ) -> Result<PreparedValidatedTensorLease<'witness, 'witness, 'consumer, T>> {
230        let operation = "CuvsResources::activate_validated_uniform_list_tensor";
231        self.validate_exact_binding(operation)?;
232        require_element::<T>(validated.element(), operation)?;
233        let export = validated.export();
234        require_device_export(export, operation)?;
235        if validated.binding_id() != self.binding_id() {
236            return Err(invalid_input(
237                operation,
238                "prepared tensor and cuVS resources must use the same exact runtime binding",
239            )
240            .with_context("tensor_binding_id", validated.binding_id())
241            .with_context("resource_binding_id", self.binding_id()));
242        }
243        let binding = validated.child_values_binding();
244        let binding_device = require_binding_device(binding, self.device().ordinal() as i32, operation)?;
245
246        let runtime_binding = self.runtime().execution_binding().map_err(|error| from_interop(operation, error))?;
247        let stream = runtime_binding.cuda_stream();
248        let lease = export
249            .try_activate_read_lease(
250                ReadStartAuthority::ConsumerStream { stream },
251                CompletionAuthority::SameDeviceStream { stream },
252            )
253            .map_err(|error| from_interop(operation, error))?;
254        let tensor = DeviceTensor::from_parts(
255            binding.bytes().base_ptr() as *mut std::ffi::c_void,
256            binding_device,
257            &[validated.rows(), validated.dimension()],
258            None,
259            binding.bytes().byte_len(),
260            validated.child_byte_offset(),
261            validated.binding_id(),
262            operation,
263        )
264        .map(ValidatedDeviceTensor::new)?;
265        Ok(PreparedValidatedTensorLease { tensor, lease: Some(lease), _consumer: PhantomData })
266    }
267
268    /// Copy a validated numeric `FixedSizeList` child into a private
269    /// matrix in the requested layout on this exact runtime binding.
270    pub fn pack_validated_fixed_size_list<T: DeviceElement>(
271        &self,
272        validated: &ValidatedFixedSizeListExport<'_>,
273        layout: MatrixLayout,
274    ) -> Result<OwnedMatrix<T>> {
275        let operation = "CuvsResources::pack_validated_fixed_size_list";
276        self.validate_exact_binding(operation)?;
277        require_element::<T>(validated.element(), operation)?;
278        if validated.binding_id() != self.binding_id() {
279            return Err(invalid_input(
280                operation,
281                "validated list input and cuVS resources must use the same exact runtime binding",
282            )
283            .with_context("tensor_binding_id", validated.binding_id())
284            .with_context("resource_binding_id", self.binding_id()));
285        }
286        let packed = validated
287            .pack_active_child_into_runtime(self.runtime(), layout.into())
288            .map_err(|error| from_cudf(operation, error))?;
289        OwnedMatrix::from_packed(self, packed)
290    }
291
292    /// Copy a validated runtime-uniform numeric `List` child into a
293    /// private matrix in the requested layout on this exact runtime binding.
294    pub fn pack_validated_uniform_list<T: DeviceElement>(
295        &self,
296        validated: &ValidatedUniformListExport<'_>,
297        layout: MatrixLayout,
298    ) -> Result<OwnedMatrix<T>> {
299        let operation = "CuvsResources::pack_validated_uniform_list";
300        self.validate_exact_binding(operation)?;
301        require_element::<T>(validated.element(), operation)?;
302        if validated.binding_id() != self.binding_id() {
303            return Err(invalid_input(
304                operation,
305                "validated list input and cuVS resources must use the same exact runtime binding",
306            )
307            .with_context("tensor_binding_id", validated.binding_id())
308            .with_context("resource_binding_id", self.binding_id()));
309        }
310        let packed = validated
311            .pack_active_child_into_runtime(self.runtime(), layout.into())
312            .map_err(|error| from_cudf(operation, error))?;
313        OwnedMatrix::from_packed(self, packed)
314    }
315
316    /// Atomically import one prepared table root and pack all top-level columns
317    /// into an owned numeric matrix on this exact runtime binding.
318    pub fn pack_prepared_table<T: DeviceElement>(
319        &self,
320        export: &PreparedTableExport<'_>,
321        layout: MatrixLayout,
322    ) -> Result<OwnedMatrix<T>> {
323        let operation = "CuvsResources::pack_prepared_table";
324        self.validate_exact_binding(operation)?;
325        let table = Table::try_import_from_prepared_table_export_into_runtime(export, self.runtime())
326            .map_err(|error| from_interop(operation, error))?;
327        let columns = table.into_columns().map_err(|error| from_cudf(operation, error))?;
328        let column_refs = columns.iter().collect::<Vec<_>>();
329        let packed =
330            self.runtime().pack_columns(&column_refs, layout.into()).map_err(|error| from_cudf(operation, error))?;
331        OwnedMatrix::from_packed(self, packed)
332    }
333}
334
335fn require_element<T: DeviceElement>(element: TypeId, operation: &'static str) -> Result<()> {
336    if element != T::TYPE_ID {
337        return Err(invalid_input(operation, "validated element does not match requested cuVS element")
338            .with_context("validated_element", format!("{element:?}"))
339            .with_context("requested_element", format!("{:?}", T::TYPE_ID)));
340    }
341    Ok(())
342}
343
344fn require_device_export(export: &PreparedColumnExport<'_>, operation: &'static str) -> Result<()> {
345    if export.profile() != PreparedExportProfile::Device {
346        return Err(invalid_input(operation, "cuVS requires a prepared device export"));
347    }
348    Ok(())
349}
350
351fn require_binding_device(binding: &BufferBinding<'_>, expected_device: i32, operation: &'static str) -> Result<i32> {
352    let device = match binding.placement() {
353        Placement::CudaDevice { device } => *device,
354        Placement::Host => {
355            return Err(invalid_input(operation, "cuVS prepared tensor binding must be CUDA device memory"));
356        },
357    };
358    if device != expected_device {
359        return Err(invalid_input(operation, "prepared tensor and cuVS resources must use the same CUDA device"));
360    }
361    Ok(device)
362}