1use std::ffi::c_void;
4use std::sync::atomic::Ordering;
5
6use crate::{
7 cuda_graph::{CapturedCudaGraph, CsmCudaGraphKey, CudaGraphNodeKind},
8 AsKernelParam, DeviceSlice, LaunchAsync, LaunchConfig,
9};
10use xlog_core::{Result, ScalarType, Schema, XlogError};
11
12use super::{
13 dedup_kernels, filter_kernels, ilp_kernels, join_kernels, pack_kernels, scan_kernels,
14 set_ops_kernels, sort_kernels, CsmCudaGraphEntry, CsmCudaGraphNodes, HashTableU64,
15 JoinHashTableV2, JoinIndexV2, JoinType, PackedKeyData, RadixSortScratch, DEDUP_MODULE,
16 DEFAULT_JOIN_MAX_OUTPUT, FILTER_MODULE, ILP_MODULE, JOIN_MODULE, NESTED_LOOP_TOTAL_THRESHOLD,
17 PACK_MODULE, SCAN_MODULE, SET_OPS_MODULE, SORT_MODULE,
18};
19use crate::device_runtime::{Access, BlockId, StreamId};
20use crate::launch::LaunchRecorder;
21use crate::memory::{CudaColumn, TrackedCudaSlice};
22use crate::CudaBuffer;
23
24const XLOG_TY_U32: u8 = 0;
29const XLOG_TY_U64: u8 = 1;
30const XLOG_TY_I32: u8 = 2;
31const XLOG_TY_I64: u8 = 3;
32const XLOG_TY_F32: u8 = 4;
33const XLOG_TY_F64: u8 = 5;
34const XLOG_TY_BOOL: u8 = 6;
35const XLOG_TY_SYMBOL: u8 = 7;
36const SMALL_FULL_ROW_SORT_MAX_ROWS: usize = 1024;
37const UNION_MANY_CHUNK_BYTES: usize = 1 << 30;
43
44#[inline]
45fn scalar_type_code_dedup(ty: ScalarType) -> u8 {
46 match ty {
47 ScalarType::U32 => XLOG_TY_U32,
48 ScalarType::U64 => XLOG_TY_U64,
49 ScalarType::I32 => XLOG_TY_I32,
50 ScalarType::I64 => XLOG_TY_I64,
51 ScalarType::F32 => XLOG_TY_F32,
52 ScalarType::F64 => XLOG_TY_F64,
53 ScalarType::Bool => XLOG_TY_BOOL,
54 ScalarType::Symbol => XLOG_TY_SYMBOL,
55 }
56}
57
58impl super::CudaKernelProvider {
59 pub fn hash_join(
77 &self,
78 left: &CudaBuffer,
79 right: &CudaBuffer,
80 left_keys: &[usize],
81 right_keys: &[usize],
82 ) -> Result<CudaBuffer> {
83 self.hash_join_with_limit(left, right, left_keys, right_keys, None)
84 }
85
86 pub fn hash_join_with_limit(
105 &self,
106 left: &CudaBuffer,
107 right: &CudaBuffer,
108 left_keys: &[usize],
109 right_keys: &[usize],
110 max_output: Option<usize>,
111 ) -> Result<CudaBuffer> {
112 let max_output_limit = max_output.unwrap_or(DEFAULT_JOIN_MAX_OUTPUT);
113
114 if left_keys.is_empty() || right_keys.is_empty() {
116 return Err(XlogError::Kernel(
117 "Join requires at least one key column".to_string(),
118 ));
119 }
120 if left_keys.len() != right_keys.len() {
121 return Err(XlogError::Kernel(
122 "Left and right key columns must have same length".to_string(),
123 ));
124 }
125 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
126 if left_idx >= left.arity() {
127 return Err(XlogError::Kernel(format!(
128 "Left key column index {} out of bounds (arity {})",
129 left_idx,
130 left.arity()
131 )));
132 }
133 if right_idx >= right.arity() {
134 return Err(XlogError::Kernel(format!(
135 "Right key column index {} out of bounds (arity {})",
136 right_idx,
137 right.arity()
138 )));
139 }
140 }
141
142 let right_key_set: std::collections::HashSet<usize> = right_keys.iter().copied().collect();
144 let mut result_columns_schema = left.schema().columns.clone();
145 let mut result_sort_labels = left.schema().sort_labels().to_vec();
146 for (idx, col) in right.schema().columns.iter().enumerate() {
147 if !right_key_set.contains(&idx) {
148 result_columns_schema.push(col.clone());
149 result_sort_labels.push(
150 right
151 .schema()
152 .column_sort_label(idx)
153 .unwrap_or(&col.0)
154 .to_string(),
155 );
156 }
157 }
158 let result_schema = Schema::new(result_columns_schema)
159 .with_sort_labels(result_sort_labels)
160 .expect("natural join sort labels match result schema arity");
161
162 if left.is_empty() || right.is_empty() {
164 return self.create_empty_buffer(result_schema);
165 }
166
167 let combined = self.hash_join_v2_with_limit(
169 left,
170 right,
171 left_keys,
172 right_keys,
173 JoinType::Inner,
174 Some(max_output_limit),
175 )?;
176
177 if combined.is_empty() {
178 return self.create_empty_buffer(result_schema);
179 }
180
181 let left_arity = left.arity();
182 let right_arity = right.arity();
183
184 let CudaBuffer {
185 columns: combined_columns,
186 row_cap,
187 d_num_rows,
188 schema: _,
189 ..
190 } = combined;
191
192 if combined_columns.len() != left_arity + right_arity {
193 return Err(XlogError::Kernel(format!(
194 "Join internal error: expected {} columns, got {}",
195 left_arity + right_arity,
196 combined_columns.len()
197 )));
198 }
199
200 let mut output_columns = Vec::with_capacity(result_schema.arity());
201 let mut it = combined_columns.into_iter();
202
203 for _ in 0..left_arity {
205 let col = it.next().ok_or_else(|| {
206 XlogError::Kernel("Join internal error: missing left columns".to_string())
207 })?;
208 output_columns.push(col);
209 }
210
211 for (right_col_idx, col) in it.enumerate() {
213 if !right_key_set.contains(&right_col_idx) {
214 output_columns.push(col);
215 }
216 }
217
218 Ok(CudaBuffer::from_columns(
219 output_columns,
220 row_cap,
221 d_num_rows,
222 result_schema,
223 ))
224 }
225 pub fn dedup(&self, input: &CudaBuffer, key_cols: &[usize]) -> Result<CudaBuffer> {
239 if !input.canonical_full_row_set_certified() {
240 self.validated_logical_row_count(input)?;
241 }
242 if input.is_empty() {
243 return self.create_empty_buffer(input.schema().clone());
244 }
245
246 if key_cols.is_empty() {
247 if input.arity() == 0 {
248 let rows = self.device_row_count(input)?;
251 if rows == 0 {
252 return self.create_empty_buffer(input.schema().clone());
253 }
254 let mut result = self.buffer_from_columns(Vec::new(), 1, input.schema().clone())?;
255 result.certify_canonical_full_row_set();
256 return Ok(result);
257 }
258 return Err(XlogError::Kernel(
259 "Dedup requires at least one key column".to_string(),
260 ));
261 }
262
263 if Self::is_full_row_key(key_cols, input.arity()) && input.arity() > 1 {
264 let mut result = self.dedup_full_row_deterministic(input)?;
265 result.certify_canonical_full_row_set();
266 return Ok(result);
267 }
268
269 let sorted = self.sort(input, key_cols)?;
270 let mut result = self.dedup_sorted(&sorted, key_cols)?;
271 if Self::is_full_row_key(key_cols, input.arity()) {
272 result.certify_canonical_full_row_set();
273 }
274 Ok(result)
275 }
276
277 pub fn dedup_sorted(&self, input: &CudaBuffer, key_cols: &[usize]) -> Result<CudaBuffer> {
289 if !input.canonical_full_row_set_certified() {
290 self.validated_logical_row_count(input)?;
291 }
292 if input.is_empty() {
293 return self.create_empty_buffer(input.schema().clone());
294 }
295
296 if key_cols.is_empty() {
297 if input.arity() == 0 {
298 let rows = self.device_row_count(input)?;
299 if rows == 0 {
300 return self.create_empty_buffer(input.schema().clone());
301 }
302 return self.buffer_from_columns(Vec::new(), 1, input.schema().clone());
303 }
304 return Err(XlogError::Kernel(
305 "Dedup requires at least one key column".to_string(),
306 ));
307 }
308
309 if Self::is_full_row_key(key_cols, input.arity()) && input.arity() > 1 {
310 return self.dedup_full_row_deterministic(input);
311 }
312
313 if input.num_rows() <= 1 {
314 return self.clone_buffer(input);
315 }
316
317 if input.num_rows() > u32::MAX as u64 {
318 return Err(XlogError::Kernel(format!(
319 "Dedup supports at most {} rows, got {}",
320 u32::MAX,
321 input.num_rows()
322 )));
323 }
324
325 let scalar_type_code = scalar_type_code_dedup;
329
330 let device = self.device.inner();
331 let num_rows = input.num_rows() as u32;
332
333 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(key_cols.len());
334 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(key_cols.len());
335 let mut col_types_host: Vec<u8> = Vec::with_capacity(key_cols.len());
336
337 for &key_col in key_cols {
338 if key_col >= input.arity() {
339 return Err(XlogError::Kernel(format!(
340 "Key column {} out of bounds (arity {})",
341 key_col,
342 input.arity()
343 )));
344 }
345
346 let col = input
347 .column(key_col)
348 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", key_col)))?;
349 let ty = input.schema().column_type(key_col).ok_or_else(|| {
350 XlogError::Kernel(format!("Key column {} type not found in schema", key_col))
351 })?;
352
353 let elem_size = ty.size_bytes();
354 let expected_bytes = (num_rows as usize) * elem_size;
355 if col.num_bytes() != expected_bytes {
356 return Err(XlogError::Kernel(format!(
357 "Key column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
358 key_col,
359 col.num_bytes(),
360 expected_bytes,
361 num_rows,
362 elem_size
363 )));
364 }
365
366 let ptr = *col.device_ptr();
367 col_ptrs_host.push(ptr);
368 col_sizes_host.push(elem_size as u32);
369 col_types_host.push(scalar_type_code(ty));
370 }
371
372 let num_key_cols = key_cols.len() as u32;
373 let mut d_col_ptrs = self.memory.alloc::<u64>(key_cols.len())?;
374 let mut d_col_sizes = self.memory.alloc::<u32>(key_cols.len())?;
375 let mut d_col_types = self.memory.alloc::<u8>(key_cols.len())?;
376
377 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
378 .map_err(|e| XlogError::Kernel(format!("Failed to upload key column ptrs: {}", e)))?;
379 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
380 .map_err(|e| XlogError::Kernel(format!("Failed to upload key column sizes: {}", e)))?;
381 self.htod_launch_metadata_sync_copy_into(&col_types_host, &mut d_col_types)
382 .map_err(|e| XlogError::Kernel(format!("Failed to upload key column types: {}", e)))?;
383
384 let block_size = 256u32;
385 let num_blocks = num_rows.div_ceil(block_size);
386 let config = LaunchConfig {
387 grid_dim: (num_blocks, 1, 1),
388 block_dim: (block_size, 1, 1),
389 shared_mem_bytes: 0,
390 };
391
392 let d_unique_mask = self.memory.alloc::<u8>(num_rows as usize)?;
393 let d_prefix_sum = self.memory.alloc::<u32>(num_rows as usize)?;
394 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
395
396 let mark_and_scan_fn = device
397 .get_func(DEDUP_MODULE, dedup_kernels::MARK_UNIQUE_AND_SCAN_COLUMNAR)
398 .ok_or_else(|| {
399 XlogError::Kernel("mark_unique_and_scan_columnar kernel not found".to_string())
400 })?;
401
402 unsafe {
407 mark_and_scan_fn.clone().launch(
408 config,
409 (
410 &d_col_ptrs,
411 &d_col_sizes,
412 &d_col_types,
413 num_key_cols,
414 input.num_rows_device(),
415 num_rows,
416 &d_unique_mask,
417 &d_prefix_sum,
418 &d_block_sums,
419 ),
420 )
421 }
422 .map_err(|e| XlogError::Kernel(format!("mark_unique_and_scan_columnar failed: {}", e)))?;
423 self.device.synchronize()?;
424
425 if num_blocks > 1 {
426 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
427
428 let phase3_fn = device
429 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
430 .ok_or_else(|| {
431 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
432 })?;
433
434 unsafe {
436 phase3_fn.clone().launch(
437 LaunchConfig {
438 grid_dim: (num_blocks, 1, 1),
439 block_dim: (block_size, 1, 1),
440 shared_mem_bytes: 0,
441 },
442 (&d_prefix_sum, &d_block_sums, num_rows),
443 )
444 }
445 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
446 self.device.synchronize()?;
447 }
448
449 self.device.synchronize()?;
450
451 let d_out_count = self.capture_compact_count(&d_prefix_sum, &d_unique_mask, num_rows)?;
452 self.compact_buffer_by_device_mask_device_count(
453 input,
454 &d_unique_mask,
455 &d_prefix_sum,
456 d_out_count,
457 )
458 }
459 pub fn union(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
471 self.union_gpu(a, b)
472 }
473
474 fn concat_buffers_gpu(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
475 if !self.schemas_type_compatible(a.schema(), b.schema()) {
476 return Err(XlogError::Kernel(format!(
477 "Concat requires compatible schemas: {:?} vs {:?}",
478 a.schema(),
479 b.schema()
480 )));
481 }
482
483 let schema = a.schema().clone();
484 let a_rows = self.device_row_count(a)? as u64;
485 let b_rows = self.device_row_count(b)? as u64;
486
487 if a_rows == 0 && b_rows == 0 {
488 return self.create_empty_buffer(schema);
489 }
490 if a_rows == 0 {
491 return self.clone_buffer(b);
492 }
493 if b_rows == 0 {
494 return self.clone_buffer(a);
495 }
496
497 let total_rows = a_rows + b_rows;
498 if total_rows > u32::MAX as u64 {
499 return Err(XlogError::Kernel(format!(
500 "Concat supports at most {} rows, got {}",
501 u32::MAX,
502 total_rows
503 )));
504 }
505
506 let device = self.device.inner();
507 let concat_fn = device
508 .get_func(SET_OPS_MODULE, set_ops_kernels::CONCAT_BYTES)
509 .ok_or_else(|| XlogError::Kernel("concat_bytes kernel not found".to_string()))?;
510
511 let block_size = 256u32;
512
513 let a_rows = usize::try_from(a_rows)
514 .map_err(|_| XlogError::Kernel(format!("Concat: a has too many rows: {}", a_rows)))?;
515 let b_rows = usize::try_from(b_rows)
516 .map_err(|_| XlogError::Kernel(format!("Concat: b has too many rows: {}", b_rows)))?;
517
518 let mut result_columns = Vec::with_capacity(schema.arity());
519 for col_idx in 0..schema.arity() {
520 let elem_size = schema
521 .column_type(col_idx)
522 .map(|t| t.size_bytes())
523 .unwrap_or(4);
524
525 let a_bytes = a_rows
526 .checked_mul(elem_size)
527 .ok_or_else(|| XlogError::Kernel("Concat: a_bytes overflow".to_string()))?;
528 let b_bytes = b_rows
529 .checked_mul(elem_size)
530 .ok_or_else(|| XlogError::Kernel("Concat: b_bytes overflow".to_string()))?;
531 let total_bytes = a_bytes
532 .checked_add(b_bytes)
533 .ok_or_else(|| XlogError::Kernel("Concat: total_bytes overflow".to_string()))?;
534
535 let a_bytes_u32 = u32::try_from(a_bytes).map_err(|_| {
536 XlogError::Kernel(format!("Concat: a_bytes too large: {}", a_bytes))
537 })?;
538 let b_bytes_u32 = u32::try_from(b_bytes).map_err(|_| {
539 XlogError::Kernel(format!("Concat: b_bytes too large: {}", b_bytes))
540 })?;
541 let total_bytes_u32 = u32::try_from(total_bytes).map_err(|_| {
542 XlogError::Kernel(format!("Concat: total_bytes too large: {}", total_bytes))
543 })?;
544
545 let a_col = a
546 .column(col_idx)
547 .ok_or_else(|| XlogError::Kernel(format!("A column {} not found", col_idx)))?;
548 let b_col = b
549 .column(col_idx)
550 .ok_or_else(|| XlogError::Kernel(format!("B column {} not found", col_idx)))?;
551
552 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
553
554 if total_bytes_u32 > 0 {
555 let grid_size = total_bytes_u32.div_ceil(block_size);
556 let config = LaunchConfig {
557 grid_dim: (grid_size, 1, 1),
558 block_dim: (block_size, 1, 1),
559 shared_mem_bytes: 0,
560 };
561
562 unsafe {
564 concat_fn.clone().launch(
565 config,
566 (a_col, a_bytes_u32, b_col, b_bytes_u32, &mut out_col),
567 )
568 }
569 .map_err(|e| XlogError::Kernel(format!("concat_bytes failed: {}", e)))?;
570 }
571
572 result_columns.push(out_col.into());
573 }
574
575 self.device.synchronize()?;
576
577 self.buffer_from_columns(result_columns, total_rows, schema)
578 }
579 pub fn diff(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
594 let num_a = self.device_row_count(a)?;
595 let num_b = self.device_row_count(b)?;
596 if num_a > u32::MAX as usize || num_b > u32::MAX as usize {
597 return Err(XlogError::Kernel(format!(
598 "Diff supports at most {} rows per side (a={}, b={})",
599 u32::MAX,
600 num_a,
601 num_b
602 )));
603 }
604
605 if num_a == 0 {
607 return self.create_empty_buffer(a.schema().clone());
608 }
609 if num_b == 0 {
610 return self.clone_buffer(a);
611 }
612
613 if !self.schemas_type_compatible(a.schema(), b.schema()) {
615 return Err(XlogError::Kernel(format!(
616 "Diff requires compatible schemas: {:?} vs {:?}",
617 a.schema(),
618 b.schema()
619 )));
620 }
621
622 if a.arity() == 0 {
624 return Err(XlogError::Kernel(
625 "Diff requires at least one column".to_string(),
626 ));
627 }
628
629 let num_b = num_b as u32;
630 let num_a = num_a as u32;
631
632 let hash_table_size = (num_b as usize * 2).max(1024) as u32;
634 let hash_table_alloc_size = (hash_table_size * 3) as usize;
635 let mut hash_table = self.memory.alloc::<u32>(hash_table_alloc_size)?;
636 let mut next_ptrs = self.memory.alloc::<u32>(num_b as usize)?;
637
638 let init_val = 0xFFFFFFFFu32;
640 self.device
641 .inner()
642 .htod_sync_copy_into(&vec![init_val; hash_table_alloc_size], &mut hash_table)
643 .map_err(|e| XlogError::Kernel(format!("Failed to init hash table: {}", e)))?;
644 self.device
645 .inner()
646 .htod_sync_copy_into(&vec![init_val; num_b as usize], &mut next_ptrs)
647 .map_err(|e| XlogError::Kernel(format!("Failed to init next pointers: {}", e)))?;
648
649 let build_func = self
651 .device
652 .inner()
653 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUILD)
654 .ok_or_else(|| XlogError::Kernel("hash_join_build kernel not found".to_string()))?;
655
656 let b_key_col = b
657 .column(0)
658 .ok_or_else(|| XlogError::Kernel("B key column not found".to_string()))?;
659 let b_keys_view = self.column_as_u32_view(b_key_col, num_b as usize)?;
660
661 let block_size = 256u32;
662 let build_grid = num_b.div_ceil(block_size);
663 let build_config = LaunchConfig {
664 grid_dim: (build_grid, 1, 1),
665 block_dim: (block_size, 1, 1),
666 shared_mem_bytes: 0,
667 };
668
669 unsafe {
671 build_func
672 .clone()
673 .launch(
674 build_config,
675 (
676 &b_keys_view,
677 &b_keys_view, num_b,
679 &hash_table,
680 &next_ptrs,
681 hash_table_size,
682 ),
683 )
684 .map_err(|e| XlogError::Kernel(format!("Build kernel failed: {}", e)))?;
685 }
686
687 self.device.synchronize()?;
689
690 let a_key_col = a
692 .column(0)
693 .ok_or_else(|| XlogError::Kernel("A key column not found".to_string()))?;
694
695 let mut a_keys_host = vec![0u8; (num_a as usize) * 4];
697 self.dtoh_sync_copy_into_tracked(a_key_col, &mut a_keys_host)
698 .map_err(|e| XlogError::Kernel(format!("Failed to read a keys: {}", e)))?;
699
700 let mut b_keys_host = vec![0u8; (num_b as usize) * 4];
701 self.dtoh_sync_copy_into_tracked(b_key_col, &mut b_keys_host)
702 .map_err(|e| XlogError::Kernel(format!("Failed to read b keys: {}", e)))?;
703
704 let b_keys_set: std::collections::HashSet<u32> = b_keys_host
706 .chunks_exact(4)
707 .map(|chunk| u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
708 .collect();
709
710 let diff_indices: Vec<usize> = a_keys_host
712 .chunks_exact(4)
713 .enumerate()
714 .map(|(i, chunk)| {
715 (
716 i,
717 u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]),
718 )
719 })
720 .filter(|(_, k)| !b_keys_set.contains(k))
721 .map(|(i, _)| i)
722 .collect();
723
724 let diff_count = diff_indices.len() as u64;
725
726 if diff_count == 0 {
727 return self.create_empty_buffer(a.schema().clone());
728 }
729
730 let schema = a.schema().clone();
732 let mut result_columns = Vec::with_capacity(schema.arity());
733
734 for col_idx in 0..schema.arity() {
735 let col_type_size = schema
736 .column_type(col_idx)
737 .map(|t| t.size_bytes())
738 .unwrap_or(4);
739 let result_bytes = (diff_count as usize) * col_type_size;
740
741 if let Some(a_col) = a.column(col_idx) {
742 let a_col_bytes = (num_a as usize) * col_type_size;
744 let mut a_col_host = vec![0u8; a_col_bytes];
745 self.dtoh_sync_copy_into_tracked(a_col, &mut a_col_host)
746 .map_err(|e| XlogError::Kernel(format!("Failed to read column: {}", e)))?;
747
748 let mut result_host = Vec::with_capacity(result_bytes);
750 for &idx in &diff_indices {
751 let start = idx * col_type_size;
752 let end = start + col_type_size;
753 result_host.extend_from_slice(&a_col_host[start..end]);
754 }
755
756 let mut result_col = self.memory.alloc::<u8>(result_bytes)?;
758 self.device
759 .inner()
760 .htod_sync_copy_into(&result_host, &mut result_col)
761 .map_err(|e| XlogError::Kernel(format!("Failed to upload result: {}", e)))?;
762
763 result_columns.push(result_col.into());
764 }
765 }
766
767 self.buffer_from_columns(result_columns, diff_count, schema)
768 }
769 fn ensure_column_bytes_kernel_indexable(&self, input: &CudaBuffer) -> Result<()> {
774 let rows = self.device_row_count(input)?;
775 for col_idx in 0..input.arity() {
776 let elem_size = input
777 .schema()
778 .column_type(col_idx)
779 .map(|t| t.size_bytes())
780 .unwrap_or(4);
781 let col_bytes = rows
782 .checked_mul(elem_size)
783 .ok_or_else(|| XlogError::Kernel("Sort: column byte size overflow".to_string()))?;
784 if u32::try_from(col_bytes).is_err() {
785 return Err(XlogError::Kernel(format!(
786 "Sort supports at most {} bytes per column, got {} (column {})",
787 u32::MAX,
788 col_bytes,
789 col_idx
790 )));
791 }
792 }
793 Ok(())
794 }
795
796 fn union_many_chunk_bytes() -> usize {
800 std::env::var("XLOG_UNION_CHUNK_BYTES")
801 .ok()
802 .and_then(|value| value.parse::<usize>().ok())
803 .filter(|&value| value > 0)
804 .unwrap_or(UNION_MANY_CHUNK_BYTES)
805 }
806
807 pub fn union_gpu(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
823 self.union_many_gpu(&[a, b])
824 }
825
826 pub fn union_many_gpu(&self, inputs: &[&CudaBuffer]) -> Result<CudaBuffer> {
860 let first = inputs.first().ok_or_else(|| {
861 XlogError::Kernel("union_many_gpu requires at least one input".to_string())
862 })?;
863 for other in &inputs[1..] {
865 if !self.schemas_type_compatible(first.schema(), other.schema()) {
866 return Err(XlogError::Kernel(format!(
867 "Union requires compatible schemas: {:?} vs {:?}",
868 first.schema(),
869 other.schema()
870 )));
871 }
872 }
873 for input in inputs {
874 if !input.canonical_full_row_set_certified() {
875 self.validated_logical_row_count(input)?;
876 }
877 }
878
879 let schema = first.schema().clone();
880
881 let possible_non_empty: Vec<&CudaBuffer> = inputs
882 .iter()
883 .copied()
884 .filter(|input| input.cached_row_count() != Some(0))
885 .collect();
886 if possible_non_empty.is_empty() {
887 return self.create_empty_buffer(schema);
888 }
889 if schema.arity() > 0
893 && possible_non_empty.len() == 1
894 && possible_non_empty[0].canonical_full_row_set_certified()
895 {
896 return self.clone_buffer(possible_non_empty[0]);
897 }
898
899 let mut non_empty: Vec<(&CudaBuffer, usize)> = Vec::with_capacity(possible_non_empty.len());
900 for input in possible_non_empty {
901 let rows = match input.cached_row_count() {
902 Some(rows) => rows as usize,
903 None => self.device_row_count(input)?,
904 };
905 if rows > 0 {
906 non_empty.push((input, rows));
907 }
908 }
909
910 if schema.arity() == 0 {
911 if non_empty.is_empty() {
913 return self.create_empty_buffer(schema);
914 }
915 return self.buffer_from_columns(Vec::new(), 1, schema);
916 }
917
918 if non_empty.is_empty() {
919 return self.create_empty_buffer(schema);
920 }
921
922 let key_cols: Vec<usize> = (0..schema.arity()).collect();
924 if non_empty.len() == 1 {
925 let input = non_empty[0].0;
926 if input.canonical_full_row_set_certified() {
927 return self.clone_buffer(input);
928 }
929 let mut result = self.dedup(input, &key_cols)?;
930 result.certify_canonical_full_row_set();
931 return Ok(result);
932 }
933
934 let row_bytes: usize = (0..schema.arity())
935 .map(|c| schema.column_type(c).map(|t| t.size_bytes()).unwrap_or(4))
936 .sum::<usize>()
937 .max(1);
938 let budget_rows = (Self::union_many_chunk_bytes() / row_bytes).max(1);
939
940 let mut acc: Option<CudaBuffer> = None;
945 let mut idx = 0usize;
946 while idx < non_empty.len() {
947 let mut chunk: Vec<&CudaBuffer> = Vec::new();
948 let mut chunk_rows = 0usize;
949 if let Some(acc_buf) = acc.as_ref() {
950 chunk_rows = self.device_row_count(acc_buf)?;
951 chunk.push(acc_buf);
952 }
953 let mut taken = 0usize;
956 while idx < non_empty.len() {
957 let (input, rows) = non_empty[idx];
958 if taken > 0 && chunk_rows.saturating_add(rows) > budget_rows {
959 break;
960 }
961 chunk.push(input);
962 chunk_rows = chunk_rows.saturating_add(rows);
963 idx += 1;
964 taken += 1;
965 }
966 acc = Some(self.union_chunk_gpu(&chunk, chunk_rows, &key_cols)?);
967 }
968 let mut result = acc.expect("at least one non-empty input was folded");
969 result.certify_canonical_full_row_set();
970 Ok(result)
971 }
972
973 fn union_chunk_gpu(
977 &self,
978 inputs: &[&CudaBuffer],
979 chunk_rows: usize,
980 key_cols: &[usize],
981 ) -> Result<CudaBuffer> {
982 if inputs.len() == 1 {
983 return self.dedup(inputs[0], key_cols);
984 }
985 let concat = if inputs.len() == 2 {
986 self.concat_buffers_gpu(inputs[0], inputs[1])?
987 } else {
988 self.concat_many_buffers_gpu(inputs, chunk_rows)?
989 };
990 if inputs[0].schema().arity() > 1 {
991 return self.dedup_full_row_deterministic(&concat);
996 }
997 let sorted = self.sort(&concat, key_cols)?;
998 self.dedup_sorted(&sorted, key_cols)
999 }
1000
1001 fn concat_many_buffers_gpu(
1018 &self,
1019 inputs: &[&CudaBuffer],
1020 total_rows: usize,
1021 ) -> Result<CudaBuffer> {
1022 let schema = inputs[0].schema().clone();
1023 if total_rows > u32::MAX as usize {
1024 return Err(XlogError::Kernel(format!(
1025 "Concat supports at most {} rows, got {}",
1026 u32::MAX,
1027 total_rows
1028 )));
1029 }
1030
1031 let mut input_rows = Vec::with_capacity(inputs.len());
1032 for input in inputs {
1033 input_rows.push(self.device_row_count(input)?);
1034 }
1035
1036 let device = self.device.inner();
1037 let mut result_columns = Vec::with_capacity(schema.arity());
1038 let enqueue_result = (|| -> Result<()> {
1039 for col_idx in 0..schema.arity() {
1040 let elem_size = schema
1041 .column_type(col_idx)
1042 .map(|t| t.size_bytes())
1043 .unwrap_or(4);
1044 let total_bytes = total_rows
1045 .checked_mul(elem_size)
1046 .ok_or_else(|| XlogError::Kernel("Concat: total_bytes overflow".to_string()))?;
1047 u32::try_from(total_bytes).map_err(|_| {
1048 XlogError::Kernel(format!("Concat: total_bytes too large: {}", total_bytes))
1049 })?;
1050
1051 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
1052 let mut offset = 0usize;
1053 for (input, &rows) in inputs.iter().zip(&input_rows) {
1054 let col_bytes = rows.checked_mul(elem_size).ok_or_else(|| {
1055 XlogError::Kernel("Concat: col_bytes overflow".to_string())
1056 })?;
1057 u32::try_from(col_bytes).map_err(|_| {
1058 XlogError::Kernel(format!("Concat: col_bytes too large: {}", col_bytes))
1059 })?;
1060 let col = input.column(col_idx).ok_or_else(|| {
1061 XlogError::Kernel(format!("Concat: column {} not found", col_idx))
1062 })?;
1063 let src = self.column_bytes_view(col, col_bytes)?;
1064 let mut dst = out_col.slice_mut(offset..offset + col_bytes);
1065 device.dtod_copy_async(&src, &mut dst).map_err(|e| {
1066 XlogError::Kernel(format!("Concat: failed to copy column: {}", e))
1067 })?;
1068 offset += col_bytes;
1069 }
1070
1071 result_columns.push(out_col.into());
1072 }
1073 Ok(())
1074 })();
1075
1076 self.device.synchronize()?;
1080 enqueue_result?;
1081
1082 self.buffer_from_columns(result_columns, total_rows as u64, schema)
1083 }
1084
1085 pub fn diff_gpu(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1103 let num_a = self.device_row_count(a)?;
1104 let num_b = self.device_row_count(b)?;
1105 if num_a > u32::MAX as usize || num_b > u32::MAX as usize {
1106 return Err(XlogError::Kernel(format!(
1107 "Diff supports at most {} rows per side (a={}, b={})",
1108 u32::MAX,
1109 num_a,
1110 num_b
1111 )));
1112 }
1113
1114 if num_a == 0 {
1115 return self.create_empty_buffer(a.schema().clone());
1116 }
1117
1118 if !self.schemas_type_compatible(a.schema(), b.schema()) {
1120 return Err(XlogError::Kernel(format!(
1121 "Diff requires compatible schemas: {:?} vs {:?}",
1122 a.schema(),
1123 b.schema()
1124 )));
1125 }
1126
1127 if a.arity() == 0 {
1128 if num_b == 0 {
1130 return self.buffer_from_columns(Vec::new(), 1, a.schema().clone());
1131 }
1132 return self.create_empty_buffer(a.schema().clone());
1133 }
1134
1135 let col_type = a
1136 .schema()
1137 .column_type(0)
1138 .ok_or_else(|| XlogError::Kernel("No columns".to_string()))?;
1139
1140 if a.arity() == 1 && matches!(col_type, ScalarType::U32) && num_b != 0 {
1143 return self.diff_gpu_u32(a, b);
1144 }
1145
1146 self.diff_via_deterministic_set(a, b)
1147 }
1148
1149 fn diff_gpu_u32(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1151 if a.arity() != 1 {
1153 return self.diff_via_deterministic_set(a, b);
1154 }
1155
1156 let sorted_a = self.sort(a, &[0])?;
1158 let deduped_a = self.dedup_sorted(&sorted_a, &[0])?;
1159
1160 let sorted_b = self.sort(b, &[0])?;
1161 let deduped_b = self.dedup_sorted(&sorted_b, &[0])?;
1162
1163 let num_a = self.device_row_count(&deduped_a)?;
1164 let num_b = self.device_row_count(&deduped_b)?;
1165 if num_a > u32::MAX as usize || num_b > u32::MAX as usize {
1166 return Err(XlogError::Kernel(format!(
1167 "Diff supports at most {} rows per side (a={}, b={})",
1168 u32::MAX,
1169 num_a,
1170 num_b
1171 )));
1172 }
1173
1174 if num_a == 0 {
1175 return self.create_empty_buffer(a.schema().clone());
1176 }
1177
1178 let num_a = num_a as u32;
1179 let num_b = num_b as u32;
1180
1181 let diff_mark_fn = self
1183 .device
1184 .inner()
1185 .get_func(SET_OPS_MODULE, set_ops_kernels::SORTED_DIFF_MARK)
1186 .ok_or_else(|| XlogError::Kernel("sorted_diff_mark kernel not found".to_string()))?;
1187
1188 let a_col = deduped_a
1190 .column(0)
1191 .ok_or_else(|| XlogError::Kernel("A column 0 not found".to_string()))?;
1192 let b_col = deduped_b
1193 .column(0)
1194 .ok_or_else(|| XlogError::Kernel("B column 0 not found".to_string()))?;
1195
1196 let a_view = self.column_as_u32_view(a_col, num_a as usize)?;
1197 let b_view = self.column_as_u32_view(b_col, num_b as usize)?;
1198
1199 let diff_mask = self.memory.alloc::<u8>(num_a as usize)?;
1201
1202 let block_size = 256u32;
1204 let grid_size = num_a.div_ceil(block_size);
1205 let config = LaunchConfig {
1206 grid_dim: (grid_size, 1, 1),
1207 block_dim: (block_size, 1, 1),
1208 shared_mem_bytes: 0,
1209 };
1210
1211 unsafe {
1215 diff_mark_fn.clone().launch(
1216 config,
1217 (
1218 &a_view,
1219 deduped_a.num_rows_device(),
1220 num_a,
1221 &b_view,
1222 deduped_b.num_rows_device(),
1223 num_b,
1224 &diff_mask,
1225 ),
1226 )
1227 }
1228 .map_err(|e| XlogError::Kernel(format!("sorted_diff_mark failed: {}", e)))?;
1229
1230 let device = self.device.inner();
1232 let num_blocks = grid_size;
1233 let d_prefix_sum = self.memory.alloc::<u32>(num_a as usize)?;
1234 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1235
1236 let phase1_fn = device
1237 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1238 .ok_or_else(|| {
1239 XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1240 })?;
1241
1242 unsafe {
1244 phase1_fn.clone().launch(
1245 LaunchConfig {
1246 grid_dim: (num_blocks, 1, 1),
1247 block_dim: (block_size, 1, 1),
1248 shared_mem_bytes: 0,
1249 },
1250 (&diff_mask, &d_prefix_sum, &d_block_sums, num_a),
1251 )
1252 }
1253 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
1254
1255 if num_blocks > 1 {
1256 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1257
1258 let phase3_fn = device
1259 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1260 .ok_or_else(|| {
1261 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1262 })?;
1263
1264 unsafe {
1266 phase3_fn.clone().launch(
1267 LaunchConfig {
1268 grid_dim: (num_blocks, 1, 1),
1269 block_dim: (block_size, 1, 1),
1270 shared_mem_bytes: 0,
1271 },
1272 (&d_prefix_sum, &d_block_sums, num_a),
1273 )
1274 }
1275 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
1276 }
1277
1278 self.device.synchronize()?;
1279
1280 let d_out_count = self.capture_compact_count(&d_prefix_sum, &diff_mask, num_a)?;
1281 self.compact_buffer_by_device_mask_device_count(
1282 &deduped_a,
1283 &diff_mask,
1284 &d_prefix_sum,
1285 d_out_count,
1286 )
1287 }
1288
1289 fn diff_via_deterministic_set(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1312 let deduped_a = self.dedup_full_row_deterministic(a)?;
1320 let deduped_b = self.dedup_full_row_deterministic(b)?;
1321
1322 let a_rows = self.device_row_count(&deduped_a)? as u32;
1323 let b_rows = self.device_row_count(&deduped_b)? as u32;
1324 if a_rows == 0 {
1325 return self.create_empty_buffer(a.schema().clone());
1326 }
1327 if b_rows == 0 {
1328 return Ok(deduped_a);
1329 }
1330 let arity = deduped_a.arity();
1331 if arity == 0 {
1332 return self.create_empty_buffer(a.schema().clone());
1334 }
1335
1336 let sorted_b = deduped_b;
1342
1343 let schema = deduped_a.schema().clone();
1347 let device = self.device.inner();
1348
1349 let mut a_col_ptrs: Vec<u64> = Vec::with_capacity(arity);
1350 let mut b_col_ptrs: Vec<u64> = Vec::with_capacity(arity);
1351 let mut col_sizes: Vec<u32> = Vec::with_capacity(arity);
1352 let mut col_types: Vec<u8> = Vec::with_capacity(arity);
1353 for col_idx in 0..arity {
1354 let a_col = deduped_a.column(col_idx).ok_or_else(|| {
1355 XlogError::Kernel(format!("diff_full_row: a column {} missing", col_idx))
1356 })?;
1357 let b_col = sorted_b.column(col_idx).ok_or_else(|| {
1358 XlogError::Kernel(format!("diff_full_row: b column {} missing", col_idx))
1359 })?;
1360 let ty = schema.column_type(col_idx).ok_or_else(|| {
1361 XlogError::Kernel(format!("diff_full_row: column {} type missing", col_idx))
1362 })?;
1363 a_col_ptrs.push(*a_col.device_ptr());
1364 b_col_ptrs.push(*b_col.device_ptr());
1365 col_sizes.push(ty.size_bytes() as u32);
1366 col_types.push(scalar_type_code_dedup(ty));
1367 }
1368
1369 let mut d_a_ptrs = self.memory.alloc::<u64>(arity)?;
1370 let mut d_b_ptrs = self.memory.alloc::<u64>(arity)?;
1371 let mut d_sizes = self.memory.alloc::<u32>(arity)?;
1372 let mut d_types = self.memory.alloc::<u8>(arity)?;
1373 self.htod_launch_metadata_sync_copy_into(&a_col_ptrs, &mut d_a_ptrs)
1374 .map_err(|e| XlogError::Kernel(format!("diff_full_row a ptr upload: {}", e)))?;
1375 self.htod_launch_metadata_sync_copy_into(&b_col_ptrs, &mut d_b_ptrs)
1376 .map_err(|e| XlogError::Kernel(format!("diff_full_row b ptr upload: {}", e)))?;
1377 self.htod_launch_metadata_sync_copy_into(&col_sizes, &mut d_sizes)
1378 .map_err(|e| XlogError::Kernel(format!("diff_full_row size upload: {}", e)))?;
1379 self.htod_launch_metadata_sync_copy_into(&col_types, &mut d_types)
1380 .map_err(|e| XlogError::Kernel(format!("diff_full_row type upload: {}", e)))?;
1381
1382 let block_size = 256u32;
1383 let grid = a_rows.div_ceil(block_size);
1384 let cfg = LaunchConfig {
1385 grid_dim: (grid, 1, 1),
1386 block_dim: (block_size, 1, 1),
1387 shared_mem_bytes: 0,
1388 };
1389
1390 let d_keep_mask = self.memory.alloc::<u8>(a_rows as usize)?;
1391 let diff_fn = device
1392 .get_func(DEDUP_MODULE, dedup_kernels::MARK_DIFF_FULL_ROW_TYPED_SORTED)
1393 .ok_or_else(|| {
1394 XlogError::Kernel("mark_diff_full_row_typed_sorted kernel not found".to_string())
1395 })?;
1396 unsafe {
1400 diff_fn.clone().launch(
1401 cfg,
1402 (
1403 &d_a_ptrs,
1404 &d_b_ptrs,
1405 &d_sizes,
1406 &d_types,
1407 arity as u32,
1408 deduped_a.num_rows_device(),
1409 b_rows,
1410 a_rows,
1411 &d_keep_mask,
1412 ),
1413 )
1414 }
1415 .map_err(|e| XlogError::Kernel(format!("mark_diff_full_row_typed_sorted launch: {}", e)))?;
1416 self.device.synchronize()?;
1417
1418 let (d_prefix_sum, d_out_count) =
1420 self.scan_mask_to_prefix_with_count(&d_keep_mask, a_rows)?;
1421
1422 self.compact_buffer_by_device_mask_device_count(
1423 &deduped_a,
1424 &d_keep_mask,
1425 &d_prefix_sum,
1426 d_out_count,
1427 )
1428 }
1429
1430 pub fn dedup_full_row(&self, input: &CudaBuffer) -> Result<CudaBuffer> {
1459 if !input.canonical_full_row_set_certified() {
1460 self.validated_logical_row_count(input)?;
1461 }
1462 if Self::use_recorded_dedup_env() && input.num_rows() > 1 && input.arity() > 0 {
1466 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
1467 let recorded_compatible = (0..input.arity()).all(|c| {
1468 matches!(
1469 input.schema.column_type(c),
1470 Some(ScalarType::U32) | Some(ScalarType::Symbol)
1471 )
1472 });
1473 if recorded_compatible {
1474 let mut result = self.dedup_full_row_recorded(input, launch_stream)?;
1475 result.certify_canonical_full_row_set();
1476 return Ok(result);
1477 }
1478 }
1479 }
1480 let mut result = self.dedup_full_row_deterministic(input)?;
1481 result.certify_canonical_full_row_set();
1482 Ok(result)
1483 }
1484
1485 pub fn diff_full_row(&self, a: &CudaBuffer, b: &CudaBuffer) -> Result<CudaBuffer> {
1490 self.diff_gpu(a, b)
1494 }
1495
1496 fn read_join_output_count_metadata(&self, d_count: &TrackedCudaSlice<u32>) -> Result<u32> {
1518 self.dtoh_scalar_untracked::<u32>(d_count, 0)
1525 .map_err(|e| match e {
1526 XlogError::Kernel(message) => {
1527 XlogError::Kernel(format!("Failed to read output count: {}", message))
1528 }
1529 other => XlogError::Kernel(format!("Failed to read output count: {}", other)),
1530 })
1531 }
1532
1533 fn is_full_row_key(key_cols: &[usize], arity: usize) -> bool {
1534 key_cols.len() == arity
1535 && key_cols
1536 .iter()
1537 .copied()
1538 .enumerate()
1539 .all(|(expected, actual)| expected == actual)
1540 }
1541
1542 fn dedup_full_row_deterministic(&self, input: &CudaBuffer) -> Result<CudaBuffer> {
1559 self.ensure_column_bytes_kernel_indexable(input)?;
1563 let row_count = if input.canonical_full_row_set_certified() {
1564 self.device_row_count(input)?
1565 } else {
1566 self.validated_logical_row_count(input)?
1567 };
1568 if row_count == 0 {
1569 return self.create_empty_buffer(input.schema().clone());
1570 }
1571 if row_count == 1 {
1572 return self.clone_buffer(input);
1573 }
1574 if row_count > u32::MAX as usize {
1575 return Err(XlogError::Kernel(format!(
1576 "dedup_full_row supports at most {} rows, got {}",
1577 u32::MAX,
1578 row_count
1579 )));
1580 }
1581 let arity = input.arity();
1582 if arity == 0 {
1583 return self.buffer_from_columns(Vec::new(), 1, input.schema().clone());
1585 }
1586
1587 let sorted = if Self::use_csm_cuda_graph_env() && row_count <= SMALL_FULL_ROW_SORT_MAX_ROWS
1590 {
1591 self.small_sort_full_row_deterministic(input, row_count)?
1592 } else {
1593 let all_cols: Vec<usize> = (0..arity).collect();
1594 self.sort(input, &all_cols)?
1595 };
1596
1597 let n = self.device_row_count(&sorted)? as u32;
1599 if n <= 1 {
1600 return Ok(sorted);
1601 }
1602
1603 let device = self.device.inner();
1604 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(arity);
1605 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(arity);
1606 for col_idx in 0..arity {
1607 let col = sorted
1608 .column(col_idx)
1609 .ok_or_else(|| XlogError::Kernel(format!("Sorted column {} not found", col_idx)))?;
1610 let ty = sorted.schema().column_type(col_idx).ok_or_else(|| {
1611 XlogError::Kernel(format!("Sorted column {} type missing", col_idx))
1612 })?;
1613 col_ptrs_host.push(*col.device_ptr());
1614 col_sizes_host.push(ty.size_bytes() as u32);
1615 }
1616
1617 let mut d_col_ptrs = self.memory.alloc::<u64>(arity)?;
1618 let mut d_col_sizes = self.memory.alloc::<u32>(arity)?;
1619 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
1620 .map_err(|e| XlogError::Kernel(format!("dedup_full_row_gpu col ptr upload: {}", e)))?;
1621 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
1622 .map_err(|e| XlogError::Kernel(format!("dedup_full_row_gpu col size upload: {}", e)))?;
1623
1624 let block_size = 256u32;
1625 let grid = n.div_ceil(block_size);
1626 let cfg = LaunchConfig {
1627 grid_dim: (grid, 1, 1),
1628 block_dim: (block_size, 1, 1),
1629 shared_mem_bytes: 0,
1630 };
1631
1632 let d_unique_mask = self.memory.alloc::<u8>(n as usize)?;
1633 let mark_fn = device
1634 .get_func(DEDUP_MODULE, dedup_kernels::MARK_UNIQUE_FULL_ROW_BYTEWISE)
1635 .ok_or_else(|| {
1636 XlogError::Kernel("mark_unique_full_row_bytewise kernel not found".to_string())
1637 })?;
1638
1639 unsafe {
1643 mark_fn.clone().launch(
1644 cfg,
1645 (
1646 &d_col_ptrs,
1647 &d_col_sizes,
1648 arity as u32,
1649 sorted.num_rows_device(),
1650 n,
1651 &d_unique_mask,
1652 ),
1653 )
1654 }
1655 .map_err(|e| {
1656 XlogError::Kernel(format!(
1657 "mark_unique_full_row_bytewise launch failed: {}",
1658 e
1659 ))
1660 })?;
1661 self.device.synchronize()?;
1662
1663 let (d_prefix_sum, d_out_count) = self.scan_mask_to_prefix_with_count(&d_unique_mask, n)?;
1665
1666 let mut result = self.compact_buffer_by_device_mask_device_count(
1670 &sorted,
1671 &d_unique_mask,
1672 &d_prefix_sum,
1673 d_out_count,
1674 )?;
1675 result.certify_canonical_full_row_set();
1676 Ok(result)
1677 }
1678
1679 fn small_sort_full_row_deterministic(
1680 &self,
1681 input: &CudaBuffer,
1682 row_count: usize,
1683 ) -> Result<CudaBuffer> {
1684 if row_count > SMALL_FULL_ROW_SORT_MAX_ROWS {
1685 return Err(XlogError::Kernel(format!(
1686 "small full-row sort supports at most {} rows, got {}",
1687 SMALL_FULL_ROW_SORT_MAX_ROWS, row_count
1688 )));
1689 }
1690 if row_count == 0 {
1691 return self.create_empty_buffer(input.schema().clone());
1692 }
1693 if row_count == 1 {
1694 return self.clone_buffer(input);
1695 }
1696
1697 let arity = input.arity();
1698 let device = self.device.inner();
1699 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(arity);
1700 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(arity);
1701 let mut col_types_host: Vec<u8> = Vec::with_capacity(arity);
1702 for col_idx in 0..arity {
1703 let col = input.column(col_idx).ok_or_else(|| {
1704 XlogError::Kernel(format!("small full-row sort: column {} missing", col_idx))
1705 })?;
1706 let ty = input.schema().column_type(col_idx).ok_or_else(|| {
1707 XlogError::Kernel(format!(
1708 "small full-row sort: column {} type missing",
1709 col_idx
1710 ))
1711 })?;
1712 let elem_size = ty.size_bytes();
1713 let expected_bytes_u64 =
1714 input
1715 .num_rows()
1716 .checked_mul(elem_size as u64)
1717 .ok_or_else(|| {
1718 XlogError::Kernel(
1719 "small full-row sort: column byte-size overflow".to_string(),
1720 )
1721 })?;
1722 let expected_bytes = usize::try_from(expected_bytes_u64).map_err(|_| {
1723 XlogError::Kernel(format!(
1724 "small full-row sort: expected byte size {} exceeds usize::MAX",
1725 expected_bytes_u64
1726 ))
1727 })?;
1728 if col.num_bytes() != expected_bytes {
1729 return Err(XlogError::Kernel(format!(
1730 "small full-row sort: column {} has {} bytes but expected {}",
1731 col_idx,
1732 col.num_bytes(),
1733 expected_bytes
1734 )));
1735 }
1736 col_ptrs_host.push(*col.device_ptr());
1737 col_sizes_host.push(elem_size as u32);
1738 col_types_host.push(scalar_type_code_dedup(ty));
1739 }
1740
1741 let mut d_col_ptrs = self.memory.alloc::<u64>(arity)?;
1742 let mut d_col_sizes = self.memory.alloc::<u32>(arity)?;
1743 let mut d_col_types = self.memory.alloc::<u8>(arity)?;
1744 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
1745 .map_err(|e| XlogError::Kernel(format!("small full-row sort ptr upload: {}", e)))?;
1746 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
1747 .map_err(|e| XlogError::Kernel(format!("small full-row sort size upload: {}", e)))?;
1748 self.htod_launch_metadata_sync_copy_into(&col_types_host, &mut d_col_types)
1749 .map_err(|e| XlogError::Kernel(format!("small full-row sort type upload: {}", e)))?;
1750
1751 let mut d_indices = self.memory.alloc::<u32>(row_count)?;
1752 let sort_fn = device
1753 .get_func(
1754 DEDUP_MODULE,
1755 dedup_kernels::SMALL_SORT_FULL_ROW_INDICES_TYPED,
1756 )
1757 .ok_or_else(|| {
1758 XlogError::Kernel("small_sort_full_row_indices_typed kernel not found".to_string())
1759 })?;
1760 let cfg = LaunchConfig {
1761 grid_dim: (1, 1, 1),
1762 block_dim: (SMALL_FULL_ROW_SORT_MAX_ROWS as u32, 1, 1),
1763 shared_mem_bytes: 0,
1764 };
1765
1766 unsafe {
1770 sort_fn.clone().launch(
1771 cfg,
1772 (
1773 &d_col_ptrs,
1774 &d_col_sizes,
1775 &d_col_types,
1776 arity as u32,
1777 input.num_rows_device(),
1778 row_count as u32,
1779 &mut d_indices,
1780 ),
1781 )
1782 }
1783 .map_err(|e| {
1784 XlogError::Kernel(format!(
1785 "small_sort_full_row_indices_typed launch failed: {}",
1786 e
1787 ))
1788 })?;
1789 self.device.synchronize()?;
1790 self.small_full_row_sort_invocations
1791 .fetch_add(1, Ordering::Relaxed);
1792
1793 self.gather_buffer_by_indices(input, &d_indices, row_count as u32)
1794 }
1795
1796 fn scan_mask_to_prefix_with_count(
1801 &self,
1802 d_mask: &cudarc::driver::CudaSlice<u8>,
1803 n: u32,
1804 ) -> Result<(
1805 crate::memory::TrackedCudaSlice<u32>,
1806 crate::memory::TrackedCudaSlice<u32>,
1807 )> {
1808 let device = self.device.inner();
1809 let block_size = 256u32;
1810 let num_blocks = n.div_ceil(block_size);
1811
1812 let d_prefix_sum = self.memory.alloc::<u32>(n as usize)?;
1813 let mut d_block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
1814
1815 let phase1_fn = device
1816 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
1817 .ok_or_else(|| {
1818 XlogError::Kernel("Failed to get multiblock_scan_phase1 kernel".to_string())
1819 })?;
1820 unsafe {
1822 phase1_fn.clone().launch(
1823 LaunchConfig {
1824 grid_dim: (num_blocks, 1, 1),
1825 block_dim: (block_size, 1, 1),
1826 shared_mem_bytes: 0,
1827 },
1828 (d_mask, &d_prefix_sum, &d_block_sums, n),
1829 )
1830 }
1831 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
1832
1833 if num_blocks > 1 {
1834 self.multiblock_scan_u32_inplace(&mut d_block_sums, num_blocks)?;
1835
1836 let phase3_fn = device
1837 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
1838 .ok_or_else(|| {
1839 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
1840 })?;
1841 unsafe {
1843 phase3_fn.clone().launch(
1844 LaunchConfig {
1845 grid_dim: (num_blocks, 1, 1),
1846 block_dim: (block_size, 1, 1),
1847 shared_mem_bytes: 0,
1848 },
1849 (&d_prefix_sum, &d_block_sums, n),
1850 )
1851 }
1852 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
1853 }
1854 self.device.synchronize()?;
1855
1856 let d_out_count = self.capture_compact_count(&d_prefix_sum, d_mask, n)?;
1857 Ok((d_prefix_sum, d_out_count))
1858 }
1859
1860 pub(super) const SORT_BLOCK_SIZE: u32 = 256;
1863
1864 pub fn sort(&self, input: &CudaBuffer, key_cols: &[usize]) -> Result<CudaBuffer> {
1882 self.ensure_column_bytes_kernel_indexable(input)?;
1887
1888 if Self::use_recorded_sort_env() && !key_cols.is_empty() && input.num_rows() > 0 {
1893 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
1894 let recorded_compatible = key_cols.iter().all(|&k| {
1895 matches!(
1896 input.schema.column_type(k),
1897 Some(ScalarType::U32) | Some(ScalarType::Symbol)
1898 )
1899 });
1900 if recorded_compatible {
1901 return self.sort_recorded(input, key_cols, launch_stream);
1902 }
1903 }
1904 }
1905
1906 if input.num_rows() == 0 {
1907 return self.create_empty_buffer(input.schema.clone());
1908 }
1909
1910 if key_cols.is_empty() {
1911 return Err(XlogError::Kernel(
1912 "Sort requires at least one key column".to_string(),
1913 ));
1914 }
1915
1916 if input.num_rows() > u32::MAX as u64 {
1917 return Err(XlogError::Kernel(format!(
1918 "Sort supports at most {} rows, got {}",
1919 u32::MAX,
1920 input.num_rows()
1921 )));
1922 }
1923
1924 for &key_col in key_cols {
1925 if key_col >= input.arity() {
1926 return Err(XlogError::Kernel(format!(
1927 "Key column index {} out of bounds (arity {})",
1928 key_col,
1929 input.arity()
1930 )));
1931 }
1932 }
1933
1934 let n = input.num_rows() as u32;
1935 let d_num_rows = input.num_rows_device();
1936 let device = self.device.inner();
1937
1938 let block_size = Self::SORT_BLOCK_SIZE;
1939 let grid_size = n.div_ceil(block_size);
1940 let launch_config = LaunchConfig {
1941 grid_dim: (grid_size, 1, 1),
1942 block_dim: (block_size, 1, 1),
1943 shared_mem_bytes: 0,
1944 };
1945
1946 let init_fn = device
1948 .get_func(SORT_MODULE, sort_kernels::INIT_INDICES)
1949 .ok_or_else(|| XlogError::Kernel("init_indices kernel not found".to_string()))?;
1950
1951 let mut indices_a = self.memory.alloc::<u32>(n as usize)?;
1952 let mut indices_b = self.memory.alloc::<u32>(n as usize)?;
1953
1954 unsafe {
1956 init_fn
1957 .clone()
1958 .launch(launch_config, (&mut indices_a, d_num_rows, n))
1959 }
1960 .map_err(|e| XlogError::Kernel(format!("init_indices failed: {}", e)))?;
1961 self.device.synchronize()?;
1962
1963 let mut keys_a = self.memory.alloc::<u32>(n as usize)?;
1965 let mut keys_b = self.memory.alloc::<u32>(n as usize)?;
1966
1967 let mut d_hist = self.memory.alloc::<u32>((grid_size as usize) * 16)?;
1969 let mut d_prefix = self.memory.alloc::<u32>(16)?;
1970 let mut d_ranks = self.memory.alloc::<u32>(n as usize)?;
1971
1972 for &col_idx in key_cols.iter().rev() {
1974 let ty = input.schema.column_type(col_idx).ok_or_else(|| {
1975 XlogError::Kernel(format!("Key column {} type not found in schema", col_idx))
1976 })?;
1977
1978 let col = input
1979 .column(col_idx)
1980 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
1981
1982 match ty {
1983 ScalarType::U32 | ScalarType::Symbol => {
1984 let col_view = self.column_as_u32_view(col, n as usize)?;
1985 let gather_fn = device
1986 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_U32)
1987 .ok_or_else(|| {
1988 XlogError::Kernel("apply_permutation_u32 kernel not found".to_string())
1989 })?;
1990
1991 unsafe {
1993 gather_fn.clone().launch(
1994 launch_config,
1995 (&col_view, &mut keys_a, &indices_a, d_num_rows, n),
1996 )
1997 }
1998 .map_err(|e| {
1999 XlogError::Kernel(format!("apply_permutation_u32 failed: {}", e))
2000 })?;
2001
2002 self.radix_sort_u32_pairs_with_scratch(
2003 &mut keys_a,
2004 &mut keys_b,
2005 &mut indices_a,
2006 &mut indices_b,
2007 &mut d_hist,
2008 &mut d_prefix,
2009 &mut d_ranks,
2010 d_num_rows,
2011 n,
2012 )?;
2013 }
2014 ScalarType::I32 => {
2015 let col_bits = self.column_as_u32_view(col, n as usize)?;
2016 let gather_fn = device
2017 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_I32_ORDERED_U32)
2018 .ok_or_else(|| {
2019 XlogError::Kernel(
2020 "gather_keys_i32_ordered_u32 kernel not found".to_string(),
2021 )
2022 })?;
2023
2024 unsafe {
2026 gather_fn.clone().launch(
2027 launch_config,
2028 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2029 )
2030 }
2031 .map_err(|e| {
2032 XlogError::Kernel(format!("gather_keys_i32_ordered_u32 failed: {}", e))
2033 })?;
2034
2035 self.radix_sort_u32_pairs_with_scratch(
2036 &mut keys_a,
2037 &mut keys_b,
2038 &mut indices_a,
2039 &mut indices_b,
2040 &mut d_hist,
2041 &mut d_prefix,
2042 &mut d_ranks,
2043 d_num_rows,
2044 n,
2045 )?;
2046 }
2047 ScalarType::F32 => {
2048 let col_bits = self.column_as_u32_view(col, n as usize)?;
2049 let gather_fn = device
2050 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_F32_ORDERED_U32)
2051 .ok_or_else(|| {
2052 XlogError::Kernel(
2053 "gather_keys_f32_ordered_u32 kernel not found".to_string(),
2054 )
2055 })?;
2056
2057 unsafe {
2059 gather_fn.clone().launch(
2060 launch_config,
2061 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2062 )
2063 }
2064 .map_err(|e| {
2065 XlogError::Kernel(format!("gather_keys_f32_ordered_u32 failed: {}", e))
2066 })?;
2067
2068 self.radix_sort_u32_pairs_with_scratch(
2069 &mut keys_a,
2070 &mut keys_b,
2071 &mut indices_a,
2072 &mut indices_b,
2073 &mut d_hist,
2074 &mut d_prefix,
2075 &mut d_ranks,
2076 d_num_rows,
2077 n,
2078 )?;
2079 }
2080 ScalarType::Bool => {
2081 if col.num_bytes() < n as usize {
2082 return Err(XlogError::Kernel(format!(
2083 "Bool column {} has {} bytes but expected {}",
2084 col_idx,
2085 col.num_bytes(),
2086 n
2087 )));
2088 }
2089
2090 let gather_fn = device
2091 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_BOOL_ORDERED_U32)
2092 .ok_or_else(|| {
2093 XlogError::Kernel(
2094 "gather_keys_bool_ordered_u32 kernel not found".to_string(),
2095 )
2096 })?;
2097
2098 unsafe {
2100 gather_fn
2101 .clone()
2102 .launch(launch_config, (col, &indices_a, d_num_rows, n, &mut keys_a))
2103 }
2104 .map_err(|e| {
2105 XlogError::Kernel(format!("gather_keys_bool_ordered_u32 failed: {}", e))
2106 })?;
2107
2108 self.radix_sort_u32_pairs_with_scratch(
2109 &mut keys_a,
2110 &mut keys_b,
2111 &mut indices_a,
2112 &mut indices_b,
2113 &mut d_hist,
2114 &mut d_prefix,
2115 &mut d_ranks,
2116 d_num_rows,
2117 n,
2118 )?;
2119 }
2120 ScalarType::U64 => {
2121 let col_bits = self.column_as_u64_view(col, n as usize)?;
2122 for &word in &[
2123 sort_kernels::GATHER_KEYS_U64_LO_U32,
2124 sort_kernels::GATHER_KEYS_U64_HI_U32,
2125 ] {
2126 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
2127 XlogError::Kernel(format!("{} kernel not found", word))
2128 })?;
2129
2130 unsafe {
2132 gather_fn.clone().launch(
2133 launch_config,
2134 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2135 )
2136 }
2137 .map_err(|e| XlogError::Kernel(format!("{} failed: {}", word, e)))?;
2138
2139 self.radix_sort_u32_pairs_with_scratch(
2140 &mut keys_a,
2141 &mut keys_b,
2142 &mut indices_a,
2143 &mut indices_b,
2144 &mut d_hist,
2145 &mut d_prefix,
2146 &mut d_ranks,
2147 d_num_rows,
2148 n,
2149 )?;
2150 }
2151 }
2152 ScalarType::I64 => {
2153 let col_bits = self.column_as_u64_view(col, n as usize)?;
2154 for &word in &[
2155 sort_kernels::GATHER_KEYS_I64_LO_U32,
2156 sort_kernels::GATHER_KEYS_I64_HI_U32,
2157 ] {
2158 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
2159 XlogError::Kernel(format!("{} kernel not found", word))
2160 })?;
2161
2162 unsafe {
2164 gather_fn.clone().launch(
2165 launch_config,
2166 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2167 )
2168 }
2169 .map_err(|e| XlogError::Kernel(format!("{} failed: {}", word, e)))?;
2170
2171 self.radix_sort_u32_pairs_with_scratch(
2172 &mut keys_a,
2173 &mut keys_b,
2174 &mut indices_a,
2175 &mut indices_b,
2176 &mut d_hist,
2177 &mut d_prefix,
2178 &mut d_ranks,
2179 d_num_rows,
2180 n,
2181 )?;
2182 }
2183 }
2184 ScalarType::F64 => {
2185 let col_bits = self.column_as_u64_view(col, n as usize)?;
2186 for &word in &[
2187 sort_kernels::GATHER_KEYS_F64_LO_U32,
2188 sort_kernels::GATHER_KEYS_F64_HI_U32,
2189 ] {
2190 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
2191 XlogError::Kernel(format!("{} kernel not found", word))
2192 })?;
2193
2194 unsafe {
2196 gather_fn.clone().launch(
2197 launch_config,
2198 (&col_bits, &indices_a, d_num_rows, n, &mut keys_a),
2199 )
2200 }
2201 .map_err(|e| XlogError::Kernel(format!("{} failed: {}", word, e)))?;
2202
2203 self.radix_sort_u32_pairs_with_scratch(
2204 &mut keys_a,
2205 &mut keys_b,
2206 &mut indices_a,
2207 &mut indices_b,
2208 &mut d_hist,
2209 &mut d_prefix,
2210 &mut d_ranks,
2211 d_num_rows,
2212 n,
2213 )?;
2214 }
2215 }
2216 }
2217 }
2218
2219 self.apply_permutation_gpu(input, &indices_a)
2220 }
2221
2222 #[allow(clippy::too_many_arguments)]
2223 fn radix_sort_u32_pairs_with_scratch(
2224 &self,
2225 keys_a: &mut crate::memory::TrackedCudaSlice<u32>,
2226 keys_b: &mut crate::memory::TrackedCudaSlice<u32>,
2227 indices_a: &mut crate::memory::TrackedCudaSlice<u32>,
2228 indices_b: &mut crate::memory::TrackedCudaSlice<u32>,
2229 hist: &mut crate::memory::TrackedCudaSlice<u32>,
2230 prefix: &mut crate::memory::TrackedCudaSlice<u32>,
2231 ranks: &mut crate::memory::TrackedCudaSlice<u32>,
2232 num_rows_device: &crate::memory::TrackedCudaSlice<u32>,
2233 row_cap: u32,
2234 ) -> Result<()> {
2235 if row_cap == 0 {
2236 return Ok(());
2237 }
2238 self.device.synchronize()?;
2239
2240 let device = self.device.inner();
2241 let block_size = Self::SORT_BLOCK_SIZE;
2242 let grid_size = row_cap.div_ceil(block_size);
2243
2244 let sort_config = LaunchConfig {
2245 grid_dim: (grid_size, 1, 1),
2246 block_dim: (block_size, 1, 1),
2247 shared_mem_bytes: 0,
2248 };
2249
2250 let histogram_fn = device
2251 .get_func(SORT_MODULE, sort_kernels::RADIX_HISTOGRAM)
2252 .ok_or_else(|| XlogError::Kernel("radix_histogram kernel not found".to_string()))?;
2253 let prefix_fn = device
2254 .get_func(SORT_MODULE, sort_kernels::COMPUTE_DIGIT_PREFIX_SUMS)
2255 .ok_or_else(|| {
2256 XlogError::Kernel("compute_digit_prefix_sums kernel not found".to_string())
2257 })?;
2258 let ranks_fn = device
2259 .get_func(SORT_MODULE, sort_kernels::COMPUTE_RANKS)
2260 .ok_or_else(|| XlogError::Kernel("compute_ranks kernel not found".to_string()))?;
2261 let scatter_fn = device
2262 .get_func(SORT_MODULE, sort_kernels::RADIX_SCATTER_STABLE)
2263 .ok_or_else(|| {
2264 XlogError::Kernel("radix_scatter_stable kernel not found".to_string())
2265 })?;
2266
2267 let prefix_config = LaunchConfig {
2268 grid_dim: (1, 1, 1),
2269 block_dim: (256, 1, 1),
2270 shared_mem_bytes: 0,
2271 };
2272
2273 let mut in_a = true;
2274 for pass in 0..8u32 {
2275 let shift = pass * 4;
2276
2277 let (keys_in, indices_in, keys_out, indices_out) = if in_a {
2278 (&*keys_a, &*indices_a, &mut *keys_b, &mut *indices_b)
2279 } else {
2280 (&*keys_b, &*indices_b, &mut *keys_a, &mut *indices_a)
2281 };
2282
2283 unsafe {
2286 histogram_fn.clone().launch(
2287 sort_config,
2288 (keys_in, num_rows_device, row_cap, &mut *hist, shift),
2289 )
2290 }
2291 .map_err(|e| XlogError::Kernel(format!("radix_histogram failed: {}", e)))?;
2292 self.device.synchronize()?;
2293
2294 unsafe {
2297 prefix_fn
2298 .clone()
2299 .launch(prefix_config, (&*hist, grid_size, &mut *prefix))
2300 }
2301 .map_err(|e| XlogError::Kernel(format!("compute_digit_prefix_sums failed: {}", e)))?;
2302 self.device.synchronize()?;
2303
2304 for digit in 0..16u32 {
2306 let start = (digit * grid_size) as usize;
2307 let end = start + (grid_size as usize);
2308 let mut digit_slice = hist.slice_mut(start..end);
2309 self.multiblock_scan_u32_view_inplace(&mut digit_slice, grid_size)?;
2310 }
2311 self.device.synchronize()?;
2312
2313 unsafe {
2316 ranks_fn.clone().launch(
2317 sort_config,
2318 (keys_in, num_rows_device, row_cap, &mut *ranks, shift),
2319 )
2320 }
2321 .map_err(|e| XlogError::Kernel(format!("compute_ranks failed: {}", e)))?;
2322 self.device.synchronize()?;
2323
2324 unsafe {
2327 scatter_fn.clone().launch(
2328 sort_config,
2329 (
2330 keys_in,
2331 indices_in,
2332 &*ranks,
2333 keys_out,
2334 indices_out,
2335 &*prefix,
2336 &*hist,
2337 num_rows_device,
2338 row_cap,
2339 shift,
2340 ),
2341 )
2342 }
2343 .map_err(|e| XlogError::Kernel(format!("radix_scatter_stable failed: {}", e)))?;
2344 self.device.synchronize()?;
2345
2346 in_a = !in_a;
2347 }
2348
2349 if !in_a {
2351 return Err(XlogError::Kernel(
2352 "Unexpected radix-sort buffer parity (expected even number of passes)".to_string(),
2353 ));
2354 }
2355
2356 Ok(())
2357 }
2358 pub fn init_indices(
2360 &self,
2361 indices: &mut crate::memory::TrackedCudaSlice<u32>,
2362 n: u32,
2363 ) -> Result<()> {
2364 if n == 0 {
2365 return Ok(());
2366 }
2367 if n as usize > indices.len() {
2368 return Err(XlogError::Kernel(format!(
2369 "init_indices: n={} exceeds indices len={}",
2370 n,
2371 indices.len()
2372 )));
2373 }
2374 let device = self.device.inner();
2375 let block_size = Self::SORT_BLOCK_SIZE;
2376 let grid_size = n.div_ceil(block_size);
2377 let config = LaunchConfig {
2378 grid_dim: (grid_size, 1, 1),
2379 block_dim: (block_size, 1, 1),
2380 shared_mem_bytes: 0,
2381 };
2382 let init_fn = device
2383 .get_func(SORT_MODULE, sort_kernels::INIT_INDICES)
2384 .ok_or_else(|| XlogError::Kernel("init_indices kernel not found".to_string()))?;
2385 let d_num_rows = self.upload_device_row_count(n)?;
2386 unsafe {
2388 init_fn
2389 .clone()
2390 .launch(config, (&mut *indices, &d_num_rows, n))
2391 }
2392 .map_err(|e| XlogError::Kernel(format!("init_indices failed: {}", e)))?;
2393 Ok(())
2394 }
2395
2396 pub fn gather_u32_by_indices(
2398 &self,
2399 input: &crate::memory::TrackedCudaSlice<u32>,
2400 indices: &crate::memory::TrackedCudaSlice<u32>,
2401 output: &mut crate::memory::TrackedCudaSlice<u32>,
2402 n: u32,
2403 ) -> Result<()> {
2404 if n == 0 {
2405 return Ok(());
2406 }
2407 if n as usize > output.len() {
2408 return Err(XlogError::Kernel(format!(
2409 "gather_u32_by_indices: n={} exceeds output len={}",
2410 n,
2411 output.len()
2412 )));
2413 }
2414 let device = self.device.inner();
2415 let block_size = Self::SORT_BLOCK_SIZE;
2416 let grid_size = n.div_ceil(block_size);
2417 let config = LaunchConfig {
2418 grid_dim: (grid_size, 1, 1),
2419 block_dim: (block_size, 1, 1),
2420 shared_mem_bytes: 0,
2421 };
2422 let gather_fn = device
2423 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_U32)
2424 .ok_or_else(|| {
2425 XlogError::Kernel("apply_permutation_u32 kernel not found".to_string())
2426 })?;
2427 let d_num_rows = self.upload_device_row_count(n)?;
2428 unsafe {
2430 gather_fn
2431 .clone()
2432 .launch(config, (input, output, indices, &d_num_rows, n))
2433 }
2434 .map_err(|e| XlogError::Kernel(format!("gather_u32_by_indices failed: {}", e)))?;
2435 Ok(())
2436 }
2437
2438 pub fn gather_u8_by_indices(
2440 &self,
2441 input: &crate::memory::TrackedCudaSlice<u8>,
2442 indices: &crate::memory::TrackedCudaSlice<u32>,
2443 output: &mut crate::memory::TrackedCudaSlice<u8>,
2444 n: u32,
2445 ) -> Result<()> {
2446 if n == 0 {
2447 return Ok(());
2448 }
2449 if n as usize > output.len() {
2450 return Err(XlogError::Kernel(format!(
2451 "gather_u8_by_indices: n={} exceeds output len={}",
2452 n,
2453 output.len()
2454 )));
2455 }
2456 let device = self.device.inner();
2457 let block_size = Self::SORT_BLOCK_SIZE;
2458 let grid_size = n.div_ceil(block_size);
2459 let config = LaunchConfig {
2460 grid_dim: (grid_size, 1, 1),
2461 block_dim: (block_size, 1, 1),
2462 shared_mem_bytes: 0,
2463 };
2464 let gather_fn = device
2465 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
2466 .ok_or_else(|| {
2467 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
2468 })?;
2469 let d_num_rows = self.upload_device_row_count(n)?;
2470 unsafe {
2472 gather_fn
2473 .clone()
2474 .launch(config, (input, output, indices, &d_num_rows, n, 1u32))
2475 }
2476 .map_err(|e| XlogError::Kernel(format!("gather_u8_by_indices failed: {}", e)))?;
2477 Ok(())
2478 }
2479
2480 pub fn gather_u64_lo_by_indices(
2482 &self,
2483 input: &crate::memory::TrackedCudaSlice<u64>,
2484 indices: &crate::memory::TrackedCudaSlice<u32>,
2485 output: &mut crate::memory::TrackedCudaSlice<u32>,
2486 n: u32,
2487 ) -> Result<()> {
2488 if n == 0 {
2489 return Ok(());
2490 }
2491 let device = self.device.inner();
2492 let block_size = Self::SORT_BLOCK_SIZE;
2493 let grid_size = n.div_ceil(block_size);
2494 let config = LaunchConfig {
2495 grid_dim: (grid_size, 1, 1),
2496 block_dim: (block_size, 1, 1),
2497 shared_mem_bytes: 0,
2498 };
2499 let gather_fn = device
2500 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_U64_LO_U32)
2501 .ok_or_else(|| XlogError::Kernel("gather_keys_u64_lo_u32 not found".to_string()))?;
2502 let d_num_rows = self.upload_device_row_count(n)?;
2503 unsafe {
2505 gather_fn
2506 .clone()
2507 .launch(config, (input, indices, &d_num_rows, n, output))
2508 }
2509 .map_err(|e| XlogError::Kernel(format!("gather_u64_lo_by_indices failed: {}", e)))?;
2510 Ok(())
2511 }
2512
2513 pub fn gather_u64_hi_by_indices(
2515 &self,
2516 input: &crate::memory::TrackedCudaSlice<u64>,
2517 indices: &crate::memory::TrackedCudaSlice<u32>,
2518 output: &mut crate::memory::TrackedCudaSlice<u32>,
2519 n: u32,
2520 ) -> Result<()> {
2521 if n == 0 {
2522 return Ok(());
2523 }
2524 let device = self.device.inner();
2525 let block_size = Self::SORT_BLOCK_SIZE;
2526 let grid_size = n.div_ceil(block_size);
2527 let config = LaunchConfig {
2528 grid_dim: (grid_size, 1, 1),
2529 block_dim: (block_size, 1, 1),
2530 shared_mem_bytes: 0,
2531 };
2532 let gather_fn = device
2533 .get_func(SORT_MODULE, sort_kernels::GATHER_KEYS_U64_HI_U32)
2534 .ok_or_else(|| XlogError::Kernel("gather_keys_u64_hi_u32 not found".to_string()))?;
2535 let d_num_rows = self.upload_device_row_count(n)?;
2536 unsafe {
2538 gather_fn
2539 .clone()
2540 .launch(config, (input, indices, &d_num_rows, n, output))
2541 }
2542 .map_err(|e| XlogError::Kernel(format!("gather_u64_hi_by_indices failed: {}", e)))?;
2543 Ok(())
2544 }
2545
2546 pub fn radix_sort_u32_pairs(
2548 &self,
2549 keys: &mut crate::memory::TrackedCudaSlice<u32>,
2550 values: &mut crate::memory::TrackedCudaSlice<u32>,
2551 n: u32,
2552 scratch: &mut RadixSortScratch,
2553 ) -> Result<()> {
2554 if n == 0 {
2555 return Ok(());
2556 }
2557 scratch.ensure_capacity(self, n)?;
2558 let d_num_rows = self.upload_device_row_count(n)?;
2559 self.radix_sort_u32_pairs_with_scratch(
2560 keys,
2561 &mut scratch.keys_b,
2562 values,
2563 &mut scratch.values_b,
2564 &mut scratch.hist,
2565 &mut scratch.prefix,
2566 &mut scratch.ranks,
2567 &d_num_rows,
2568 n,
2569 )
2570 }
2571 pub fn scan_u8_mask_device(
2573 &self,
2574 mask: &crate::memory::TrackedCudaSlice<u8>,
2575 n: u32,
2576 ) -> Result<crate::memory::TrackedCudaSlice<u32>> {
2577 if n == 0 {
2578 return self.memory.alloc::<u32>(0);
2579 }
2580 if n as usize > mask.len() {
2581 return Err(XlogError::Kernel(format!(
2582 "scan_u8_mask_device: n={} exceeds mask len={}",
2583 n,
2584 mask.len()
2585 )));
2586 }
2587 let device = self.device.inner();
2588 let block_size = 256u32;
2589 let num_blocks = n.div_ceil(block_size);
2590
2591 let mut prefix_sum = self.memory.alloc::<u32>(n as usize)?;
2592 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2593
2594 let phase1_fn = device
2595 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE1)
2596 .ok_or_else(|| {
2597 XlogError::Kernel("multiblock_scan_phase1 kernel not found".to_string())
2598 })?;
2599
2600 unsafe {
2602 phase1_fn.clone().launch(
2603 LaunchConfig {
2604 grid_dim: (num_blocks, 1, 1),
2605 block_dim: (block_size, 1, 1),
2606 shared_mem_bytes: 0,
2607 },
2608 (mask, &mut prefix_sum, &mut block_sums, n),
2609 )
2610 }
2611 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase1 failed: {}", e)))?;
2612
2613 if num_blocks > 1 {
2614 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2615
2616 let phase3_fn = device
2617 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2618 .ok_or_else(|| {
2619 XlogError::Kernel("multiblock_scan_phase3 kernel not found".to_string())
2620 })?;
2621
2622 unsafe {
2624 phase3_fn.clone().launch(
2625 LaunchConfig {
2626 grid_dim: (num_blocks, 1, 1),
2627 block_dim: (block_size, 1, 1),
2628 shared_mem_bytes: 0,
2629 },
2630 (&mut prefix_sum, &block_sums, n),
2631 )
2632 }
2633 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2634 }
2635
2636 Ok(prefix_sum)
2637 }
2638
2639 pub fn count_mask_device(
2643 &self,
2644 mask: &crate::memory::TrackedCudaSlice<u8>,
2645 n: u32,
2646 ) -> Result<crate::memory::TrackedCudaSlice<u32>> {
2647 let mut d_count = self.memory.alloc::<u32>(1)?;
2648 self.htod_launch_metadata_sync_copy_into(&[0u32], &mut d_count)
2649 .map_err(|e| {
2650 XlogError::Kernel(format!("count_mask_device: zero init failed: {}", e))
2651 })?;
2652
2653 if n == 0 {
2654 return Ok(d_count);
2655 }
2656
2657 let device = self.device.inner();
2658 let block_size = 256u32;
2659 let grid_size = n.div_ceil(block_size);
2660
2661 let count_fn = device
2662 .get_func(SCAN_MODULE, scan_kernels::COUNT_MASK)
2663 .ok_or_else(|| XlogError::Kernel("count_mask kernel not found".to_string()))?;
2664
2665 unsafe {
2667 count_fn.clone().launch(
2668 LaunchConfig {
2669 grid_dim: (grid_size, 1, 1),
2670 block_dim: (block_size, 1, 1),
2671 shared_mem_bytes: 0,
2672 },
2673 (mask, n, &mut d_count),
2674 )
2675 }
2676 .map_err(|e| XlogError::Kernel(format!("count_mask kernel failed: {}", e)))?;
2677
2678 self.device.synchronize()?;
2679
2680 Ok(d_count)
2681 }
2682
2683 pub fn count_mask_into_slot(
2692 &self,
2693 mask: &crate::memory::TrackedCudaSlice<u8>,
2694 n: u32,
2695 task_counts: &mut crate::memory::TrackedCudaSlice<u32>,
2696 slot_idx: usize,
2697 ) -> Result<()> {
2698 if n == 0 {
2699 return Ok(());
2701 }
2702 if slot_idx >= task_counts.len() {
2703 return Err(XlogError::Kernel(format!(
2704 "count_mask_into_slot: slot_idx={} >= len={}",
2705 slot_idx,
2706 task_counts.len()
2707 )));
2708 }
2709
2710 let device = self.device.inner();
2711 let block_size = 256u32;
2712 let grid_size = n.div_ceil(block_size);
2713
2714 let count_fn = device
2715 .get_func(SCAN_MODULE, scan_kernels::COUNT_MASK)
2716 .ok_or_else(|| XlogError::Kernel("count_mask kernel not found".to_string()))?;
2717
2718 let mut slot = task_counts.slice_mut(slot_idx..slot_idx + 1);
2720
2721 unsafe {
2725 count_fn.clone().launch(
2726 LaunchConfig {
2727 grid_dim: (grid_size, 1, 1),
2728 block_dim: (block_size, 1, 1),
2729 shared_mem_bytes: 0,
2730 },
2731 (mask, n, &mut slot),
2732 )
2733 }
2734 .map_err(|e| XlogError::Kernel(format!("count_mask_into_slot kernel failed: {}", e)))?;
2735
2736 Ok(())
2737 }
2738 fn apply_permutation_gpu(
2740 &self,
2741 input: &CudaBuffer,
2742 permutation: &cudarc::driver::CudaSlice<u32>,
2743 ) -> Result<CudaBuffer> {
2744 let row_cap = input.num_rows() as u32;
2745 let d_num_rows = input.num_rows_device();
2746 let device = self.device.inner();
2747
2748 let grid_size = row_cap.div_ceil(Self::SORT_BLOCK_SIZE);
2749 let launch_config = LaunchConfig {
2750 grid_dim: (grid_size, 1, 1),
2751 block_dim: (Self::SORT_BLOCK_SIZE, 1, 1),
2752 shared_mem_bytes: 0,
2753 };
2754
2755 let apply_perm_fn = device
2756 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
2757 .ok_or_else(|| {
2758 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
2759 })?;
2760
2761 let mut new_columns = Vec::with_capacity(input.columns.len());
2762
2763 for col_idx in 0..input.columns.len() {
2764 let src_col = input
2765 .column(col_idx)
2766 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
2767
2768 let elem_size = input
2769 .schema
2770 .column_type(col_idx)
2771 .ok_or_else(|| {
2772 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
2773 })?
2774 .size_bytes() as u32;
2775
2776 let output_bytes = (row_cap as usize) * (elem_size as usize);
2777 if src_col.num_bytes() != output_bytes {
2778 return Err(XlogError::Kernel(format!(
2779 "Column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
2780 col_idx,
2781 src_col.num_bytes(),
2782 output_bytes,
2783 row_cap,
2784 elem_size
2785 )));
2786 }
2787 let dst_col = self.memory.alloc::<u8>(output_bytes)?;
2788
2789 unsafe {
2791 apply_perm_fn.clone().launch(
2792 launch_config,
2793 (
2794 src_col,
2795 &dst_col,
2796 permutation,
2797 d_num_rows,
2798 row_cap,
2799 elem_size,
2800 ),
2801 )
2802 }
2803 .map_err(|e| XlogError::Kernel(format!("apply_permutation_bytes failed: {}", e)))?;
2804
2805 new_columns.push(dst_col.into());
2806 }
2807
2808 self.device.synchronize()?;
2809
2810 self.buffer_from_columns_with_device_count(
2811 new_columns,
2812 input.num_rows(),
2813 input.schema.clone(),
2814 input,
2815 )
2816 }
2817
2818 fn gather_buffer_by_indices(
2823 &self,
2824 input: &CudaBuffer,
2825 indices: &cudarc::driver::CudaSlice<u32>,
2826 output_rows: u32,
2827 ) -> Result<CudaBuffer> {
2828 if output_rows == 0 {
2829 return self.create_empty_buffer(input.schema().clone());
2830 }
2831
2832 if input.num_rows() > u32::MAX as u64 {
2833 return Err(XlogError::Kernel(format!(
2834 "GPU gather supports at most {} input rows, got {}",
2835 u32::MAX,
2836 input.num_rows()
2837 )));
2838 }
2839
2840 let d_output_rows = self.upload_device_row_count(output_rows)?;
2841 let device = self.device.inner();
2842 let block_size = 256u32;
2843 let grid_size = output_rows.div_ceil(block_size);
2844 let launch_config = LaunchConfig {
2845 grid_dim: (grid_size, 1, 1),
2846 block_dim: (block_size, 1, 1),
2847 shared_mem_bytes: 0,
2848 };
2849
2850 let gather_fn = device
2851 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
2852 .ok_or_else(|| {
2853 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
2854 })?;
2855
2856 let mut new_columns = Vec::with_capacity(input.columns.len());
2857 for col_idx in 0..input.columns.len() {
2858 let src_col = input
2859 .column(col_idx)
2860 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
2861
2862 let elem_size = input
2863 .schema
2864 .column_type(col_idx)
2865 .ok_or_else(|| {
2866 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
2867 })?
2868 .size_bytes() as u32;
2869
2870 let expected_src_bytes = (input.num_rows() as usize) * (elem_size as usize);
2871 if src_col.num_bytes() != expected_src_bytes {
2872 return Err(XlogError::Kernel(format!(
2873 "Column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
2874 col_idx,
2875 src_col.num_bytes(),
2876 expected_src_bytes,
2877 input.num_rows(),
2878 elem_size
2879 )));
2880 }
2881
2882 let dst_bytes = (output_rows as usize) * (elem_size as usize);
2883 let dst_col = self.memory.alloc::<u8>(dst_bytes)?;
2884
2885 unsafe {
2887 gather_fn.clone().launch(
2888 launch_config,
2889 (
2890 src_col,
2891 &dst_col,
2892 indices,
2893 &d_output_rows,
2894 output_rows,
2895 elem_size,
2896 ),
2897 )
2898 }
2899 .map_err(|e| XlogError::Kernel(format!("apply_permutation_bytes failed: {}", e)))?;
2900
2901 new_columns.push(dst_col.into());
2902 }
2903
2904 self.device.synchronize()?;
2905
2906 Ok(CudaBuffer::from_columns(
2907 new_columns,
2908 output_rows as u64,
2909 d_output_rows,
2910 input.schema.clone(),
2911 ))
2912 }
2913 pub fn hash_join_v2(
2927 &self,
2928 left: &CudaBuffer,
2929 right: &CudaBuffer,
2930 left_keys: &[usize],
2931 right_keys: &[usize],
2932 join_type: JoinType,
2933 ) -> Result<CudaBuffer> {
2934 self.hash_join_v2_with_limit(left, right, left_keys, right_keys, join_type, None)
2935 }
2936
2937 pub fn hash_join_v2_with_limit(
2953 &self,
2954 left: &CudaBuffer,
2955 right: &CudaBuffer,
2956 left_keys: &[usize],
2957 right_keys: &[usize],
2958 join_type: JoinType,
2959 max_output: Option<usize>,
2960 ) -> Result<CudaBuffer> {
2961 if Self::use_recorded_hash_join_env()
2966 && !left_keys.is_empty()
2967 && left_keys.len() == right_keys.len()
2968 && left_keys.len() <= 4
2969 {
2970 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
2971 return self.hash_join_v2_recorded(
2972 left,
2973 right,
2974 left_keys,
2975 right_keys,
2976 join_type,
2977 max_output,
2978 launch_stream,
2979 );
2980 }
2981 }
2982 match join_type {
2983 JoinType::Inner => {
2984 self.hash_join_inner_v2(left, right, left_keys, right_keys, max_output)
2985 }
2986 JoinType::Semi => self.hash_join_semi_impl(left, right, left_keys, right_keys),
2987 JoinType::Anti => self.hash_join_anti_impl(left, right, left_keys, right_keys),
2988 JoinType::LeftOuter => {
2989 self.hash_join_left_outer_impl(left, right, left_keys, right_keys, max_output)
2990 }
2991 }
2992 }
2993
2994 pub fn nested_loop_join_v2_inner_u32_1key(
3042 &self,
3043 left: &CudaBuffer,
3044 right: &CudaBuffer,
3045 left_key: usize,
3046 right_key: usize,
3047 ) -> Result<CudaBuffer> {
3048 let num_left = self.device_row_count(left)?;
3050 let num_right = self.device_row_count(right)?;
3051
3052 if num_left == 0 || num_right == 0 {
3054 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3055 return self.create_empty_buffer(combined_schema);
3056 }
3057
3058 if left.arity() <= left_key {
3060 return Err(XlogError::Kernel(format!(
3061 "nested_loop: left_key={} out of bounds (arity={})",
3062 left_key,
3063 left.arity()
3064 )));
3065 }
3066 if right.arity() <= right_key {
3067 return Err(XlogError::Kernel(format!(
3068 "nested_loop: right_key={} out of bounds (arity={})",
3069 right_key,
3070 right.arity()
3071 )));
3072 }
3073 let lt = left.schema().column_type(left_key);
3074 let rt = right.schema().column_type(right_key);
3075 if lt != rt || !matches!(lt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3076 return Err(XlogError::Kernel(format!(
3077 "nested_loop: key types must be equal U32/Symbol; got left={:?} right={:?}",
3078 lt, rt
3079 )));
3080 }
3081 let left_col = left
3082 .column(left_key)
3083 .ok_or_else(|| XlogError::Kernel(format!("nested_loop: left.column({})", left_key)))?;
3084 let right_col = right.column(right_key).ok_or_else(|| {
3085 XlogError::Kernel(format!("nested_loop: right.column({})", right_key))
3086 })?;
3087 let required_left_bytes = num_left
3102 .checked_mul(4)
3103 .ok_or_else(|| XlogError::Kernel("nested_loop: left byte-count overflow".into()))?;
3104 let required_right_bytes = num_right
3105 .checked_mul(4)
3106 .ok_or_else(|| XlogError::Kernel("nested_loop: right byte-count overflow".into()))?;
3107 if left_col.num_bytes() < required_left_bytes {
3108 return Err(XlogError::Kernel(format!(
3109 "nested_loop: left key column has {} bytes; \
3110 require at least {} ({} rows × 4) — buffer allocation \
3111 is smaller than logical row count",
3112 left_col.num_bytes(),
3113 required_left_bytes,
3114 num_left
3115 )));
3116 }
3117 if right_col.num_bytes() < required_right_bytes {
3118 return Err(XlogError::Kernel(format!(
3119 "nested_loop: right key column has {} bytes; \
3120 require at least {} ({} rows × 4) — buffer allocation \
3121 is smaller than logical row count",
3122 right_col.num_bytes(),
3123 required_right_bytes,
3124 num_right
3125 )));
3126 }
3127
3128 let upper_bound: u64 = (num_left as u64)
3130 .checked_mul(num_right as u64)
3131 .ok_or_else(|| XlogError::Kernel("nested_loop: row-count product overflow".into()))?;
3132 if upper_bound > NESTED_LOOP_TOTAL_THRESHOLD {
3133 return Err(XlogError::Kernel(format!(
3134 "nested_loop: caller violated eligibility threshold: \
3135 num_left * num_right = {} > {} (NESTED_LOOP_TOTAL_THRESHOLD)",
3136 upper_bound, NESTED_LOOP_TOTAL_THRESHOLD
3137 )));
3138 }
3139
3140 let upper_bound_usize = upper_bound as usize;
3142 let mut d_output_left_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3143 let mut d_output_right_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3144 let mut d_output_count = self.memory.alloc::<u32>(1)?;
3145 self.device
3146 .inner()
3147 .memset_zeros(&mut d_output_count)
3148 .map_err(|e| XlogError::Kernel(format!("nested_loop: counter zero failed: {}", e)))?;
3149
3150 let func = self
3152 .device
3153 .inner()
3154 .get_func(
3155 JOIN_MODULE,
3156 join_kernels::NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS,
3157 )
3158 .ok_or_else(|| {
3159 XlogError::Kernel("nested_loop_join_inner_u32_1key_pairs kernel not found".into())
3160 })?;
3161
3162 let num_left_u32 = num_left as u32;
3163 let num_right_u32 = num_right as u32;
3164 let upper_bound_u32 = upper_bound as u32;
3165 let block_size = 256u32;
3166 let grid_size = num_left_u32.div_ceil(block_size);
3167 let config = LaunchConfig {
3168 grid_dim: (grid_size, 1, 1),
3169 block_dim: (block_size, 1, 1),
3170 shared_mem_bytes: 0,
3171 };
3172
3173 unsafe {
3183 func.clone()
3184 .launch(
3185 config,
3186 (
3187 left_col,
3188 right_col,
3189 num_left_u32,
3190 num_right_u32,
3191 &mut d_output_left_idx,
3192 &mut d_output_right_idx,
3193 &mut d_output_count,
3194 upper_bound_u32,
3195 ),
3196 )
3197 .map_err(|e| XlogError::Kernel(format!("nested_loop launch failed: {}", e)))?;
3198 }
3199
3200 self.device.synchronize()?;
3201
3202 let output_rows = self.dtoh_scalar_untracked(&d_output_count, 0)?;
3204 if (output_rows as u64) > upper_bound {
3208 return Err(XlogError::Kernel(format!(
3209 "nested_loop: kernel reported {} output rows > upper_bound {}",
3210 output_rows, upper_bound
3211 )));
3212 }
3213
3214 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left_idx, output_rows)?;
3216 let gathered_right =
3217 self.gather_buffer_by_indices(right, &d_output_right_idx, output_rows)?;
3218
3219 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3221 let mut result_columns = Vec::with_capacity(combined_schema.arity());
3222 result_columns.extend(gathered_left.columns);
3223 result_columns.extend(gathered_right.columns);
3224 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
3227 }
3228
3229 pub fn is_sorted_ascending_u32(&self, buf: &CudaBuffer, key_col: usize) -> Result<bool> {
3259 let n = self.device_row_count(buf)?;
3261 if n < 2 {
3262 return Ok(true);
3263 }
3264
3265 if buf.arity() <= key_col {
3267 return Err(XlogError::Kernel(format!(
3268 "is_sorted_ascending_u32: key_col={} out of bounds (arity={})",
3269 key_col,
3270 buf.arity()
3271 )));
3272 }
3273 let kt = buf.schema().column_type(key_col);
3274 if !matches!(kt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3275 return Err(XlogError::Kernel(format!(
3276 "is_sorted_ascending_u32: key column must be U32 or Symbol; got {:?}",
3277 kt
3278 )));
3279 }
3280 let key_column = buf.column(key_col).ok_or_else(|| {
3281 XlogError::Kernel(format!(
3282 "is_sorted_ascending_u32: column({}) missing",
3283 key_col
3284 ))
3285 })?;
3286 let required_bytes = n
3287 .checked_mul(4)
3288 .ok_or_else(|| XlogError::Kernel("is_sorted_ascending_u32: byte overflow".into()))?;
3289 if key_column.num_bytes() < required_bytes {
3290 return Err(XlogError::Kernel(format!(
3291 "is_sorted_ascending_u32: key column has {} bytes; require at least {} ({} rows × 4)",
3292 key_column.num_bytes(),
3293 required_bytes,
3294 n
3295 )));
3296 }
3297
3298 let mut d_result = self.memory.alloc::<u32>(1)?;
3302 self.htod_launch_metadata_sync_copy_into(&[1u32], &mut d_result)
3303 .map_err(|e| {
3304 XlogError::Kernel(format!("is_sorted_ascending_u32: htod result init: {}", e))
3305 })?;
3306
3307 let func = self
3309 .device
3310 .inner()
3311 .get_func(SORT_MODULE, sort_kernels::CHECK_ASCENDING_SORTED_U32)
3312 .ok_or_else(|| {
3313 XlogError::Kernel("check_ascending_sorted_u32 kernel not found".into())
3314 })?;
3315 let n_u32 = n as u32;
3316 let block_size = 256u32;
3317 let grid_size = n_u32.div_ceil(block_size);
3318 let config = LaunchConfig {
3319 grid_dim: (grid_size, 1, 1),
3320 block_dim: (block_size, 1, 1),
3321 shared_mem_bytes: 0,
3322 };
3323
3324 unsafe {
3332 func.clone()
3333 .launch(config, (key_column, n_u32, &mut d_result))
3334 .map_err(|e| {
3335 XlogError::Kernel(format!("check_ascending_sorted_u32 launch: {}", e))
3336 })?;
3337 }
3338
3339 self.device.synchronize()?;
3340 let result = self.dtoh_scalar_untracked(&d_result, 0)?;
3341 Ok(result == 1)
3342 }
3343
3344 pub fn sort_merge_join_v2_inner_u32_1key(
3402 &self,
3403 left: &CudaBuffer,
3404 right: &CudaBuffer,
3405 left_key: usize,
3406 right_key: usize,
3407 ) -> Result<CudaBuffer> {
3408 let num_left = self.device_row_count(left)?;
3410 let num_right = self.device_row_count(right)?;
3411
3412 if num_left == 0 || num_right == 0 {
3414 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3415 return self.create_empty_buffer(combined_schema);
3416 }
3417
3418 if left.arity() <= left_key {
3420 return Err(XlogError::Kernel(format!(
3421 "sort_merge: left_key={} out of bounds (arity={})",
3422 left_key,
3423 left.arity()
3424 )));
3425 }
3426 if right.arity() <= right_key {
3427 return Err(XlogError::Kernel(format!(
3428 "sort_merge: right_key={} out of bounds (arity={})",
3429 right_key,
3430 right.arity()
3431 )));
3432 }
3433 let lt = left.schema().column_type(left_key);
3434 let rt = right.schema().column_type(right_key);
3435 if lt != rt || !matches!(lt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3436 return Err(XlogError::Kernel(format!(
3437 "sort_merge: key types must be equal U32/Symbol; got left={:?} right={:?}",
3438 lt, rt
3439 )));
3440 }
3441 let left_col = left
3442 .column(left_key)
3443 .ok_or_else(|| XlogError::Kernel(format!("sort_merge: left.column({})", left_key)))?;
3444 let right_col = right
3445 .column(right_key)
3446 .ok_or_else(|| XlogError::Kernel(format!("sort_merge: right.column({})", right_key)))?;
3447 let required_left_bytes = num_left
3448 .checked_mul(4)
3449 .ok_or_else(|| XlogError::Kernel("sort_merge: left byte overflow".into()))?;
3450 let required_right_bytes = num_right
3451 .checked_mul(4)
3452 .ok_or_else(|| XlogError::Kernel("sort_merge: right byte overflow".into()))?;
3453 if left_col.num_bytes() < required_left_bytes {
3454 return Err(XlogError::Kernel(format!(
3455 "sort_merge: left key column has {} bytes; \
3456 require at least {} ({} rows × 4)",
3457 left_col.num_bytes(),
3458 required_left_bytes,
3459 num_left
3460 )));
3461 }
3462 if right_col.num_bytes() < required_right_bytes {
3463 return Err(XlogError::Kernel(format!(
3464 "sort_merge: right key column has {} bytes; \
3465 require at least {} ({} rows × 4)",
3466 right_col.num_bytes(),
3467 required_right_bytes,
3468 num_right
3469 )));
3470 }
3471
3472 let upper_bound: u64 = (num_left as u64)
3474 .checked_mul(num_right as u64)
3475 .ok_or_else(|| XlogError::Kernel("sort_merge: row-count product overflow".into()))?;
3476 if upper_bound > NESTED_LOOP_TOTAL_THRESHOLD {
3477 return Err(XlogError::Kernel(format!(
3478 "sort_merge: caller violated eligibility threshold: \
3479 num_left * num_right = {} > {} (NESTED_LOOP_TOTAL_THRESHOLD)",
3480 upper_bound, NESTED_LOOP_TOTAL_THRESHOLD
3481 )));
3482 }
3483
3484 let upper_bound_usize = upper_bound as usize;
3486 let mut d_output_left_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3487 let mut d_output_right_idx = self.memory.alloc::<u32>(upper_bound_usize)?;
3488 let mut d_output_count = self.memory.alloc::<u32>(1)?;
3489 self.device
3490 .inner()
3491 .memset_zeros(&mut d_output_count)
3492 .map_err(|e| XlogError::Kernel(format!("sort_merge: counter zero: {}", e)))?;
3493
3494 let func = self
3496 .device
3497 .inner()
3498 .get_func(
3499 JOIN_MODULE,
3500 join_kernels::SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS,
3501 )
3502 .ok_or_else(|| {
3503 XlogError::Kernel("sort_merge_join_inner_u32_1key_pairs kernel not found".into())
3504 })?;
3505
3506 let num_left_u32 = num_left as u32;
3507 let num_right_u32 = num_right as u32;
3508 let upper_bound_u32 = upper_bound as u32;
3509 let block_size = 256u32;
3510 let grid_size = num_left_u32.div_ceil(block_size);
3511 let config = LaunchConfig {
3512 grid_dim: (grid_size, 1, 1),
3513 block_dim: (block_size, 1, 1),
3514 shared_mem_bytes: 0,
3515 };
3516
3517 unsafe {
3530 func.clone()
3531 .launch(
3532 config,
3533 (
3534 left_col,
3535 right_col,
3536 num_left_u32,
3537 num_right_u32,
3538 &mut d_output_left_idx,
3539 &mut d_output_right_idx,
3540 &mut d_output_count,
3541 upper_bound_u32,
3542 ),
3543 )
3544 .map_err(|e| XlogError::Kernel(format!("sort_merge launch: {}", e)))?;
3545 }
3546
3547 self.device.synchronize()?;
3548
3549 let output_rows = self.dtoh_scalar_untracked(&d_output_count, 0)?;
3551 if (output_rows as u64) > upper_bound {
3554 return Err(XlogError::Kernel(format!(
3555 "sort_merge: kernel reported {} output rows > upper_bound {}",
3556 output_rows, upper_bound
3557 )));
3558 }
3559
3560 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left_idx, output_rows)?;
3562 let gathered_right =
3563 self.gather_buffer_by_indices(right, &d_output_right_idx, output_rows)?;
3564
3565 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3567 let mut result_columns = Vec::with_capacity(combined_schema.arity());
3568 result_columns.extend(gathered_left.columns);
3569 result_columns.extend(gathered_right.columns);
3570 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
3571 }
3572
3573 pub fn sort_merge_join_v2_inner_u32_1key_bounded(
3583 &self,
3584 left: &CudaBuffer,
3585 right: &CudaBuffer,
3586 left_key: usize,
3587 right_key: usize,
3588 output_capacity: usize,
3589 ) -> Result<CudaBuffer> {
3590 let num_left = self.device_row_count(left)?;
3591 let num_right = self.device_row_count(right)?;
3592
3593 if num_left == 0 || num_right == 0 {
3594 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3595 return self.create_empty_buffer(combined_schema);
3596 }
3597 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
3598 return Err(XlogError::Kernel(format!(
3599 "sort_merge_bounded: row counts exceed u32 surface: left={} right={}",
3600 num_left, num_right
3601 )));
3602 }
3603 if output_capacity == 0 || output_capacity > u32::MAX as usize {
3604 return Err(XlogError::Kernel(format!(
3605 "sort_merge_bounded: invalid output capacity {}",
3606 output_capacity
3607 )));
3608 }
3609 if left.arity() <= left_key {
3610 return Err(XlogError::Kernel(format!(
3611 "sort_merge_bounded: left_key={} out of bounds (arity={})",
3612 left_key,
3613 left.arity()
3614 )));
3615 }
3616 if right.arity() <= right_key {
3617 return Err(XlogError::Kernel(format!(
3618 "sort_merge_bounded: right_key={} out of bounds (arity={})",
3619 right_key,
3620 right.arity()
3621 )));
3622 }
3623 let lt = left.schema().column_type(left_key);
3624 let rt = right.schema().column_type(right_key);
3625 if lt != rt || !matches!(lt, Some(ScalarType::U32) | Some(ScalarType::Symbol)) {
3626 return Err(XlogError::Kernel(format!(
3627 "sort_merge_bounded: key types must be equal U32/Symbol; got left={:?} right={:?}",
3628 lt, rt
3629 )));
3630 }
3631
3632 let left_col = left.column(left_key).ok_or_else(|| {
3633 XlogError::Kernel(format!("sort_merge_bounded: left.column({})", left_key))
3634 })?;
3635 let right_col = right.column(right_key).ok_or_else(|| {
3636 XlogError::Kernel(format!("sort_merge_bounded: right.column({})", right_key))
3637 })?;
3638 let required_left_bytes = num_left
3639 .checked_mul(4)
3640 .ok_or_else(|| XlogError::Kernel("sort_merge_bounded: left byte overflow".into()))?;
3641 let required_right_bytes = num_right
3642 .checked_mul(4)
3643 .ok_or_else(|| XlogError::Kernel("sort_merge_bounded: right byte overflow".into()))?;
3644 if left_col.num_bytes() < required_left_bytes {
3645 return Err(XlogError::Kernel(format!(
3646 "sort_merge_bounded: left key column has {} bytes; require at least {}",
3647 left_col.num_bytes(),
3648 required_left_bytes
3649 )));
3650 }
3651 if right_col.num_bytes() < required_right_bytes {
3652 return Err(XlogError::Kernel(format!(
3653 "sort_merge_bounded: right key column has {} bytes; require at least {}",
3654 right_col.num_bytes(),
3655 required_right_bytes
3656 )));
3657 }
3658
3659 let mut d_output_left_idx = self.memory.alloc::<u32>(output_capacity)?;
3660 let mut d_output_right_idx = self.memory.alloc::<u32>(output_capacity)?;
3661 let mut d_output_count = self.memory.alloc::<u32>(1)?;
3662 self.device
3663 .inner()
3664 .memset_zeros(&mut d_output_count)
3665 .map_err(|e| XlogError::Kernel(format!("sort_merge_bounded: counter zero: {}", e)))?;
3666
3667 let func = self
3668 .device
3669 .inner()
3670 .get_func(
3671 JOIN_MODULE,
3672 join_kernels::SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS,
3673 )
3674 .ok_or_else(|| {
3675 XlogError::Kernel("sort_merge_join_inner_u32_1key_pairs kernel not found".into())
3676 })?;
3677
3678 let num_left_u32 = num_left as u32;
3679 let num_right_u32 = num_right as u32;
3680 let output_capacity_u32 = output_capacity as u32;
3681 let block_size = 256u32;
3682 let grid_size = num_left_u32.div_ceil(block_size);
3683 let config = LaunchConfig {
3684 grid_dim: (grid_size, 1, 1),
3685 block_dim: (block_size, 1, 1),
3686 shared_mem_bytes: 0,
3687 };
3688
3689 unsafe {
3690 func.clone()
3691 .launch(
3692 config,
3693 (
3694 left_col,
3695 right_col,
3696 num_left_u32,
3697 num_right_u32,
3698 &mut d_output_left_idx,
3699 &mut d_output_right_idx,
3700 &mut d_output_count,
3701 output_capacity_u32,
3702 ),
3703 )
3704 .map_err(|e| XlogError::Kernel(format!("sort_merge_bounded launch: {}", e)))?;
3705 }
3706
3707 self.device.synchronize()?;
3708 let output_rows = self.dtoh_scalar_untracked(&d_output_count, 0)?;
3709 if output_rows as usize > output_capacity {
3710 return Err(XlogError::Kernel(format!(
3711 "sort_merge_bounded: output {} exceeded bounded capacity {}",
3712 output_rows, output_capacity
3713 )));
3714 }
3715
3716 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left_idx, output_rows)?;
3717 let gathered_right =
3718 self.gather_buffer_by_indices(right, &d_output_right_idx, output_rows)?;
3719
3720 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3721 let mut result_columns = Vec::with_capacity(combined_schema.arity());
3722 result_columns.extend(gathered_left.columns);
3723 result_columns.extend(gathered_right.columns);
3724 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
3725 }
3726
3727 pub fn build_join_index_v2(
3729 &self,
3730 right: &CudaBuffer,
3731 right_keys: &[usize],
3732 ) -> Result<JoinIndexV2> {
3733 let num_right = self.device_row_count(right)?;
3734 if num_right == 0 {
3735 return Err(XlogError::Kernel(
3736 "Cannot build join index for empty relation".to_string(),
3737 ));
3738 }
3739 if num_right > u32::MAX as usize {
3740 return Err(XlogError::Kernel(format!(
3741 "Join index supports at most {} rows, got {}",
3742 u32::MAX,
3743 num_right
3744 )));
3745 }
3746 if right_keys.is_empty() {
3747 return Err(XlogError::Kernel(
3748 "Join requires at least one key column".to_string(),
3749 ));
3750 }
3751 for &k in right_keys {
3752 if k >= right.arity() {
3753 return Err(XlogError::Kernel(format!(
3754 "Right key column index {} out of bounds (arity {})",
3755 k,
3756 right.arity()
3757 )));
3758 }
3759 }
3760
3761 let num_right = num_right as u32;
3762 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
3763 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
3764
3765 Ok(JoinIndexV2 {
3766 right_num_rows: num_right,
3767 right_keys: right_keys.to_vec(),
3768 key_bytes: right_packed.key_bytes,
3769 packed_keys: right_packed.packed_keys,
3770 table,
3771 })
3772 }
3773
3774 pub fn build_join_index_v2_background(
3781 &self,
3782 right: &CudaBuffer,
3783 right_keys: &[usize],
3784 ) -> Result<JoinIndexV2> {
3785 if Self::use_recorded_hash_join_env()
3786 && !right_keys.is_empty()
3787 && right_keys.len() <= 4
3788 && right.num_rows() > 0
3789 {
3790 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
3791 return self.build_join_index_v2_recorded(right, right_keys, launch_stream);
3792 }
3793 }
3794
3795 self.build_join_index_v2(right, right_keys)
3796 }
3797
3798 pub fn build_join_index_v2_recorded(
3804 &self,
3805 right: &CudaBuffer,
3806 right_keys: &[usize],
3807 launch_stream: StreamId,
3808 ) -> Result<JoinIndexV2> {
3809 let runtime = self.memory.runtime().ok_or_else(|| {
3810 XlogError::Kernel(
3811 "build_join_index_v2_recorded requires a runtime-backed GpuMemoryManager"
3812 .to_string(),
3813 )
3814 })?;
3815 let cu_stream = runtime
3816 .stream_pool()
3817 .resolve(launch_stream)
3818 .ok_or_else(|| {
3819 XlogError::Kernel(format!(
3820 "build_join_index_v2_recorded: launch_stream StreamId({}) does not resolve",
3821 launch_stream.0
3822 ))
3823 })?;
3824
3825 let num_right = self.device_row_count(right)?;
3826 if num_right == 0 {
3827 return Err(XlogError::Kernel(
3828 "Cannot build join index for empty relation".to_string(),
3829 ));
3830 }
3831 if num_right > u32::MAX as usize {
3832 return Err(XlogError::Kernel(format!(
3833 "Join index supports at most {} rows, got {}",
3834 u32::MAX,
3835 num_right
3836 )));
3837 }
3838 if right_keys.is_empty() {
3839 return Err(XlogError::Kernel(
3840 "Join requires at least one key column".to_string(),
3841 ));
3842 }
3843 if right_keys.len() > 4 {
3844 return Err(XlogError::Kernel(
3845 "build_join_index_v2_recorded: max 4 key columns supported".to_string(),
3846 ));
3847 }
3848 for &k in right_keys {
3849 if k >= right.arity() {
3850 return Err(XlogError::Kernel(format!(
3851 "Right key column index {} out of bounds (arity {})",
3852 k,
3853 right.arity()
3854 )));
3855 }
3856 }
3857
3858 let num_right = num_right as u32;
3859 let right_packed =
3860 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
3861 let table = self.build_hash_table_v2_on_stream(
3862 &right_packed.hashes,
3863 num_right,
3864 &cu_stream,
3865 launch_stream,
3866 runtime,
3867 )?;
3868
3869 Ok(JoinIndexV2 {
3870 right_num_rows: num_right,
3871 right_keys: right_keys.to_vec(),
3872 key_bytes: right_packed.key_bytes,
3873 packed_keys: right_packed.packed_keys,
3874 table,
3875 })
3876 }
3877
3878 #[allow(clippy::too_many_arguments)]
3882 pub fn hash_join_v2_with_index(
3883 &self,
3884 left: &CudaBuffer,
3885 right: &CudaBuffer,
3886 left_keys: &[usize],
3887 right_keys: &[usize],
3888 join_type: JoinType,
3889 index: &JoinIndexV2,
3890 max_output: Option<usize>,
3891 ) -> Result<CudaBuffer> {
3892 if Self::use_recorded_hash_join_env()
3895 && !left_keys.is_empty()
3896 && left_keys.len() == right_keys.len()
3897 && left_keys.len() <= 4
3898 {
3899 if let Some(launch_stream) = self.recorded_op_stream_or_init() {
3900 return self.hash_join_v2_with_index_recorded(
3901 left,
3902 right,
3903 left_keys,
3904 right_keys,
3905 join_type,
3906 index,
3907 max_output,
3908 launch_stream,
3909 );
3910 }
3911 }
3912 let left_rows = self.device_row_count(left)?;
3913 let right_rows = self.device_row_count(right)?;
3914 if left_rows > u32::MAX as usize || right_rows > u32::MAX as usize {
3915 return Err(XlogError::Kernel(format!(
3916 "Join supports at most {} rows per side (left={}, right={})",
3917 u32::MAX,
3918 left_rows,
3919 right_rows
3920 )));
3921 }
3922
3923 if left_rows == 0 {
3925 return match join_type {
3926 JoinType::Inner | JoinType::LeftOuter => {
3927 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3928 self.create_empty_buffer(combined_schema)
3929 }
3930 JoinType::Semi | JoinType::Anti => self.create_empty_buffer(left.schema().clone()),
3931 };
3932 }
3933 if right_rows == 0 {
3934 return match join_type {
3935 JoinType::Inner => {
3936 let combined_schema = self.combine_schemas(left.schema(), right.schema());
3937 self.create_empty_buffer(combined_schema)
3938 }
3939 JoinType::Semi => self.create_empty_buffer(left.schema().clone()),
3940 JoinType::Anti => self.clone_buffer(left),
3941 JoinType::LeftOuter => self.left_outer_with_nulls(left, right),
3942 };
3943 }
3944
3945 if left_keys.is_empty() || right_keys.is_empty() {
3947 return Err(XlogError::Kernel(
3948 "Join requires at least one key column".to_string(),
3949 ));
3950 }
3951 if left_keys.len() != right_keys.len() {
3952 return Err(XlogError::Kernel(
3953 "Left and right key columns must have same length".to_string(),
3954 ));
3955 }
3956 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
3957 if left_idx >= left.arity() {
3958 return Err(XlogError::Kernel(format!(
3959 "Left key column index {} out of bounds (arity {})",
3960 left_idx,
3961 left.arity()
3962 )));
3963 }
3964 if right_idx >= right.arity() {
3965 return Err(XlogError::Kernel(format!(
3966 "Right key column index {} out of bounds (arity {})",
3967 right_idx,
3968 right.arity()
3969 )));
3970 }
3971 let left_type = left.schema().column_type(left_idx);
3972 let right_type = right.schema().column_type(right_idx);
3973 if left_type != right_type {
3974 return Err(XlogError::Kernel(format!(
3975 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
3976 left_idx, left_type, right_idx, right_type
3977 )));
3978 }
3979 }
3980
3981 if index.right_num_rows != right_rows as u32 {
3983 return Err(XlogError::Kernel(
3984 "Join index row count does not match right relation".to_string(),
3985 ));
3986 }
3987 if index.right_keys.as_slice() != right_keys {
3988 return Err(XlogError::Kernel(
3989 "Join index key columns do not match requested right_keys".to_string(),
3990 ));
3991 }
3992
3993 match join_type {
3994 JoinType::Inner => {
3995 self.hash_join_inner_v2_indexed(left, right, left_keys, index, max_output)
3996 }
3997 JoinType::Semi => self.hash_join_semi_indexed(left, left_keys, index),
3998 JoinType::Anti => self.hash_join_anti_indexed(left, right, left_keys, index),
3999 JoinType::LeftOuter => {
4000 self.hash_join_left_outer_indexed(left, right, left_keys, index, max_output)
4001 }
4002 }
4003 }
4004
4005 fn pack_keys_gpu(&self, buffer: &CudaBuffer, key_cols: &[usize]) -> Result<PackedKeyData> {
4025 if key_cols.is_empty() {
4026 return Err(XlogError::Kernel(
4027 "pack_keys_gpu: no key columns specified".into(),
4028 ));
4029 }
4030 if key_cols.len() > 4 {
4031 return Err(XlogError::Kernel(
4032 "pack_keys_gpu: max 4 key columns supported".into(),
4033 ));
4034 }
4035
4036 let num_rows = self.device_row_count(buffer)?;
4037 if num_rows > u32::MAX as usize {
4038 return Err(XlogError::Kernel(format!(
4039 "pack_keys_gpu supports at most {} rows, got {}",
4040 u32::MAX,
4041 num_rows
4042 )));
4043 }
4044 let num_rows = num_rows as u32;
4045 if num_rows == 0 {
4046 return Ok(PackedKeyData {
4048 hashes: self.memory.alloc::<u64>(0)?,
4049 packed_keys: self.memory.alloc::<u8>(0)?,
4050 key_bytes: 0,
4051 });
4052 }
4053
4054 let mut col_sizes: Vec<u32> = Vec::with_capacity(key_cols.len());
4056 let mut row_size: u32 = 0;
4057 for &col_idx in key_cols {
4058 let col_type = buffer
4059 .schema()
4060 .column_type(col_idx)
4061 .ok_or_else(|| XlogError::Kernel(format!("Invalid column index: {}", col_idx)))?;
4062 let size = col_type.size_bytes() as u32;
4063 col_sizes.push(size);
4064 row_size += size;
4065 }
4066
4067 let packed_bytes = (num_rows as u64) * (row_size as u64);
4069 let packed_slice = self.memory.alloc::<u8>(packed_bytes as usize)?;
4070 let hash_slice = self.memory.alloc::<u64>(num_rows as usize)?;
4071
4072 let mut col_ptrs: [u64; 4] = [0; 4];
4075 for (i, &col_idx) in key_cols.iter().enumerate() {
4076 let col = buffer
4077 .column(col_idx)
4078 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
4079 col_ptrs[i] = *col.device_ptr();
4081 }
4082 let mut packed_col_sizes = 0u64;
4083 for (i, size) in col_sizes.iter().copied().enumerate() {
4084 if size > u16::MAX as u32 {
4085 return Err(XlogError::Kernel(format!(
4086 "pack_keys_gpu: column element size {} exceeds 16-bit kernel argument",
4087 size
4088 )));
4089 }
4090 packed_col_sizes |= (size as u64) << (i * 16);
4091 }
4092
4093 let func = self
4095 .device
4096 .inner()
4097 .get_func(PACK_MODULE, pack_kernels::PACK_AND_HASH_KEYS)
4098 .ok_or_else(|| XlogError::Kernel("pack_and_hash_keys kernel not found".to_string()))?;
4099
4100 let block_size = 256u32;
4102 let grid_size = num_rows.div_ceil(block_size);
4103 let config = LaunchConfig {
4104 grid_dim: (grid_size, 1, 1),
4105 block_dim: (block_size, 1, 1),
4106 shared_mem_bytes: 0,
4107 };
4108
4109 unsafe {
4116 func.clone()
4117 .launch(
4118 config,
4119 (
4120 col_ptrs[0],
4121 col_ptrs[1],
4122 col_ptrs[2],
4123 col_ptrs[3],
4124 packed_col_sizes,
4125 key_cols.len() as u32,
4126 num_rows,
4127 row_size,
4128 &packed_slice,
4129 &hash_slice,
4130 ),
4131 )
4132 .map_err(|e| {
4133 XlogError::Kernel(format!("pack_and_hash_keys launch failed: {}", e))
4134 })?;
4135 }
4136
4137 self.device.synchronize()?;
4138
4139 Ok(PackedKeyData {
4140 hashes: hash_slice,
4141 packed_keys: packed_slice,
4142 key_bytes: row_size,
4143 })
4144 }
4145
4146 fn pack_keys_gpu_generic(
4148 &self,
4149 buffer: &CudaBuffer,
4150 key_cols: &[usize],
4151 ) -> Result<PackedKeyData> {
4152 if key_cols.is_empty() {
4153 return Err(XlogError::Kernel(
4154 "pack_keys_gpu_generic: no key columns specified".into(),
4155 ));
4156 }
4157
4158 let num_rows = self.device_row_count(buffer)?;
4159 if num_rows > u32::MAX as usize {
4160 return Err(XlogError::Kernel(format!(
4161 "pack_keys_gpu_generic supports at most {} rows, got {}",
4162 u32::MAX,
4163 num_rows
4164 )));
4165 }
4166 let num_rows = num_rows as u32;
4167 if num_rows == 0 {
4168 return Ok(PackedKeyData {
4169 hashes: self.memory.alloc::<u64>(0)?,
4170 packed_keys: self.memory.alloc::<u8>(0)?,
4171 key_bytes: 0,
4172 });
4173 }
4174
4175 let mut col_sizes: Vec<u32> = Vec::with_capacity(key_cols.len());
4176 let mut col_ptrs: Vec<u64> = Vec::with_capacity(key_cols.len());
4177 let mut row_size: u32 = 0;
4178
4179 for &col_idx in key_cols {
4180 let col_type = buffer
4181 .schema()
4182 .column_type(col_idx)
4183 .ok_or_else(|| XlogError::Kernel(format!("Invalid column index: {}", col_idx)))?;
4184 let size = col_type.size_bytes() as u32;
4185 row_size = row_size
4186 .checked_add(size)
4187 .ok_or_else(|| XlogError::Kernel("Row size overflow".to_string()))?;
4188 col_sizes.push(size);
4189
4190 let col = buffer
4191 .column(col_idx)
4192 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
4193 col_ptrs.push(*col.device_ptr());
4194 }
4195
4196 let packed_bytes = (num_rows as u64)
4197 .checked_mul(row_size as u64)
4198 .ok_or_else(|| XlogError::Kernel("Packed key byte size overflow".to_string()))?;
4199 let packed_slice = self.memory.alloc::<u8>(packed_bytes as usize)?;
4200 let hash_slice = self.memory.alloc::<u64>(num_rows as usize)?;
4201
4202 let mut d_col_sizes = self.memory.alloc::<u32>(col_sizes.len())?;
4203 self.htod_sync_copy_into_tracked(&col_sizes, &mut d_col_sizes)
4204 .map_err(|e| XlogError::Kernel(format!("Failed to upload col_sizes: {}", e)))?;
4205
4206 let mut d_col_ptrs = self.memory.alloc::<u64>(col_ptrs.len())?;
4207 self.htod_sync_copy_into_tracked(&col_ptrs, &mut d_col_ptrs)
4208 .map_err(|e| XlogError::Kernel(format!("Failed to upload col_ptrs: {}", e)))?;
4209
4210 let func = self
4211 .device
4212 .inner()
4213 .get_func(PACK_MODULE, pack_kernels::PACK_AND_HASH_KEYS_GENERIC)
4214 .ok_or_else(|| {
4215 XlogError::Kernel("pack_and_hash_keys_generic kernel not found".to_string())
4216 })?;
4217
4218 let block_size = 256u32;
4219 let grid_size = num_rows.div_ceil(block_size);
4220 let config = LaunchConfig {
4221 grid_dim: (grid_size, 1, 1),
4222 block_dim: (block_size, 1, 1),
4223 shared_mem_bytes: 0,
4224 };
4225
4226 unsafe {
4228 func.clone()
4229 .launch(
4230 config,
4231 (
4232 &d_col_ptrs,
4233 &d_col_sizes,
4234 key_cols.len() as u32,
4235 num_rows,
4236 row_size,
4237 &packed_slice,
4238 &hash_slice,
4239 ),
4240 )
4241 .map_err(|e| {
4242 XlogError::Kernel(format!("pack_and_hash_keys_generic launch failed: {}", e))
4243 })?;
4244 }
4245
4246 self.device.synchronize()?;
4247
4248 Ok(PackedKeyData {
4249 hashes: hash_slice,
4250 packed_keys: packed_slice,
4251 key_bytes: row_size,
4252 })
4253 }
4254
4255 pub(super) fn compute_hashes_and_pack_keys(
4263 &self,
4264 buffer: &CudaBuffer,
4265 key_cols: &[usize],
4266 ) -> Result<PackedKeyData> {
4267 if key_cols.is_empty() {
4268 return Err(XlogError::Kernel(
4269 "compute_hashes_and_pack_keys: no key columns specified".to_string(),
4270 ));
4271 }
4272
4273 if key_cols.len() <= 4 {
4274 self.pack_keys_gpu(buffer, key_cols)
4275 } else {
4276 self.pack_keys_gpu_generic(buffer, key_cols)
4277 }
4278 }
4279
4280 fn build_hash_table_v2(
4285 &self,
4286 hashes: &cudarc::driver::CudaSlice<u64>,
4287 num_rows: u32,
4288 ) -> Result<JoinHashTableV2> {
4289 let device = self.device.inner();
4290
4291 let target = (num_rows as u64).saturating_mul(2).max(1024);
4293 let num_buckets_u64 = target.next_power_of_two();
4294 let num_buckets = u32::try_from(num_buckets_u64).map_err(|_| {
4295 XlogError::Kernel(format!(
4296 "Join hash table too large: num_buckets={}",
4297 num_buckets_u64
4298 ))
4299 })?;
4300 let bucket_mask = num_buckets
4301 .checked_sub(1)
4302 .ok_or_else(|| XlogError::Kernel("Join hash table size underflow".to_string()))?;
4303
4304 let mut bucket_counts = self.memory.alloc::<u32>(num_buckets as usize)?;
4305 if num_buckets > 0 {
4306 device
4307 .memset_zeros(&mut bucket_counts)
4308 .map_err(|e| XlogError::Kernel(format!("Failed to zero bucket_counts: {}", e)))?;
4309 self.device.synchronize()?;
4310 }
4311
4312 let block_size = 256u32;
4313 let grid_size = num_rows.div_ceil(block_size);
4314 let config = LaunchConfig {
4315 grid_dim: (grid_size, 1, 1),
4316 block_dim: (block_size, 1, 1),
4317 shared_mem_bytes: 0,
4318 };
4319
4320 let count_fn = device
4321 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUCKET_COUNT_V2)
4322 .ok_or_else(|| {
4323 XlogError::Kernel("hash_join_bucket_count_v2 kernel not found".to_string())
4324 })?;
4325
4326 unsafe {
4328 count_fn
4329 .clone()
4330 .launch(config, (hashes, num_rows, &bucket_counts, bucket_mask))
4331 .map_err(|e| {
4332 XlogError::Kernel(format!("hash_join_bucket_count_v2 failed: {}", e))
4333 })?;
4334 }
4335 self.device.synchronize()?;
4336
4337 let mut bucket_offsets = self.memory.alloc::<u32>(num_buckets as usize)?;
4339 if num_buckets > 0 {
4340 device
4341 .dtod_copy(&bucket_counts, &mut bucket_offsets)
4342 .map_err(|e| XlogError::Kernel(format!("Failed to copy bucket_counts: {}", e)))?;
4343 self.device.synchronize()?;
4344 self.multiblock_scan_u32_inplace(&mut bucket_offsets, num_buckets)?;
4345 self.device.synchronize()?;
4346 }
4347
4348 let mut bucket_cursors = self.memory.alloc::<u32>(num_buckets as usize)?;
4350 if num_buckets > 0 {
4351 device
4352 .dtod_copy(&bucket_offsets, &mut bucket_cursors)
4353 .map_err(|e| XlogError::Kernel(format!("Failed to copy bucket_offsets: {}", e)))?;
4354 self.device.synchronize()?;
4355 }
4356
4357 let bucket_entries = self.memory.alloc::<u32>(num_rows as usize)?;
4358 let bucket_entry_hashes = self.memory.alloc::<u64>(num_rows as usize)?;
4359
4360 let scatter_fn = device
4361 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SCATTER_V2)
4362 .ok_or_else(|| {
4363 XlogError::Kernel("hash_join_scatter_v2 kernel not found".to_string())
4364 })?;
4365
4366 unsafe {
4368 scatter_fn
4369 .clone()
4370 .launch(
4371 config,
4372 (
4373 hashes,
4374 num_rows,
4375 &bucket_cursors,
4376 bucket_mask,
4377 &bucket_entries,
4378 &bucket_entry_hashes,
4379 ),
4380 )
4381 .map_err(|e| XlogError::Kernel(format!("hash_join_scatter_v2 failed: {}", e)))?;
4382 }
4383
4384 self.device.synchronize()?;
4385 Ok(JoinHashTableV2 {
4386 bucket_counts,
4387 bucket_offsets,
4388 bucket_entries,
4389 bucket_entry_hashes,
4390 bucket_mask,
4391 })
4392 }
4393
4394 pub fn build_hash_table_u64(
4396 &self,
4397 hashes: &crate::memory::TrackedCudaSlice<u64>,
4398 num_rows: u32,
4399 ) -> Result<HashTableU64> {
4400 let JoinHashTableV2 {
4401 bucket_counts,
4402 bucket_offsets,
4403 bucket_entries,
4404 bucket_entry_hashes,
4405 bucket_mask,
4406 } = self.build_hash_table_v2(hashes, num_rows)?;
4407 Ok(HashTableU64 {
4408 bucket_counts,
4409 bucket_offsets,
4410 bucket_entries,
4411 bucket_entry_hashes,
4412 bucket_mask,
4413 })
4414 }
4415
4416 fn hash_join_inner_v2(
4418 &self,
4419 left: &CudaBuffer,
4420 right: &CudaBuffer,
4421 left_keys: &[usize],
4422 right_keys: &[usize],
4423 max_output: Option<usize>,
4424 ) -> Result<CudaBuffer> {
4425 let num_left = self.device_row_count(left)?;
4426 let num_right = self.device_row_count(right)?;
4427 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
4428 return Err(XlogError::Kernel(format!(
4429 "Join supports at most {} rows per side (left={}, right={})",
4430 u32::MAX,
4431 num_left,
4432 num_right
4433 )));
4434 }
4435
4436 if num_left == 0 || num_right == 0 {
4438 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4439 return self.create_empty_buffer(combined_schema);
4440 }
4441
4442 if left_keys.is_empty() || right_keys.is_empty() {
4444 return Err(XlogError::Kernel(
4445 "Join requires at least one key column".to_string(),
4446 ));
4447 }
4448 if left_keys.len() != right_keys.len() {
4449 return Err(XlogError::Kernel(
4450 "Left and right key columns must have same length".to_string(),
4451 ));
4452 }
4453
4454 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
4456 let left_type = left.schema().column_type(left_idx);
4457 let right_type = right.schema().column_type(right_idx);
4458 if left_type != right_type {
4459 return Err(XlogError::Kernel(format!(
4460 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
4461 left_idx, left_type, right_idx, right_type
4462 )));
4463 }
4464 }
4465
4466 let num_left = num_left as u32;
4467 let num_right = num_right as u32;
4468
4469 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
4471 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
4472
4473 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
4475
4476 let probe_func = self
4481 .device
4482 .inner()
4483 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
4484 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
4485
4486 let block_size = 256u32;
4487 let probe_grid = num_left.div_ceil(block_size);
4488 let probe_config = LaunchConfig {
4489 grid_dim: (probe_grid, 1, 1),
4490 block_dim: (block_size, 1, 1),
4491 shared_mem_bytes: 0,
4492 };
4493
4494 let mut d_count_only = self.memory.alloc::<u32>(1)?;
4495 self.device
4496 .inner()
4497 .memset_zeros(&mut d_count_only)
4498 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4499 self.device.synchronize()?;
4500 let d_dummy_left = self.memory.alloc::<u32>(1)?;
4501 let d_dummy_right = self.memory.alloc::<u32>(1)?;
4502 let max_output_count_only = 0u32;
4503
4504 unsafe {
4511 let mut params: Vec<*mut c_void> = vec![
4512 (&left_packed.hashes).as_kernel_param(),
4513 num_left.as_kernel_param(),
4514 (&table.bucket_offsets).as_kernel_param(),
4515 (&table.bucket_counts).as_kernel_param(),
4516 (&table.bucket_entries).as_kernel_param(),
4517 (&table.bucket_entry_hashes).as_kernel_param(),
4518 table.bucket_mask.as_kernel_param(),
4519 (&left_packed.packed_keys).as_kernel_param(),
4520 (&right_packed.packed_keys).as_kernel_param(),
4521 left_packed.key_bytes.as_kernel_param(),
4522 (&d_dummy_left).as_kernel_param(),
4523 (&d_dummy_right).as_kernel_param(),
4524 (&d_count_only).as_kernel_param(),
4525 max_output_count_only.as_kernel_param(),
4526 ];
4527 probe_func
4528 .clone()
4529 .launch(probe_config, &mut params)
4530 .map_err(|e| {
4531 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
4532 })?;
4533 }
4534
4535 self.device.synchronize()?;
4536
4537 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
4541 let requested = max_output
4542 .map(|limit| (limit as u64).min(full_count))
4543 .unwrap_or(full_count);
4544
4545 if requested == 0 {
4546 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4547 return self.create_empty_buffer(combined_schema);
4548 }
4549
4550 if requested > u32::MAX as u64 {
4551 return Err(XlogError::Kernel(format!(
4552 "Join produced {} rows which exceeds the u32 index limit",
4553 requested
4554 )));
4555 }
4556
4557 let max_output = requested as u32;
4559 let d_output_left = self.memory.alloc::<u32>(max_output as usize)?;
4560 let d_output_right = self.memory.alloc::<u32>(max_output as usize)?;
4561 let mut d_output_count = self.memory.alloc::<u32>(1)?;
4562 self.device
4563 .inner()
4564 .memset_zeros(&mut d_output_count)
4565 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4566 self.device.synchronize()?;
4567
4568 unsafe {
4575 let mut params: Vec<*mut c_void> = vec![
4576 (&left_packed.hashes).as_kernel_param(),
4577 num_left.as_kernel_param(),
4578 (&table.bucket_offsets).as_kernel_param(),
4579 (&table.bucket_counts).as_kernel_param(),
4580 (&table.bucket_entries).as_kernel_param(),
4581 (&table.bucket_entry_hashes).as_kernel_param(),
4582 table.bucket_mask.as_kernel_param(),
4583 (&left_packed.packed_keys).as_kernel_param(),
4584 (&right_packed.packed_keys).as_kernel_param(),
4585 left_packed.key_bytes.as_kernel_param(),
4586 (&d_output_left).as_kernel_param(),
4587 (&d_output_right).as_kernel_param(),
4588 (&d_output_count).as_kernel_param(),
4589 max_output.as_kernel_param(),
4590 ];
4591 probe_func
4592 .clone()
4593 .launch(probe_config, &mut params)
4594 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
4595 }
4596
4597 self.device.synchronize()?;
4598
4599 let result_count =
4606 (self.read_join_output_count_metadata(&d_output_count)? as u64).min(max_output as u64);
4607
4608 if result_count == 0 {
4609 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4610 return self.create_empty_buffer(combined_schema);
4611 }
4612
4613 let output_rows = result_count as u32;
4614
4615 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left, output_rows)?;
4617 let gathered_right = self.gather_buffer_by_indices(right, &d_output_right, output_rows)?;
4618
4619 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4620 let mut result_columns = Vec::with_capacity(combined_schema.arity());
4621 result_columns.extend(gathered_left.columns);
4622 result_columns.extend(gathered_right.columns);
4623
4624 self.buffer_from_columns(result_columns, result_count, combined_schema)
4625 }
4626
4627 fn hash_join_inner_v2_indexed(
4628 &self,
4629 left: &CudaBuffer,
4630 right: &CudaBuffer,
4631 left_keys: &[usize],
4632 index: &JoinIndexV2,
4633 max_output: Option<usize>,
4634 ) -> Result<CudaBuffer> {
4635 let num_left = self.device_row_count(left)?;
4636 let num_right = self.device_row_count(right)?;
4637 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
4638 return Err(XlogError::Kernel(format!(
4639 "Join supports at most {} rows per side (left={}, right={})",
4640 u32::MAX,
4641 num_left,
4642 num_right
4643 )));
4644 }
4645
4646 if num_left == 0 || num_right == 0 {
4648 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4649 return self.create_empty_buffer(combined_schema);
4650 }
4651
4652 let num_left = num_left as u32;
4653
4654 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
4656 if left_packed.key_bytes != index.key_bytes {
4657 return Err(XlogError::Kernel(
4658 "Join key byte width mismatch between probe and cached index".to_string(),
4659 ));
4660 }
4661
4662 let table = &index.table;
4663
4664 let probe_func = self
4666 .device
4667 .inner()
4668 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
4669 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
4670
4671 let block_size = 256u32;
4672 let probe_grid = num_left.div_ceil(block_size);
4673 let probe_config = LaunchConfig {
4674 grid_dim: (probe_grid, 1, 1),
4675 block_dim: (block_size, 1, 1),
4676 shared_mem_bytes: 0,
4677 };
4678
4679 let mut d_count_only = self.memory.alloc::<u32>(1)?;
4680 self.device
4681 .inner()
4682 .memset_zeros(&mut d_count_only)
4683 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4684 self.device.synchronize()?;
4685 let d_dummy_left = self.memory.alloc::<u32>(1)?;
4686 let d_dummy_right = self.memory.alloc::<u32>(1)?;
4687 let max_output_count_only = 0u32;
4688
4689 unsafe {
4691 let mut params: Vec<*mut c_void> = vec![
4692 (&left_packed.hashes).as_kernel_param(),
4693 num_left.as_kernel_param(),
4694 (&table.bucket_offsets).as_kernel_param(),
4695 (&table.bucket_counts).as_kernel_param(),
4696 (&table.bucket_entries).as_kernel_param(),
4697 (&table.bucket_entry_hashes).as_kernel_param(),
4698 table.bucket_mask.as_kernel_param(),
4699 (&left_packed.packed_keys).as_kernel_param(),
4700 (&index.packed_keys).as_kernel_param(),
4701 index.key_bytes.as_kernel_param(),
4702 (&d_dummy_left).as_kernel_param(),
4703 (&d_dummy_right).as_kernel_param(),
4704 (&d_count_only).as_kernel_param(),
4705 max_output_count_only.as_kernel_param(),
4706 ];
4707 probe_func
4708 .clone()
4709 .launch(probe_config, &mut params)
4710 .map_err(|e| {
4711 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
4712 })?;
4713 }
4714
4715 self.device.synchronize()?;
4716
4717 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
4721 let requested = max_output
4722 .map(|limit| (limit as u64).min(full_count))
4723 .unwrap_or(full_count);
4724
4725 if requested == 0 {
4726 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4727 return self.create_empty_buffer(combined_schema);
4728 }
4729
4730 if requested > u32::MAX as u64 {
4731 return Err(XlogError::Kernel(format!(
4732 "Join produced {} rows which exceeds the u32 index limit",
4733 requested
4734 )));
4735 }
4736
4737 let max_output = requested as u32;
4739 let d_output_left = self.memory.alloc::<u32>(max_output as usize)?;
4740 let d_output_right = self.memory.alloc::<u32>(max_output as usize)?;
4741 let mut d_output_count = self.memory.alloc::<u32>(1)?;
4742 self.device
4743 .inner()
4744 .memset_zeros(&mut d_output_count)
4745 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
4746 self.device.synchronize()?;
4747
4748 unsafe {
4750 let mut params: Vec<*mut c_void> = vec![
4751 (&left_packed.hashes).as_kernel_param(),
4752 num_left.as_kernel_param(),
4753 (&table.bucket_offsets).as_kernel_param(),
4754 (&table.bucket_counts).as_kernel_param(),
4755 (&table.bucket_entries).as_kernel_param(),
4756 (&table.bucket_entry_hashes).as_kernel_param(),
4757 table.bucket_mask.as_kernel_param(),
4758 (&left_packed.packed_keys).as_kernel_param(),
4759 (&index.packed_keys).as_kernel_param(),
4760 index.key_bytes.as_kernel_param(),
4761 (&d_output_left).as_kernel_param(),
4762 (&d_output_right).as_kernel_param(),
4763 (&d_output_count).as_kernel_param(),
4764 max_output.as_kernel_param(),
4765 ];
4766 probe_func
4767 .clone()
4768 .launch(probe_config, &mut params)
4769 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
4770 }
4771
4772 self.device.synchronize()?;
4773
4774 let result_count =
4777 (self.read_join_output_count_metadata(&d_output_count)? as u64).min(max_output as u64);
4778
4779 if result_count == 0 {
4780 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4781 return self.create_empty_buffer(combined_schema);
4782 }
4783
4784 let output_rows = result_count as u32;
4785
4786 let gathered_left = self.gather_buffer_by_indices(left, &d_output_left, output_rows)?;
4787 let gathered_right = self.gather_buffer_by_indices(right, &d_output_right, output_rows)?;
4788
4789 let combined_schema = self.combine_schemas(left.schema(), right.schema());
4790 let mut result_columns = Vec::with_capacity(combined_schema.arity());
4791 result_columns.extend(gathered_left.columns);
4792 result_columns.extend(gathered_right.columns);
4793
4794 self.buffer_from_columns(result_columns, result_count, combined_schema)
4795 }
4796
4797 fn hash_join_semi_impl(
4799 &self,
4800 left: &CudaBuffer,
4801 right: &CudaBuffer,
4802 left_keys: &[usize],
4803 right_keys: &[usize],
4804 ) -> Result<CudaBuffer> {
4805 let num_left = self.device_row_count(left)?;
4806 let num_right = self.device_row_count(right)?;
4807 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
4808 return Err(XlogError::Kernel(format!(
4809 "Join supports at most {} rows per side (left={}, right={})",
4810 u32::MAX,
4811 num_left,
4812 num_right
4813 )));
4814 }
4815
4816 if num_left == 0 {
4818 return self.create_empty_buffer(left.schema().clone());
4819 }
4820 if num_right == 0 {
4821 return self.create_empty_buffer(left.schema().clone());
4823 }
4824
4825 if left_keys.is_empty() || right_keys.is_empty() {
4827 return Err(XlogError::Kernel(
4828 "Join requires at least one key column".to_string(),
4829 ));
4830 }
4831 if left_keys.len() != right_keys.len() {
4832 return Err(XlogError::Kernel(
4833 "Left and right key columns must have same length".to_string(),
4834 ));
4835 }
4836
4837 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
4839 let left_type = left.schema().column_type(left_idx);
4840 let right_type = right.schema().column_type(right_idx);
4841 if left_type != right_type {
4842 return Err(XlogError::Kernel(format!(
4843 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
4844 left_idx, left_type, right_idx, right_type
4845 )));
4846 }
4847 }
4848
4849 let num_left = num_left as u32;
4850 let num_right = num_right as u32;
4851
4852 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
4854 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
4855
4856 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
4858
4859 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
4861
4862 let semi_func = self
4864 .device
4865 .inner()
4866 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
4867 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
4868
4869 let block_size = 256u32;
4870 let grid_size = num_left.div_ceil(block_size);
4871 let config = LaunchConfig {
4872 grid_dim: (grid_size, 1, 1),
4873 block_dim: (block_size, 1, 1),
4874 shared_mem_bytes: 0,
4875 };
4876
4877 unsafe {
4882 semi_func
4883 .clone()
4884 .launch(
4885 config,
4886 (
4887 &left_packed.hashes,
4888 num_left,
4889 &table.bucket_offsets,
4890 &table.bucket_counts,
4891 &table.bucket_entries,
4892 &table.bucket_entry_hashes,
4893 table.bucket_mask,
4894 &left_packed.packed_keys,
4895 &right_packed.packed_keys,
4896 left_packed.key_bytes,
4897 &d_has_match,
4898 ),
4899 )
4900 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
4901 }
4902
4903 self.device.synchronize()?;
4904 self.filter_by_device_mask(left, &d_has_match)
4905 }
4906
4907 pub fn membership_mask_device(
4912 &self,
4913 probe: &CudaBuffer,
4914 build: &CudaBuffer,
4915 probe_keys: &[usize],
4916 build_keys: &[usize],
4917 ) -> Result<TrackedCudaSlice<u8>> {
4918 let num_probe = self.device_row_count(probe)?;
4919 let num_build = self.device_row_count(build)?;
4920
4921 if num_probe == 0 {
4923 return self.memory.alloc::<u8>(0);
4924 }
4925
4926 if num_build == 0 {
4928 let mut d_mask = self.memory.alloc::<u8>(num_probe)?;
4929 self.device.inner().memset_zeros(&mut d_mask).map_err(|e| {
4930 XlogError::Kernel(format!(
4931 "Failed to zero membership mask for empty build: {}",
4932 e
4933 ))
4934 })?;
4935 return Ok(d_mask);
4936 }
4937
4938 if num_probe > u32::MAX as usize || num_build > u32::MAX as usize {
4939 return Err(XlogError::Kernel(format!(
4940 "membership_mask supports at most {} rows per side (probe={}, build={})",
4941 u32::MAX,
4942 num_probe,
4943 num_build
4944 )));
4945 }
4946
4947 if probe_keys.is_empty() || build_keys.is_empty() {
4949 return Err(XlogError::Kernel(
4950 "membership_mask requires at least one key column".to_string(),
4951 ));
4952 }
4953 if probe_keys.len() != build_keys.len() {
4954 return Err(XlogError::Kernel(
4955 "Probe and build key columns must have same length".to_string(),
4956 ));
4957 }
4958
4959 for (&p_idx, &b_idx) in probe_keys.iter().zip(build_keys.iter()) {
4961 let p_type = probe.schema().column_type(p_idx);
4962 let b_type = build.schema().column_type(b_idx);
4963 if p_type != b_type {
4964 return Err(XlogError::Kernel(format!(
4965 "Key column type mismatch: probe[{}]={:?}, build[{}]={:?}",
4966 p_idx, p_type, b_idx, b_type
4967 )));
4968 }
4969 }
4970
4971 let num_probe_u32 = num_probe as u32;
4972 let num_build_u32 = num_build as u32;
4973
4974 let probe_packed = self.compute_hashes_and_pack_keys(probe, probe_keys)?;
4976 let build_packed = self.compute_hashes_and_pack_keys(build, build_keys)?;
4977
4978 let table = self.build_hash_table_v2(&build_packed.hashes, num_build_u32)?;
4980
4981 let d_has_match = self.memory.alloc::<u8>(num_probe)?;
4983
4984 let semi_func = self
4986 .device
4987 .inner()
4988 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
4989 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
4990
4991 let block_size = 256u32;
4992 let grid_size = num_probe_u32.div_ceil(block_size);
4993 let config = LaunchConfig {
4994 grid_dim: (grid_size, 1, 1),
4995 block_dim: (block_size, 1, 1),
4996 shared_mem_bytes: 0,
4997 };
4998
4999 unsafe {
5005 semi_func
5006 .clone()
5007 .launch(
5008 config,
5009 (
5010 &probe_packed.hashes,
5011 num_probe_u32,
5012 &table.bucket_offsets,
5013 &table.bucket_counts,
5014 &table.bucket_entries,
5015 &table.bucket_entry_hashes,
5016 table.bucket_mask,
5017 &probe_packed.packed_keys,
5018 &build_packed.packed_keys,
5019 probe_packed.key_bytes,
5020 &d_has_match,
5021 ),
5022 )
5023 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5024 }
5025
5026 Ok(d_has_match)
5027 }
5028
5029 pub fn membership_mask(
5039 &self,
5040 probe: &CudaBuffer,
5041 build: &CudaBuffer,
5042 probe_keys: &[usize],
5043 build_keys: &[usize],
5044 ) -> Result<Vec<bool>> {
5045 let d_has_match = self.membership_mask_device(probe, build, probe_keys, build_keys)?;
5046 let num_probe = d_has_match.len();
5047 if num_probe == 0 {
5048 return Ok(Vec::new());
5049 }
5050 let mut host_mask = vec![0u8; num_probe];
5051 self.dtoh_sync_copy_into_tracked(&d_has_match, &mut host_mask)?;
5052 Ok(host_mask.into_iter().map(|b| b != 0).collect())
5053 }
5054
5055 fn hash_join_semi_indexed(
5056 &self,
5057 left: &CudaBuffer,
5058 left_keys: &[usize],
5059 index: &JoinIndexV2,
5060 ) -> Result<CudaBuffer> {
5061 let num_left = self.device_row_count(left)?;
5062 if num_left > u32::MAX as usize {
5063 return Err(XlogError::Kernel(format!(
5064 "Join supports at most {} rows on left side (left={})",
5065 u32::MAX,
5066 num_left
5067 )));
5068 }
5069
5070 if num_left == 0 {
5072 return self.create_empty_buffer(left.schema().clone());
5073 }
5074 if index.right_num_rows == 0 {
5075 return self.create_empty_buffer(left.schema().clone());
5076 }
5077
5078 let num_left = num_left as u32;
5079
5080 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5081 if left_packed.key_bytes != index.key_bytes {
5082 return Err(XlogError::Kernel(
5083 "Join key byte width mismatch between probe and cached index".to_string(),
5084 ));
5085 }
5086
5087 let table = &index.table;
5088
5089 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
5091
5092 let semi_func = self
5093 .device
5094 .inner()
5095 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5096 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5097
5098 let block_size = 256u32;
5099 let grid_size = num_left.div_ceil(block_size);
5100 let config = LaunchConfig {
5101 grid_dim: (grid_size, 1, 1),
5102 block_dim: (block_size, 1, 1),
5103 shared_mem_bytes: 0,
5104 };
5105
5106 unsafe {
5108 semi_func
5109 .clone()
5110 .launch(
5111 config,
5112 (
5113 &left_packed.hashes,
5114 num_left,
5115 &table.bucket_offsets,
5116 &table.bucket_counts,
5117 &table.bucket_entries,
5118 &table.bucket_entry_hashes,
5119 table.bucket_mask,
5120 &left_packed.packed_keys,
5121 &index.packed_keys,
5122 index.key_bytes,
5123 &d_has_match,
5124 ),
5125 )
5126 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5127 }
5128
5129 self.device.synchronize()?;
5130 self.filter_by_device_mask(left, &d_has_match)
5131 }
5132
5133 fn hash_join_anti_impl(
5135 &self,
5136 left: &CudaBuffer,
5137 right: &CudaBuffer,
5138 left_keys: &[usize],
5139 right_keys: &[usize],
5140 ) -> Result<CudaBuffer> {
5141 let num_left = self.device_row_count(left)?;
5142 let num_right = self.device_row_count(right)?;
5143 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5144 return Err(XlogError::Kernel(format!(
5145 "Join supports at most {} rows per side (left={}, right={})",
5146 u32::MAX,
5147 num_left,
5148 num_right
5149 )));
5150 }
5151
5152 if num_left == 0 {
5154 return self.create_empty_buffer(left.schema().clone());
5155 }
5156 if num_right == 0 {
5157 return self.clone_buffer(left);
5159 }
5160
5161 if left_keys.is_empty() || right_keys.is_empty() {
5163 return Err(XlogError::Kernel(
5164 "Join requires at least one key column".to_string(),
5165 ));
5166 }
5167 if left_keys.len() != right_keys.len() {
5168 return Err(XlogError::Kernel(
5169 "Left and right key columns must have same length".to_string(),
5170 ));
5171 }
5172
5173 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
5175 let left_type = left.schema().column_type(left_idx);
5176 let right_type = right.schema().column_type(right_idx);
5177 if left_type != right_type {
5178 return Err(XlogError::Kernel(format!(
5179 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
5180 left_idx, left_type, right_idx, right_type
5181 )));
5182 }
5183 }
5184
5185 let num_left = num_left as u32;
5186 let num_right = num_right as u32;
5187
5188 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5190 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
5191
5192 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
5194
5195 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5197
5198 let anti_func = self
5200 .device
5201 .inner()
5202 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_ANTI)
5203 .ok_or_else(|| XlogError::Kernel("hash_join_anti kernel not found".to_string()))?;
5204
5205 let block_size = 256u32;
5206 let grid_size = num_left.div_ceil(block_size);
5207 let config = LaunchConfig {
5208 grid_dim: (grid_size, 1, 1),
5209 block_dim: (block_size, 1, 1),
5210 shared_mem_bytes: 0,
5211 };
5212
5213 unsafe {
5218 anti_func
5219 .clone()
5220 .launch(
5221 config,
5222 (
5223 &left_packed.hashes,
5224 num_left,
5225 &table.bucket_offsets,
5226 &table.bucket_counts,
5227 &table.bucket_entries,
5228 &table.bucket_entry_hashes,
5229 table.bucket_mask,
5230 &left_packed.packed_keys,
5231 &right_packed.packed_keys,
5232 left_packed.key_bytes,
5233 &d_no_match,
5234 ),
5235 )
5236 .map_err(|e| XlogError::Kernel(format!("hash_join_anti failed: {}", e)))?;
5237 }
5238
5239 self.device.synchronize()?;
5240 self.filter_by_device_mask(left, &d_no_match)
5241 }
5242
5243 fn hash_join_anti_indexed(
5244 &self,
5245 left: &CudaBuffer,
5246 right: &CudaBuffer,
5247 left_keys: &[usize],
5248 index: &JoinIndexV2,
5249 ) -> Result<CudaBuffer> {
5250 let num_left = self.device_row_count(left)?;
5251 let num_right = self.device_row_count(right)?;
5252 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5253 return Err(XlogError::Kernel(format!(
5254 "Join supports at most {} rows per side (left={}, right={})",
5255 u32::MAX,
5256 num_left,
5257 num_right
5258 )));
5259 }
5260 if num_left == 0 {
5261 return self.create_empty_buffer(left.schema().clone());
5262 }
5263 if num_right == 0 {
5264 return self.clone_buffer(left);
5265 }
5266
5267 let num_left = num_left as u32;
5268
5269 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5270 if left_packed.key_bytes != index.key_bytes {
5271 return Err(XlogError::Kernel(
5272 "Join key byte width mismatch between probe and cached index".to_string(),
5273 ));
5274 }
5275
5276 let table = &index.table;
5277
5278 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5279
5280 let anti_func = self
5281 .device
5282 .inner()
5283 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_ANTI)
5284 .ok_or_else(|| XlogError::Kernel("hash_join_anti kernel not found".to_string()))?;
5285
5286 let block_size = 256u32;
5287 let grid_size = num_left.div_ceil(block_size);
5288 let config = LaunchConfig {
5289 grid_dim: (grid_size, 1, 1),
5290 block_dim: (block_size, 1, 1),
5291 shared_mem_bytes: 0,
5292 };
5293
5294 unsafe {
5296 anti_func
5297 .clone()
5298 .launch(
5299 config,
5300 (
5301 &left_packed.hashes,
5302 num_left,
5303 &table.bucket_offsets,
5304 &table.bucket_counts,
5305 &table.bucket_entries,
5306 &table.bucket_entry_hashes,
5307 table.bucket_mask,
5308 &left_packed.packed_keys,
5309 &index.packed_keys,
5310 index.key_bytes,
5311 &d_no_match,
5312 ),
5313 )
5314 .map_err(|e| XlogError::Kernel(format!("hash_join_anti failed: {}", e)))?;
5315 }
5316
5317 self.device.synchronize()?;
5318 self.filter_by_device_mask(left, &d_no_match)
5319 }
5320
5321 fn hash_join_left_outer_indexed(
5322 &self,
5323 left: &CudaBuffer,
5324 right: &CudaBuffer,
5325 left_keys: &[usize],
5326 index: &JoinIndexV2,
5327 max_output: Option<usize>,
5328 ) -> Result<CudaBuffer> {
5329 let num_left = self.device_row_count(left)?;
5330 let num_right = self.device_row_count(right)?;
5331 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5332 return Err(XlogError::Kernel(format!(
5333 "Join supports at most {} rows per side (left={}, right={})",
5334 u32::MAX,
5335 num_left,
5336 num_right
5337 )));
5338 }
5339
5340 if num_left == 0 {
5342 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5343 return self.create_empty_buffer(combined_schema);
5344 }
5345 if num_right == 0 {
5347 return self.left_outer_with_nulls(left, right);
5348 }
5349
5350 let num_left = num_left as u32;
5351
5352 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5353 if left_packed.key_bytes != index.key_bytes {
5354 return Err(XlogError::Kernel(
5355 "Join key byte width mismatch between probe and cached index".to_string(),
5356 ));
5357 }
5358
5359 let table = &index.table;
5360
5361 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
5363
5364 let semi_func = self
5365 .device
5366 .inner()
5367 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5368 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5369
5370 let block_size = 256u32;
5371 let grid_size = num_left.div_ceil(block_size);
5372 let config = LaunchConfig {
5373 grid_dim: (grid_size, 1, 1),
5374 block_dim: (block_size, 1, 1),
5375 shared_mem_bytes: 0,
5376 };
5377
5378 unsafe {
5380 semi_func
5381 .clone()
5382 .launch(
5383 config,
5384 (
5385 &left_packed.hashes,
5386 num_left,
5387 &table.bucket_offsets,
5388 &table.bucket_counts,
5389 &table.bucket_entries,
5390 &table.bucket_entry_hashes,
5391 table.bucket_mask,
5392 &left_packed.packed_keys,
5393 &index.packed_keys,
5394 index.key_bytes,
5395 &d_has_match,
5396 ),
5397 )
5398 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5399 }
5400
5401 let probe_func = self
5402 .device
5403 .inner()
5404 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
5405 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
5406
5407 let mut d_count_only = self.memory.alloc::<u32>(1)?;
5409 self.device
5410 .inner()
5411 .memset_zeros(&mut d_count_only)
5412 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5413 let d_dummy_left = self.memory.alloc::<u32>(1)?;
5414 let d_dummy_right = self.memory.alloc::<u32>(1)?;
5415 let max_output_count_only = 0u32;
5416
5417 unsafe {
5419 let mut params: Vec<*mut c_void> = vec![
5420 (&left_packed.hashes).as_kernel_param(),
5421 num_left.as_kernel_param(),
5422 (&table.bucket_offsets).as_kernel_param(),
5423 (&table.bucket_counts).as_kernel_param(),
5424 (&table.bucket_entries).as_kernel_param(),
5425 (&table.bucket_entry_hashes).as_kernel_param(),
5426 table.bucket_mask.as_kernel_param(),
5427 (&left_packed.packed_keys).as_kernel_param(),
5428 (&index.packed_keys).as_kernel_param(),
5429 index.key_bytes.as_kernel_param(),
5430 (&d_dummy_left).as_kernel_param(),
5431 (&d_dummy_right).as_kernel_param(),
5432 (&d_count_only).as_kernel_param(),
5433 max_output_count_only.as_kernel_param(),
5434 ];
5435 probe_func
5436 .clone()
5437 .launch(config, &mut params)
5438 .map_err(|e| {
5439 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
5440 })?;
5441 }
5442
5443 self.device.synchronize()?;
5444
5445 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
5448 let requested_inner = max_output
5449 .map(|limit| (limit as u64).min(full_inner))
5450 .unwrap_or(full_inner);
5451
5452 if requested_inner > u32::MAX as u64 {
5453 return Err(XlogError::Kernel(format!(
5454 "Join produced {} rows which exceeds the u32 index limit",
5455 requested_inner
5456 )));
5457 }
5458
5459 let max_output = requested_inner as u32;
5460 let alloc_len = (requested_inner.max(1)) as usize;
5461 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
5462 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
5463 let mut d_output_count = self.memory.alloc::<u32>(1)?;
5464 self.device
5465 .inner()
5466 .memset_zeros(&mut d_output_count)
5467 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5468
5469 unsafe {
5471 let mut params: Vec<*mut c_void> = vec![
5472 (&left_packed.hashes).as_kernel_param(),
5473 num_left.as_kernel_param(),
5474 (&table.bucket_offsets).as_kernel_param(),
5475 (&table.bucket_counts).as_kernel_param(),
5476 (&table.bucket_entries).as_kernel_param(),
5477 (&table.bucket_entry_hashes).as_kernel_param(),
5478 table.bucket_mask.as_kernel_param(),
5479 (&left_packed.packed_keys).as_kernel_param(),
5480 (&index.packed_keys).as_kernel_param(),
5481 index.key_bytes.as_kernel_param(),
5482 (&d_output_left).as_kernel_param(),
5483 (&d_output_right).as_kernel_param(),
5484 (&d_output_count).as_kernel_param(),
5485 max_output.as_kernel_param(),
5486 ];
5487 probe_func
5488 .clone()
5489 .launch(config, &mut params)
5490 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
5491 }
5492
5493 self.device.synchronize()?;
5494
5495 let device = self.device.inner();
5496
5497 let inner_count = self
5502 .read_join_output_count_metadata(&d_output_count)?
5503 .min(max_output);
5504
5505 let mask_not_fn = device
5506 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
5507 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
5508
5509 let mut d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5510
5511 unsafe {
5513 mask_not_fn
5514 .clone()
5515 .launch(config, (&d_has_match, &mut d_no_match, num_left))
5516 }
5517 .map_err(|e| XlogError::Kernel(format!("mask_not failed: {}", e)))?;
5518
5519 let unmatched_left = self.filter_by_device_mask(left, &d_no_match)?;
5520
5521 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
5522 let total_rows = (inner_count as u64) + unmatched_rows;
5523
5524 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5525
5526 if total_rows == 0 {
5527 return self.create_empty_buffer(combined_schema);
5528 }
5529
5530 let inner_left = self.gather_buffer_by_indices(left, &d_output_left, inner_count)?;
5531 let inner_right = self.gather_buffer_by_indices(right, &d_output_right, inner_count)?;
5532
5533 if unmatched_rows == 0 {
5534 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5535 result_columns.extend(inner_left.columns);
5536 result_columns.extend(inner_right.columns);
5537 return self.buffer_from_columns(result_columns, inner_count as u64, combined_schema);
5538 }
5539
5540 if inner_count == 0 {
5541 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5542 result_columns.extend(unmatched_left.columns);
5543
5544 for col_idx in 0..right.arity() {
5545 let elem_size = right
5546 .schema()
5547 .column_type(col_idx)
5548 .map(|t| t.size_bytes())
5549 .unwrap_or(4);
5550
5551 let bytes = (unmatched_rows as usize)
5552 .checked_mul(elem_size)
5553 .ok_or_else(|| {
5554 XlogError::Kernel(
5555 "Left outer join: right column byte size overflow".to_string(),
5556 )
5557 })?;
5558
5559 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
5560 if bytes > 0 {
5561 device.memset_zeros(&mut dst_col).map_err(|e| {
5562 XlogError::Kernel(format!("Failed to zero null right column: {}", e))
5563 })?;
5564 }
5565 result_columns.push(dst_col.into());
5566 }
5567
5568 self.device.synchronize()?;
5569 return self.buffer_from_columns(result_columns, unmatched_rows, combined_schema);
5570 }
5571
5572 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5573 let inner_rows = inner_count as u64;
5574
5575 for (col_idx, (inner_col, unmatched_col)) in inner_left
5576 .columns
5577 .into_iter()
5578 .zip(unmatched_left.columns)
5579 .enumerate()
5580 {
5581 let elem_size = left
5582 .schema()
5583 .column_type(col_idx)
5584 .map(|t| t.size_bytes())
5585 .unwrap_or(4);
5586
5587 let inner_bytes = (inner_rows as usize)
5588 .checked_mul(elem_size)
5589 .ok_or_else(|| {
5590 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
5591 })?;
5592 let unmatched_bytes = (unmatched_rows as usize)
5593 .checked_mul(elem_size)
5594 .ok_or_else(|| {
5595 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
5596 })?;
5597 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
5598 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
5599 })?;
5600
5601 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
5602
5603 if inner_bytes > 0 {
5604 let mut out_view = out_col.slice_mut(0..inner_bytes);
5605 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
5606 XlogError::Kernel(format!("Failed to copy inner left column: {}", e))
5607 })?;
5608 }
5609 if unmatched_bytes > 0 {
5610 let mut out_view = out_col.slice_mut(inner_bytes..total_bytes);
5611 let unmatched_view = self.column_bytes_view(&unmatched_col, unmatched_bytes)?;
5612 device
5613 .dtod_copy(&unmatched_view, &mut out_view)
5614 .map_err(|e| {
5615 XlogError::Kernel(format!("Failed to copy unmatched left column: {}", e))
5616 })?;
5617 }
5618
5619 result_columns.push(out_col.into());
5620 }
5621
5622 for (col_idx, inner_col) in inner_right.columns.into_iter().enumerate() {
5623 let elem_size = right
5624 .schema()
5625 .column_type(col_idx)
5626 .map(|t| t.size_bytes())
5627 .unwrap_or(4);
5628
5629 let inner_bytes = (inner_rows as usize)
5630 .checked_mul(elem_size)
5631 .ok_or_else(|| {
5632 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
5633 })?;
5634 let unmatched_bytes = (unmatched_rows as usize)
5635 .checked_mul(elem_size)
5636 .ok_or_else(|| {
5637 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
5638 })?;
5639 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
5640 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
5641 })?;
5642
5643 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
5644
5645 if total_bytes > 0 {
5646 device.memset_zeros(&mut out_col).map_err(|e| {
5647 XlogError::Kernel(format!("Failed to zero right outer column: {}", e))
5648 })?;
5649 }
5650
5651 if inner_bytes > 0 {
5652 let mut out_view = out_col.slice_mut(0..inner_bytes);
5653 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
5654 XlogError::Kernel(format!("Failed to copy inner right column: {}", e))
5655 })?;
5656 }
5657
5658 result_columns.push(out_col.into());
5659 }
5660
5661 self.device.synchronize()?;
5662
5663 self.buffer_from_columns(result_columns, total_rows, combined_schema)
5664 }
5665
5666 fn hash_join_left_outer_impl(
5668 &self,
5669 left: &CudaBuffer,
5670 right: &CudaBuffer,
5671 left_keys: &[usize],
5672 right_keys: &[usize],
5673 max_output: Option<usize>,
5674 ) -> Result<CudaBuffer> {
5675 let num_left = self.device_row_count(left)?;
5676 let num_right = self.device_row_count(right)?;
5677 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
5678 return Err(XlogError::Kernel(format!(
5679 "Join supports at most {} rows per side (left={}, right={})",
5680 u32::MAX,
5681 num_left,
5682 num_right
5683 )));
5684 }
5685
5686 if num_left == 0 {
5688 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5689 return self.create_empty_buffer(combined_schema);
5690 }
5691
5692 if num_right == 0 {
5694 return self.left_outer_with_nulls(left, right);
5695 }
5696
5697 if left_keys.is_empty() || right_keys.is_empty() {
5699 return Err(XlogError::Kernel(
5700 "Join requires at least one key column".to_string(),
5701 ));
5702 }
5703 if left_keys.len() != right_keys.len() {
5704 return Err(XlogError::Kernel(
5705 "Left and right key columns must have same length".to_string(),
5706 ));
5707 }
5708
5709 for (&left_idx, &right_idx) in left_keys.iter().zip(right_keys.iter()) {
5711 let left_type = left.schema().column_type(left_idx);
5712 let right_type = right.schema().column_type(right_idx);
5713 if left_type != right_type {
5714 return Err(XlogError::Kernel(format!(
5715 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
5716 left_idx, left_type, right_idx, right_type
5717 )));
5718 }
5719 }
5720
5721 let num_left = num_left as u32;
5722 let num_right = num_right as u32;
5723
5724 let left_packed = self.compute_hashes_and_pack_keys(left, left_keys)?;
5726 let right_packed = self.compute_hashes_and_pack_keys(right, right_keys)?;
5727
5728 let table = self.build_hash_table_v2(&right_packed.hashes, num_right)?;
5730
5731 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
5733
5734 let semi_func = self
5736 .device
5737 .inner()
5738 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
5739 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
5740
5741 let block_size = 256u32;
5742 let grid_size = num_left.div_ceil(block_size);
5743 let config = LaunchConfig {
5744 grid_dim: (grid_size, 1, 1),
5745 block_dim: (block_size, 1, 1),
5746 shared_mem_bytes: 0,
5747 };
5748
5749 unsafe {
5754 semi_func
5755 .clone()
5756 .launch(
5757 config,
5758 (
5759 &left_packed.hashes,
5760 num_left,
5761 &table.bucket_offsets,
5762 &table.bucket_counts,
5763 &table.bucket_entries,
5764 &table.bucket_entry_hashes,
5765 table.bucket_mask,
5766 &left_packed.packed_keys,
5767 &right_packed.packed_keys,
5768 left_packed.key_bytes,
5769 &d_has_match,
5770 ),
5771 )
5772 .map_err(|e| XlogError::Kernel(format!("hash_join_semi failed: {}", e)))?;
5773 }
5774
5775 let probe_func = self
5776 .device
5777 .inner()
5778 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
5779 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
5780
5781 let mut d_count_only = self.memory.alloc::<u32>(1)?;
5783 self.device
5784 .inner()
5785 .memset_zeros(&mut d_count_only)
5786 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5787 let d_dummy_left = self.memory.alloc::<u32>(1)?;
5788 let d_dummy_right = self.memory.alloc::<u32>(1)?;
5789 let max_output_count_only = 0u32;
5790
5791 unsafe {
5793 let mut params: Vec<*mut c_void> = vec![
5794 (&left_packed.hashes).as_kernel_param(),
5795 num_left.as_kernel_param(),
5796 (&table.bucket_offsets).as_kernel_param(),
5797 (&table.bucket_counts).as_kernel_param(),
5798 (&table.bucket_entries).as_kernel_param(),
5799 (&table.bucket_entry_hashes).as_kernel_param(),
5800 table.bucket_mask.as_kernel_param(),
5801 (&left_packed.packed_keys).as_kernel_param(),
5802 (&right_packed.packed_keys).as_kernel_param(),
5803 left_packed.key_bytes.as_kernel_param(),
5804 (&d_dummy_left).as_kernel_param(),
5805 (&d_dummy_right).as_kernel_param(),
5806 (&d_count_only).as_kernel_param(),
5807 max_output_count_only.as_kernel_param(),
5808 ];
5809 probe_func
5810 .clone()
5811 .launch(config, &mut params)
5812 .map_err(|e| {
5813 XlogError::Kernel(format!("hash_join_probe_v2 (count) failed: {}", e))
5814 })?;
5815 }
5816
5817 self.device.synchronize()?;
5818
5819 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
5822 let requested_inner = max_output
5823 .map(|limit| (limit as u64).min(full_inner))
5824 .unwrap_or(full_inner);
5825
5826 if requested_inner > u32::MAX as u64 {
5827 return Err(XlogError::Kernel(format!(
5828 "Join produced {} rows which exceeds the u32 index limit",
5829 requested_inner
5830 )));
5831 }
5832
5833 let max_output = requested_inner as u32;
5834 let alloc_len = (requested_inner.max(1)) as usize;
5835 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
5836 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
5837 let mut d_output_count = self.memory.alloc::<u32>(1)?;
5838 self.device
5839 .inner()
5840 .memset_zeros(&mut d_output_count)
5841 .map_err(|e| XlogError::Kernel(format!("Failed to zero output count: {}", e)))?;
5842
5843 unsafe {
5850 let mut params: Vec<*mut c_void> = vec![
5851 (&left_packed.hashes).as_kernel_param(),
5852 num_left.as_kernel_param(),
5853 (&table.bucket_offsets).as_kernel_param(),
5854 (&table.bucket_counts).as_kernel_param(),
5855 (&table.bucket_entries).as_kernel_param(),
5856 (&table.bucket_entry_hashes).as_kernel_param(),
5857 table.bucket_mask.as_kernel_param(),
5858 (&left_packed.packed_keys).as_kernel_param(),
5859 (&right_packed.packed_keys).as_kernel_param(),
5860 left_packed.key_bytes.as_kernel_param(),
5861 (&d_output_left).as_kernel_param(),
5862 (&d_output_right).as_kernel_param(),
5863 (&d_output_count).as_kernel_param(),
5864 max_output.as_kernel_param(),
5865 ];
5866 probe_func
5867 .clone()
5868 .launch(config, &mut params)
5869 .map_err(|e| XlogError::Kernel(format!("hash_join_probe_v2 failed: {}", e)))?;
5870 }
5871
5872 self.device.synchronize()?;
5873
5874 let device = self.device.inner();
5875
5876 let inner_count = self
5881 .read_join_output_count_metadata(&d_output_count)?
5882 .min(max_output);
5883
5884 let mask_not_fn = device
5886 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
5887 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
5888
5889 let mut d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
5890
5891 unsafe {
5893 mask_not_fn
5894 .clone()
5895 .launch(config, (&d_has_match, &mut d_no_match, num_left))
5896 }
5897 .map_err(|e| XlogError::Kernel(format!("mask_not failed: {}", e)))?;
5898
5899 let unmatched_left = self.filter_by_device_mask(left, &d_no_match)?;
5900
5901 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
5902 let total_rows = (inner_count as u64) + unmatched_rows;
5903
5904 let combined_schema = self.combine_schemas(left.schema(), right.schema());
5905
5906 if total_rows == 0 {
5907 return self.create_empty_buffer(combined_schema);
5908 }
5909
5910 let inner_left = self.gather_buffer_by_indices(left, &d_output_left, inner_count)?;
5912 let inner_right = self.gather_buffer_by_indices(right, &d_output_right, inner_count)?;
5913
5914 if unmatched_rows == 0 {
5915 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5916 result_columns.extend(inner_left.columns);
5917 result_columns.extend(inner_right.columns);
5918 return self.buffer_from_columns(result_columns, inner_count as u64, combined_schema);
5919 }
5920
5921 if inner_count == 0 {
5922 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5923 result_columns.extend(unmatched_left.columns);
5924
5925 for col_idx in 0..right.arity() {
5926 let elem_size = right
5927 .schema()
5928 .column_type(col_idx)
5929 .map(|t| t.size_bytes())
5930 .unwrap_or(4);
5931
5932 let bytes = (unmatched_rows as usize)
5933 .checked_mul(elem_size)
5934 .ok_or_else(|| {
5935 XlogError::Kernel(
5936 "Left outer join: right column byte size overflow".to_string(),
5937 )
5938 })?;
5939
5940 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
5941 if bytes > 0 {
5942 device.memset_zeros(&mut dst_col).map_err(|e| {
5943 XlogError::Kernel(format!("Failed to zero null right column: {}", e))
5944 })?;
5945 }
5946 result_columns.push(dst_col.into());
5947 }
5948
5949 self.device.synchronize()?;
5950 return self.buffer_from_columns(result_columns, unmatched_rows, combined_schema);
5951 }
5952
5953 let mut result_columns = Vec::with_capacity(combined_schema.arity());
5955 let inner_rows = inner_count as u64;
5956
5957 for (col_idx, (inner_col, unmatched_col)) in inner_left
5959 .columns
5960 .into_iter()
5961 .zip(unmatched_left.columns)
5962 .enumerate()
5963 {
5964 let elem_size = left
5965 .schema()
5966 .column_type(col_idx)
5967 .map(|t| t.size_bytes())
5968 .unwrap_or(4);
5969
5970 let inner_bytes = (inner_rows as usize)
5971 .checked_mul(elem_size)
5972 .ok_or_else(|| {
5973 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
5974 })?;
5975 let unmatched_bytes = (unmatched_rows as usize)
5976 .checked_mul(elem_size)
5977 .ok_or_else(|| {
5978 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
5979 })?;
5980 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
5981 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
5982 })?;
5983
5984 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
5985
5986 if inner_bytes > 0 {
5987 let mut out_view = out_col.slice_mut(0..inner_bytes);
5988 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
5989 XlogError::Kernel(format!("Failed to copy inner left column: {}", e))
5990 })?;
5991 }
5992 if unmatched_bytes > 0 {
5993 let mut out_view = out_col.slice_mut(inner_bytes..total_bytes);
5994 let unmatched_view = self.column_bytes_view(&unmatched_col, unmatched_bytes)?;
5995 device
5996 .dtod_copy(&unmatched_view, &mut out_view)
5997 .map_err(|e| {
5998 XlogError::Kernel(format!("Failed to copy unmatched left column: {}", e))
5999 })?;
6000 }
6001
6002 result_columns.push(out_col.into());
6003 }
6004
6005 for (col_idx, inner_col) in inner_right.columns.into_iter().enumerate() {
6007 let elem_size = right
6008 .schema()
6009 .column_type(col_idx)
6010 .map(|t| t.size_bytes())
6011 .unwrap_or(4);
6012
6013 let inner_bytes = (inner_rows as usize)
6014 .checked_mul(elem_size)
6015 .ok_or_else(|| {
6016 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
6017 })?;
6018 let unmatched_bytes = (unmatched_rows as usize)
6019 .checked_mul(elem_size)
6020 .ok_or_else(|| {
6021 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
6022 })?;
6023 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
6024 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
6025 })?;
6026
6027 let mut out_col = self.memory.alloc::<u8>(total_bytes)?;
6028
6029 if total_bytes > 0 {
6030 device.memset_zeros(&mut out_col).map_err(|e| {
6031 XlogError::Kernel(format!("Failed to zero right outer column: {}", e))
6032 })?;
6033 }
6034
6035 if inner_bytes > 0 {
6036 let mut out_view = out_col.slice_mut(0..inner_bytes);
6037 device.dtod_copy(&inner_col, &mut out_view).map_err(|e| {
6038 XlogError::Kernel(format!("Failed to copy inner right column: {}", e))
6039 })?;
6040 }
6041
6042 result_columns.push(out_col.into());
6043 }
6044
6045 self.device.synchronize()?;
6046
6047 self.buffer_from_columns(result_columns, total_rows, combined_schema)
6048 }
6049
6050 fn left_outer_with_nulls(&self, left: &CudaBuffer, right: &CudaBuffer) -> Result<CudaBuffer> {
6052 let combined_schema = self.combine_schemas(left.schema(), right.schema());
6053 let num_rows = self.device_row_count(left)? as u64;
6054 if num_rows == 0 {
6055 return self.create_empty_buffer(combined_schema);
6056 }
6057 let device = self.device.inner();
6058
6059 let mut result_columns = Vec::with_capacity(combined_schema.arity());
6060
6061 for col_idx in 0..left.arity() {
6063 let col = left
6064 .column(col_idx)
6065 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
6066
6067 let elem_size = left
6068 .schema()
6069 .column_type(col_idx)
6070 .map(|t| t.size_bytes())
6071 .unwrap_or(4);
6072
6073 let bytes = (num_rows as usize) * elem_size;
6074 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
6075 if bytes > 0 {
6076 let src_view = self.column_bytes_view(col, bytes)?;
6077 device
6078 .dtod_copy(&src_view, &mut dst_col)
6079 .map_err(|e| XlogError::Kernel(format!("Failed to copy left column: {}", e)))?;
6080 }
6081
6082 result_columns.push(dst_col.into());
6083 }
6084
6085 for col_idx in 0..right.arity() {
6087 let elem_size = right
6088 .schema()
6089 .column_type(col_idx)
6090 .map(|t| t.size_bytes())
6091 .unwrap_or(4);
6092
6093 let bytes = (num_rows as usize) * elem_size;
6094 let mut dst_col = self.memory.alloc::<u8>(bytes)?;
6095 if bytes > 0 {
6096 device
6097 .memset_zeros(&mut dst_col)
6098 .map_err(|e| XlogError::Kernel(format!("Failed to zero null column: {}", e)))?;
6099 }
6100
6101 result_columns.push(dst_col.into());
6102 }
6103
6104 self.device.synchronize()?;
6105
6106 self.buffer_from_columns(result_columns, num_rows, combined_schema)
6107 }
6108
6109 pub fn clone_buffer(&self, buffer: &CudaBuffer) -> Result<CudaBuffer> {
6114 let verify = {
6119 static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6120 *ENABLED.get_or_init(|| {
6121 std::env::var("XLOG_DEBUG_VERIFY_CLONES").map(|v| v == "1") == Ok(true)
6122 })
6123 };
6124
6125 let mut result_columns = Vec::with_capacity(buffer.arity());
6126 let device = self.device.inner();
6127
6128 for col_idx in 0..buffer.arity() {
6129 let src_col = buffer
6130 .column(col_idx)
6131 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6132 let mut dst_col = self.memory.alloc::<u8>(src_col.len())?;
6133 if !src_col.is_empty() {
6134 device
6135 .dtod_copy(src_col, &mut dst_col)
6136 .map_err(|e| XlogError::Kernel(format!("Failed to clone column: {}", e)))?;
6137 }
6138 if verify && !src_col.is_empty() {
6139 self.device.synchronize()?;
6140 let mut src_host = vec![0u8; src_col.len()];
6141 let mut dst_host = vec![0u8; dst_col.len()];
6142 device
6143 .dtoh_sync_copy_into(src_col, &mut src_host)
6144 .map_err(|e| XlogError::Kernel(format!("verify src dtoh: {}", e)))?;
6145 device
6146 .dtoh_sync_copy_into(&dst_col, &mut dst_host)
6147 .map_err(|e| XlogError::Kernel(format!("verify dst dtoh: {}", e)))?;
6148 if src_host != dst_host {
6149 let first_diff = src_host
6150 .iter()
6151 .zip(dst_host.iter())
6152 .position(|(a, b)| a != b)
6153 .unwrap_or(0);
6154 return Err(XlogError::Kernel(format!(
6155 "CLONE VERIFY FAILED: column {} differs from source at byte {} of {} (clone is wrong at birth)",
6156 col_idx,
6157 first_diff,
6158 src_col.len(),
6159 )));
6160 }
6161 }
6162 result_columns.push(dst_col.into());
6163 }
6164
6165 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
6166 device
6167 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
6168 .map_err(|e| XlogError::Kernel(format!("Failed to clone row count: {}", e)))?;
6169
6170 let mut cloned = CudaBuffer::from_columns(
6171 result_columns,
6172 buffer.row_cap,
6173 d_num_rows,
6174 buffer.schema().clone(),
6175 );
6176 if let Some(cached) = buffer.cached_row_count() {
6179 cloned.set_cached_row_count_if_unset(cached);
6180 }
6181 if buffer.canonical_full_row_set_certified() {
6182 cloned.certify_canonical_full_row_set();
6183 }
6184 Ok(cloned)
6185 }
6186 pub fn extract_column(&self, buffer: &CudaBuffer, col_idx: usize) -> Result<CudaBuffer> {
6197 if buffer.is_empty() {
6198 let col_type = buffer
6199 .schema()
6200 .column_type(col_idx)
6201 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6202 let schema = Schema::new(vec![("col".to_string(), col_type)]);
6203 return self.create_empty_buffer(schema);
6204 }
6205
6206 let col_type = buffer
6207 .schema()
6208 .column_type(col_idx)
6209 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6210 let src_col = buffer
6211 .column(col_idx)
6212 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found in buffer", col_idx)))?;
6213 let mut dst_col = self.memory.alloc::<u8>(src_col.len())?;
6214 let device = self.device.inner();
6215 if !src_col.is_empty() {
6216 device
6217 .dtod_copy(src_col, &mut dst_col)
6218 .map_err(|e| XlogError::Kernel(format!("Failed to copy column: {}", e)))?;
6219 }
6220
6221 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
6222 device
6223 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
6224 .map_err(|e| XlogError::Kernel(format!("Failed to copy row count: {}", e)))?;
6225 self.device.synchronize()?;
6226
6227 let schema = Schema::new(vec![("col".to_string(), col_type)]);
6228 Ok(CudaBuffer::from_columns(
6229 vec![dst_col.into()],
6230 buffer.row_cap,
6231 d_num_rows,
6232 schema,
6233 ))
6234 }
6235
6236 pub fn extract_active_rule_indices(
6239 &self,
6240 mask_hard: &CudaBuffer,
6241 mask_soft: &CudaBuffer,
6242 n: usize,
6243 max_active: usize,
6244 ) -> Result<Vec<(u32, u32, u32)>> {
6245 let total = n * n * n;
6246 let block_size = 256usize;
6247 let grid_size = total.div_ceil(block_size);
6248
6249 let mut out_i = self.memory().alloc::<u32>(total)?;
6250 let mut out_j = self.memory().alloc::<u32>(total)?;
6251 let mut out_k = self.memory().alloc::<u32>(total)?;
6252 let mut out_p = self.memory().alloc::<f32>(total)?;
6253 let mut count = self.memory().alloc::<u32>(1)?;
6254
6255 self.htod_launch_metadata_sync_copy_into(&[0u32], &mut count)
6256 .map_err(|e| XlogError::Kernel(format!("ILP htod count: {}", e)))?;
6257
6258 let hard_col = mask_hard
6259 .column(0)
6260 .ok_or_else(|| XlogError::Kernel("ILP hard mask has no column".into()))?;
6261 let soft_col = mask_soft
6262 .column(0)
6263 .ok_or_else(|| XlogError::Kernel("ILP soft mask has no column".into()))?;
6264
6265 let kernel = self
6266 .device()
6267 .inner()
6268 .get_func(ILP_MODULE, ilp_kernels::EXTRACT_NONZERO_INDICES)
6269 .ok_or_else(|| XlogError::Kernel("extract_nonzero_indices kernel not found".into()))?;
6270
6271 let hard_bytes = total * std::mem::size_of::<f32>();
6272 let soft_bytes = total * std::mem::size_of::<f32>();
6273 let hard_view = self.column_bytes_view(hard_col, hard_bytes)?;
6274 let soft_view = self.column_bytes_view(soft_col, soft_bytes)?;
6275
6276 unsafe {
6278 kernel
6279 .clone()
6280 .launch(
6281 cudarc::driver::LaunchConfig {
6282 grid_dim: (grid_size as u32, 1, 1),
6283 block_dim: (block_size as u32, 1, 1),
6284 shared_mem_bytes: 0,
6285 },
6286 (
6287 &hard_view, &soft_view, n as u32, &mut out_i, &mut out_j, &mut out_k,
6288 &mut out_p, &mut count,
6289 ),
6290 )
6291 .map_err(|e| {
6292 XlogError::Kernel(format!("Failed to launch extract_nonzero_indices: {}", e))
6293 })?;
6294 }
6295
6296 let mut count_host = [0u32];
6297 self.device()
6298 .inner()
6299 .dtoh_sync_copy_into(&count, &mut count_host)
6300 .map_err(|e| XlogError::Kernel(format!("ILP dtoh count: {}", e)))?;
6301 let active_count = count_host[0] as usize;
6302
6303 if active_count == 0 {
6304 return Ok(Vec::new());
6305 }
6306
6307 let mut i_host = vec![0u32; active_count];
6308 let mut j_host = vec![0u32; active_count];
6309 let mut k_host = vec![0u32; active_count];
6310 let mut p_host = vec![0f32; active_count];
6311
6312 let out_i_view = out_i
6313 .try_slice(0..active_count)
6314 .ok_or_else(|| XlogError::Kernel("ILP slice i out of bounds".into()))?;
6315 let out_j_view = out_j
6316 .try_slice(0..active_count)
6317 .ok_or_else(|| XlogError::Kernel("ILP slice j out of bounds".into()))?;
6318 let out_k_view = out_k
6319 .try_slice(0..active_count)
6320 .ok_or_else(|| XlogError::Kernel("ILP slice k out of bounds".into()))?;
6321 let out_p_view = out_p
6322 .try_slice(0..active_count)
6323 .ok_or_else(|| XlogError::Kernel("ILP slice p out of bounds".into()))?;
6324
6325 self.device()
6326 .inner()
6327 .dtoh_sync_copy_into(&out_i_view, &mut i_host)
6328 .map_err(|e| XlogError::Kernel(format!("ILP dtoh i: {}", e)))?;
6329 self.device()
6330 .inner()
6331 .dtoh_sync_copy_into(&out_j_view, &mut j_host)
6332 .map_err(|e| XlogError::Kernel(format!("ILP dtoh j: {}", e)))?;
6333 self.device()
6334 .inner()
6335 .dtoh_sync_copy_into(&out_k_view, &mut k_host)
6336 .map_err(|e| XlogError::Kernel(format!("ILP dtoh k: {}", e)))?;
6337 self.device()
6338 .inner()
6339 .dtoh_sync_copy_into(&out_p_view, &mut p_host)
6340 .map_err(|e| XlogError::Kernel(format!("ILP dtoh p: {}", e)))?;
6341
6342 let mut indices: Vec<(f32, u32, u32, u32)> = (0..active_count)
6343 .map(|idx| (p_host[idx], i_host[idx], j_host[idx], k_host[idx]))
6344 .collect();
6345 indices.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
6346 indices.truncate(max_active);
6347
6348 Ok(indices.into_iter().map(|(_, i, j, k)| (i, j, k)).collect())
6349 }
6350
6351 #[allow(clippy::too_many_arguments)]
6377 fn radix_sort_u32_pairs_with_scratch_on_stream(
6378 &self,
6379 keys_a: &mut TrackedCudaSlice<u32>,
6380 keys_b: &mut TrackedCudaSlice<u32>,
6381 indices_a: &mut TrackedCudaSlice<u32>,
6382 indices_b: &mut TrackedCudaSlice<u32>,
6383 hist: &mut TrackedCudaSlice<u32>,
6384 prefix: &mut TrackedCudaSlice<u32>,
6385 ranks: &mut TrackedCudaSlice<u32>,
6386 num_rows_device: &TrackedCudaSlice<u32>,
6387 row_cap: u32,
6388 cu_stream: &cudarc::driver::CudaStream,
6389 launch_stream: StreamId,
6390 runtime: &crate::device_runtime::XlogDeviceRuntime,
6391 ) -> Result<()> {
6392 if row_cap == 0 {
6393 return Ok(());
6394 }
6395 let device = self.device.inner();
6396 let block_size = Self::SORT_BLOCK_SIZE;
6397 let grid_size = row_cap.div_ceil(block_size);
6398 let sort_config = LaunchConfig {
6399 grid_dim: (grid_size, 1, 1),
6400 block_dim: (block_size, 1, 1),
6401 shared_mem_bytes: 0,
6402 };
6403
6404 let histogram_fn = device
6405 .get_func(SORT_MODULE, sort_kernels::RADIX_HISTOGRAM)
6406 .ok_or_else(|| XlogError::Kernel("radix_histogram kernel not found".to_string()))?;
6407 let prefix_fn = device
6408 .get_func(SORT_MODULE, sort_kernels::COMPUTE_DIGIT_PREFIX_SUMS)
6409 .ok_or_else(|| {
6410 XlogError::Kernel("compute_digit_prefix_sums kernel not found".to_string())
6411 })?;
6412 let ranks_fn = device
6413 .get_func(SORT_MODULE, sort_kernels::COMPUTE_RANKS)
6414 .ok_or_else(|| XlogError::Kernel("compute_ranks kernel not found".to_string()))?;
6415 let scatter_fn = device
6416 .get_func(SORT_MODULE, sort_kernels::RADIX_SCATTER_STABLE)
6417 .ok_or_else(|| {
6418 XlogError::Kernel("radix_scatter_stable kernel not found".to_string())
6419 })?;
6420 let prefix_config = LaunchConfig {
6421 grid_dim: (1, 1, 1),
6422 block_dim: (256, 1, 1),
6423 shared_mem_bytes: 0,
6424 };
6425
6426 let mut in_a = true;
6427 for pass in 0..8u32 {
6428 let shift = pass * 4;
6429 let (keys_in, indices_in, keys_out, indices_out) = if in_a {
6430 (&*keys_a, &*indices_a, &mut *keys_b, &mut *indices_b)
6431 } else {
6432 (&*keys_b, &*indices_b, &mut *keys_a, &mut *indices_a)
6433 };
6434
6435 unsafe {
6437 histogram_fn.clone().launch_on_stream(
6438 cu_stream,
6439 sort_config,
6440 (keys_in, num_rows_device, row_cap, &mut *hist, shift),
6441 )
6442 }
6443 .map_err(|e| XlogError::Kernel(format!("radix_histogram (on_stream) failed: {}", e)))?;
6444
6445 unsafe {
6447 prefix_fn.clone().launch_on_stream(
6448 cu_stream,
6449 prefix_config,
6450 (&*hist, grid_size, &mut *prefix),
6451 )
6452 }
6453 .map_err(|e| {
6454 XlogError::Kernel(format!(
6455 "compute_digit_prefix_sums (on_stream) failed: {}",
6456 e
6457 ))
6458 })?;
6459
6460 for digit in 0..16u32 {
6463 let start = (digit * grid_size) as usize;
6464 let end = start + (grid_size as usize);
6465 let mut digit_slice = hist.slice_mut(start..end);
6466 self.multiblock_scan_u32_view_inplace_on_stream(
6467 &mut digit_slice,
6468 grid_size,
6469 cu_stream,
6470 launch_stream,
6471 runtime,
6472 )?;
6473 }
6474
6475 unsafe {
6477 ranks_fn.clone().launch_on_stream(
6478 cu_stream,
6479 sort_config,
6480 (keys_in, num_rows_device, row_cap, &mut *ranks, shift),
6481 )
6482 }
6483 .map_err(|e| XlogError::Kernel(format!("compute_ranks (on_stream) failed: {}", e)))?;
6484
6485 unsafe {
6488 scatter_fn.clone().launch_on_stream(
6489 cu_stream,
6490 sort_config,
6491 (
6492 keys_in,
6493 indices_in,
6494 &*ranks,
6495 keys_out,
6496 indices_out,
6497 &*prefix,
6498 &*hist,
6499 num_rows_device,
6500 row_cap,
6501 shift,
6502 ),
6503 )
6504 }
6505 .map_err(|e| {
6506 XlogError::Kernel(format!("radix_scatter_stable (on_stream) failed: {}", e))
6507 })?;
6508
6509 in_a = !in_a;
6510 }
6511
6512 if !in_a {
6513 return Err(XlogError::Kernel(
6514 "Unexpected radix-sort buffer parity (expected even number of passes)".to_string(),
6515 ));
6516 }
6517 Ok(())
6518 }
6519
6520 fn apply_permutation_gpu_on_stream(
6526 &self,
6527 input: &CudaBuffer,
6528 permutation: &TrackedCudaSlice<u32>,
6529 dst_cols: &mut [TrackedCudaSlice<u8>],
6530 cu_stream: &cudarc::driver::CudaStream,
6531 ) -> Result<()> {
6532 let row_cap = input.num_rows() as u32;
6533 let d_num_rows = input.num_rows_device();
6534 let device = self.device.inner();
6535
6536 let grid_size = row_cap.div_ceil(Self::SORT_BLOCK_SIZE);
6537 let launch_config = LaunchConfig {
6538 grid_dim: (grid_size, 1, 1),
6539 block_dim: (Self::SORT_BLOCK_SIZE, 1, 1),
6540 shared_mem_bytes: 0,
6541 };
6542
6543 let apply_perm_fn = device
6544 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
6545 .ok_or_else(|| {
6546 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
6547 })?;
6548
6549 if dst_cols.len() != input.columns.len() {
6550 return Err(XlogError::Kernel(format!(
6551 "apply_permutation_gpu_on_stream: dst_cols.len()={} mismatches input.cols={}",
6552 dst_cols.len(),
6553 input.columns.len()
6554 )));
6555 }
6556
6557 for (col_idx, dst_col) in dst_cols.iter_mut().enumerate() {
6558 let src_col = input
6559 .column(col_idx)
6560 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6561 let elem_size = input
6562 .schema
6563 .column_type(col_idx)
6564 .ok_or_else(|| {
6565 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
6566 })?
6567 .size_bytes() as u32;
6568 let output_bytes = (row_cap as usize) * (elem_size as usize);
6569 if src_col.num_bytes() != output_bytes {
6570 return Err(XlogError::Kernel(format!(
6571 "Column {} has {} bytes but expected {} (num_rows={}, elem_size={})",
6572 col_idx,
6573 src_col.num_bytes(),
6574 output_bytes,
6575 row_cap,
6576 elem_size
6577 )));
6578 }
6579 unsafe {
6582 apply_perm_fn.clone().launch_on_stream(
6583 cu_stream,
6584 launch_config,
6585 (
6586 src_col,
6587 &mut *dst_col,
6588 permutation,
6589 d_num_rows,
6590 row_cap,
6591 elem_size,
6592 ),
6593 )
6594 }
6595 .map_err(|e| {
6596 XlogError::Kernel(format!("apply_permutation_bytes (on_stream) failed: {}", e))
6597 })?;
6598 }
6599 Ok(())
6600 }
6601
6602 pub fn sort_recorded(
6620 &self,
6621 input: &CudaBuffer,
6622 key_cols: &[usize],
6623 launch_stream: StreamId,
6624 ) -> Result<CudaBuffer> {
6625 let runtime = self.memory.runtime().ok_or_else(|| {
6626 XlogError::Kernel(
6627 "sort_recorded requires a runtime-backed GpuMemoryManager (with_runtime)"
6628 .to_string(),
6629 )
6630 })?;
6631 let cu_stream = runtime
6632 .stream_pool()
6633 .resolve(launch_stream)
6634 .ok_or_else(|| {
6635 XlogError::Kernel(format!(
6636 "sort_recorded: launch_stream StreamId({}) does not resolve",
6637 launch_stream.0
6638 ))
6639 })?;
6640
6641 if input.num_rows() == 0 {
6642 return self.create_empty_buffer(input.schema.clone());
6643 }
6644 if key_cols.is_empty() {
6645 return Err(XlogError::Kernel(
6646 "Sort requires at least one key column".to_string(),
6647 ));
6648 }
6649 if input.num_rows() > u32::MAX as u64 {
6650 return Err(XlogError::Kernel(format!(
6651 "Sort supports at most {} rows, got {}",
6652 u32::MAX,
6653 input.num_rows()
6654 )));
6655 }
6656 for &k in key_cols {
6657 if k >= input.arity() {
6658 return Err(XlogError::Kernel(format!(
6659 "Key column index {} out of bounds (arity {})",
6660 k,
6661 input.arity()
6662 )));
6663 }
6664 let ty = input.schema.column_type(k).ok_or_else(|| {
6665 XlogError::Kernel(format!("Key column {} type not found in schema", k))
6666 })?;
6667 if !matches!(ty, ScalarType::U32 | ScalarType::Symbol | ScalarType::U64) {
6668 return Err(XlogError::Kernel(format!(
6669 "sort_recorded supports only U32 / Symbol / U64 key columns; \
6670 got {:?} for column {}",
6671 ty, k
6672 )));
6673 }
6674 }
6675
6676 let n = input.num_rows() as u32;
6677 let block_size = Self::SORT_BLOCK_SIZE;
6678 let grid_size = n.div_ceil(block_size);
6679 let device = self.device.inner();
6680 let launch_config = LaunchConfig {
6681 grid_dim: (grid_size, 1, 1),
6682 block_dim: (block_size, 1, 1),
6683 shared_mem_bytes: 0,
6684 };
6685
6686 let mut indices_a = self.memory.alloc::<u32>(n as usize)?;
6689 let mut indices_b = self.memory.alloc::<u32>(n as usize)?;
6690 let mut keys_a = self.memory.alloc::<u32>(n as usize)?;
6691 let mut keys_b = self.memory.alloc::<u32>(n as usize)?;
6692 let mut d_hist = self.memory.alloc::<u32>((grid_size as usize) * 16)?;
6693 let mut d_prefix = self.memory.alloc::<u32>(16)?;
6694 let mut d_ranks = self.memory.alloc::<u32>(n as usize)?;
6695 let output_d_num_rows = self.memory.alloc::<u32>(1)?;
6701
6702 let mut dst_cols: Vec<TrackedCudaSlice<u8>> = Vec::with_capacity(input.columns.len());
6703 for col_idx in 0..input.columns.len() {
6704 let elem_size = input
6705 .schema
6706 .column_type(col_idx)
6707 .ok_or_else(|| {
6708 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
6709 })?
6710 .size_bytes();
6711 dst_cols.push(self.memory.alloc::<u8>((n as usize) * elem_size)?);
6712 }
6713
6714 let mut rec = LaunchRecorder::new_strict(launch_stream);
6715 rec.read(input.num_rows_device());
6716 for col_idx in 0..input.columns.len() {
6717 let c = input
6718 .column(col_idx)
6719 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
6720 rec.read_column(c);
6721 }
6722 rec.write(&indices_a);
6726 rec.write(&indices_b);
6727 rec.write(&keys_a);
6728 rec.write(&keys_b);
6729 rec.write(&d_hist);
6730 rec.write(&d_prefix);
6731 rec.write(&d_ranks);
6732 rec.write(&output_d_num_rows);
6733 for dst_col in &dst_cols {
6734 rec.write(dst_col);
6735 }
6736 rec.preflight(runtime)
6737 .map_err(|e| XlogError::Kernel(format!("sort_recorded: preflight failed: {}", e)))?;
6738
6739 let init_fn = device
6741 .get_func(SORT_MODULE, sort_kernels::INIT_INDICES)
6742 .ok_or_else(|| XlogError::Kernel("init_indices kernel not found".to_string()))?;
6743 unsafe {
6745 init_fn.clone().launch_on_stream(
6746 &cu_stream,
6747 launch_config,
6748 (&mut indices_a, input.num_rows_device(), n),
6749 )
6750 }
6751 .map_err(|e| XlogError::Kernel(format!("init_indices (on_stream) failed: {}", e)))?;
6752
6753 for &col_idx in key_cols.iter().rev() {
6760 let col = input
6761 .column(col_idx)
6762 .ok_or_else(|| XlogError::Kernel(format!("Key column {} not found", col_idx)))?;
6763 let ty = input.schema.column_type(col_idx).ok_or_else(|| {
6764 XlogError::Kernel(format!("Key column {} type not found in schema", col_idx))
6765 })?;
6766 match ty {
6767 ScalarType::U32 | ScalarType::Symbol => {
6768 let col_view = self.column_as_u32_view(col, n as usize)?;
6769 let gather_fn = device
6770 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_U32)
6771 .ok_or_else(|| {
6772 XlogError::Kernel("apply_permutation_u32 kernel not found".to_string())
6773 })?;
6774 unsafe {
6777 gather_fn.clone().launch_on_stream(
6778 &cu_stream,
6779 launch_config,
6780 (
6781 &col_view,
6782 &mut keys_a,
6783 &indices_a,
6784 input.num_rows_device(),
6785 n,
6786 ),
6787 )
6788 }
6789 .map_err(|e| {
6790 XlogError::Kernel(format!(
6791 "apply_permutation_u32 (on_stream) failed: {}",
6792 e
6793 ))
6794 })?;
6795
6796 self.radix_sort_u32_pairs_with_scratch_on_stream(
6797 &mut keys_a,
6798 &mut keys_b,
6799 &mut indices_a,
6800 &mut indices_b,
6801 &mut d_hist,
6802 &mut d_prefix,
6803 &mut d_ranks,
6804 input.num_rows_device(),
6805 n,
6806 &cu_stream,
6807 launch_stream,
6808 runtime,
6809 )?;
6810 }
6811 ScalarType::U64 => {
6812 let col_view = self.column_as_u64_view(col, n as usize)?;
6813 for &word in &[
6814 sort_kernels::GATHER_KEYS_U64_LO_U32,
6815 sort_kernels::GATHER_KEYS_U64_HI_U32,
6816 ] {
6817 let gather_fn = device.get_func(SORT_MODULE, word).ok_or_else(|| {
6818 XlogError::Kernel(format!("{} kernel not found", word))
6819 })?;
6820 unsafe {
6823 gather_fn.clone().launch_on_stream(
6824 &cu_stream,
6825 launch_config,
6826 (
6827 &col_view,
6828 &indices_a,
6829 input.num_rows_device(),
6830 n,
6831 &mut keys_a,
6832 ),
6833 )
6834 }
6835 .map_err(|e| {
6836 XlogError::Kernel(format!("{} (on_stream) failed: {}", word, e))
6837 })?;
6838
6839 self.radix_sort_u32_pairs_with_scratch_on_stream(
6840 &mut keys_a,
6841 &mut keys_b,
6842 &mut indices_a,
6843 &mut indices_b,
6844 &mut d_hist,
6845 &mut d_prefix,
6846 &mut d_ranks,
6847 input.num_rows_device(),
6848 n,
6849 &cu_stream,
6850 launch_stream,
6851 runtime,
6852 )?;
6853 }
6854 }
6855 other => {
6856 return Err(XlogError::Kernel(format!(
6857 "sort_recorded: column {} unexpected type {:?} after guard",
6858 col_idx, other
6859 )));
6860 }
6861 }
6862 }
6863
6864 self.apply_permutation_gpu_on_stream(input, &indices_a, &mut dst_cols, &cu_stream)?;
6866
6867 unsafe {
6874 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
6875 *output_d_num_rows.device_ptr(),
6876 *input.num_rows_device().device_ptr(),
6877 std::mem::size_of::<u32>(),
6878 cu_stream.cu_stream(),
6879 );
6880 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
6881 return Err(XlogError::Kernel(format!(
6882 "sort_recorded: cuMemcpyDtoDAsync (output_d_num_rows) failed: {:?}",
6883 res
6884 )));
6885 }
6886 }
6887
6888 rec.commit(runtime)
6889 .map_err(|e| XlogError::Kernel(format!("sort_recorded: commit failed: {}", e)))?;
6890
6891 let new_columns: Vec<CudaColumn> = dst_cols.into_iter().map(|s| s.into()).collect();
6892 Ok(CudaBuffer::from_columns(
6893 new_columns,
6894 input.num_rows(),
6895 output_d_num_rows,
6896 input.schema.clone(),
6897 ))
6898 }
6899
6900 pub fn dedup_full_row_recorded(
6911 &self,
6912 input: &CudaBuffer,
6913 launch_stream: StreamId,
6914 ) -> Result<CudaBuffer> {
6915 if !input.canonical_full_row_set_certified() {
6916 self.validated_logical_row_count(input)?;
6917 }
6918 let runtime = self.memory.runtime().ok_or_else(|| {
6919 XlogError::Kernel(
6920 "dedup_full_row_recorded requires a runtime-backed GpuMemoryManager".to_string(),
6921 )
6922 })?;
6923 let cu_stream = runtime
6924 .stream_pool()
6925 .resolve(launch_stream)
6926 .ok_or_else(|| {
6927 XlogError::Kernel(format!(
6928 "dedup_full_row_recorded: launch_stream StreamId({}) does not resolve",
6929 launch_stream.0
6930 ))
6931 })?;
6932
6933 let row_count = input.num_rows() as usize;
6934 if row_count == 0 {
6935 return self.create_empty_buffer(input.schema().clone());
6936 }
6937 if row_count == 1 {
6938 return self.clone_buffer(input);
6939 }
6940 if row_count > u32::MAX as usize {
6941 return Err(XlogError::Kernel(format!(
6942 "dedup_full_row_recorded supports at most {} rows, got {}",
6943 u32::MAX,
6944 row_count
6945 )));
6946 }
6947 let arity = input.arity();
6948 if arity == 0 {
6949 return self.buffer_from_columns(Vec::new(), 1, input.schema().clone());
6950 }
6951 for col_idx in 0..arity {
6952 let ty = input.schema.column_type(col_idx).ok_or_else(|| {
6953 XlogError::Kernel(format!("Column {} type not found in schema", col_idx))
6954 })?;
6955 if !matches!(ty, ScalarType::U32 | ScalarType::Symbol | ScalarType::U64) {
6956 return Err(XlogError::Kernel(format!(
6957 "dedup_full_row_recorded supports only U32 / Symbol / U64 columns; \
6958 got {:?} for column {}",
6959 ty, col_idx
6960 )));
6961 }
6962 }
6963
6964 let all_cols: Vec<usize> = (0..arity).collect();
6966 let sorted = self.sort_recorded(input, &all_cols, launch_stream)?;
6967 let n = sorted.num_rows() as u32;
6968 if n <= 1 {
6969 return Ok(sorted);
6970 }
6971
6972 let device = self.device.inner();
6978 let mut col_ptrs_host: Vec<u64> = Vec::with_capacity(arity);
6979 let mut col_sizes_host: Vec<u32> = Vec::with_capacity(arity);
6980 for col_idx in 0..arity {
6981 let c = sorted
6982 .column(col_idx)
6983 .ok_or_else(|| XlogError::Kernel(format!("Sorted column {} not found", col_idx)))?;
6984 let ty = sorted.schema().column_type(col_idx).ok_or_else(|| {
6985 XlogError::Kernel(format!("Sorted column {} type missing", col_idx))
6986 })?;
6987 col_ptrs_host.push(*c.device_ptr());
6988 col_sizes_host.push(ty.size_bytes() as u32);
6989 }
6990 let mut d_col_ptrs = self.memory.alloc::<u64>(arity)?;
6991 let mut d_col_sizes = self.memory.alloc::<u32>(arity)?;
6992 self.htod_launch_metadata_sync_copy_into(&col_ptrs_host, &mut d_col_ptrs)
6993 .map_err(|e| {
6994 XlogError::Kernel(format!("dedup_full_row_recorded col ptr upload: {}", e))
6995 })?;
6996 self.htod_launch_metadata_sync_copy_into(&col_sizes_host, &mut d_col_sizes)
6997 .map_err(|e| {
6998 XlogError::Kernel(format!("dedup_full_row_recorded col size upload: {}", e))
6999 })?;
7000 let d_unique_mask = self.memory.alloc::<u8>(n as usize)?;
7001
7002 let mut rec = LaunchRecorder::new_strict(launch_stream);
7003 for col_idx in 0..arity {
7004 let c = sorted
7005 .column(col_idx)
7006 .ok_or_else(|| XlogError::Kernel(format!("Sorted column {} not found", col_idx)))?;
7007 rec.read_column(c);
7008 }
7009 rec.read(sorted.num_rows_device());
7010 rec.write(&d_col_ptrs);
7011 rec.write(&d_col_sizes);
7012 rec.write(&d_unique_mask);
7013 rec.preflight(runtime).map_err(|e| {
7014 XlogError::Kernel(format!(
7015 "dedup_full_row_recorded: mark_unique preflight failed: {}",
7016 e
7017 ))
7018 })?;
7019
7020 let block_size = 256u32;
7021 let grid = n.div_ceil(block_size);
7022 let cfg = LaunchConfig {
7023 grid_dim: (grid, 1, 1),
7024 block_dim: (block_size, 1, 1),
7025 shared_mem_bytes: 0,
7026 };
7027 let mark_fn = device
7028 .get_func(DEDUP_MODULE, dedup_kernels::MARK_UNIQUE_FULL_ROW_BYTEWISE)
7029 .ok_or_else(|| {
7030 XlogError::Kernel("mark_unique_full_row_bytewise kernel not found".to_string())
7031 })?;
7032 unsafe {
7035 mark_fn.clone().launch_on_stream(
7036 &cu_stream,
7037 cfg,
7038 (
7039 &d_col_ptrs,
7040 &d_col_sizes,
7041 arity as u32,
7042 sorted.num_rows_device(),
7043 n,
7044 &d_unique_mask,
7045 ),
7046 )
7047 }
7048 .map_err(|e| {
7049 XlogError::Kernel(format!(
7050 "mark_unique_full_row_bytewise (on_stream) failed: {}",
7051 e
7052 ))
7053 })?;
7054
7055 rec.commit(runtime).map_err(|e| {
7056 XlogError::Kernel(format!(
7057 "dedup_full_row_recorded: mark_unique commit failed: {}",
7058 e
7059 ))
7060 })?;
7061
7062 let mut result = self.compact_buffer_by_device_mask_counted_recorded(
7064 &sorted,
7065 &d_unique_mask,
7066 launch_stream,
7067 )?;
7068 result.certify_canonical_full_row_set();
7069 Ok(result)
7070 }
7071
7072 fn build_hash_table_v2_on_stream(
7106 &self,
7107 hashes: &TrackedCudaSlice<u64>,
7108 num_rows: u32,
7109 cu_stream: &cudarc::driver::CudaStream,
7110 launch_stream: StreamId,
7111 runtime: &crate::device_runtime::XlogDeviceRuntime,
7112 ) -> Result<crate::provider::JoinHashTableV2> {
7113 let device = self.device.inner();
7114
7115 let target = (num_rows as u64).saturating_mul(2).max(1024);
7116 let num_buckets_u64 = target.next_power_of_two();
7117 let num_buckets = u32::try_from(num_buckets_u64).map_err(|_| {
7118 XlogError::Kernel(format!(
7119 "Join hash table too large: num_buckets={}",
7120 num_buckets_u64
7121 ))
7122 })?;
7123 let bucket_mask = num_buckets
7124 .checked_sub(1)
7125 .ok_or_else(|| XlogError::Kernel("Join hash table size underflow".to_string()))?;
7126
7127 let bucket_counts = self.memory.alloc::<u32>(num_buckets as usize)?;
7128 runtime
7135 .prepare_first_use(&bucket_counts, launch_stream, Access::Write)
7136 .map_err(|e| {
7137 XlogError::Kernel(format!(
7138 "build_hash_table_v2_on_stream: prepare bucket_counts failed: {}",
7139 e
7140 ))
7141 })?;
7142 if num_buckets > 0 {
7144 unsafe {
7148 let res = cudarc::driver::sys::cuMemsetD8Async(
7149 *bucket_counts.device_ptr(),
7150 0,
7151 (num_buckets as usize) * std::mem::size_of::<u32>(),
7152 cu_stream.cu_stream(),
7153 );
7154 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7155 return Err(XlogError::Kernel(format!(
7156 "cuMemsetD8Async (bucket_counts) failed: {:?}",
7157 res
7158 )));
7159 }
7160 }
7161 }
7162
7163 let block_size = 256u32;
7164 let grid_size = num_rows.div_ceil(block_size);
7165 let cfg = LaunchConfig {
7166 grid_dim: (grid_size, 1, 1),
7167 block_dim: (block_size, 1, 1),
7168 shared_mem_bytes: 0,
7169 };
7170
7171 let count_fn = device
7172 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUCKET_COUNT_V2)
7173 .ok_or_else(|| {
7174 XlogError::Kernel("hash_join_bucket_count_v2 kernel not found".to_string())
7175 })?;
7176 unsafe {
7178 count_fn.clone().launch_on_stream(
7179 cu_stream,
7180 cfg,
7181 (hashes, num_rows, &bucket_counts, bucket_mask),
7182 )
7183 }
7184 .map_err(|e| {
7185 XlogError::Kernel(format!(
7186 "hash_join_bucket_count_v2 (on_stream) failed: {}",
7187 e
7188 ))
7189 })?;
7190
7191 let mut bucket_offsets = self.memory.alloc::<u32>(num_buckets as usize)?;
7192 runtime
7195 .prepare_first_use(&bucket_offsets, launch_stream, Access::Write)
7196 .map_err(|e| {
7197 XlogError::Kernel(format!(
7198 "build_hash_table_v2_on_stream: prepare bucket_offsets failed: {}",
7199 e
7200 ))
7201 })?;
7202 if num_buckets > 0 {
7203 unsafe {
7207 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7208 *bucket_offsets.device_ptr(),
7209 *bucket_counts.device_ptr(),
7210 (num_buckets as usize) * std::mem::size_of::<u32>(),
7211 cu_stream.cu_stream(),
7212 );
7213 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7214 return Err(XlogError::Kernel(format!(
7215 "cuMemcpyDtoDAsync (bucket_counts → bucket_offsets) failed: {:?}",
7216 res
7217 )));
7218 }
7219 }
7220 self.multiblock_scan_u32_inplace_on_stream(
7221 &mut bucket_offsets,
7222 num_buckets,
7223 cu_stream,
7224 launch_stream,
7225 runtime,
7226 )?;
7227 }
7228
7229 let bucket_cursors = self.memory.alloc::<u32>(num_buckets as usize)?;
7230 runtime
7233 .prepare_first_use(&bucket_cursors, launch_stream, Access::Write)
7234 .map_err(|e| {
7235 XlogError::Kernel(format!(
7236 "build_hash_table_v2_on_stream: prepare bucket_cursors failed: {}",
7237 e
7238 ))
7239 })?;
7240 if num_buckets > 0 {
7241 unsafe {
7244 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
7245 *bucket_cursors.device_ptr(),
7246 *bucket_offsets.device_ptr(),
7247 (num_buckets as usize) * std::mem::size_of::<u32>(),
7248 cu_stream.cu_stream(),
7249 );
7250 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7251 return Err(XlogError::Kernel(format!(
7252 "cuMemcpyDtoDAsync (bucket_offsets → bucket_cursors) failed: {:?}",
7253 res
7254 )));
7255 }
7256 }
7257 }
7258
7259 let bucket_entries = self.memory.alloc::<u32>(num_rows as usize)?;
7260 let bucket_entry_hashes = self.memory.alloc::<u64>(num_rows as usize)?;
7261 runtime
7264 .prepare_first_use(&bucket_entries, launch_stream, Access::Write)
7265 .map_err(|e| {
7266 XlogError::Kernel(format!(
7267 "build_hash_table_v2_on_stream: prepare bucket_entries failed: {}",
7268 e
7269 ))
7270 })?;
7271 runtime
7272 .prepare_first_use(&bucket_entry_hashes, launch_stream, Access::Write)
7273 .map_err(|e| {
7274 XlogError::Kernel(format!(
7275 "build_hash_table_v2_on_stream: prepare bucket_entry_hashes failed: {}",
7276 e
7277 ))
7278 })?;
7279
7280 let scatter_fn = device
7281 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SCATTER_V2)
7282 .ok_or_else(|| {
7283 XlogError::Kernel("hash_join_scatter_v2 kernel not found".to_string())
7284 })?;
7285 unsafe {
7287 scatter_fn.clone().launch_on_stream(
7288 cu_stream,
7289 cfg,
7290 (
7291 hashes,
7292 num_rows,
7293 &bucket_cursors,
7294 bucket_mask,
7295 &bucket_entries,
7296 &bucket_entry_hashes,
7297 ),
7298 )
7299 }
7300 .map_err(|e| {
7301 XlogError::Kernel(format!("hash_join_scatter_v2 (on_stream) failed: {}", e))
7302 })?;
7303
7304 for blk in [
7312 bucket_counts.runtime_block(),
7313 bucket_offsets.runtime_block(),
7314 bucket_cursors.runtime_block(),
7315 bucket_entries.runtime_block(),
7316 bucket_entry_hashes.runtime_block(),
7317 ] {
7318 if let Some(b) = blk {
7319 runtime
7320 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
7321 .map_err(|e| {
7322 XlogError::Kernel(format!(
7323 "build_hash_table_v2_on_stream: finish_block_use failed: {}",
7324 e
7325 ))
7326 })?;
7327 } else {
7328 return Err(XlogError::Kernel(
7329 "build_hash_table_v2_on_stream: buffer has no runtime block — \
7330 caller must use a runtime-backed manager"
7331 .to_string(),
7332 ));
7333 }
7334 }
7335
7336 Ok(crate::provider::JoinHashTableV2 {
7337 bucket_counts,
7338 bucket_offsets,
7339 bucket_entries,
7340 bucket_entry_hashes,
7341 bucket_mask,
7342 })
7343 }
7344
7345 fn gather_buffer_by_indices_on_stream(
7353 &self,
7354 input: &CudaBuffer,
7355 indices: &TrackedCudaSlice<u32>,
7356 output_rows: u32,
7357 cu_stream: &cudarc::driver::CudaStream,
7358 launch_stream: StreamId,
7359 runtime: &crate::device_runtime::XlogDeviceRuntime,
7360 ) -> Result<CudaBuffer> {
7361 if output_rows == 0 {
7362 return self.create_empty_buffer(input.schema().clone());
7363 }
7364 if input.num_rows() > u32::MAX as u64 {
7365 return Err(XlogError::Kernel(format!(
7366 "GPU gather supports at most {} input rows, got {}",
7367 u32::MAX,
7368 input.num_rows()
7369 )));
7370 }
7371
7372 let d_output_rows = self.upload_device_row_count(output_rows)?;
7373 runtime
7380 .finish_first_use(&d_output_rows, StreamId::DEFAULT, Access::Write)
7381 .map_err(|e| {
7382 XlogError::Kernel(format!(
7383 "gather_buffer_by_indices_on_stream: record d_output_rows upload failed: {}",
7384 e
7385 ))
7386 })?;
7387 runtime
7388 .prepare_first_use(&d_output_rows, launch_stream, Access::Read)
7389 .map_err(|e| {
7390 XlogError::Kernel(format!(
7391 "gather_buffer_by_indices_on_stream: prepare d_output_rows failed: {}",
7392 e
7393 ))
7394 })?;
7395 let device = self.device.inner();
7396 let block_size = 256u32;
7397 let grid_size = output_rows.div_ceil(block_size);
7398 let launch_config = LaunchConfig {
7399 grid_dim: (grid_size, 1, 1),
7400 block_dim: (block_size, 1, 1),
7401 shared_mem_bytes: 0,
7402 };
7403
7404 let gather_fn = device
7405 .get_func(SORT_MODULE, sort_kernels::APPLY_PERMUTATION_BYTES)
7406 .ok_or_else(|| {
7407 XlogError::Kernel("apply_permutation_bytes kernel not found".to_string())
7408 })?;
7409
7410 let mut dst_cols: Vec<TrackedCudaSlice<u8>> = Vec::with_capacity(input.columns.len());
7411 for col_idx in 0..input.columns.len() {
7412 let elem_size = input
7413 .schema
7414 .column_type(col_idx)
7415 .ok_or_else(|| {
7416 XlogError::Kernel(format!("Schema type for column {} not found", col_idx))
7417 })?
7418 .size_bytes() as u32;
7419 let dst_bytes = (output_rows as usize) * (elem_size as usize);
7420 let dst = self.memory.alloc::<u8>(dst_bytes)?;
7421 runtime
7424 .prepare_first_use(&dst, launch_stream, Access::Write)
7425 .map_err(|e| {
7426 XlogError::Kernel(format!(
7427 "gather_buffer_by_indices_on_stream: prepare dst_col {} failed: {}",
7428 col_idx, e
7429 ))
7430 })?;
7431 dst_cols.push(dst);
7432 }
7433
7434 for (col_idx, dst_col) in dst_cols.iter_mut().enumerate() {
7435 let src_col = input
7436 .column(col_idx)
7437 .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;
7438 let elem_size = input
7439 .schema
7440 .column_type(col_idx)
7441 .map(|t| t.size_bytes() as u32)
7442 .unwrap_or(4);
7443 unsafe {
7445 gather_fn.clone().launch_on_stream(
7446 cu_stream,
7447 launch_config,
7448 (
7449 src_col,
7450 &mut *dst_col,
7451 indices,
7452 &d_output_rows,
7453 output_rows,
7454 elem_size,
7455 ),
7456 )
7457 }
7458 .map_err(|e| {
7459 XlogError::Kernel(format!("apply_permutation_bytes (on_stream) failed: {}", e))
7460 })?;
7461 }
7462
7463 runtime
7464 .finish_first_use(&d_output_rows, launch_stream, Access::Read)
7465 .map_err(|e| {
7466 XlogError::Kernel(format!(
7467 "gather_buffer_by_indices_on_stream: record d_output_rows read failed: {}",
7468 e
7469 ))
7470 })?;
7471
7472 for dst_col in &dst_cols {
7477 if let Some(b) = dst_col.runtime_block() {
7478 runtime
7479 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
7480 .map_err(|e| {
7481 XlogError::Kernel(format!(
7482 "gather_buffer_by_indices_on_stream: finish_block_use \
7483 (dst_col) failed: {}",
7484 e
7485 ))
7486 })?;
7487 } else {
7488 return Err(XlogError::Kernel(
7489 "gather_buffer_by_indices_on_stream: dst_col has no runtime block".to_string(),
7490 ));
7491 }
7492 }
7493
7494 let new_columns: Vec<CudaColumn> = dst_cols.into_iter().map(|s| s.into()).collect();
7495 Ok(CudaBuffer::from_columns(
7496 new_columns,
7497 output_rows as u64,
7498 d_output_rows,
7499 input.schema.clone(),
7500 ))
7501 }
7502
7503 pub fn hash_join_inner_v2_recorded(
7510 &self,
7511 left: &CudaBuffer,
7512 right: &CudaBuffer,
7513 left_keys: &[usize],
7514 right_keys: &[usize],
7515 max_output: Option<usize>,
7516 launch_stream: StreamId,
7517 ) -> Result<CudaBuffer> {
7518 use crate::launch::LaunchRecorder;
7519
7520 let runtime = self.memory.runtime().ok_or_else(|| {
7521 XlogError::Kernel(
7522 "hash_join_inner_v2_recorded requires a runtime-backed GpuMemoryManager"
7523 .to_string(),
7524 )
7525 })?;
7526 let cu_stream = runtime
7527 .stream_pool()
7528 .resolve(launch_stream)
7529 .ok_or_else(|| {
7530 XlogError::Kernel(format!(
7531 "hash_join_inner_v2_recorded: launch_stream StreamId({}) does not resolve",
7532 launch_stream.0
7533 ))
7534 })?;
7535
7536 let num_left = self.device_row_count(left)?;
7537 let num_right = self.device_row_count(right)?;
7538 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
7539 return Err(XlogError::Kernel(format!(
7540 "Join supports at most {} rows per side (left={}, right={})",
7541 u32::MAX,
7542 num_left,
7543 num_right
7544 )));
7545 }
7546 if num_left == 0 || num_right == 0 {
7547 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7548 return self.create_empty_buffer(combined_schema);
7549 }
7550 if left_keys.is_empty() || right_keys.is_empty() {
7551 return Err(XlogError::Kernel(
7552 "Join requires at least one key column".to_string(),
7553 ));
7554 }
7555 if left_keys.len() != right_keys.len() {
7556 return Err(XlogError::Kernel(
7557 "Left and right key columns must have same length".to_string(),
7558 ));
7559 }
7560 if left_keys.len() > 4 {
7561 return Err(XlogError::Kernel(
7562 "hash_join_inner_v2_recorded: max 4 key columns supported (pack_keys constraint)"
7563 .to_string(),
7564 ));
7565 }
7566 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
7567 let lt = left.schema().column_type(l);
7568 let rt = right.schema().column_type(r);
7569 if lt != rt {
7570 return Err(XlogError::Kernel(format!(
7571 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
7572 l, lt, r, rt
7573 )));
7574 }
7575 }
7576
7577 let num_left = num_left as u32;
7578 let num_right = num_right as u32;
7579
7580 let left_packed =
7582 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
7583 let right_packed =
7584 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
7585
7586 let table = self.build_hash_table_v2_on_stream(
7588 &right_packed.hashes,
7589 num_right,
7590 &cu_stream,
7591 launch_stream,
7592 runtime,
7593 )?;
7594
7595 let probe_func = self
7596 .device
7597 .inner()
7598 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
7599 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
7600 let block_size = 256u32;
7601 let probe_grid = num_left.div_ceil(block_size);
7602 let probe_config = LaunchConfig {
7603 grid_dim: (probe_grid, 1, 1),
7604 block_dim: (block_size, 1, 1),
7605 shared_mem_bytes: 0,
7606 };
7607
7608 let d_count_only = self.memory.alloc::<u32>(1)?;
7613 let d_dummy_left = self.memory.alloc::<u32>(1)?;
7614 let d_dummy_right = self.memory.alloc::<u32>(1)?;
7615
7616 let max_output_count_only = 0u32;
7622 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
7623 rec_count.read(&left_packed.hashes);
7624 rec_count.read(&left_packed.packed_keys);
7625 rec_count.read(&right_packed.packed_keys);
7626 rec_count.read(&table.bucket_offsets);
7627 rec_count.read(&table.bucket_counts);
7628 rec_count.read(&table.bucket_entries);
7629 rec_count.read(&table.bucket_entry_hashes);
7630 rec_count.write(&d_count_only);
7631 rec_count.write(&d_dummy_left);
7632 rec_count.write(&d_dummy_right);
7633 rec_count.preflight(runtime).map_err(|e| {
7634 XlogError::Kernel(format!(
7635 "hash_join_inner_v2_recorded: count-pass preflight failed: {}",
7636 e
7637 ))
7638 })?;
7639
7640 unsafe {
7646 let res = cudarc::driver::sys::cuMemsetD8Async(
7647 *d_count_only.device_ptr(),
7648 0,
7649 std::mem::size_of::<u32>(),
7650 cu_stream.cu_stream(),
7651 );
7652 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7653 return Err(XlogError::Kernel(format!(
7654 "cuMemsetD8Async (d_count_only) failed: {:?}",
7655 res
7656 )));
7657 }
7658 }
7659
7660 unsafe {
7665 let mut params: Vec<*mut c_void> = vec![
7666 (&left_packed.hashes).as_kernel_param(),
7667 num_left.as_kernel_param(),
7668 (&table.bucket_offsets).as_kernel_param(),
7669 (&table.bucket_counts).as_kernel_param(),
7670 (&table.bucket_entries).as_kernel_param(),
7671 (&table.bucket_entry_hashes).as_kernel_param(),
7672 table.bucket_mask.as_kernel_param(),
7673 (&left_packed.packed_keys).as_kernel_param(),
7674 (&right_packed.packed_keys).as_kernel_param(),
7675 left_packed.key_bytes.as_kernel_param(),
7676 (&d_dummy_left).as_kernel_param(),
7677 (&d_dummy_right).as_kernel_param(),
7678 (&d_count_only).as_kernel_param(),
7679 max_output_count_only.as_kernel_param(),
7680 ];
7681 probe_func
7682 .clone()
7683 .launch_on_stream(&cu_stream, probe_config, &mut params)
7684 .map_err(|e| {
7685 XlogError::Kernel(format!(
7686 "hash_join_probe_v2 (count, on_stream) failed: {}",
7687 e
7688 ))
7689 })?;
7690 }
7691
7692 rec_count.commit(runtime).map_err(|e| {
7693 XlogError::Kernel(format!(
7694 "hash_join_inner_v2_recorded: count-pass commit failed: {}",
7695 e
7696 ))
7697 })?;
7698
7699 cu_stream.synchronize().map_err(|e| {
7701 XlogError::Kernel(format!(
7702 "hash_join_inner_v2_recorded: launch_stream sync (count read) failed: {}",
7703 e
7704 ))
7705 })?;
7706 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
7707 let requested = max_output
7708 .map(|limit| (limit as u64).min(full_count))
7709 .unwrap_or(full_count);
7710 if requested == 0 {
7711 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7712 return self.create_empty_buffer(combined_schema);
7713 }
7714 if requested > u32::MAX as u64 {
7715 return Err(XlogError::Kernel(format!(
7716 "Join produced {} rows which exceeds the u32 index limit",
7717 requested
7718 )));
7719 }
7720 let max_output_u32 = requested as u32;
7721
7722 let d_output_left = self.memory.alloc::<u32>(max_output_u32 as usize)?;
7726 let d_output_right = self.memory.alloc::<u32>(max_output_u32 as usize)?;
7727 let d_output_count = self.memory.alloc::<u32>(1)?;
7728
7729 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
7730 rec_mat.read(&left_packed.hashes);
7731 rec_mat.read(&left_packed.packed_keys);
7732 rec_mat.read(&right_packed.packed_keys);
7733 rec_mat.read(&table.bucket_offsets);
7734 rec_mat.read(&table.bucket_counts);
7735 rec_mat.read(&table.bucket_entries);
7736 rec_mat.read(&table.bucket_entry_hashes);
7737 rec_mat.write(&d_output_left);
7738 rec_mat.write(&d_output_right);
7739 rec_mat.write(&d_output_count);
7740 rec_mat.preflight(runtime).map_err(|e| {
7741 XlogError::Kernel(format!(
7742 "hash_join_inner_v2_recorded: materialize-pass preflight failed: {}",
7743 e
7744 ))
7745 })?;
7746
7747 unsafe {
7751 let res = cudarc::driver::sys::cuMemsetD8Async(
7752 *d_output_count.device_ptr(),
7753 0,
7754 std::mem::size_of::<u32>(),
7755 cu_stream.cu_stream(),
7756 );
7757 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
7758 return Err(XlogError::Kernel(format!(
7759 "cuMemsetD8Async (d_output_count) failed: {:?}",
7760 res
7761 )));
7762 }
7763 }
7764
7765 unsafe {
7767 let mut params: Vec<*mut c_void> = vec![
7768 (&left_packed.hashes).as_kernel_param(),
7769 num_left.as_kernel_param(),
7770 (&table.bucket_offsets).as_kernel_param(),
7771 (&table.bucket_counts).as_kernel_param(),
7772 (&table.bucket_entries).as_kernel_param(),
7773 (&table.bucket_entry_hashes).as_kernel_param(),
7774 table.bucket_mask.as_kernel_param(),
7775 (&left_packed.packed_keys).as_kernel_param(),
7776 (&right_packed.packed_keys).as_kernel_param(),
7777 left_packed.key_bytes.as_kernel_param(),
7778 (&d_output_left).as_kernel_param(),
7779 (&d_output_right).as_kernel_param(),
7780 (&d_output_count).as_kernel_param(),
7781 max_output_u32.as_kernel_param(),
7782 ];
7783 probe_func
7784 .clone()
7785 .launch_on_stream(&cu_stream, probe_config, &mut params)
7786 .map_err(|e| {
7787 XlogError::Kernel(format!(
7788 "hash_join_probe_v2 (materialize, on_stream) failed: {}",
7789 e
7790 ))
7791 })?;
7792 }
7793
7794 rec_mat.commit(runtime).map_err(|e| {
7795 XlogError::Kernel(format!(
7796 "hash_join_inner_v2_recorded: materialize-pass commit failed: {}",
7797 e
7798 ))
7799 })?;
7800
7801 cu_stream.synchronize().map_err(|e| {
7803 XlogError::Kernel(format!(
7804 "hash_join_inner_v2_recorded: launch_stream sync (mat read) failed: {}",
7805 e
7806 ))
7807 })?;
7808 let result_count = (self.read_join_output_count_metadata(&d_output_count)? as u64)
7809 .min(max_output_u32 as u64);
7810 if result_count == 0 {
7811 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7812 return self.create_empty_buffer(combined_schema);
7813 }
7814 let output_rows = result_count as u32;
7815
7816 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
7820 for col_idx in 0..left.columns.len() {
7821 let c = left
7822 .column(col_idx)
7823 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
7824 rec_gather.read_column(c);
7825 }
7826 for col_idx in 0..right.columns.len() {
7827 let c = right
7828 .column(col_idx)
7829 .ok_or_else(|| XlogError::Kernel(format!("Right column {} not found", col_idx)))?;
7830 rec_gather.read_column(c);
7831 }
7832 rec_gather.read(&d_output_left);
7833 rec_gather.read(&d_output_right);
7834 rec_gather.preflight(runtime).map_err(|e| {
7835 XlogError::Kernel(format!(
7836 "hash_join_inner_v2_recorded: gather preflight failed: {}",
7837 e
7838 ))
7839 })?;
7840
7841 let gathered_left = self.gather_buffer_by_indices_on_stream(
7842 left,
7843 &d_output_left,
7844 output_rows,
7845 &cu_stream,
7846 launch_stream,
7847 runtime,
7848 )?;
7849 let gathered_right = self.gather_buffer_by_indices_on_stream(
7850 right,
7851 &d_output_right,
7852 output_rows,
7853 &cu_stream,
7854 launch_stream,
7855 runtime,
7856 )?;
7857
7858 rec_gather.commit(runtime).map_err(|e| {
7859 XlogError::Kernel(format!(
7860 "hash_join_inner_v2_recorded: gather commit failed: {}",
7861 e
7862 ))
7863 })?;
7864
7865 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7866 let mut result_columns = Vec::with_capacity(combined_schema.arity());
7867 result_columns.extend(gathered_left.columns);
7868 result_columns.extend(gathered_right.columns);
7869 self.buffer_from_columns(result_columns, result_count, combined_schema)
7870 }
7871
7872 pub fn hash_join_inner_v2_count_scan_materialize_recorded(
7903 &self,
7904 left: &CudaBuffer,
7905 right: &CudaBuffer,
7906 left_keys: &[usize],
7907 right_keys: &[usize],
7908 max_output: Option<usize>,
7909 launch_stream: StreamId,
7910 ) -> Result<CudaBuffer> {
7911 if Self::use_csm_cuda_graph_env() {
7912 if let Some(result) = self
7913 .hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded(
7914 left,
7915 right,
7916 left_keys,
7917 right_keys,
7918 max_output,
7919 launch_stream,
7920 )?
7921 {
7922 return Ok(result);
7923 }
7924 self.csm_cuda_graph_fallbacks
7925 .fetch_add(1, Ordering::Relaxed);
7926 }
7927
7928 use crate::launch::LaunchRecorder;
7929
7930 let runtime = self.memory.runtime().ok_or_else(|| {
7931 XlogError::Kernel(
7932 "hash_join_inner_v2_count_scan_materialize_recorded requires a \
7933 runtime-backed GpuMemoryManager"
7934 .to_string(),
7935 )
7936 })?;
7937 let cu_stream = runtime
7938 .stream_pool()
7939 .resolve(launch_stream)
7940 .ok_or_else(|| {
7941 XlogError::Kernel(format!(
7942 "hash_join_inner_v2_count_scan_materialize_recorded: launch_stream \
7943 StreamId({}) does not resolve",
7944 launch_stream.0
7945 ))
7946 })?;
7947
7948 let num_left = self.device_row_count(left)?;
7950 let num_right = self.device_row_count(right)?;
7951 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
7952 return Err(XlogError::Kernel(format!(
7953 "Join supports at most {} rows per side (left={}, right={})",
7954 u32::MAX,
7955 num_left,
7956 num_right
7957 )));
7958 }
7959 if num_left == 0 || num_right == 0 {
7960 let combined_schema = self.combine_schemas(left.schema(), right.schema());
7961 return self.create_empty_buffer(combined_schema);
7962 }
7963 if left_keys.is_empty() || right_keys.is_empty() {
7964 return Err(XlogError::Kernel(
7965 "Join requires at least one key column".to_string(),
7966 ));
7967 }
7968 if left_keys.len() != right_keys.len() {
7969 return Err(XlogError::Kernel(
7970 "Left and right key columns must have same length".to_string(),
7971 ));
7972 }
7973 if left_keys.len() > 4 {
7974 return Err(XlogError::Kernel(
7975 "hash_join_inner_v2_count_scan_materialize_recorded: max 4 key \
7976 columns supported (pack_keys constraint)"
7977 .to_string(),
7978 ));
7979 }
7980 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
7981 let lt = left.schema().column_type(l);
7982 let rt = right.schema().column_type(r);
7983 if lt != rt {
7984 return Err(XlogError::Kernel(format!(
7985 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
7986 l, lt, r, rt
7987 )));
7988 }
7989 }
7990
7991 let _num_left = num_left as u32;
7992 let probe_cap = left.num_rows() as u32;
7993
7994 let left_packed =
7996 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
7997 let right_packed =
7998 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
7999 let table = self.build_hash_table_v2_on_stream(
8000 &right_packed.hashes,
8001 num_right as u32,
8002 &cu_stream,
8003 launch_stream,
8004 runtime,
8005 )?;
8006
8007 let device = self.device.inner();
8008 let block_size = 256u32;
8009 let probe_grid = probe_cap.div_ceil(block_size);
8010 let probe_config = LaunchConfig {
8011 grid_dim: (probe_grid, 1, 1),
8012 block_dim: (block_size, 1, 1),
8013 shared_mem_bytes: 0,
8014 };
8015
8016 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
8018 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
8019 let d_logical_count = self.memory.alloc::<u32>(1)?;
8020 let d_overflow = self.memory.alloc::<u8>(1)?;
8021 runtime
8027 .prepare_first_use(&d_overflow, launch_stream, Access::Write)
8028 .map_err(|e| {
8029 XlogError::Kernel(format!(
8030 "hash_join_inner_v2_count_scan_materialize_recorded: prepare d_overflow \
8031 failed: {}",
8032 e
8033 ))
8034 })?;
8035 runtime
8036 .prepare_first_use(&d_logical_count, launch_stream, Access::Write)
8037 .map_err(|e| {
8038 XlogError::Kernel(format!(
8039 "hash_join_inner_v2_count_scan_materialize_recorded: prepare d_logical_count \
8040 failed: {}",
8041 e
8042 ))
8043 })?;
8044 unsafe {
8047 let res = cudarc::driver::sys::cuMemsetD8Async(
8048 *d_overflow.device_ptr(),
8049 0,
8050 1,
8051 cu_stream.cu_stream(),
8052 );
8053 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8054 return Err(XlogError::Kernel(format!(
8055 "cuMemsetD8Async (d_overflow init) failed: {:?}",
8056 res
8057 )));
8058 }
8059 let res = cudarc::driver::sys::cuMemsetD8Async(
8060 *d_logical_count.device_ptr(),
8061 0,
8062 std::mem::size_of::<u32>(),
8063 cu_stream.cu_stream(),
8064 );
8065 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8066 return Err(XlogError::Kernel(format!(
8067 "cuMemsetD8Async (d_logical_count init) failed: {:?}",
8068 res
8069 )));
8070 }
8071 }
8072
8073 let count_func = device
8079 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
8080 .ok_or_else(|| {
8081 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
8082 })?;
8083 let total_func = device
8084 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
8085 .ok_or_else(|| {
8086 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
8087 })?;
8088
8089 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
8090 rec_count.read(&left_packed.hashes);
8091 rec_count.read(&left_packed.packed_keys);
8092 rec_count.read(&right_packed.packed_keys);
8093 rec_count.read(&table.bucket_offsets);
8094 rec_count.read(&table.bucket_counts);
8095 rec_count.read(&table.bucket_entries);
8096 rec_count.read(&table.bucket_entry_hashes);
8097 rec_count.read(left.num_rows_device());
8098 rec_count.write(&per_probe_count);
8099 rec_count.write(&per_probe_offsets);
8100 rec_count.write(&d_logical_count);
8101 rec_count.write(&d_overflow);
8102 rec_count.preflight(runtime).map_err(|e| {
8103 XlogError::Kernel(format!("csm inner: count/scan preflight failed: {}", e))
8104 })?;
8105
8106 unsafe {
8109 count_func.clone().launch_on_stream(
8110 &cu_stream,
8111 probe_config,
8112 (
8113 &left_packed.hashes,
8114 left.num_rows_device(),
8115 probe_cap,
8116 &table.bucket_offsets,
8117 &table.bucket_counts,
8118 &table.bucket_entries,
8119 &table.bucket_entry_hashes,
8120 table.bucket_mask,
8121 &left_packed.packed_keys,
8122 &right_packed.packed_keys,
8123 left_packed.key_bytes,
8124 &per_probe_count,
8125 ),
8126 )
8127 }
8128 .map_err(|e| {
8129 XlogError::Kernel(format!(
8130 "hash_join_probe_v2_count_per_row (on_stream) failed: {}",
8131 e
8132 ))
8133 })?;
8134
8135 unsafe {
8139 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
8140 *per_probe_offsets.device_ptr(),
8141 *per_probe_count.device_ptr(),
8142 (probe_cap as usize) * std::mem::size_of::<u32>(),
8143 cu_stream.cu_stream(),
8144 );
8145 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8146 return Err(XlogError::Kernel(format!(
8147 "csm inner: cuMemcpyDtoDAsync (per_probe_count → offsets) failed: {:?}",
8148 res
8149 )));
8150 }
8151 }
8152 self.multiblock_scan_u32_inplace_on_stream(
8153 &mut per_probe_offsets,
8154 probe_cap,
8155 &cu_stream,
8156 launch_stream,
8157 runtime,
8158 )?;
8159
8160 let materialize_capacity_bound: u64 = (probe_cap as u64).saturating_mul(num_right as u64);
8167 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
8168 unsafe {
8169 total_func.clone().launch_on_stream(
8170 &cu_stream,
8171 LaunchConfig {
8172 grid_dim: (1, 1, 1),
8173 block_dim: (1, 1, 1),
8174 shared_mem_bytes: 0,
8175 },
8176 (
8177 &per_probe_offsets,
8178 &per_probe_count,
8179 left.num_rows_device(),
8180 probe_cap,
8181 materialize_capacity_u32,
8182 &d_logical_count,
8183 &d_overflow,
8184 ),
8185 )
8186 }
8187 .map_err(|e| {
8188 XlogError::Kernel(format!(
8189 "hash_join_total_from_scan (on_stream) failed: {}",
8190 e
8191 ))
8192 })?;
8193
8194 rec_count.commit(runtime).map_err(|e| {
8195 XlogError::Kernel(format!("csm inner: count/scan commit failed: {}", e))
8196 })?;
8197
8198 cu_stream.synchronize().map_err(|e| {
8201 XlogError::Kernel(format!("csm inner: sync (total read) failed: {}", e))
8202 })?;
8203 let total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
8204 let requested = max_output
8205 .map(|limit| (limit as u64).min(total))
8206 .unwrap_or(total);
8207 if requested == 0 {
8208 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8209 return self.create_empty_buffer(combined_schema);
8210 }
8211 if requested > u32::MAX as u64 {
8212 return Err(XlogError::Kernel(format!(
8213 "Join produced {} rows which exceeds the u32 index limit",
8214 requested
8215 )));
8216 }
8217 let output_capacity = requested as u32;
8218
8219 let d_output_left = self.memory.alloc::<u32>(output_capacity as usize)?;
8228 let d_output_right = self.memory.alloc::<u32>(output_capacity as usize)?;
8229
8230 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
8231 rec_mat.read(&left_packed.hashes);
8232 rec_mat.read(&left_packed.packed_keys);
8233 rec_mat.read(&right_packed.packed_keys);
8234 rec_mat.read(&table.bucket_offsets);
8235 rec_mat.read(&table.bucket_counts);
8236 rec_mat.read(&table.bucket_entries);
8237 rec_mat.read(&table.bucket_entry_hashes);
8238 rec_mat.read(&per_probe_offsets);
8239 rec_mat.read(left.num_rows_device());
8240 rec_mat.write(&d_output_left);
8241 rec_mat.write(&d_output_right);
8242 rec_mat.write(&d_overflow);
8244 rec_mat.preflight(runtime).map_err(|e| {
8245 XlogError::Kernel(format!("csm inner: materialize preflight failed: {}", e))
8246 })?;
8247
8248 let materialize_func = device
8249 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
8250 .ok_or_else(|| {
8251 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
8252 })?;
8253 unsafe {
8256 let mut params: Vec<*mut c_void> = vec![
8257 (&left_packed.hashes).as_kernel_param(),
8258 left.num_rows_device().as_kernel_param(),
8259 probe_cap.as_kernel_param(),
8260 (&table.bucket_offsets).as_kernel_param(),
8261 (&table.bucket_counts).as_kernel_param(),
8262 (&table.bucket_entries).as_kernel_param(),
8263 (&table.bucket_entry_hashes).as_kernel_param(),
8264 table.bucket_mask.as_kernel_param(),
8265 (&left_packed.packed_keys).as_kernel_param(),
8266 (&right_packed.packed_keys).as_kernel_param(),
8267 left_packed.key_bytes.as_kernel_param(),
8268 (&per_probe_offsets).as_kernel_param(),
8269 output_capacity.as_kernel_param(),
8270 (&d_output_left).as_kernel_param(),
8271 (&d_output_right).as_kernel_param(),
8272 (&d_overflow).as_kernel_param(),
8273 ];
8274 materialize_func
8275 .clone()
8276 .launch_on_stream(&cu_stream, probe_config, &mut params)
8277 .map_err(|e| {
8278 XlogError::Kernel(format!(
8279 "hash_join_probe_v2_materialize (on_stream) failed: {}",
8280 e
8281 ))
8282 })?;
8283 }
8284
8285 rec_mat.commit(runtime).map_err(|e| {
8286 XlogError::Kernel(format!("csm inner: materialize commit failed: {}", e))
8287 })?;
8288
8289 cu_stream.synchronize().map_err(|e| {
8290 XlogError::Kernel(format!("csm inner: sync (post-materialize) failed: {}", e))
8291 })?;
8292
8293 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
8295 for col_idx in 0..left.columns.len() {
8296 let c = left
8297 .column(col_idx)
8298 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
8299 rec_gather.read_column(c);
8300 }
8301 for col_idx in 0..right.columns.len() {
8302 let c = right
8303 .column(col_idx)
8304 .ok_or_else(|| XlogError::Kernel(format!("Right column {} not found", col_idx)))?;
8305 rec_gather.read_column(c);
8306 }
8307 rec_gather.read(&d_output_left);
8308 rec_gather.read(&d_output_right);
8309 rec_gather
8310 .preflight(runtime)
8311 .map_err(|e| XlogError::Kernel(format!("csm inner: gather preflight failed: {}", e)))?;
8312 let gathered_left = self.gather_buffer_by_indices_on_stream(
8313 left,
8314 &d_output_left,
8315 output_capacity,
8316 &cu_stream,
8317 launch_stream,
8318 runtime,
8319 )?;
8320 let gathered_right = self.gather_buffer_by_indices_on_stream(
8321 right,
8322 &d_output_right,
8323 output_capacity,
8324 &cu_stream,
8325 launch_stream,
8326 runtime,
8327 )?;
8328 rec_gather
8329 .commit(runtime)
8330 .map_err(|e| XlogError::Kernel(format!("csm inner: gather commit failed: {}", e)))?;
8331
8332 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8333 let mut result_columns = Vec::with_capacity(combined_schema.arity());
8334 result_columns.extend(gathered_left.columns);
8335 result_columns.extend(gathered_right.columns);
8336 self.buffer_from_columns(result_columns, output_capacity as u64, combined_schema)
8337 }
8338
8339 fn hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded(
8340 &self,
8341 left: &CudaBuffer,
8342 right: &CudaBuffer,
8343 left_keys: &[usize],
8344 right_keys: &[usize],
8345 max_output: Option<usize>,
8346 launch_stream: StreamId,
8347 ) -> Result<Option<CudaBuffer>> {
8348 let runtime = self.memory.runtime().ok_or_else(|| {
8349 XlogError::Kernel(
8350 "hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded requires a \
8351 runtime-backed GpuMemoryManager"
8352 .to_string(),
8353 )
8354 })?;
8355 let cu_stream = runtime
8356 .stream_pool()
8357 .resolve(launch_stream)
8358 .ok_or_else(|| {
8359 XlogError::Kernel(format!(
8360 "hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded: \
8361 launch_stream StreamId({}) does not resolve",
8362 launch_stream.0
8363 ))
8364 })?;
8365
8366 let num_left = self.device_row_count(left)?;
8367 let num_right = self.device_row_count(right)?;
8368 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
8369 return Err(XlogError::Kernel(format!(
8370 "Join supports at most {} rows per side (left={}, right={})",
8371 u32::MAX,
8372 num_left,
8373 num_right
8374 )));
8375 }
8376 if num_left == 0 || num_right == 0 || max_output == Some(0) {
8377 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8378 return self.create_empty_buffer(combined_schema).map(Some);
8379 }
8380 if left_keys.is_empty() || right_keys.is_empty() {
8381 return Err(XlogError::Kernel(
8382 "Join requires at least one key column".to_string(),
8383 ));
8384 }
8385 if left_keys.len() != right_keys.len() {
8386 return Err(XlogError::Kernel(
8387 "Left and right key columns must have same length".to_string(),
8388 ));
8389 }
8390 if left_keys.len() > 4 {
8391 return Err(XlogError::Kernel(
8392 "hash_join_inner_v2_count_scan_materialize_cuda_graph_recorded: max 4 key \
8393 columns supported (pack_keys constraint)"
8394 .to_string(),
8395 ));
8396 }
8397 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
8398 let lt = left.schema().column_type(l);
8399 let rt = right.schema().column_type(r);
8400 if lt != rt {
8401 return Err(XlogError::Kernel(format!(
8402 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
8403 l, lt, r, rt
8404 )));
8405 }
8406 }
8407
8408 let logical_probe_cap = left.num_rows() as u32;
8409 let probe_cap = crate::cuda_graph::graph_capacity_class_u32(logical_probe_cap);
8410 let Some(output_capacity) =
8411 Self::csm_cuda_graph_output_capacity(logical_probe_cap, num_right as u32, max_output)?
8412 else {
8413 return Ok(None);
8414 };
8415 if output_capacity == 0 {
8416 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8417 return self.create_empty_buffer(combined_schema).map(Some);
8418 }
8419
8420 let left_packed =
8421 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
8422 let right_packed =
8423 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
8424 let graph_key = CsmCudaGraphKey::inner(
8425 left_keys.len(),
8426 left_packed.key_bytes,
8427 probe_cap,
8428 output_capacity,
8429 )?;
8430 let table = self.build_hash_table_v2_on_stream(
8431 &right_packed.hashes,
8432 num_right as u32,
8433 &cu_stream,
8434 launch_stream,
8435 runtime,
8436 )?;
8437
8438 let device = self.device.inner();
8439 let block_size = 256u32;
8440 let probe_grid = probe_cap.div_ceil(block_size);
8441 let probe_config = LaunchConfig {
8442 grid_dim: (probe_grid, 1, 1),
8443 block_dim: (block_size, 1, 1),
8444 shared_mem_bytes: 0,
8445 };
8446
8447 let materialize_capacity_bound: u64 = (probe_cap as u64).saturating_mul(num_right as u64);
8448 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
8449
8450 {
8451 let mut cache = self.csm_cuda_graph_cache.lock().map_err(|e| {
8452 XlogError::Kernel(format!("csm CUDA Graph cache lock poisoned: {}", e))
8453 })?;
8454 if let Some(entry) = cache.get_mut(&graph_key) {
8455 let result = self.launch_csm_cuda_graph_entry(
8456 entry,
8457 left,
8458 right,
8459 &left_packed,
8460 &right_packed,
8461 &table,
8462 max_output,
8463 materialize_capacity_u32,
8464 probe_config,
8465 &cu_stream,
8466 launch_stream,
8467 runtime,
8468 )?;
8469 self.csm_cuda_graph_cache_hits
8470 .fetch_add(1, Ordering::Relaxed);
8471 return Ok(Some(result));
8472 }
8473 }
8474
8475 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
8476 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
8477 let d_logical_count = self.memory.alloc::<u32>(1)?;
8478 let d_overflow = self.memory.alloc::<u8>(1)?;
8479 let d_output_left = self.memory.alloc::<u32>(output_capacity as usize)?;
8480 let d_output_right = self.memory.alloc::<u32>(output_capacity as usize)?;
8481 let mut scan_scratch = self.multiblock_scan_u32_scratch_for_len(probe_cap)?;
8482
8483 let count_func = device
8484 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
8485 .ok_or_else(|| {
8486 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
8487 })?;
8488 let total_func = device
8489 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
8490 .ok_or_else(|| {
8491 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
8492 })?;
8493 let materialize_func = device
8494 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
8495 .ok_or_else(|| {
8496 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
8497 })?;
8498
8499 let graph = CapturedCudaGraph::capture_on_stream(&cu_stream, || {
8500 unsafe {
8503 let res = cudarc::driver::sys::cuMemsetD8Async(
8504 *d_overflow.device_ptr(),
8505 0,
8506 1,
8507 cu_stream.cu_stream(),
8508 );
8509 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8510 return Err(XlogError::Kernel(format!(
8511 "csm inner graph: cuMemsetD8Async (d_overflow) failed: {:?}",
8512 res
8513 )));
8514 }
8515 let res = cudarc::driver::sys::cuMemsetD8Async(
8516 *d_logical_count.device_ptr(),
8517 0,
8518 std::mem::size_of::<u32>(),
8519 cu_stream.cu_stream(),
8520 );
8521 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8522 return Err(XlogError::Kernel(format!(
8523 "csm inner graph: cuMemsetD8Async (d_logical_count) failed: {:?}",
8524 res
8525 )));
8526 }
8527 }
8528
8529 unsafe {
8531 count_func.clone().launch_on_stream(
8532 &cu_stream,
8533 probe_config,
8534 (
8535 &left_packed.hashes,
8536 left.num_rows_device(),
8537 probe_cap,
8538 &table.bucket_offsets,
8539 &table.bucket_counts,
8540 &table.bucket_entries,
8541 &table.bucket_entry_hashes,
8542 table.bucket_mask,
8543 &left_packed.packed_keys,
8544 &right_packed.packed_keys,
8545 left_packed.key_bytes,
8546 &per_probe_count,
8547 ),
8548 )
8549 }
8550 .map_err(|e| {
8551 XlogError::Kernel(format!("csm inner graph: count_per_row failed: {}", e))
8552 })?;
8553
8554 unsafe {
8556 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
8557 *per_probe_offsets.device_ptr(),
8558 *per_probe_count.device_ptr(),
8559 (probe_cap as usize) * std::mem::size_of::<u32>(),
8560 cu_stream.cu_stream(),
8561 );
8562 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
8563 return Err(XlogError::Kernel(format!(
8564 "csm inner graph: cuMemcpyDtoDAsync (count -> offsets) failed: {:?}",
8565 res
8566 )));
8567 }
8568 }
8569 self.multiblock_scan_u32_inplace_on_stream_with_scratch(
8570 &mut per_probe_offsets,
8571 probe_cap,
8572 &cu_stream,
8573 &mut scan_scratch,
8574 )?;
8575
8576 unsafe {
8578 total_func.clone().launch_on_stream(
8579 &cu_stream,
8580 LaunchConfig {
8581 grid_dim: (1, 1, 1),
8582 block_dim: (1, 1, 1),
8583 shared_mem_bytes: 0,
8584 },
8585 (
8586 &per_probe_offsets,
8587 &per_probe_count,
8588 left.num_rows_device(),
8589 probe_cap,
8590 materialize_capacity_u32,
8591 &d_logical_count,
8592 &d_overflow,
8593 ),
8594 )
8595 }
8596 .map_err(|e| XlogError::Kernel(format!("csm inner graph: total failed: {}", e)))?;
8597
8598 unsafe {
8600 let mut params: Vec<*mut c_void> = vec![
8601 (&left_packed.hashes).as_kernel_param(),
8602 left.num_rows_device().as_kernel_param(),
8603 probe_cap.as_kernel_param(),
8604 (&table.bucket_offsets).as_kernel_param(),
8605 (&table.bucket_counts).as_kernel_param(),
8606 (&table.bucket_entries).as_kernel_param(),
8607 (&table.bucket_entry_hashes).as_kernel_param(),
8608 table.bucket_mask.as_kernel_param(),
8609 (&left_packed.packed_keys).as_kernel_param(),
8610 (&right_packed.packed_keys).as_kernel_param(),
8611 left_packed.key_bytes.as_kernel_param(),
8612 (&per_probe_offsets).as_kernel_param(),
8613 output_capacity.as_kernel_param(),
8614 (&d_output_left).as_kernel_param(),
8615 (&d_output_right).as_kernel_param(),
8616 (&d_overflow).as_kernel_param(),
8617 ];
8618 materialize_func
8619 .clone()
8620 .launch_on_stream(&cu_stream, probe_config, &mut params)
8621 .map_err(|e| {
8622 XlogError::Kernel(format!("csm inner graph: materialize failed: {}", e))
8623 })?;
8624 }
8625 Ok(())
8626 })?;
8627 let nodes = Self::csm_cuda_graph_nodes(&graph)?;
8628 let mut entry = CsmCudaGraphEntry {
8629 graph,
8630 nodes,
8631 per_probe_count,
8632 per_probe_offsets,
8633 d_logical_count,
8634 d_overflow,
8635 d_output_left,
8636 d_output_right,
8637 scan_scratch,
8638 probe_capacity: probe_cap,
8639 output_capacity,
8640 };
8641 self.csm_cuda_graph_captures.fetch_add(1, Ordering::Relaxed);
8642
8643 let result = self.launch_csm_cuda_graph_entry(
8644 &mut entry,
8645 left,
8646 right,
8647 &left_packed,
8648 &right_packed,
8649 &table,
8650 max_output,
8651 materialize_capacity_u32,
8652 probe_config,
8653 &cu_stream,
8654 launch_stream,
8655 runtime,
8656 )?;
8657 self.csm_cuda_graph_cache
8658 .lock()
8659 .map_err(|e| XlogError::Kernel(format!("csm CUDA Graph cache lock poisoned: {}", e)))?
8660 .insert(graph_key, entry);
8661 Ok(Some(result))
8662 }
8663
8664 #[allow(clippy::too_many_arguments)]
8665 fn launch_csm_cuda_graph_entry(
8666 &self,
8667 entry: &mut CsmCudaGraphEntry,
8668 left: &CudaBuffer,
8669 right: &CudaBuffer,
8670 left_packed: &PackedKeyData,
8671 right_packed: &PackedKeyData,
8672 table: &JoinHashTableV2,
8673 max_output: Option<usize>,
8674 materialize_capacity_u32: u32,
8675 probe_config: LaunchConfig,
8676 cu_stream: &cudarc::driver::CudaStream,
8677 launch_stream: StreamId,
8678 runtime: &crate::device_runtime::XlogDeviceRuntime,
8679 ) -> Result<CudaBuffer> {
8680 let mut rec_graph = LaunchRecorder::new_strict(launch_stream);
8681 rec_graph.read(&left_packed.hashes);
8682 rec_graph.read(&left_packed.packed_keys);
8683 rec_graph.read(&right_packed.packed_keys);
8684 rec_graph.read(&table.bucket_offsets);
8685 rec_graph.read(&table.bucket_counts);
8686 rec_graph.read(&table.bucket_entries);
8687 rec_graph.read(&table.bucket_entry_hashes);
8688 rec_graph.read(left.num_rows_device());
8689 rec_graph.read_write(&entry.per_probe_count);
8690 rec_graph.read_write(&entry.per_probe_offsets);
8691 rec_graph.read_write(&entry.d_logical_count);
8692 rec_graph.read_write(&entry.d_overflow);
8693 rec_graph.write(&entry.d_output_left);
8694 rec_graph.write(&entry.d_output_right);
8695 for level in entry.scan_scratch.levels() {
8696 rec_graph.read_write(level);
8697 }
8698 rec_graph
8699 .preflight(runtime)
8700 .map_err(|e| XlogError::Kernel(format!("csm inner graph: preflight failed: {}", e)))?;
8701
8702 let probe_cap = entry.probe_capacity;
8703 let output_capacity = entry.output_capacity;
8704 if probe_config.grid_dim.0 != probe_cap.div_ceil(probe_config.block_dim.0) {
8705 return Err(XlogError::Kernel(format!(
8706 "csm CUDA Graph replay probe grid mismatch: graph probe_cap={}, grid={:?}",
8707 probe_cap, probe_config.grid_dim
8708 )));
8709 }
8710 if entry.nodes.node_count < 5 {
8711 return Err(XlogError::Kernel(format!(
8712 "csm CUDA Graph replay node inventory too small: {}",
8713 entry.nodes.node_count
8714 )));
8715 }
8716
8717 let mut count_params = entry.graph.kernel_node_params(entry.nodes.count)?;
8718 let mut total_params = entry.graph.kernel_node_params(entry.nodes.total)?;
8719 let mut materialize_params = entry.graph.kernel_node_params(entry.nodes.materialize)?;
8720 let mut count_args: Vec<*mut c_void> = vec![
8721 (&left_packed.hashes).as_kernel_param(),
8722 left.num_rows_device().as_kernel_param(),
8723 probe_cap.as_kernel_param(),
8724 (&table.bucket_offsets).as_kernel_param(),
8725 (&table.bucket_counts).as_kernel_param(),
8726 (&table.bucket_entries).as_kernel_param(),
8727 (&table.bucket_entry_hashes).as_kernel_param(),
8728 table.bucket_mask.as_kernel_param(),
8729 (&left_packed.packed_keys).as_kernel_param(),
8730 (&right_packed.packed_keys).as_kernel_param(),
8731 left_packed.key_bytes.as_kernel_param(),
8732 (&entry.per_probe_count).as_kernel_param(),
8733 ];
8734 let mut total_args: Vec<*mut c_void> = vec![
8735 (&entry.per_probe_offsets).as_kernel_param(),
8736 (&entry.per_probe_count).as_kernel_param(),
8737 left.num_rows_device().as_kernel_param(),
8738 probe_cap.as_kernel_param(),
8739 materialize_capacity_u32.as_kernel_param(),
8740 (&entry.d_logical_count).as_kernel_param(),
8741 (&entry.d_overflow).as_kernel_param(),
8742 ];
8743 let mut materialize_args: Vec<*mut c_void> = vec![
8744 (&left_packed.hashes).as_kernel_param(),
8745 left.num_rows_device().as_kernel_param(),
8746 probe_cap.as_kernel_param(),
8747 (&table.bucket_offsets).as_kernel_param(),
8748 (&table.bucket_counts).as_kernel_param(),
8749 (&table.bucket_entries).as_kernel_param(),
8750 (&table.bucket_entry_hashes).as_kernel_param(),
8751 table.bucket_mask.as_kernel_param(),
8752 (&left_packed.packed_keys).as_kernel_param(),
8753 (&right_packed.packed_keys).as_kernel_param(),
8754 left_packed.key_bytes.as_kernel_param(),
8755 (&entry.per_probe_offsets).as_kernel_param(),
8756 output_capacity.as_kernel_param(),
8757 (&entry.d_output_left).as_kernel_param(),
8758 (&entry.d_output_right).as_kernel_param(),
8759 (&entry.d_overflow).as_kernel_param(),
8760 ];
8761 count_params.kernelParams = count_args.as_mut_ptr();
8762 count_params.extra = std::ptr::null_mut();
8763 total_params.kernelParams = total_args.as_mut_ptr();
8764 total_params.extra = std::ptr::null_mut();
8765 materialize_params.kernelParams = materialize_args.as_mut_ptr();
8766 materialize_params.extra = std::ptr::null_mut();
8767 unsafe {
8768 entry
8769 .graph
8770 .set_kernel_node_params(entry.nodes.count, &count_params)?;
8771 entry
8772 .graph
8773 .set_kernel_node_params(entry.nodes.total, &total_params)?;
8774 entry
8775 .graph
8776 .set_kernel_node_params(entry.nodes.materialize, &materialize_params)?;
8777 }
8778
8779 entry.graph.launch(cu_stream)?;
8780 self.csm_cuda_graph_launches.fetch_add(1, Ordering::Relaxed);
8781 rec_graph
8782 .commit(runtime)
8783 .map_err(|e| XlogError::Kernel(format!("csm inner graph: commit failed: {}", e)))?;
8784
8785 cu_stream.synchronize().map_err(|e| {
8786 XlogError::Kernel(format!("csm inner graph: sync (total read) failed: {}", e))
8787 })?;
8788 let total = self.read_join_output_count_metadata(&entry.d_logical_count)? as u64;
8789 let requested = max_output
8790 .map(|limit| (limit as u64).min(total))
8791 .unwrap_or(total);
8792 if requested == 0 {
8793 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8794 return self.create_empty_buffer(combined_schema);
8795 }
8796 if requested > output_capacity as u64 {
8797 return Err(XlogError::Kernel(format!(
8798 "csm inner graph produced {} rows but graph output capacity is {}",
8799 requested, output_capacity
8800 )));
8801 }
8802 let output_rows = requested as u32;
8803
8804 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
8805 for col_idx in 0..left.columns.len() {
8806 let c = left
8807 .column(col_idx)
8808 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
8809 rec_gather.read_column(c);
8810 }
8811 for col_idx in 0..right.columns.len() {
8812 let c = right
8813 .column(col_idx)
8814 .ok_or_else(|| XlogError::Kernel(format!("Right column {} not found", col_idx)))?;
8815 rec_gather.read_column(c);
8816 }
8817 rec_gather.read(&entry.d_output_left);
8818 rec_gather.read(&entry.d_output_right);
8819 rec_gather.preflight(runtime).map_err(|e| {
8820 XlogError::Kernel(format!("csm inner graph: gather preflight failed: {}", e))
8821 })?;
8822 let gathered_left = self.gather_buffer_by_indices_on_stream(
8823 left,
8824 &entry.d_output_left,
8825 output_rows,
8826 cu_stream,
8827 launch_stream,
8828 runtime,
8829 )?;
8830 let gathered_right = self.gather_buffer_by_indices_on_stream(
8831 right,
8832 &entry.d_output_right,
8833 output_rows,
8834 cu_stream,
8835 launch_stream,
8836 runtime,
8837 )?;
8838 rec_gather.commit(runtime).map_err(|e| {
8839 XlogError::Kernel(format!("csm inner graph: gather commit failed: {}", e))
8840 })?;
8841
8842 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8843 let mut result_columns = Vec::with_capacity(combined_schema.arity());
8844 result_columns.extend(gathered_left.columns);
8845 result_columns.extend(gathered_right.columns);
8846 self.buffer_from_columns(result_columns, output_rows as u64, combined_schema)
8847 }
8848
8849 fn csm_cuda_graph_nodes(graph: &CapturedCudaGraph) -> Result<CsmCudaGraphNodes> {
8850 let nodes = graph.nodes()?;
8851 if nodes.len() < 5 {
8852 return Err(XlogError::Kernel(format!(
8853 "csm inner graph captured too few nodes: {}",
8854 nodes.len()
8855 )));
8856 }
8857 let kernel_nodes: Vec<_> = nodes
8858 .iter()
8859 .copied()
8860 .filter(|n| n.kind == CudaGraphNodeKind::Kernel)
8861 .collect();
8862 if kernel_nodes.len() < 3 {
8863 return Err(XlogError::Kernel(format!(
8864 "csm inner graph captured too few kernel nodes: {}",
8865 kernel_nodes.len()
8866 )));
8867 }
8868 Ok(CsmCudaGraphNodes {
8869 count: kernel_nodes[0],
8870 total: kernel_nodes[kernel_nodes.len() - 2],
8871 materialize: kernel_nodes[kernel_nodes.len() - 1],
8872 node_count: nodes.len(),
8873 })
8874 }
8875
8876 fn csm_cuda_graph_output_capacity(
8877 probe_cap: u32,
8878 num_right: u32,
8879 max_output: Option<usize>,
8880 ) -> Result<Option<u32>> {
8881 if let Some(limit) = max_output {
8882 let limit = u32::try_from(limit).map_err(|_| {
8883 XlogError::Kernel(format!(
8884 "csm CUDA Graph max_output {} exceeds u32::MAX",
8885 limit
8886 ))
8887 })?;
8888 return Ok(Some(crate::cuda_graph::graph_capacity_class_u32(limit)));
8889 }
8890
8891 let worst_case = (probe_cap as u64).saturating_mul(num_right as u64);
8892 if worst_case > u32::MAX as u64 {
8893 return Ok(None);
8894 }
8895 let auto_cap = std::env::var("XLOG_CSM_CUDA_GRAPH_AUTO_OUTPUT_CAP")
8896 .ok()
8897 .and_then(|v| v.parse::<u64>().ok())
8898 .unwrap_or(1_000_000);
8899 if worst_case <= auto_cap {
8900 Ok(Some(crate::cuda_graph::graph_capacity_class_u32(
8901 worst_case as u32,
8902 )))
8903 } else {
8904 Ok(None)
8905 }
8906 }
8907
8908 pub fn hash_join_left_outer_v2_count_scan_materialize_recorded(
8935 &self,
8936 left: &CudaBuffer,
8937 right: &CudaBuffer,
8938 left_keys: &[usize],
8939 right_keys: &[usize],
8940 max_output: Option<usize>,
8941 launch_stream: StreamId,
8942 ) -> Result<CudaBuffer> {
8943 use crate::launch::LaunchRecorder;
8944
8945 let runtime = self.memory.runtime().ok_or_else(|| {
8946 XlogError::Kernel(
8947 "hash_join_left_outer_v2_count_scan_materialize_recorded requires a \
8948 runtime-backed GpuMemoryManager"
8949 .to_string(),
8950 )
8951 })?;
8952 let cu_stream = runtime
8953 .stream_pool()
8954 .resolve(launch_stream)
8955 .ok_or_else(|| {
8956 XlogError::Kernel(format!(
8957 "csm left_outer: launch_stream StreamId({}) does not resolve",
8958 launch_stream.0
8959 ))
8960 })?;
8961
8962 let num_left = self.device_row_count(left)?;
8964 let num_right = self.device_row_count(right)?;
8965 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
8966 return Err(XlogError::Kernel(format!(
8967 "Join supports at most {} rows per side (left={}, right={})",
8968 u32::MAX,
8969 num_left,
8970 num_right
8971 )));
8972 }
8973 if num_left == 0 {
8974 let combined_schema = self.combine_schemas(left.schema(), right.schema());
8975 return self.create_empty_buffer(combined_schema);
8976 }
8977 if num_right == 0 {
8978 return self.left_outer_with_nulls(left, right);
8983 }
8984 if left_keys.is_empty() || right_keys.is_empty() {
8985 return Err(XlogError::Kernel(
8986 "Join requires at least one key column".to_string(),
8987 ));
8988 }
8989 if left_keys.len() != right_keys.len() {
8990 return Err(XlogError::Kernel(
8991 "Left and right key columns must have same length".to_string(),
8992 ));
8993 }
8994 if left_keys.len() > 4 {
8995 return Err(XlogError::Kernel(
8996 "csm left_outer: max 4 key columns supported (pack_keys constraint)".to_string(),
8997 ));
8998 }
8999 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
9000 let lt = left.schema().column_type(l);
9001 let rt = right.schema().column_type(r);
9002 if lt != rt {
9003 return Err(XlogError::Kernel(format!(
9004 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
9005 l, lt, r, rt
9006 )));
9007 }
9008 }
9009
9010 let probe_cap = u32::try_from(num_left).map_err(|_| {
9018 XlogError::Kernel("csm left_outer: left row count exceeds u32::MAX".to_string())
9019 })?;
9020 let num_right_u32 = u32::try_from(num_right).map_err(|_| {
9021 XlogError::Kernel("csm left_outer: right row count exceeds u32::MAX".to_string())
9022 })?;
9023
9024 let left_packed =
9026 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
9027 let right_packed =
9028 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
9029 let table = self.build_hash_table_v2_on_stream(
9030 &right_packed.hashes,
9031 num_right_u32,
9032 &cu_stream,
9033 launch_stream,
9034 runtime,
9035 )?;
9036
9037 let device = self.device.inner();
9038 let block_size = 256u32;
9039 let probe_grid = probe_cap.div_ceil(block_size);
9040 let probe_config = LaunchConfig {
9041 grid_dim: (probe_grid, 1, 1),
9042 block_dim: (block_size, 1, 1),
9043 shared_mem_bytes: 0,
9044 };
9045
9046 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
9048 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
9049 let d_logical_count = self.memory.alloc::<u32>(1)?;
9050 let d_overflow = self.memory.alloc::<u8>(1)?;
9051 runtime
9055 .prepare_first_use(&d_overflow, launch_stream, Access::Write)
9056 .map_err(|e| {
9057 XlogError::Kernel(format!("csm left_outer: prepare d_overflow failed: {}", e))
9058 })?;
9059 runtime
9060 .prepare_first_use(&d_logical_count, launch_stream, Access::Write)
9061 .map_err(|e| {
9062 XlogError::Kernel(format!(
9063 "csm left_outer: prepare d_logical_count failed: {}",
9064 e
9065 ))
9066 })?;
9067 unsafe {
9070 let res = cudarc::driver::sys::cuMemsetD8Async(
9071 *d_overflow.device_ptr(),
9072 0,
9073 1,
9074 cu_stream.cu_stream(),
9075 );
9076 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9077 return Err(XlogError::Kernel(format!(
9078 "csm left_outer: cuMemsetD8Async (d_overflow) failed: {:?}",
9079 res
9080 )));
9081 }
9082 let res = cudarc::driver::sys::cuMemsetD8Async(
9083 *d_logical_count.device_ptr(),
9084 0,
9085 std::mem::size_of::<u32>(),
9086 cu_stream.cu_stream(),
9087 );
9088 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9089 return Err(XlogError::Kernel(format!(
9090 "csm left_outer: cuMemsetD8Async (d_logical_count) failed: {:?}",
9091 res
9092 )));
9093 }
9094 }
9095
9096 let count_func = device
9097 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
9098 .ok_or_else(|| {
9099 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
9100 })?;
9101 let total_func = device
9102 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
9103 .ok_or_else(|| {
9104 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
9105 })?;
9106
9107 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
9108 rec_count.read(&left_packed.hashes);
9109 rec_count.read(&left_packed.packed_keys);
9110 rec_count.read(&right_packed.packed_keys);
9111 rec_count.read(&table.bucket_offsets);
9112 rec_count.read(&table.bucket_counts);
9113 rec_count.read(&table.bucket_entries);
9114 rec_count.read(&table.bucket_entry_hashes);
9115 rec_count.read(left.num_rows_device());
9116 rec_count.write(&per_probe_count);
9117 rec_count.write(&per_probe_offsets);
9118 rec_count.write(&d_logical_count);
9119 rec_count.write(&d_overflow);
9120 rec_count.preflight(runtime).map_err(|e| {
9121 XlogError::Kernel(format!(
9122 "csm left_outer: count/scan preflight failed: {}",
9123 e
9124 ))
9125 })?;
9126
9127 unsafe {
9130 count_func.clone().launch_on_stream(
9131 &cu_stream,
9132 probe_config,
9133 (
9134 &left_packed.hashes,
9135 left.num_rows_device(),
9136 probe_cap,
9137 &table.bucket_offsets,
9138 &table.bucket_counts,
9139 &table.bucket_entries,
9140 &table.bucket_entry_hashes,
9141 table.bucket_mask,
9142 &left_packed.packed_keys,
9143 &right_packed.packed_keys,
9144 left_packed.key_bytes,
9145 &per_probe_count,
9146 ),
9147 )
9148 }
9149 .map_err(|e| {
9150 XlogError::Kernel(format!(
9151 "hash_join_probe_v2_count_per_row (csm left_outer) failed: {}",
9152 e
9153 ))
9154 })?;
9155
9156 unsafe {
9160 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9161 *per_probe_offsets.device_ptr(),
9162 *per_probe_count.device_ptr(),
9163 (probe_cap as usize) * std::mem::size_of::<u32>(),
9164 cu_stream.cu_stream(),
9165 );
9166 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9167 return Err(XlogError::Kernel(format!(
9168 "csm left_outer: cuMemcpyDtoDAsync (count → offsets) failed: {:?}",
9169 res
9170 )));
9171 }
9172 }
9173 self.multiblock_scan_u32_inplace_on_stream(
9174 &mut per_probe_offsets,
9175 probe_cap,
9176 &cu_stream,
9177 launch_stream,
9178 runtime,
9179 )?;
9180
9181 let materialize_capacity_bound: u64 = (probe_cap as u64).saturating_mul(num_right as u64);
9183 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
9184 unsafe {
9186 total_func.clone().launch_on_stream(
9187 &cu_stream,
9188 LaunchConfig {
9189 grid_dim: (1, 1, 1),
9190 block_dim: (1, 1, 1),
9191 shared_mem_bytes: 0,
9192 },
9193 (
9194 &per_probe_offsets,
9195 &per_probe_count,
9196 left.num_rows_device(),
9197 probe_cap,
9198 materialize_capacity_u32,
9199 &d_logical_count,
9200 &d_overflow,
9201 ),
9202 )
9203 }
9204 .map_err(|e| {
9205 XlogError::Kernel(format!(
9206 "hash_join_total_from_scan (csm left_outer) failed: {}",
9207 e
9208 ))
9209 })?;
9210
9211 rec_count.commit(runtime).map_err(|e| {
9212 XlogError::Kernel(format!("csm left_outer: count/scan commit failed: {}", e))
9213 })?;
9214
9215 cu_stream.synchronize().map_err(|e| {
9216 XlogError::Kernel(format!("csm left_outer: sync (count read) failed: {}", e))
9217 })?;
9218 let inner_total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
9219 let inner_clamped = max_output
9220 .map(|limit| (limit as u64).min(inner_total))
9221 .unwrap_or(inner_total);
9222 if inner_clamped > u32::MAX as u64 {
9223 return Err(XlogError::Kernel(format!(
9224 "Join produced {} matched rows which exceeds the u32 index limit",
9225 inner_clamped
9226 )));
9227 }
9228 let inner_count_u32 = inner_clamped as u32;
9229
9230 let materialize_func = device
9232 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
9233 .ok_or_else(|| {
9234 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
9235 })?;
9236 let d_output_left = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
9237 let d_output_right = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
9238
9239 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
9240 rec_mat.read(&left_packed.hashes);
9241 rec_mat.read(&left_packed.packed_keys);
9242 rec_mat.read(&right_packed.packed_keys);
9243 rec_mat.read(&table.bucket_offsets);
9244 rec_mat.read(&table.bucket_counts);
9245 rec_mat.read(&table.bucket_entries);
9246 rec_mat.read(&table.bucket_entry_hashes);
9247 rec_mat.read(&per_probe_offsets);
9248 rec_mat.read(left.num_rows_device());
9249 rec_mat.write(&d_output_left);
9250 rec_mat.write(&d_output_right);
9251 rec_mat.write(&d_overflow);
9253 rec_mat.preflight(runtime).map_err(|e| {
9254 XlogError::Kernel(format!(
9255 "csm left_outer: materialize preflight failed: {}",
9256 e
9257 ))
9258 })?;
9259 if inner_count_u32 > 0 {
9260 unsafe {
9262 let mut params: Vec<*mut c_void> = vec![
9263 (&left_packed.hashes).as_kernel_param(),
9264 left.num_rows_device().as_kernel_param(),
9265 probe_cap.as_kernel_param(),
9266 (&table.bucket_offsets).as_kernel_param(),
9267 (&table.bucket_counts).as_kernel_param(),
9268 (&table.bucket_entries).as_kernel_param(),
9269 (&table.bucket_entry_hashes).as_kernel_param(),
9270 table.bucket_mask.as_kernel_param(),
9271 (&left_packed.packed_keys).as_kernel_param(),
9272 (&right_packed.packed_keys).as_kernel_param(),
9273 left_packed.key_bytes.as_kernel_param(),
9274 (&per_probe_offsets).as_kernel_param(),
9275 inner_count_u32.as_kernel_param(),
9276 (&d_output_left).as_kernel_param(),
9277 (&d_output_right).as_kernel_param(),
9278 (&d_overflow).as_kernel_param(),
9279 ];
9280 materialize_func
9281 .clone()
9282 .launch_on_stream(&cu_stream, probe_config, &mut params)
9283 .map_err(|e| {
9284 XlogError::Kernel(format!(
9285 "hash_join_probe_v2_materialize (csm left_outer) failed: {}",
9286 e
9287 ))
9288 })?;
9289 }
9290 }
9291 rec_mat.commit(runtime).map_err(|e| {
9292 XlogError::Kernel(format!("csm left_outer: materialize commit failed: {}", e))
9293 })?;
9294
9295 let d_unmatched_mask = self.memory.alloc::<u8>(probe_cap as usize)?;
9297 let unmatched_mask_func = device
9298 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_CSM_UNMATCHED_MASK)
9299 .ok_or_else(|| {
9300 XlogError::Kernel("hash_join_csm_unmatched_mask kernel not found".to_string())
9301 })?;
9302 let mut rec_um = LaunchRecorder::new_strict(launch_stream);
9303 rec_um.read(&per_probe_count);
9304 rec_um.read(left.num_rows_device());
9305 rec_um.write(&d_unmatched_mask);
9306 rec_um.preflight(runtime).map_err(|e| {
9307 XlogError::Kernel(format!(
9308 "csm left_outer: unmatched mask preflight failed: {}",
9309 e
9310 ))
9311 })?;
9312 unsafe {
9314 unmatched_mask_func.clone().launch_on_stream(
9315 &cu_stream,
9316 probe_config,
9317 (
9318 &per_probe_count,
9319 left.num_rows_device(),
9320 probe_cap,
9321 &d_unmatched_mask,
9322 ),
9323 )
9324 }
9325 .map_err(|e| {
9326 XlogError::Kernel(format!(
9327 "hash_join_csm_unmatched_mask (on_stream) failed: {}",
9328 e
9329 ))
9330 })?;
9331 rec_um.commit(runtime).map_err(|e| {
9332 XlogError::Kernel(format!(
9333 "csm left_outer: unmatched mask commit failed: {}",
9334 e
9335 ))
9336 })?;
9337
9338 let unmatched_left = self.compact_buffer_by_device_mask_counted_recorded(
9339 left,
9340 &d_unmatched_mask,
9341 launch_stream,
9342 )?;
9343 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
9344 let total_rows = (inner_count_u32 as u64) + unmatched_rows;
9345
9346 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9347 if total_rows == 0 {
9348 return self.create_empty_buffer(combined_schema);
9349 }
9350
9351 let inner_left_buf;
9353 let inner_right_buf;
9354 if inner_count_u32 > 0 {
9355 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
9356 for col_idx in 0..left.columns.len() {
9357 let c = left.column(col_idx).ok_or_else(|| {
9358 XlogError::Kernel(format!("Left column {} not found", col_idx))
9359 })?;
9360 rec_gather.read_column(c);
9361 }
9362 for col_idx in 0..right.columns.len() {
9363 let c = right.column(col_idx).ok_or_else(|| {
9364 XlogError::Kernel(format!("Right column {} not found", col_idx))
9365 })?;
9366 rec_gather.read_column(c);
9367 }
9368 rec_gather.read(&d_output_left);
9369 rec_gather.read(&d_output_right);
9370 rec_gather.preflight(runtime).map_err(|e| {
9371 XlogError::Kernel(format!("csm left_outer: gather preflight failed: {}", e))
9372 })?;
9373 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
9374 left,
9375 &d_output_left,
9376 inner_count_u32,
9377 &cu_stream,
9378 launch_stream,
9379 runtime,
9380 )?);
9381 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
9382 right,
9383 &d_output_right,
9384 inner_count_u32,
9385 &cu_stream,
9386 launch_stream,
9387 runtime,
9388 )?);
9389 rec_gather.commit(runtime).map_err(|e| {
9390 XlogError::Kernel(format!("csm left_outer: gather commit failed: {}", e))
9391 })?;
9392 } else {
9393 inner_left_buf = None;
9394 inner_right_buf = None;
9395 }
9396
9397 let mut rec_d = LaunchRecorder::new_strict(launch_stream);
9400 for col_idx in 0..unmatched_left.columns.len() {
9401 let c = unmatched_left.column(col_idx).ok_or_else(|| {
9402 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
9403 })?;
9404 rec_d.read_column(c);
9405 }
9406 if let Some(b) = inner_left_buf.as_ref() {
9407 for col_idx in 0..b.columns.len() {
9408 let c = b.column(col_idx).ok_or_else(|| {
9409 XlogError::Kernel(format!("inner_left col {} not found", col_idx))
9410 })?;
9411 rec_d.read_column(c);
9412 }
9413 }
9414 if let Some(b) = inner_right_buf.as_ref() {
9415 for col_idx in 0..b.columns.len() {
9416 let c = b.column(col_idx).ok_or_else(|| {
9417 XlogError::Kernel(format!("inner_right col {} not found", col_idx))
9418 })?;
9419 rec_d.read_column(c);
9420 }
9421 }
9422 rec_d.preflight(runtime).map_err(|e| {
9423 XlogError::Kernel(format!("csm left_outer: phase-E preflight failed: {}", e))
9424 })?;
9425
9426 let inner_rows = inner_count_u32 as u64;
9427 let mut result_columns: Vec<CudaColumn> = Vec::with_capacity(combined_schema.arity());
9428
9429 for col_idx in 0..left.arity() {
9431 let elem_size = left
9432 .schema()
9433 .column_type(col_idx)
9434 .map(|t| t.size_bytes())
9435 .unwrap_or(4);
9436 let inner_bytes = (inner_rows as usize)
9437 .checked_mul(elem_size)
9438 .ok_or_else(|| XlogError::Kernel("csm left_outer: inner_bytes overflow".into()))?;
9439 let unmatched_bytes = (unmatched_rows as usize)
9440 .checked_mul(elem_size)
9441 .ok_or_else(|| {
9442 XlogError::Kernel("csm left_outer: unmatched_bytes overflow".into())
9443 })?;
9444 let total_bytes = inner_bytes
9445 .checked_add(unmatched_bytes)
9446 .ok_or_else(|| XlogError::Kernel("csm left_outer: total_bytes overflow".into()))?;
9447 let out_col = self.memory.alloc::<u8>(total_bytes)?;
9448 let dst_ptr = *out_col.device_ptr();
9449 runtime
9451 .prepare_first_use(&out_col, launch_stream, Access::Write)
9452 .map_err(|e| {
9453 XlogError::Kernel(format!(
9454 "csm left_outer: prepare left out_col {} failed: {}",
9455 col_idx, e
9456 ))
9457 })?;
9458 if inner_bytes > 0 {
9459 let src_col = inner_left_buf
9460 .as_ref()
9461 .expect("inner_count > 0")
9462 .column(col_idx)
9463 .ok_or_else(|| XlogError::Kernel("inner_left col missing".into()))?;
9464 unsafe {
9466 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9467 dst_ptr,
9468 *src_col.device_ptr(),
9469 inner_bytes,
9470 cu_stream.cu_stream(),
9471 );
9472 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9473 return Err(XlogError::Kernel(format!(
9474 "csm left_outer: dtod inner_left col {} failed: {:?}",
9475 col_idx, res
9476 )));
9477 }
9478 }
9479 }
9480 if unmatched_bytes > 0 {
9481 let src_col = unmatched_left.column(col_idx).ok_or_else(|| {
9482 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
9483 })?;
9484 unsafe {
9486 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9487 dst_ptr + inner_bytes as u64,
9488 *src_col.device_ptr(),
9489 unmatched_bytes,
9490 cu_stream.cu_stream(),
9491 );
9492 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9493 return Err(XlogError::Kernel(format!(
9494 "csm left_outer: dtod unmatched col {} failed: {:?}",
9495 col_idx, res
9496 )));
9497 }
9498 }
9499 }
9500 if let Some(b) = out_col.runtime_block() {
9501 runtime
9502 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
9503 .map_err(|e| {
9504 XlogError::Kernel(format!(
9505 "csm left_outer: finish_block_use (left col {}) failed: {}",
9506 col_idx, e
9507 ))
9508 })?;
9509 }
9510 result_columns.push(out_col.into());
9511 }
9512
9513 for col_idx in 0..right.arity() {
9515 let elem_size = right
9516 .schema()
9517 .column_type(col_idx)
9518 .map(|t| t.size_bytes())
9519 .unwrap_or(4);
9520 let inner_bytes = (inner_rows as usize)
9521 .checked_mul(elem_size)
9522 .ok_or_else(|| {
9523 XlogError::Kernel("csm left_outer: right inner_bytes overflow".into())
9524 })?;
9525 let unmatched_bytes = (unmatched_rows as usize)
9526 .checked_mul(elem_size)
9527 .ok_or_else(|| {
9528 XlogError::Kernel("csm left_outer: right unmatched_bytes overflow".into())
9529 })?;
9530 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
9531 XlogError::Kernel("csm left_outer: right total_bytes overflow".into())
9532 })?;
9533 let out_col = self.memory.alloc::<u8>(total_bytes)?;
9534 let dst_ptr = *out_col.device_ptr();
9535 runtime
9537 .prepare_first_use(&out_col, launch_stream, Access::Write)
9538 .map_err(|e| {
9539 XlogError::Kernel(format!(
9540 "csm left_outer: prepare right out_col {} failed: {}",
9541 col_idx, e
9542 ))
9543 })?;
9544 if total_bytes > 0 {
9545 unsafe {
9547 let res = cudarc::driver::sys::cuMemsetD8Async(
9548 dst_ptr,
9549 0,
9550 total_bytes,
9551 cu_stream.cu_stream(),
9552 );
9553 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9554 return Err(XlogError::Kernel(format!(
9555 "csm left_outer: zero-fill right col {} failed: {:?}",
9556 col_idx, res
9557 )));
9558 }
9559 }
9560 }
9561 if inner_bytes > 0 {
9562 let src_col = inner_right_buf
9563 .as_ref()
9564 .expect("inner_count > 0")
9565 .column(col_idx)
9566 .ok_or_else(|| XlogError::Kernel("inner_right col missing".into()))?;
9567 unsafe {
9569 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9570 dst_ptr,
9571 *src_col.device_ptr(),
9572 inner_bytes,
9573 cu_stream.cu_stream(),
9574 );
9575 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9576 return Err(XlogError::Kernel(format!(
9577 "csm left_outer: dtod inner_right col {} failed: {:?}",
9578 col_idx, res
9579 )));
9580 }
9581 }
9582 }
9583 if let Some(b) = out_col.runtime_block() {
9584 runtime
9585 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
9586 .map_err(|e| {
9587 XlogError::Kernel(format!(
9588 "csm left_outer: finish_block_use (right col {}) failed: {}",
9589 col_idx, e
9590 ))
9591 })?;
9592 }
9593 result_columns.push(out_col.into());
9594 }
9595
9596 rec_d.commit(runtime).map_err(|e| {
9597 XlogError::Kernel(format!("csm left_outer: phase-E commit failed: {}", e))
9598 })?;
9599
9600 if total_rows > u32::MAX as u64 {
9609 return Err(XlogError::Kernel(format!(
9610 "csm left_outer: output row count {} exceeds u32::MAX",
9611 total_rows
9612 )));
9613 }
9614 let total_rows_u32 = total_rows as u32;
9615 let d_num_rows = self.upload_device_row_count(total_rows_u32)?;
9616 Ok(CudaBuffer::from_columns_with_host_count(
9617 result_columns,
9618 total_rows,
9619 d_num_rows,
9620 combined_schema,
9621 total_rows_u32,
9622 ))
9623 }
9624
9625 #[allow(clippy::too_many_arguments)]
9650 pub fn hash_join_inner_v2_with_index_count_scan_materialize_recorded(
9651 &self,
9652 left: &CudaBuffer,
9653 right: &CudaBuffer,
9654 left_keys: &[usize],
9655 right_keys: &[usize],
9656 index: &crate::provider::JoinIndexV2,
9657 max_output: Option<usize>,
9658 launch_stream: StreamId,
9659 ) -> Result<CudaBuffer> {
9660 use crate::launch::LaunchRecorder;
9661
9662 let runtime = self.memory.runtime().ok_or_else(|| {
9663 XlogError::Kernel(
9664 "hash_join_inner_v2_with_index_count_scan_materialize_recorded requires \
9665 a runtime-backed GpuMemoryManager"
9666 .to_string(),
9667 )
9668 })?;
9669 let cu_stream = runtime
9670 .stream_pool()
9671 .resolve(launch_stream)
9672 .ok_or_else(|| {
9673 XlogError::Kernel(format!(
9674 "indexed CSM inner: launch_stream StreamId({}) does not resolve",
9675 launch_stream.0
9676 ))
9677 })?;
9678
9679 let left_rows = self.device_row_count(left)?;
9682 let right_rows = self.device_row_count(right)?;
9683 if left_rows > u32::MAX as usize || right_rows > u32::MAX as usize {
9684 return Err(XlogError::Kernel(format!(
9685 "Join supports at most {} rows per side (left={}, right={})",
9686 u32::MAX,
9687 left_rows,
9688 right_rows
9689 )));
9690 }
9691 if left_rows == 0 || right_rows == 0 {
9692 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9693 return self.create_empty_buffer(combined_schema);
9694 }
9695 if left_keys.is_empty() || right_keys.is_empty() {
9696 return Err(XlogError::Kernel(
9697 "Join requires at least one key column".to_string(),
9698 ));
9699 }
9700 if left_keys.len() != right_keys.len() {
9701 return Err(XlogError::Kernel(
9702 "Left and right key columns must have same length".to_string(),
9703 ));
9704 }
9705 if left_keys.len() > 4 {
9706 return Err(XlogError::Kernel(
9707 "indexed CSM inner: max 4 key columns supported (pack_keys constraint)".to_string(),
9708 ));
9709 }
9710 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
9711 if l >= left.arity() {
9712 return Err(XlogError::Kernel(format!(
9713 "Left key column index {} out of bounds (arity {})",
9714 l,
9715 left.arity()
9716 )));
9717 }
9718 if r >= right.arity() {
9719 return Err(XlogError::Kernel(format!(
9720 "Right key column index {} out of bounds (arity {})",
9721 r,
9722 right.arity()
9723 )));
9724 }
9725 let lt = left.schema().column_type(l);
9726 let rt = right.schema().column_type(r);
9727 if lt != rt {
9728 return Err(XlogError::Kernel(format!(
9729 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
9730 l, lt, r, rt
9731 )));
9732 }
9733 }
9734 if index.right_num_rows() != right_rows as u32 {
9735 return Err(XlogError::Kernel(
9736 "Join index row count does not match right relation".to_string(),
9737 ));
9738 }
9739 if index.right_keys() != right_keys {
9740 return Err(XlogError::Kernel(
9741 "Join index key columns do not match requested right_keys".to_string(),
9742 ));
9743 }
9744
9745 let probe_cap = left.num_rows() as u32;
9746 let table = &index.table;
9747
9748 let left_packed =
9751 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
9752 if left_packed.key_bytes != index.key_bytes {
9753 return Err(XlogError::Kernel(
9754 "Join key byte width mismatch between probe and cached index".to_string(),
9755 ));
9756 }
9757
9758 let device = self.device.inner();
9759 let block_size = 256u32;
9760 let probe_grid = probe_cap.div_ceil(block_size);
9761 let probe_config = LaunchConfig {
9762 grid_dim: (probe_grid, 1, 1),
9763 block_dim: (block_size, 1, 1),
9764 shared_mem_bytes: 0,
9765 };
9766
9767 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
9769 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
9770 let d_logical_count = self.memory.alloc::<u32>(1)?;
9771 let d_overflow = self.memory.alloc::<u8>(1)?;
9772 runtime
9774 .prepare_first_use(&d_overflow, launch_stream, Access::Write)
9775 .map_err(|e| {
9776 XlogError::Kernel(format!(
9777 "indexed CSM inner: prepare d_overflow failed: {}",
9778 e
9779 ))
9780 })?;
9781 runtime
9782 .prepare_first_use(&d_logical_count, launch_stream, Access::Write)
9783 .map_err(|e| {
9784 XlogError::Kernel(format!(
9785 "indexed CSM inner: prepare d_logical_count failed: {}",
9786 e
9787 ))
9788 })?;
9789 unsafe {
9792 let res = cudarc::driver::sys::cuMemsetD8Async(
9793 *d_overflow.device_ptr(),
9794 0,
9795 1,
9796 cu_stream.cu_stream(),
9797 );
9798 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9799 return Err(XlogError::Kernel(format!(
9800 "indexed CSM inner: cuMemsetD8Async (d_overflow) failed: {:?}",
9801 res
9802 )));
9803 }
9804 let res = cudarc::driver::sys::cuMemsetD8Async(
9805 *d_logical_count.device_ptr(),
9806 0,
9807 std::mem::size_of::<u32>(),
9808 cu_stream.cu_stream(),
9809 );
9810 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9811 return Err(XlogError::Kernel(format!(
9812 "indexed CSM inner: cuMemsetD8Async (d_logical_count) failed: {:?}",
9813 res
9814 )));
9815 }
9816 }
9817
9818 let count_func = device
9823 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
9824 .ok_or_else(|| {
9825 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
9826 })?;
9827 let total_func = device
9828 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
9829 .ok_or_else(|| {
9830 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
9831 })?;
9832
9833 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
9834 rec_count.read(&left_packed.hashes);
9835 rec_count.read(&left_packed.packed_keys);
9836 rec_count.read(&index.packed_keys);
9837 rec_count.read(&table.bucket_offsets);
9838 rec_count.read(&table.bucket_counts);
9839 rec_count.read(&table.bucket_entries);
9840 rec_count.read(&table.bucket_entry_hashes);
9841 rec_count.read(left.num_rows_device());
9842 rec_count.write(&per_probe_count);
9843 rec_count.write(&per_probe_offsets);
9844 rec_count.write(&d_logical_count);
9845 rec_count.write(&d_overflow);
9846 rec_count.preflight(runtime).map_err(|e| {
9847 XlogError::Kernel(format!(
9848 "indexed CSM inner: count/scan preflight failed: {}",
9849 e
9850 ))
9851 })?;
9852
9853 unsafe {
9856 count_func.clone().launch_on_stream(
9857 &cu_stream,
9858 probe_config,
9859 (
9860 &left_packed.hashes,
9861 left.num_rows_device(),
9862 probe_cap,
9863 &table.bucket_offsets,
9864 &table.bucket_counts,
9865 &table.bucket_entries,
9866 &table.bucket_entry_hashes,
9867 table.bucket_mask,
9868 &left_packed.packed_keys,
9869 &index.packed_keys,
9870 index.key_bytes,
9871 &per_probe_count,
9872 ),
9873 )
9874 }
9875 .map_err(|e| {
9876 XlogError::Kernel(format!(
9877 "hash_join_probe_v2_count_per_row (on_stream, indexed) failed: {}",
9878 e
9879 ))
9880 })?;
9881
9882 unsafe {
9886 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
9887 *per_probe_offsets.device_ptr(),
9888 *per_probe_count.device_ptr(),
9889 (probe_cap as usize) * std::mem::size_of::<u32>(),
9890 cu_stream.cu_stream(),
9891 );
9892 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
9893 return Err(XlogError::Kernel(format!(
9894 "indexed CSM inner: cuMemcpyDtoDAsync (count → offsets) failed: {:?}",
9895 res
9896 )));
9897 }
9898 }
9899 self.multiblock_scan_u32_inplace_on_stream(
9900 &mut per_probe_offsets,
9901 probe_cap,
9902 &cu_stream,
9903 launch_stream,
9904 runtime,
9905 )?;
9906
9907 let materialize_capacity_bound: u64 = (probe_cap as u64).saturating_mul(right_rows as u64);
9909 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
9910 unsafe {
9912 total_func.clone().launch_on_stream(
9913 &cu_stream,
9914 LaunchConfig {
9915 grid_dim: (1, 1, 1),
9916 block_dim: (1, 1, 1),
9917 shared_mem_bytes: 0,
9918 },
9919 (
9920 &per_probe_offsets,
9921 &per_probe_count,
9922 left.num_rows_device(),
9923 probe_cap,
9924 materialize_capacity_u32,
9925 &d_logical_count,
9926 &d_overflow,
9927 ),
9928 )
9929 }
9930 .map_err(|e| {
9931 XlogError::Kernel(format!(
9932 "hash_join_total_from_scan (on_stream, indexed) failed: {}",
9933 e
9934 ))
9935 })?;
9936
9937 rec_count.commit(runtime).map_err(|e| {
9938 XlogError::Kernel(format!(
9939 "indexed CSM inner: count/scan commit failed: {}",
9940 e
9941 ))
9942 })?;
9943
9944 cu_stream.synchronize().map_err(|e| {
9945 XlogError::Kernel(format!(
9946 "indexed CSM inner: sync (total read) failed: {}",
9947 e
9948 ))
9949 })?;
9950 let total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
9951 let requested = max_output
9952 .map(|limit| (limit as u64).min(total))
9953 .unwrap_or(total);
9954 if requested == 0 {
9955 let combined_schema = self.combine_schemas(left.schema(), right.schema());
9956 return self.create_empty_buffer(combined_schema);
9957 }
9958 if requested > u32::MAX as u64 {
9959 return Err(XlogError::Kernel(format!(
9960 "Join produced {} rows which exceeds the u32 index limit",
9961 requested
9962 )));
9963 }
9964 let output_capacity = requested as u32;
9965
9966 let d_output_left = self.memory.alloc::<u32>(output_capacity as usize)?;
9968 let d_output_right = self.memory.alloc::<u32>(output_capacity as usize)?;
9969
9970 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
9971 rec_mat.read(&left_packed.hashes);
9972 rec_mat.read(&left_packed.packed_keys);
9973 rec_mat.read(&index.packed_keys);
9974 rec_mat.read(&table.bucket_offsets);
9975 rec_mat.read(&table.bucket_counts);
9976 rec_mat.read(&table.bucket_entries);
9977 rec_mat.read(&table.bucket_entry_hashes);
9978 rec_mat.read(&per_probe_offsets);
9979 rec_mat.read(left.num_rows_device());
9980 rec_mat.write(&d_output_left);
9981 rec_mat.write(&d_output_right);
9982 rec_mat.write(&d_overflow);
9984 rec_mat.preflight(runtime).map_err(|e| {
9985 XlogError::Kernel(format!(
9986 "indexed CSM inner: materialize preflight failed: {}",
9987 e
9988 ))
9989 })?;
9990
9991 let materialize_func = device
9992 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
9993 .ok_or_else(|| {
9994 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
9995 })?;
9996 unsafe {
9998 let mut params: Vec<*mut c_void> = vec![
9999 (&left_packed.hashes).as_kernel_param(),
10000 left.num_rows_device().as_kernel_param(),
10001 probe_cap.as_kernel_param(),
10002 (&table.bucket_offsets).as_kernel_param(),
10003 (&table.bucket_counts).as_kernel_param(),
10004 (&table.bucket_entries).as_kernel_param(),
10005 (&table.bucket_entry_hashes).as_kernel_param(),
10006 table.bucket_mask.as_kernel_param(),
10007 (&left_packed.packed_keys).as_kernel_param(),
10008 (&index.packed_keys).as_kernel_param(),
10009 index.key_bytes.as_kernel_param(),
10010 (&per_probe_offsets).as_kernel_param(),
10011 output_capacity.as_kernel_param(),
10012 (&d_output_left).as_kernel_param(),
10013 (&d_output_right).as_kernel_param(),
10014 (&d_overflow).as_kernel_param(),
10015 ];
10016 materialize_func
10017 .clone()
10018 .launch_on_stream(&cu_stream, probe_config, &mut params)
10019 .map_err(|e| {
10020 XlogError::Kernel(format!(
10021 "hash_join_probe_v2_materialize (on_stream, indexed) failed: {}",
10022 e
10023 ))
10024 })?;
10025 }
10026
10027 rec_mat.commit(runtime).map_err(|e| {
10028 XlogError::Kernel(format!(
10029 "indexed CSM inner: materialize commit failed: {}",
10030 e
10031 ))
10032 })?;
10033
10034 cu_stream.synchronize().map_err(|e| {
10035 XlogError::Kernel(format!(
10036 "indexed CSM inner: sync (post-materialize) failed: {}",
10037 e
10038 ))
10039 })?;
10040
10041 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
10043 for col_idx in 0..left.columns.len() {
10044 let c = left
10045 .column(col_idx)
10046 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
10047 rec_gather.read_column(c);
10048 }
10049 for col_idx in 0..right.columns.len() {
10050 let c = right
10051 .column(col_idx)
10052 .ok_or_else(|| XlogError::Kernel(format!("Right column {} not found", col_idx)))?;
10053 rec_gather.read_column(c);
10054 }
10055 rec_gather.read(&d_output_left);
10056 rec_gather.read(&d_output_right);
10057 rec_gather.preflight(runtime).map_err(|e| {
10058 XlogError::Kernel(format!("indexed CSM inner: gather preflight failed: {}", e))
10059 })?;
10060 let gathered_left = self.gather_buffer_by_indices_on_stream(
10061 left,
10062 &d_output_left,
10063 output_capacity,
10064 &cu_stream,
10065 launch_stream,
10066 runtime,
10067 )?;
10068 let gathered_right = self.gather_buffer_by_indices_on_stream(
10069 right,
10070 &d_output_right,
10071 output_capacity,
10072 &cu_stream,
10073 launch_stream,
10074 runtime,
10075 )?;
10076 rec_gather.commit(runtime).map_err(|e| {
10077 XlogError::Kernel(format!("indexed CSM inner: gather commit failed: {}", e))
10078 })?;
10079
10080 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10081 let mut result_columns = Vec::with_capacity(combined_schema.arity());
10082 result_columns.extend(gathered_left.columns);
10083 result_columns.extend(gathered_right.columns);
10084 self.buffer_from_columns(result_columns, output_capacity as u64, combined_schema)
10085 }
10086
10087 #[allow(clippy::too_many_arguments)]
10118 pub fn hash_join_left_outer_v2_with_index_count_scan_materialize_recorded(
10119 &self,
10120 left: &CudaBuffer,
10121 right: &CudaBuffer,
10122 left_keys: &[usize],
10123 right_keys: &[usize],
10124 index: &crate::provider::JoinIndexV2,
10125 max_output: Option<usize>,
10126 launch_stream: StreamId,
10127 ) -> Result<CudaBuffer> {
10128 use crate::launch::LaunchRecorder;
10129
10130 let runtime = self.memory.runtime().ok_or_else(|| {
10131 XlogError::Kernel(
10132 "hash_join_left_outer_v2_with_index_count_scan_materialize_recorded requires \
10133 a runtime-backed GpuMemoryManager"
10134 .to_string(),
10135 )
10136 })?;
10137 let cu_stream = runtime
10138 .stream_pool()
10139 .resolve(launch_stream)
10140 .ok_or_else(|| {
10141 XlogError::Kernel(format!(
10142 "indexed csm left_outer: launch_stream StreamId({}) does not resolve",
10143 launch_stream.0
10144 ))
10145 })?;
10146
10147 let num_left = self.device_row_count(left)?;
10150 let num_right = self.device_row_count(right)?;
10151 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
10152 return Err(XlogError::Kernel(format!(
10153 "Join supports at most {} rows per side (left={}, right={})",
10154 u32::MAX,
10155 num_left,
10156 num_right
10157 )));
10158 }
10159 if num_left == 0 {
10160 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10161 return self.create_empty_buffer(combined_schema);
10162 }
10163 if num_right == 0 {
10164 return self.left_outer_with_nulls(left, right);
10169 }
10170 if left_keys.is_empty() || right_keys.is_empty() {
10171 return Err(XlogError::Kernel(
10172 "Join requires at least one key column".to_string(),
10173 ));
10174 }
10175 if left_keys.len() != right_keys.len() {
10176 return Err(XlogError::Kernel(
10177 "Left and right key columns must have same length".to_string(),
10178 ));
10179 }
10180 if left_keys.len() > 4 {
10181 return Err(XlogError::Kernel(
10182 "indexed csm left_outer: max 4 key columns supported (pack_keys constraint)"
10183 .to_string(),
10184 ));
10185 }
10186 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
10187 if l >= left.arity() {
10188 return Err(XlogError::Kernel(format!(
10189 "Left key column index {} out of bounds (arity {})",
10190 l,
10191 left.arity()
10192 )));
10193 }
10194 if r >= right.arity() {
10195 return Err(XlogError::Kernel(format!(
10196 "Right key column index {} out of bounds (arity {})",
10197 r,
10198 right.arity()
10199 )));
10200 }
10201 let lt = left.schema().column_type(l);
10202 let rt = right.schema().column_type(r);
10203 if lt != rt {
10204 return Err(XlogError::Kernel(format!(
10205 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
10206 l, lt, r, rt
10207 )));
10208 }
10209 }
10210 if index.right_num_rows() != num_right as u32 {
10211 return Err(XlogError::Kernel(
10212 "Join index row count does not match right relation".to_string(),
10213 ));
10214 }
10215 if index.right_keys() != right_keys {
10216 return Err(XlogError::Kernel(
10217 "Join index key columns do not match requested right_keys".to_string(),
10218 ));
10219 }
10220
10221 let probe_cap = u32::try_from(num_left).map_err(|_| {
10223 XlogError::Kernel("indexed csm left_outer: left row count exceeds u32::MAX".to_string())
10224 })?;
10225
10226 let table = &index.table;
10227
10228 let left_packed =
10231 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
10232 if left_packed.key_bytes != index.key_bytes {
10233 return Err(XlogError::Kernel(
10234 "Join key byte width mismatch between probe and cached index".to_string(),
10235 ));
10236 }
10237
10238 let device = self.device.inner();
10239 let block_size = 256u32;
10240 let probe_grid = probe_cap.div_ceil(block_size);
10241 let probe_config = LaunchConfig {
10242 grid_dim: (probe_grid, 1, 1),
10243 block_dim: (block_size, 1, 1),
10244 shared_mem_bytes: 0,
10245 };
10246
10247 let per_probe_count = self.memory.alloc::<u32>(probe_cap as usize)?;
10249 let mut per_probe_offsets = self.memory.alloc::<u32>(probe_cap as usize)?;
10250 let d_logical_count = self.memory.alloc::<u32>(1)?;
10251 let d_overflow = self.memory.alloc::<u8>(1)?;
10252 runtime
10253 .prepare_first_use(&d_overflow, launch_stream, Access::Write)
10254 .map_err(|e| {
10255 XlogError::Kernel(format!(
10256 "indexed csm left_outer: prepare d_overflow failed: {}",
10257 e
10258 ))
10259 })?;
10260 runtime
10261 .prepare_first_use(&d_logical_count, launch_stream, Access::Write)
10262 .map_err(|e| {
10263 XlogError::Kernel(format!(
10264 "indexed csm left_outer: prepare d_logical_count failed: {}",
10265 e
10266 ))
10267 })?;
10268 unsafe {
10271 let res = cudarc::driver::sys::cuMemsetD8Async(
10272 *d_overflow.device_ptr(),
10273 0,
10274 1,
10275 cu_stream.cu_stream(),
10276 );
10277 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10278 return Err(XlogError::Kernel(format!(
10279 "indexed csm left_outer: cuMemsetD8Async (d_overflow) failed: {:?}",
10280 res
10281 )));
10282 }
10283 let res = cudarc::driver::sys::cuMemsetD8Async(
10284 *d_logical_count.device_ptr(),
10285 0,
10286 std::mem::size_of::<u32>(),
10287 cu_stream.cu_stream(),
10288 );
10289 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10290 return Err(XlogError::Kernel(format!(
10291 "indexed csm left_outer: cuMemsetD8Async (d_logical_count) failed: {:?}",
10292 res
10293 )));
10294 }
10295 }
10296
10297 let count_func = device
10298 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_COUNT_PER_ROW)
10299 .ok_or_else(|| {
10300 XlogError::Kernel("hash_join_probe_v2_count_per_row kernel not found".to_string())
10301 })?;
10302 let total_func = device
10303 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_TOTAL_FROM_SCAN)
10304 .ok_or_else(|| {
10305 XlogError::Kernel("hash_join_total_from_scan kernel not found".to_string())
10306 })?;
10307
10308 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
10309 rec_count.read(&left_packed.hashes);
10310 rec_count.read(&left_packed.packed_keys);
10311 rec_count.read(&index.packed_keys);
10312 rec_count.read(&table.bucket_offsets);
10313 rec_count.read(&table.bucket_counts);
10314 rec_count.read(&table.bucket_entries);
10315 rec_count.read(&table.bucket_entry_hashes);
10316 rec_count.read(left.num_rows_device());
10317 rec_count.write(&per_probe_count);
10318 rec_count.write(&per_probe_offsets);
10319 rec_count.write(&d_logical_count);
10320 rec_count.write(&d_overflow);
10321 rec_count.preflight(runtime).map_err(|e| {
10322 XlogError::Kernel(format!(
10323 "indexed csm left_outer: count/scan preflight failed: {}",
10324 e
10325 ))
10326 })?;
10327
10328 unsafe {
10331 count_func.clone().launch_on_stream(
10332 &cu_stream,
10333 probe_config,
10334 (
10335 &left_packed.hashes,
10336 left.num_rows_device(),
10337 probe_cap,
10338 &table.bucket_offsets,
10339 &table.bucket_counts,
10340 &table.bucket_entries,
10341 &table.bucket_entry_hashes,
10342 table.bucket_mask,
10343 &left_packed.packed_keys,
10344 &index.packed_keys,
10345 index.key_bytes,
10346 &per_probe_count,
10347 ),
10348 )
10349 }
10350 .map_err(|e| {
10351 XlogError::Kernel(format!(
10352 "hash_join_probe_v2_count_per_row (indexed csm left_outer) failed: {}",
10353 e
10354 ))
10355 })?;
10356
10357 unsafe {
10360 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
10361 *per_probe_offsets.device_ptr(),
10362 *per_probe_count.device_ptr(),
10363 (probe_cap as usize) * std::mem::size_of::<u32>(),
10364 cu_stream.cu_stream(),
10365 );
10366 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10367 return Err(XlogError::Kernel(format!(
10368 "indexed csm left_outer: cuMemcpyDtoDAsync (count → offsets) failed: {:?}",
10369 res
10370 )));
10371 }
10372 }
10373 self.multiblock_scan_u32_inplace_on_stream(
10374 &mut per_probe_offsets,
10375 probe_cap,
10376 &cu_stream,
10377 launch_stream,
10378 runtime,
10379 )?;
10380
10381 let materialize_capacity_bound: u64 = (probe_cap as u64).saturating_mul(num_right as u64);
10383 let materialize_capacity_u32 = materialize_capacity_bound.min(u32::MAX as u64) as u32;
10384 unsafe {
10386 total_func.clone().launch_on_stream(
10387 &cu_stream,
10388 LaunchConfig {
10389 grid_dim: (1, 1, 1),
10390 block_dim: (1, 1, 1),
10391 shared_mem_bytes: 0,
10392 },
10393 (
10394 &per_probe_offsets,
10395 &per_probe_count,
10396 left.num_rows_device(),
10397 probe_cap,
10398 materialize_capacity_u32,
10399 &d_logical_count,
10400 &d_overflow,
10401 ),
10402 )
10403 }
10404 .map_err(|e| {
10405 XlogError::Kernel(format!(
10406 "hash_join_total_from_scan (indexed csm left_outer) failed: {}",
10407 e
10408 ))
10409 })?;
10410
10411 rec_count.commit(runtime).map_err(|e| {
10412 XlogError::Kernel(format!(
10413 "indexed csm left_outer: count/scan commit failed: {}",
10414 e
10415 ))
10416 })?;
10417
10418 cu_stream.synchronize().map_err(|e| {
10419 XlogError::Kernel(format!(
10420 "indexed csm left_outer: sync (count read) failed: {}",
10421 e
10422 ))
10423 })?;
10424 let inner_total = self.read_join_output_count_metadata(&d_logical_count)? as u64;
10425 let inner_clamped = max_output
10426 .map(|limit| (limit as u64).min(inner_total))
10427 .unwrap_or(inner_total);
10428 if inner_clamped > u32::MAX as u64 {
10429 return Err(XlogError::Kernel(format!(
10430 "Join produced {} matched rows which exceeds the u32 index limit",
10431 inner_clamped
10432 )));
10433 }
10434 let inner_count_u32 = inner_clamped as u32;
10435
10436 let materialize_func = device
10438 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2_MATERIALIZE)
10439 .ok_or_else(|| {
10440 XlogError::Kernel("hash_join_probe_v2_materialize kernel not found".to_string())
10441 })?;
10442 let d_output_left = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
10443 let d_output_right = self.memory.alloc::<u32>(inner_count_u32.max(1) as usize)?;
10444
10445 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
10446 rec_mat.read(&left_packed.hashes);
10447 rec_mat.read(&left_packed.packed_keys);
10448 rec_mat.read(&index.packed_keys);
10449 rec_mat.read(&table.bucket_offsets);
10450 rec_mat.read(&table.bucket_counts);
10451 rec_mat.read(&table.bucket_entries);
10452 rec_mat.read(&table.bucket_entry_hashes);
10453 rec_mat.read(&per_probe_offsets);
10454 rec_mat.read(left.num_rows_device());
10455 rec_mat.write(&d_output_left);
10456 rec_mat.write(&d_output_right);
10457 rec_mat.write(&d_overflow);
10459 rec_mat.preflight(runtime).map_err(|e| {
10460 XlogError::Kernel(format!(
10461 "indexed csm left_outer: materialize preflight failed: {}",
10462 e
10463 ))
10464 })?;
10465 if inner_count_u32 > 0 {
10466 unsafe {
10468 let mut params: Vec<*mut c_void> = vec![
10469 (&left_packed.hashes).as_kernel_param(),
10470 left.num_rows_device().as_kernel_param(),
10471 probe_cap.as_kernel_param(),
10472 (&table.bucket_offsets).as_kernel_param(),
10473 (&table.bucket_counts).as_kernel_param(),
10474 (&table.bucket_entries).as_kernel_param(),
10475 (&table.bucket_entry_hashes).as_kernel_param(),
10476 table.bucket_mask.as_kernel_param(),
10477 (&left_packed.packed_keys).as_kernel_param(),
10478 (&index.packed_keys).as_kernel_param(),
10479 index.key_bytes.as_kernel_param(),
10480 (&per_probe_offsets).as_kernel_param(),
10481 inner_count_u32.as_kernel_param(),
10482 (&d_output_left).as_kernel_param(),
10483 (&d_output_right).as_kernel_param(),
10484 (&d_overflow).as_kernel_param(),
10485 ];
10486 materialize_func
10487 .clone()
10488 .launch_on_stream(&cu_stream, probe_config, &mut params)
10489 .map_err(|e| {
10490 XlogError::Kernel(format!(
10491 "hash_join_probe_v2_materialize (indexed csm left_outer) failed: {}",
10492 e
10493 ))
10494 })?;
10495 }
10496 }
10497 rec_mat.commit(runtime).map_err(|e| {
10498 XlogError::Kernel(format!(
10499 "indexed csm left_outer: materialize commit failed: {}",
10500 e
10501 ))
10502 })?;
10503
10504 let d_unmatched_mask = self.memory.alloc::<u8>(probe_cap as usize)?;
10506 let unmatched_mask_func = device
10507 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_CSM_UNMATCHED_MASK)
10508 .ok_or_else(|| {
10509 XlogError::Kernel("hash_join_csm_unmatched_mask kernel not found".to_string())
10510 })?;
10511 let mut rec_um = LaunchRecorder::new_strict(launch_stream);
10512 rec_um.read(&per_probe_count);
10513 rec_um.read(left.num_rows_device());
10514 rec_um.write(&d_unmatched_mask);
10515 rec_um.preflight(runtime).map_err(|e| {
10516 XlogError::Kernel(format!(
10517 "indexed csm left_outer: unmatched mask preflight failed: {}",
10518 e
10519 ))
10520 })?;
10521 unsafe {
10523 unmatched_mask_func.clone().launch_on_stream(
10524 &cu_stream,
10525 probe_config,
10526 (
10527 &per_probe_count,
10528 left.num_rows_device(),
10529 probe_cap,
10530 &d_unmatched_mask,
10531 ),
10532 )
10533 }
10534 .map_err(|e| {
10535 XlogError::Kernel(format!(
10536 "hash_join_csm_unmatched_mask (indexed csm left_outer) failed: {}",
10537 e
10538 ))
10539 })?;
10540 rec_um.commit(runtime).map_err(|e| {
10541 XlogError::Kernel(format!(
10542 "indexed csm left_outer: unmatched mask commit failed: {}",
10543 e
10544 ))
10545 })?;
10546
10547 let unmatched_left = self.compact_buffer_by_device_mask_counted_recorded(
10548 left,
10549 &d_unmatched_mask,
10550 launch_stream,
10551 )?;
10552 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
10553 let total_rows = (inner_count_u32 as u64) + unmatched_rows;
10554
10555 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10556 if total_rows == 0 {
10557 return self.create_empty_buffer(combined_schema);
10558 }
10559
10560 let inner_left_buf;
10562 let inner_right_buf;
10563 if inner_count_u32 > 0 {
10564 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
10565 for col_idx in 0..left.columns.len() {
10566 let c = left.column(col_idx).ok_or_else(|| {
10567 XlogError::Kernel(format!("Left column {} not found", col_idx))
10568 })?;
10569 rec_gather.read_column(c);
10570 }
10571 for col_idx in 0..right.columns.len() {
10572 let c = right.column(col_idx).ok_or_else(|| {
10573 XlogError::Kernel(format!("Right column {} not found", col_idx))
10574 })?;
10575 rec_gather.read_column(c);
10576 }
10577 rec_gather.read(&d_output_left);
10578 rec_gather.read(&d_output_right);
10579 rec_gather.preflight(runtime).map_err(|e| {
10580 XlogError::Kernel(format!(
10581 "indexed csm left_outer: gather preflight failed: {}",
10582 e
10583 ))
10584 })?;
10585 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
10586 left,
10587 &d_output_left,
10588 inner_count_u32,
10589 &cu_stream,
10590 launch_stream,
10591 runtime,
10592 )?);
10593 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
10594 right,
10595 &d_output_right,
10596 inner_count_u32,
10597 &cu_stream,
10598 launch_stream,
10599 runtime,
10600 )?);
10601 rec_gather.commit(runtime).map_err(|e| {
10602 XlogError::Kernel(format!(
10603 "indexed csm left_outer: gather commit failed: {}",
10604 e
10605 ))
10606 })?;
10607 } else {
10608 inner_left_buf = None;
10609 inner_right_buf = None;
10610 }
10611
10612 let mut rec_d = LaunchRecorder::new_strict(launch_stream);
10614 for col_idx in 0..unmatched_left.columns.len() {
10615 let c = unmatched_left.column(col_idx).ok_or_else(|| {
10616 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
10617 })?;
10618 rec_d.read_column(c);
10619 }
10620 if let Some(b) = inner_left_buf.as_ref() {
10621 for col_idx in 0..b.columns.len() {
10622 let c = b.column(col_idx).ok_or_else(|| {
10623 XlogError::Kernel(format!("inner_left col {} not found", col_idx))
10624 })?;
10625 rec_d.read_column(c);
10626 }
10627 }
10628 if let Some(b) = inner_right_buf.as_ref() {
10629 for col_idx in 0..b.columns.len() {
10630 let c = b.column(col_idx).ok_or_else(|| {
10631 XlogError::Kernel(format!("inner_right col {} not found", col_idx))
10632 })?;
10633 rec_d.read_column(c);
10634 }
10635 }
10636 rec_d.preflight(runtime).map_err(|e| {
10637 XlogError::Kernel(format!(
10638 "indexed csm left_outer: phase-E preflight failed: {}",
10639 e
10640 ))
10641 })?;
10642
10643 let inner_rows = inner_count_u32 as u64;
10644 let mut result_columns: Vec<CudaColumn> = Vec::with_capacity(combined_schema.arity());
10645
10646 for col_idx in 0..left.arity() {
10648 let elem_size = left
10649 .schema()
10650 .column_type(col_idx)
10651 .map(|t| t.size_bytes())
10652 .unwrap_or(4);
10653 let inner_bytes = (inner_rows as usize)
10654 .checked_mul(elem_size)
10655 .ok_or_else(|| {
10656 XlogError::Kernel("indexed csm left_outer: inner_bytes overflow".into())
10657 })?;
10658 let unmatched_bytes = (unmatched_rows as usize)
10659 .checked_mul(elem_size)
10660 .ok_or_else(|| {
10661 XlogError::Kernel("indexed csm left_outer: unmatched_bytes overflow".into())
10662 })?;
10663 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
10664 XlogError::Kernel("indexed csm left_outer: total_bytes overflow".into())
10665 })?;
10666 let out_col = self.memory.alloc::<u8>(total_bytes)?;
10667 let dst_ptr = *out_col.device_ptr();
10668 runtime
10669 .prepare_first_use(&out_col, launch_stream, Access::Write)
10670 .map_err(|e| {
10671 XlogError::Kernel(format!(
10672 "indexed csm left_outer: prepare left out_col {} failed: {}",
10673 col_idx, e
10674 ))
10675 })?;
10676 if inner_bytes > 0 {
10677 let src_col = inner_left_buf
10678 .as_ref()
10679 .expect("inner_count > 0")
10680 .column(col_idx)
10681 .ok_or_else(|| XlogError::Kernel("inner_left col missing".into()))?;
10682 unsafe {
10684 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
10685 dst_ptr,
10686 *src_col.device_ptr(),
10687 inner_bytes,
10688 cu_stream.cu_stream(),
10689 );
10690 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10691 return Err(XlogError::Kernel(format!(
10692 "indexed csm left_outer: dtod inner_left col {} failed: {:?}",
10693 col_idx, res
10694 )));
10695 }
10696 }
10697 }
10698 if unmatched_bytes > 0 {
10699 let src_col = unmatched_left.column(col_idx).ok_or_else(|| {
10700 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
10701 })?;
10702 unsafe {
10704 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
10705 dst_ptr + inner_bytes as u64,
10706 *src_col.device_ptr(),
10707 unmatched_bytes,
10708 cu_stream.cu_stream(),
10709 );
10710 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10711 return Err(XlogError::Kernel(format!(
10712 "indexed csm left_outer: dtod unmatched col {} failed: {:?}",
10713 col_idx, res
10714 )));
10715 }
10716 }
10717 }
10718 if let Some(b) = out_col.runtime_block() {
10719 runtime
10720 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
10721 .map_err(|e| {
10722 XlogError::Kernel(format!(
10723 "indexed csm left_outer: finish_block_use (left col {}) failed: {}",
10724 col_idx, e
10725 ))
10726 })?;
10727 }
10728 result_columns.push(out_col.into());
10729 }
10730
10731 for col_idx in 0..right.arity() {
10733 let elem_size = right
10734 .schema()
10735 .column_type(col_idx)
10736 .map(|t| t.size_bytes())
10737 .unwrap_or(4);
10738 let inner_bytes = (inner_rows as usize)
10739 .checked_mul(elem_size)
10740 .ok_or_else(|| {
10741 XlogError::Kernel("indexed csm left_outer: right inner_bytes overflow".into())
10742 })?;
10743 let unmatched_bytes = (unmatched_rows as usize)
10744 .checked_mul(elem_size)
10745 .ok_or_else(|| {
10746 XlogError::Kernel(
10747 "indexed csm left_outer: right unmatched_bytes overflow".into(),
10748 )
10749 })?;
10750 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
10751 XlogError::Kernel("indexed csm left_outer: right total_bytes overflow".into())
10752 })?;
10753 let out_col = self.memory.alloc::<u8>(total_bytes)?;
10754 let dst_ptr = *out_col.device_ptr();
10755 runtime
10756 .prepare_first_use(&out_col, launch_stream, Access::Write)
10757 .map_err(|e| {
10758 XlogError::Kernel(format!(
10759 "indexed csm left_outer: prepare right out_col {} failed: {}",
10760 col_idx, e
10761 ))
10762 })?;
10763 if total_bytes > 0 {
10764 unsafe {
10766 let res = cudarc::driver::sys::cuMemsetD8Async(
10767 dst_ptr,
10768 0,
10769 total_bytes,
10770 cu_stream.cu_stream(),
10771 );
10772 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10773 return Err(XlogError::Kernel(format!(
10774 "indexed csm left_outer: zero-fill right col {} failed: {:?}",
10775 col_idx, res
10776 )));
10777 }
10778 }
10779 }
10780 if inner_bytes > 0 {
10781 let src_col = inner_right_buf
10782 .as_ref()
10783 .expect("inner_count > 0")
10784 .column(col_idx)
10785 .ok_or_else(|| XlogError::Kernel("inner_right col missing".into()))?;
10786 unsafe {
10788 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
10789 dst_ptr,
10790 *src_col.device_ptr(),
10791 inner_bytes,
10792 cu_stream.cu_stream(),
10793 );
10794 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
10795 return Err(XlogError::Kernel(format!(
10796 "indexed csm left_outer: dtod inner_right col {} failed: {:?}",
10797 col_idx, res
10798 )));
10799 }
10800 }
10801 }
10802 if let Some(b) = out_col.runtime_block() {
10803 runtime
10804 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
10805 .map_err(|e| {
10806 XlogError::Kernel(format!(
10807 "indexed csm left_outer: finish_block_use (right col {}) failed: {}",
10808 col_idx, e
10809 ))
10810 })?;
10811 }
10812 result_columns.push(out_col.into());
10813 }
10814
10815 rec_d.commit(runtime).map_err(|e| {
10816 XlogError::Kernel(format!(
10817 "indexed csm left_outer: phase-E commit failed: {}",
10818 e
10819 ))
10820 })?;
10821
10822 if total_rows > u32::MAX as u64 {
10824 return Err(XlogError::Kernel(format!(
10825 "indexed csm left_outer: output row count {} exceeds u32::MAX",
10826 total_rows
10827 )));
10828 }
10829 let total_rows_u32 = total_rows as u32;
10830 let d_num_rows = self.upload_device_row_count(total_rows_u32)?;
10831 Ok(CudaBuffer::from_columns_with_host_count(
10832 result_columns,
10833 total_rows,
10834 d_num_rows,
10835 combined_schema,
10836 total_rows_u32,
10837 ))
10838 }
10839
10840 #[allow(clippy::too_many_arguments)]
10855 pub fn hash_join_v2_recorded(
10856 &self,
10857 left: &CudaBuffer,
10858 right: &CudaBuffer,
10859 left_keys: &[usize],
10860 right_keys: &[usize],
10861 join_type: JoinType,
10862 max_output: Option<usize>,
10863 launch_stream: StreamId,
10864 ) -> Result<CudaBuffer> {
10865 let csm_on = Self::use_recorded_csm_env();
10866 match join_type {
10867 JoinType::Inner => {
10868 if csm_on {
10869 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
10870 self.hash_join_inner_v2_count_scan_materialize_recorded(
10871 left,
10872 right,
10873 left_keys,
10874 right_keys,
10875 max_output,
10876 launch_stream,
10877 )
10878 } else {
10879 self.hash_join_inner_v2_recorded(
10880 left,
10881 right,
10882 left_keys,
10883 right_keys,
10884 max_output,
10885 launch_stream,
10886 )
10887 }
10888 }
10889 JoinType::Semi => self.hash_join_semi_or_anti_v2_recorded(
10890 left,
10891 right,
10892 left_keys,
10893 right_keys,
10894 false,
10895 launch_stream,
10896 ),
10897 JoinType::Anti => self.hash_join_semi_or_anti_v2_recorded(
10898 left,
10899 right,
10900 left_keys,
10901 right_keys,
10902 true,
10903 launch_stream,
10904 ),
10905 JoinType::LeftOuter => {
10906 if csm_on {
10907 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
10908 self.hash_join_left_outer_v2_count_scan_materialize_recorded(
10909 left,
10910 right,
10911 left_keys,
10912 right_keys,
10913 max_output,
10914 launch_stream,
10915 )
10916 } else {
10917 self.hash_join_left_outer_v2_recorded(
10918 left,
10919 right,
10920 left_keys,
10921 right_keys,
10922 max_output,
10923 launch_stream,
10924 )
10925 }
10926 }
10927 }
10928 }
10929
10930 fn hash_join_left_outer_v2_recorded(
10951 &self,
10952 left: &CudaBuffer,
10953 right: &CudaBuffer,
10954 left_keys: &[usize],
10955 right_keys: &[usize],
10956 max_output: Option<usize>,
10957 launch_stream: StreamId,
10958 ) -> Result<CudaBuffer> {
10959 use crate::launch::LaunchRecorder;
10960
10961 let runtime = self.memory.runtime().ok_or_else(|| {
10962 XlogError::Kernel(
10963 "hash_join_v2_recorded (left_outer) requires a runtime-backed GpuMemoryManager"
10964 .to_string(),
10965 )
10966 })?;
10967 let cu_stream = runtime
10968 .stream_pool()
10969 .resolve(launch_stream)
10970 .ok_or_else(|| {
10971 XlogError::Kernel(format!(
10972 "hash_join_v2_recorded (left_outer): launch_stream StreamId({}) does not resolve",
10973 launch_stream.0
10974 ))
10975 })?;
10976
10977 let num_left = self.device_row_count(left)?;
10978 let num_right = self.device_row_count(right)?;
10979 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
10980 return Err(XlogError::Kernel(format!(
10981 "Join supports at most {} rows per side (left={}, right={})",
10982 u32::MAX,
10983 num_left,
10984 num_right
10985 )));
10986 }
10987 if num_left == 0 {
10988 let combined_schema = self.combine_schemas(left.schema(), right.schema());
10989 return self.create_empty_buffer(combined_schema);
10990 }
10991 if num_right == 0 {
10992 return self.left_outer_with_nulls(left, right);
10997 }
10998 if left_keys.is_empty() || right_keys.is_empty() {
10999 return Err(XlogError::Kernel(
11000 "Join requires at least one key column".to_string(),
11001 ));
11002 }
11003 if left_keys.len() != right_keys.len() {
11004 return Err(XlogError::Kernel(
11005 "Left and right key columns must have same length".to_string(),
11006 ));
11007 }
11008 if left_keys.len() > 4 {
11009 return Err(XlogError::Kernel(
11010 "hash_join_v2_recorded (left_outer): max 4 key columns supported \
11011 (pack_keys constraint)"
11012 .to_string(),
11013 ));
11014 }
11015 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
11016 let lt = left.schema().column_type(l);
11017 let rt = right.schema().column_type(r);
11018 if lt != rt {
11019 return Err(XlogError::Kernel(format!(
11020 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
11021 l, lt, r, rt
11022 )));
11023 }
11024 }
11025
11026 let num_left = num_left as u32;
11027 let num_right = num_right as u32;
11028
11029 let left_packed =
11030 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
11031 let right_packed =
11032 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
11033 let table = self.build_hash_table_v2_on_stream(
11034 &right_packed.hashes,
11035 num_right,
11036 &cu_stream,
11037 launch_stream,
11038 runtime,
11039 )?;
11040
11041 let device = self.device.inner();
11042 let block_size = 256u32;
11043 let grid_size = num_left.div_ceil(block_size);
11044 let cfg = LaunchConfig {
11045 grid_dim: (grid_size, 1, 1),
11046 block_dim: (block_size, 1, 1),
11047 shared_mem_bytes: 0,
11048 };
11049
11050 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
11054 let d_count_only = self.memory.alloc::<u32>(1)?;
11055 let d_dummy_left = self.memory.alloc::<u32>(1)?;
11056 let d_dummy_right = self.memory.alloc::<u32>(1)?;
11057 runtime
11061 .prepare_first_use(&d_count_only, launch_stream, Access::Write)
11062 .map_err(|e| {
11063 XlogError::Kernel(format!(
11064 "left_outer recorded: prepare d_count_only failed: {}",
11065 e
11066 ))
11067 })?;
11068 unsafe {
11070 let res = cudarc::driver::sys::cuMemsetD8Async(
11071 *d_count_only.device_ptr(),
11072 0,
11073 std::mem::size_of::<u32>(),
11074 cu_stream.cu_stream(),
11075 );
11076 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11077 return Err(XlogError::Kernel(format!(
11078 "cuMemsetD8Async (left_outer d_count_only) failed: {:?}",
11079 res
11080 )));
11081 }
11082 }
11083
11084 let semi_func = device
11085 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
11086 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
11087 let probe_func = device
11088 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
11089 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
11090
11091 let mut rec_a = LaunchRecorder::new_strict(launch_stream);
11092 rec_a.read(&left_packed.hashes);
11093 rec_a.read(&left_packed.packed_keys);
11094 rec_a.read(&right_packed.packed_keys);
11095 rec_a.read(&table.bucket_offsets);
11096 rec_a.read(&table.bucket_counts);
11097 rec_a.read(&table.bucket_entries);
11098 rec_a.read(&table.bucket_entry_hashes);
11099 rec_a.write(&d_has_match);
11100 rec_a.write(&d_count_only);
11101 rec_a.write(&d_dummy_left);
11102 rec_a.write(&d_dummy_right);
11103 rec_a.preflight(runtime).map_err(|e| {
11104 XlogError::Kernel(format!(
11105 "hash_join_v2_recorded (left_outer): semi/count preflight failed: {}",
11106 e
11107 ))
11108 })?;
11109
11110 unsafe {
11112 semi_func.clone().launch_on_stream(
11113 &cu_stream,
11114 cfg,
11115 (
11116 &left_packed.hashes,
11117 num_left,
11118 &table.bucket_offsets,
11119 &table.bucket_counts,
11120 &table.bucket_entries,
11121 &table.bucket_entry_hashes,
11122 table.bucket_mask,
11123 &left_packed.packed_keys,
11124 &right_packed.packed_keys,
11125 left_packed.key_bytes,
11126 &d_has_match,
11127 ),
11128 )
11129 }
11130 .map_err(|e| XlogError::Kernel(format!("hash_join_semi (on_stream) failed: {}", e)))?;
11131
11132 let max_output_count_only = 0u32;
11133 unsafe {
11135 let mut params: Vec<*mut c_void> = vec![
11136 (&left_packed.hashes).as_kernel_param(),
11137 num_left.as_kernel_param(),
11138 (&table.bucket_offsets).as_kernel_param(),
11139 (&table.bucket_counts).as_kernel_param(),
11140 (&table.bucket_entries).as_kernel_param(),
11141 (&table.bucket_entry_hashes).as_kernel_param(),
11142 table.bucket_mask.as_kernel_param(),
11143 (&left_packed.packed_keys).as_kernel_param(),
11144 (&right_packed.packed_keys).as_kernel_param(),
11145 left_packed.key_bytes.as_kernel_param(),
11146 (&d_dummy_left).as_kernel_param(),
11147 (&d_dummy_right).as_kernel_param(),
11148 (&d_count_only).as_kernel_param(),
11149 max_output_count_only.as_kernel_param(),
11150 ];
11151 probe_func
11152 .clone()
11153 .launch_on_stream(&cu_stream, cfg, &mut params)
11154 .map_err(|e| {
11155 XlogError::Kernel(format!(
11156 "hash_join_probe_v2 (count, on_stream, left_outer) failed: {}",
11157 e
11158 ))
11159 })?;
11160 }
11161
11162 rec_a.commit(runtime).map_err(|e| {
11163 XlogError::Kernel(format!(
11164 "hash_join_v2_recorded (left_outer): semi/count commit failed: {}",
11165 e
11166 ))
11167 })?;
11168
11169 cu_stream.synchronize().map_err(|e| {
11171 XlogError::Kernel(format!(
11172 "hash_join_v2_recorded (left_outer): sync (count read) failed: {}",
11173 e
11174 ))
11175 })?;
11176 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
11177 let requested_inner = max_output
11178 .map(|limit| (limit as u64).min(full_inner))
11179 .unwrap_or(full_inner);
11180 if requested_inner > u32::MAX as u64 {
11181 return Err(XlogError::Kernel(format!(
11182 "Join produced {} rows which exceeds the u32 index limit",
11183 requested_inner
11184 )));
11185 }
11186 let max_output_u32 = requested_inner as u32;
11187 let alloc_len = (requested_inner.max(1)) as usize;
11188
11189 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
11190 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
11191 let d_output_count = self.memory.alloc::<u32>(1)?;
11192 runtime
11195 .prepare_first_use(&d_output_count, launch_stream, Access::Write)
11196 .map_err(|e| {
11197 XlogError::Kernel(format!(
11198 "left_outer recorded: prepare d_output_count failed: {}",
11199 e
11200 ))
11201 })?;
11202 unsafe {
11204 let res = cudarc::driver::sys::cuMemsetD8Async(
11205 *d_output_count.device_ptr(),
11206 0,
11207 std::mem::size_of::<u32>(),
11208 cu_stream.cu_stream(),
11209 );
11210 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11211 return Err(XlogError::Kernel(format!(
11212 "cuMemsetD8Async (left_outer d_output_count) failed: {:?}",
11213 res
11214 )));
11215 }
11216 }
11217
11218 let mut rec_b = LaunchRecorder::new_strict(launch_stream);
11219 rec_b.read(&left_packed.hashes);
11220 rec_b.read(&left_packed.packed_keys);
11221 rec_b.read(&right_packed.packed_keys);
11222 rec_b.read(&table.bucket_offsets);
11223 rec_b.read(&table.bucket_counts);
11224 rec_b.read(&table.bucket_entries);
11225 rec_b.read(&table.bucket_entry_hashes);
11226 rec_b.write(&d_output_left);
11227 rec_b.write(&d_output_right);
11228 rec_b.write(&d_output_count);
11229 rec_b.preflight(runtime).map_err(|e| {
11230 XlogError::Kernel(format!(
11231 "hash_join_v2_recorded (left_outer): materialize preflight failed: {}",
11232 e
11233 ))
11234 })?;
11235
11236 unsafe {
11238 let mut params: Vec<*mut c_void> = vec![
11239 (&left_packed.hashes).as_kernel_param(),
11240 num_left.as_kernel_param(),
11241 (&table.bucket_offsets).as_kernel_param(),
11242 (&table.bucket_counts).as_kernel_param(),
11243 (&table.bucket_entries).as_kernel_param(),
11244 (&table.bucket_entry_hashes).as_kernel_param(),
11245 table.bucket_mask.as_kernel_param(),
11246 (&left_packed.packed_keys).as_kernel_param(),
11247 (&right_packed.packed_keys).as_kernel_param(),
11248 left_packed.key_bytes.as_kernel_param(),
11249 (&d_output_left).as_kernel_param(),
11250 (&d_output_right).as_kernel_param(),
11251 (&d_output_count).as_kernel_param(),
11252 max_output_u32.as_kernel_param(),
11253 ];
11254 probe_func
11255 .clone()
11256 .launch_on_stream(&cu_stream, cfg, &mut params)
11257 .map_err(|e| {
11258 XlogError::Kernel(format!(
11259 "hash_join_probe_v2 (materialize, on_stream, left_outer) failed: {}",
11260 e
11261 ))
11262 })?;
11263 }
11264
11265 rec_b.commit(runtime).map_err(|e| {
11266 XlogError::Kernel(format!(
11267 "hash_join_v2_recorded (left_outer): materialize commit failed: {}",
11268 e
11269 ))
11270 })?;
11271
11272 cu_stream.synchronize().map_err(|e| {
11273 XlogError::Kernel(format!(
11274 "hash_join_v2_recorded (left_outer): sync (materialize read) failed: {}",
11275 e
11276 ))
11277 })?;
11278 let inner_count = self
11279 .read_join_output_count_metadata(&d_output_count)?
11280 .min(max_output_u32);
11281
11282 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
11285 let mask_not_fn = device
11286 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
11287 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
11288
11289 let mut rec_c = LaunchRecorder::new_strict(launch_stream);
11290 rec_c.read(&d_has_match);
11291 rec_c.write(&d_no_match);
11292 rec_c.preflight(runtime).map_err(|e| {
11293 XlogError::Kernel(format!(
11294 "hash_join_v2_recorded (left_outer): mask_not preflight failed: {}",
11295 e
11296 ))
11297 })?;
11298 unsafe {
11300 mask_not_fn.clone().launch_on_stream(
11301 &cu_stream,
11302 cfg,
11303 (&d_has_match, &d_no_match, num_left),
11304 )
11305 }
11306 .map_err(|e| XlogError::Kernel(format!("mask_not (on_stream) failed: {}", e)))?;
11307 rec_c.commit(runtime).map_err(|e| {
11308 XlogError::Kernel(format!(
11309 "hash_join_v2_recorded (left_outer): mask_not commit failed: {}",
11310 e
11311 ))
11312 })?;
11313
11314 let unmatched_left =
11315 self.compact_buffer_by_device_mask_counted_recorded(left, &d_no_match, launch_stream)?;
11316 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
11317 let total_rows = (inner_count as u64) + unmatched_rows;
11318
11319 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11320 if total_rows == 0 {
11321 return self.create_empty_buffer(combined_schema);
11322 }
11323
11324 let inner_count_u32 = inner_count;
11333 let inner_left_buf;
11334 let inner_right_buf;
11335 if inner_count > 0 {
11336 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
11337 for col_idx in 0..left.columns.len() {
11338 let c = left.column(col_idx).ok_or_else(|| {
11339 XlogError::Kernel(format!("Left column {} not found", col_idx))
11340 })?;
11341 rec_gather.read_column(c);
11342 }
11343 for col_idx in 0..right.columns.len() {
11344 let c = right.column(col_idx).ok_or_else(|| {
11345 XlogError::Kernel(format!("Right column {} not found", col_idx))
11346 })?;
11347 rec_gather.read_column(c);
11348 }
11349 rec_gather.read(&d_output_left);
11350 rec_gather.read(&d_output_right);
11351 rec_gather.preflight(runtime).map_err(|e| {
11352 XlogError::Kernel(format!(
11353 "hash_join_v2_recorded (left_outer): gather preflight failed: {}",
11354 e
11355 ))
11356 })?;
11357 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
11358 left,
11359 &d_output_left,
11360 inner_count_u32,
11361 &cu_stream,
11362 launch_stream,
11363 runtime,
11364 )?);
11365 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
11366 right,
11367 &d_output_right,
11368 inner_count_u32,
11369 &cu_stream,
11370 launch_stream,
11371 runtime,
11372 )?);
11373 rec_gather.commit(runtime).map_err(|e| {
11374 XlogError::Kernel(format!(
11375 "hash_join_v2_recorded (left_outer): gather commit failed: {}",
11376 e
11377 ))
11378 })?;
11379 } else {
11380 inner_left_buf = None;
11381 inner_right_buf = None;
11382 }
11383
11384 let mut rec_d = LaunchRecorder::new_strict(launch_stream);
11401 for col_idx in 0..unmatched_left.columns.len() {
11402 let c = unmatched_left.column(col_idx).ok_or_else(|| {
11403 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
11404 })?;
11405 rec_d.read_column(c);
11406 }
11407 if let Some(b) = inner_left_buf.as_ref() {
11408 for col_idx in 0..b.columns.len() {
11409 let c = b.column(col_idx).ok_or_else(|| {
11410 XlogError::Kernel(format!("inner_left col {} not found", col_idx))
11411 })?;
11412 rec_d.read_column(c);
11413 }
11414 }
11415 if let Some(b) = inner_right_buf.as_ref() {
11416 for col_idx in 0..b.columns.len() {
11417 let c = b.column(col_idx).ok_or_else(|| {
11418 XlogError::Kernel(format!("inner_right col {} not found", col_idx))
11419 })?;
11420 rec_d.read_column(c);
11421 }
11422 }
11423 rec_d.preflight(runtime).map_err(|e| {
11424 XlogError::Kernel(format!(
11425 "hash_join_v2_recorded (left_outer): step-D preflight failed: {}",
11426 e
11427 ))
11428 })?;
11429
11430 let mut result_columns: Vec<CudaColumn> = Vec::with_capacity(combined_schema.arity());
11431 let inner_rows = inner_count as u64;
11432
11433 for col_idx in 0..left.arity() {
11435 let elem_size = left
11436 .schema()
11437 .column_type(col_idx)
11438 .map(|t| t.size_bytes())
11439 .unwrap_or(4);
11440 let inner_bytes = (inner_rows as usize)
11441 .checked_mul(elem_size)
11442 .ok_or_else(|| {
11443 XlogError::Kernel("Left outer join: inner_bytes overflow".to_string())
11444 })?;
11445 let unmatched_bytes = (unmatched_rows as usize)
11446 .checked_mul(elem_size)
11447 .ok_or_else(|| {
11448 XlogError::Kernel("Left outer join: unmatched_bytes overflow".to_string())
11449 })?;
11450 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
11451 XlogError::Kernel("Left outer join: total_bytes overflow".to_string())
11452 })?;
11453
11454 let out_col = self.memory.alloc::<u8>(total_bytes)?;
11455 let dst_ptr = *out_col.device_ptr();
11456 runtime
11459 .prepare_first_use(&out_col, launch_stream, Access::Write)
11460 .map_err(|e| {
11461 XlogError::Kernel(format!(
11462 "left_outer recorded: prepare left out_col {} failed: {}",
11463 col_idx, e
11464 ))
11465 })?;
11466
11467 if inner_bytes > 0 {
11468 let src_col = inner_left_buf
11469 .as_ref()
11470 .expect("inner_count > 0 but inner_left_buf is None")
11471 .column(col_idx)
11472 .ok_or_else(|| {
11473 XlogError::Kernel(format!("inner_left col {} not found", col_idx))
11474 })?;
11475 unsafe {
11477 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
11478 dst_ptr,
11479 *src_col.device_ptr(),
11480 inner_bytes,
11481 cu_stream.cu_stream(),
11482 );
11483 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11484 return Err(XlogError::Kernel(format!(
11485 "cuMemcpyDtoDAsync (left_outer inner_left col {}) failed: {:?}",
11486 col_idx, res
11487 )));
11488 }
11489 }
11490 }
11491 if unmatched_bytes > 0 {
11492 let src_col = unmatched_left.column(col_idx).ok_or_else(|| {
11493 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
11494 })?;
11495 unsafe {
11498 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
11499 dst_ptr + inner_bytes as u64,
11500 *src_col.device_ptr(),
11501 unmatched_bytes,
11502 cu_stream.cu_stream(),
11503 );
11504 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11505 return Err(XlogError::Kernel(format!(
11506 "cuMemcpyDtoDAsync (left_outer unmatched_left col {}) failed: {:?}",
11507 col_idx, res
11508 )));
11509 }
11510 }
11511 }
11512
11513 if let Some(b) = out_col.runtime_block() {
11517 runtime
11518 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
11519 .map_err(|e| {
11520 XlogError::Kernel(format!(
11521 "hash_join_v2_recorded (left_outer): finish_block_use \
11522 (left col {}) failed: {}",
11523 col_idx, e
11524 ))
11525 })?;
11526 }
11527 result_columns.push(out_col.into());
11528 }
11529
11530 for col_idx in 0..right.arity() {
11532 let elem_size = right
11533 .schema()
11534 .column_type(col_idx)
11535 .map(|t| t.size_bytes())
11536 .unwrap_or(4);
11537 let inner_bytes = (inner_rows as usize)
11538 .checked_mul(elem_size)
11539 .ok_or_else(|| {
11540 XlogError::Kernel("Left outer join: right inner_bytes overflow".to_string())
11541 })?;
11542 let unmatched_bytes = (unmatched_rows as usize)
11543 .checked_mul(elem_size)
11544 .ok_or_else(|| {
11545 XlogError::Kernel("Left outer join: right unmatched_bytes overflow".to_string())
11546 })?;
11547 let total_bytes = inner_bytes.checked_add(unmatched_bytes).ok_or_else(|| {
11548 XlogError::Kernel("Left outer join: right total_bytes overflow".to_string())
11549 })?;
11550
11551 let out_col = self.memory.alloc::<u8>(total_bytes)?;
11552 let dst_ptr = *out_col.device_ptr();
11553 runtime
11556 .prepare_first_use(&out_col, launch_stream, Access::Write)
11557 .map_err(|e| {
11558 XlogError::Kernel(format!(
11559 "left_outer recorded: prepare right out_col {} failed: {}",
11560 col_idx, e
11561 ))
11562 })?;
11563
11564 if total_bytes > 0 {
11568 unsafe {
11570 let res = cudarc::driver::sys::cuMemsetD8Async(
11571 dst_ptr,
11572 0,
11573 total_bytes,
11574 cu_stream.cu_stream(),
11575 );
11576 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11577 return Err(XlogError::Kernel(format!(
11578 "cuMemsetD8Async (left_outer right col {}) failed: {:?}",
11579 col_idx, res
11580 )));
11581 }
11582 }
11583 }
11584 if inner_bytes > 0 {
11585 let src_col = inner_right_buf
11586 .as_ref()
11587 .expect("inner_count > 0 but inner_right_buf is None")
11588 .column(col_idx)
11589 .ok_or_else(|| {
11590 XlogError::Kernel(format!("inner_right col {} not found", col_idx))
11591 })?;
11592 unsafe {
11594 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
11595 dst_ptr,
11596 *src_col.device_ptr(),
11597 inner_bytes,
11598 cu_stream.cu_stream(),
11599 );
11600 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
11601 return Err(XlogError::Kernel(format!(
11602 "cuMemcpyDtoDAsync (left_outer inner_right col {}) failed: {:?}",
11603 col_idx, res
11604 )));
11605 }
11606 }
11607 }
11608
11609 if let Some(b) = out_col.runtime_block() {
11610 runtime
11611 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
11612 .map_err(|e| {
11613 XlogError::Kernel(format!(
11614 "hash_join_v2_recorded (left_outer): finish_block_use \
11615 (right col {}) failed: {}",
11616 col_idx, e
11617 ))
11618 })?;
11619 }
11620 result_columns.push(out_col.into());
11621 }
11622
11623 rec_d.commit(runtime).map_err(|e| {
11629 XlogError::Kernel(format!(
11630 "hash_join_v2_recorded (left_outer): step-D commit failed: {}",
11631 e
11632 ))
11633 })?;
11634
11635 let d_num_rows = self.upload_device_row_count(total_rows as u32)?;
11638 Ok(CudaBuffer::from_columns_with_host_count(
11639 result_columns,
11640 total_rows,
11641 d_num_rows,
11642 combined_schema,
11643 total_rows as u32,
11644 ))
11645 }
11646
11647 fn hash_join_semi_or_anti_v2_recorded(
11658 &self,
11659 left: &CudaBuffer,
11660 right: &CudaBuffer,
11661 left_keys: &[usize],
11662 right_keys: &[usize],
11663 anti: bool,
11664 launch_stream: StreamId,
11665 ) -> Result<CudaBuffer> {
11666 use crate::launch::LaunchRecorder;
11667
11668 let runtime = self.memory.runtime().ok_or_else(|| {
11669 XlogError::Kernel(
11670 "hash_join_v2_recorded (semi/anti) requires a runtime-backed GpuMemoryManager"
11671 .to_string(),
11672 )
11673 })?;
11674 let cu_stream = runtime
11675 .stream_pool()
11676 .resolve(launch_stream)
11677 .ok_or_else(|| {
11678 XlogError::Kernel(format!(
11679 "hash_join_v2_recorded (semi/anti): launch_stream StreamId({}) does not resolve",
11680 launch_stream.0
11681 ))
11682 })?;
11683
11684 let num_left = self.device_row_count(left)?;
11685 let num_right = self.device_row_count(right)?;
11686 if num_left > u32::MAX as usize || num_right > u32::MAX as usize {
11687 return Err(XlogError::Kernel(format!(
11688 "Join supports at most {} rows per side (left={}, right={})",
11689 u32::MAX,
11690 num_left,
11691 num_right
11692 )));
11693 }
11694 if num_left == 0 {
11695 return self.create_empty_buffer(left.schema().clone());
11696 }
11697 if num_right == 0 {
11698 return if anti {
11711 self.clone_buffer(left)
11712 } else {
11713 self.create_empty_buffer(left.schema().clone())
11714 };
11715 }
11716 if left_keys.is_empty() || right_keys.is_empty() {
11717 return Err(XlogError::Kernel(
11718 "Join requires at least one key column".to_string(),
11719 ));
11720 }
11721 if left_keys.len() != right_keys.len() {
11722 return Err(XlogError::Kernel(
11723 "Left and right key columns must have same length".to_string(),
11724 ));
11725 }
11726 if left_keys.len() > 4 {
11727 return Err(XlogError::Kernel(
11728 "hash_join_v2_recorded (semi/anti): max 4 key columns supported (pack_keys constraint)"
11729 .to_string(),
11730 ));
11731 }
11732 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
11733 let lt = left.schema().column_type(l);
11734 let rt = right.schema().column_type(r);
11735 if lt != rt {
11736 return Err(XlogError::Kernel(format!(
11737 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
11738 l, lt, r, rt
11739 )));
11740 }
11741 }
11742
11743 let num_left = num_left as u32;
11744 let num_right = num_right as u32;
11745
11746 let left_packed =
11747 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
11748 let right_packed =
11749 self.pack_keys_gpu_on_stream(right, right_keys, &cu_stream, launch_stream, runtime)?;
11750 let table = self.build_hash_table_v2_on_stream(
11751 &right_packed.hashes,
11752 num_right,
11753 &cu_stream,
11754 launch_stream,
11755 runtime,
11756 )?;
11757
11758 let d_mask = self.memory.alloc::<u8>(num_left as usize)?;
11759
11760 let kernel_name = if anti {
11761 join_kernels::HASH_JOIN_ANTI
11762 } else {
11763 join_kernels::HASH_JOIN_SEMI
11764 };
11765 let func = self
11766 .device
11767 .inner()
11768 .get_func(JOIN_MODULE, kernel_name)
11769 .ok_or_else(|| XlogError::Kernel(format!("{} kernel not found", kernel_name)))?;
11770
11771 let block_size = 256u32;
11772 let grid_size = num_left.div_ceil(block_size);
11773 let cfg = LaunchConfig {
11774 grid_dim: (grid_size, 1, 1),
11775 block_dim: (block_size, 1, 1),
11776 shared_mem_bytes: 0,
11777 };
11778
11779 let mut rec = LaunchRecorder::new_strict(launch_stream);
11780 rec.read(&left_packed.hashes);
11781 rec.read(&left_packed.packed_keys);
11782 rec.read(&right_packed.packed_keys);
11783 rec.read(&table.bucket_offsets);
11784 rec.read(&table.bucket_counts);
11785 rec.read(&table.bucket_entries);
11786 rec.read(&table.bucket_entry_hashes);
11787 rec.write(&d_mask);
11788 rec.preflight(runtime).map_err(|e| {
11789 XlogError::Kernel(format!(
11790 "hash_join_v2_recorded (semi/anti): preflight failed: {}",
11791 e
11792 ))
11793 })?;
11794
11795 unsafe {
11800 func.clone().launch_on_stream(
11801 &cu_stream,
11802 cfg,
11803 (
11804 &left_packed.hashes,
11805 num_left,
11806 &table.bucket_offsets,
11807 &table.bucket_counts,
11808 &table.bucket_entries,
11809 &table.bucket_entry_hashes,
11810 table.bucket_mask,
11811 &left_packed.packed_keys,
11812 &right_packed.packed_keys,
11813 left_packed.key_bytes,
11814 &d_mask,
11815 ),
11816 )
11817 }
11818 .map_err(|e| XlogError::Kernel(format!("{} (on_stream) failed: {}", kernel_name, e)))?;
11819
11820 rec.commit(runtime).map_err(|e| {
11821 XlogError::Kernel(format!(
11822 "hash_join_v2_recorded (semi/anti): commit failed: {}",
11823 e
11824 ))
11825 })?;
11826
11827 self.compact_buffer_by_device_mask_counted_recorded(left, &d_mask, launch_stream)
11834 }
11835
11836 #[allow(clippy::too_many_arguments)]
11863 pub fn hash_join_v2_with_index_recorded(
11864 &self,
11865 left: &CudaBuffer,
11866 right: &CudaBuffer,
11867 left_keys: &[usize],
11868 right_keys: &[usize],
11869 join_type: JoinType,
11870 index: &crate::provider::JoinIndexV2,
11871 max_output: Option<usize>,
11872 launch_stream: StreamId,
11873 ) -> Result<CudaBuffer> {
11874 let runtime = self.memory.runtime().ok_or_else(|| {
11875 XlogError::Kernel(
11876 "hash_join_v2_with_index_recorded requires a runtime-backed GpuMemoryManager"
11877 .to_string(),
11878 )
11879 })?;
11880 runtime
11882 .stream_pool()
11883 .resolve(launch_stream)
11884 .ok_or_else(|| {
11885 XlogError::Kernel(format!(
11886 "hash_join_v2_with_index_recorded: launch_stream StreamId({}) does not resolve",
11887 launch_stream.0
11888 ))
11889 })?;
11890
11891 let left_rows = self.device_row_count(left)?;
11893 let right_rows = self.device_row_count(right)?;
11894 if left_rows > u32::MAX as usize || right_rows > u32::MAX as usize {
11895 return Err(XlogError::Kernel(format!(
11896 "Join supports at most {} rows per side (left={}, right={})",
11897 u32::MAX,
11898 left_rows,
11899 right_rows
11900 )));
11901 }
11902 if left_rows == 0 {
11903 return match join_type {
11904 JoinType::Inner | JoinType::LeftOuter => {
11905 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11906 self.create_empty_buffer(combined_schema)
11907 }
11908 JoinType::Semi | JoinType::Anti => self.create_empty_buffer(left.schema().clone()),
11909 };
11910 }
11911 if right_rows == 0 {
11912 return match join_type {
11913 JoinType::Inner => {
11914 let combined_schema = self.combine_schemas(left.schema(), right.schema());
11915 self.create_empty_buffer(combined_schema)
11916 }
11917 JoinType::Semi => self.create_empty_buffer(left.schema().clone()),
11918 JoinType::Anti => self.clone_buffer(left),
11919 JoinType::LeftOuter => self.left_outer_with_nulls(left, right),
11920 };
11921 }
11922 if left_keys.is_empty() || right_keys.is_empty() {
11923 return Err(XlogError::Kernel(
11924 "Join requires at least one key column".to_string(),
11925 ));
11926 }
11927 if left_keys.len() != right_keys.len() {
11928 return Err(XlogError::Kernel(
11929 "Left and right key columns must have same length".to_string(),
11930 ));
11931 }
11932 if left_keys.len() > 4 {
11933 return Err(XlogError::Kernel(
11934 "hash_join_v2_with_index_recorded: max 4 key columns supported \
11935 (pack_keys constraint)"
11936 .to_string(),
11937 ));
11938 }
11939 for (&l, &r) in left_keys.iter().zip(right_keys.iter()) {
11940 if l >= left.arity() {
11941 return Err(XlogError::Kernel(format!(
11942 "Left key column index {} out of bounds (arity {})",
11943 l,
11944 left.arity()
11945 )));
11946 }
11947 if r >= right.arity() {
11948 return Err(XlogError::Kernel(format!(
11949 "Right key column index {} out of bounds (arity {})",
11950 r,
11951 right.arity()
11952 )));
11953 }
11954 let lt = left.schema().column_type(l);
11955 let rt = right.schema().column_type(r);
11956 if lt != rt {
11957 return Err(XlogError::Kernel(format!(
11958 "Key column type mismatch: left[{}]={:?}, right[{}]={:?}",
11959 l, lt, r, rt
11960 )));
11961 }
11962 }
11963 if index.right_num_rows() != right_rows as u32 {
11964 return Err(XlogError::Kernel(
11965 "Join index row count does not match right relation".to_string(),
11966 ));
11967 }
11968 if index.right_keys() != right_keys {
11969 return Err(XlogError::Kernel(
11970 "Join index key columns do not match requested right_keys".to_string(),
11971 ));
11972 }
11973
11974 let csm_on = Self::use_recorded_csm_env();
11975 match join_type {
11976 JoinType::Inner => {
11977 if csm_on {
11978 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
11979 self.hash_join_inner_v2_with_index_count_scan_materialize_recorded(
11980 left,
11981 right,
11982 left_keys,
11983 right_keys,
11984 index,
11985 max_output,
11986 launch_stream,
11987 )
11988 } else {
11989 self.hash_join_inner_v2_with_index_recorded(
11990 left,
11991 right,
11992 left_keys,
11993 index,
11994 max_output,
11995 launch_stream,
11996 )
11997 }
11998 }
11999 JoinType::Semi => self.hash_join_semi_or_anti_v2_with_index_recorded(
12000 left,
12001 left_keys,
12002 index,
12003 false,
12004 launch_stream,
12005 ),
12006 JoinType::Anti => self.hash_join_semi_or_anti_v2_with_index_recorded(
12007 left,
12008 left_keys,
12009 index,
12010 true,
12011 launch_stream,
12012 ),
12013 JoinType::LeftOuter => {
12014 if csm_on {
12015 self.csm_invocations.fetch_add(1, Ordering::Relaxed);
12016 self.hash_join_left_outer_v2_with_index_count_scan_materialize_recorded(
12017 left,
12018 right,
12019 left_keys,
12020 right_keys,
12021 index,
12022 max_output,
12023 launch_stream,
12024 )
12025 } else {
12026 self.hash_join_left_outer_v2_with_index_recorded(
12027 left,
12028 right,
12029 left_keys,
12030 index,
12031 max_output,
12032 launch_stream,
12033 )
12034 }
12035 }
12036 }
12037 }
12038
12039 fn hash_join_inner_v2_with_index_recorded(
12045 &self,
12046 left: &CudaBuffer,
12047 right: &CudaBuffer,
12048 left_keys: &[usize],
12049 index: &crate::provider::JoinIndexV2,
12050 max_output: Option<usize>,
12051 launch_stream: StreamId,
12052 ) -> Result<CudaBuffer> {
12053 use crate::launch::LaunchRecorder;
12054
12055 let runtime = self.memory.runtime().ok_or_else(|| {
12056 XlogError::Kernel(
12057 "hash_join_v2_with_index_recorded (inner) requires runtime-backed manager"
12058 .to_string(),
12059 )
12060 })?;
12061 let cu_stream = runtime
12062 .stream_pool()
12063 .resolve(launch_stream)
12064 .ok_or_else(|| {
12065 XlogError::Kernel("indexed inner: launch_stream does not resolve".to_string())
12066 })?;
12067
12068 let num_left = left.num_rows() as u32;
12069 let table = &index.table;
12070
12071 let left_packed =
12073 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
12074 if left_packed.key_bytes != index.key_bytes {
12075 return Err(XlogError::Kernel(
12076 "Join key byte width mismatch between probe and cached index".to_string(),
12077 ));
12078 }
12079
12080 let probe_func = self
12081 .device
12082 .inner()
12083 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
12084 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
12085 let block_size = 256u32;
12086 let probe_grid = num_left.div_ceil(block_size);
12087 let probe_config = LaunchConfig {
12088 grid_dim: (probe_grid, 1, 1),
12089 block_dim: (block_size, 1, 1),
12090 shared_mem_bytes: 0,
12091 };
12092
12093 let d_count_only = self.memory.alloc::<u32>(1)?;
12095 let d_dummy_left = self.memory.alloc::<u32>(1)?;
12096 let d_dummy_right = self.memory.alloc::<u32>(1)?;
12097 runtime
12100 .prepare_first_use(&d_count_only, launch_stream, Access::Write)
12101 .map_err(|e| {
12102 XlogError::Kernel(format!("indexed inner: prepare d_count_only failed: {}", e))
12103 })?;
12104 unsafe {
12106 let res = cudarc::driver::sys::cuMemsetD8Async(
12107 *d_count_only.device_ptr(),
12108 0,
12109 std::mem::size_of::<u32>(),
12110 cu_stream.cu_stream(),
12111 );
12112 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12113 return Err(XlogError::Kernel(format!(
12114 "cuMemsetD8Async (indexed inner d_count_only) failed: {:?}",
12115 res
12116 )));
12117 }
12118 }
12119
12120 let max_output_count_only = 0u32;
12121 let mut rec_count = LaunchRecorder::new_strict(launch_stream);
12122 rec_count.read(&left_packed.hashes);
12123 rec_count.read(&left_packed.packed_keys);
12124 rec_count.read(&index.packed_keys);
12125 rec_count.read(&table.bucket_offsets);
12126 rec_count.read(&table.bucket_counts);
12127 rec_count.read(&table.bucket_entries);
12128 rec_count.read(&table.bucket_entry_hashes);
12129 rec_count.write(&d_count_only);
12130 rec_count.write(&d_dummy_left);
12131 rec_count.write(&d_dummy_right);
12132 rec_count.preflight(runtime).map_err(|e| {
12133 XlogError::Kernel(format!("indexed inner: count-pass preflight failed: {}", e))
12134 })?;
12135 unsafe {
12137 let mut params: Vec<*mut c_void> = vec![
12138 (&left_packed.hashes).as_kernel_param(),
12139 num_left.as_kernel_param(),
12140 (&table.bucket_offsets).as_kernel_param(),
12141 (&table.bucket_counts).as_kernel_param(),
12142 (&table.bucket_entries).as_kernel_param(),
12143 (&table.bucket_entry_hashes).as_kernel_param(),
12144 table.bucket_mask.as_kernel_param(),
12145 (&left_packed.packed_keys).as_kernel_param(),
12146 (&index.packed_keys).as_kernel_param(),
12147 index.key_bytes.as_kernel_param(),
12148 (&d_dummy_left).as_kernel_param(),
12149 (&d_dummy_right).as_kernel_param(),
12150 (&d_count_only).as_kernel_param(),
12151 max_output_count_only.as_kernel_param(),
12152 ];
12153 probe_func
12154 .clone()
12155 .launch_on_stream(&cu_stream, probe_config, &mut params)
12156 .map_err(|e| {
12157 XlogError::Kernel(format!(
12158 "hash_join_probe_v2 (indexed count, on_stream) failed: {}",
12159 e
12160 ))
12161 })?;
12162 }
12163 rec_count.commit(runtime).map_err(|e| {
12164 XlogError::Kernel(format!("indexed inner: count-pass commit failed: {}", e))
12165 })?;
12166
12167 cu_stream.synchronize().map_err(|e| {
12168 XlogError::Kernel(format!("indexed inner: sync (count read) failed: {}", e))
12169 })?;
12170 let full_count = self.read_join_output_count_metadata(&d_count_only)? as u64;
12171 let requested = max_output
12172 .map(|limit| (limit as u64).min(full_count))
12173 .unwrap_or(full_count);
12174 if requested == 0 {
12175 let combined_schema = self.combine_schemas(left.schema(), right.schema());
12176 return self.create_empty_buffer(combined_schema);
12177 }
12178 if requested > u32::MAX as u64 {
12179 return Err(XlogError::Kernel(format!(
12180 "Join produced {} rows which exceeds the u32 index limit",
12181 requested
12182 )));
12183 }
12184 let max_output_u32 = requested as u32;
12185
12186 let d_output_left = self.memory.alloc::<u32>(max_output_u32 as usize)?;
12188 let d_output_right = self.memory.alloc::<u32>(max_output_u32 as usize)?;
12189 let d_output_count = self.memory.alloc::<u32>(1)?;
12190 runtime
12193 .prepare_first_use(&d_output_count, launch_stream, Access::Write)
12194 .map_err(|e| {
12195 XlogError::Kernel(format!(
12196 "indexed inner: prepare d_output_count failed: {}",
12197 e
12198 ))
12199 })?;
12200 unsafe {
12202 let res = cudarc::driver::sys::cuMemsetD8Async(
12203 *d_output_count.device_ptr(),
12204 0,
12205 std::mem::size_of::<u32>(),
12206 cu_stream.cu_stream(),
12207 );
12208 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12209 return Err(XlogError::Kernel(format!(
12210 "cuMemsetD8Async (indexed inner d_output_count) failed: {:?}",
12211 res
12212 )));
12213 }
12214 }
12215
12216 let mut rec_mat = LaunchRecorder::new_strict(launch_stream);
12217 rec_mat.read(&left_packed.hashes);
12218 rec_mat.read(&left_packed.packed_keys);
12219 rec_mat.read(&index.packed_keys);
12220 rec_mat.read(&table.bucket_offsets);
12221 rec_mat.read(&table.bucket_counts);
12222 rec_mat.read(&table.bucket_entries);
12223 rec_mat.read(&table.bucket_entry_hashes);
12224 rec_mat.write(&d_output_left);
12225 rec_mat.write(&d_output_right);
12226 rec_mat.write(&d_output_count);
12227 rec_mat.preflight(runtime).map_err(|e| {
12228 XlogError::Kernel(format!(
12229 "indexed inner: materialize preflight failed: {}",
12230 e
12231 ))
12232 })?;
12233 unsafe {
12235 let mut params: Vec<*mut c_void> = vec![
12236 (&left_packed.hashes).as_kernel_param(),
12237 num_left.as_kernel_param(),
12238 (&table.bucket_offsets).as_kernel_param(),
12239 (&table.bucket_counts).as_kernel_param(),
12240 (&table.bucket_entries).as_kernel_param(),
12241 (&table.bucket_entry_hashes).as_kernel_param(),
12242 table.bucket_mask.as_kernel_param(),
12243 (&left_packed.packed_keys).as_kernel_param(),
12244 (&index.packed_keys).as_kernel_param(),
12245 index.key_bytes.as_kernel_param(),
12246 (&d_output_left).as_kernel_param(),
12247 (&d_output_right).as_kernel_param(),
12248 (&d_output_count).as_kernel_param(),
12249 max_output_u32.as_kernel_param(),
12250 ];
12251 probe_func
12252 .clone()
12253 .launch_on_stream(&cu_stream, probe_config, &mut params)
12254 .map_err(|e| {
12255 XlogError::Kernel(format!(
12256 "hash_join_probe_v2 (indexed mat, on_stream) failed: {}",
12257 e
12258 ))
12259 })?;
12260 }
12261 rec_mat.commit(runtime).map_err(|e| {
12262 XlogError::Kernel(format!("indexed inner: materialize commit failed: {}", e))
12263 })?;
12264
12265 cu_stream.synchronize().map_err(|e| {
12266 XlogError::Kernel(format!("indexed inner: sync (mat read) failed: {}", e))
12267 })?;
12268 let result_count = (self.read_join_output_count_metadata(&d_output_count)? as u64)
12269 .min(max_output_u32 as u64);
12270 if result_count == 0 {
12271 let combined_schema = self.combine_schemas(left.schema(), right.schema());
12272 return self.create_empty_buffer(combined_schema);
12273 }
12274 let output_rows = result_count as u32;
12275
12276 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
12278 for col_idx in 0..left.columns.len() {
12279 let c = left
12280 .column(col_idx)
12281 .ok_or_else(|| XlogError::Kernel(format!("Left column {} not found", col_idx)))?;
12282 rec_gather.read_column(c);
12283 }
12284 for col_idx in 0..right.columns.len() {
12285 let c = right
12286 .column(col_idx)
12287 .ok_or_else(|| XlogError::Kernel(format!("Right column {} not found", col_idx)))?;
12288 rec_gather.read_column(c);
12289 }
12290 rec_gather.read(&d_output_left);
12291 rec_gather.read(&d_output_right);
12292 rec_gather.preflight(runtime).map_err(|e| {
12293 XlogError::Kernel(format!("indexed inner: gather preflight failed: {}", e))
12294 })?;
12295 let gathered_left = self.gather_buffer_by_indices_on_stream(
12296 left,
12297 &d_output_left,
12298 output_rows,
12299 &cu_stream,
12300 launch_stream,
12301 runtime,
12302 )?;
12303 let gathered_right = self.gather_buffer_by_indices_on_stream(
12304 right,
12305 &d_output_right,
12306 output_rows,
12307 &cu_stream,
12308 launch_stream,
12309 runtime,
12310 )?;
12311 rec_gather.commit(runtime).map_err(|e| {
12312 XlogError::Kernel(format!("indexed inner: gather commit failed: {}", e))
12313 })?;
12314
12315 let combined_schema = self.combine_schemas(left.schema(), right.schema());
12316 let mut result_columns = Vec::with_capacity(combined_schema.arity());
12317 result_columns.extend(gathered_left.columns);
12318 result_columns.extend(gathered_right.columns);
12319 self.buffer_from_columns(result_columns, result_count, combined_schema)
12320 }
12321
12322 fn hash_join_semi_or_anti_v2_with_index_recorded(
12328 &self,
12329 left: &CudaBuffer,
12330 left_keys: &[usize],
12331 index: &crate::provider::JoinIndexV2,
12332 anti: bool,
12333 launch_stream: StreamId,
12334 ) -> Result<CudaBuffer> {
12335 use crate::launch::LaunchRecorder;
12336
12337 let runtime = self.memory.runtime().ok_or_else(|| {
12338 XlogError::Kernel(
12339 "hash_join_v2_with_index_recorded (semi/anti) requires runtime-backed manager"
12340 .to_string(),
12341 )
12342 })?;
12343 let cu_stream = runtime
12344 .stream_pool()
12345 .resolve(launch_stream)
12346 .ok_or_else(|| {
12347 XlogError::Kernel("indexed semi/anti: launch_stream does not resolve".to_string())
12348 })?;
12349
12350 let num_left = left.num_rows() as u32;
12351 let table = &index.table;
12352
12353 let left_packed =
12354 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
12355 if left_packed.key_bytes != index.key_bytes {
12356 return Err(XlogError::Kernel(
12357 "Join key byte width mismatch between probe and cached index".to_string(),
12358 ));
12359 }
12360
12361 let d_mask = self.memory.alloc::<u8>(num_left as usize)?;
12362 let kernel_name = if anti {
12363 join_kernels::HASH_JOIN_ANTI
12364 } else {
12365 join_kernels::HASH_JOIN_SEMI
12366 };
12367 let func = self
12368 .device
12369 .inner()
12370 .get_func(JOIN_MODULE, kernel_name)
12371 .ok_or_else(|| XlogError::Kernel(format!("{} kernel not found", kernel_name)))?;
12372 let block_size = 256u32;
12373 let grid_size = num_left.div_ceil(block_size);
12374 let cfg = LaunchConfig {
12375 grid_dim: (grid_size, 1, 1),
12376 block_dim: (block_size, 1, 1),
12377 shared_mem_bytes: 0,
12378 };
12379
12380 let mut rec = LaunchRecorder::new_strict(launch_stream);
12381 rec.read(&left_packed.hashes);
12382 rec.read(&left_packed.packed_keys);
12383 rec.read(&index.packed_keys);
12384 rec.read(&table.bucket_offsets);
12385 rec.read(&table.bucket_counts);
12386 rec.read(&table.bucket_entries);
12387 rec.read(&table.bucket_entry_hashes);
12388 rec.write(&d_mask);
12389 rec.preflight(runtime).map_err(|e| {
12390 XlogError::Kernel(format!("indexed semi/anti: preflight failed: {}", e))
12391 })?;
12392 unsafe {
12394 func.clone().launch_on_stream(
12395 &cu_stream,
12396 cfg,
12397 (
12398 &left_packed.hashes,
12399 num_left,
12400 &table.bucket_offsets,
12401 &table.bucket_counts,
12402 &table.bucket_entries,
12403 &table.bucket_entry_hashes,
12404 table.bucket_mask,
12405 &left_packed.packed_keys,
12406 &index.packed_keys,
12407 index.key_bytes,
12408 &d_mask,
12409 ),
12410 )
12411 }
12412 .map_err(|e| {
12413 XlogError::Kernel(format!(
12414 "{} (on_stream, indexed) failed: {}",
12415 kernel_name, e
12416 ))
12417 })?;
12418 rec.commit(runtime)
12419 .map_err(|e| XlogError::Kernel(format!("indexed semi/anti: commit failed: {}", e)))?;
12420
12421 self.compact_buffer_by_device_mask_counted_recorded(left, &d_mask, launch_stream)
12422 }
12423
12424 fn hash_join_left_outer_v2_with_index_recorded(
12430 &self,
12431 left: &CudaBuffer,
12432 right: &CudaBuffer,
12433 left_keys: &[usize],
12434 index: &crate::provider::JoinIndexV2,
12435 max_output: Option<usize>,
12436 launch_stream: StreamId,
12437 ) -> Result<CudaBuffer> {
12438 use crate::launch::LaunchRecorder;
12439
12440 let runtime = self.memory.runtime().ok_or_else(|| {
12441 XlogError::Kernel(
12442 "hash_join_v2_with_index_recorded (left_outer) requires runtime-backed manager"
12443 .to_string(),
12444 )
12445 })?;
12446 let cu_stream = runtime
12447 .stream_pool()
12448 .resolve(launch_stream)
12449 .ok_or_else(|| {
12450 XlogError::Kernel("indexed left_outer: launch_stream does not resolve".to_string())
12451 })?;
12452
12453 let num_left = left.num_rows() as u32;
12454 let table = &index.table;
12455
12456 let left_packed =
12457 self.pack_keys_gpu_on_stream(left, left_keys, &cu_stream, launch_stream, runtime)?;
12458 if left_packed.key_bytes != index.key_bytes {
12459 return Err(XlogError::Kernel(
12460 "Join key byte width mismatch between probe and cached index".to_string(),
12461 ));
12462 }
12463
12464 let device = self.device.inner();
12465 let block_size = 256u32;
12466 let grid_size = num_left.div_ceil(block_size);
12467 let cfg = LaunchConfig {
12468 grid_dim: (grid_size, 1, 1),
12469 block_dim: (block_size, 1, 1),
12470 shared_mem_bytes: 0,
12471 };
12472
12473 let d_has_match = self.memory.alloc::<u8>(num_left as usize)?;
12475 let d_count_only = self.memory.alloc::<u32>(1)?;
12476 let d_dummy_left = self.memory.alloc::<u32>(1)?;
12477 let d_dummy_right = self.memory.alloc::<u32>(1)?;
12478 runtime
12481 .prepare_first_use(&d_count_only, launch_stream, Access::Write)
12482 .map_err(|e| {
12483 XlogError::Kernel(format!(
12484 "indexed left_outer: prepare d_count_only failed: {}",
12485 e
12486 ))
12487 })?;
12488 unsafe {
12490 let res = cudarc::driver::sys::cuMemsetD8Async(
12491 *d_count_only.device_ptr(),
12492 0,
12493 std::mem::size_of::<u32>(),
12494 cu_stream.cu_stream(),
12495 );
12496 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12497 return Err(XlogError::Kernel(format!(
12498 "cuMemsetD8Async (indexed left_outer d_count_only) failed: {:?}",
12499 res
12500 )));
12501 }
12502 }
12503
12504 let semi_func = device
12505 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_SEMI)
12506 .ok_or_else(|| XlogError::Kernel("hash_join_semi kernel not found".to_string()))?;
12507 let probe_func = device
12508 .get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE_V2)
12509 .ok_or_else(|| XlogError::Kernel("hash_join_probe_v2 kernel not found".to_string()))?;
12510
12511 let mut rec_a = LaunchRecorder::new_strict(launch_stream);
12512 rec_a.read(&left_packed.hashes);
12513 rec_a.read(&left_packed.packed_keys);
12514 rec_a.read(&index.packed_keys);
12515 rec_a.read(&table.bucket_offsets);
12516 rec_a.read(&table.bucket_counts);
12517 rec_a.read(&table.bucket_entries);
12518 rec_a.read(&table.bucket_entry_hashes);
12519 rec_a.write(&d_has_match);
12520 rec_a.write(&d_count_only);
12521 rec_a.write(&d_dummy_left);
12522 rec_a.write(&d_dummy_right);
12523 rec_a.preflight(runtime).map_err(|e| {
12524 XlogError::Kernel(format!(
12525 "indexed left_outer: semi/count preflight failed: {}",
12526 e
12527 ))
12528 })?;
12529 unsafe {
12531 semi_func.clone().launch_on_stream(
12532 &cu_stream,
12533 cfg,
12534 (
12535 &left_packed.hashes,
12536 num_left,
12537 &table.bucket_offsets,
12538 &table.bucket_counts,
12539 &table.bucket_entries,
12540 &table.bucket_entry_hashes,
12541 table.bucket_mask,
12542 &left_packed.packed_keys,
12543 &index.packed_keys,
12544 index.key_bytes,
12545 &d_has_match,
12546 ),
12547 )
12548 }
12549 .map_err(|e| {
12550 XlogError::Kernel(format!(
12551 "hash_join_semi (on_stream, indexed left_outer) failed: {}",
12552 e
12553 ))
12554 })?;
12555
12556 let max_output_count_only = 0u32;
12557 unsafe {
12559 let mut params: Vec<*mut c_void> = vec![
12560 (&left_packed.hashes).as_kernel_param(),
12561 num_left.as_kernel_param(),
12562 (&table.bucket_offsets).as_kernel_param(),
12563 (&table.bucket_counts).as_kernel_param(),
12564 (&table.bucket_entries).as_kernel_param(),
12565 (&table.bucket_entry_hashes).as_kernel_param(),
12566 table.bucket_mask.as_kernel_param(),
12567 (&left_packed.packed_keys).as_kernel_param(),
12568 (&index.packed_keys).as_kernel_param(),
12569 index.key_bytes.as_kernel_param(),
12570 (&d_dummy_left).as_kernel_param(),
12571 (&d_dummy_right).as_kernel_param(),
12572 (&d_count_only).as_kernel_param(),
12573 max_output_count_only.as_kernel_param(),
12574 ];
12575 probe_func
12576 .clone()
12577 .launch_on_stream(&cu_stream, cfg, &mut params)
12578 .map_err(|e| {
12579 XlogError::Kernel(format!(
12580 "hash_join_probe_v2 (count, on_stream, indexed left_outer) failed: {}",
12581 e
12582 ))
12583 })?;
12584 }
12585 rec_a.commit(runtime).map_err(|e| {
12586 XlogError::Kernel(format!(
12587 "indexed left_outer: semi/count commit failed: {}",
12588 e
12589 ))
12590 })?;
12591
12592 cu_stream.synchronize().map_err(|e| {
12593 XlogError::Kernel(format!(
12594 "indexed left_outer: sync (count read) failed: {}",
12595 e
12596 ))
12597 })?;
12598 let full_inner = self.read_join_output_count_metadata(&d_count_only)? as u64;
12599 let requested_inner = max_output
12600 .map(|limit| (limit as u64).min(full_inner))
12601 .unwrap_or(full_inner);
12602 if requested_inner > u32::MAX as u64 {
12603 return Err(XlogError::Kernel(format!(
12604 "Join produced {} rows which exceeds the u32 index limit",
12605 requested_inner
12606 )));
12607 }
12608 let max_output_u32 = requested_inner as u32;
12609 let alloc_len = (requested_inner.max(1)) as usize;
12610
12611 let d_output_left = self.memory.alloc::<u32>(alloc_len)?;
12613 let d_output_right = self.memory.alloc::<u32>(alloc_len)?;
12614 let d_output_count = self.memory.alloc::<u32>(1)?;
12615 runtime
12618 .prepare_first_use(&d_output_count, launch_stream, Access::Write)
12619 .map_err(|e| {
12620 XlogError::Kernel(format!(
12621 "indexed left_outer: prepare d_output_count failed: {}",
12622 e
12623 ))
12624 })?;
12625 unsafe {
12627 let res = cudarc::driver::sys::cuMemsetD8Async(
12628 *d_output_count.device_ptr(),
12629 0,
12630 std::mem::size_of::<u32>(),
12631 cu_stream.cu_stream(),
12632 );
12633 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12634 return Err(XlogError::Kernel(format!(
12635 "cuMemsetD8Async (indexed left_outer d_output_count) failed: {:?}",
12636 res
12637 )));
12638 }
12639 }
12640
12641 let mut rec_b = LaunchRecorder::new_strict(launch_stream);
12642 rec_b.read(&left_packed.hashes);
12643 rec_b.read(&left_packed.packed_keys);
12644 rec_b.read(&index.packed_keys);
12645 rec_b.read(&table.bucket_offsets);
12646 rec_b.read(&table.bucket_counts);
12647 rec_b.read(&table.bucket_entries);
12648 rec_b.read(&table.bucket_entry_hashes);
12649 rec_b.write(&d_output_left);
12650 rec_b.write(&d_output_right);
12651 rec_b.write(&d_output_count);
12652 rec_b.preflight(runtime).map_err(|e| {
12653 XlogError::Kernel(format!(
12654 "indexed left_outer: materialize preflight failed: {}",
12655 e
12656 ))
12657 })?;
12658 unsafe {
12660 let mut params: Vec<*mut c_void> = vec![
12661 (&left_packed.hashes).as_kernel_param(),
12662 num_left.as_kernel_param(),
12663 (&table.bucket_offsets).as_kernel_param(),
12664 (&table.bucket_counts).as_kernel_param(),
12665 (&table.bucket_entries).as_kernel_param(),
12666 (&table.bucket_entry_hashes).as_kernel_param(),
12667 table.bucket_mask.as_kernel_param(),
12668 (&left_packed.packed_keys).as_kernel_param(),
12669 (&index.packed_keys).as_kernel_param(),
12670 index.key_bytes.as_kernel_param(),
12671 (&d_output_left).as_kernel_param(),
12672 (&d_output_right).as_kernel_param(),
12673 (&d_output_count).as_kernel_param(),
12674 max_output_u32.as_kernel_param(),
12675 ];
12676 probe_func
12677 .clone()
12678 .launch_on_stream(&cu_stream, cfg, &mut params)
12679 .map_err(|e| {
12680 XlogError::Kernel(format!(
12681 "hash_join_probe_v2 (mat, on_stream, indexed left_outer) failed: {}",
12682 e
12683 ))
12684 })?;
12685 }
12686 rec_b.commit(runtime).map_err(|e| {
12687 XlogError::Kernel(format!(
12688 "indexed left_outer: materialize commit failed: {}",
12689 e
12690 ))
12691 })?;
12692
12693 cu_stream.synchronize().map_err(|e| {
12694 XlogError::Kernel(format!("indexed left_outer: sync (mat read) failed: {}", e))
12695 })?;
12696 let inner_count = self
12697 .read_join_output_count_metadata(&d_output_count)?
12698 .min(max_output_u32);
12699
12700 let d_no_match = self.memory.alloc::<u8>(num_left as usize)?;
12702 let mask_not_fn = device
12703 .get_func(FILTER_MODULE, filter_kernels::MASK_NOT)
12704 .ok_or_else(|| XlogError::Kernel("mask_not kernel not found".to_string()))?;
12705 let mut rec_c = LaunchRecorder::new_strict(launch_stream);
12706 rec_c.read(&d_has_match);
12707 rec_c.write(&d_no_match);
12708 rec_c.preflight(runtime).map_err(|e| {
12709 XlogError::Kernel(format!(
12710 "indexed left_outer: mask_not preflight failed: {}",
12711 e
12712 ))
12713 })?;
12714 unsafe {
12716 mask_not_fn.clone().launch_on_stream(
12717 &cu_stream,
12718 cfg,
12719 (&d_has_match, &d_no_match, num_left),
12720 )
12721 }
12722 .map_err(|e| {
12723 XlogError::Kernel(format!(
12724 "mask_not (on_stream, indexed left_outer) failed: {}",
12725 e
12726 ))
12727 })?;
12728 rec_c.commit(runtime).map_err(|e| {
12729 XlogError::Kernel(format!("indexed left_outer: mask_not commit failed: {}", e))
12730 })?;
12731
12732 let unmatched_left =
12733 self.compact_buffer_by_device_mask_counted_recorded(left, &d_no_match, launch_stream)?;
12734 let unmatched_rows = self.device_row_count(&unmatched_left)? as u64;
12735 let total_rows = (inner_count as u64) + unmatched_rows;
12736
12737 let combined_schema = self.combine_schemas(left.schema(), right.schema());
12738 if total_rows == 0 {
12739 return self.create_empty_buffer(combined_schema);
12740 }
12741
12742 let inner_count_u32 = inner_count;
12748 let inner_left_buf;
12749 let inner_right_buf;
12750 if inner_count > 0 {
12751 let mut rec_gather = LaunchRecorder::new_strict(launch_stream);
12752 for col_idx in 0..left.columns.len() {
12753 let c = left.column(col_idx).ok_or_else(|| {
12754 XlogError::Kernel(format!("Left column {} not found", col_idx))
12755 })?;
12756 rec_gather.read_column(c);
12757 }
12758 for col_idx in 0..right.columns.len() {
12759 let c = right.column(col_idx).ok_or_else(|| {
12760 XlogError::Kernel(format!("Right column {} not found", col_idx))
12761 })?;
12762 rec_gather.read_column(c);
12763 }
12764 rec_gather.read(&d_output_left);
12765 rec_gather.read(&d_output_right);
12766 rec_gather.preflight(runtime).map_err(|e| {
12767 XlogError::Kernel(format!(
12768 "indexed left_outer: gather preflight failed: {}",
12769 e
12770 ))
12771 })?;
12772 inner_left_buf = Some(self.gather_buffer_by_indices_on_stream(
12773 left,
12774 &d_output_left,
12775 inner_count_u32,
12776 &cu_stream,
12777 launch_stream,
12778 runtime,
12779 )?);
12780 inner_right_buf = Some(self.gather_buffer_by_indices_on_stream(
12781 right,
12782 &d_output_right,
12783 inner_count_u32,
12784 &cu_stream,
12785 launch_stream,
12786 runtime,
12787 )?);
12788 rec_gather.commit(runtime).map_err(|e| {
12789 XlogError::Kernel(format!("indexed left_outer: gather commit failed: {}", e))
12790 })?;
12791 } else {
12792 inner_left_buf = None;
12793 inner_right_buf = None;
12794 }
12795
12796 let mut rec_d = LaunchRecorder::new_strict(launch_stream);
12802 for col_idx in 0..unmatched_left.columns.len() {
12803 let c = unmatched_left.column(col_idx).ok_or_else(|| {
12804 XlogError::Kernel(format!("unmatched_left col {} not found", col_idx))
12805 })?;
12806 rec_d.read_column(c);
12807 }
12808 if let Some(b) = inner_left_buf.as_ref() {
12809 for col_idx in 0..b.columns.len() {
12810 let c = b.column(col_idx).ok_or_else(|| {
12811 XlogError::Kernel(format!("inner_left col {} not found", col_idx))
12812 })?;
12813 rec_d.read_column(c);
12814 }
12815 }
12816 if let Some(b) = inner_right_buf.as_ref() {
12817 for col_idx in 0..b.columns.len() {
12818 let c = b.column(col_idx).ok_or_else(|| {
12819 XlogError::Kernel(format!("inner_right col {} not found", col_idx))
12820 })?;
12821 rec_d.read_column(c);
12822 }
12823 }
12824 rec_d.preflight(runtime).map_err(|e| {
12825 XlogError::Kernel(format!(
12826 "indexed left_outer: step-D preflight failed: {}",
12827 e
12828 ))
12829 })?;
12830
12831 let mut result_columns: Vec<CudaColumn> = Vec::with_capacity(combined_schema.arity());
12832 let inner_rows = inner_count as u64;
12833
12834 for col_idx in 0..left.arity() {
12835 let elem_size = left
12836 .schema()
12837 .column_type(col_idx)
12838 .map(|t| t.size_bytes())
12839 .unwrap_or(4);
12840 let inner_bytes = (inner_rows as usize)
12841 .checked_mul(elem_size)
12842 .ok_or_else(|| XlogError::Kernel("inner_bytes overflow".to_string()))?;
12843 let unmatched_bytes = (unmatched_rows as usize)
12844 .checked_mul(elem_size)
12845 .ok_or_else(|| XlogError::Kernel("unmatched_bytes overflow".to_string()))?;
12846 let total_bytes = inner_bytes
12847 .checked_add(unmatched_bytes)
12848 .ok_or_else(|| XlogError::Kernel("total_bytes overflow".to_string()))?;
12849 let out_col = self.memory.alloc::<u8>(total_bytes)?;
12850 let dst_ptr = *out_col.device_ptr();
12851 runtime
12853 .prepare_first_use(&out_col, launch_stream, Access::Write)
12854 .map_err(|e| {
12855 XlogError::Kernel(format!(
12856 "indexed left_outer: prepare left out_col {} failed: {}",
12857 col_idx, e
12858 ))
12859 })?;
12860 if inner_bytes > 0 {
12861 let src_col = inner_left_buf
12862 .as_ref()
12863 .expect("inner_count > 0")
12864 .column(col_idx)
12865 .ok_or_else(|| XlogError::Kernel("inner_left col missing".to_string()))?;
12866 unsafe {
12868 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
12869 dst_ptr,
12870 *src_col.device_ptr(),
12871 inner_bytes,
12872 cu_stream.cu_stream(),
12873 );
12874 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12875 return Err(XlogError::Kernel(format!(
12876 "indexed left_outer: dtod copy inner_left col {} failed: {:?}",
12877 col_idx, res
12878 )));
12879 }
12880 }
12881 }
12882 if unmatched_bytes > 0 {
12883 let src_col = unmatched_left
12884 .column(col_idx)
12885 .ok_or_else(|| XlogError::Kernel("unmatched col missing".to_string()))?;
12886 unsafe {
12888 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
12889 dst_ptr + inner_bytes as u64,
12890 *src_col.device_ptr(),
12891 unmatched_bytes,
12892 cu_stream.cu_stream(),
12893 );
12894 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12895 return Err(XlogError::Kernel(format!(
12896 "indexed left_outer: dtod copy unmatched col {} failed: {:?}",
12897 col_idx, res
12898 )));
12899 }
12900 }
12901 }
12902 if let Some(b) = out_col.runtime_block() {
12903 runtime
12904 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
12905 .map_err(|e| {
12906 XlogError::Kernel(format!(
12907 "indexed left_outer: finish_block_use (left col {}) failed: {}",
12908 col_idx, e
12909 ))
12910 })?;
12911 }
12912 result_columns.push(out_col.into());
12913 }
12914
12915 for col_idx in 0..right.arity() {
12916 let elem_size = right
12917 .schema()
12918 .column_type(col_idx)
12919 .map(|t| t.size_bytes())
12920 .unwrap_or(4);
12921 let inner_bytes = (inner_rows as usize)
12922 .checked_mul(elem_size)
12923 .ok_or_else(|| XlogError::Kernel("right inner_bytes overflow".to_string()))?;
12924 let unmatched_bytes = (unmatched_rows as usize)
12925 .checked_mul(elem_size)
12926 .ok_or_else(|| XlogError::Kernel("right unmatched_bytes overflow".to_string()))?;
12927 let total_bytes = inner_bytes
12928 .checked_add(unmatched_bytes)
12929 .ok_or_else(|| XlogError::Kernel("right total_bytes overflow".to_string()))?;
12930 let out_col = self.memory.alloc::<u8>(total_bytes)?;
12931 let dst_ptr = *out_col.device_ptr();
12932 runtime
12934 .prepare_first_use(&out_col, launch_stream, Access::Write)
12935 .map_err(|e| {
12936 XlogError::Kernel(format!(
12937 "indexed left_outer: prepare right out_col {} failed: {}",
12938 col_idx, e
12939 ))
12940 })?;
12941 if total_bytes > 0 {
12942 unsafe {
12944 let res = cudarc::driver::sys::cuMemsetD8Async(
12945 dst_ptr,
12946 0,
12947 total_bytes,
12948 cu_stream.cu_stream(),
12949 );
12950 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12951 return Err(XlogError::Kernel(format!(
12952 "indexed left_outer: zero-fill right col {} failed: {:?}",
12953 col_idx, res
12954 )));
12955 }
12956 }
12957 }
12958 if inner_bytes > 0 {
12959 let src_col = inner_right_buf
12960 .as_ref()
12961 .expect("inner_count > 0")
12962 .column(col_idx)
12963 .ok_or_else(|| XlogError::Kernel("inner_right col missing".to_string()))?;
12964 unsafe {
12966 let res = cudarc::driver::sys::cuMemcpyDtoDAsync_v2(
12967 dst_ptr,
12968 *src_col.device_ptr(),
12969 inner_bytes,
12970 cu_stream.cu_stream(),
12971 );
12972 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
12973 return Err(XlogError::Kernel(format!(
12974 "indexed left_outer: dtod copy inner_right col {} failed: {:?}",
12975 col_idx, res
12976 )));
12977 }
12978 }
12979 }
12980 if let Some(b) = out_col.runtime_block() {
12981 runtime
12982 .finish_block_use(BlockId::from_block(b), launch_stream, Access::Write)
12983 .map_err(|e| {
12984 XlogError::Kernel(format!(
12985 "indexed left_outer: finish_block_use (right col {}) failed: {}",
12986 col_idx, e
12987 ))
12988 })?;
12989 }
12990 result_columns.push(out_col.into());
12991 }
12992
12993 rec_d.commit(runtime).map_err(|e| {
12996 XlogError::Kernel(format!("indexed left_outer: step-D commit failed: {}", e))
12997 })?;
12998
12999 let d_num_rows = self.upload_device_row_count(total_rows as u32)?;
13000 Ok(CudaBuffer::from_columns_with_host_count(
13001 result_columns,
13002 total_rows,
13003 d_num_rows,
13004 combined_schema,
13005 total_rows as u32,
13006 ))
13007 }
13008}