1use std::collections::HashMap;
7use std::marker::PhantomData;
8use std::path::PathBuf;
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::sync::{Arc, Mutex, OnceLock};
11
12use std::ffi::c_void;
13use xlog_core::{Result, Schema, XlogError};
14
15use crate::{
16 cuda_compat::{
17 AsKernelParam, DeviceParamStorage, DevicePtr, DeviceRepr, DeviceSlice,
18 IntoKernelParamStorage, LaunchAsync, LaunchConfig,
19 },
20 cuda_graph::{CapturedCudaGraph, CsmCudaGraphKey, CudaGraphNode},
21 memory::{validate_logical_row_count, CudaColumn, TrackedCudaSlice},
22 CudaBuffer, CudaDevice, CudaStream, CudaViewMut, GpuMemoryManager,
23};
24
25static NEXT_PROVIDER_IDENTITY: AtomicU64 = AtomicU64::new(1);
26
27mod arithmetic;
28mod filter;
29mod fj;
30mod fj_delta;
31mod fj_delta_sparse;
32mod groupby;
33mod ilp;
34mod ilp_exact;
35mod ilp_exact_nary;
36mod io;
37mod kernel_loading;
38pub mod kernel_paths;
39mod launch_safe;
40mod probabilistic;
41mod relational;
42pub mod resident_filter_project;
43pub mod resident_relational;
44pub mod resident_schedule;
45mod transfer;
46mod wcoj;
47mod wcoj_metadata;
48mod wcoj_project;
49
50pub use fj::{FjNode, FjPlan, FjSubAtom};
51pub use fj_delta::{FjDeltaCols, FJ_DELTA_MAX_DOMAIN};
52pub use ilp_exact_nary::{IlpExactNaryPatterns, IlpExactNaryRequest};
53
54#[derive(Debug, Clone, Default)]
56pub struct PtxLoadProfile {
57 pub total_sec: f64,
58 pub per_module_sec: Vec<(String, f64)>,
59 pub cubin_loaded: u32,
60 pub ptx_fallback: u32,
61}
62
63fn warmup_profiling_enabled() -> bool {
64 std::env::var("XLOG_WARMUP_PROFILE")
65 .map(|v| v == "1")
66 .unwrap_or(false)
67}
68
69pub(crate) fn detect_compute_capability(device: &Arc<CudaDevice>) -> Result<u32> {
71 let major = device
72 .inner()
73 .attribute(
74 cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
75 )
76 .map_err(|e| XlogError::Kernel(format!("Failed to query SM major: {}", e)))?;
77 let minor = device
78 .inner()
79 .attribute(
80 cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
81 )
82 .map_err(|e| XlogError::Kernel(format!("Failed to query SM minor: {}", e)))?;
83 Ok((major as u32) * 10 + (minor as u32))
84}
85
86#[cfg(test)]
87fn resolve_module_path(name: &str, cc: u32) -> Option<(std::path::PathBuf, bool)> {
88 kernel_paths::KernelArtifactLocator::from_env().resolve_module_path(name, cc)
89}
90
91#[derive(Debug)]
92pub(crate) enum KernelModuleSource {
93 File { path: PathBuf, is_cubin: bool },
94 EmbeddedPortablePtx { ptx: &'static str },
95}
96
97pub(crate) fn resolve_module_sources_with_locator(
98 name: &str,
99 cc: u32,
100 locator: &kernel_paths::KernelArtifactLocator,
101) -> Vec<KernelModuleSource> {
102 let mut sources: Vec<KernelModuleSource> = locator
103 .resolve_module_paths(name, cc)
104 .into_iter()
105 .filter(|(path, _)| !staged_artifact_is_stale(path))
110 .map(|(path, is_cubin)| KernelModuleSource::File { path, is_cubin })
111 .collect();
112
113 if let Some(ptx) = crate::embedded_kernel_data::portable_ptx(name) {
120 sources.push(KernelModuleSource::EmbeddedPortablePtx { ptx });
121 }
122 sources
123}
124
125fn staged_artifact_is_stale(path: &std::path::Path) -> bool {
134 let Some(file_name) = path.file_name().and_then(|n| n.to_str()) else {
135 return false;
136 };
137 let Some(expected) = crate::embedded_kernel_data::canonical_artifact_hash(file_name) else {
138 return false;
139 };
140 match std::fs::read(path) {
141 Ok(bytes) => fnv1a_64(&bytes) != expected,
142 Err(_) => false,
143 }
144}
145
146fn fnv1a_64(bytes: &[u8]) -> u64 {
148 let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
149 for &byte in bytes {
150 hash ^= byte as u64;
151 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
152 }
153 hash
154}
155
156#[cfg(test)]
157mod kernel_source_resolution_tests {
158 use super::{
159 kernel_paths::KernelArtifactLocator, resolve_module_sources_with_locator,
160 KernelModuleSource,
161 };
162 use std::fs;
163
164 #[test]
165 fn keeps_portable_ptx_fallback_when_cubin_exists() {
166 let root = std::env::temp_dir().join(format!(
167 "xlog-kernel-fallback-{}-{}",
168 std::process::id(),
169 std::time::SystemTime::now()
170 .duration_since(std::time::UNIX_EPOCH)
171 .expect("system clock before UNIX_EPOCH")
172 .as_nanos()
173 ));
174 let kernels = root.join("kernels");
175 fs::create_dir_all(&kernels).expect("create kernels dir");
176 fs::write(kernels.join("fakekernel.sm_86.cubin"), b"cubin").expect("write cubin");
181 fs::write(kernels.join("fakekernel.portable.ptx"), b"ptx").expect("write ptx");
182 let expected_cubin = kernels.join("fakekernel.sm_86.cubin");
183 let expected_ptx = kernels.join("fakekernel.portable.ptx");
184
185 let locator = KernelArtifactLocator::new(None, Some(kernels.clone()), None);
186 let sources = resolve_module_sources_with_locator("fakekernel", 86, &locator);
187
188 assert_eq!(sources.len(), 2);
189 assert!(matches!(
190 &sources[0],
191 KernelModuleSource::File {
192 path,
193 is_cubin: true
194 } if path == &expected_cubin
195 ));
196 assert!(matches!(
197 &sources[1],
198 KernelModuleSource::File {
199 path,
200 is_cubin: false
201 } if path == &expected_ptx
202 ));
203
204 fs::remove_dir_all(root).expect("remove temp kernels");
205 }
206
207 #[test]
212 fn fnv1a_64_matches_known_vectors() {
213 assert_eq!(super::fnv1a_64(b""), 0xcbf2_9ce4_8422_2325);
214 assert_eq!(super::fnv1a_64(b"a"), 0xaf63_dc4c_8601_ec8c);
215 assert_eq!(super::fnv1a_64(b"foobar"), 0x85944171_f73967e8);
216 }
217
218 #[test]
222 fn staged_artifact_not_stale_without_canonical_hash() {
223 let root = std::env::temp_dir().join(format!(
224 "xlog-kernel-stale-{}-{}",
225 std::process::id(),
226 std::time::SystemTime::now()
227 .duration_since(std::time::UNIX_EPOCH)
228 .expect("system clock before UNIX_EPOCH")
229 .as_nanos()
230 ));
231 fs::create_dir_all(&root).expect("create dir");
232 let unknown = root.join("definitely_not_a_real_kernel.sm_86.cubin");
233 fs::write(&unknown, b"bytes").expect("write");
234 assert!(!super::staged_artifact_is_stale(&unknown));
235 assert!(!super::staged_artifact_is_stale(
236 &root.join("missing.portable.ptx")
237 ));
238 fs::remove_dir_all(root).expect("remove temp dir");
239 }
240}
241
242pub(crate) fn load_module_sources(name: &str, cc: u32) -> Result<Vec<KernelModuleSource>> {
247 debug_assert!(
248 crate::kernel_manifest_data::KERNEL_CU_NAMES.contains(&name),
249 "kernel module '{name}' is not in KERNEL_CU_NAMES manifest — update kernel_manifest_data.rs"
250 );
251 let locator = kernel_paths::KernelArtifactLocator::from_env();
252 let sources = resolve_module_sources_with_locator(name, cc, &locator);
253 if sources.is_empty() {
254 Err(XlogError::Kernel(format!(
255 "{name}: no cubin, sidecar portable PTX, or embedded portable PTX found"
256 )))
257 } else {
258 Ok(sources)
259 }
260}
261
262#[derive(Clone)]
269pub(crate) struct RawCudaView<'a, T> {
270 ptr: cudarc::driver::sys::CUdeviceptr,
271 len: usize,
272 stream: Arc<CudaStream>,
273 _marker: PhantomData<&'a [T]>,
274}
275
276pub(crate) struct MultiblockScanScratchU32 {
283 levels: Vec<TrackedCudaSlice<u32>>,
284}
285
286impl MultiblockScanScratchU32 {
287 pub(crate) fn levels(&self) -> &[TrackedCudaSlice<u32>] {
288 &self.levels
289 }
290}
291
292pub(crate) struct CsmCudaGraphNodes {
293 pub(crate) count: CudaGraphNode,
294 pub(crate) total: CudaGraphNode,
295 pub(crate) materialize: CudaGraphNode,
296 pub(crate) node_count: usize,
297}
298
299pub(crate) struct CsmCudaGraphEntry {
300 pub(crate) graph: CapturedCudaGraph,
301 pub(crate) nodes: CsmCudaGraphNodes,
302 pub(crate) per_probe_count: TrackedCudaSlice<u32>,
303 pub(crate) per_probe_offsets: TrackedCudaSlice<u32>,
304 pub(crate) d_logical_count: TrackedCudaSlice<u32>,
305 pub(crate) d_overflow: TrackedCudaSlice<u8>,
306 pub(crate) d_output_left: TrackedCudaSlice<u32>,
307 pub(crate) d_output_right: TrackedCudaSlice<u32>,
308 pub(crate) scan_scratch: MultiblockScanScratchU32,
309 pub(crate) probe_capacity: u32,
310 pub(crate) output_capacity: u32,
311}
312
313impl<'a, T> DeviceSlice<T> for RawCudaView<'a, T> {
314 fn len(&self) -> usize {
315 self.len
316 }
317
318 fn stream(&self) -> &Arc<CudaStream> {
319 &self.stream
320 }
321}
322
323impl<'a, T> DevicePtr<T> for RawCudaView<'a, T> {
324 fn device_ptr<'b>(
325 &'b self,
326 _stream: &'b CudaStream,
327 ) -> (
328 cudarc::driver::sys::CUdeviceptr,
329 cudarc::driver::SyncOnDrop<'b>,
330 ) {
331 (self.ptr, cudarc::driver::SyncOnDrop::Sync(None))
332 }
333}
334
335impl<'a, T> RawCudaView<'a, T> {
336 pub fn device_ptr(&self) -> &cudarc::driver::sys::CUdeviceptr {
337 &self.ptr
338 }
339}
340
341impl<'a, T: DeviceRepr> AsKernelParam for &RawCudaView<'a, T> {
342 fn as_kernel_param(&self) -> *mut c_void {
343 ((*self).device_ptr() as *const cudarc::driver::sys::CUdeviceptr)
344 .cast_mut()
345 .cast()
346 }
347}
348
349impl<'a, T: DeviceRepr> IntoKernelParamStorage for &'a RawCudaView<'a, T> {
350 type Storage = DeviceParamStorage<'a>;
351
352 fn into_kernel_param_storage(self) -> Self::Storage {
353 DeviceParamStorage::unsynced(self.ptr)
354 }
355}
356
357pub struct RadixSortScratch {
359 keys_b: TrackedCudaSlice<u32>,
360 values_b: TrackedCudaSlice<u32>,
361 hist: TrackedCudaSlice<u32>,
362 prefix: TrackedCudaSlice<u32>,
363 ranks: TrackedCudaSlice<u32>,
364 len: u32,
365}
366
367impl RadixSortScratch {
368 pub fn new(provider: &CudaKernelProvider, n: u32) -> Result<Self> {
369 let memory = provider.memory();
370 let len = n.max(1);
371 let keys_b = memory.alloc::<u32>(len as usize)?;
372 let values_b = memory.alloc::<u32>(len as usize)?;
373 let ranks = memory.alloc::<u32>(len as usize)?;
374 let block_size = CudaKernelProvider::SORT_BLOCK_SIZE;
375 let grid_size = len.div_ceil(block_size).max(1);
376 let hist = memory.alloc::<u32>((grid_size as usize) * 16)?;
377 let prefix = memory.alloc::<u32>(16)?;
378 Ok(Self {
379 keys_b,
380 values_b,
381 hist,
382 prefix,
383 ranks,
384 len,
385 })
386 }
387
388 pub fn ensure_capacity(&mut self, provider: &CudaKernelProvider, n: u32) -> Result<()> {
389 if n <= self.len {
390 return Ok(());
391 }
392 *self = Self::new(provider, n)?;
393 Ok(())
394 }
395}
396
397pub const JOIN_MODULE: &str = "xlog_join";
399pub const DEDUP_MODULE: &str = "xlog_dedup";
400pub const GROUPBY_MODULE: &str = "xlog_groupby";
401pub const SCAN_MODULE: &str = "xlog_scan";
402pub const SORT_MODULE: &str = "xlog_sort";
403pub const FILTER_MODULE: &str = "xlog_filter";
404pub const SET_OPS_MODULE: &str = "xlog_set_ops";
405pub const PACK_MODULE: &str = "xlog_pack";
406pub const CIRCUIT_MODULE: &str = "xlog_circuit";
407pub const MC_SAMPLE_MODULE: &str = "xlog_mc_sample";
408pub const MC_EVAL_MODULE: &str = "xlog_mc_eval";
409pub const MC_RESIDENT_MODULE: &str = "xlog_mc_resident";
410pub const ARITH_MODULE: &str = "xlog_arith";
411pub const SAT_MODULE: &str = "xlog_sat";
412pub const D4_MODULE: &str = "xlog_d4";
413pub const NEURAL_MODULE: &str = "xlog_neural";
414pub const PIR_MODULE: &str = "xlog_pir";
415pub const CNF_MODULE: &str = "xlog_cnf";
416pub const CACHE_MODULE: &str = "xlog_cache";
417pub const WEIGHTS_MODULE: &str = "xlog_weights";
418pub const ILP_MODULE: &str = "xlog_ilp";
419pub const ILP_CREDIT_MODULE: &str = "xlog_ilp_credit";
420pub const ILP_EXACT_MODULE: &str = "xlog_ilp_exact";
421pub const ILP_EXACT_NARY_MODULE: &str = "xlog_ilp_exact_nary";
422pub const EPISTEMIC_MODULE: &str = "xlog_epistemic";
423pub const WCOJ_MODULE: &str = "xlog_wcoj";
424pub const JOINT_SOLVE_MODULE: &str = "xlog_joint_solve";
425
426const _: () = assert!(crate::kernel_manifest_data::KERNEL_CU_NAMES.len() == 30);
428
429pub mod wcoj_kernels {
431 pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U32: &str =
432 "wcoj_build_metadata_mark_boundaries_u32";
433 pub const WCOJ_BUILD_METADATA_MARK_BOUNDARIES_U64: &str =
434 "wcoj_build_metadata_mark_boundaries_u64";
435 pub const WCOJ_BUILD_METADATA_SCATTER_U32: &str = "wcoj_build_metadata_scatter_u32";
436 pub const WCOJ_BUILD_METADATA_SCATTER_U64: &str = "wcoj_build_metadata_scatter_u64";
437 pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_triangle_build_hg_work_plan_u32";
438 pub const WCOJ_TRIANGLE_COUNT_HG_U32: &str = "wcoj_triangle_count_hg_u32";
439 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U32: &str =
440 "wcoj_triangle_groupby_root_count_hg_u32";
441 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_triangle_groupby_root_sum_hg_u32";
442 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_triangle_groupby_root_min_hg_u32";
443 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_triangle_groupby_root_max_hg_u32";
444 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U32: &str = "wcoj_triangle_materialize_hg_u32";
445 pub const WCOJ_TRIANGLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_triangle_build_hg_work_plan_u64";
446 pub const WCOJ_TRIANGLE_COUNT_HG_U64: &str = "wcoj_triangle_count_hg_u64";
447 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_COUNT_HG_U64: &str =
448 "wcoj_triangle_groupby_root_count_hg_u64";
449 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_SUM_HG_U64: &str = "wcoj_triangle_groupby_root_sum_hg_u64";
450 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MIN_HG_U64: &str = "wcoj_triangle_groupby_root_min_hg_u64";
451 pub const WCOJ_TRIANGLE_GROUPBY_ROOT_MAX_HG_U64: &str = "wcoj_triangle_groupby_root_max_hg_u64";
452 pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_COUNTS_U32: &str =
453 "wcoj_groupby_root_segment_sum_counts_u32";
454 pub const WCOJ_GROUPBY_ROOT_SEGMENT_SUM_VALUES_U64: &str =
455 "wcoj_groupby_root_segment_sum_values_u64";
456 pub const WCOJ_GROUPBY_ROOT_SEGMENT_MIN_VALUES_U64: &str =
457 "wcoj_groupby_root_segment_min_values_u64";
458 pub const WCOJ_GROUPBY_ROOT_SEGMENT_MAX_VALUES_U64: &str =
459 "wcoj_groupby_root_segment_max_values_u64";
460 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_U64: &str = "wcoj_triangle_materialize_hg_u64";
461 pub const WCOJ_TRIANGLE_COUNT_HG_CACHED_U32: &str = "wcoj_triangle_count_hg_cached_u32";
462 pub const WCOJ_TRIANGLE_MATERIALIZE_HG_CACHED_U32: &str =
463 "wcoj_triangle_materialize_hg_cached_u32";
464 pub const WCOJ_SCAN_HG_BLOCK_COUNTS_U32: &str = "wcoj_scan_hg_block_counts_u32";
465 pub const WCOJ_COMPUTE_TOTAL: &str = "wcoj_compute_total";
466 pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U32: &str = "wcoj_layout_check_sorted_unique_u32";
467 pub const WCOJ_LAYOUT_CHECK_SORTED_UNIQUE_U64: &str = "wcoj_layout_check_sorted_unique_u64";
468 pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U32: &str = "wcoj_4cycle_build_e2_work_prefix_u32";
469 pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U32: &str = "wcoj_4cycle_build_hg_work_plan_u32";
470 pub const WCOJ_4CYCLE_COUNT_HG_U32: &str = "wcoj_4cycle_count_hg_u32";
471 pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U32: &str = "wcoj_4cycle_groupby_root_count_hg_u32";
472 pub const WCOJ_4CYCLE_GROUPBY_ROOT_SUM_HG_U32: &str = "wcoj_4cycle_groupby_root_sum_hg_u32";
473 pub const WCOJ_4CYCLE_GROUPBY_ROOT_MIN_HG_U32: &str = "wcoj_4cycle_groupby_root_min_hg_u32";
474 pub const WCOJ_4CYCLE_GROUPBY_ROOT_MAX_HG_U32: &str = "wcoj_4cycle_groupby_root_max_hg_u32";
475 pub const WCOJ_4CYCLE_MATERIALIZE_HG_U32: &str = "wcoj_4cycle_materialize_hg_u32";
476 pub const WCOJ_4CYCLE_BUILD_E2_WORK_PREFIX_U64: &str = "wcoj_4cycle_build_e2_work_prefix_u64";
477 pub const WCOJ_4CYCLE_BUILD_HG_WORK_PLAN_U64: &str = "wcoj_4cycle_build_hg_work_plan_u64";
478 pub const WCOJ_4CYCLE_COUNT_HG_U64: &str = "wcoj_4cycle_count_hg_u64";
479 pub const WCOJ_4CYCLE_GROUPBY_ROOT_COUNT_HG_U64: &str = "wcoj_4cycle_groupby_root_count_hg_u64";
480 pub const WCOJ_4CYCLE_MATERIALIZE_HG_U64: &str = "wcoj_4cycle_materialize_hg_u64";
481 pub const WCOJ_CLIQUE5_COUNT_HG_U32: &str = "wcoj_clique5_count_hg_u32";
483 pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U32: &str = "wcoj_clique5_materialize_hg_u32";
484 pub const WCOJ_CLIQUE5_COUNT_HG_U64: &str = "wcoj_clique5_count_hg_u64";
485 pub const WCOJ_CLIQUE5_MATERIALIZE_HG_U64: &str = "wcoj_clique5_materialize_hg_u64";
486 pub const WCOJ_CLIQUE6_COUNT_HG_U32: &str = "wcoj_clique6_count_hg_u32";
487 pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U32: &str = "wcoj_clique6_materialize_hg_u32";
488 pub const WCOJ_CLIQUE6_COUNT_HG_U64: &str = "wcoj_clique6_count_hg_u64";
489 pub const WCOJ_CLIQUE6_MATERIALIZE_HG_U64: &str = "wcoj_clique6_materialize_hg_u64";
490 pub const WCOJ_CLIQUE7_COUNT_HG_U32: &str = "wcoj_clique7_count_hg_u32";
491 pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U32: &str = "wcoj_clique7_materialize_hg_u32";
492 pub const WCOJ_CLIQUE7_COUNT_HG_U64: &str = "wcoj_clique7_count_hg_u64";
493 pub const WCOJ_CLIQUE7_MATERIALIZE_HG_U64: &str = "wcoj_clique7_materialize_hg_u64";
494 pub const WCOJ_CLIQUE8_COUNT_HG_U32: &str = "wcoj_clique8_count_hg_u32";
495 pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U32: &str = "wcoj_clique8_materialize_hg_u32";
496 pub const WCOJ_CLIQUE8_COUNT_HG_U64: &str = "wcoj_clique8_count_hg_u64";
497 pub const WCOJ_CLIQUE8_MATERIALIZE_HG_U64: &str = "wcoj_clique8_materialize_hg_u64";
498 pub const WCOJ_CLIQUE5_GROUPBY_ROOT_COUNT_HG_U32: &str =
499 "wcoj_clique5_groupby_root_count_hg_u32";
500 pub const WCOJ_CLIQUE6_GROUPBY_ROOT_COUNT_HG_U32: &str =
501 "wcoj_clique6_groupby_root_count_hg_u32";
502 pub const FJ_EXPAND_WORK_PREFIX_U32: &str = "fj_expand_work_prefix_u32";
506 pub const FJ_EXPAND_COUNT_U32: &str = "fj_expand_count_u32";
507 pub const FJ_EXPAND_EMIT_U32: &str = "fj_expand_emit_u32";
508 pub const FJ_PROBE_REFINE_U32: &str = "fj_probe_refine_u32";
509 pub const FJ_EXPAND_COUNT_U64: &str = "fj_expand_count_u64";
510 pub const FJ_EXPAND_EMIT_U64: &str = "fj_expand_emit_u64";
511 pub const FJ_PROBE_REFINE_U64: &str = "fj_probe_refine_u64";
512 pub const FJ_COUNT_MULTIPLICITY: &str = "fj_count_multiplicity";
513 pub const FJ_DELTA_RANGE_U32: &str = "fj_delta_range_u32";
515 pub const FJ_DELTA_MARK_U32: &str = "fj_delta_mark_u32";
516 pub const FJ_DELTA_SUBTRACT_U32: &str = "fj_delta_subtract_u32";
517 pub const FJ_DELTA_POPCOUNT: &str = "fj_delta_popcount";
518 pub const FJ_DELTA_EMIT_U32: &str = "fj_delta_emit_u32";
519 pub const FJ_DELTA_MAX_U32: &str = "fj_delta_max_u32";
520 pub const FJ_DELTA_SPARSE_ESTIMATE: &str = "fj_delta_sparse_estimate";
521 pub const FJ_DELTA_SPARSE_LOAD_R: &str = "fj_delta_sparse_load_r";
522 pub const FJ_DELTA_SPARSE_INSERT_CANDIDATES: &str = "fj_delta_sparse_insert_candidates";
523 pub const FJ_DELTA_SPARSE_MARK: &str = "fj_delta_sparse_mark";
524 pub const FJ_DELTA_SPARSE_EMIT: &str = "fj_delta_sparse_emit";
525}
526
527pub mod mc_sample_kernels {
529 pub const MC_SAMPLE_BERNOULLI: &str = "mc_sample_bernoulli";
530}
531
532pub mod mc_eval_kernels {
534 pub const MC_EVAL_MASK_VAR: &str = "mc_eval_mask_var";
535 pub const MC_EVAL_MASK_AD: &str = "mc_eval_mask_ad_choice";
536 pub const MC_EVAL_QUERY_EVIDENCE_TRUTH: &str = "mc_eval_query_evidence_truth";
537 pub const MC_EVAL_ACCUMULATE_COUNTS: &str = "mc_accumulate_counts";
538}
539
540pub mod mc_resident_kernels {
542 pub const MC_RESIDENT_ENGINE: &str = "mc_resident_engine";
545}
546
547pub mod arith_kernels {
549 pub const ARITH_BINARY_I64: &str = "arith_binary_i64";
550 pub const ARITH_BINARY_I32: &str = "arith_binary_i32";
551 pub const ARITH_BINARY_U64: &str = "arith_binary_u64";
552 pub const ARITH_BINARY_U32: &str = "arith_binary_u32";
553 pub const ARITH_BINARY_F64: &str = "arith_binary_f64";
554 pub const ARITH_BINARY_F32: &str = "arith_binary_f32";
555 pub const ARITH_ABS_I64: &str = "arith_abs_i64";
556 pub const ARITH_ABS_I32: &str = "arith_abs_i32";
557 pub const ARITH_ABS_F64: &str = "arith_abs_f64";
558 pub const ARITH_ABS_F32: &str = "arith_abs_f32";
559 pub const ARITH_POW_F64: &str = "arith_pow_f64";
560 pub const ARITH_CAST: &str = "arith_cast";
561 pub const ARITH_FILL_CONST_U32: &str = "arith_fill_const_u32";
562 pub const ARITH_FILL_CONST_U64: &str = "arith_fill_const_u64";
563 pub const ARITH_FILL_CONST_I64: &str = "arith_fill_const_i64";
564 pub const ARITH_FILL_CONST_I32: &str = "arith_fill_const_i32";
565 pub const ARITH_FILL_CONST_F64: &str = "arith_fill_const_f64";
566 pub const ARITH_FILL_CONST_F32: &str = "arith_fill_const_f32";
567 pub const ARITH_FILL_CONST_U8: &str = "arith_fill_const_u8";
568 pub const ARITH_SELECT_I64: &str = "arith_select_i64";
570 pub const ARITH_SELECT_I32: &str = "arith_select_i32";
571 pub const ARITH_SELECT_U64: &str = "arith_select_u64";
572 pub const ARITH_SELECT_U32: &str = "arith_select_u32";
573 pub const ARITH_SELECT_F64: &str = "arith_select_f64";
574 pub const ARITH_SELECT_F32: &str = "arith_select_f32";
575}
576
577pub mod epistemic_kernels {
579 pub const EPISTEMIC_GENERATE_CANDIDATE_ASSUMPTIONS_U8: &str =
581 "epistemic_generate_candidate_assumptions_u8";
582 pub const EPISTEMIC_PROPAGATE_CANDIDATES_U8: &str = "epistemic_propagate_candidates_u8";
584 pub const EPISTEMIC_VALIDATE_CANDIDATE_BITS_U8: &str = "epistemic_validate_candidate_bits_u8";
586 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_U8: &str =
588 "epistemic_populate_model_membership_u8";
589 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_U8: &str =
591 "epistemic_populate_model_membership_from_tuple_source_u8";
592 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY1_U8: &str =
594 "epistemic_populate_model_membership_from_tuple_source_arity1_u8";
595 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY2_U8: &str =
597 "epistemic_populate_model_membership_from_tuple_source_arity2_u8";
598 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY3_U8: &str =
600 "epistemic_populate_model_membership_from_tuple_source_arity3_u8";
601 pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY_N_U8: &str =
603 "epistemic_populate_model_membership_from_tuple_source_arity_n_u8";
604 pub const EPISTEMIC_VALIDATE_WORLD_VIEWS_U8: &str = "epistemic_validate_world_views_u8";
606 pub const EPISTEMIC_VALIDATE_CONSTRAINTS_U8: &str = "epistemic_validate_constraints_u8";
608 pub const EPISTEMIC_MATERIALIZE_ACCEPTED_CANDIDATES_U8: &str =
610 "epistemic_materialize_accepted_candidates_u8";
611
612 pub const EPISTEMIC_MATERIALIZE_FINAL_RESULT_FLAGS_U8: &str =
614 "epistemic_materialize_final_result_flags_u8";
615 pub const EPISTEMIC_MATERIALIZE_FINAL_TUPLE_COLUMN_U8: &str =
617 "epistemic_materialize_final_tuple_column_u8";
618 pub const EPISTEMIC_BUILD_FINAL_TUPLE_ROW_MAP_U8: &str =
620 "epistemic_build_final_tuple_row_map_u8";
621 pub const EPISTEMIC_CLOSE_FINAL_TUPLE_REJECTIONS_U8: &str =
623 "epistemic_close_final_tuple_rejections_u8";
624}
625
626pub mod neural_kernels {
628 pub const NEURAL_FILL_AD_CHAIN_F32: &str = "neural_fill_ad_chain_f32";
629 pub const NEURAL_SCATTER_AD_CHAIN_GRADS_F32: &str = "neural_scatter_ad_chain_grads_f32";
630}
631
632pub mod ilp_kernels {
634 pub const EXTRACT_NONZERO_INDICES: &str = "extract_nonzero_indices";
635 pub const ILP_MARK_SELECTED_IDS_U32: &str = "ilp_mark_selected_ids_u32";
636 pub const ILP_MARK_SELECTED_IDS_I32: &str = "ilp_mark_selected_ids_i32";
637 pub const ILP_MARK_SELECTED_IDS_I64: &str = "ilp_mark_selected_ids_i64";
638 pub const ILP_MARK_SELECTED_IDS_U64: &str = "ilp_mark_selected_ids_u64";
639 pub const ILP_VALIDATE_SELECTED_IDS_U32: &str = "ilp_validate_selected_ids_u32";
640 pub const ILP_VALIDATE_SELECTED_IDS_I32: &str = "ilp_validate_selected_ids_i32";
641 pub const ILP_VALIDATE_SELECTED_IDS_I64: &str = "ilp_validate_selected_ids_i64";
642 pub const ILP_VALIDATE_SELECTED_IDS_U64: &str = "ilp_validate_selected_ids_u64";
643 pub const ILP_BROADCAST_CANDIDATE_FLAG: &str = "ilp_broadcast_candidate_flag";
644 pub const ILP_COO_FILL_FROM_MASK: &str = "ilp_coo_fill_from_mask";
645 pub const ILP_CSR_HISTOGRAM: &str = "ilp_csr_histogram";
646 pub const ILP_REDUCE_SUM_F32: &str = "ilp_reduce_sum_f32";
647 pub const ILP_REDUCE_SUM_F64: &str = "ilp_reduce_sum_f64";
648}
649
650pub mod ilp_credit_kernels {
652 pub const ILP_COO_FILL: &str = "ilp_coo_fill";
653 pub const ILP_CREDIT_FORWARD_F32: &str = "ilp_credit_forward_f32";
654 pub const ILP_CREDIT_FORWARD_F64: &str = "ilp_credit_forward_f64";
655 pub const ILP_CREDIT_BACKWARD_F32: &str = "ilp_credit_backward_f32";
656 pub const ILP_CREDIT_BACKWARD_F64: &str = "ilp_credit_backward_f64";
657}
658
659pub mod ilp_exact_kernels {
661 pub const ILP_EXACT_SCORE: &str = "ilp_exact_score";
662 pub const ILP_EXACT_SCORE_U32: &str = "ilp_exact_score_u32";
663 pub const ILP_EXACT_SCORE_CHAIN_SMEM: &str = "ilp_exact_score_chain_smem";
664 pub const ILP_EXACT_SCORE_CHAIN_SMEM_U32: &str = "ilp_exact_score_chain_smem_u32";
665 pub const ILP_EXACT_SELECT_TOPK: &str = "ilp_exact_select_topk";
666}
667
668pub mod ilp_exact_nary_kernels {
670 pub const ILP_EXACT_NARY_SCORE: &str = "ilp_exact_nary_score";
671}
672
673pub mod pir_kernels {
675 pub const PIR_PACK_KEYS: &str = "pir_pack_keys";
676 pub const PIR_HASH_KEYS: &str = "pir_hash_keys";
677 pub const PIR_MARK_UNIQUE: &str = "pir_mark_unique";
678 pub const PIR_FIND_EXISTING: &str = "pir_find_existing";
679 pub const PIR_MARK_NEW_GROUPS: &str = "pir_mark_new_groups";
680 pub const PIR_BUILD_GROUP_IDS: &str = "pir_build_group_ids";
681 pub const PIR_FILL_CHILD_PARENTS: &str = "pir_fill_child_parents";
682 pub const PIR_MARK_UNIQUE_PAIRS: &str = "pir_mark_unique_pairs";
683 pub const PIR_COMPACT_PAIRS: &str = "pir_compact_pairs";
684 pub const PIR_COUNT_CHILDREN: &str = "pir_count_children";
685 pub const PIR_WRITE_CHILD_OFFSETS: &str = "pir_write_child_offsets";
686 pub const PIR_GATHER_CHILDREN: &str = "pir_gather_children";
687 pub const PIR_BUILD_GRAPH_CHILD_COUNTS: &str = "pir_build_graph_child_counts";
688 pub const PIR_SUM_COUNTS: &str = "pir_sum_counts";
689 pub const PIR_EMIT_NODES_AND_IDS: &str = "pir_emit_nodes_and_ids";
690 pub const PIR_UPDATE_COUNTS: &str = "pir_update_counts";
691}
692
693pub mod cnf_kernels {
695 pub const CNF_REACHABILITY_INIT: &str = "cnf_reachability_init";
696 pub const CNF_REACHABILITY_BFS: &str = "cnf_reachability_bfs";
697 pub const CNF_MARK_LEAF_CHOICE: &str = "cnf_mark_leaf_choice";
698 pub const CNF_ASSIGN_LEAF_VAR: &str = "cnf_assign_leaf_var";
699 pub const CNF_ASSIGN_CHOICE_VAR: &str = "cnf_assign_choice_var";
700 pub const CNF_MARK_NODE_VARS: &str = "cnf_mark_node_vars";
701 pub const CNF_COUNT_CLAUSES: &str = "cnf_count_clauses";
702 pub const CNF_CAPTURE_LAST_COUNTS: &str = "cnf_capture_last_counts";
703 pub const CNF_COMPUTE_LEAF_CHOICE_TOTALS: &str = "cnf_compute_leaf_choice_totals";
704 pub const CNF_COMPUTE_TOTALS: &str = "cnf_compute_totals";
705 pub const CNF_ASSIGN_NODE_VAR: &str = "cnf_assign_node_var";
706 pub const CNF_EMIT_CLAUSES: &str = "cnf_emit_clauses";
707 pub const CNF_SET_CLAUSE_END: &str = "cnf_set_clause_end";
708}
709
710pub mod weights_kernels {
712 pub const WEIGHTS_FILL_LEAF: &str = "weights_fill_leaf";
713 pub const WEIGHTS_FILL_CHOICE: &str = "weights_fill_choice";
714 pub const WEIGHTS_COUNT_LIFT_EXACT: &str = "weights_count_lift_exact";
715 pub const WEIGHTS_SET_EVIDENCE_FROM_NODES: &str = "weights_set_evidence_from_nodes";
716 pub const WEIGHTS_APPLY_EVIDENCE: &str = "weights_apply_evidence";
717 pub const WEIGHTS_MAP_NODES_TO_VARS: &str = "weights_map_nodes_to_vars";
718 pub const WEIGHTS_FORCE_VAR_FALSE: &str = "weights_force_var_false";
719 pub const WEIGHTS_RESTORE_VAR_FALSE: &str = "weights_restore_var_false";
720 pub const WEIGHTS_FORCE_VAR_TRUE: &str = "weights_force_var_true";
721 pub const WEIGHTS_RESTORE_VAR_TRUE: &str = "weights_restore_var_true";
722 pub const WEIGHTS_COPY_SLOT_TO_BATCH: &str = "weights_copy_slot_to_batch";
723 pub const WEIGHTS_APPLY_QUERY_VARS: &str = "weights_apply_query_vars";
724 pub const WEIGHTS_RESTORE_QUERY_VARS: &str = "weights_restore_query_vars";
725 pub const WEIGHTS_APPLY_QUERY_VARS_FALSE_BATCHED: &str =
726 "weights_apply_query_vars_false_batched";
727 pub const WEIGHTS_RESTORE_QUERY_VARS_FALSE_BATCHED: &str =
728 "weights_restore_query_vars_false_batched";
729 pub const WEIGHTS_APPLY_QUERY_VARS_TRUE_BATCHED: &str = "weights_apply_query_vars_true_batched";
730 pub const WEIGHTS_RESTORE_QUERY_VARS_TRUE_BATCHED: &str =
731 "weights_restore_query_vars_true_batched";
732}
733
734pub mod d4_kernels {
737 pub const D4_VALIDATE_CNF: &str = "d4_validate_cnf";
738 pub const D4_LEVELIZE_COUNTS: &str = "d4_levelize_counts";
739 pub const D4_LEVELIZE_EMIT: &str = "d4_levelize_emit";
740 pub const D4_FRONTIER_PREPARE: &str = "d4_frontier_prepare";
742 pub const D4_FRONTIER_EXPAND: &str = "d4_frontier_expand";
743 pub const D4_COMPILE_COUNT: &str = "d4_compile_count";
745 pub const D4_COMPILE_EMIT: &str = "d4_compile_emit";
746 pub const D4_CAPTURE_EMIT_META: &str = "d4_capture_emit_meta";
747 pub const D4_SUPPORT_LEVEL: &str = "d4_support_level";
749 pub const D4_SUPPORT_SET_ROOT_BITS: &str = "d4_support_set_root_bits";
750 pub const D4_SMOOTH_COUNT: &str = "d4_smooth_count";
751 pub const D4_SMOOTH_WRAPPER_COUNTS: &str = "d4_smooth_wrapper_counts";
752 pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_OR: &str = "d4_smooth_wrapper_edge_counts_or";
753 pub const D4_SMOOTH_WRAPPER_EDGE_COUNTS_DEC: &str = "d4_smooth_wrapper_edge_counts_dec";
754 pub const D4_SMOOTH_INIT_NODES: &str = "d4_smooth_init_nodes";
755 pub const D4_SMOOTH_EMIT_LEVEL: &str = "d4_smooth_emit_level";
756 pub const D4_SMOOTH_CHECK_EDGE_CAP: &str = "d4_smooth_check_edge_cap";
757 pub const D4_MARK_VARS_IN_CLAUSES: &str = "d4_mark_vars_in_clauses";
759 pub const D4_MARK_VARS_IN_CIRCUIT: &str = "d4_mark_vars_in_circuit";
760 pub const D4_BUILD_FREE_VAR_MASK: &str = "d4_build_free_var_mask";
761 pub const D4_ASSERT_U32_EQ: &str = "d4_assert_u32_eq";
763 pub const D4_ASSERT_BITSET_VAR: &str = "d4_assert_bitset_var";
764 pub const D4_ASSERT_LEAF_ROOT_AND_DEGREE: &str = "d4_assert_leaf_root_and_degree";
765}
766
767pub mod join_kernels {
769 pub const HASH_JOIN_BUILD: &str = "hash_join_build";
770 pub const HASH_JOIN_PROBE: &str = "hash_join_probe";
771 pub const COMPUTE_COMPOSITE_HASH: &str = "compute_composite_hash";
773 pub const HASH_JOIN_BUCKET_COUNT_V2: &str = "hash_join_bucket_count_v2";
774 pub const HASH_JOIN_SCATTER_V2: &str = "hash_join_scatter_v2";
775 pub const HASH_JOIN_PROBE_V2: &str = "hash_join_probe_v2";
776 pub const HASH_JOIN_PROBE_V2_COUNT_PER_ROW: &str = "hash_join_probe_v2_count_per_row";
777 pub const HASH_JOIN_PROBE_V2_MATERIALIZE: &str = "hash_join_probe_v2_materialize";
778 pub const HASH_JOIN_TOTAL_FROM_SCAN: &str = "hash_join_total_from_scan";
779 pub const HASH_JOIN_CSM_UNMATCHED_MASK: &str = "hash_join_csm_unmatched_mask";
780 pub const HASH_JOIN_SEMI: &str = "hash_join_semi";
781 pub const HASH_JOIN_ANTI: &str = "hash_join_anti";
782 pub const INIT_HASH_TABLE: &str = "init_hash_table";
783 pub const NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS: &str = "nested_loop_join_inner_u32_1key_pairs";
789 pub const SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS: &str = "sort_merge_join_inner_u32_1key_pairs";
797}
798
799pub mod dedup_kernels {
801 pub const MARK_DUPLICATES: &str = "mark_duplicates";
802 pub const MARK_UNIQUE_COLUMNAR: &str = "mark_unique_columnar";
803 pub const MARK_UNIQUE_AND_SCAN_COLUMNAR: &str = "mark_unique_and_scan_columnar";
804 pub const COMPACT_ROWS: &str = "compact_rows";
805 pub const MARK_UNIQUE_FULL_ROW_BYTEWISE: &str = "mark_unique_full_row_bytewise";
806 pub const MARK_DIFF_FULL_ROW_TYPED_SORTED: &str = "mark_diff_full_row_typed_sorted";
807 pub const SMALL_SORT_FULL_ROW_INDICES_TYPED: &str = "small_sort_full_row_indices_typed";
808}
809
810pub mod groupby_kernels {
812 pub const DETECT_GROUP_BOUNDARIES: &str = "detect_group_boundaries";
813 pub const DETECT_BOUNDARIES: &str = "detect_boundaries";
814 pub const EXTRACT_GROUP_KEYS: &str = "extract_group_keys";
815 pub const GROUP_IDS_FROM_BOUNDARIES: &str = "group_ids_from_boundaries";
816 pub const GROUP_START_INDICES: &str = "group_start_indices";
817 pub const CAPTURE_NUM_GROUPS: &str = "capture_num_groups";
818 pub const GROUPBY_COUNT: &str = "groupby_count";
819 pub const GROUPBY_SUM: &str = "groupby_sum";
820 pub const GROUPBY_SUM_U64: &str = "groupby_sum_u64";
821 pub const GROUPBY_MIN: &str = "groupby_min";
822 pub const GROUPBY_MIN_U64: &str = "groupby_min_u64";
823 pub const GROUPBY_MAX: &str = "groupby_max";
824 pub const GROUPBY_MAX_U64: &str = "groupby_max_u64";
825 pub const GROUPBY_LOGSUMEXP_MAX: &str = "groupby_logsumexp_max";
826 pub const GROUPBY_LOGSUMEXP_SUMEXP: &str = "groupby_logsumexp_sumexp";
827 pub const GROUPBY_LOGSUMEXP_FINAL: &str = "groupby_logsumexp_final";
828}
829
830pub mod scan_kernels {
832 pub const BLOCK_INCLUSIVE_SCAN: &str = "block_inclusive_scan";
833 pub const ADD_BLOCK_OFFSETS: &str = "add_block_offsets";
834 pub const EXCLUSIVE_SCAN_MASK: &str = "exclusive_scan_mask";
835 pub const COUNT_MASK: &str = "count_mask";
836 pub const MULTIBLOCK_SCAN_PHASE1: &str = "multiblock_scan_phase1";
838 pub const MULTIBLOCK_SCAN_U32_PHASE1: &str = "multiblock_scan_u32_phase1";
839 pub const MULTIBLOCK_SCAN_PHASE2: &str = "multiblock_scan_phase2";
840 pub const MULTIBLOCK_SCAN_PHASE3: &str = "multiblock_scan_phase3";
841}
842
843pub mod sort_kernels {
845 pub const RADIX_HISTOGRAM: &str = "radix_histogram";
846 pub const RADIX_SCATTER: &str = "radix_scatter";
847 pub const COMPUTE_RANKS: &str = "compute_ranks";
848 pub const RADIX_SCATTER_STABLE: &str = "radix_scatter_stable";
849 pub const COMPUTE_DIGIT_PREFIX_SUMS: &str = "compute_digit_prefix_sums";
850 pub const INIT_INDICES: &str = "init_indices";
851 pub const APPLY_PERMUTATION_U32: &str = "apply_permutation_u32";
852 pub const APPLY_PERMUTATION_BYTES: &str = "apply_permutation_bytes";
853
854 pub const GATHER_KEYS_I32_ORDERED_U32: &str = "gather_keys_i32_ordered_u32";
855 pub const GATHER_KEYS_F32_ORDERED_U32: &str = "gather_keys_f32_ordered_u32";
856 pub const GATHER_KEYS_BOOL_ORDERED_U32: &str = "gather_keys_bool_ordered_u32";
857
858 pub const GATHER_KEYS_U64_LO_U32: &str = "gather_keys_u64_lo_u32";
859 pub const GATHER_KEYS_U64_HI_U32: &str = "gather_keys_u64_hi_u32";
860
861 pub const GATHER_KEYS_I64_LO_U32: &str = "gather_keys_i64_lo_u32";
862 pub const GATHER_KEYS_I64_HI_U32: &str = "gather_keys_i64_hi_u32";
863
864 pub const GATHER_KEYS_F64_LO_U32: &str = "gather_keys_f64_lo_u32";
865 pub const GATHER_KEYS_F64_HI_U32: &str = "gather_keys_f64_hi_u32";
866 pub const CHECK_ASCENDING_SORTED_U32: &str = "check_ascending_sorted_u32";
874}
875
876pub mod filter_kernels {
878 pub const FILTER_COMPARE_U32: &str = "filter_compare_u32";
879 pub const FILTER_COMPARE_I64: &str = "filter_compare_i64";
880 pub const FILTER_COMPARE_F64: &str = "filter_compare_f64";
881 pub const FILTER_COMPARE_I32: &str = "filter_compare_i32";
882 pub const FILTER_COMPARE_U64: &str = "filter_compare_u64";
883 pub const FILTER_COMPARE_F32: &str = "filter_compare_f32";
884 pub const FILTER_COMPARE_U8: &str = "filter_compare_u8";
885 pub const FILTER_COMPARE_U32_SCAN_PHASE1: &str = "filter_compare_u32_scan_phase1";
886 pub const FILTER_COMPARE_F64_SCAN_PHASE1: &str = "filter_compare_f64_scan_phase1";
887 pub const FILTER_COMPARE_F32_SCAN_PHASE1: &str = "filter_compare_f32_scan_phase1";
888 pub const FILTER_COMPARE_U32_COL: &str = "filter_compare_u32_col";
889 pub const FILTER_COMPARE_I32_COL: &str = "filter_compare_i32_col";
890 pub const FILTER_COMPARE_I64_COL: &str = "filter_compare_i64_col";
891 pub const FILTER_COMPARE_U64_COL: &str = "filter_compare_u64_col";
892 pub const FILTER_COMPARE_F32_COL: &str = "filter_compare_f32_col";
893 pub const FILTER_COMPARE_F64_COL: &str = "filter_compare_f64_col";
894 pub const FILTER_COMPARE_U8_COL: &str = "filter_compare_u8_col";
895 pub const FILL_U32_IOTA: &str = "fill_u32_iota";
896 pub const FILL_U32_CONST: &str = "fill_u32_const";
897 pub const MARK_RANDOM_VARS: &str = "mark_random_vars";
898 pub const RANDOM_VAR_TO_BIT_FROM_LIST: &str = "random_var_to_bit_from_list";
899 pub const CHECK_RANDOM_VAR_COUNT: &str = "check_random_var_count";
900 pub const COMPACT_U32_BY_MASK: &str = "compact_u32_by_mask";
901 pub const COMPACT_I64_BY_MASK: &str = "compact_i64_by_mask";
902 pub const COMPACT_F64_BY_MASK: &str = "compact_f64_by_mask";
903 pub const COMPACT_BYTES_BY_MASK: &str = "compact_bytes_by_mask";
904 pub const CAPTURE_COMPACT_COUNT: &str = "capture_compact_count";
905 pub const MASK_CLAMP_ROWS: &str = "mask_clamp_rows";
906 pub const MASK_AND: &str = "mask_and";
907 pub const MASK_OR: &str = "mask_or";
908 pub const MASK_NOT: &str = "mask_not";
909}
910
911pub mod set_ops_kernels {
913 pub const CONCAT_U32: &str = "concat_u32";
914 pub const CONCAT_BYTES: &str = "concat_bytes";
915 pub const SORTED_DIFF_MARK: &str = "sorted_diff_mark";
916}
917
918pub mod pack_kernels {
920 pub const PACK_KEYS: &str = "pack_keys";
922 pub const HASH_PACKED_KEYS: &str = "hash_packed_keys";
924 pub const PACK_AND_HASH_KEYS: &str = "pack_and_hash_keys";
926 pub const PACK_AND_HASH_KEYS_GENERIC: &str = "pack_and_hash_keys_generic";
928 pub const PACK_KEYS_ALIGNED: &str = "pack_keys_aligned";
930 pub const UNPACK_COLUMN: &str = "unpack_column";
932 pub const UNPACK_COLUMN_COUNTED: &str = "unpack_column_counted";
934 pub const GATHER_PACKED_ROWS: &str = "gather_packed_rows";
936 pub const GATHER_PACKED_ROWS_COUNTED: &str = "gather_packed_rows_counted";
938 pub const SCATTER_PACKED_ROWS: &str = "scatter_packed_rows";
940 pub const COMPARE_PACKED_KEYS: &str = "compare_packed_keys";
942 pub const PACK_BOOLS_TO_BITMAP: &str = "pack_bools_to_bitmap";
944}
945
946pub mod circuit_kernels {
948 pub const XGCF_FORWARD_LEVEL: &str = "xgcf_forward_level";
949 pub const XGCF_BACKWARD_LEVEL_PROPAGATE: &str = "xgcf_backward_level_propagate";
950 pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD: &str = "xgcf_backward_level_decision_grad";
951 pub const XGCF_BACKWARD_LEVEL_LIT_GRAD: &str = "xgcf_backward_level_lit_grad";
952 pub const XGCF_FREE_VAR_APPLY_GRAD: &str = "xgcf_free_var_apply_grad";
953 pub const XGCF_FREE_VAR_REDUCE_STAGE: &str = "xgcf_free_var_reduce_stage";
954 pub const XGCF_ADD_SCALAR: &str = "xgcf_add_scalar";
955 pub const XGCF_FORWARD_LEVEL_CACHED: &str = "xgcf_forward_level_cached";
956 pub const XGCF_EVAL_ALL_LEVELS_CACHED: &str = "xgcf_eval_all_levels_cached";
957 pub const XGCF_EVAL_ALL_LEVELS_CACHED_BATCHED: &str = "xgcf_eval_all_levels_cached_batched";
958 pub const XGCF_BACKWARD_LEVEL_PROPAGATE_CACHED: &str = "xgcf_backward_level_propagate_cached";
959 pub const XGCF_BACKWARD_LEVEL_DECISION_GRAD_CACHED: &str =
960 "xgcf_backward_level_decision_grad_cached";
961 pub const XGCF_BACKWARD_LEVEL_LIT_GRAD_CACHED: &str = "xgcf_backward_level_lit_grad_cached";
962 pub const XGCF_BACKWARD_ALL_LEVELS_CACHED: &str = "xgcf_backward_all_levels_cached";
963 pub const XGCF_BACKWARD_ALL_LEVELS_CACHED_BATCHED: &str =
964 "xgcf_backward_all_levels_cached_batched";
965 pub const XGCF_FREE_VAR_APPLY_GRAD_CACHED: &str = "xgcf_free_var_apply_grad_cached";
966 pub const XGCF_FREE_VAR_REDUCE_STAGE_CACHED: &str = "xgcf_free_var_reduce_stage_cached";
967 pub const XGCF_ADD_SCALAR_CACHED: &str = "xgcf_add_scalar_cached";
968 pub const XGCF_SET_ROOT_ADJ_CACHED_BATCHED: &str = "xgcf_set_root_adj_cached_batched";
969 pub const XGCF_COPY_ROOT_CACHED: &str = "xgcf_copy_root_cached";
970 pub const XGCF_COPY_ROOT_CACHED_META: &str = "xgcf_copy_root_cached_meta";
971 pub const XGCF_COPY_ROOT_CACHED_META_BATCHED: &str = "xgcf_copy_root_cached_meta_batched";
972}
973
974pub mod cache_kernels {
976 pub const CACHE_CNF_HASH: &str = "cache_cnf_hash";
977 pub const CACHE_LOOKUP_OR_INSERT: &str = "cache_lookup_or_insert";
978 pub const CACHE_EVICT_LRU: &str = "cache_evict_lru";
979 pub const CACHE_STORE_U8: &str = "cache_store_u8";
980 pub const CACHE_STORE_U32: &str = "cache_store_u32";
981 pub const CACHE_STORE_I32: &str = "cache_store_i32";
982 pub const CACHE_STORE_F64: &str = "cache_store_f64";
983 pub const CACHE_STORE_META: &str = "cache_store_meta";
984}
985
986pub mod sat_kernels {
988 pub const SAT_CDCL_SOLVE: &str = "sat_cdcl_solve";
989 pub const SAT_CHECK_MODEL: &str = "sat_check_model";
990 pub const SAT_PROOF_MARK_NEEDED: &str = "sat_proof_mark_needed";
991 pub const SAT_PROOF_CHECK: &str = "sat_proof_check";
992 pub const SAT_ASSERT_STATUS: &str = "sat_assert_status";
993 pub const SAT_ASSERT_OK: &str = "sat_assert_ok";
994 pub const SAT_XGCF_CNF_COUNTS: &str = "sat_xgcf_cnf_counts";
995 pub const SAT_XGCF_CNF_EMIT: &str = "sat_xgcf_cnf_emit";
996 pub const SAT_XGCF_CNF_CAPTURE_LAST_COUNTS: &str = "sat_xgcf_cnf_capture_last_counts";
997 pub const SAT_XGCF_CNF_COMPUTE_TOTALS: &str = "sat_xgcf_cnf_compute_totals";
998 pub const SAT_CNF_WRITE_TERMINATOR: &str = "sat_cnf_write_terminator";
999 pub const SAT_CNF_COPY_INTO: &str = "sat_cnf_copy_into";
1000 pub const SAT_SHIFT_OFFSETS: &str = "sat_shift_offsets";
1001 pub const SAT_XGCF_WRITE_ROOT_UNIT_CLAUSE: &str = "sat_xgcf_write_root_unit_clause";
1002 pub const SAT_NOT_PHI_COUNTS: &str = "sat_not_phi_counts";
1003 pub const SAT_EMIT_NOT_PHI: &str = "sat_emit_not_phi";
1004}
1005
1006pub const DEFAULT_JOIN_MAX_OUTPUT: usize = 1_000_000;
1009
1010pub const NESTED_LOOP_TOTAL_THRESHOLD: u64 = 4_000_000;
1031
1032#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1034#[repr(u8)]
1035pub enum CompareOp {
1036 Eq = 0,
1037 Ne = 1,
1038 Lt = 2,
1039 Le = 3,
1040 Gt = 4,
1041 Ge = 5,
1042}
1043
1044#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1046pub enum JoinType {
1047 Inner,
1049 Semi,
1051 Anti,
1053 LeftOuter,
1055}
1056
1057struct PackedKeyData {
1059 hashes: crate::memory::TrackedCudaSlice<u64>,
1061 packed_keys: crate::memory::TrackedCudaSlice<u8>,
1063 key_bytes: u32,
1065}
1066
1067struct JoinHashTableV2 {
1068 bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1069 bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1070 bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1071 bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1072 bucket_mask: u32,
1073}
1074
1075pub struct HashTableU64 {
1077 pub bucket_counts: crate::memory::TrackedCudaSlice<u32>,
1078 pub bucket_offsets: crate::memory::TrackedCudaSlice<u32>,
1079 pub bucket_entries: crate::memory::TrackedCudaSlice<u32>,
1080 pub bucket_entry_hashes: crate::memory::TrackedCudaSlice<u64>,
1081 pub bucket_mask: u32,
1082}
1083
1084pub struct JoinIndexV2 {
1089 right_num_rows: u32,
1090 right_keys: Vec<usize>,
1091 key_bytes: u32,
1092 packed_keys: crate::memory::TrackedCudaSlice<u8>,
1093 table: JoinHashTableV2,
1094}
1095
1096impl JoinIndexV2 {
1097 pub fn right_keys(&self) -> &[usize] {
1099 &self.right_keys
1100 }
1101
1102 pub fn right_num_rows(&self) -> u32 {
1104 self.right_num_rows
1105 }
1106
1107 pub fn estimated_bytes(&self) -> u64 {
1109 let mut bytes = 0u64;
1110 bytes = bytes.saturating_add(self.packed_keys.len() as u64);
1111 bytes = bytes.saturating_add(self.table.bucket_counts.len() as u64 * 4);
1112 bytes = bytes.saturating_add(self.table.bucket_offsets.len() as u64 * 4);
1113 bytes = bytes.saturating_add(self.table.bucket_entries.len() as u64 * 4);
1114 bytes = bytes.saturating_add(self.table.bucket_entry_hashes.len() as u64 * 8);
1115 bytes
1116 }
1117}
1118
1119pub struct CudaKernelProvider {
1140 provider_identity: u64,
1142 device: Arc<CudaDevice>,
1144 memory: Arc<GpuMemoryManager>,
1146 transfer_tracker: HostTransferTracker,
1148 final_observation_transfer_tracker: HostTransferTracker,
1150 final_observation_pinned_receipts: AtomicU64,
1152 ptx_load_profile: Option<PtxLoadProfile>,
1154 d2h_transfer_count: AtomicU64,
1156 untracked_metadata_dtoh_count: AtomicU64,
1163 strict_deterministic_d2h: AtomicBool,
1169 deterministic_d2h_violations: AtomicU64,
1173 recorded_op_stream: OnceLock<crate::device_runtime::StreamId>,
1182 csm_invocations: AtomicU64,
1190 csm_cuda_graph_captures: AtomicU64,
1192 csm_cuda_graph_launches: AtomicU64,
1194 csm_cuda_graph_fallbacks: AtomicU64,
1196 csm_cuda_graph_cache_hits: AtomicU64,
1198 small_full_row_sort_invocations: AtomicU64,
1202 csm_cuda_graph_cache: Mutex<HashMap<CsmCudaGraphKey, CsmCudaGraphEntry>>,
1204 wcoj_layout_fast_path_hit_count: AtomicU64,
1211 wcoj_layout_sort_invocation_count: AtomicU64,
1216 kclique_metadata_build_count: AtomicU64,
1218 kclique_metadata_build_nanos: AtomicU64,
1221 wcoj_triangle_hg_dispatch_count: AtomicU64,
1224 #[cfg(feature = "wcoj-phase-timing")]
1231 last_triangle_phase_timing:
1232 std::sync::Mutex<Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming>>,
1233}
1234
1235#[derive(Default)]
1236struct HostTransferTracker {
1237 dtoh_bytes: AtomicU64,
1238 htod_bytes: AtomicU64,
1239 dtoh_calls: AtomicU64,
1240 htod_calls: AtomicU64,
1241 launch_metadata_htod_bytes: AtomicU64,
1242 launch_metadata_htod_calls: AtomicU64,
1243}
1244
1245#[derive(Debug, Clone, Copy)]
1246pub struct HostTransferStats {
1247 pub dtoh_bytes: u64,
1248 pub htod_bytes: u64,
1249 pub dtoh_calls: u64,
1250 pub htod_calls: u64,
1251}
1252
1253#[derive(Debug, Clone, Copy, Default)]
1255pub struct FinalObservationTransferStats {
1256 pub dtoh_bytes: u64,
1258 pub dtoh_calls: u64,
1260 pub pinned_receipts: u64,
1262}
1263
1264#[derive(Debug, Clone, Copy, Default)]
1265pub struct HostLaunchMetadataTransferStats {
1266 pub htod_bytes: u64,
1267 pub htod_calls: u64,
1268}
1269
1270impl HostTransferTracker {
1271 fn record_dtoh(&self, bytes: u64) {
1272 self.dtoh_calls.fetch_add(1, Ordering::Relaxed);
1273 self.dtoh_bytes.fetch_add(bytes, Ordering::Relaxed);
1274 }
1275
1276 fn record_htod(&self, bytes: u64) {
1277 self.htod_calls.fetch_add(1, Ordering::Relaxed);
1278 self.htod_bytes.fetch_add(bytes, Ordering::Relaxed);
1279 }
1280
1281 fn record_htod_launch_metadata(&self, bytes: u64) {
1282 self.launch_metadata_htod_calls
1283 .fetch_add(1, Ordering::Relaxed);
1284 self.launch_metadata_htod_bytes
1285 .fetch_add(bytes, Ordering::Relaxed);
1286 }
1287
1288 fn snapshot(&self) -> HostTransferStats {
1289 HostTransferStats {
1290 dtoh_bytes: self.dtoh_bytes.load(Ordering::Relaxed),
1291 htod_bytes: self.htod_bytes.load(Ordering::Relaxed),
1292 dtoh_calls: self.dtoh_calls.load(Ordering::Relaxed),
1293 htod_calls: self.htod_calls.load(Ordering::Relaxed),
1294 }
1295 }
1296
1297 fn launch_metadata_snapshot(&self) -> HostLaunchMetadataTransferStats {
1298 HostLaunchMetadataTransferStats {
1299 htod_bytes: self.launch_metadata_htod_bytes.load(Ordering::Relaxed),
1300 htod_calls: self.launch_metadata_htod_calls.load(Ordering::Relaxed),
1301 }
1302 }
1303
1304 fn reset(&self) {
1305 self.dtoh_bytes.store(0, Ordering::Relaxed);
1306 self.htod_bytes.store(0, Ordering::Relaxed);
1307 self.dtoh_calls.store(0, Ordering::Relaxed);
1308 self.htod_calls.store(0, Ordering::Relaxed);
1309 self.launch_metadata_htod_bytes.store(0, Ordering::Relaxed);
1310 self.launch_metadata_htod_calls.store(0, Ordering::Relaxed);
1311 }
1312}
1313
1314impl CudaKernelProvider {
1315 pub fn new(device: Arc<CudaDevice>, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1334 let profiling = warmup_profiling_enabled();
1335 let ptx_load_profile = Self::load_all_kernel_modules(&device, profiling)?;
1336
1337 Ok(Self::from_loaded_device(device, memory, ptx_load_profile))
1338 }
1339
1340 fn from_loaded_device(
1343 device: Arc<CudaDevice>,
1344 memory: Arc<GpuMemoryManager>,
1345 ptx_load_profile: Option<PtxLoadProfile>,
1346 ) -> Self {
1347 Self {
1348 provider_identity: NEXT_PROVIDER_IDENTITY.fetch_add(1, Ordering::Relaxed),
1349 device,
1350 memory,
1351 transfer_tracker: HostTransferTracker::default(),
1352 final_observation_transfer_tracker: HostTransferTracker::default(),
1353 final_observation_pinned_receipts: AtomicU64::new(0),
1354 ptx_load_profile,
1355 d2h_transfer_count: AtomicU64::new(0),
1356 untracked_metadata_dtoh_count: AtomicU64::new(0),
1357 strict_deterministic_d2h: AtomicBool::new(false),
1358 deterministic_d2h_violations: AtomicU64::new(0),
1359 recorded_op_stream: OnceLock::new(),
1360 csm_invocations: AtomicU64::new(0),
1361 csm_cuda_graph_captures: AtomicU64::new(0),
1362 csm_cuda_graph_launches: AtomicU64::new(0),
1363 csm_cuda_graph_fallbacks: AtomicU64::new(0),
1364 csm_cuda_graph_cache_hits: AtomicU64::new(0),
1365 small_full_row_sort_invocations: AtomicU64::new(0),
1366 csm_cuda_graph_cache: Mutex::new(HashMap::new()),
1367 wcoj_layout_fast_path_hit_count: AtomicU64::new(0),
1368 wcoj_layout_sort_invocation_count: AtomicU64::new(0),
1369 kclique_metadata_build_count: AtomicU64::new(0),
1370 kclique_metadata_build_nanos: AtomicU64::new(0),
1371 wcoj_triangle_hg_dispatch_count: AtomicU64::new(0),
1372 #[cfg(feature = "wcoj-phase-timing")]
1373 last_triangle_phase_timing: std::sync::Mutex::new(None),
1374 }
1375 }
1376
1377 pub fn with_runtime(device: Arc<CudaDevice>, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1425 if memory.runtime().is_none() {
1426 return Err(XlogError::Kernel(
1427 "CudaKernelProvider::with_runtime requires a GpuMemoryManager built via \
1428 GpuMemoryManager::with_runtime; got a manager with no runtime attached"
1429 .to_string(),
1430 ));
1431 }
1432 Self::new(device, memory)
1433 }
1434
1435 pub fn with_runtime_memory_view(&self, memory: Arc<GpuMemoryManager>) -> Result<Self> {
1443 let runtime = memory.runtime().ok_or_else(|| {
1444 XlogError::Kernel(
1445 "CudaKernelProvider::with_runtime_memory_view requires a GpuMemoryManager with an attached runtime"
1446 .to_string(),
1447 )
1448 })?;
1449 if !Arc::ptr_eq(&self.device, memory.device()) {
1450 return Err(XlogError::Kernel(
1451 "CudaKernelProvider::with_runtime_memory_view requires the memory manager to share the provider's exact CUDA device handle"
1452 .to_string(),
1453 ));
1454 }
1455 if !Arc::ptr_eq(&self.device, runtime.device()) {
1456 return Err(XlogError::Kernel(
1457 "CudaKernelProvider::with_runtime_memory_view requires the runtime to share the provider's exact CUDA device handle"
1458 .to_string(),
1459 ));
1460 }
1461 let device_ordinal = u32::try_from(self.device.ordinal()).map_err(|_| {
1462 XlogError::Kernel(format!(
1463 "CUDA device ordinal {} is not representable as u32",
1464 self.device.ordinal()
1465 ))
1466 })?;
1467 if runtime.device_ordinal() != device_ordinal {
1468 return Err(XlogError::Kernel(format!(
1469 "CudaKernelProvider::with_runtime_memory_view device ordinal mismatch: provider={} runtime={}",
1470 device_ordinal,
1471 runtime.device_ordinal()
1472 )));
1473 }
1474 if !runtime.supports_block_use_tracking() {
1475 return Err(XlogError::Kernel(
1476 "CudaKernelProvider::with_runtime_memory_view requires a runtime with cross-stream block-use tracking"
1477 .to_string(),
1478 ));
1479 }
1480
1481 Ok(Self::from_loaded_device(
1482 Arc::clone(&self.device),
1483 memory,
1484 None,
1485 ))
1486 }
1487
1488 fn env_flag(name: &str) -> bool {
1491 std::env::var(name)
1492 .map(|v| !v.is_empty() && v != "0")
1493 .unwrap_or(false)
1494 }
1495
1496 pub(crate) fn use_recorded_filters_env() -> bool {
1508 Self::env_flag("XLOG_USE_RECORDED_FILTERS") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1509 }
1510
1511 pub(crate) fn use_recorded_sort_env() -> bool {
1518 Self::env_flag("XLOG_USE_RECORDED_SORT") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1519 }
1520
1521 pub(crate) fn use_recorded_dedup_env() -> bool {
1526 Self::env_flag("XLOG_USE_RECORDED_DEDUP") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1527 }
1528
1529 pub(crate) fn use_recorded_groupby_env() -> bool {
1535 Self::env_flag("XLOG_USE_RECORDED_GROUPBY") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1536 }
1537
1538 pub(crate) fn use_recorded_hash_join_env() -> bool {
1546 Self::env_flag("XLOG_USE_RECORDED_HASH_JOIN") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1547 }
1548
1549 pub(crate) fn use_recorded_csm_env() -> bool {
1557 Self::env_flag("XLOG_USE_RECORDED_CSM") || Self::env_flag("XLOG_USE_RECORDED_OPS")
1558 }
1559
1560 pub(crate) fn use_csm_cuda_graph_env() -> bool {
1566 Self::env_flag("XLOG_USE_CSM_CUDA_GRAPH") || Self::env_flag("XLOG_USE_CUDA_GRAPHS")
1567 }
1568
1569 #[doc(hidden)]
1583 pub fn csm_invocations(&self) -> u64 {
1584 self.csm_invocations.load(Ordering::Relaxed)
1585 }
1586
1587 #[doc(hidden)]
1588 pub fn csm_cuda_graph_captures(&self) -> u64 {
1589 self.csm_cuda_graph_captures.load(Ordering::Relaxed)
1590 }
1591
1592 #[doc(hidden)]
1593 pub fn csm_cuda_graph_launches(&self) -> u64 {
1594 self.csm_cuda_graph_launches.load(Ordering::Relaxed)
1595 }
1596
1597 #[doc(hidden)]
1598 pub fn csm_cuda_graph_fallbacks(&self) -> u64 {
1599 self.csm_cuda_graph_fallbacks.load(Ordering::Relaxed)
1600 }
1601
1602 #[doc(hidden)]
1603 pub fn csm_cuda_graph_cache_hits(&self) -> u64 {
1604 self.csm_cuda_graph_cache_hits.load(Ordering::Relaxed)
1605 }
1606
1607 #[doc(hidden)]
1608 pub fn small_full_row_sort_invocations(&self) -> u64 {
1609 self.small_full_row_sort_invocations.load(Ordering::Relaxed)
1610 }
1611
1612 pub(crate) fn recorded_op_stream_or_init(&self) -> Option<crate::device_runtime::StreamId> {
1632 if let Some(s) = self.recorded_op_stream.get() {
1633 return Some(*s);
1634 }
1635 let runtime = self.memory.runtime()?;
1636 let stream = runtime.stream_pool().acquire().ok()?;
1637 let _ = self.recorded_op_stream.set(stream);
1638 self.recorded_op_stream.get().copied()
1639 }
1640
1641 #[cfg(feature = "wcoj-phase-timing")]
1651 pub fn take_wcoj_triangle_phase_timing(
1652 &self,
1653 ) -> Option<crate::wcoj_phase_timing::WcojTrianglePhaseTiming> {
1654 self.last_triangle_phase_timing
1655 .lock()
1656 .ok()
1657 .and_then(|mut g| g.take())
1658 }
1659
1660 #[cfg(feature = "wcoj-phase-timing")]
1664 #[allow(dead_code)]
1665 pub(crate) fn put_wcoj_triangle_phase_timing(
1666 &self,
1667 timing: crate::wcoj_phase_timing::WcojTrianglePhaseTiming,
1668 ) {
1669 if let Ok(mut g) = self.last_triangle_phase_timing.lock() {
1670 *g = Some(timing);
1671 }
1672 }
1673
1674 pub fn wcoj_layout_fast_path_hit_count(&self) -> u64 {
1681 self.wcoj_layout_fast_path_hit_count.load(Ordering::Relaxed)
1682 }
1683
1684 pub fn wcoj_triangle_hg_dispatch_count(&self) -> u64 {
1687 self.wcoj_triangle_hg_dispatch_count.load(Ordering::Relaxed)
1688 }
1689
1690 pub fn reset_wcoj_layout_fast_path_hit_count(&self) {
1693 self.wcoj_layout_fast_path_hit_count
1694 .store(0, Ordering::Relaxed);
1695 }
1696
1697 pub fn wcoj_layout_sort_invocation_count(&self) -> u64 {
1700 self.wcoj_layout_sort_invocation_count
1701 .load(Ordering::Relaxed)
1702 }
1703
1704 pub fn reset_wcoj_layout_sort_invocation_count(&self) {
1706 self.wcoj_layout_sort_invocation_count
1707 .store(0, Ordering::Relaxed);
1708 }
1709
1710 pub fn kclique_metadata_build_count(&self) -> u64 {
1713 self.kclique_metadata_build_count.load(Ordering::Relaxed)
1714 }
1715
1716 pub fn kclique_metadata_build_nanos(&self) -> u64 {
1719 self.kclique_metadata_build_nanos.load(Ordering::Relaxed)
1720 }
1721
1722 pub fn reset_kclique_metadata_build_metrics(&self) {
1724 self.kclique_metadata_build_count
1725 .store(0, Ordering::Relaxed);
1726 self.kclique_metadata_build_nanos
1727 .store(0, Ordering::Relaxed);
1728 }
1729
1730 pub(crate) fn record_wcoj_layout_fast_path_hit(&self) {
1734 self.wcoj_layout_fast_path_hit_count
1735 .fetch_add(1, Ordering::Relaxed);
1736 }
1737
1738 pub(crate) fn record_wcoj_layout_sort_invocation(&self) {
1740 self.wcoj_layout_sort_invocation_count
1741 .fetch_add(1, Ordering::Relaxed);
1742 }
1743
1744 pub(crate) fn record_kclique_metadata_build_nanos(&self, nanos: u128) {
1746 self.kclique_metadata_build_count
1747 .fetch_add(1, Ordering::Relaxed);
1748 let nanos = u64::try_from(nanos).unwrap_or(u64::MAX);
1749 self.kclique_metadata_build_nanos
1750 .fetch_add(nanos, Ordering::Relaxed);
1751 }
1752
1753 #[doc(hidden)]
1756 pub fn record_wcoj_triangle_hg_dispatch(&self) {
1757 self.wcoj_triangle_hg_dispatch_count
1758 .fetch_add(1, Ordering::Relaxed);
1759 }
1760
1761 pub fn device(&self) -> &Arc<CudaDevice> {
1763 &self.device
1764 }
1765
1766 pub(crate) fn provider_identity(&self) -> u64 {
1767 self.provider_identity
1768 }
1769
1770 pub fn memory(&self) -> &Arc<GpuMemoryManager> {
1772 &self.memory
1773 }
1774
1775 pub fn ptx_load_profile(&self) -> Option<&PtxLoadProfile> {
1777 self.ptx_load_profile.as_ref()
1778 }
1779
1780 pub fn reset_host_transfer_stats(&self) {
1782 self.transfer_tracker.reset();
1783 }
1784
1785 pub fn host_transfer_stats(&self) -> HostTransferStats {
1787 self.transfer_tracker.snapshot()
1788 }
1789
1790 pub fn reset_final_observation_transfer_stats(&self) {
1792 self.final_observation_transfer_tracker.reset();
1793 self.final_observation_pinned_receipts
1794 .store(0, Ordering::Relaxed);
1795 }
1796
1797 pub fn final_observation_transfer_stats(&self) -> FinalObservationTransferStats {
1799 let transfers = self.final_observation_transfer_tracker.snapshot();
1800 FinalObservationTransferStats {
1801 dtoh_bytes: transfers.dtoh_bytes,
1802 dtoh_calls: transfers.dtoh_calls,
1803 pinned_receipts: self
1804 .final_observation_pinned_receipts
1805 .load(Ordering::Relaxed),
1806 }
1807 }
1808
1809 #[doc(hidden)]
1815 pub fn record_final_observation_transfer(&self, bytes: u64) {
1816 self.final_observation_transfer_tracker.record_dtoh(bytes);
1817 self.final_observation_pinned_receipts
1818 .fetch_add(1, Ordering::Relaxed);
1819 }
1820
1821 #[doc(hidden)]
1828 pub fn finalize_resident_logical_counts(&self, entries: &[(&CudaBuffer, u32)]) -> Result<()> {
1829 for (buffer, count) in entries {
1830 if u64::from(*count) > buffer.num_rows() {
1831 return Err(XlogError::Kernel(format!(
1832 "resident receipt logical row count {} exceeds buffer capacity {}",
1833 count,
1834 buffer.num_rows()
1835 )));
1836 }
1837 if let Some(cached) = buffer.cached_row_count() {
1838 if cached != *count {
1839 return Err(XlogError::Kernel(format!(
1840 "resident receipt logical row count {} conflicts with cached count {}",
1841 count, cached
1842 )));
1843 }
1844 }
1845 }
1846 for (buffer, count) in entries {
1847 buffer.set_cached_row_count_if_unset(*count);
1848 }
1849 Ok(())
1850 }
1851
1852 pub fn host_launch_metadata_transfer_stats(&self) -> HostLaunchMetadataTransferStats {
1855 self.transfer_tracker.launch_metadata_snapshot()
1856 }
1857
1858 pub fn d2h_transfer_count(&self) -> u64 {
1864 self.d2h_transfer_count.load(Ordering::Relaxed)
1865 }
1866
1867 pub fn reset_d2h_transfer_count(&self) {
1869 self.d2h_transfer_count.store(0, Ordering::Relaxed);
1870 }
1871
1872 pub fn untracked_metadata_dtoh_count(&self) -> u64 {
1875 self.untracked_metadata_dtoh_count.load(Ordering::Relaxed)
1876 }
1877
1878 pub fn reset_untracked_metadata_dtoh_count(&self) {
1880 self.untracked_metadata_dtoh_count
1881 .store(0, Ordering::Relaxed);
1882 }
1883
1884 pub fn enable_strict_deterministic_d2h(&self) {
1901 self.strict_deterministic_d2h.store(true, Ordering::Relaxed);
1902 }
1903
1904 pub fn disable_strict_deterministic_d2h(&self) {
1906 self.strict_deterministic_d2h
1907 .store(false, Ordering::Relaxed);
1908 }
1909
1910 pub fn strict_deterministic_d2h_enabled(&self) -> bool {
1912 self.strict_deterministic_d2h.load(Ordering::Relaxed)
1913 }
1914
1915 pub fn deterministic_d2h_violation_count(&self) -> u64 {
1917 self.deterministic_d2h_violations.load(Ordering::Relaxed)
1918 }
1919
1920 pub fn reset_deterministic_d2h_violations(&self) {
1922 self.deterministic_d2h_violations
1923 .store(0, Ordering::Relaxed);
1924 }
1925
1926 pub(crate) fn check_deterministic_d2h(&self, op: &'static str, bytes: u64) -> Result<()> {
1932 if self.strict_deterministic_d2h.load(Ordering::Relaxed) {
1933 self.deterministic_d2h_violations
1934 .fetch_add(1, Ordering::Relaxed);
1935 return Err(XlogError::Execution(format!(
1936 "deterministic D2H gate: {} attempted to copy {} bytes from device to host",
1937 op, bytes
1938 )));
1939 }
1940 Ok(())
1941 }
1942
1943 fn dtoh_sync_copy_into_tracked<T: DeviceRepr, Src: DevicePtr<T>>(
1944 &self,
1945 src: &Src,
1946 dst: &mut [T],
1947 ) -> Result<()> {
1948 let bytes = std::mem::size_of::<T>()
1949 .checked_mul(dst.len())
1950 .ok_or_else(|| XlogError::Kernel("dtoh size overflow".to_string()))?;
1951 self.check_deterministic_d2h("dtoh_sync_copy_into_tracked", bytes as u64)?;
1952 self.transfer_tracker.record_dtoh(bytes as u64);
1953 self.device
1954 .inner()
1955 .dtoh_sync_copy_into(src, dst)
1956 .map_err(|e| XlogError::Kernel(format!("Failed to copy from device: {}", e)))
1957 }
1958
1959 pub const DTOH_SMALL_METADATA_MAX_BYTES: usize = 4096;
1964
1965 pub fn dtoh_small_metadata_untracked<T: DeviceRepr + Default + Copy>(
1990 &self,
1991 src: &crate::memory::TrackedCudaSlice<T>,
1992 count: usize,
1993 ) -> Result<Vec<T>> {
1994 let bytes = count.checked_mul(std::mem::size_of::<T>()).ok_or_else(|| {
1995 XlogError::Kernel("dtoh_small_metadata_untracked: byte size overflow".to_string())
1996 })?;
1997 if bytes > Self::DTOH_SMALL_METADATA_MAX_BYTES {
1998 return Err(XlogError::Kernel(format!(
1999 "dtoh_small_metadata_untracked: requested {} bytes exceeds metadata cap of {} bytes \
2000 (this is metadata-only; use download_column* for data-plane transfers)",
2001 bytes,
2002 Self::DTOH_SMALL_METADATA_MAX_BYTES
2003 )));
2004 }
2005 if count > src.len() {
2006 return Err(XlogError::Kernel(format!(
2007 "dtoh_small_metadata_untracked: count={count} > src.len={}",
2008 src.len()
2009 )));
2010 }
2011 if count == 0 {
2012 return Ok(Vec::new());
2013 }
2014 let slice = src.try_slice(0..count).ok_or_else(|| {
2015 XlogError::Kernel(format!(
2016 "dtoh_small_metadata_untracked: try_slice(0..{count}) failed"
2017 ))
2018 })?;
2019 let mut buf: Vec<T> = vec![T::default(); count];
2020 self.untracked_metadata_dtoh_count
2021 .fetch_add(1, Ordering::Relaxed);
2022 self.device
2023 .inner()
2024 .dtoh_sync_copy_into(&slice, &mut buf)
2025 .map_err(|e| {
2026 XlogError::Kernel(format!("dtoh_small_metadata_untracked: copy failed: {}", e))
2027 })?;
2028 Ok(buf)
2029 }
2030
2031 pub fn dtoh_scalar_untracked<T: DeviceRepr + Default + Copy>(
2039 &self,
2040 src: &crate::memory::TrackedCudaSlice<T>,
2041 index: usize,
2042 ) -> Result<T> {
2043 if index >= src.len() {
2044 return Err(XlogError::Kernel(format!(
2045 "dtoh_scalar_untracked: index={} >= len={}",
2046 index,
2047 src.len()
2048 )));
2049 }
2050 let slice = src.try_slice(index..index + 1).ok_or_else(|| {
2051 XlogError::Kernel(format!(
2052 "dtoh_scalar_untracked: slice failed at index={}",
2053 index
2054 ))
2055 })?;
2056 let mut buf = [T::default()];
2057 self.untracked_metadata_dtoh_count
2058 .fetch_add(1, Ordering::Relaxed);
2059 self.device
2060 .inner()
2061 .dtoh_sync_copy_into(&slice, &mut buf)
2062 .map_err(|e| XlogError::Kernel(format!("dtoh_scalar_untracked: copy failed: {}", e)))?;
2063 Ok(buf[0])
2064 }
2065
2066 pub fn htod_sync_copy_into_tracked<T: DeviceRepr, Dst: cudarc::driver::DevicePtrMut<T>>(
2068 &self,
2069 src: &[T],
2070 dst: &mut Dst,
2071 ) -> Result<()> {
2072 let bytes = std::mem::size_of::<T>()
2073 .checked_mul(src.len())
2074 .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
2075 self.transfer_tracker.record_htod(bytes as u64);
2076 self.device
2077 .inner()
2078 .htod_sync_copy_into(src, dst)
2079 .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
2080 }
2081
2082 pub fn htod_sync_copy_tracked<T: DeviceRepr>(
2085 &self,
2086 src: &[T],
2087 ) -> Result<cudarc::driver::CudaSlice<T>> {
2088 let bytes = std::mem::size_of::<T>()
2089 .checked_mul(src.len())
2090 .ok_or_else(|| XlogError::Kernel("htod size overflow".to_string()))?;
2091 self.transfer_tracker.record_htod(bytes as u64);
2092 self.device
2093 .inner()
2094 .htod_sync_copy(src)
2095 .map_err(|e| XlogError::Kernel(format!("Failed to copy to device: {}", e)))
2096 }
2097
2098 pub fn htod_launch_metadata_sync_copy_into<
2101 T: DeviceRepr,
2102 Dst: cudarc::driver::DevicePtrMut<T>,
2103 >(
2104 &self,
2105 src: &[T],
2106 dst: &mut Dst,
2107 ) -> Result<()> {
2108 let bytes = std::mem::size_of::<T>()
2109 .checked_mul(src.len())
2110 .ok_or_else(|| XlogError::Kernel("launch metadata htod size overflow".to_string()))?;
2111 self.transfer_tracker
2112 .record_htod_launch_metadata(bytes as u64);
2113 self.device
2114 .inner()
2115 .htod_sync_copy_into(src, dst)
2116 .map_err(|e| {
2117 XlogError::Kernel(format!("Failed to copy launch metadata to device: {}", e))
2118 })
2119 }
2120
2121 pub(crate) fn htod_launch_metadata_async_copy_one<T: DeviceRepr>(
2124 &self,
2125 src: &T,
2126 dst: &TrackedCudaSlice<T>,
2127 stream: &CudaStream,
2128 context: &str,
2129 ) -> Result<()> {
2130 let bytes = std::mem::size_of::<T>();
2131 self.transfer_tracker
2132 .record_htod_launch_metadata(bytes as u64);
2133 unsafe {
2134 let res = cudarc::driver::sys::cuMemcpyHtoDAsync_v2(
2135 *dst.device_ptr(),
2136 src as *const T as *const c_void,
2137 bytes,
2138 stream.cu_stream(),
2139 );
2140 if res != cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS {
2141 return Err(XlogError::Kernel(format!(
2142 "{context}: launch metadata H2D failed: {res:?}"
2143 )));
2144 }
2145 }
2146 Ok(())
2147 }
2148
2149 pub fn exclusive_scan_u32_inplace(
2178 &self,
2179 data: &mut crate::memory::TrackedCudaSlice<u32>,
2180 n: u32,
2181 ) -> Result<()> {
2182 if n as usize > data.len() {
2183 return Err(XlogError::Kernel(format!(
2184 "exclusive_scan_u32_inplace: n={} exceeds slice len={}",
2185 n,
2186 data.len()
2187 )));
2188 }
2189 self.multiblock_scan_u32_inplace(data, n)
2190 }
2191
2192 fn multiblock_scan_u32_inplace(
2193 &self,
2194 data: &mut crate::memory::TrackedCudaSlice<u32>,
2195 n: u32,
2196 ) -> Result<()> {
2197 if n == 0 {
2198 return Ok(());
2199 }
2200
2201 let device = self.device.inner();
2202 let block_size = 256u32;
2203
2204 if n <= block_size {
2205 let phase2_fn = device
2206 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2207 .ok_or_else(|| {
2208 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2209 })?;
2210
2211 unsafe {
2213 phase2_fn.clone().launch(
2214 LaunchConfig {
2215 grid_dim: (1, 1, 1),
2216 block_dim: (block_size, 1, 1),
2217 shared_mem_bytes: 0,
2218 },
2219 (&mut *data, n),
2220 )
2221 }
2222 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2223
2224 self.device.synchronize()?;
2225 return Ok(());
2226 }
2227
2228 let num_blocks = n.div_ceil(block_size);
2229 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2230
2231 let phase1_u32_fn = device
2232 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2233 .ok_or_else(|| {
2234 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2235 })?;
2236
2237 unsafe {
2239 phase1_u32_fn.clone().launch(
2240 LaunchConfig {
2241 grid_dim: (num_blocks, 1, 1),
2242 block_dim: (block_size, 1, 1),
2243 shared_mem_bytes: 0,
2244 },
2245 (&mut *data, &mut block_sums, n),
2246 )
2247 }
2248 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2249 self.device.synchronize()?;
2250
2251 if num_blocks > 1 {
2252 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2253 }
2254
2255 let phase3_fn = device
2256 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2257 .ok_or_else(|| {
2258 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2259 })?;
2260
2261 unsafe {
2263 phase3_fn.clone().launch(
2264 LaunchConfig {
2265 grid_dim: (num_blocks, 1, 1),
2266 block_dim: (block_size, 1, 1),
2267 shared_mem_bytes: 0,
2268 },
2269 (&mut *data, &block_sums, n),
2270 )
2271 }
2272 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2273
2274 self.device.synchronize()?;
2275 Ok(())
2276 }
2277
2278 pub(crate) fn multiblock_scan_u32_inplace_on_stream(
2293 &self,
2294 data: &mut crate::memory::TrackedCudaSlice<u32>,
2295 n: u32,
2296 cu_stream: &cudarc::driver::CudaStream,
2297 launch_stream: crate::device_runtime::StreamId,
2298 runtime: &crate::device_runtime::XlogDeviceRuntime,
2299 ) -> Result<()> {
2300 if n == 0 {
2301 return Ok(());
2302 }
2303 let device = self.device.inner();
2304 let block_size = 256u32;
2305
2306 if n <= block_size {
2307 let phase2_fn = device
2308 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2309 .ok_or_else(|| {
2310 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2311 })?;
2312 unsafe {
2314 phase2_fn.clone().launch_on_stream(
2315 cu_stream,
2316 LaunchConfig {
2317 grid_dim: (1, 1, 1),
2318 block_dim: (block_size, 1, 1),
2319 shared_mem_bytes: 0,
2320 },
2321 (&mut *data, n),
2322 )
2323 }
2324 .map_err(|e| {
2325 XlogError::Kernel(format!("multiblock_scan_phase2 (on_stream) failed: {}", e))
2326 })?;
2327 return Ok(());
2328 }
2329
2330 let num_blocks = n.div_ceil(block_size);
2331 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2332 runtime
2339 .prepare_first_use(
2340 &block_sums,
2341 launch_stream,
2342 crate::device_runtime::Access::Write,
2343 )
2344 .map_err(|e| {
2345 XlogError::Kernel(format!(
2346 "multiblock_scan_u32_inplace_on_stream: prepare block_sums failed: {}",
2347 e
2348 ))
2349 })?;
2350
2351 let phase1_u32_fn = device
2352 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2353 .ok_or_else(|| {
2354 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2355 })?;
2356 unsafe {
2358 phase1_u32_fn.clone().launch_on_stream(
2359 cu_stream,
2360 LaunchConfig {
2361 grid_dim: (num_blocks, 1, 1),
2362 block_dim: (block_size, 1, 1),
2363 shared_mem_bytes: 0,
2364 },
2365 (&mut *data, &mut block_sums, n),
2366 )
2367 }
2368 .map_err(|e| {
2369 XlogError::Kernel(format!(
2370 "multiblock_scan_u32_phase1 (on_stream) failed: {}",
2371 e
2372 ))
2373 })?;
2374
2375 if num_blocks > 1 {
2376 self.multiblock_scan_u32_inplace_on_stream(
2377 &mut block_sums,
2378 num_blocks,
2379 cu_stream,
2380 launch_stream,
2381 runtime,
2382 )?;
2383 }
2384
2385 let phase3_fn = device
2386 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2387 .ok_or_else(|| {
2388 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2389 })?;
2390 unsafe {
2392 phase3_fn.clone().launch_on_stream(
2393 cu_stream,
2394 LaunchConfig {
2395 grid_dim: (num_blocks, 1, 1),
2396 block_dim: (block_size, 1, 1),
2397 shared_mem_bytes: 0,
2398 },
2399 (&mut *data, &block_sums, n),
2400 )
2401 }
2402 .map_err(|e| {
2403 XlogError::Kernel(format!("multiblock_scan_phase3 (on_stream) failed: {}", e))
2404 })?;
2405
2406 if let Some(b) = block_sums.runtime_block() {
2412 runtime
2413 .finish_block_use(
2414 crate::device_runtime::BlockId::from_block(b),
2415 launch_stream,
2416 crate::device_runtime::Access::Write,
2417 )
2418 .map_err(|e| {
2419 XlogError::Kernel(format!(
2420 "multiblock_scan_u32_inplace_on_stream: finish_block_use \
2421 for intermediate block_sums failed: {}",
2422 e
2423 ))
2424 })?;
2425 } else {
2426 return Err(XlogError::Kernel(
2427 "multiblock_scan_u32_inplace_on_stream: intermediate block_sums has no \
2428 runtime block — caller must use a runtime-backed manager"
2429 .to_string(),
2430 ));
2431 }
2432 Ok(())
2433 }
2434
2435 pub(crate) fn multiblock_scan_u32_scratch_for_len(
2438 &self,
2439 mut n: u32,
2440 ) -> Result<MultiblockScanScratchU32> {
2441 let block_size = 256u32;
2442 let mut levels = Vec::new();
2443 while n > block_size {
2444 let num_blocks = n.div_ceil(block_size);
2445 levels.push(self.memory.alloc::<u32>(num_blocks as usize)?);
2446 n = num_blocks;
2447 }
2448 Ok(MultiblockScanScratchU32 { levels })
2449 }
2450
2451 pub(crate) fn multiblock_scan_u32_inplace_on_stream_with_scratch(
2458 &self,
2459 data: &mut crate::memory::TrackedCudaSlice<u32>,
2460 n: u32,
2461 cu_stream: &cudarc::driver::CudaStream,
2462 scratch: &mut MultiblockScanScratchU32,
2463 ) -> Result<()> {
2464 self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2465 data,
2466 n,
2467 cu_stream,
2468 &mut scratch.levels,
2469 )
2470 }
2471
2472 fn multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2473 &self,
2474 data: &mut crate::memory::TrackedCudaSlice<u32>,
2475 n: u32,
2476 cu_stream: &cudarc::driver::CudaStream,
2477 scratch_levels: &mut [TrackedCudaSlice<u32>],
2478 ) -> Result<()> {
2479 if n == 0 {
2480 return Ok(());
2481 }
2482 let device = self.device.inner();
2483 let block_size = 256u32;
2484
2485 if n <= block_size {
2486 let phase2_fn = device
2487 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2488 .ok_or_else(|| {
2489 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2490 })?;
2491 unsafe {
2493 phase2_fn.clone().launch_on_stream(
2494 cu_stream,
2495 LaunchConfig {
2496 grid_dim: (1, 1, 1),
2497 block_dim: (block_size, 1, 1),
2498 shared_mem_bytes: 0,
2499 },
2500 (&mut *data, n),
2501 )
2502 }
2503 .map_err(|e| {
2504 XlogError::Kernel(format!(
2505 "multiblock_scan_phase2 (graph scratch) failed: {}",
2506 e
2507 ))
2508 })?;
2509 return Ok(());
2510 }
2511
2512 let num_blocks = n.div_ceil(block_size);
2513 let (block_sums, rest) = scratch_levels.split_first_mut().ok_or_else(|| {
2514 XlogError::Kernel(format!(
2515 "multiblock_scan_u32_inplace_on_stream_with_scratch: missing scratch level \
2516 for n={n}, num_blocks={num_blocks}"
2517 ))
2518 })?;
2519 if block_sums.len() < num_blocks as usize {
2520 return Err(XlogError::Kernel(format!(
2521 "multiblock_scan_u32_inplace_on_stream_with_scratch: scratch level too small \
2522 (have {}, need {})",
2523 block_sums.len(),
2524 num_blocks
2525 )));
2526 }
2527
2528 let phase1_u32_fn = device
2529 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2530 .ok_or_else(|| {
2531 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2532 })?;
2533 unsafe {
2535 phase1_u32_fn.clone().launch_on_stream(
2536 cu_stream,
2537 LaunchConfig {
2538 grid_dim: (num_blocks, 1, 1),
2539 block_dim: (block_size, 1, 1),
2540 shared_mem_bytes: 0,
2541 },
2542 (&mut *data, &mut *block_sums, n),
2543 )
2544 }
2545 .map_err(|e| {
2546 XlogError::Kernel(format!(
2547 "multiblock_scan_u32_phase1 (graph scratch) failed: {}",
2548 e
2549 ))
2550 })?;
2551
2552 if num_blocks > 1 {
2553 self.multiblock_scan_u32_inplace_on_stream_with_scratch_levels(
2554 block_sums, num_blocks, cu_stream, rest,
2555 )?;
2556 }
2557
2558 let phase3_fn = device
2559 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2560 .ok_or_else(|| {
2561 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2562 })?;
2563 unsafe {
2565 phase3_fn.clone().launch_on_stream(
2566 cu_stream,
2567 LaunchConfig {
2568 grid_dim: (num_blocks, 1, 1),
2569 block_dim: (block_size, 1, 1),
2570 shared_mem_bytes: 0,
2571 },
2572 (&mut *data, &*block_sums, n),
2573 )
2574 }
2575 .map_err(|e| {
2576 XlogError::Kernel(format!(
2577 "multiblock_scan_phase3 (graph scratch) failed: {}",
2578 e
2579 ))
2580 })?;
2581 Ok(())
2582 }
2583
2584 pub(crate) fn multiblock_scan_u32_view_inplace_on_stream(
2592 &self,
2593 data: &mut CudaViewMut<'_, u32>,
2594 n: u32,
2595 cu_stream: &cudarc::driver::CudaStream,
2596 launch_stream: crate::device_runtime::StreamId,
2597 runtime: &crate::device_runtime::XlogDeviceRuntime,
2598 ) -> Result<()> {
2599 if n == 0 {
2600 return Ok(());
2601 }
2602 let device = self.device.inner();
2603 let block_size = 256u32;
2604
2605 if n <= block_size {
2606 let phase2_fn = device
2607 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2608 .ok_or_else(|| {
2609 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2610 })?;
2611 unsafe {
2613 phase2_fn.clone().launch_on_stream(
2614 cu_stream,
2615 LaunchConfig {
2616 grid_dim: (1, 1, 1),
2617 block_dim: (block_size, 1, 1),
2618 shared_mem_bytes: 0,
2619 },
2620 (data, n),
2621 )
2622 }
2623 .map_err(|e| {
2624 XlogError::Kernel(format!(
2625 "multiblock_scan_phase2 (view on_stream) failed: {}",
2626 e
2627 ))
2628 })?;
2629 return Ok(());
2630 }
2631
2632 let num_blocks = n.div_ceil(block_size);
2633 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2634 runtime
2638 .prepare_first_use(
2639 &block_sums,
2640 launch_stream,
2641 crate::device_runtime::Access::Write,
2642 )
2643 .map_err(|e| {
2644 XlogError::Kernel(format!(
2645 "multiblock_scan_u32_view_inplace_on_stream: prepare block_sums failed: {}",
2646 e
2647 ))
2648 })?;
2649
2650 let phase1_u32_fn = device
2651 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2652 .ok_or_else(|| {
2653 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2654 })?;
2655 unsafe {
2657 phase1_u32_fn.clone().launch_on_stream(
2658 cu_stream,
2659 LaunchConfig {
2660 grid_dim: (num_blocks, 1, 1),
2661 block_dim: (block_size, 1, 1),
2662 shared_mem_bytes: 0,
2663 },
2664 (&mut *data, &mut block_sums, n),
2665 )
2666 }
2667 .map_err(|e| {
2668 XlogError::Kernel(format!(
2669 "multiblock_scan_u32_phase1 (view on_stream) failed: {}",
2670 e
2671 ))
2672 })?;
2673
2674 if num_blocks > 1 {
2675 self.multiblock_scan_u32_inplace_on_stream(
2676 &mut block_sums,
2677 num_blocks,
2678 cu_stream,
2679 launch_stream,
2680 runtime,
2681 )?;
2682 }
2683
2684 let phase3_fn = device
2685 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2686 .ok_or_else(|| {
2687 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2688 })?;
2689 unsafe {
2691 phase3_fn.clone().launch_on_stream(
2692 cu_stream,
2693 LaunchConfig {
2694 grid_dim: (num_blocks, 1, 1),
2695 block_dim: (block_size, 1, 1),
2696 shared_mem_bytes: 0,
2697 },
2698 (&mut *data, &block_sums, n),
2699 )
2700 }
2701 .map_err(|e| {
2702 XlogError::Kernel(format!(
2703 "multiblock_scan_phase3 (view on_stream) failed: {}",
2704 e
2705 ))
2706 })?;
2707
2708 if let Some(b) = block_sums.runtime_block() {
2710 runtime
2711 .finish_block_use(
2712 crate::device_runtime::BlockId::from_block(b),
2713 launch_stream,
2714 crate::device_runtime::Access::Write,
2715 )
2716 .map_err(|e| {
2717 XlogError::Kernel(format!(
2718 "multiblock_scan_u32_view_inplace_on_stream: finish_block_use \
2719 for intermediate block_sums failed: {}",
2720 e
2721 ))
2722 })?;
2723 } else {
2724 return Err(XlogError::Kernel(
2725 "multiblock_scan_u32_view_inplace_on_stream: intermediate block_sums has no \
2726 runtime block — caller must use a runtime-backed manager"
2727 .to_string(),
2728 ));
2729 }
2730 Ok(())
2731 }
2732
2733 fn multiblock_scan_u32_view_inplace(
2734 &self,
2735 data: &mut CudaViewMut<'_, u32>,
2736 n: u32,
2737 ) -> Result<()> {
2738 if n == 0 {
2739 return Ok(());
2740 }
2741
2742 let device = self.device.inner();
2743 let block_size = 256u32;
2744
2745 if n <= block_size {
2746 let phase2_fn = device
2747 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE2)
2748 .ok_or_else(|| {
2749 XlogError::Kernel("Failed to get multiblock_scan_phase2 kernel".to_string())
2750 })?;
2751
2752 unsafe {
2754 phase2_fn.clone().launch(
2755 LaunchConfig {
2756 grid_dim: (1, 1, 1),
2757 block_dim: (block_size, 1, 1),
2758 shared_mem_bytes: 0,
2759 },
2760 (data, n),
2761 )
2762 }
2763 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase2 failed: {}", e)))?;
2764
2765 self.device.synchronize()?;
2766 return Ok(());
2767 }
2768
2769 let num_blocks = n.div_ceil(block_size);
2770 let mut block_sums = self.memory.alloc::<u32>(num_blocks as usize)?;
2771
2772 let phase1_u32_fn = device
2773 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_U32_PHASE1)
2774 .ok_or_else(|| {
2775 XlogError::Kernel("Failed to get multiblock_scan_u32_phase1 kernel".to_string())
2776 })?;
2777
2778 unsafe {
2780 phase1_u32_fn.clone().launch(
2781 LaunchConfig {
2782 grid_dim: (num_blocks, 1, 1),
2783 block_dim: (block_size, 1, 1),
2784 shared_mem_bytes: 0,
2785 },
2786 (&mut *data, &mut block_sums, n),
2787 )
2788 }
2789 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_u32_phase1 failed: {}", e)))?;
2790 self.device.synchronize()?;
2791
2792 if num_blocks > 1 {
2793 self.multiblock_scan_u32_inplace(&mut block_sums, num_blocks)?;
2794 }
2795
2796 let phase3_fn = device
2797 .get_func(SCAN_MODULE, scan_kernels::MULTIBLOCK_SCAN_PHASE3)
2798 .ok_or_else(|| {
2799 XlogError::Kernel("Failed to get multiblock_scan_phase3 kernel".to_string())
2800 })?;
2801
2802 unsafe {
2804 phase3_fn.clone().launch(
2805 LaunchConfig {
2806 grid_dim: (num_blocks, 1, 1),
2807 block_dim: (block_size, 1, 1),
2808 shared_mem_bytes: 0,
2809 },
2810 (&mut *data, &block_sums, n),
2811 )
2812 }
2813 .map_err(|e| XlogError::Kernel(format!("multiblock_scan_phase3 failed: {}", e)))?;
2814
2815 self.device.synchronize()?;
2816 Ok(())
2817 }
2818
2819 pub fn device_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2824 if let Some(n) = buffer.cached_row_count() {
2825 return Ok(n as usize);
2826 }
2827 let host_rows = self.dtoh_small_metadata_untracked(buffer.num_rows_device(), 1)?;
2828 buffer.set_cached_row_count_if_unset(host_rows[0]);
2829 Ok(host_rows[0] as usize)
2830 }
2831
2832 pub fn validated_logical_row_count(&self, buffer: &CudaBuffer) -> Result<usize> {
2837 let logical_rows = self.device_row_count(buffer)?;
2838 validate_logical_row_count(buffer.num_rows(), logical_rows)
2839 }
2840
2841 fn clone_device_row_count(&self, buffer: &CudaBuffer) -> Result<TrackedCudaSlice<u32>> {
2842 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2843 self.device
2844 .inner()
2845 .dtod_copy(buffer.num_rows_device(), &mut d_num_rows)
2846 .map_err(|e| XlogError::Kernel(format!("Failed to copy row count: {}", e)))?;
2847 Ok(d_num_rows)
2848 }
2849
2850 fn upload_device_row_count(&self, row_count: u32) -> Result<TrackedCudaSlice<u32>> {
2851 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
2852 self.htod_launch_metadata_sync_copy_into(&[row_count], &mut d_num_rows)
2853 .map_err(|e| XlogError::Kernel(format!("Failed to upload row count: {}", e)))?;
2854 Ok(d_num_rows)
2855 }
2856
2857 fn buffer_from_columns_with_device_count(
2858 &self,
2859 columns: Vec<CudaColumn>,
2860 row_cap: u64,
2861 schema: Schema,
2862 src: &CudaBuffer,
2863 ) -> Result<CudaBuffer> {
2864 let d_num_rows = self.clone_device_row_count(src)?;
2865 Ok(match src.cached_row_count() {
2866 Some(row_count) => CudaBuffer::from_columns_with_host_count(
2867 columns, row_cap, d_num_rows, schema, row_count,
2868 ),
2869 None => CudaBuffer::from_columns(columns, row_cap, d_num_rows, schema),
2870 })
2871 }
2872
2873 fn column_bytes_view<'a>(
2874 &self,
2875 col: &'a CudaColumn,
2876 num_bytes: usize,
2877 ) -> Result<RawCudaView<'a, u8>> {
2878 if col.num_bytes() < num_bytes {
2879 return Err(XlogError::Kernel(format!(
2880 "Column has {} bytes but {} required",
2881 col.num_bytes(),
2882 num_bytes
2883 )));
2884 }
2885 let ptr = *col.device_ptr();
2886 Ok(RawCudaView {
2887 ptr,
2888 len: num_bytes,
2889 stream: col.stream().clone(),
2890 _marker: PhantomData,
2891 })
2892 }
2893
2894 fn bytes_as_u32_view<'a>(
2895 &self,
2896 bytes: &'a TrackedCudaSlice<u8>,
2897 num_elements: usize,
2898 ) -> Result<RawCudaView<'a, u32>> {
2899 let required_bytes = num_elements * std::mem::size_of::<u32>();
2900 if bytes.len() < required_bytes {
2901 return Err(XlogError::Kernel(format!(
2902 "Packed keys have {} bytes but {} required for {} u32 elements",
2903 bytes.len(),
2904 required_bytes,
2905 num_elements
2906 )));
2907 }
2908 let ptr = *bytes.device_ptr();
2909 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2910 return Err(XlogError::Kernel(
2911 "Packed keys device pointer is not u32-aligned".to_string(),
2912 ));
2913 }
2914 Ok(RawCudaView {
2915 ptr,
2916 len: num_elements,
2917 stream: bytes.stream().clone(),
2918 _marker: PhantomData,
2919 })
2920 }
2921
2922 fn column_as_u32_view<'a>(
2924 &self,
2925 col: &'a CudaColumn,
2926 num_elements: usize,
2927 ) -> Result<RawCudaView<'a, u32>> {
2928 let required_bytes = num_elements * std::mem::size_of::<u32>();
2929 if col.num_bytes() < required_bytes {
2930 return Err(XlogError::Kernel(format!(
2931 "Column has {} bytes but {} required for {} u32 elements",
2932 col.num_bytes(),
2933 required_bytes,
2934 num_elements
2935 )));
2936 }
2937 let ptr = *col.device_ptr();
2938 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u32>()) {
2939 return Err(XlogError::Kernel(
2940 "Column device pointer is not u32-aligned".to_string(),
2941 ));
2942 }
2943 Ok(RawCudaView {
2944 ptr,
2945 len: num_elements,
2946 stream: col.stream().clone(),
2947 _marker: PhantomData,
2948 })
2949 }
2950
2951 fn column_as_u64_view<'a>(
2952 &self,
2953 col: &'a CudaColumn,
2954 num_elements: usize,
2955 ) -> Result<RawCudaView<'a, u64>> {
2956 let required_bytes = num_elements * std::mem::size_of::<u64>();
2957 if col.num_bytes() < required_bytes {
2958 return Err(XlogError::Kernel(format!(
2959 "Column has {} bytes but {} required for {} u64 elements",
2960 col.num_bytes(),
2961 required_bytes,
2962 num_elements
2963 )));
2964 }
2965 let ptr = *col.device_ptr();
2966 if !(ptr as usize).is_multiple_of(std::mem::align_of::<u64>()) {
2967 return Err(XlogError::Kernel(
2968 "Column device pointer is not u64-aligned".to_string(),
2969 ));
2970 }
2971 Ok(RawCudaView {
2972 ptr,
2973 len: num_elements,
2974 stream: col.stream().clone(),
2975 _marker: PhantomData,
2976 })
2977 }
2978
2979 fn column_as_f64_view<'a>(
2981 &self,
2982 col: &'a CudaColumn,
2983 num_elements: usize,
2984 ) -> Result<RawCudaView<'a, f64>> {
2985 let required_bytes = num_elements * std::mem::size_of::<f64>();
2986 if col.num_bytes() < required_bytes {
2987 return Err(XlogError::Kernel(format!(
2988 "Column has {} bytes but {} required for {} f64 elements",
2989 col.num_bytes(),
2990 required_bytes,
2991 num_elements
2992 )));
2993 }
2994 let ptr = *col.device_ptr();
2995 if !(ptr as usize).is_multiple_of(std::mem::align_of::<f64>()) {
2996 return Err(XlogError::Kernel(
2997 "Column device pointer is not f64-aligned".to_string(),
2998 ));
2999 }
3000 Ok(RawCudaView {
3001 ptr,
3002 len: num_elements,
3003 stream: col.stream().clone(),
3004 _marker: PhantomData,
3005 })
3006 }
3007
3008 pub fn create_empty_buffer(&self, schema: Schema) -> Result<CudaBuffer> {
3019 let mut columns = Vec::with_capacity(schema.arity());
3020 for _ in 0..schema.arity() {
3021 columns.push(self.memory.alloc::<u8>(0)?.into());
3023 }
3024 self.buffer_from_columns(columns, 0, schema)
3025 }
3026
3027 pub fn create_zero_arity_buffer(&self, schema: Schema, rows: u32) -> Result<CudaBuffer> {
3035 debug_assert_eq!(
3036 schema.arity(),
3037 0,
3038 "create_zero_arity_buffer requires arity 0"
3039 );
3040 self.buffer_from_columns(Vec::new(), u64::from(rows), schema)
3041 }
3042
3043 pub(crate) fn buffer_from_columns(
3044 &self,
3045 columns: Vec<CudaColumn>,
3046 row_cap: u64,
3047 schema: Schema,
3048 ) -> Result<CudaBuffer> {
3049 let row_u32 = u32::try_from(row_cap)
3050 .map_err(|_| XlogError::Kernel(format!("Row capacity {} exceeds u32::MAX", row_cap)))?;
3051 let mut d_num_rows = self.memory.alloc::<u32>(1)?;
3052 self.htod_launch_metadata_sync_copy_into(&[row_u32], &mut d_num_rows)
3053 .map_err(|e| XlogError::Kernel(format!("Failed to set row count: {}", e)))?;
3054 let mut result =
3055 CudaBuffer::from_columns_with_host_count(columns, row_cap, d_num_rows, schema, row_u32);
3056 if row_u32 <= 1 {
3057 result.certify_canonical_full_row_set();
3058 }
3059 Ok(result)
3060 }
3061
3062 fn combine_schemas(&self, left: &Schema, right: &Schema) -> Schema {
3064 let mut columns = left.columns.clone();
3065 columns.extend(right.columns.iter().cloned());
3066 let mut sort_labels = left.sort_labels().to_vec();
3067 sort_labels.extend(right.sort_labels().iter().cloned());
3068 Schema::new(columns)
3069 .with_sort_labels(sort_labels)
3070 .expect("combined schema sort labels match column arity")
3071 }
3072
3073 fn schemas_type_compatible(&self, a: &Schema, b: &Schema) -> bool {
3078 if a.arity() != b.arity() {
3079 return false;
3080 }
3081 for i in 0..a.arity() {
3082 if a.column_type(i) != b.column_type(i) {
3083 return false;
3084 }
3085 }
3086 true
3087 }
3088}
3089
3090#[cfg(test)]
3091mod tests {
3092 use super::*;
3093 use crate::device_runtime::{
3094 AsyncCudaResource, DeviceMemoryResource, GlobalDeviceBudget, LoggingResource, NullSink,
3095 StreamPool, XlogDeviceRuntime,
3096 };
3097 use xlog_core::{AggOp, MemoryBudget, ScalarType};
3098
3099 fn has_cuda_device() -> bool {
3100 CudaDevice::new(0).is_ok()
3101 }
3102
3103 #[test]
3104 fn test_kernel_artifact_locator_precedence_order() {
3105 use super::kernel_paths::KernelArtifactLocator;
3106 use std::fs;
3107 use std::path::PathBuf;
3108
3109 let root = std::env::temp_dir().join(format!(
3110 "xlog-kernel-paths-{}-{}",
3111 std::process::id(),
3112 std::time::SystemTime::now()
3113 .duration_since(std::time::UNIX_EPOCH)
3114 .expect("system clock before UNIX_EPOCH")
3115 .as_nanos()
3116 ));
3117 let cubin_dir = root.join("cubin");
3118 let package_dir = root.join("bin").join("kernels");
3119 let out_dir = root.join("out");
3120 fs::create_dir_all(&cubin_dir).expect("create cubin dir");
3121 fs::create_dir_all(&package_dir).expect("create package kernels dir");
3122 fs::create_dir_all(&out_dir).expect("create out dir");
3123
3124 let name = "xlog_join";
3125 let cc = 75;
3126 let cubin_path = cubin_dir.join(format!("{name}.sm_{cc}.cubin"));
3127 let package_path = package_dir.join(format!("{name}.sm_{cc}.cubin"));
3128 let out_path = out_dir.join(format!("{name}.sm_{cc}.cubin"));
3129 fs::write(&cubin_path, b"cubin").expect("write cubin file");
3130 fs::write(&package_path, b"package").expect("write package file");
3131 fs::write(&out_path, b"out").expect("write out file");
3132
3133 let locator = KernelArtifactLocator::new(
3134 Some(cubin_dir.clone()),
3135 Some(package_dir.clone()),
3136 Some(out_dir.clone()),
3137 );
3138
3139 let (path, is_cubin) = locator
3140 .resolve_module_path(name, cc)
3141 .expect("expected a kernel artifact");
3142 assert_eq!(path, cubin_path);
3143 assert!(is_cubin);
3144
3145 fs::remove_file(&cubin_path).expect("remove cubin file");
3146 let (path, is_cubin) = locator
3147 .resolve_module_path(name, cc)
3148 .expect("expected package kernel artifact");
3149 assert_eq!(path, package_path);
3150 assert!(is_cubin);
3151
3152 fs::remove_file(&package_path).expect("remove package file");
3153 let (path, is_cubin) = locator
3154 .resolve_module_path(name, cc)
3155 .expect("expected out dir kernel artifact");
3156 assert_eq!(path, out_path);
3157 assert!(is_cubin);
3158
3159 let _ = fs::remove_dir_all(PathBuf::from(&root));
3160 }
3161
3162 #[test]
3163 fn test_module_resolution_finds_portable_ptx() {
3164 for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
3167 let result = resolve_module_path(name, 999);
3168 assert!(
3169 result.is_some(),
3170 "resolve_module_path({name}, 999) should find portable PTX"
3171 );
3172 let (path, is_cubin) = result.unwrap();
3173 assert!(
3174 !is_cubin,
3175 "{name}: expected portable PTX fallback, got cubin"
3176 );
3177 assert!(
3178 path.to_str().unwrap().ends_with(".portable.ptx"),
3179 "{name}: path should end with .portable.ptx, got {:?}",
3180 path
3181 );
3182 }
3183 }
3184
3185 #[test]
3186 fn test_module_resolution_falls_back_to_embedded_portable_ptx() {
3187 use super::kernel_paths::KernelArtifactLocator;
3188
3189 let locator = KernelArtifactLocator::new(None, None, None);
3190 for name in crate::kernel_manifest_data::KERNEL_CU_NAMES {
3191 let sources = resolve_module_sources_with_locator(name, 999, &locator);
3192 assert_eq!(
3193 sources.len(),
3194 1,
3195 "{name}: expected only embedded portable PTX fallback"
3196 );
3197
3198 match &sources[0] {
3199 KernelModuleSource::EmbeddedPortablePtx { ptx } => {
3200 assert!(
3201 ptx.contains(".entry"),
3202 "{name}: embedded PTX should contain CUDA entry points"
3203 );
3204 }
3205 KernelModuleSource::File { path, .. } => {
3206 panic!(
3207 "{name}: expected embedded portable PTX fallback, got file {}",
3208 path.display()
3209 );
3210 }
3211 }
3212 }
3213 }
3214
3215 #[test]
3216 fn test_embedded_portable_ptx_manifest_matches_kernel_manifest() {
3217 let embedded_names: std::collections::BTreeSet<_> =
3218 crate::embedded_kernel_data::EMBEDDED_PORTABLE_PTX
3219 .iter()
3220 .map(|artifact| artifact.name)
3221 .collect();
3222 let manifest_names: std::collections::BTreeSet<_> =
3223 crate::kernel_manifest_data::KERNEL_CU_NAMES
3224 .iter()
3225 .copied()
3226 .collect();
3227
3228 assert_eq!(
3229 embedded_names, manifest_names,
3230 "embedded portable PTX table should cover every runtime kernel module"
3231 );
3232 }
3233
3234 #[test]
3235 fn test_kernel_provider_creation() {
3236 if !has_cuda_device() {
3237 eprintln!("Skipping test: no CUDA device available");
3238 return;
3239 }
3240
3241 let device = Arc::new(CudaDevice::new(0).expect("Failed to create device"));
3242 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024); let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3244
3245 let provider = CudaKernelProvider::new(device.clone(), memory.clone());
3246 assert!(
3247 provider.is_ok(),
3248 "Failed to create kernel provider: {:?}",
3249 provider.err()
3250 );
3251
3252 let provider = provider.unwrap();
3253 assert!(Arc::ptr_eq(provider.device(), &device));
3254 assert!(Arc::ptr_eq(provider.memory(), &memory));
3255 }
3256
3257 #[test]
3258 fn test_kernel_functions_accessible() {
3259 if !has_cuda_device() {
3260 eprintln!("Skipping test: no CUDA device available");
3261 return;
3262 }
3263
3264 let device = Arc::new(CudaDevice::new(0).expect("Failed to create device"));
3265 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3266 let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3267
3268 let _provider =
3269 CudaKernelProvider::new(device.clone(), memory).expect("Failed to create provider");
3270
3271 let inner = device.inner();
3273
3274 let build_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_BUILD);
3276 assert!(
3277 build_fn.is_some(),
3278 "hash_join_build function should be accessible"
3279 );
3280
3281 let probe_fn = inner.get_func(JOIN_MODULE, join_kernels::HASH_JOIN_PROBE);
3282 assert!(
3283 probe_fn.is_some(),
3284 "hash_join_probe function should be accessible"
3285 );
3286
3287 let mark_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::MARK_DUPLICATES);
3289 assert!(
3290 mark_fn.is_some(),
3291 "mark_duplicates function should be accessible"
3292 );
3293
3294 let compact_fn = inner.get_func(DEDUP_MODULE, dedup_kernels::COMPACT_ROWS);
3295 assert!(
3296 compact_fn.is_some(),
3297 "compact_rows function should be accessible"
3298 );
3299
3300 let boundaries_fn =
3302 inner.get_func(GROUPBY_MODULE, groupby_kernels::DETECT_GROUP_BOUNDARIES);
3303 assert!(
3304 boundaries_fn.is_some(),
3305 "detect_group_boundaries function should be accessible"
3306 );
3307
3308 let count_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_COUNT);
3309 assert!(
3310 count_fn.is_some(),
3311 "groupby_count function should be accessible"
3312 );
3313
3314 let sum_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_SUM);
3315 assert!(
3316 sum_fn.is_some(),
3317 "groupby_sum function should be accessible"
3318 );
3319
3320 let min_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MIN);
3321 assert!(
3322 min_fn.is_some(),
3323 "groupby_min function should be accessible"
3324 );
3325
3326 let max_fn = inner.get_func(GROUPBY_MODULE, groupby_kernels::GROUPBY_MAX);
3327 assert!(
3328 max_fn.is_some(),
3329 "groupby_max function should be accessible"
3330 );
3331
3332 let xgcf_forward = inner.get_func(CIRCUIT_MODULE, "xgcf_forward_level");
3334 assert!(
3335 xgcf_forward.is_some(),
3336 "xgcf_forward_level function should be accessible"
3337 );
3338
3339 let xgcf_backward_propagate =
3340 inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_propagate");
3341 assert!(
3342 xgcf_backward_propagate.is_some(),
3343 "xgcf_backward_level_propagate function should be accessible"
3344 );
3345
3346 let xgcf_backward_decision_grad =
3347 inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_decision_grad");
3348 assert!(
3349 xgcf_backward_decision_grad.is_some(),
3350 "xgcf_backward_level_decision_grad function should be accessible"
3351 );
3352
3353 let xgcf_backward_lit_grad = inner.get_func(CIRCUIT_MODULE, "xgcf_backward_level_lit_grad");
3354 assert!(
3355 xgcf_backward_lit_grad.is_some(),
3356 "xgcf_backward_level_lit_grad function should be accessible"
3357 );
3358
3359 let neural_fill = inner.get_func("xlog_neural", "neural_fill_ad_chain_f32");
3361 assert!(
3362 neural_fill.is_some(),
3363 "neural_fill_ad_chain_f32 function should be accessible"
3364 );
3365 let neural_scatter = inner.get_func("xlog_neural", "neural_scatter_ad_chain_grads_f32");
3366 assert!(
3367 neural_scatter.is_some(),
3368 "neural_scatter_ad_chain_grads_f32 function should be accessible"
3369 );
3370 }
3371
3372 #[test]
3373 fn test_module_names_unique() {
3374 assert_ne!(JOIN_MODULE, DEDUP_MODULE);
3376 assert_ne!(JOIN_MODULE, GROUPBY_MODULE);
3377 assert_ne!(DEDUP_MODULE, GROUPBY_MODULE);
3378 }
3379
3380 fn create_test_provider() -> Option<CudaKernelProvider> {
3382 if !has_cuda_device() {
3383 return None;
3384 }
3385 let device = Arc::new(CudaDevice::new(0).ok()?);
3386 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
3387 let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
3388 CudaKernelProvider::new(device, memory).ok()
3389 }
3390
3391 fn create_test_provider_with_runtime() -> Option<(CudaKernelProvider, Arc<XlogDeviceRuntime>)> {
3392 if !has_cuda_device() {
3393 return None;
3394 }
3395 let device = Arc::new(CudaDevice::new(0).ok()?);
3396 let pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
3397 let sink = Arc::new(NullSink::new());
3398 let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
3399 AsyncCudaResource::new(Arc::clone(&device), 0, Arc::clone(&pool)),
3400 );
3401 let logging: Box<dyn DeviceMemoryResource + Send + Sync> =
3402 Box::new(LoggingResource::new(async_resource, sink));
3403 let budget: Box<dyn DeviceMemoryResource + Send + Sync> =
3404 Box::new(GlobalDeviceBudget::new(logging, 1024 * 1024 * 1024));
3405 let runtime = Arc::new(XlogDeviceRuntime::with_resource(
3406 Arc::clone(&device),
3407 0,
3408 pool,
3409 budget,
3410 ));
3411 let memory = Arc::new(GpuMemoryManager::with_runtime(
3412 Arc::clone(&device),
3413 MemoryBudget::with_limit(1024 * 1024 * 1024),
3414 Arc::clone(&runtime),
3415 ));
3416 let provider = CudaKernelProvider::with_runtime(device, memory).ok()?;
3417 Some((provider, runtime))
3418 }
3419
3420 #[test]
3421 fn test_recorded_join_index_build_runs_on_runtime_stream() {
3422 let (provider, runtime) = match create_test_provider_with_runtime() {
3423 Some(fixture) => fixture,
3424 None => {
3425 eprintln!("Skipping test: no CUDA device available");
3426 return;
3427 }
3428 };
3429 let stream = runtime.stream_pool().acquire().expect("recorded stream");
3430 let left = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3431 let right = create_test_buffer(&provider, &[1, 2, 3, 4], "key");
3432
3433 let index = provider
3434 .build_join_index_v2_recorded(&right, &[0], stream)
3435 .expect("recorded join-index build");
3436 let joined = provider
3437 .hash_join_v2_with_index_recorded(
3438 &left,
3439 &right,
3440 &[0],
3441 &[0],
3442 JoinType::Inner,
3443 &index,
3444 None,
3445 stream,
3446 )
3447 .expect("recorded indexed join consumes recorded build");
3448 runtime
3449 .stream_pool()
3450 .resolve(stream)
3451 .expect("stream resolves")
3452 .synchronize()
3453 .expect("recorded stream synchronized");
3454
3455 assert_eq!(index.right_num_rows(), 4);
3456 assert_eq!(index.right_keys(), &[0]);
3457 assert_eq!(provider.device_row_count(&joined).expect("joined rows"), 4);
3458 }
3459
3460 fn create_test_buffer(
3462 provider: &CudaKernelProvider,
3463 data: &[u32],
3464 col_name: &str,
3465 ) -> CudaBuffer {
3466 let schema = Schema::new(vec![(col_name.to_string(), ScalarType::U32)]);
3467 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3468
3469 let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3470 provider
3471 .device()
3472 .inner()
3473 .htod_sync_copy_into(&bytes, &mut col)
3474 .expect("htod");
3475
3476 provider
3477 .buffer_from_columns(vec![col.into()], data.len() as u64, schema)
3478 .expect("buffer")
3479 }
3480
3481 fn create_empty_test_buffer(provider: &CudaKernelProvider, schema: Schema) -> CudaBuffer {
3483 let mut columns = Vec::with_capacity(schema.arity());
3484 for _ in 0..schema.arity() {
3485 columns.push(provider.memory().alloc::<u8>(0).expect("alloc").into());
3486 }
3487 provider
3488 .buffer_from_columns(columns, 0, schema)
3489 .expect("buffer")
3490 }
3491
3492 fn read_buffer_u32(provider: &CudaKernelProvider, buffer: &CudaBuffer, col: usize) -> Vec<u32> {
3494 if buffer.is_empty() || buffer.column(col).is_none() {
3495 return vec![];
3496 }
3497 let num_rows = buffer.num_rows() as usize;
3498 let mut bytes = vec![0u8; num_rows * 4];
3499 provider
3500 .device()
3501 .inner()
3502 .dtoh_sync_copy_into(buffer.column(col).unwrap(), &mut bytes)
3503 .expect("dtoh");
3504 bytes
3505 .chunks_exact(4)
3506 .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
3507 .collect()
3508 }
3509
3510 #[test]
3511 fn test_compact_device_mask_respects_mask_len_smaller_than_row_cap() {
3512 let provider = match create_test_provider() {
3513 Some(p) => p,
3514 None => {
3515 eprintln!("Skipping test: no CUDA device available");
3516 return;
3517 }
3518 };
3519
3520 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3521 let base = create_test_buffer(&provider, &[1, 2, 3, 4, 5, 6, 7, 8], "id");
3522
3523 let row_cap = 16u64;
3524 let data: Vec<u32> = (0..row_cap as u32).collect();
3525 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
3526 let mut col = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
3527 provider
3528 .device()
3529 .inner()
3530 .htod_sync_copy_into(&bytes, &mut col)
3531 .expect("htod");
3532 let expanded = provider
3533 .buffer_from_columns_with_device_count(vec![col.into()], row_cap, schema, &base)
3534 .expect("buffer");
3535
3536 let mask: Vec<u8> = vec![1, 0, 1, 0, 1, 0, 1, 0];
3537 let (prefix_sum, count) = provider.prefix_sum_mask(&mask).expect("prefix sum");
3538
3539 let mut d_mask = provider.memory().alloc::<u8>(mask.len()).expect("alloc");
3540 provider
3541 .device()
3542 .inner()
3543 .htod_sync_copy_into(&mask, &mut d_mask)
3544 .expect("mask htod");
3545
3546 let mut d_prefix = provider
3547 .memory()
3548 .alloc::<u32>(prefix_sum.len())
3549 .expect("alloc");
3550 provider
3551 .device()
3552 .inner()
3553 .htod_sync_copy_into(&prefix_sum, &mut d_prefix)
3554 .expect("prefix htod");
3555
3556 let mut d_out_count = provider.memory().alloc::<u32>(1).expect("alloc");
3557 provider
3558 .device()
3559 .inner()
3560 .htod_sync_copy_into(&[count], &mut d_out_count)
3561 .expect("count htod");
3562
3563 let compacted = provider
3564 .compact_buffer_by_device_mask_device_count(&expanded, &d_mask, &d_prefix, d_out_count)
3565 .expect("compact");
3566
3567 assert_eq!(compacted.num_rows(), mask.len() as u64);
3568 let device_rows = provider.device_row_count(&compacted).expect("row count");
3569 assert_eq!(device_rows as u32, count);
3570 }
3571
3572 #[test]
3573 fn test_clone_buffer_preserves_device_count() {
3574 let provider = match create_test_provider() {
3575 Some(p) => p,
3576 None => {
3577 eprintln!("Skipping test: no CUDA device available");
3578 return;
3579 }
3580 };
3581
3582 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3583 let ids: Vec<u32> = vec![10, 20, 30];
3584 let buffer = provider
3585 .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3586 .unwrap();
3587
3588 let cloned = provider.clone_buffer(&buffer).unwrap();
3589
3590 let mut host_count = [0u32];
3591 provider
3592 .device()
3593 .inner()
3594 .dtoh_sync_copy_into(cloned.num_rows_device(), &mut host_count)
3595 .unwrap();
3596 assert_eq!(host_count[0], 3);
3597 }
3598
3599 #[test]
3608 fn test_clone_buffer_preserves_cached_row_count() {
3609 let provider = match create_test_provider() {
3610 Some(p) => p,
3611 None => {
3612 eprintln!("Skipping test: no CUDA device available");
3613 return;
3614 }
3615 };
3616
3617 let schema = Schema::new(vec![("id".to_string(), ScalarType::U32)]);
3618 let ids: Vec<u32> = vec![7, 11, 13, 17];
3619 let source = provider
3620 .create_buffer_from_slices(&[bytemuck::cast_slice(&ids)], schema)
3621 .unwrap();
3622 assert_eq!(
3626 source.cached_row_count(),
3627 Some(4),
3628 "source buffer should have its cached row count populated by \
3629 create_buffer_from_slices"
3630 );
3631
3632 let cloned = provider.clone_buffer(&source).unwrap();
3633
3634 assert_eq!(
3635 cloned.cached_row_count(),
3636 Some(4),
3637 "clone_buffer must propagate cached_row_count from source to clone",
3638 );
3639 }
3640
3641 #[test]
3642 fn certified_singleton_union_avoids_cold_row_count_read() {
3643 let provider = match create_test_provider() {
3644 Some(p) => p,
3645 None => {
3646 eprintln!("Skipping test: no CUDA device available");
3647 return;
3648 }
3649 };
3650
3651 let source = create_test_buffer(&provider, &[1, 2, 3], "id");
3652 let CudaBuffer {
3653 columns,
3654 row_cap,
3655 d_num_rows,
3656 schema,
3657 ..
3658 } = source;
3659 let mut cold_set = CudaBuffer::from_columns(columns, row_cap, d_num_rows, schema);
3660 cold_set.certify_canonical_full_row_set();
3661 assert_eq!(cold_set.cached_row_count(), None);
3662 let empty = provider
3663 .create_empty_buffer(cold_set.schema().clone())
3664 .unwrap();
3665
3666 provider.reset_host_transfer_stats();
3667 provider.reset_untracked_metadata_dtoh_count();
3668 provider.memory().reset_alloc_count();
3669 let result = provider
3670 .union_many_gpu(&[&empty, &cold_set, &empty])
3671 .unwrap();
3672
3673 assert_eq!(
3674 cold_set.cached_row_count(),
3675 None,
3676 "the certified fast path must not populate the cold count via D2H"
3677 );
3678 assert_eq!(provider.host_transfer_stats().dtoh_calls, 0);
3679 assert_eq!(provider.untracked_metadata_dtoh_count(), 0);
3680 assert_eq!(provider.memory().alloc_count(), 2);
3681 assert!(result.canonical_full_row_set_certified());
3682 assert_eq!(read_buffer_u32(&provider, &result, 0), vec![1, 2, 3]);
3683 }
3684
3685 #[test]
3686 fn cold_row_count_read_uses_audited_metadata_counter_once() {
3687 let provider = match create_test_provider() {
3688 Some(p) => p,
3689 None => {
3690 eprintln!("Skipping test: no CUDA device available");
3691 return;
3692 }
3693 };
3694
3695 let source = create_test_buffer(&provider, &[1, 2, 3], "id");
3696 let CudaBuffer {
3697 columns,
3698 row_cap,
3699 d_num_rows,
3700 schema,
3701 ..
3702 } = source;
3703 let cold_bag = CudaBuffer::from_columns(columns, row_cap, d_num_rows, schema);
3704
3705 provider.reset_untracked_metadata_dtoh_count();
3706 assert_eq!(provider.device_row_count(&cold_bag).unwrap(), 3);
3707
3708 assert_eq!(provider.untracked_metadata_dtoh_count(), 1);
3709 assert_eq!(cold_bag.cached_row_count(), Some(3));
3710 provider.reset_untracked_metadata_dtoh_count();
3711 let result = provider.union_many_gpu(&[&cold_bag]).unwrap();
3712 assert_eq!(provider.untracked_metadata_dtoh_count(), 1);
3713 assert_eq!(read_buffer_u32(&provider, &result, 0), vec![1, 2, 3]);
3714 }
3715
3716 #[test]
3717 fn zero_arity_certificate_does_not_bypass_unit_set_semantics() {
3718 let provider = match create_test_provider() {
3719 Some(p) => p,
3720 None => {
3721 eprintln!("Skipping test: no CUDA device available");
3722 return;
3723 }
3724 };
3725
3726 let d_num_rows = provider.upload_device_row_count(3).unwrap();
3727 let mut multiplicity = CudaBuffer::from_columns(
3728 Vec::new(),
3729 3,
3730 d_num_rows,
3731 xlog_core::Schema::new(Vec::new()),
3732 );
3733 multiplicity.certify_canonical_full_row_set();
3734
3735 let unit = provider.union_many_gpu(&[&multiplicity]).unwrap();
3736 assert_eq!(provider.device_row_count(&unit).unwrap(), 1);
3737 }
3738
3739 #[test]
3742 fn test_hash_join_empty_inputs() {
3743 let provider = match create_test_provider() {
3744 Some(p) => p,
3745 None => {
3746 eprintln!("Skipping test: no CUDA device available");
3747 return;
3748 }
3749 };
3750
3751 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3752 let empty = create_empty_test_buffer(&provider, schema.clone());
3753
3754 let result = provider.hash_join(&empty, &empty, &[0], &[0]);
3756 assert!(result.is_ok());
3757 assert!(result.unwrap().is_empty());
3758 }
3759
3760 #[test]
3761 fn test_hash_join_validation() {
3762 let provider = match create_test_provider() {
3763 Some(p) => p,
3764 None => {
3765 eprintln!("Skipping test: no CUDA device available");
3766 return;
3767 }
3768 };
3769
3770 let left = create_test_buffer(&provider, &[1, 2, 3], "left_key");
3771 let right = create_test_buffer(&provider, &[2, 3, 4], "right_key");
3772
3773 let result = provider.hash_join(&left, &right, &[], &[0]);
3775 assert!(result.is_err());
3776
3777 let result = provider.hash_join(&left, &right, &[0], &[0, 0]);
3779 assert!(result.is_err());
3780 }
3781
3782 #[test]
3785 fn test_dedup_empty_input() {
3786 let provider = match create_test_provider() {
3787 Some(p) => p,
3788 None => {
3789 eprintln!("Skipping test: no CUDA device available");
3790 return;
3791 }
3792 };
3793
3794 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3795 let empty = create_empty_test_buffer(&provider, schema);
3796
3797 let result = provider.dedup(&empty, &[0]);
3798 assert!(result.is_ok());
3799 assert!(result.unwrap().is_empty());
3800 }
3801
3802 #[test]
3803 fn test_dedup_validation() {
3804 let provider = match create_test_provider() {
3805 Some(p) => p,
3806 None => {
3807 eprintln!("Skipping test: no CUDA device available");
3808 return;
3809 }
3810 };
3811
3812 let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
3813
3814 let result = provider.dedup(&buffer, &[]);
3816 assert!(result.is_err());
3817 }
3818
3819 #[test]
3820 fn test_dedup_with_duplicates() {
3821 let provider = match create_test_provider() {
3822 Some(p) => p,
3823 None => {
3824 eprintln!("Skipping test: no CUDA device available");
3825 return;
3826 }
3827 };
3828
3829 let buffer = create_test_buffer(&provider, &[3, 1, 2, 1, 3, 2], "key");
3831 let deduped = provider.dedup(&buffer, &[0]).unwrap();
3832
3833 let dedup_count = provider
3834 .device_row_count(&deduped)
3835 .expect("read dedup row count");
3836 assert_eq!(dedup_count, 3, "Should have 3 unique values");
3837
3838 let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3839 assert_eq!(result, vec![1, 2, 3]);
3841 }
3842
3843 #[test]
3844 fn test_dedup_larger_input() {
3845 let provider = match create_test_provider() {
3846 Some(p) => p,
3847 None => {
3848 eprintln!("Skipping test: no CUDA device available");
3849 return;
3850 }
3851 };
3852
3853 let a: Vec<u32> = (0..500).collect();
3855 let b: Vec<u32> = (250..750).collect();
3856 let input: Vec<u32> = a.iter().chain(b.iter()).copied().collect();
3857
3858 let buffer = create_test_buffer(&provider, &input, "key");
3859 let deduped = provider.dedup(&buffer, &[0]).unwrap();
3860
3861 let dedup_count = provider
3862 .device_row_count(&deduped)
3863 .expect("read dedup row count");
3864 assert_eq!(dedup_count, 750, "Should have 750 unique values (0..750)");
3865
3866 let result = provider.download_column::<u32>(&deduped, 0).unwrap();
3868 let is_sorted = result.windows(2).all(|w| w[0] <= w[1]);
3869 assert!(is_sorted, "Output should be sorted");
3870
3871 let expected: Vec<u32> = (0..750).collect();
3873 assert_eq!(result, expected);
3874 }
3875
3876 #[test]
3879 fn test_union_empty_inputs() {
3880 let provider = match create_test_provider() {
3881 Some(p) => p,
3882 None => {
3883 eprintln!("Skipping test: no CUDA device available");
3884 return;
3885 }
3886 };
3887
3888 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3889 let empty = create_empty_test_buffer(&provider, schema.clone());
3890
3891 let result = provider.union(&empty, &empty);
3893 assert!(result.is_ok());
3894 assert!(result.unwrap().is_empty());
3895
3896 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3898 let empty2 = create_empty_test_buffer(&provider, schema);
3899 let result = provider.union(&a, &empty2);
3900 assert!(result.is_ok());
3901 let result = result.unwrap();
3902 assert_eq!(result.num_rows(), 3);
3903 }
3904
3905 #[test]
3906 fn test_union_schema_type_mismatch() {
3907 let provider = match create_test_provider() {
3908 Some(p) => p,
3909 None => {
3910 eprintln!("Skipping test: no CUDA device available");
3911 return;
3912 }
3913 };
3914
3915 let a = create_test_buffer(&provider, &[1, 2], "col_a");
3916 let b = create_test_buffer(&provider, &[3, 4], "col_b");
3917
3918 let result = provider.union(&a, &b);
3920 assert!(result.is_ok());
3921
3922 let two_col_schema = Schema::new(vec![
3924 ("x".to_string(), ScalarType::U32),
3925 ("y".to_string(), ScalarType::U32),
3926 ]);
3927 let c = provider
3928 .create_buffer_from_u32_columns(&[&[1, 2], &[3, 4]], two_col_schema)
3929 .unwrap();
3930 let result = provider.union(&a, &c);
3931 assert!(result.is_err());
3932 }
3933
3934 #[test]
3937 fn test_diff_empty_inputs() {
3938 let provider = match create_test_provider() {
3939 Some(p) => p,
3940 None => {
3941 eprintln!("Skipping test: no CUDA device available");
3942 return;
3943 }
3944 };
3945
3946 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
3947 let empty = create_empty_test_buffer(&provider, schema.clone());
3948
3949 let result = provider.diff(&empty, &empty);
3951 assert!(result.is_ok());
3952 assert!(result.unwrap().is_empty());
3953
3954 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3956 let empty2 = create_empty_test_buffer(&provider, schema);
3957 let result = provider.diff(&a, &empty2);
3958 assert!(result.is_ok());
3959 let result = result.unwrap();
3960 assert_eq!(result.num_rows(), 3);
3961 }
3962
3963 #[test]
3964 fn test_diff_basic() {
3965 let provider = match create_test_provider() {
3966 Some(p) => p,
3967 None => {
3968 eprintln!("Skipping test: no CUDA device available");
3969 return;
3970 }
3971 };
3972
3973 let a = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3974 let b = create_test_buffer(&provider, &[2, 4], "key");
3975
3976 let result = provider.diff(&a, &b);
3977 assert!(result.is_ok());
3978 let result = result.unwrap();
3979 assert_eq!(result.num_rows(), 3); let values = read_buffer_u32(&provider, &result, 0);
3982 assert_eq!(values, vec![1, 3, 5]);
3983 }
3984
3985 #[test]
3986 fn test_diff_all_filtered_out() {
3987 let provider = match create_test_provider() {
3988 Some(p) => p,
3989 None => {
3990 eprintln!("Skipping test: no CUDA device available");
3991 return;
3992 }
3993 };
3994
3995 let a = create_test_buffer(&provider, &[1, 2, 3], "key");
3996 let b = create_test_buffer(&provider, &[1, 2, 3, 4, 5], "key");
3997
3998 let result = provider.diff(&a, &b);
3999 assert!(result.is_ok());
4000 assert!(result.unwrap().is_empty());
4001 }
4002
4003 #[test]
4004 fn test_diff_schema_mismatch() {
4005 let provider = match create_test_provider() {
4006 Some(p) => p,
4007 None => {
4008 eprintln!("Skipping test: no CUDA device available");
4009 return;
4010 }
4011 };
4012
4013 let a = create_test_buffer(&provider, &[1, 2], "col_a");
4015 let b = create_test_buffer(&provider, &[1, 2], "col_b");
4016 let result = provider.diff(&a, &b);
4017 assert!(
4018 result.is_ok(),
4019 "Same types with different names should succeed"
4020 );
4021
4022 let schema_2col = Schema::new(vec![
4024 ("c0".to_string(), ScalarType::U32),
4025 ("c1".to_string(), ScalarType::U32),
4026 ]);
4027
4028 let bytes_2col: Vec<u8> = [1u32, 2, 3, 4]
4029 .iter()
4030 .flat_map(|v| v.to_le_bytes())
4031 .collect();
4032 let mut col0 = provider
4033 .memory()
4034 .alloc::<u8>(bytes_2col.len() / 2)
4035 .expect("alloc");
4036 let mut col1 = provider
4037 .memory()
4038 .alloc::<u8>(bytes_2col.len() / 2)
4039 .expect("alloc");
4040 provider
4041 .device()
4042 .inner()
4043 .htod_sync_copy_into(&bytes_2col[..8], &mut col0)
4044 .expect("htod");
4045 provider
4046 .device()
4047 .inner()
4048 .htod_sync_copy_into(&bytes_2col[8..], &mut col1)
4049 .expect("htod");
4050 let buffer_2col = provider
4051 .buffer_from_columns(vec![col0.into(), col1.into()], 2, schema_2col)
4052 .expect("buffer");
4053
4054 let buffer_1col = create_test_buffer(&provider, &[1, 2], "c0");
4055
4056 let result = provider.diff(&buffer_2col, &buffer_1col);
4057 assert!(result.is_err(), "Different arities should fail");
4058 }
4059
4060 #[test]
4063 fn test_groupby_empty_input() {
4064 let provider = match create_test_provider() {
4065 Some(p) => p,
4066 None => {
4067 eprintln!("Skipping test: no CUDA device available");
4068 return;
4069 }
4070 };
4071
4072 let schema = Schema::new(vec![("key".to_string(), ScalarType::U32)]);
4073 let empty = create_empty_test_buffer(&provider, schema);
4074
4075 let result = provider.groupby_agg(&empty, &[0], AggOp::Count, 0);
4076 assert!(result.is_ok());
4077 assert!(result.unwrap().is_empty());
4078 }
4079
4080 #[test]
4081 fn test_groupby_validation() {
4082 let provider = match create_test_provider() {
4083 Some(p) => p,
4084 None => {
4085 eprintln!("Skipping test: no CUDA device available");
4086 return;
4087 }
4088 };
4089
4090 let buffer = create_test_buffer(&provider, &[1, 1, 2, 2, 3], "key");
4091
4092 let result = provider.groupby_agg(&buffer, &[], AggOp::Count, 0);
4094 assert!(result.is_err());
4095
4096 let result = provider.groupby_agg(&buffer, &[0], AggOp::Count, 5);
4098 assert!(result.is_err());
4099 }
4100
4101 #[test]
4102 fn test_groupby_logsumexp() {
4103 let provider = match create_test_provider() {
4104 Some(p) => p,
4105 None => {
4106 eprintln!("Skipping test: no CUDA device available");
4107 return;
4108 }
4109 };
4110
4111 let keys: Vec<u32> = vec![1, 1, 2, 2];
4115 let values: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0];
4116
4117 let schema = Schema::new(vec![
4118 ("key".to_string(), ScalarType::U32),
4119 ("value".to_string(), ScalarType::F64),
4120 ]);
4121
4122 let key_bytes: Vec<u8> = keys.iter().flat_map(|v| v.to_le_bytes()).collect();
4124 let mut key_col = provider
4125 .memory()
4126 .alloc::<u8>(key_bytes.len())
4127 .expect("alloc key");
4128 provider
4129 .device()
4130 .inner()
4131 .htod_sync_copy_into(&key_bytes, &mut key_col)
4132 .expect("upload key");
4133
4134 let val_bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
4136 let mut val_col = provider
4137 .memory()
4138 .alloc::<u8>(val_bytes.len())
4139 .expect("alloc val");
4140 provider
4141 .device()
4142 .inner()
4143 .htod_sync_copy_into(&val_bytes, &mut val_col)
4144 .expect("upload val");
4145
4146 let buffer = provider
4147 .buffer_from_columns(vec![key_col.into(), val_col.into()], 4, schema)
4148 .expect("buffer");
4149
4150 let result = provider.groupby_agg(&buffer, &[0], AggOp::LogSumExp, 1);
4152 assert!(
4153 result.is_ok(),
4154 "groupby_agg with LogSumExp should succeed: {:?}",
4155 result.err()
4156 );
4157
4158 let result = result.unwrap();
4159 let group_count = provider
4160 .device_row_count(&result)
4161 .expect("read group count");
4162 assert_eq!(group_count, 2, "Should have 2 groups");
4163
4164 let result_values = provider
4166 .download_column::<f64>(&result, 1)
4167 .expect("download result");
4168
4169 let expected_0 = 2.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); let expected_1 = 4.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); let tolerance = 1e-5;
4176 assert!(
4177 (result_values[0] - expected_0).abs() < tolerance,
4178 "Group 0 logsumexp mismatch: got {}, expected {}",
4179 result_values[0],
4180 expected_0
4181 );
4182 assert!(
4183 (result_values[1] - expected_1).abs() < tolerance,
4184 "Group 1 logsumexp mismatch: got {}, expected {}",
4185 result_values[1],
4186 expected_1
4187 );
4188 }
4189
4190 #[test]
4193 fn test_combine_schemas() {
4194 let provider = match create_test_provider() {
4195 Some(p) => p,
4196 None => {
4197 eprintln!("Skipping test: no CUDA device available");
4198 return;
4199 }
4200 };
4201
4202 let left = Schema::new(vec![("a".to_string(), ScalarType::U32)]);
4203 let right = Schema::new(vec![("b".to_string(), ScalarType::U64)]);
4204
4205 let combined = provider.combine_schemas(&left, &right);
4206 assert_eq!(combined.arity(), 2);
4207 assert_eq!(combined.column_type(0), Some(ScalarType::U32));
4208 assert_eq!(combined.column_type(1), Some(ScalarType::U64));
4209 }
4210
4211 #[test]
4212 fn test_groupby_result_schema() {
4213 let provider = match create_test_provider() {
4214 Some(p) => p,
4215 None => {
4216 eprintln!("Skipping test: no CUDA device available");
4217 return;
4218 }
4219 };
4220
4221 let input = Schema::new(vec![
4222 ("key".to_string(), ScalarType::U32),
4223 ("value".to_string(), ScalarType::U32),
4224 ]);
4225
4226 let count_schema =
4228 provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Count)]);
4229 assert_eq!(count_schema.arity(), 2);
4230 assert_eq!(count_schema.column_type(1), Some(ScalarType::U64));
4231
4232 let sum_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Sum)]);
4234 assert_eq!(sum_schema.arity(), 2);
4235 assert_eq!(sum_schema.column_type(1), Some(ScalarType::U64));
4236
4237 let min_schema = provider.groupby_multi_agg_result_schema(&input, &[0], &[(1, AggOp::Min)]);
4239 assert_eq!(min_schema.arity(), 2);
4240 assert_eq!(min_schema.column_type(1), Some(ScalarType::U32));
4241 }
4242
4243 #[test]
4244 fn test_groupby_multi_agg_sum_returns_u64_schema() {
4245 let provider = match create_test_provider() {
4246 Some(p) => p,
4247 None => {
4248 eprintln!("Skipping test: no CUDA device");
4249 return;
4250 }
4251 };
4252
4253 let schema = Schema::new(vec![
4254 ("key".to_string(), ScalarType::U32),
4255 ("val".to_string(), ScalarType::U32),
4256 ]);
4257
4258 let result_schema =
4259 provider.groupby_multi_agg_result_schema(&schema, &[0], &[(1, AggOp::Sum)]);
4260
4261 assert_eq!(
4263 result_schema.column_type(1),
4264 Some(ScalarType::U64),
4265 "Sum aggregation should return U64 type, not U32"
4266 );
4267 }
4268
4269 #[test]
4270 fn test_join_custom_max_output() {
4271 let provider = match create_test_provider() {
4272 Some(p) => p,
4273 None => {
4274 eprintln!("Skipping test: no CUDA device available");
4275 return;
4276 }
4277 };
4278
4279 let left = create_test_buffer(&provider, &[1, 1, 1, 1, 2, 2, 2, 2], "left_key");
4284 let right = create_test_buffer(&provider, &[1, 1, 1, 2, 2, 2], "right_key");
4285
4286 let result_limited = provider
4288 .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, Some(10))
4289 .expect("join with limit should succeed");
4290 assert!(
4291 result_limited.num_rows() <= 10,
4292 "With limit 10, got {} rows but expected at most 10",
4293 result_limited.num_rows()
4294 );
4295
4296 let result_unlimited = provider
4298 .hash_join_v2_with_limit(&left, &right, &[0], &[0], JoinType::Inner, None)
4299 .expect("join without limit should succeed");
4300 assert_eq!(
4301 result_unlimited.num_rows(),
4302 24,
4303 "Without limit, expected 24 rows but got {}",
4304 result_unlimited.num_rows()
4305 );
4306
4307 let result_legacy = provider
4309 .hash_join_v2(&left, &right, &[0], &[0], JoinType::Inner)
4310 .expect("legacy hash_join_v2 should succeed");
4311 assert_eq!(
4312 result_legacy.num_rows(),
4313 24,
4314 "Legacy API without limit, expected 24 rows but got {}",
4315 result_legacy.num_rows()
4316 );
4317 }
4318
4319 fn create_arith_test_provider() -> Option<CudaKernelProvider> {
4323 if !has_cuda_device() {
4324 return None;
4325 }
4326 let device = Arc::new(CudaDevice::new(0).ok()?);
4327 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
4328 let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
4329 CudaKernelProvider::new(device, memory).ok()
4330 }
4331
4332 fn create_i64_buffer(provider: &CudaKernelProvider, data: &[i64]) -> CudaBuffer {
4334 let schema = Schema::new(vec![("col".to_string(), ScalarType::I64)]);
4335 provider
4336 .create_buffer_from_slice::<i64>(data, schema)
4337 .unwrap()
4338 }
4339
4340 fn create_f64_buffer(provider: &CudaKernelProvider, data: &[f64]) -> CudaBuffer {
4342 let schema = Schema::new(vec![("col".to_string(), ScalarType::F64)]);
4343 provider
4344 .create_buffer_from_slice::<f64>(data, schema)
4345 .unwrap()
4346 }
4347
4348 #[test]
4349 fn test_add_columns_i64() {
4350 let Some(provider) = create_arith_test_provider() else {
4351 eprintln!("Skipping test: no CUDA device available");
4352 return;
4353 };
4354
4355 let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4356 let b = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4357
4358 let result = provider.add_columns(&a, &b).unwrap();
4359 let values = provider.download_column::<i64>(&result, 0).unwrap();
4360
4361 assert_eq!(values, vec![11, 22, 33, 44, 55]);
4362 }
4363
4364 #[test]
4365 fn test_sub_columns_i64() {
4366 let Some(provider) = create_arith_test_provider() else {
4367 eprintln!("Skipping test: no CUDA device available");
4368 return;
4369 };
4370
4371 let a = create_i64_buffer(&provider, &[10, 20, 30, 40, 50]);
4372 let b = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4373
4374 let result = provider.sub_columns(&a, &b).unwrap();
4375 let values = provider.download_column::<i64>(&result, 0).unwrap();
4376
4377 assert_eq!(values, vec![9, 18, 27, 36, 45]);
4378 }
4379
4380 #[test]
4381 fn test_mul_columns_i64() {
4382 let Some(provider) = create_arith_test_provider() else {
4383 eprintln!("Skipping test: no CUDA device available");
4384 return;
4385 };
4386
4387 let a = create_i64_buffer(&provider, &[2, 3, 4, 5, 6]);
4388 let b = create_i64_buffer(&provider, &[3, 4, 5, 6, 7]);
4389
4390 let result = provider.mul_columns(&a, &b).unwrap();
4391 let values = provider.download_column::<i64>(&result, 0).unwrap();
4392
4393 assert_eq!(values, vec![6, 12, 20, 30, 42]);
4394 }
4395
4396 #[test]
4397 fn test_div_columns_i64() {
4398 let Some(provider) = create_arith_test_provider() else {
4399 eprintln!("Skipping test: no CUDA device available");
4400 return;
4401 };
4402
4403 let a = create_i64_buffer(&provider, &[100, 200, 300, 400]);
4404 let b = create_i64_buffer(&provider, &[10, 20, 30, 40]);
4405
4406 let result = provider.div_columns(&a, &b).unwrap();
4407 let values = provider.download_column::<i64>(&result, 0).unwrap();
4408
4409 assert_eq!(values, vec![10, 10, 10, 10]);
4410 }
4411
4412 #[test]
4413 fn test_div_columns_by_zero() {
4414 let Some(provider) = create_arith_test_provider() else {
4415 eprintln!("Skipping test: no CUDA device available");
4416 return;
4417 };
4418
4419 let a = create_i64_buffer(&provider, &[10, 20, 30]);
4420 let b = create_i64_buffer(&provider, &[2, 0, 3]); let result = provider.div_columns(&a, &b).unwrap();
4423 let values = provider.download_column::<i64>(&result, 0).unwrap();
4424
4425 assert_eq!(values, vec![5, i64::MAX, 10]);
4427 }
4428
4429 #[test]
4430 fn test_mod_columns_i64() {
4431 let Some(provider) = create_arith_test_provider() else {
4432 eprintln!("Skipping test: no CUDA device available");
4433 return;
4434 };
4435
4436 let a = create_i64_buffer(&provider, &[17, 23, 100, 7]);
4437 let b = create_i64_buffer(&provider, &[5, 7, 30, 3]);
4438
4439 let result = provider.mod_columns(&a, &b).unwrap();
4440 let values = provider.download_column::<i64>(&result, 0).unwrap();
4441
4442 assert_eq!(values, vec![2, 2, 10, 1]);
4443 }
4444
4445 #[test]
4446 fn test_mod_columns_by_zero() {
4447 let Some(provider) = create_arith_test_provider() else {
4448 eprintln!("Skipping test: no CUDA device available");
4449 return;
4450 };
4451
4452 let a = create_i64_buffer(&provider, &[10, 20]);
4453 let b = create_i64_buffer(&provider, &[3, 0]); let result = provider.mod_columns(&a, &b).unwrap();
4456 let values = provider.download_column::<i64>(&result, 0).unwrap();
4457
4458 assert_eq!(values, vec![1, 0]);
4460 }
4461
4462 #[test]
4463 fn test_abs_column_i64() {
4464 let Some(provider) = create_arith_test_provider() else {
4465 eprintln!("Skipping test: no CUDA device available");
4466 return;
4467 };
4468
4469 let a = create_i64_buffer(&provider, &[-5, 10, -15, 20, 0]);
4470
4471 let result = provider.abs_column(&a).unwrap();
4472 let values = provider.download_column::<i64>(&result, 0).unwrap();
4473
4474 assert_eq!(values, vec![5, 10, 15, 20, 0]);
4475 }
4476
4477 #[test]
4478 fn test_min_columns_i64() {
4479 let Some(provider) = create_arith_test_provider() else {
4480 eprintln!("Skipping test: no CUDA device available");
4481 return;
4482 };
4483
4484 let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4485 let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4486
4487 let result = provider.min_columns(&a, &b).unwrap();
4488 let values = provider.download_column::<i64>(&result, 0).unwrap();
4489
4490 assert_eq!(values, vec![3, 10, 10, 20]);
4491 }
4492
4493 #[test]
4494 fn test_max_columns_i64() {
4495 let Some(provider) = create_arith_test_provider() else {
4496 eprintln!("Skipping test: no CUDA device available");
4497 return;
4498 };
4499
4500 let a = create_i64_buffer(&provider, &[5, 10, 15, 20]);
4501 let b = create_i64_buffer(&provider, &[3, 12, 10, 25]);
4502
4503 let result = provider.max_columns(&a, &b).unwrap();
4504 let values = provider.download_column::<i64>(&result, 0).unwrap();
4505
4506 assert_eq!(values, vec![5, 12, 15, 25]);
4507 }
4508
4509 #[test]
4510 fn test_add_columns_f64() {
4511 let Some(provider) = create_arith_test_provider() else {
4512 eprintln!("Skipping test: no CUDA device available");
4513 return;
4514 };
4515
4516 let a = create_f64_buffer(&provider, &[1.5, 2.5, 3.5]);
4517 let b = create_f64_buffer(&provider, &[0.5, 1.5, 2.5]);
4518
4519 let result = provider.add_columns(&a, &b).unwrap();
4520 let values = provider.download_column::<f64>(&result, 0).unwrap();
4521
4522 assert_eq!(values, vec![2.0, 4.0, 6.0]);
4523 }
4524
4525 #[test]
4526 fn test_mul_columns_f64() {
4527 let Some(provider) = create_arith_test_provider() else {
4528 eprintln!("Skipping test: no CUDA device available");
4529 return;
4530 };
4531
4532 let a = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4533 let b = create_f64_buffer(&provider, &[1.5, 2.0, 2.5]);
4534
4535 let result = provider.mul_columns(&a, &b).unwrap();
4536 let values = provider.download_column::<f64>(&result, 0).unwrap();
4537
4538 assert_eq!(values, vec![3.0, 6.0, 10.0]);
4539 }
4540
4541 #[test]
4542 fn test_div_columns_f64_by_zero() {
4543 let Some(provider) = create_arith_test_provider() else {
4544 eprintln!("Skipping test: no CUDA device available");
4545 return;
4546 };
4547
4548 let a = create_f64_buffer(&provider, &[1.0, -1.0, 0.0]);
4549 let b = create_f64_buffer(&provider, &[0.0, 0.0, 0.0]);
4550
4551 let result = provider.div_columns(&a, &b).unwrap();
4552 let values = provider.download_column::<f64>(&result, 0).unwrap();
4553
4554 assert!(values[0].is_infinite() && values[0].is_sign_positive());
4556 assert!(values[1].is_infinite() && values[1].is_sign_negative());
4557 assert!(values[2].is_nan());
4558 }
4559
4560 #[test]
4561 fn test_pow_columns() {
4562 let Some(provider) = create_arith_test_provider() else {
4563 eprintln!("Skipping test: no CUDA device available");
4564 return;
4565 };
4566
4567 let base = create_i64_buffer(&provider, &[2, 3, 4, 5]);
4568 let exp = create_i64_buffer(&provider, &[3, 2, 2, 1]);
4569
4570 let result = provider.pow_columns(&base, &exp).unwrap();
4571 let values = provider.download_column::<f64>(&result, 0).unwrap();
4572
4573 assert_eq!(values, vec![8.0, 9.0, 16.0, 5.0]);
4575 }
4576
4577 #[test]
4578 fn test_pow_columns_fractional_exp() {
4579 let Some(provider) = create_arith_test_provider() else {
4580 eprintln!("Skipping test: no CUDA device available");
4581 return;
4582 };
4583
4584 let base = create_f64_buffer(&provider, &[4.0, 9.0, 27.0]);
4585 let exp = create_f64_buffer(&provider, &[0.5, 0.5, 1.0 / 3.0]);
4586
4587 let result = provider.pow_columns(&base, &exp).unwrap();
4588 let values = provider.download_column::<f64>(&result, 0).unwrap();
4589
4590 assert!((values[0] - 2.0).abs() < 1e-10);
4592 assert!((values[1] - 3.0).abs() < 1e-10);
4593 assert!((values[2] - 3.0).abs() < 1e-10);
4594 }
4595
4596 #[test]
4597 fn test_cast_i64_to_f64() {
4598 let Some(provider) = create_arith_test_provider() else {
4599 eprintln!("Skipping test: no CUDA device available");
4600 return;
4601 };
4602
4603 let a = create_i64_buffer(&provider, &[1, 2, 3, 4, 5]);
4604
4605 let result = provider.cast_column(&a, ScalarType::F64).unwrap();
4606 let values = provider.download_column::<f64>(&result, 0).unwrap();
4607
4608 assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
4609 }
4610
4611 #[test]
4612 fn test_cast_f64_to_i64() {
4613 let Some(provider) = create_arith_test_provider() else {
4614 eprintln!("Skipping test: no CUDA device available");
4615 return;
4616 };
4617
4618 let a = create_f64_buffer(&provider, &[1.9, 2.1, 3.5, 4.0, 5.7]);
4619
4620 let result = provider.cast_column(&a, ScalarType::I64).unwrap();
4621 let values = provider.download_column::<i64>(&result, 0).unwrap();
4622
4623 assert_eq!(values, vec![1, 2, 3, 4, 5]);
4625 }
4626
4627 #[test]
4628 fn test_cast_i64_to_i32() {
4629 let Some(provider) = create_arith_test_provider() else {
4630 eprintln!("Skipping test: no CUDA device available");
4631 return;
4632 };
4633
4634 let a = create_i64_buffer(&provider, &[1, 2, 3, 100, 200]);
4635
4636 let result = provider.cast_column(&a, ScalarType::I32).unwrap();
4637 let values = provider.download_column::<i32>(&result, 0).unwrap();
4638
4639 assert_eq!(values, vec![1, 2, 3, 100, 200]);
4640 }
4641
4642 #[test]
4643 fn test_arithmetic_row_count_mismatch() {
4644 let Some(provider) = create_arith_test_provider() else {
4645 eprintln!("Skipping test: no CUDA device available");
4646 return;
4647 };
4648
4649 let a = create_i64_buffer(&provider, &[1, 2, 3]);
4650 let b = create_i64_buffer(&provider, &[1, 2]); let result = provider.add_columns(&a, &b);
4653 assert!(result.is_err());
4654 let err = result.err().unwrap();
4655 assert!(err.to_string().contains("Row count mismatch"));
4656 }
4657
4658 #[test]
4659 fn test_arithmetic_empty_buffers() {
4660 let Some(provider) = create_arith_test_provider() else {
4661 eprintln!("Skipping test: no CUDA device available");
4662 return;
4663 };
4664
4665 let a = create_i64_buffer(&provider, &[]);
4666 let b = create_i64_buffer(&provider, &[]);
4667
4668 let result = provider.add_columns(&a, &b).unwrap();
4669 let values = provider.download_column::<i64>(&result, 0).unwrap();
4670
4671 assert_eq!(values, Vec::<i64>::new());
4672 }
4673
4674 #[test]
4675 fn test_wrapping_arithmetic_overflow() {
4676 let Some(provider) = create_arith_test_provider() else {
4677 eprintln!("Skipping test: no CUDA device available");
4678 return;
4679 };
4680
4681 let a = create_i64_buffer(&provider, &[i64::MAX, i64::MIN]);
4682 let b = create_i64_buffer(&provider, &[1, -1]);
4683
4684 let add_result = provider.add_columns(&a, &b).unwrap();
4686 let add_values = provider.download_column::<i64>(&add_result, 0).unwrap();
4687 assert_eq!(add_values[0], i64::MIN); assert_eq!(add_values[1], i64::MAX); }
4690
4691 #[test]
4692 fn test_abs_column_f64() {
4693 let Some(provider) = create_arith_test_provider() else {
4694 eprintln!("Skipping test: no CUDA device available");
4695 return;
4696 };
4697
4698 let a = create_f64_buffer(&provider, &[-1.5, 2.5, -3.5, 0.0]);
4699
4700 let result = provider.abs_column(&a).unwrap();
4701 let values = provider.download_column::<f64>(&result, 0).unwrap();
4702
4703 assert_eq!(values, vec![1.5, 2.5, 3.5, 0.0]);
4704 }
4705
4706 #[test]
4707 fn test_min_max_columns_f64() {
4708 let Some(provider) = create_arith_test_provider() else {
4709 eprintln!("Skipping test: no CUDA device available");
4710 return;
4711 };
4712
4713 let a = create_f64_buffer(&provider, &[1.5, 5.0, 3.0]);
4714 let b = create_f64_buffer(&provider, &[2.0, 3.0, 4.0]);
4715
4716 let min_result = provider.min_columns(&a, &b).unwrap();
4717 let min_values = provider.download_column::<f64>(&min_result, 0).unwrap();
4718 assert_eq!(min_values, vec![1.5, 3.0, 3.0]);
4719
4720 let max_result = provider.max_columns(&a, &b).unwrap();
4721 let max_values = provider.download_column::<f64>(&max_result, 0).unwrap();
4722 assert_eq!(max_values, vec![2.0, 5.0, 4.0]);
4723 }
4724}