1use 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 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 pub fn filter<T: GpuScalar>(
51 &self,
52 input: &CudaBuffer,
53 col: usize,
54 value: T,
55 op: CompareOp,
56 ) -> Result<CudaBuffer> {
57 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 self.filter_fused_scan::<T>(input, col, value, op)
74 } else {
75 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 unsafe { col.device_view().slice(..required_bytes).cast() }
1188 .ok_or_else(|| XlogError::Kernel("device column cast is invalid".into()))
1189 }
1190
1191 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 if n <= 256 {
1211 return self.prefix_sum_mask_cpu(mask);
1212 }
1213
1214 self.prefix_sum_mask_gpu_multiblock(mask)
1216 }
1217
1218 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 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 let d_mask = self
1242 .htod_sync_copy_tracked(mask)
1243 .map_err(|e| XlogError::Kernel(format!("Failed to upload mask: {}", e)))?;
1244
1245 let d_prefix_sum = self.memory.alloc::<u32>(n)?;
1247
1248 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1250
1251 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 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 if num_blocks > 1 {
1277 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1278
1279 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 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 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 let count = prefix_sum[n - 1] + mask[n - 1] as u32;
1313
1314 Ok((prefix_sum, count))
1315 }
1316
1317 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 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 self.compact_buffer_by_mask(input, mask, &prefix_sum, count as u64)
1348 }
1349
1350 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}