Skip to main content

xlog_cuda/provider/
filter.rs

1//! Filter and compare operations on GPU buffers.
2//!
3//! Generic versions of `filter`, `compare_columns`, plus the shared helpers
4//! `compare_const_mask` and `compare_columns_mask` (moved from mod.rs).
5
6use crate::{DeviceRepr, DeviceSlice, KernelScalar, LaunchAsync, LaunchConfig};
7use xlog_core::{Result, ScalarType, XlogError};
8
9use super::{filter_kernels, scan_kernels, FILTER_MODULE, SCAN_MODULE};
10use crate::device_runtime::StreamId;
11use crate::launch::LaunchRecorder;
12use crate::memory::{CudaColumn, DeviceMemoryView, DeviceRead, TrackedCudaSlice};
13use crate::type_seam::GpuScalar;
14use crate::{CompareOp, CudaBuffer};
15
16impl super::CudaKernelProvider {
17    // ------------------------------------------------------------------
18    // Generic public API
19    // ------------------------------------------------------------------
20
21    /// Generic compare-columns: produce a device mask for `left <op> right`.
22    ///
23    /// Replaces: `compare_columns_u32`, `compare_columns_i32`, `compare_columns_i64`,
24    /// `compare_columns_u64`, `compare_columns_f32`, `compare_columns_f64`,
25    /// `compare_columns_u8`.
26    pub fn compare_columns<T: GpuScalar>(
27        &self,
28        input: &CudaBuffer,
29        left: usize,
30        right: usize,
31        op: CompareOp,
32    ) -> Result<TrackedCudaSlice<u8>> {
33        self.compare_columns_mask::<T>(
34            input,
35            left,
36            right,
37            op,
38            T::allowed_scalar_types(),
39            T::compare_col_kernel(),
40        )
41    }
42
43    /// Generic filter: keep rows where `column[col] <op> value`.
44    ///
45    /// Dispatches between fused compare+scan+compact (u32, f64) and
46    /// mask+compact (all other types).
47    ///
48    /// Replaces: `filter_u32`, `filter_f64`, `filter_i32`, `filter_u64`,
49    /// `filter_f32`, `filter_bool`.
50    pub fn filter<T: GpuScalar>(
51        &self,
52        input: &CudaBuffer,
53        col: usize,
54        value: T,
55        op: CompareOp,
56    ) -> Result<CudaBuffer> {
57        // Opt-in dispatch: route through the recorded filter path
58        // when the env var is set AND the manager is runtime-backed
59        // AND a launch stream can be acquired from the pool.
60        // Default behavior is unchanged (legacy fused / mask+compact).
61        if Self::use_recorded_filters_env()? {
62            if let Some(launch_stream) = self.recorded_op_stream_or_init() {
63                return self.filter_recorded::<T>(input, col, value, op, launch_stream);
64            }
65        }
66
67        if input.is_empty() {
68            return self.create_empty_buffer(input.schema.clone());
69        }
70
71        if T::filter_scan_phase1_kernel().is_some() {
72            // Fused compare+scan+compact path (u32, f64).
73            self.filter_fused_scan::<T>(input, col, value, op)
74        } else {
75            // Mask + compact path (i32, u64, f32, i64, u8, bool).
76            let mask = self.compare_const_mask::<T>(
77                input,
78                col,
79                value,
80                op,
81                T::allowed_scalar_types(),
82                T::filter_compare_kernel(),
83            )?;
84            self.filter_by_device_mask(input, &mask)
85        }
86    }
87
88    /// Strict-recorder, end-to-end variant of [`Self::filter`] —
89    /// the first composed migrated DATA path.
90    ///
91    /// Composes [`Self::compare_const_mask_recorded`] and
92    /// [`Self::compact_buffer_by_device_mask_counted_recorded`]
93    /// on a single `launch_stream`. Each primitive builds its
94    /// own [`crate::launch::LaunchRecorder`], records its uses,
95    /// preflights, runs its kernels, and commits independently.
96    ///
97    /// Composition correctness rests on the runtime's
98    /// "record-all, wait-all" semantics: every
99    /// `record_block_use` call APPENDS a fresh event to the
100    /// live entry's `last_use_events: Vec<CudaEvent>`, and
101    /// `deallocate` waits on EVERY event in that vector before
102    /// queueing `cuMemFreeAsync`. So the compare's commit and
103    /// the compact's later commit each push their own event
104    /// for `input.column[i]` (and other shared buffers), and
105    /// the deallocate gates the free behind both — closing
106    /// the cross-stream lifetime gap end-to-end. (Latest-event
107    /// coalescing per `(block, launch_stream)` is a possible
108    /// future optimization; today every recorded use is
109    /// retained and waited on.)
110    ///
111    /// # Dispatch
112    /// For types with a fused `filter_compare_*_scan_phase1`
113    /// kernel (`u32`, `f64`), routes to
114    /// [`Self::filter_fused_scan_recorded`] — single-pass
115    /// compare+scan+compact mirror of the legacy fast path.
116    /// For all other types, composes
117    /// [`Self::compare_const_mask_recorded`] +
118    /// [`Self::compact_buffer_by_device_mask_counted_recorded`].
119    ///
120    /// # Errors
121    /// Propagates the structured `XlogError::Kernel` errors
122    /// produced by either underlying recorded primitive
123    /// (legacy manager, unresolved launch_stream, external
124    /// column, preflight / commit failures, kernel launch
125    /// failures, `cu_stream.synchronize()` before host scalar
126    /// read).
127    pub fn filter_recorded<T: GpuScalar>(
128        &self,
129        input: &CudaBuffer,
130        col: usize,
131        value: T,
132        op: CompareOp,
133        launch_stream: StreamId,
134    ) -> Result<CudaBuffer> {
135        if input.is_empty() {
136            return self.create_empty_buffer(input.schema.clone());
137        }
138        if T::filter_scan_phase1_kernel().is_some() {
139            return self.filter_fused_scan_recorded::<T>(input, col, value, op, launch_stream);
140        }
141        let d_mask = self.compare_const_mask_recorded::<T>(input, col, value, op, launch_stream)?;
142        self.compact_buffer_by_device_mask_counted_recorded(input, &d_mask, launch_stream)
143    }
144
145    /// Strict-recorder variant of [`Self::filter_fused_scan`] —
146    /// the migrated fused compare+scan+compact fast path for
147    /// `u32` and `f64`.
148    ///
149    /// Mirrors the legacy chain on a single explicit
150    /// `launch_stream`:
151    ///   1. `filter_compare_T_scan_phase1` — fused compare +
152    ///      block-local scan that produces `d_mask`,
153    ///      `d_prefix_sum`, `d_block_sums` in one launch.
154    ///   2. When `num_blocks > 1`,
155    ///      `multiblock_scan_u32_inplace_on_stream` on
156    ///      `d_block_sums` followed by
157    ///      `multiblock_scan_phase3` to propagate block offsets
158    ///      into `d_prefix_sum`.
159    ///   3. `capture_compact_count` — writes `d_out_count` for
160    ///      the masked total.
161    ///   4. Per-input-column `compact_bytes_by_mask` on the
162    ///      same `launch_stream`, including an all-zero mask.
163    ///   5. Commit the complete chain, then read `d_out_count`
164    ///      through its published dependency to obtain `output_rows`.
165    ///
166    /// This method returns a host row count and excludes stream capture before
167    /// allocating or enqueueing work. The exclusion lasts through readback.
168    ///
169    /// # Strict-mode contract
170    /// Identical to
171    /// [`Self::compact_buffer_by_device_mask_counted_recorded`]:
172    /// `input.num_rows_device()` and every `input.column(i)`
173    /// recorded as reads BEFORE preflight; every fresh
174    /// runtime-backed allocation (`d_mask`, `d_prefix_sum`,
175    /// `d_block_sums`, `d_out_count`, each `dst_col`) recorded
176    /// via `write` BEFORE preflight; the recorder snapshots
177    /// block identity at record time and drops the source
178    /// borrow, so kernel `&mut` borrows after preflight remain
179    /// valid before the kernels enqueue.
180    ///
181    /// # Panics
182    /// `T::filter_scan_phase1_kernel()` must be `Some` —
183    /// callers should only reach this method for `u32` / `f64`.
184    pub fn filter_fused_scan_recorded<T: GpuScalar>(
185        &self,
186        input: &CudaBuffer,
187        col: usize,
188        value: T,
189        op: CompareOp,
190        launch_stream: StreamId,
191    ) -> Result<CudaBuffer> {
192        let runtime = self.memory.runtime().ok_or_else(|| {
193            XlogError::Kernel(
194                "filter_fused_scan_recorded requires a runtime-backed GpuMemoryManager \
195                 (constructed via with_runtime)"
196                    .to_string(),
197            )
198        })?;
199        let pool = runtime.stream_pool();
200        let cu_stream = pool.resolve(launch_stream).ok_or_else(|| {
201            XlogError::Kernel(format!(
202                "filter_fused_scan_recorded: launch_stream StreamId({}) does not resolve",
203                launch_stream.0
204            ))
205        })?;
206        let _capture_exclusion =
207            crate::cuda_graph::reserve_uncaptured_stream(&cu_stream).map_err(|error| {
208                XlogError::Kernel(format!(
209                    "filter_fused_scan_recorded requires an uncaptured stream: {error}"
210                ))
211            })?;
212
213        if input.num_rows() > u32::MAX as u64 {
214            return Err(XlogError::Kernel(format!(
215                "filter supports at most {} rows, got {}",
216                u32::MAX,
217                input.num_rows()
218            )));
219        }
220        if col >= input.arity() {
221            return Err(XlogError::Kernel(format!(
222                "Column index {} out of bounds (arity {})",
223                col,
224                input.arity()
225            )));
226        }
227        let col_type = input
228            .schema()
229            .column_type(col)
230            .ok_or_else(|| XlogError::Kernel("Missing column type".into()))?;
231        if !T::allowed_scalar_types().contains(&col_type) {
232            return Err(XlogError::Kernel(format!(
233                "Column {} is {:?} (expected one of {:?})",
234                col,
235                col_type,
236                T::allowed_scalar_types()
237            )));
238        }
239        if input.is_empty() {
240            return self.create_empty_buffer(input.schema.clone());
241        }
242
243        let n = input.num_rows() as usize;
244        let num_rows = input.num_rows() as u32;
245        let block_size = 256u32;
246        let num_blocks = num_rows.div_ceil(block_size);
247        let row_cap = u64::from(num_rows);
248        let device = self.device.inner();
249
250        let col_data = input
251            .column(col)
252            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col)))?;
253        let col_view = Self::column_as_typed_view::<T>(col_data, n)?;
254
255        // Allocate ALL fresh runtime-backed buffers BEFORE the
256        // recorder (Rust drop order — these are recorded as
257        // pre-launch writes via the standard `write` API; the
258        // recorder snapshots block identity at record time so
259        // the kernel `&mut` borrow after preflight is unaffected).
260        let d_mask = self.memory.alloc::<u8>(n)?;
261        let d_prefix_sum = self.memory.alloc::<u32>(n)?;
262        let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
263        let mut d_out_count = self.memory.alloc::<u32>(1)?;
264
265        let mut dst_cols: Vec<TrackedCudaSlice<u8>> = Vec::with_capacity(input.columns.len());
266        for col_idx in 0..input.columns.len() {
267            let elem_size = input
268                .schema
269                .column_type(col_idx)
270                .map(|t| t.size_bytes())
271                .unwrap_or(4);
272            let output_bytes = (row_cap as usize) * elem_size;
273            dst_cols.push(self.memory.alloc::<u8>(output_bytes)?);
274        }
275
276        let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(num_blocks)?;
277        let mut rec = LaunchRecorder::new_strict(launch_stream);
278        rec.read(input.num_rows_device());
279        for col_idx in 0..input.columns.len() {
280            let src_col = input
281                .column(col_idx)
282                .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
283            rec.read_column(src_col);
284        }
285        rec.write(&d_mask);
286        rec.write(&d_prefix_sum);
287        rec.write(&d_block_sums);
288        rec.write(&d_out_count);
289        for dst_col in &dst_cols {
290            rec.write(dst_col);
291        }
292        for level in scan_scratch.levels() {
293            rec.read_write(level);
294        }
295        rec.preflight(runtime).map_err(|e| {
296            XlogError::Kernel(format!(
297                "filter_fused_scan_recorded: launch recorder preflight failed: {}",
298                e
299            ))
300        })?;
301
302        // Step 1: fused compare + block-local scan.
303        let scan_kernel_name = T::filter_scan_phase1_kernel()
304            .expect("filter_fused_scan_recorded called without scan phase1 kernel");
305        let filter_scan_fn = device
306            .get_func(FILTER_MODULE, scan_kernel_name)
307            .ok_or_else(|| XlogError::Kernel(format!("{} kernel not found", scan_kernel_name)))?;
308        let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
309            // SAFETY: filter_compare_*_scan_phase1(column, constant, num_rows,
310            // num_rows_device, op, mask, prefix_sum, block_sums)
311            unsafe {
312                filter_scan_fn.clone().launch_in(
313                    stream,
314                    LaunchConfig {
315                        grid_dim: (num_blocks, 1, 1),
316                        block_dim: (block_size, 1, 1),
317                        shared_mem_bytes: 0,
318                    },
319                    (
320                        &col_view,
321                        value,
322                        num_rows,
323                        input.num_rows_device(),
324                        op as u8,
325                        &d_mask,
326                        &d_prefix_sum,
327                        &d_block_sums,
328                    ),
329                )
330            }
331            .map_err(|e| {
332                XlogError::Kernel(format!("{} (on_stream) failed: {}", scan_kernel_name, e))
333            })?;
334
335            // Step 2: multi-block scan propagation (only when there
336            // is more than one block).
337            if num_blocks > 1 {
338                self.multiblock_scan_u32_inplace_on_stream(
339                    &mut d_block_sums,
340                    num_blocks,
341                    stream,
342                    &mut scan_scratch,
343                )?;
344
345                let phase3_fn = device
346                    .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
347                    .ok_or_else(|| {
348                        XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
349                    })?;
350                // SAFETY: multiblock_scan_phase3(prefix_sum, block_offsets, n)
351                unsafe {
352                    phase3_fn.clone().launch_in(
353                        stream,
354                        LaunchConfig {
355                            grid_dim: (num_blocks, 1, 1),
356                            block_dim: (block_size, 1, 1),
357                            shared_mem_bytes: 0,
358                        },
359                        (&d_prefix_sum, &d_block_sums, num_rows),
360                    )
361                }
362                .map_err(|e| {
363                    XlogError::Kernel(format!("multiblock_scan_phase3 (on_stream) failed: {}", e))
364                })?;
365            }
366
367            // Step 3: capture compact count on launch_stream.
368            let capture_fn = device
369                .get_func(FILTER_MODULE, filter_kernels::CAPTURE_COMPACT_COUNT)
370                .ok_or_else(|| {
371                    XlogError::Kernel("capture_compact_count kernel not found".to_string())
372                })?;
373            // SAFETY: capture_compact_count(prefix_sum, mask, n, out_count)
374            unsafe {
375                capture_fn.clone().launch_in(
376                    stream,
377                    LaunchConfig {
378                        grid_dim: (1, 1, 1),
379                        block_dim: (1, 1, 1),
380                        shared_mem_bytes: 0,
381                    },
382                    (&d_prefix_sum, &d_mask, num_rows, &mut d_out_count),
383                )
384            }
385            .map_err(|e| {
386                XlogError::Kernel(format!("capture_compact_count (on_stream) failed: {}", e))
387            })?;
388
389            // Step 4: compact every column on this stream. The kernel writes
390            // only rows whose mask is nonzero, including an all-zero result.
391            let compact_fn = device
392                .get_func(FILTER_MODULE, filter_kernels::COMPACT_BYTES_BY_MASK)
393                .ok_or_else(|| {
394                    XlogError::Kernel("compact_bytes_by_mask kernel not found".to_string())
395                })?;
396            let cfg = LaunchConfig {
397                grid_dim: (num_blocks, 1, 1),
398                block_dim: (block_size, 1, 1),
399                shared_mem_bytes: 0,
400            };
401            for (col_idx, dst_col) in dst_cols.iter().enumerate() {
402                let src_col = input
403                    .column(col_idx)
404                    .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
405                let elem_size = input
406                    .schema
407                    .column_type(col_idx)
408                    .map(|t| t.size_bytes())
409                    .unwrap_or(4) as u32;
410                // SAFETY: compact_bytes_by_mask(input, mask, prefix_sum, n, elem_size, output)
411                unsafe {
412                    compact_fn.clone().launch_in(
413                        stream,
414                        cfg,
415                        (
416                            src_col,
417                            &d_mask,
418                            &d_prefix_sum,
419                            num_rows,
420                            elem_size,
421                            dst_col,
422                        ),
423                    )
424                }
425                .map_err(|e| {
426                    XlogError::Kernel(format!("compact_bytes_by_mask (on_stream) failed: {}", e))
427                })?;
428            }
429
430            Ok(())
431        };
432        // SAFETY: the callback uses the preflighted stream and only buffers
433        // owned by this recorder; it completes before their borrows end.
434        let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
435            .map_err(|error| error.into_xlog_error())?;
436
437        rec.commit().map_err(|e| {
438            XlogError::Kernel(format!(
439                "filter_fused_scan_recorded: launch recorder commit failed: {}",
440                e
441            ))
442        })?;
443
444        // Commit publishes the complete producer chain before the separate
445        // metadata read admits its default-stream read dependency.
446        let output_rows = self.dtoh_scalar_untracked(&d_out_count, 0)?;
447
448        let new_columns: Vec<CudaColumn> = dst_cols.into_iter().map(|s| s.into()).collect();
449        Ok(CudaBuffer::from_columns_with_host_count(
450            new_columns,
451            row_cap,
452            d_out_count,
453            input.schema.clone(),
454            output_rows,
455        ))
456    }
457
458    /// Strict-recorder, end-to-end variant of column-column
459    /// filter: keep rows where `column[left] <op> column[right]`.
460    ///
461    /// Composes [`Self::compare_columns_mask_recorded`] and
462    /// [`Self::compact_buffer_by_device_mask_counted_recorded`]
463    /// on a single `launch_stream`. Same composition contract
464    /// as [`Self::filter_recorded`]: each primitive builds its
465    /// own recorder and commits independently; the runtime
466    /// appends every recorded event to `last_use_events`, and
467    /// `deallocate` waits on every event, so input columns
468    /// referenced by BOTH the compare AND the per-column
469    /// compacts are correctly gated end-to-end.
470    ///
471    /// # Errors
472    /// Propagates the structured `XlogError::Kernel` errors
473    /// produced by either underlying recorded primitive
474    /// (legacy manager, unresolved launch_stream, external
475    /// column on either input side, preflight / commit
476    /// failures, kernel launch failures,
477    /// `cu_stream.synchronize()` before host scalar read).
478    pub fn filter_columns_recorded<T: GpuScalar>(
479        &self,
480        input: &CudaBuffer,
481        left: usize,
482        right: usize,
483        op: CompareOp,
484        launch_stream: StreamId,
485    ) -> Result<CudaBuffer> {
486        if input.is_empty() {
487            return self.create_empty_buffer(input.schema.clone());
488        }
489        let d_mask =
490            self.compare_columns_mask_recorded::<T>(input, left, right, op, launch_stream)?;
491        self.compact_buffer_by_device_mask_counted_recorded(input, &d_mask, launch_stream)
492    }
493
494    // ------------------------------------------------------------------
495    // Private helpers (moved from mod.rs)
496    // ------------------------------------------------------------------
497
498    /// Generate a device mask by comparing a column against a constant value.
499    ///
500    /// Each output byte is 1 if the comparison holds, 0 otherwise.
501    pub(crate) fn compare_const_mask<T: KernelScalar>(
502        &self,
503        input: &CudaBuffer,
504        col: usize,
505        value: T,
506        op: CompareOp,
507        allowed_types: &[ScalarType],
508        kernel: &str,
509    ) -> Result<TrackedCudaSlice<u8>> {
510        if input.num_rows() > u32::MAX as u64 {
511            return Err(XlogError::Kernel(format!(
512                "Filter supports at most {} rows, got {}",
513                u32::MAX,
514                input.num_rows()
515            )));
516        }
517        if col >= input.arity() {
518            return Err(XlogError::Kernel(format!(
519                "Column index {} out of bounds (arity {})",
520                col,
521                input.arity()
522            )));
523        }
524
525        if input.is_empty() {
526            return self.memory.alloc::<u8>(0);
527        }
528
529        let col_type = input
530            .schema()
531            .column_type(col)
532            .ok_or_else(|| XlogError::Kernel("Missing column type".into()))?;
533        if !allowed_types.contains(&col_type) {
534            return Err(XlogError::Kernel(format!(
535                "Column {} is {:?} (expected {:?})",
536                col, col_type, allowed_types
537            )));
538        }
539
540        let num_rows = input.num_rows() as u32;
541        let expected_bytes = (num_rows as usize)
542            .checked_mul(std::mem::size_of::<T>())
543            .ok_or_else(|| XlogError::Kernel("filter compare size overflow".into()))?;
544        let col_data = input
545            .column(col)
546            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col)))?;
547        if col_data.num_bytes() != expected_bytes {
548            return Err(XlogError::Kernel(format!(
549                "Column {} has {} bytes but expected {} for {} rows",
550                col,
551                col_data.num_bytes(),
552                expected_bytes,
553                input.num_rows()
554            )));
555        }
556
557        let block_size = 256u32;
558        let num_blocks = num_rows.div_ceil(block_size);
559        let config = LaunchConfig {
560            grid_dim: (num_blocks, 1, 1),
561            block_dim: (block_size, 1, 1),
562            shared_mem_bytes: 0,
563        };
564
565        let mut d_mask = self.memory.alloc::<u8>(num_rows as usize)?;
566        let func = self
567            .device
568            .inner()
569            .get_func(FILTER_MODULE, kernel)
570            .ok_or_else(|| XlogError::Kernel("filter compare kernel not found".into()))?;
571
572        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
573        unsafe {
574            func.clone()
575                .launch(config, (col_data, value, num_rows, op as u8, &mut d_mask))
576        }
577        .map_err(|e| XlogError::Kernel(format!("filter compare failed: {}", e)))?;
578
579        Ok(d_mask)
580    }
581
582    /// Strict-recorder variant of [`Self::compare_const_mask`].
583    ///
584    /// Runs the filter compare kernel on the caller-supplied
585    /// `launch_stream` and threads the column read through the
586    /// runtime via [`LaunchRecorder`]. This is the second
587    /// migrated launch path (after `memset_recorded`) and the
588    /// first kernel-driven one — it is intentionally a sibling
589    /// of the legacy [`Self::compare_const_mask`] rather than a
590    /// replacement. Existing callers stay on the legacy path
591    /// until the broader filter migration lands.
592    ///
593    /// # Strict-mode contract
594    /// * Requires the provider's manager to be built via
595    ///   [`crate::GpuMemoryManager::with_runtime`]; otherwise
596    ///   returns `XlogError::Kernel` before any allocation.
597    /// * `input.column(col)` is recorded as a read; external
598    ///   (`CudaColumn::Dlpack` / `CudaColumn::ArrowDevice`)
599    ///   columns are rejected at preflight, before the kernel
600    ///   is enqueued.
601    /// * `d_mask` is freshly allocated through the same
602    ///   runtime-backed manager. By construction its
603    ///   `runtime_block()` is `Some`. Its owning slice is
604    ///   registered as a write before preflight, which snapshots
605    ///   the allocation identity without retaining the borrow.
606    ///   Preflight queues its dependency waits, the kernel is
607    ///   then enqueued with `&mut d_mask`, and commit afterward
608    ///   finalizes the use by publishing the write event. Future
609    ///   recorded launchers must likewise register each owning
610    ///   slice or column before preflight.
611    ///
612    /// # Errors
613    ///   * `XlogError::Kernel` if the manager has no runtime,
614    ///     or if `launch_stream` does not resolve.
615    ///   * `XlogError::Kernel` from preflight (external column,
616    ///     unsupported active resource).
617    ///   * `XlogError::Kernel` from the underlying CUDA launch.
618    ///   * `XlogError::Kernel` from commit on transient
619    ///     `record_block_use` failure.
620    pub fn compare_const_mask_recorded<T: GpuScalar>(
621        &self,
622        input: &CudaBuffer,
623        col: usize,
624        value: T,
625        op: CompareOp,
626        launch_stream: StreamId,
627    ) -> Result<TrackedCudaSlice<u8>> {
628        let allowed_types = T::allowed_scalar_types();
629        let kernel = T::filter_compare_kernel();
630        let runtime = self.memory.runtime().ok_or_else(|| {
631            XlogError::Kernel(
632                "compare_const_mask_recorded requires a runtime-backed GpuMemoryManager \
633                 (constructed via with_runtime)"
634                    .to_string(),
635            )
636        })?;
637        let pool = runtime.stream_pool();
638        let cu_stream = pool.resolve(launch_stream).ok_or_else(|| {
639            XlogError::Kernel(format!(
640                "compare_const_mask_recorded: launch_stream StreamId({}) does not resolve",
641                launch_stream.0
642            ))
643        })?;
644
645        if input.num_rows() > u32::MAX as u64 {
646            return Err(XlogError::Kernel(format!(
647                "Filter supports at most {} rows, got {}",
648                u32::MAX,
649                input.num_rows()
650            )));
651        }
652        if col >= input.arity() {
653            return Err(XlogError::Kernel(format!(
654                "Column index {} out of bounds (arity {})",
655                col,
656                input.arity()
657            )));
658        }
659
660        if input.is_empty() {
661            return self.memory.alloc::<u8>(0);
662        }
663
664        let col_type = input
665            .schema()
666            .column_type(col)
667            .ok_or_else(|| XlogError::Kernel("Missing column type".into()))?;
668        if !allowed_types.contains(&col_type) {
669            return Err(XlogError::Kernel(format!(
670                "Column {} is {:?} (expected {:?})",
671                col, col_type, allowed_types
672            )));
673        }
674
675        let num_rows = input.num_rows() as u32;
676        let expected_bytes = (num_rows as usize)
677            .checked_mul(std::mem::size_of::<T>())
678            .ok_or_else(|| XlogError::Kernel("filter compare size overflow".into()))?;
679        let col_data = input
680            .column(col)
681            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col)))?;
682        if col_data.num_bytes() != expected_bytes {
683            return Err(XlogError::Kernel(format!(
684                "Column {} has {} bytes but expected {} for {} rows",
685                col,
686                col_data.num_bytes(),
687                expected_bytes,
688                input.num_rows()
689            )));
690        }
691
692        let block_size = 256u32;
693        let num_blocks = num_rows.div_ceil(block_size);
694        let config = LaunchConfig {
695            grid_dim: (num_blocks, 1, 1),
696            block_dim: (block_size, 1, 1),
697            shared_mem_bytes: 0,
698        };
699
700        let mut d_mask = self.memory.alloc::<u8>(num_rows as usize)?;
701        let func = self
702            .device
703            .inner()
704            .get_func(FILTER_MODULE, kernel)
705            .ok_or_else(|| XlogError::Kernel("filter compare kernel not found".into()))?;
706
707        // Strict recorder + PREFLIGHT before the kernel queues.
708        // External columns (DLPack / Arrow) and unsupported
709        // active resources are caught here without any CUDA
710        // work in flight.
711        let mut rec = LaunchRecorder::new_strict(launch_stream);
712        rec.read_column(col_data);
713        rec.write(&d_mask);
714        rec.preflight(runtime).map_err(|e| {
715            XlogError::Kernel(format!(
716                "compare_const_mask_recorded: launch recorder preflight failed: {}",
717                e
718            ))
719        })?;
720
721        // SAFETY: PTX kernel signature matches the params tuple;
722        // col_data is validated above and lives through the
723        // launch (held by `input`); d_mask was allocated by the
724        // same runtime-backed manager and matches `num_rows`.
725        // launch_on_stream queues on `cu_stream` and returns
726        // immediately.
727        let rec = unsafe {
728            rec.enqueue_prepared_with(&cu_stream, |stream| {
729                func.clone()
730                    .launch_in(
731                        stream,
732                        config,
733                        (col_data, value, num_rows, op as u8, &mut d_mask),
734                    )
735                    .map_err(|e| {
736                        XlogError::Kernel(format!(
737                            "compare_const_mask_recorded launch failed: {}",
738                            e
739                        ))
740                    })
741            })
742        }
743        .map_err(|error| error.into_xlog_error())?;
744
745        // Finalize the preflighted write registration AFTER the
746        // launch enqueues: commit publishes the write event for
747        // future dependent uses. See the "Strict-mode contract"
748        // on this method.
749        rec.commit().map_err(|e| {
750            XlogError::Kernel(format!(
751                "compare_const_mask_recorded: launch recorder commit failed: {}",
752                e
753            ))
754        })?;
755
756        Ok(d_mask)
757    }
758
759    /// Generate a device mask by comparing two columns element-wise.
760    pub(crate) fn compare_columns_mask<T: DeviceRepr>(
761        &self,
762        input: &CudaBuffer,
763        left: usize,
764        right: usize,
765        op: CompareOp,
766        allowed_types: &[ScalarType],
767        kernel: &str,
768    ) -> Result<TrackedCudaSlice<u8>> {
769        if input.num_rows() > u32::MAX as u64 {
770            return Err(XlogError::Kernel(format!(
771                "Filter supports at most {} rows, got {}",
772                u32::MAX,
773                input.num_rows()
774            )));
775        }
776        if left >= input.arity() || right >= input.arity() {
777            return Err(XlogError::Kernel(format!(
778                "Column indices {} or {} out of bounds (arity {})",
779                left,
780                right,
781                input.arity()
782            )));
783        }
784
785        if input.is_empty() {
786            return self.memory.alloc::<u8>(0);
787        }
788
789        let left_type = input
790            .schema()
791            .column_type(left)
792            .ok_or_else(|| XlogError::Kernel("Missing left column type".into()))?;
793        let right_type = input
794            .schema()
795            .column_type(right)
796            .ok_or_else(|| XlogError::Kernel("Missing right column type".into()))?;
797
798        if left_type != right_type {
799            return Err(XlogError::Kernel(
800                "Column-column compare requires matching types".into(),
801            ));
802        }
803        if !allowed_types.contains(&left_type) {
804            return Err(XlogError::Kernel(format!(
805                "Column type {:?} not supported for compare",
806                left_type
807            )));
808        }
809
810        let num_rows = input.num_rows() as u32;
811        let expected_bytes = (num_rows as usize)
812            .checked_mul(std::mem::size_of::<T>())
813            .ok_or_else(|| XlogError::Kernel("compare columns size overflow".into()))?;
814
815        let left_col = input
816            .column(left)
817            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", left)))?;
818        let right_col = input
819            .column(right)
820            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", right)))?;
821
822        if left_col.num_bytes() != expected_bytes || right_col.num_bytes() != expected_bytes {
823            return Err(XlogError::Kernel(format!(
824                "Compare columns expect {} bytes per column for {} rows",
825                expected_bytes,
826                input.num_rows()
827            )));
828        }
829
830        let block_size = 256u32;
831        let num_blocks = num_rows.div_ceil(block_size);
832        let config = LaunchConfig {
833            grid_dim: (num_blocks, 1, 1),
834            block_dim: (block_size, 1, 1),
835            shared_mem_bytes: 0,
836        };
837
838        let mut d_mask = self.memory.alloc::<u8>(num_rows as usize)?;
839        let func = self
840            .device
841            .inner()
842            .get_func(FILTER_MODULE, kernel)
843            .ok_or_else(|| XlogError::Kernel("filter compare kernel not found".into()))?;
844
845        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
846        unsafe {
847            func.clone().launch(
848                config,
849                (left_col, right_col, num_rows, op as u8, &mut d_mask),
850            )
851        }
852        .map_err(|e| XlogError::Kernel(format!("filter compare failed: {}", e)))?;
853
854        Ok(d_mask)
855    }
856
857    /// Strict-recorder variant of [`Self::compare_columns_mask`].
858    ///
859    /// Runs the column-column compare kernel on the
860    /// caller-supplied `launch_stream` and threads BOTH column
861    /// reads through the runtime via [`LaunchRecorder`]. Sibling
862    /// of the legacy [`Self::compare_columns_mask`]; existing
863    /// callers stay on the legacy path.
864    ///
865    /// # Strict-mode contract
866    /// * Requires the provider's manager to be built via
867    ///   [`crate::GpuMemoryManager::with_runtime`]; otherwise
868    ///   returns `XlogError::Kernel` before any allocation.
869    /// * `input.column(left)` and `input.column(right)` are both
870    ///   recorded as reads BEFORE preflight. External (DLPack /
871    ///   Arrow) columns on either side are rejected at preflight,
872    ///   before the kernel is enqueued.
873    /// * `d_mask` is freshly allocated by the same runtime-backed
874    ///   manager; its write is recorded via the standard `write`
875    ///   API BEFORE preflight (the recorder snapshots block
876    ///   identity, so the kernel `&mut d_mask` borrow after
877    ///   preflight is unaffected).
878    ///
879    /// # Errors
880    ///   * `XlogError::Kernel` if the manager has no runtime,
881    ///     or if `launch_stream` does not resolve.
882    ///   * `XlogError::Kernel` from preflight (external column
883    ///     on either side, unsupported active resource).
884    ///   * `XlogError::Kernel` from the underlying CUDA launch.
885    ///   * `XlogError::Kernel` from commit on transient
886    ///     `record_block_use` failure.
887    pub fn compare_columns_mask_recorded<T: GpuScalar>(
888        &self,
889        input: &CudaBuffer,
890        left: usize,
891        right: usize,
892        op: CompareOp,
893        launch_stream: StreamId,
894    ) -> Result<TrackedCudaSlice<u8>> {
895        let allowed_types = T::allowed_scalar_types();
896        let kernel = T::compare_col_kernel();
897
898        let runtime = self.memory.runtime().ok_or_else(|| {
899            XlogError::Kernel(
900                "compare_columns_mask_recorded requires a runtime-backed GpuMemoryManager \
901                 (constructed via with_runtime)"
902                    .to_string(),
903            )
904        })?;
905        let pool = runtime.stream_pool();
906        let cu_stream = pool.resolve(launch_stream).ok_or_else(|| {
907            XlogError::Kernel(format!(
908                "compare_columns_mask_recorded: launch_stream StreamId({}) does not resolve",
909                launch_stream.0
910            ))
911        })?;
912
913        if input.num_rows() > u32::MAX as u64 {
914            return Err(XlogError::Kernel(format!(
915                "Filter supports at most {} rows, got {}",
916                u32::MAX,
917                input.num_rows()
918            )));
919        }
920        if left >= input.arity() || right >= input.arity() {
921            return Err(XlogError::Kernel(format!(
922                "Column indices {} or {} out of bounds (arity {})",
923                left,
924                right,
925                input.arity()
926            )));
927        }
928
929        if input.is_empty() {
930            return self.memory.alloc::<u8>(0);
931        }
932
933        let left_type = input
934            .schema()
935            .column_type(left)
936            .ok_or_else(|| XlogError::Kernel("Missing left column type".into()))?;
937        let right_type = input
938            .schema()
939            .column_type(right)
940            .ok_or_else(|| XlogError::Kernel("Missing right column type".into()))?;
941        if left_type != right_type {
942            return Err(XlogError::Kernel(
943                "Column-column compare requires matching types".into(),
944            ));
945        }
946        if !allowed_types.contains(&left_type) {
947            return Err(XlogError::Kernel(format!(
948                "Column type {:?} not supported for compare",
949                left_type
950            )));
951        }
952
953        let num_rows = input.num_rows() as u32;
954        let expected_bytes = (num_rows as usize)
955            .checked_mul(std::mem::size_of::<T>())
956            .ok_or_else(|| XlogError::Kernel("compare columns size overflow".into()))?;
957        let left_col = input
958            .column(left)
959            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", left)))?;
960        let right_col = input
961            .column(right)
962            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", right)))?;
963        if left_col.num_bytes() != expected_bytes || right_col.num_bytes() != expected_bytes {
964            return Err(XlogError::Kernel(format!(
965                "Compare columns expect {} bytes per column for {} rows",
966                expected_bytes,
967                input.num_rows()
968            )));
969        }
970
971        let block_size = 256u32;
972        let num_blocks = num_rows.div_ceil(block_size);
973        let config = LaunchConfig {
974            grid_dim: (num_blocks, 1, 1),
975            block_dim: (block_size, 1, 1),
976            shared_mem_bytes: 0,
977        };
978
979        let mut d_mask = self.memory.alloc::<u8>(num_rows as usize)?;
980        let func = self
981            .device
982            .inner()
983            .get_func(FILTER_MODULE, kernel)
984            .ok_or_else(|| XlogError::Kernel("filter compare kernel not found".into()))?;
985
986        // Record BOTH column reads BEFORE preflight. Strict mode
987        // catches external columns on either side here, before
988        // any CUDA work is queued.
989        let mut rec = LaunchRecorder::new_strict(launch_stream);
990        rec.read_column(left_col);
991        rec.read_column(right_col);
992        rec.write(&d_mask);
993        rec.preflight(runtime).map_err(|e| {
994            XlogError::Kernel(format!(
995                "compare_columns_mask_recorded: launch recorder preflight failed: {}",
996                e
997            ))
998        })?;
999
1000        // SAFETY: PTX kernel signature matches the params tuple;
1001        // both columns are validated above and live through the
1002        // launch (held by `input`); d_mask was allocated by the
1003        // same runtime-backed manager and matches `num_rows`.
1004        // launch_on_stream queues on `cu_stream` and returns
1005        // immediately.
1006        let rec = unsafe {
1007            rec.enqueue_prepared_with(&cu_stream, |stream| {
1008                func.clone()
1009                    .launch_in(
1010                        stream,
1011                        config,
1012                        (left_col, right_col, num_rows, op as u8, &mut d_mask),
1013                    )
1014                    .map_err(|e| {
1015                        XlogError::Kernel(format!(
1016                            "compare_columns_mask_recorded launch failed: {}",
1017                            e
1018                        ))
1019                    })
1020            })
1021        }
1022        .map_err(|error| error.into_xlog_error())?;
1023
1024        // d_mask was registered as a write before preflight;
1025        // preflight waited for dependencies, and the kernel is now
1026        // enqueued. Commit publishes the write event for future uses.
1027        rec.commit().map_err(|e| {
1028            XlogError::Kernel(format!(
1029                "compare_columns_mask_recorded: launch recorder commit failed: {}",
1030                e
1031            ))
1032        })?;
1033
1034        Ok(d_mask)
1035    }
1036
1037    // ------------------------------------------------------------------
1038    // Fused compare+scan+compact path (generic over T)
1039    // ------------------------------------------------------------------
1040
1041    /// Fused compare+scan+compact filter. Used for types with a dedicated
1042    /// `filter_scan_phase1` kernel (u32 and f64).
1043    fn filter_fused_scan<T: GpuScalar>(
1044        &self,
1045        input: &CudaBuffer,
1046        col: usize,
1047        value: T,
1048        op: CompareOp,
1049    ) -> Result<CudaBuffer> {
1050        if input.num_rows() > u32::MAX as u64 {
1051            return Err(XlogError::Kernel(format!(
1052                "filter supports at most {} rows, got {}",
1053                u32::MAX,
1054                input.num_rows()
1055            )));
1056        }
1057
1058        let n = input.num_rows() as usize;
1059        let num_rows = input.num_rows() as u32;
1060        let device = self.device.inner();
1061
1062        // Validate column index
1063        if col >= input.arity() {
1064            return Err(XlogError::Kernel(format!(
1065                "Column index {} out of bounds (arity {})",
1066                col,
1067                input.arity()
1068            )));
1069        }
1070
1071        // Validate column type
1072        let col_type = input
1073            .schema()
1074            .column_type(col)
1075            .ok_or_else(|| XlogError::Kernel("Missing column type".into()))?;
1076        if !T::allowed_scalar_types().contains(&col_type) {
1077            return Err(XlogError::Kernel(format!(
1078                "Column {} is {:?} (expected one of {:?})",
1079                col,
1080                col_type,
1081                T::allowed_scalar_types()
1082            )));
1083        }
1084
1085        // Get the filter column as a typed view
1086        let col_data = input
1087            .column(col)
1088            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col)))?;
1089        let col_view = Self::column_as_typed_view::<T>(col_data, n)?;
1090
1091        let block_size = 256u32;
1092        let num_blocks = num_rows.div_ceil(block_size);
1093        let config = LaunchConfig {
1094            grid_dim: (num_blocks, 1, 1),
1095            block_dim: (block_size, 1, 1),
1096            shared_mem_bytes: 0,
1097        };
1098
1099        // Fused compare + scan phase1.
1100        let d_mask = self.memory.alloc::<u8>(n)?;
1101        let d_prefix_sum = self.memory.alloc::<u32>(n)?;
1102        let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1103
1104        let scan_kernel_name = T::filter_scan_phase1_kernel()
1105            .expect("filter_fused_scan called without scan phase1 kernel");
1106        let filter_scan_fn = device
1107            .get_func(FILTER_MODULE, scan_kernel_name)
1108            .ok_or_else(|| XlogError::Kernel(format!("{} kernel not found", scan_kernel_name)))?;
1109
1110        // SAFETY: filter_compare_*_scan_phase1(column, constant, num_rows, num_rows_device, op, mask, prefix_sum, block_sums)
1111        unsafe {
1112            filter_scan_fn.clone().launch(
1113                config,
1114                (
1115                    &col_view,
1116                    value,
1117                    num_rows,
1118                    input.num_rows_device(),
1119                    op as u8,
1120                    &d_mask,
1121                    &d_prefix_sum,
1122                    &d_block_sums,
1123                ),
1124            )
1125        }
1126        .map_err(|e| XlogError::Kernel(format!("{} failed: {}", scan_kernel_name, e)))?;
1127
1128        if num_blocks > 1 {
1129            self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1130
1131            let phase3_fn = device
1132                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1133                .ok_or_else(|| {
1134                    XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1135                })?;
1136
1137            // SAFETY: multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
1138            unsafe {
1139                phase3_fn.clone().launch(
1140                    LaunchConfig {
1141                        grid_dim: (num_blocks, 1, 1),
1142                        block_dim: (block_size, 1, 1),
1143                        shared_mem_bytes: 0,
1144                    },
1145                    (&d_prefix_sum, &d_block_sums, num_rows),
1146                )
1147            }
1148            .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
1149        }
1150
1151        self.device.synchronize()?;
1152
1153        let d_out_count = self.capture_compact_count(&d_prefix_sum, &d_mask, num_rows)?;
1154        self.compact_buffer_by_device_mask_device_count(input, &d_mask, &d_prefix_sum, d_out_count)
1155    }
1156
1157    // ------------------------------------------------------------------
1158    // Generic column view helper
1159    // ------------------------------------------------------------------
1160
1161    /// Reinterpret a `CudaColumn` as a typed owner-bearing view for kernel access.
1162    ///
1163    /// This is the generic equivalent of `column_as_u32_view`, `column_as_f64_view`, etc.
1164    fn column_as_typed_view<T: GpuScalar>(
1165        col: &CudaColumn,
1166        num_elements: usize,
1167    ) -> Result<DeviceMemoryView<T>> {
1168        let required_bytes = num_elements * T::BYTE_WIDTH;
1169        if col.num_bytes() < required_bytes {
1170            return Err(XlogError::Kernel(format!(
1171                "Column has {} bytes but {} required for {} elements of size {}",
1172                col.num_bytes(),
1173                required_bytes,
1174                num_elements,
1175                T::BYTE_WIDTH,
1176            )));
1177        }
1178        let ptr = *col.device_ptr();
1179        if T::BYTE_WIDTH > 1 && !(ptr as usize).is_multiple_of(T::BYTE_WIDTH) {
1180            return Err(XlogError::Kernel(format!(
1181                "Column device pointer is not {}-byte aligned",
1182                T::BYTE_WIDTH,
1183            )));
1184        }
1185        // SAFETY: GpuScalar defines the column representation; the cast
1186        // validates alignment and keeps the original storage owner.
1187        unsafe { col.device_view().slice(..required_bytes).cast() }
1188            .ok_or_else(|| XlogError::Kernel("device column cast is invalid".into()))
1189    }
1190
1191    // ------------------------------------------------------------------
1192    // Host-mask prefix sum + filter/compact infrastructure
1193    // (moved from mod.rs)
1194    // ------------------------------------------------------------------
1195
1196    pub fn prefix_sum_mask(&self, mask: &[u8]) -> Result<(Vec<u32>, u32)> {
1197        if mask.is_empty() {
1198            return Ok((vec![], 0));
1199        }
1200
1201        let n = mask.len();
1202        if n > u32::MAX as usize {
1203            return Err(XlogError::Kernel(format!(
1204                "Mask length {} exceeds u32::MAX",
1205                n
1206            )));
1207        }
1208
1209        // For small inputs, use CPU scan (faster than kernel launch overhead)
1210        if n <= 256 {
1211            return self.prefix_sum_mask_cpu(mask);
1212        }
1213
1214        // For larger inputs, use multi-block GPU scan
1215        self.prefix_sum_mask_gpu_multiblock(mask)
1216    }
1217
1218    /// CPU implementation for small prefix sums (avoids kernel launch overhead)
1219    fn prefix_sum_mask_cpu(&self, mask: &[u8]) -> Result<(Vec<u32>, u32)> {
1220        let mut prefix_sum = Vec::with_capacity(mask.len());
1221        let mut sum = 0u32;
1222        for &m in mask {
1223            prefix_sum.push(sum);
1224            sum += m as u32;
1225        }
1226        Ok((prefix_sum, sum))
1227    }
1228
1229    /// Multi-block GPU implementation for large prefix sums
1230    /// Uses three-phase algorithm:
1231    /// 1. Each block computes local exclusive scan and outputs block total
1232    /// 2. Scan the block totals to get block offsets
1233    /// 3. Add block offsets to each element
1234    fn prefix_sum_mask_gpu_multiblock(&self, mask: &[u8]) -> Result<(Vec<u32>, u32)> {
1235        let n = mask.len();
1236        let device = self.device.inner();
1237        let block_size = 256u32;
1238        let num_blocks = (n as u32).div_ceil(block_size);
1239
1240        // Upload mask to GPU
1241        let d_mask = self
1242            .htod_sync_copy_tracked(mask)
1243            .map_err(|e| XlogError::Kernel(format!("Failed to upload mask: {}", e)))?;
1244
1245        // Allocate output for prefix sum (using memory manager for budget enforcement)
1246        let d_prefix_sum = self.memory.alloc::<u32>(n)?;
1247
1248        // Allocate block sums array (using memory manager for budget enforcement)
1249        let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1250
1251        // Phase 1: Block-level exclusive scans + collect block totals
1252        let phase1_fn = device
1253            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1254            .ok_or_else(|| {
1255                XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1256            })?;
1257
1258        // SAFETY: Kernel parameters match expected signature:
1259        // multiblock_scan_phase1(const uint8_t* mask, uint32_t* prefix_sum, uint32_t* block_sums, uint32_t n)
1260        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
1261        unsafe {
1262            phase1_fn.clone().launch(
1263                LaunchConfig {
1264                    grid_dim: (num_blocks, 1, 1),
1265                    block_dim: (block_size, 1, 1),
1266                    shared_mem_bytes: 0,
1267                },
1268                (&d_mask, &d_prefix_sum, &d_block_sums, n as u32),
1269            )
1270        }
1271        .map_err(|e| {
1272            XlogError::Kernel(format!("Failed to launch multiblock_scan_phase1: {}", e))
1273        })?;
1274
1275        // Phase 2: Scan block sums (only if we have more than 1 block)
1276        if num_blocks > 1 {
1277            self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1278
1279            // Phase 3: Add block offsets
1280            let phase3_fn = device
1281                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1282                .ok_or_else(|| {
1283                    XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1284                })?;
1285
1286            // SAFETY: Kernel parameters match expected signature:
1287            // multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
1288            // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
1289            unsafe {
1290                phase3_fn.clone().launch(
1291                    LaunchConfig {
1292                        grid_dim: (num_blocks, 1, 1),
1293                        block_dim: (block_size, 1, 1),
1294                        shared_mem_bytes: 0,
1295                    },
1296                    (&d_prefix_sum, &d_block_sums, n as u32),
1297                )
1298            }
1299            .map_err(|e| {
1300                XlogError::Kernel(format!("Failed to launch multiblock_scan_phase3: {}", e))
1301            })?;
1302        }
1303
1304        // Synchronize and download results
1305        self.device.synchronize()?;
1306
1307        let prefix_sum = device
1308            .dtoh_sync_copy(&d_prefix_sum)
1309            .map_err(|e| XlogError::Kernel(format!("Failed to download prefix_sum: {}", e)))?;
1310
1311        // Compute count from the last prefix sum value + last mask value
1312        let count = prefix_sum[n - 1] + mask[n - 1] as u32;
1313
1314        Ok((prefix_sum, count))
1315    }
1316
1317    /// Filter buffer by pre-computed mask.
1318    ///
1319    /// # Arguments
1320    /// * `input` - The input buffer to filter
1321    /// * `mask` - Mask slice where non-zero means keep the row
1322    ///
1323    /// # Errors
1324    /// Returns error if mask length doesn't match buffer rows.
1325    pub fn filter_by_mask(&self, input: &CudaBuffer, mask: &[u8]) -> Result<CudaBuffer> {
1326        if input.num_rows() == 0 {
1327            return self.create_empty_buffer(input.schema.clone());
1328        }
1329
1330        let n = input.num_rows() as usize;
1331        if mask.len() != n {
1332            return Err(XlogError::Kernel(format!(
1333                "Mask length {} doesn't match buffer rows {}",
1334                mask.len(),
1335                n
1336            )));
1337        }
1338
1339        // Compute prefix sum and count
1340        let (prefix_sum, count) = self.prefix_sum_mask(mask)?;
1341
1342        if count == 0 {
1343            return self.create_empty_buffer(input.schema.clone());
1344        }
1345
1346        // Compact all columns using mask
1347        self.compact_buffer_by_mask(input, mask, &prefix_sum, count as u64)
1348    }
1349
1350    /// Compact buffer columns using mask and prefix sum indices
1351    fn compact_buffer_by_mask(
1352        &self,
1353        input: &CudaBuffer,
1354        mask: &[u8],
1355        prefix_sum: &[u32],
1356        output_count: u64,
1357    ) -> Result<CudaBuffer> {
1358        // Upload mask and prefix sum to GPU
1359        let d_mask = self
1360            .htod_sync_copy_tracked(mask)
1361            .map_err(|e| XlogError::Kernel(format!("Failed to upload mask: {}", e)))?;
1362        let d_prefix_sum = self
1363            .htod_sync_copy_tracked(prefix_sum)
1364            .map_err(|e| XlogError::Kernel(format!("Failed to upload prefix_sum: {}", e)))?;
1365
1366        self.compact_buffer_by_device_mask(input, &d_mask, &d_prefix_sum, output_count)
1367    }
1368
1369    /// Compact a buffer using a device-resident mask.
1370    ///
1371    /// Computes prefix sum and output count fully on-device.
1372    /// Strict-recorder variant of
1373    /// [`Self::compact_buffer_by_device_mask_counted`] — the
1374    /// first migrated COMPACT path.
1375    ///
1376    /// The compact pipeline is a multi-kernel chain:
1377    /// `mask_clamp_rows` → `multiblock_scan_phase1` →
1378    /// `multiblock_scan_u32_inplace_on_stream` (recursive,
1379    /// only when `num_blocks > 1`) → `multiblock_scan_phase3` →
1380    /// `capture_compact_count` → per-column `compact_bytes_by_mask`.
1381    /// **Every kernel runs on the same explicit `launch_stream`
1382    /// via `launch_on_stream`**. After commit publishes the complete chain,
1383    /// the separate host scalar read of `d_out_count` waits on that dependency.
1384    /// This host-count-returning method excludes stream capture before
1385    /// allocating or enqueueing work, and holds the exclusion through readback.
1386    ///
1387    /// # Strict-mode contract
1388    /// * Requires the provider's manager to be built via
1389    ///   [`crate::GpuMemoryManager::with_runtime`]; otherwise
1390    ///   returns `XlogError::Kernel` before any allocation.
1391    /// * `d_mask` is recorded as a read.
1392    /// * `input.num_rows_device()` is recorded as a read.
1393    /// * Each `input.column(i)` is recorded as a read; external
1394    ///   columns on any side are rejected at preflight, before
1395    ///   any CUDA work is enqueued.
1396    /// * Every fresh runtime-backed allocation that this
1397    ///   function makes (`d_mask_clamped`, `d_prefix_sum`,
1398    ///   `d_block_sums`, `d_out_count`, each `dst_col`) is
1399    ///   recorded via `write` BEFORE the
1400    ///   kernel chain enqueues. Locals that drop at end-of-scope
1401    ///   (`d_mask_clamped`, `d_prefix_sum`, `d_block_sums`)
1402    ///   stay safe because the runtime's deallocate queues
1403    ///   `cuStreamWaitEvent(alloc_stream, recorded_event)`
1404    ///   BEFORE `cuMemFreeAsync`, gating the free on the
1405    ///   launch_stream chain.
1406    /// * Intermediate `block_sums` allocations created by the
1407    ///   recursive scan helper are recorded directly inside the
1408    ///   helper (they don't outlive the helper call).
1409    ///
1410    /// # Errors
1411    ///   * `XlogError::Kernel` if the manager has no runtime,
1412    ///     or if `launch_stream` does not resolve.
1413    ///   * `XlogError::Kernel` from preflight (external column
1414    ///     on any side, unsupported active resource).
1415    ///   * `XlogError::Kernel` if capture admission is unavailable, or from any
1416    ///     underlying CUDA launch or the post-commit host scalar read.
1417    ///   * `XlogError::Kernel` from commit on transient
1418    ///     `record_block_use` failure.
1419    pub fn compact_buffer_by_device_mask_counted_recorded(
1420        &self,
1421        input: &CudaBuffer,
1422        d_mask: &TrackedCudaSlice<u8>,
1423        launch_stream: StreamId,
1424    ) -> Result<CudaBuffer> {
1425        let runtime = self.memory.runtime().ok_or_else(|| {
1426            XlogError::Kernel(
1427                "compact_buffer_by_device_mask_counted_recorded requires a \
1428                 runtime-backed GpuMemoryManager (constructed via with_runtime)"
1429                    .to_string(),
1430            )
1431        })?;
1432        let pool = runtime.stream_pool();
1433        let cu_stream = pool.resolve(launch_stream).ok_or_else(|| {
1434            XlogError::Kernel(format!(
1435                "compact_buffer_by_device_mask_counted_recorded: launch_stream \
1436                 StreamId({}) does not resolve",
1437                launch_stream.0
1438            ))
1439        })?;
1440        let _capture_exclusion = crate::cuda_graph::reserve_uncaptured_stream(&cu_stream)
1441            .map_err(|error| {
1442                XlogError::Kernel(format!(
1443                    "compact_buffer_by_device_mask_counted_recorded requires an uncaptured stream: {error}"
1444                ))
1445            })?;
1446
1447        let n = input.num_rows() as u32;
1448        if n == 0 {
1449            return self.create_empty_buffer(input.schema.clone());
1450        }
1451        // The bound `d_mask.len() >= n (row_cap)` is too
1452        // strict — `mask_clamp_rows` only reads `d_mask[i]`
1453        // for `i < num_rows_device` (logical count). Real
1454        // callers (hash_join_semi/anti recorded) pass masks
1455        // sized to logical count which can be < row_cap. We
1456        // require only that the mask is non-empty; OOB reads
1457        // are bounded by the kernel's own check against the
1458        // device-resident logical row count.
1459        if d_mask.is_empty() {
1460            return Err(XlogError::Kernel(
1461                "compact_buffer_by_device_mask_counted_recorded: empty d_mask".to_string(),
1462            ));
1463        }
1464
1465        let device = self.device.inner();
1466        let block_size = 256u32;
1467        let num_blocks = n.div_ceil(block_size);
1468        let row_cap = u64::from(n);
1469
1470        // Allocate ALL fresh runtime-backed buffers up front,
1471        // BEFORE the recorder is constructed. The recorder
1472        // snapshots each block's identity at record time and
1473        // drops the slice borrow, so the buffers can be
1474        // mutably borrowed by kernel launches after preflight.
1475        // Output column sizes are known up front from
1476        // `row_cap = n`, so this
1477        // is sound — the host scalar read of `d_out_count` only
1478        // tells us `output_rows`, which we use as metadata, not
1479        // for sizing.
1480        let mut d_mask_clamped = self.memory.alloc::<u8>(n as usize)?;
1481        let d_prefix_sum = self.memory.alloc::<u32>(n as usize)?;
1482        let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1483        let mut d_out_count = self.memory.alloc::<u32>(1)?;
1484
1485        let mut dst_cols: Vec<TrackedCudaSlice<u8>> = Vec::with_capacity(input.columns.len());
1486        for col_idx in 0..input.columns.len() {
1487            let elem_size = input
1488                .schema
1489                .column_type(col_idx)
1490                .map(|t| t.size_bytes())
1491                .unwrap_or(4);
1492            let output_bytes = (row_cap as usize) * elem_size;
1493            dst_cols.push(self.memory.alloc::<u8>(output_bytes)?);
1494        }
1495
1496        // Build recorder, record reads BEFORE preflight.
1497        let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(num_blocks)?;
1498        let mut rec = LaunchRecorder::new_strict(launch_stream);
1499        rec.read(d_mask);
1500        rec.read(input.num_rows_device());
1501        for col_idx in 0..input.columns.len() {
1502            let src_col = input
1503                .column(col_idx)
1504                .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
1505            rec.read_column(src_col);
1506        }
1507        rec.write(&d_mask_clamped);
1508        rec.write(&d_prefix_sum);
1509        rec.write(&d_block_sums);
1510        rec.write(&d_out_count);
1511        for dst_col in &dst_cols {
1512            rec.write(dst_col);
1513        }
1514        for level in scan_scratch.levels() {
1515            rec.read_write(level);
1516        }
1517        rec.preflight(runtime).map_err(|e| {
1518            XlogError::Kernel(format!(
1519                "compact_buffer_by_device_mask_counted_recorded: launch recorder \
1520                 preflight failed: {}",
1521                e
1522            ))
1523        })?;
1524
1525        // Step 1: mask_clamp_rows on launch_stream.
1526        let clamp_fn = device
1527            .get_func(FILTER_MODULE, filter_kernels::MASK_CLAMP_ROWS)
1528            .ok_or_else(|| XlogError::Kernel("mask_clamp_rows kernel not found".to_string()))?;
1529        let enqueue = |stream: &crate::launch::CudaEnqueue<'_>| -> Result<()> {
1530            // SAFETY: mask_clamp_rows(in_mask, num_rows_device, row_cap, out_mask)
1531            unsafe {
1532                clamp_fn.clone().launch_in(
1533                    stream,
1534                    LaunchConfig {
1535                        grid_dim: (num_blocks, 1, 1),
1536                        block_dim: (block_size, 1, 1),
1537                        shared_mem_bytes: 0,
1538                    },
1539                    (d_mask, input.num_rows_device(), n, &mut d_mask_clamped),
1540                )
1541            }
1542            .map_err(|e| XlogError::Kernel(format!("mask_clamp_rows (on_stream) failed: {}", e)))?;
1543
1544            // Step 2: multiblock_scan_phase1 on launch_stream.
1545            let phase1_fn = device
1546                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1547                .ok_or_else(|| {
1548                    XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1549                })?;
1550            // SAFETY: multiblock_scan_phase1(const u8 mask, u32 prefix_sum, u32 block_sums, u32 n)
1551            unsafe {
1552                phase1_fn.clone().launch_in(
1553                    stream,
1554                    LaunchConfig {
1555                        grid_dim: (num_blocks, 1, 1),
1556                        block_dim: (block_size, 1, 1),
1557                        shared_mem_bytes: 0,
1558                    },
1559                    (&d_mask_clamped, &d_prefix_sum, &d_block_sums, n),
1560                )
1561            }
1562            .map_err(|e| {
1563                XlogError::Kernel(format!("multiblock_scan_phase1 (on_stream) failed: {}", e))
1564            })?;
1565
1566            // Step 3: scan inplace on block_sums + phase3 propagate
1567            // (only when there is more than one block).
1568            if num_blocks > 1 {
1569                self.multiblock_scan_u32_inplace_on_stream(
1570                    &mut d_block_sums,
1571                    num_blocks,
1572                    stream,
1573                    &mut scan_scratch,
1574                )?;
1575
1576                let phase3_fn = device
1577                    .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1578                    .ok_or_else(|| {
1579                        XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1580                    })?;
1581                // SAFETY: multiblock_scan_phase3(prefix_sum, block_offsets, n)
1582                unsafe {
1583                    phase3_fn.clone().launch_in(
1584                        stream,
1585                        LaunchConfig {
1586                            grid_dim: (num_blocks, 1, 1),
1587                            block_dim: (block_size, 1, 1),
1588                            shared_mem_bytes: 0,
1589                        },
1590                        (&d_prefix_sum, &d_block_sums, n),
1591                    )
1592                }
1593                .map_err(|e| {
1594                    XlogError::Kernel(format!("multiblock_scan_phase3 (on_stream) failed: {}", e))
1595                })?;
1596            }
1597
1598            // Step 4: capture_compact_count on launch_stream.
1599            // (`d_out_count` was pre-allocated up front; see header.)
1600            let capture_fn = device
1601                .get_func(FILTER_MODULE, filter_kernels::CAPTURE_COMPACT_COUNT)
1602                .ok_or_else(|| {
1603                    XlogError::Kernel("capture_compact_count kernel not found".to_string())
1604                })?;
1605            // SAFETY: capture_compact_count(prefix_sum, mask, n, out_count)
1606            //
1607            // Use the clamped mask, not the caller's original mask.
1608            // Some recorded callers provide a mask sized to the
1609            // device-resident logical row count while `n` is the
1610            // buffer row capacity. `mask_clamp_rows` expanded that
1611            // shorter domain into a row-capacity-sized mask with
1612            // slack rows forced to zero; every downstream consumer
1613            // in this compaction chain must use that expanded mask.
1614            unsafe {
1615                capture_fn.clone().launch_in(
1616                    stream,
1617                    LaunchConfig {
1618                        grid_dim: (1, 1, 1),
1619                        block_dim: (1, 1, 1),
1620                        shared_mem_bytes: 0,
1621                    },
1622                    (&d_prefix_sum, &d_mask_clamped, n, &mut d_out_count),
1623                )
1624            }
1625            .map_err(|e| {
1626                XlogError::Kernel(format!("capture_compact_count (on_stream) failed: {}", e))
1627            })?;
1628
1629            // Step 5: compact every column with the clamped capacity-domain
1630            // mask. An all-zero mask writes nothing; slack rows stay excluded.
1631            let compact_fn = device
1632                .get_func(FILTER_MODULE, filter_kernels::COMPACT_BYTES_BY_MASK)
1633                .ok_or_else(|| {
1634                    XlogError::Kernel("compact_bytes_by_mask kernel not found".to_string())
1635                })?;
1636            let grid_size = n.div_ceil(block_size);
1637            let cfg = LaunchConfig {
1638                grid_dim: (grid_size, 1, 1),
1639                block_dim: (block_size, 1, 1),
1640                shared_mem_bytes: 0,
1641            };
1642            for (col_idx, dst_col) in dst_cols.iter().enumerate() {
1643                let src_col = input
1644                    .column(col_idx)
1645                    .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
1646                let elem_size = input
1647                    .schema
1648                    .column_type(col_idx)
1649                    .map(|t| t.size_bytes())
1650                    .unwrap_or(4) as u32;
1651                // SAFETY: compact_bytes_by_mask(input, mask, prefix_sum, n, elem_size, output)
1652                //
1653                // Same domain rule as capture_compact_count:
1654                // compact over the row-capacity-sized clamped
1655                // mask so rows >= logical_count are never
1656                // materialized from valid-looking slack.
1657                unsafe {
1658                    compact_fn.clone().launch_in(
1659                        stream,
1660                        cfg,
1661                        (
1662                            src_col,
1663                            &d_mask_clamped,
1664                            &d_prefix_sum,
1665                            n,
1666                            elem_size,
1667                            dst_col,
1668                        ),
1669                    )
1670                }
1671                .map_err(|e| {
1672                    XlogError::Kernel(format!("compact_bytes_by_mask (on_stream) failed: {}", e))
1673                })?;
1674            }
1675
1676            Ok(())
1677        };
1678        // SAFETY: the callback uses the preflighted stream and only buffers
1679        // owned by this recorder; it completes before their borrows end.
1680        let rec = unsafe { rec.enqueue_prepared_with(&cu_stream, enqueue) }
1681            .map_err(|error| error.into_xlog_error())?;
1682
1683        rec.commit().map_err(|e| {
1684            XlogError::Kernel(format!(
1685                "compact_buffer_by_device_mask_counted_recorded: launch recorder \
1686                 commit failed: {}",
1687                e
1688            ))
1689        })?;
1690
1691        // Commit publishes the complete producer chain before the separate
1692        // metadata read admits its default-stream read dependency.
1693        let output_rows = self.dtoh_scalar_untracked(&d_out_count, 0)?;
1694
1695        let new_columns: Vec<CudaColumn> = dst_cols.into_iter().map(|s| s.into()).collect();
1696        Ok(CudaBuffer::from_columns_with_host_count(
1697            new_columns,
1698            row_cap,
1699            d_out_count,
1700            input.schema.clone(),
1701            output_rows,
1702        ))
1703    }
1704
1705    pub fn compact_buffer_by_device_mask_counted(
1706        &self,
1707        input: &CudaBuffer,
1708        d_mask: &TrackedCudaSlice<u8>,
1709    ) -> Result<CudaBuffer> {
1710        let n = input.num_rows() as u32;
1711        if n == 0 {
1712            return self.create_empty_buffer(input.schema.clone());
1713        }
1714        if n as usize > d_mask.len() {
1715            return Err(XlogError::Kernel(format!(
1716                "compact_buffer_by_device_mask_counted: mask len {} < rows {}",
1717                d_mask.len(),
1718                n
1719            )));
1720        }
1721
1722        let device = self.device.inner();
1723        let block_size = 256u32;
1724        let num_blocks = n.div_ceil(block_size);
1725
1726        let mut d_mask_clamped = self.memory.alloc::<u8>(n as usize)?;
1727        let clamp_fn = device
1728            .get_func(FILTER_MODULE, filter_kernels::MASK_CLAMP_ROWS)
1729            .ok_or_else(|| XlogError::Kernel("mask_clamp_rows kernel not found".to_string()))?;
1730
1731        // SAFETY: mask_clamp_rows(in_mask, num_rows_device, row_cap, out_mask)
1732        unsafe {
1733            clamp_fn.clone().launch(
1734                LaunchConfig {
1735                    grid_dim: (num_blocks, 1, 1),
1736                    block_dim: (block_size, 1, 1),
1737                    shared_mem_bytes: 0,
1738                },
1739                (d_mask, input.num_rows_device(), n, &mut d_mask_clamped),
1740            )
1741        }
1742        .map_err(|e| XlogError::Kernel(format!("mask_clamp_rows failed: {}", e)))?;
1743
1744        let d_prefix_sum = self.memory.alloc::<u32>(n as usize)?;
1745        let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1746
1747        let phase1_fn = device
1748            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1749            .ok_or_else(|| {
1750                XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1751            })?;
1752
1753        // SAFETY: multiblock_scan_phase1(const uint8_t* mask, uint32_t* prefix_sum, uint32_t* block_sums, uint32_t n)
1754        unsafe {
1755            phase1_fn.clone().launch(
1756                LaunchConfig {
1757                    grid_dim: (num_blocks, 1, 1),
1758                    block_dim: (block_size, 1, 1),
1759                    shared_mem_bytes: 0,
1760                },
1761                (&d_mask_clamped, &d_prefix_sum, &d_block_sums, n),
1762            )
1763        }
1764        .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
1765
1766        if num_blocks > 1 {
1767            self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1768
1769            let phase3_fn = device
1770                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1771                .ok_or_else(|| {
1772                    XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1773                })?;
1774
1775            // SAFETY: multiblock_scan_phase3(prefix_sum, block_offsets, n)
1776            unsafe {
1777                phase3_fn.clone().launch(
1778                    LaunchConfig {
1779                        grid_dim: (num_blocks, 1, 1),
1780                        block_dim: (block_size, 1, 1),
1781                        shared_mem_bytes: 0,
1782                    },
1783                    (&d_prefix_sum, &d_block_sums, n),
1784                )
1785            }
1786            .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
1787        }
1788
1789        let d_out_count = self.capture_compact_count(&d_prefix_sum, d_mask, n)?;
1790
1791        self.compact_buffer_by_device_mask_device_count(input, d_mask, &d_prefix_sum, d_out_count)
1792    }
1793
1794    pub(crate) fn capture_compact_count(
1795        &self,
1796        d_prefix_sum: &impl DeviceRead<u32>,
1797        d_mask: &impl DeviceRead<u8>,
1798        n: u32,
1799    ) -> Result<TrackedCudaSlice<u32>> {
1800        let d_prefix_sum = &d_prefix_sum.device_view();
1801        let d_mask = &d_mask.device_view();
1802        let mut d_out_count = self.memory.alloc::<u32>(1)?;
1803        let device = self.device.inner();
1804        let capture_fn = device
1805            .get_func(FILTER_MODULE, filter_kernels::CAPTURE_COMPACT_COUNT)
1806            .ok_or_else(|| {
1807                XlogError::Kernel("capture_compact_count kernel not found".to_string())
1808            })?;
1809        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
1810        unsafe {
1811            capture_fn.clone().launch(
1812                LaunchConfig {
1813                    grid_dim: (1, 1, 1),
1814                    block_dim: (1, 1, 1),
1815                    shared_mem_bytes: 0,
1816                },
1817                (d_prefix_sum, d_mask, n, &mut d_out_count),
1818            )
1819        }
1820        .map_err(|e| XlogError::Kernel(format!("capture_compact_count failed: {}", e)))?;
1821        Ok(d_out_count)
1822    }
1823
1824    pub(crate) fn compact_buffer_by_device_mask_device_count(
1825        &self,
1826        input: &CudaBuffer,
1827        d_mask: &impl DeviceRead<u8>,
1828        d_prefix_sum: &impl DeviceRead<u32>,
1829        d_out_count: TrackedCudaSlice<u32>,
1830    ) -> Result<CudaBuffer> {
1831        let d_mask = &d_mask.device_view();
1832        let d_prefix_sum = &d_prefix_sum.device_view();
1833        let mask_len = u32::try_from(d_mask.len()).map_err(|_| {
1834            XlogError::Kernel(format!(
1835                "compact_buffer_by_device_mask_device_count: mask len {} exceeds u32::MAX",
1836                d_mask.len()
1837            ))
1838        })?;
1839        let prefix_len = u32::try_from(d_prefix_sum.len()).map_err(|_| {
1840            XlogError::Kernel(format!(
1841                "compact_buffer_by_device_mask_device_count: prefix sum len {} exceeds u32::MAX",
1842                d_prefix_sum.len()
1843            ))
1844        })?;
1845        if prefix_len < mask_len {
1846            return Err(XlogError::Kernel(format!(
1847                "compact_buffer_by_device_mask_device_count: prefix sum len {} < mask len {}",
1848                prefix_len, mask_len
1849            )));
1850        }
1851        if mask_len as u64 > input.num_rows() {
1852            return Err(XlogError::Kernel(format!(
1853                "compact_buffer_by_device_mask_device_count: mask len {} > row cap {}",
1854                mask_len,
1855                input.num_rows()
1856            )));
1857        }
1858        if mask_len == 0 {
1859            return self.create_empty_buffer(input.schema.clone());
1860        }
1861        let n = mask_len;
1862        let device = self.device.inner();
1863
1864        let compact_fn = device
1865            .get_func(FILTER_MODULE, filter_kernels::COMPACT_BYTES_BY_MASK)
1866            .ok_or_else(|| {
1867                XlogError::Kernel("compact_bytes_by_mask kernel not found".to_string())
1868            })?;
1869
1870        let block_size = 256u32;
1871        let grid_size = n.div_ceil(block_size);
1872        let config = LaunchConfig {
1873            grid_dim: (grid_size, 1, 1),
1874            block_dim: (block_size, 1, 1),
1875            shared_mem_bytes: 0,
1876        };
1877
1878        // The compact kernel only writes `output_rows` elements, but downstream
1879        // metadata must still remember the logical span covered by the mask.
1880        // This keeps host-visible `row_cap` aligned with the masked prefix while
1881        // the device-side row count tracks the actual compacted cardinality.
1882        let output_rows = self.dtoh_scalar_untracked(&d_out_count, 0)? as u64;
1883        let row_cap = u64::from(n);
1884
1885        if output_rows == 0 {
1886            let mut new_columns = Vec::with_capacity(input.columns.len());
1887            for col_idx in 0..input.columns.len() {
1888                let elem_size = input
1889                    .schema
1890                    .column_type(col_idx)
1891                    .map(|t| t.size_bytes())
1892                    .unwrap_or(4);
1893                let output_bytes = (row_cap as usize) * elem_size;
1894                new_columns.push(self.memory.alloc::<u8>(output_bytes)?.into());
1895            }
1896            let d_zero_rows = self.upload_device_row_count(0)?;
1897            return Ok(CudaBuffer::from_columns_with_host_count(
1898                new_columns,
1899                row_cap,
1900                d_zero_rows,
1901                input.schema.clone(),
1902                0,
1903            ));
1904        }
1905
1906        let mut new_columns = Vec::with_capacity(input.columns.len());
1907        for col_idx in 0..input.columns.len() {
1908            let src_col = input
1909                .column(col_idx)
1910                .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
1911
1912            let elem_size = input
1913                .schema
1914                .column_type(col_idx)
1915                .map(|t| t.size_bytes())
1916                .unwrap_or(4) as u32;
1917
1918            let output_bytes = (row_cap as usize) * (elem_size as usize);
1919            let dst_col = self.memory.alloc::<u8>(output_bytes)?;
1920
1921            // SAFETY: Kernel signature matches:
1922            // compact_bytes_by_mask(input, mask, prefix_sum, num_rows, elem_size, output)
1923            // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
1924            unsafe {
1925                compact_fn.clone().launch(
1926                    config,
1927                    (src_col, d_mask, d_prefix_sum, n, elem_size, &dst_col),
1928                )
1929            }
1930            .map_err(|e| XlogError::Kernel(format!("compact_bytes_by_mask failed: {}", e)))?;
1931
1932            new_columns.push(dst_col.into());
1933        }
1934
1935        self.device.synchronize()?;
1936
1937        Ok(CudaBuffer::from_columns_with_host_count(
1938            new_columns,
1939            row_cap,
1940            d_out_count,
1941            input.schema.clone(),
1942            output_rows as u32,
1943        ))
1944    }
1945
1946    fn compact_buffer_by_device_mask(
1947        &self,
1948        input: &CudaBuffer,
1949        d_mask: &impl DeviceRead<u8>,
1950        d_prefix_sum: &impl DeviceRead<u32>,
1951        output_count: u64,
1952    ) -> Result<CudaBuffer> {
1953        let d_mask = &d_mask.device_view();
1954        let d_prefix_sum = &d_prefix_sum.device_view();
1955        let n = input.num_rows() as u32;
1956        let device = self.device.inner();
1957
1958        // Get compact kernel
1959        let compact_fn = device
1960            .get_func(FILTER_MODULE, filter_kernels::COMPACT_BYTES_BY_MASK)
1961            .ok_or_else(|| {
1962                XlogError::Kernel("compact_bytes_by_mask kernel not found".to_string())
1963            })?;
1964
1965        let block_size = 256u32;
1966        let grid_size = n.div_ceil(block_size);
1967        let config = LaunchConfig {
1968            grid_dim: (grid_size, 1, 1),
1969            block_dim: (block_size, 1, 1),
1970            shared_mem_bytes: 0,
1971        };
1972
1973        // Compact each column
1974        let mut new_columns = Vec::with_capacity(input.columns.len());
1975        for col_idx in 0..input.columns.len() {
1976            let src_col = input
1977                .column(col_idx)
1978                .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
1979
1980            let elem_size = input
1981                .schema
1982                .column_type(col_idx)
1983                .map(|t| t.size_bytes())
1984                .unwrap_or(4) as u32;
1985
1986            let output_bytes = (output_count as usize) * (elem_size as usize);
1987            let dst_col = self.memory.alloc::<u8>(output_bytes)?;
1988
1989            // SAFETY: Kernel signature matches:
1990            // compact_bytes_by_mask(input, mask, prefix_sum, num_rows, elem_size, output)
1991            // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
1992            unsafe {
1993                compact_fn.clone().launch(
1994                    config,
1995                    (src_col, d_mask, d_prefix_sum, n, elem_size, &dst_col),
1996                )
1997            }
1998            .map_err(|e| XlogError::Kernel(format!("compact_bytes_by_mask failed: {}", e)))?;
1999
2000            new_columns.push(dst_col.into());
2001        }
2002
2003        self.device.synchronize()?;
2004
2005        self.buffer_from_columns(new_columns, output_count, input.schema.clone())
2006    }
2007
2008    pub fn filter_by_device_mask(
2009        &self,
2010        input: &CudaBuffer,
2011        d_mask: &impl DeviceRead<u8>,
2012    ) -> Result<CudaBuffer> {
2013        let d_mask = &d_mask.device_view();
2014        if input.is_empty() {
2015            return self.create_empty_buffer(input.schema().clone());
2016        }
2017
2018        if input.num_rows() > u32::MAX as u64 {
2019            return Err(XlogError::Kernel(format!(
2020                "Device-mask filtering supports at most {} rows, got {}",
2021                u32::MAX,
2022                input.num_rows()
2023            )));
2024        }
2025
2026        let n = input.num_rows() as u32;
2027        let device = self.device.inner();
2028
2029        let block_size = 256u32;
2030        let num_blocks = n.div_ceil(block_size);
2031
2032        let mut d_mask_clamped = self.memory.alloc::<u8>(n as usize)?;
2033        let clamp_fn = device
2034            .get_func(FILTER_MODULE, filter_kernels::MASK_CLAMP_ROWS)
2035            .ok_or_else(|| XlogError::Kernel("mask_clamp_rows kernel not found".to_string()))?;
2036
2037        // SAFETY: mask_clamp_rows(in_mask, num_rows_device, row_cap, out_mask)
2038        unsafe {
2039            clamp_fn.clone().launch(
2040                LaunchConfig {
2041                    grid_dim: (num_blocks, 1, 1),
2042                    block_dim: (block_size, 1, 1),
2043                    shared_mem_bytes: 0,
2044                },
2045                (d_mask, input.num_rows_device(), n, &mut d_mask_clamped),
2046            )
2047        }
2048        .map_err(|e| XlogError::Kernel(format!("mask_clamp_rows failed: {}", e)))?;
2049
2050        let d_prefix_sum = self.memory.alloc::<u32>(n as usize)?;
2051        let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2052
2053        let phase1_fn = device
2054            .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
2055            .ok_or_else(|| {
2056                XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
2057            })?;
2058
2059        // SAFETY: multiblock_scan_phase1(const uint8_t* mask, uint32_t* prefix_sum, uint32_t* block_sums, uint32_t n)
2060        unsafe {
2061            phase1_fn.clone().launch(
2062                LaunchConfig {
2063                    grid_dim: (num_blocks, 1, 1),
2064                    block_dim: (block_size, 1, 1),
2065                    shared_mem_bytes: 0,
2066                },
2067                (&d_mask_clamped, &d_prefix_sum, &d_block_sums, n),
2068            )
2069        }
2070        .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
2071
2072        if num_blocks > 1 {
2073            self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
2074
2075            let phase3_fn = device
2076                .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2077                .ok_or_else(|| {
2078                    XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2079                })?;
2080
2081            // SAFETY: multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
2082            unsafe {
2083                phase3_fn.clone().launch(
2084                    LaunchConfig {
2085                        grid_dim: (num_blocks, 1, 1),
2086                        block_dim: (block_size, 1, 1),
2087                        shared_mem_bytes: 0,
2088                    },
2089                    (&d_prefix_sum, &d_block_sums, n),
2090                )
2091            }
2092            .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2093        }
2094
2095        self.device.synchronize()?;
2096
2097        let d_out_count = self.capture_compact_count(&d_prefix_sum, &d_mask_clamped, n)?;
2098        self.compact_buffer_by_device_mask_device_count(
2099            input,
2100            &d_mask_clamped,
2101            &d_prefix_sum,
2102            d_out_count,
2103        )
2104    }
2105}