Skip to main content

xlog_cuda/provider/
mod.rs

1//! CUDA kernel provider implementation
2//!
3//! This module provides the `CudaKernelProvider` which manages pre-compiled
4//! PTX kernels for GPU execution of relational operations (join, dedup, groupby).
5
6use 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/// Per-module PTX load timing (populated only when XLOG_WARMUP_PROFILE=1).
55#[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
69/// Detect device compute capability as a two-digit number (e.g. 75, 80, 120).
70pub(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        // Skip any staged cubin/PTX whose bytes diverge from what this binary
106        // was built against. A stale staged artifact (kernel signature changed
107        // but the staged copy was never refreshed) otherwise loads "fine" and
108        // then launches a mismatched kernel into an illegal address.
109        .filter(|(path, _)| !staged_artifact_is_stale(path))
110        .map(|(path, is_cubin)| KernelModuleSource::File { path, is_cubin })
111        .collect();
112
113    // ALWAYS append the embedded portable PTX as the final fallback. It is
114    // compiled into this binary, so it can never be stale relative to the launch
115    // sites — it guarantees a signature-correct kernel even when every staged
116    // File artifact was skipped as stale or fails to load. (Previously this was
117    // suppressed whenever any portable-PTX *file* existed, which let a stale
118    // staged PTX shadow the fresh embedded one.)
119    if let Some(ptx) = crate::embedded_kernel_data::portable_ptx(name) {
120        sources.push(KernelModuleSource::EmbeddedPortablePtx { ptx });
121    }
122    sources
123}
124
125/// A staged cubin/PTX is "stale" when this binary embeds a canonical integrity
126/// hash for that artifact file name and the on-disk bytes do not match it — the
127/// staged artifact diverges from what this build produced. Loading such an
128/// artifact can launch a mismatched kernel into an illegal address, so it is
129/// skipped in favor of a fresh source. Artifacts with no embedded canonical
130/// hash (e.g. an arch this build did not produce) are NOT treated as stale — we
131/// can only validate what we built — nor are unreadable files (the loader
132/// surfaces the IO error).
133fn 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
146/// FNV-1a 64-bit, matching the build-time hash in `crates/xlog-cuda/build.rs`.
147fn 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        // Use a name this build does NOT produce, so neither file carries a
177        // canonical integrity hash (the staleness skip is exercised separately).
178        // This isolates the file-resolution precedence: cubin first, then the
179        // portable-PTX file as fallback.
180        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    // Locks the FNV-1a contract between build.rs (which embeds canonical
208    // artifact hashes) and the runtime (which re-hashes staged artifacts). If
209    // these two implementations ever diverge, every staged artifact would read
210    // as "stale" — so these canonical FNV-1a 64-bit vectors must hold.
211    #[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    // A file whose name this build did not produce has no canonical hash, so it
219    // is conservatively NOT treated as stale (we only validate what we built).
220    // A nonexistent path is likewise not "stale" — the loader surfaces IO.
221    #[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
242/// Resolve a kernel module from sidecar artifacts or embedded portable PTX.
243///
244/// Asserts (in debug builds) that `name` is present in the kernel manifest,
245/// catching name/order drift between the manifest and provider load blocks.
246pub(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/// Typed kernel-parameter view over a borrowed device allocation.
263///
264/// The view preserves the pointer, element count, stream, and Rust borrow
265/// lifetime, but it does not carry runtime allocation identity. Recorded
266/// launchers must register the owning slice or column directly with
267/// [`crate::launch::LaunchRecorder`] before preflight.
268#[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
276/// Preallocated scratch layout for graph-capturable u32 multi-block scans.
277///
278/// The legacy stream-aware scan helper allocates recursive `block_sums`
279/// buffers inside the helper. CUDA Graph capture records concrete allocation
280/// addresses, so bounded CSM CUDA Graph replay needs the scan topology and
281/// scratch buffers to be fixed before capture begins.
282pub(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
357/// Scratch buffers for stable radix sorting of u32 key/value pairs.
358pub 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
397/// Module names for loaded PTX modules
398pub 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
426// Compile-time check: kernel manifest lists exactly 30 modules.
427const _: () = assert!(crate::kernel_manifest_data::KERNEL_CU_NAMES.len() == 30);
428
429/// Kernel function names in the GPU WCOJ module.
430pub 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    // General-arity clique kernels (k=5..8 from a single template).
482    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    // Free Join frontier engine primitives. The work
503    // prefix kernel is width-agnostic (ranges are u32 row indices in
504    // every width class); count/emit/probe have u64 data twins.
505    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    // D3 S3 spike — factorized recursive delta novel-set pipeline.
514    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
527/// Kernel function names in the Monte Carlo sampling module
528pub mod mc_sample_kernels {
529    pub const MC_SAMPLE_BERNOULLI: &str = "mc_sample_bernoulli";
530}
531
532/// Kernel function names in the Monte Carlo evaluation module
533pub 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
540/// Kernel function names in the GPU-resident Datalog/MC engine module.
541pub mod mc_resident_kernels {
542    /// Single megakernel: evaluates all MC worlds to fixpoint and counts
543    /// query/evidence satisfaction with zero host interaction in-region.
544    pub const MC_RESIDENT_ENGINE: &str = "mc_resident_engine";
545}
546
547/// Kernel function names in the arithmetic module
548pub 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    // Conditional select kernels
569    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
577/// Kernel function names in the epistemic module.
578pub mod epistemic_kernels {
579    /// Device-side epistemic candidate-assumption generator.
580    pub const EPISTEMIC_GENERATE_CANDIDATE_ASSUMPTIONS_U8: &str =
581        "epistemic_generate_candidate_assumptions_u8";
582    /// Device-side epistemic candidate propagation staging kernel.
583    pub const EPISTEMIC_PROPAGATE_CANDIDATES_U8: &str = "epistemic_propagate_candidates_u8";
584    /// Device-side epistemic candidate bit validation kernel.
585    pub const EPISTEMIC_VALIDATE_CANDIDATE_BITS_U8: &str = "epistemic_validate_candidate_bits_u8";
586    /// Device-side model-membership staging kernel.
587    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_U8: &str =
588        "epistemic_populate_model_membership_u8";
589    /// Device-side tuple-source-backed model-membership kernel.
590    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_U8: &str =
591        "epistemic_populate_model_membership_from_tuple_source_u8";
592    /// Device-side arity-one tuple-key-backed model-membership kernel.
593    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY1_U8: &str =
594        "epistemic_populate_model_membership_from_tuple_source_arity1_u8";
595    /// Device-side arity-two tuple-key-backed model-membership kernel.
596    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY2_U8: &str =
597        "epistemic_populate_model_membership_from_tuple_source_arity2_u8";
598    /// Device-side arity-three tuple-key-backed model-membership kernel.
599    pub const EPISTEMIC_POPULATE_MODEL_MEMBERSHIP_FROM_TUPLE_SOURCE_ARITY3_U8: &str =
600        "epistemic_populate_model_membership_from_tuple_source_arity3_u8";
601    /// Device-side generic-arity tuple-key-backed model-membership kernel.
602    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    /// Device-side world-view validation kernel.
605    pub const EPISTEMIC_VALIDATE_WORLD_VIEWS_U8: &str = "epistemic_validate_world_views_u8";
606    /// Device-side world-view integrity-constraint validation kernel.
607    pub const EPISTEMIC_VALIDATE_CONSTRAINTS_U8: &str = "epistemic_validate_constraints_u8";
608    /// Device-side accepted-candidate materialization staging kernel.
609    pub const EPISTEMIC_MATERIALIZE_ACCEPTED_CANDIDATES_U8: &str =
610        "epistemic_materialize_accepted_candidates_u8";
611
612    /// Device-side final-result flag materialization staging kernel.
613    pub const EPISTEMIC_MATERIALIZE_FINAL_RESULT_FLAGS_U8: &str =
614        "epistemic_materialize_final_result_flags_u8";
615    /// Device-side final tuple materialization kernel.
616    pub const EPISTEMIC_MATERIALIZE_FINAL_TUPLE_COLUMN_U8: &str =
617        "epistemic_materialize_final_tuple_column_u8";
618    /// Device-side final tuple row-map kernel.
619    pub const EPISTEMIC_BUILD_FINAL_TUPLE_ROW_MAP_U8: &str =
620        "epistemic_build_final_tuple_row_map_u8";
621    /// Device-side final tuple rejection-close kernel.
622    pub const EPISTEMIC_CLOSE_FINAL_TUPLE_REJECTIONS_U8: &str =
623        "epistemic_close_final_tuple_rejections_u8";
624}
625
626/// Kernel function names in the neural fast-path module.
627pub 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
632/// Kernel function names in the ILP module.
633pub 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
650/// Kernel function names in the ILP credit module.
651pub 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
659/// Kernel function names in the native bounded exact-induction module.
660pub 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
668/// Kernel function names in the n-ary exact-induction module.
669pub mod ilp_exact_nary_kernels {
670    pub const ILP_EXACT_NARY_SCORE: &str = "ilp_exact_nary_score";
671}
672
673/// Kernel function names in the PIR interning module.
674pub 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
693/// Kernel function names in the GPU CNF encoder module.
694pub 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
710/// Kernel function names in the weights module.
711pub 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
734/// Kernel function names in the GPU Decision-DNNF compiler module
735/// (CNF validation + circuit levelization).
736pub 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    // BFS frontier expansion and unit propagation.
741    pub const D4_FRONTIER_PREPARE: &str = "d4_frontier_prepare";
742    pub const D4_FRONTIER_EXPAND: &str = "d4_frontier_expand";
743    // Per-frontier Decision-DNNF DFS worker (count+emit).
744    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    // GPU smoothing with random-variable support and wrapper emission.
748    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    // GPU free-variable mask for variables in clauses versus the circuit.
758    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    // GPU-only assertions (tests + invariant enforcement without host reads).
762    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
767/// Kernel function names in the join module
768pub mod join_kernels {
769    pub const HASH_JOIN_BUILD: &str = "hash_join_build";
770    pub const HASH_JOIN_PROBE: &str = "hash_join_probe";
771    // V2 kernels for multi-column joins
772    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    /// Nested-loop inner join (emit-pairs design). Reads
784    /// the single key column from each side; emits matched
785    /// `(left_idx, right_idx)` pairs as two parallel u32 arrays.
786    /// Payload columns are materialized after the kernel via
787    /// `gather_buffer_by_indices` in the provider fn.
788    pub const NESTED_LOOP_JOIN_INNER_U32_1KEY_PAIRS: &str = "nested_loop_join_inner_u32_1key_pairs";
789    /// Sort-merge inner join (emit-pairs design,
790    /// caller-asserted pre-sorted inputs). Reads the single
791    /// key column from each side, performs per-thread binary
792    /// search on the right side to find matched-key runs,
793    /// emits `(left_idx, right_idx)` pairs as two parallel
794    /// u32 arrays. Payload columns materialize after the
795    /// kernel via `gather_buffer_by_indices`.
796    pub const SORT_MERGE_JOIN_INNER_U32_1KEY_PAIRS: &str = "sort_merge_join_inner_u32_1key_pairs";
797}
798
799/// Kernel function names in the dedup module
800pub 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
810/// Kernel function names in the groupby module
811pub 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
830/// Kernel function names in the scan module
831pub 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    // Multi-block scan kernels for large prefix sums
837    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
843/// Kernel function names in the sort module
844pub 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    /// Sort-merge sortedness-detection kernel — single-pass adjacent-
867    /// pair check; atomically writes 0 to a u32 flag on
868    /// `keys[i] > keys[i+1]`. Caller initializes flag to 1
869    /// before launch, reads result post-launch. Used by the
870    /// dispatch-site eligibility check at `execute_join` to
871    /// validate caller-asserted sortedness before invoking
872    /// `sort_merge_join_v2_inner_u32_1key`.
873    pub const CHECK_ASCENDING_SORTED_U32: &str = "check_ascending_sorted_u32";
874}
875
876/// Kernel function names in the filter module
877pub 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
911/// Kernel function names in the set_ops module
912pub 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
918/// Kernel function names in the pack module (GPU-side key packing)
919pub mod pack_kernels {
920    /// Pack multiple columns into row-major byte array
921    pub const PACK_KEYS: &str = "pack_keys";
922    /// Compute FNV-1a hash from packed keys
923    pub const HASH_PACKED_KEYS: &str = "hash_packed_keys";
924    /// Fused pack + hash in single pass (optimal for join key preparation)
925    pub const PACK_AND_HASH_KEYS: &str = "pack_and_hash_keys";
926    /// Fused pack + hash for arbitrary key column counts
927    pub const PACK_AND_HASH_KEYS_GENERIC: &str = "pack_and_hash_keys_generic";
928    /// Vectorized pack for 8-byte aligned columns
929    pub const PACK_KEYS_ALIGNED: &str = "pack_keys_aligned";
930    /// Unpack single column from packed row data
931    pub const UNPACK_COLUMN: &str = "unpack_column";
932    /// Unpack single column with device-resident row count
933    pub const UNPACK_COLUMN_COUNTED: &str = "unpack_column_counted";
934    /// Gather rows from packed data based on index array
935    pub const GATHER_PACKED_ROWS: &str = "gather_packed_rows";
936    /// Gather rows with device-resident row count
937    pub const GATHER_PACKED_ROWS_COUNTED: &str = "gather_packed_rows_counted";
938    /// Scatter write: distribute packed rows to non-contiguous output positions
939    pub const SCATTER_PACKED_ROWS: &str = "scatter_packed_rows";
940    /// Compare packed keys for equality
941    pub const COMPARE_PACKED_KEYS: &str = "compare_packed_keys";
942    /// Pack u8 bools into Arrow bitmap bytes
943    pub const PACK_BOOLS_TO_BITMAP: &str = "pack_bools_to_bitmap";
944}
945
946/// Kernel function names in the circuit module
947pub 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
974/// Kernel function names in the cache module
975pub 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
986/// Kernel function names in the SAT module
987pub 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
1006/// Default maximum output size for join operations.
1007/// This prevents memory overflow when joining large tables with high cardinality matches.
1008pub const DEFAULT_JOIN_MAX_OUTPUT: usize = 1_000_000;
1009
1010/// Nested-loop join eligibility threshold (Cartesian product
1011/// upper bound). The dispatcher routes to nested-loop iff
1012/// `num_left * num_right <= NESTED_LOOP_TOTAL_THRESHOLD`; the
1013/// provider validates the same invariant fail-closed before any
1014/// allocation.
1015///
1016/// This is the **single source of truth** for the threshold.
1017/// `xlog-runtime`'s dispatch site imports this constant; do NOT
1018/// redeclare in xlog-runtime (would create either drift risk or
1019/// a reverse `xlog-cuda → xlog-runtime` dep cycle).
1020///
1021/// Value (`4_000_000`) is grounded in the bench-spike at
1022/// `bench-spike/w42-nested-loop` HEAD `9c0cefc6` (see
1023/// `docs/evidence/2026-05-07-w42-bench-spike/README.md`):
1024/// largest symmetric tested cell `L=R=2000` → 4M total wins by
1025/// 5.41× over hash; the algorithmic crossover is extrapolated to
1026/// ~10000×10000 = 100M; 4M leaves 6× margin to absorb
1027/// production-kernel cost asymmetry. The threshold also caps the
1028/// index-array allocation at 32 MB total (4M × 4 bytes × 2
1029/// arrays).
1030pub const NESTED_LOOP_TOTAL_THRESHOLD: u64 = 4_000_000;
1031
1032/// Comparison operators for filtering
1033#[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/// Join types for hash_join_v2
1045#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1046pub enum JoinType {
1047    /// Inner join: return rows where keys match on both sides
1048    Inner,
1049    /// Semi join: return left rows that have any match in right (no right columns)
1050    Semi,
1051    /// Anti join: return left rows that have NO match in right
1052    Anti,
1053    /// Left outer join: return all left rows, with nulls for non-matching right
1054    LeftOuter,
1055}
1056
1057/// Result of packing key columns and computing hashes for join operations
1058struct PackedKeyData {
1059    /// Computed hash values (one per row)
1060    hashes: crate::memory::TrackedCudaSlice<u64>,
1061    /// Packed key data in row-major format
1062    packed_keys: crate::memory::TrackedCudaSlice<u8>,
1063    /// Total bytes per row (key stride)
1064    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
1075/// Bucketed hash table for u64 hashes.
1076pub 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
1084/// Cached build-side join index for v2 hash join.
1085///
1086/// This captures the packed key bytes and bucketed hash table layout for the build (right) side,
1087/// enabling reuse across repeated joins on the same relation + key columns.
1088pub 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    /// Key columns (indices) this index was built for.
1098    pub fn right_keys(&self) -> &[usize] {
1099        &self.right_keys
1100    }
1101
1102    /// Row count of the build-side buffer at index build time.
1103    pub fn right_num_rows(&self) -> u32 {
1104        self.right_num_rows
1105    }
1106
1107    /// Approximate device memory used by this cached index.
1108    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
1119/// CUDA kernel provider for xlog GPU operations
1120///
1121/// Manages pre-compiled PTX modules for relational operations:
1122/// - **Join**: Hash join with build/probe phases
1123/// - **Dedup**: Sort-based deduplication with prefix-sum compaction
1124/// - **GroupBy**: Sorted-input group aggregation (count, sum, min, max)
1125///
1126/// PTX modules are loaded at construction time and stored in the CUDA device.
1127/// Kernel functions can be retrieved using `device.get_func()`.
1128///
1129/// # Example
1130/// ```ignore
1131/// use std::sync::Arc;
1132/// use xlog_cuda::{CudaDevice, GpuMemoryManager, CudaKernelProvider};
1133/// use xlog_core::MemoryBudget;
1134///
1135/// let device = Arc::new(CudaDevice::new(0)?);
1136/// let memory = Arc::new(GpuMemoryManager::new(device.clone(), MemoryBudget::default()));
1137/// let provider = CudaKernelProvider::new(device, memory)?;
1138/// ```
1139pub struct CudaKernelProvider {
1140    /// Process-local identity retained by prepared resident schedules.
1141    provider_identity: u64,
1142    /// The CUDA device with loaded PTX modules
1143    device: Arc<CudaDevice>,
1144    /// GPU memory manager for kernel allocations
1145    memory: Arc<GpuMemoryManager>,
1146    /// Tracked host transfers for diagnostics
1147    transfer_tracker: HostTransferTracker,
1148    /// Transfers for the one bounded, post-synchronization resident receipt.
1149    final_observation_transfer_tracker: HostTransferTracker,
1150    /// Number of final resident receipts copied into page-locked host memory.
1151    final_observation_pinned_receipts: AtomicU64,
1152    /// PTX load profiling data (populated only when XLOG_WARMUP_PROFILE=1)
1153    ptx_load_profile: Option<PtxLoadProfile>,
1154    /// Column-level D2H transfer counter (incremented by each download_column_* call)
1155    d2h_transfer_count: AtomicU64,
1156    /// Untracked control-plane metadata D2H read counter. Incremented by every
1157    /// `dtoh_scalar_untracked` / `dtoh_small_metadata_untracked` call. These are
1158    /// bounded metadata reads (row counts, scan totals) exempt from the
1159    /// data-plane transfer contract, but the GPU-resident MC engine's no-host
1160    /// gate must prove they are *also* zero inside the measured region — hence an
1161    /// explicit, resettable counter.
1162    untracked_metadata_dtoh_count: AtomicU64,
1163    /// Strict deterministic-Datalog D2H gate. When `true`, any data-plane D2H
1164    /// transfer (column downloads or `dtoh_sync_copy_into_tracked`) increments
1165    /// the violation counter and returns `XlogError::Execution` from the
1166    /// originating call. Metadata reads via `dtoh_scalar_untracked` are NOT
1167    /// gated. See [`CudaKernelProvider::enable_strict_deterministic_d2h`].
1168    strict_deterministic_d2h: AtomicBool,
1169    /// Cumulative count of deterministic-D2H gate violations observed since
1170    /// the last reset. Increments even on the failing path (the originating
1171    /// call still returns `Err`); kept for telemetry and tests.
1172    deterministic_d2h_violations: AtomicU64,
1173    /// Lazy-initialized non-default launch stream used by
1174    /// env-gated recorded-operator dispatch (filter, sort,
1175    /// dedup, GroupBy, hash-join). Cached for the provider's
1176    /// lifetime — the [`crate::device_runtime::StreamPool`]
1177    /// never returns streams to a free-list, so per-call
1178    /// acquire would saturate it. One stream per provider is
1179    /// sufficient because the recorder serializes work on it;
1180    /// multiple operations chain through commit-order events.
1181    recorded_op_stream: OnceLock<crate::device_runtime::StreamId>,
1182    /// Test/diagnostic-only counter for CSM (count-scan-materialize)
1183    /// invocations selected by the recorded hash-join dispatch.
1184    /// **Not part of any public stability guarantee** — its existence,
1185    /// shape, exposure, and increment semantics may change in any
1186    /// release. Used by the env-dispatch test suite to prove that CSM
1187    /// was actually selected for eligible Inner / LeftOuter cases (and
1188    /// not selected for Semi / Anti or when the env gate is off).
1189    csm_invocations: AtomicU64,
1190    /// Diagnostic counter for bounded CSM CUDA Graph captures.
1191    csm_cuda_graph_captures: AtomicU64,
1192    /// Diagnostic counter for bounded CSM CUDA Graph launches.
1193    csm_cuda_graph_launches: AtomicU64,
1194    /// Diagnostic counter for bounded CSM CUDA Graph ineligibility fallbacks.
1195    csm_cuda_graph_fallbacks: AtomicU64,
1196    /// Diagnostic counter for bounded CSM CUDA Graph cache replays.
1197    csm_cuda_graph_cache_hits: AtomicU64,
1198    /// Diagnostic counter for graph-mode small full-row set-maintenance
1199    /// sorts. This is test telemetry only; production correctness must not
1200    /// depend on the value.
1201    small_full_row_sort_invocations: AtomicU64,
1202    /// Bounded CSM CUDA Graph replay cache.
1203    csm_cuda_graph_cache: Mutex<HashMap<CsmCudaGraphKey, CsmCudaGraphEntry>>,
1204    /// Per-process counter of WCOJ layout fast-path hits. The
1205    /// fast-path skips `dedup_full_row_recorded` when the input
1206    /// is already strictly lex-sorted and full-row unique.
1207    /// Tests + the phase report binary read this counter to
1208    /// confirm the fast-path actually fired vs. silently fell
1209    /// through to the existing dedup pipeline.
1210    wcoj_layout_fast_path_hit_count: AtomicU64,
1211    /// Diagnostic counter for generic WCOJ layout-sort helper
1212    /// invocations. Used by K-clique dispatch-plan certifications to
1213    /// prove K-clique runtime dispatch no longer routes every edge
1214    /// through the old all-edge `wcoj_layout_sort_*_recorded` path.
1215    wcoj_layout_sort_invocation_count: AtomicU64,
1216    /// Diagnostic counter for K-clique leader-edge metadata builds.
1217    kclique_metadata_build_count: AtomicU64,
1218    /// Diagnostic counter for cumulative nanoseconds spent building K-clique
1219    /// leader-edge metadata.
1220    kclique_metadata_build_nanos: AtomicU64,
1221    /// Histogram-guided triangle WCOJ routing counter: successful dispatches
1222    /// accepted through the block-slice provider entry.
1223    wcoj_triangle_hg_dispatch_count: AtomicU64,
1224    /// Diagnostic-only: last WCOJ triangle dispatch's per-phase
1225    /// CUDA-event timings, populated by `wcoj_triangle_*_recorded`
1226    /// when the `wcoj-phase-timing` Cargo feature is on. Read by
1227    /// the `wcoj_phase_report` binary in xlog-integration. Field
1228    /// is absent when the feature is off, so production builds
1229    /// have zero overhead.
1230    #[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/// Separately accounted one-shot observation after resident execution ends.
1254#[derive(Debug, Clone, Copy, Default)]
1255pub struct FinalObservationTransferStats {
1256    /// Device-to-host bytes copied for terminal receipts.
1257    pub dtoh_bytes: u64,
1258    /// Device-to-host terminal receipt copy calls.
1259    pub dtoh_calls: u64,
1260    /// Receipt copies whose destination was page-locked host memory.
1261    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    /// Create a new CUDA kernel provider
1316    ///
1317    /// Loads all kernel modules into the CUDA device.
1318    /// Prefers cubin for the detected SM arch, falls back to portable PTX (sm_75+).
1319    ///
1320    /// # Arguments
1321    /// * `device` - The CUDA device to load modules into
1322    /// * `memory` - The GPU memory manager for kernel allocations
1323    ///
1324    /// # Errors
1325    /// Returns `XlogError::Kernel` if PTX loading fails
1326    ///
1327    /// # Example
1328    /// ```ignore
1329    /// let device = Arc::new(CudaDevice::new(0)?);
1330    /// let memory = Arc::new(GpuMemoryManager::new(device.clone(), MemoryBudget::default()));
1331    /// let provider = CudaKernelProvider::new(device, memory)?;
1332    /// ```
1333    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    /// Initialize provider-local state for a device whose kernel modules are
1341    /// already loaded. This must never perform module insertion or replacement.
1342    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    /// Construct a provider whose `GpuMemoryManager` must already
1378    /// have a v0.6 [`crate::device_runtime::XlogDeviceRuntime`]
1379    /// attached via [`GpuMemoryManager::with_runtime`].
1380    ///
1381    /// Equivalent to [`Self::new`] in every respect — same kernel
1382    /// loading, same field initialization — but **rejects** managers
1383    /// that lack a runtime. This guards against the misconfiguration
1384    /// in which a caller asks for runtime-routed provider semantics
1385    /// (by calling `with_runtime`) but supplies a legacy manager
1386    /// built via [`GpuMemoryManager::new`]; without the check, the
1387    /// resulting provider would silently keep using the cudarc
1388    /// default allocator and the runtime budget/logging stack would
1389    /// never observe the allocations the caller expected to be
1390    /// routed through it.
1391    ///
1392    /// Note: a runtime-routed manager passed to [`Self::new`] still
1393    /// routes correctly — `alloc::<T>` and `alloc_raw` consult
1394    /// `memory.runtime()` regardless of which provider constructor
1395    /// was used. `with_runtime` exists for callers that want the
1396    /// requirement enforced at construction time, not for
1397    /// correctness of the routing itself.
1398    ///
1399    /// This is the **opt-in** runtime entry point for providers.
1400    /// `Self::new` continues to accept managers without a runtime
1401    /// (the legacy default) and remains the production constructor
1402    /// until the runtime stack is certified end-to-end.
1403    ///
1404    /// # Errors
1405    /// Returns `XlogError::Kernel` if `memory.runtime()` is `None`,
1406    /// or anything `Self::new` would return.
1407    ///
1408    /// # Example
1409    /// ```ignore
1410    /// let device = Arc::new(CudaDevice::new(0)?);
1411    /// let runtime = Arc::new(XlogDeviceRuntime::with_resource(
1412    ///     Arc::clone(&device),
1413    ///     0,
1414    ///     Arc::new(StreamPool::with_defaults(Arc::clone(&device))),
1415    ///     Box::new(AsyncCudaResource::new(/* ... */)),
1416    /// ));
1417    /// let memory = Arc::new(GpuMemoryManager::with_runtime(
1418    ///     Arc::clone(&device),
1419    ///     MemoryBudget::default(),
1420    ///     runtime,
1421    /// ));
1422    /// let provider = CudaKernelProvider::with_runtime(device, memory)?;
1423    /// ```
1424    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    /// Create a provider-local runtime allocation view over this provider's
1436    /// already-loaded CUDA device.
1437    ///
1438    /// Unlike [`Self::with_runtime`], this method deliberately does not reload
1439    /// PTX modules. Replacing a module can invalidate functions or graphs held
1440    /// by the original provider. The supplied manager and attached runtime must
1441    /// therefore share this exact [`CudaDevice`] handle and ordinal.
1442    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    /// Internal: parse a "boolean" env var. Empty / unset / `"0"`
1489    /// → false; any other value → true.
1490    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    /// Whether the recorded filter dispatch is enabled via env.
1497    ///
1498    /// Returns `true` when either `XLOG_USE_RECORDED_FILTERS` or
1499    /// the umbrella `XLOG_USE_RECORDED_OPS` env var is set.
1500    /// Combined with a runtime-backed manager, this routes
1501    /// `filter::<T>` through the recorded launch path.
1502    ///
1503    /// Env-gated rather than default-on so the migration is
1504    /// opt-in for real callers; the existing legacy paths remain
1505    /// the production default until the runtime stack is
1506    /// certified end-to-end.
1507    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    /// Whether the recorded sort dispatch is enabled via env.
1512    /// Reads `XLOG_USE_RECORDED_SORT` or the umbrella
1513    /// `XLOG_USE_RECORDED_OPS`. The recorded-sort path is narrowed
1514    /// to U32 / Symbol keys only — the public
1515    /// `sort()` dispatcher checks both this env flag AND key
1516    /// type compatibility before routing.
1517    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    /// Whether the recorded full-row dedup dispatch is enabled
1522    /// via env. Reads `XLOG_USE_RECORDED_DEDUP` or the umbrella
1523    /// `XLOG_USE_RECORDED_OPS`. `dedup_full_row_recorded` is
1524    /// narrow to all-U32 / Symbol columns.
1525    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    /// Whether the recorded GroupBy dispatch is enabled via
1530    /// env. Reads `XLOG_USE_RECORDED_GROUPBY` or
1531    /// `XLOG_USE_RECORDED_OPS`. `groupby_multi_agg_recorded`
1532    /// supports U32 / Symbol keys + Count / Sum / Min / Max
1533    /// aggs only.
1534    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    /// Whether the recorded hash-join dispatch is enabled via
1539    /// env. Reads `XLOG_USE_RECORDED_HASH_JOIN` or
1540    /// `XLOG_USE_RECORDED_OPS`. `hash_join_v2_recorded` and
1541    /// `hash_join_v2_with_index_recorded` cover all four join
1542    /// types (Inner / Semi / Anti / LeftOuter); the only
1543    /// hard constraint inherited from `pack_keys` is `≤4`
1544    /// key columns.
1545    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    /// Whether the recorded CSM (count-scan-materialize)
1550    /// dispatch is enabled via env. Reads `XLOG_USE_RECORDED_CSM`
1551    /// or `XLOG_USE_RECORDED_OPS`. CSM is a sub-strategy of the
1552    /// recorded hash-join: it is consulted only after the
1553    /// recorded path has already been selected, and only for
1554    /// `JoinType::Inner` / `JoinType::LeftOuter` where a CSM
1555    /// implementation exists. `Semi` / `Anti` are not affected.
1556    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    /// Whether the bounded CSM CUDA Graph path is enabled.
1561    ///
1562    /// This is narrower than `XLOG_USE_RECORDED_CSM`: callers must first select
1563    /// the recorded CSM hash-join path, then opt into graph capture/replay with
1564    /// `XLOG_USE_CSM_CUDA_GRAPH=1` (or the broader `XLOG_USE_CUDA_GRAPHS=1`).
1565    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    /// Test/diagnostic-only telemetry: number of times the recorded
1570    /// hash-join dispatch routed through a CSM (count-scan-materialize)
1571    /// method since this provider was created. Increments once per
1572    /// dispatched call across all four CSM methods (Inner / LeftOuter,
1573    /// non-indexed / indexed). Used by `test_csm_env_dispatch` to
1574    /// prove dispatch selection.
1575    ///
1576    /// **Not part of any public stability guarantee.** Hidden from
1577    /// rustdoc with `#[doc(hidden)]` so it does not appear in
1578    /// generated API docs; the symbol remains callable from
1579    /// integration tests within this crate but production callers
1580    /// must not depend on it. May be renamed, gated behind a cargo
1581    /// feature, or withdrawn in any release without notice.
1582    #[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    /// Lazily acquire one non-default launch stream from the
1613    /// runtime's [`crate::device_runtime::StreamPool`] for
1614    /// recorded-operator dispatch, and cache it for this
1615    /// provider's lifetime. Shared across all env-gated
1616    /// recorded paths (filter, sort, dedup, GroupBy,
1617    /// hash-join) — a single stream is sufficient because the
1618    /// recorder serializes work on it; multiple operations
1619    /// chain naturally through commit-order events.
1620    ///
1621    /// Returns `None` when:
1622    ///   * the manager has no runtime attached
1623    ///     (`memory.runtime() == None`), or
1624    ///   * the stream pool is at capacity and `acquire` fails.
1625    ///
1626    /// On a lost race during first init the loser leaks one
1627    /// stream (the pool keeps it alive); both winners cache
1628    /// the same `StreamId`. Acceptable cost — practical pool
1629    /// sizes are large compared to the number of providers
1630    /// per process.
1631    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    /// Take the per-phase WCOJ triangle dispatch timings recorded
1642    /// by the most recent `wcoj_triangle_*_recorded` call. Reading
1643    /// clears the slot — designed for one-shot consumption by the
1644    /// `wcoj_phase_report` binary in xlog-integration. Returns
1645    /// `None` if no triangle dispatch has fired since the last
1646    /// read (or since construction).
1647    ///
1648    /// Compiled in only with the `wcoj-phase-timing` Cargo
1649    /// feature; production builds have no such method.
1650    #[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    /// Internal: store the phase timings produced by a triangle
1661    /// dispatch. Overwrites any prior unread slot — the report
1662    /// binary is expected to read after every `execute_plan`.
1663    #[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    /// Number of times `wcoj_layout_*_recorded` short-circuited
1675    /// to the fast-path (recorded clone) instead of running
1676    /// `dedup_full_row_recorded`. Increments by 1 per
1677    /// fast-path hit (3 hits per dispatch when all inputs are
1678    /// already sorted+unique). Used by tests + the phase
1679    /// report to confirm the fast-path fired.
1680    pub fn wcoj_layout_fast_path_hit_count(&self) -> u64 {
1681        self.wcoj_layout_fast_path_hit_count.load(Ordering::Relaxed)
1682    }
1683
1684    /// Histogram-guided block-slice triangle WCOJ test/diagnostic counter:
1685    /// successful dispatches that routed through the provider entry.
1686    pub fn wcoj_triangle_hg_dispatch_count(&self) -> u64 {
1687        self.wcoj_triangle_hg_dispatch_count.load(Ordering::Relaxed)
1688    }
1689
1690    /// Reset the fast-path hit counter to 0. Tests use this to
1691    /// scope counter assertions to a single dispatch.
1692    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    /// Number of calls to `wcoj_layout_sort_*_recorded` since the
1698    /// last reset. Diagnostic-only; used by dispatch-plan certification.
1699    pub fn wcoj_layout_sort_invocation_count(&self) -> u64 {
1700        self.wcoj_layout_sort_invocation_count
1701            .load(Ordering::Relaxed)
1702    }
1703
1704    /// Reset the WCOJ layout-sort invocation counter to 0.
1705    pub fn reset_wcoj_layout_sort_invocation_count(&self) {
1706        self.wcoj_layout_sort_invocation_count
1707            .store(0, Ordering::Relaxed);
1708    }
1709
1710    /// Number of K-clique leader-edge metadata builds since the
1711    /// last reset.
1712    pub fn kclique_metadata_build_count(&self) -> u64 {
1713        self.kclique_metadata_build_count.load(Ordering::Relaxed)
1714    }
1715
1716    /// Cumulative nanoseconds spent building K-clique leader-edge
1717    /// metadata since the last reset.
1718    pub fn kclique_metadata_build_nanos(&self) -> u64 {
1719        self.kclique_metadata_build_nanos.load(Ordering::Relaxed)
1720    }
1721
1722    /// Reset K-clique metadata build diagnostics.
1723    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    /// Internal: increment the fast-path counter. Called by
1731    /// `wcoj_layout_*_recorded` after a successful fast-path
1732    /// branch. Not part of any public stability guarantee.
1733    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    /// Internal: increment the generic WCOJ layout-sort counter.
1739    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    /// Internal: record a K-clique leader-edge metadata build.
1745    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    /// Runtime hook: record a successful histogram-guided block-slice triangle
1754    /// dispatch.
1755    #[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    /// Get the CUDA device
1762    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    /// Get the GPU memory manager
1771    pub fn memory(&self) -> &Arc<GpuMemoryManager> {
1772        &self.memory
1773    }
1774
1775    /// Get PTX load profiling data (only populated when XLOG_WARMUP_PROFILE=1).
1776    pub fn ptx_load_profile(&self) -> Option<&PtxLoadProfile> {
1777        self.ptx_load_profile.as_ref()
1778    }
1779
1780    /// Reset tracked host transfer statistics.
1781    pub fn reset_host_transfer_stats(&self) {
1782        self.transfer_tracker.reset();
1783    }
1784
1785    /// Snapshot tracked host transfer statistics.
1786    pub fn host_transfer_stats(&self) -> HostTransferStats {
1787        self.transfer_tracker.snapshot()
1788    }
1789
1790    /// Reset the separately accounted final resident-receipt transfer.
1791    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    /// Snapshot the separately accounted final resident-receipt transfer.
1798    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    /// Record a successfully scheduled device-to-pinned-host final receipt.
1810    ///
1811    /// This deliberately does not increment ordinary hot-loop transfer
1812    /// counters. Callers may invoke it only after the resident graph's single
1813    /// terminal synchronization.
1814    #[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    /// Validate and publish logical row counts decoded from one resident receipt.
1822    ///
1823    /// Every entry is validated before any cache is changed. This keeps receipt
1824    /// application all-or-nothing and ensures later query, constraint, export,
1825    /// and statistics consumers never mistake reserved capacity for cardinality
1826    /// or trigger another metadata download.
1827    #[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    /// Snapshot launch-parameter H2D uploads tracked separately from
1853    /// `host_transfer_stats`.
1854    pub fn host_launch_metadata_transfer_stats(&self) -> HostLaunchMetadataTransferStats {
1855        self.transfer_tracker.launch_metadata_snapshot()
1856    }
1857
1858    /// Read the column-level D2H transfer counter.
1859    ///
1860    /// This counter increments once per `download_column_*` call, enabling
1861    /// callers (e.g. the ILP trainer) to assert that no column downloads
1862    /// occurred during a performance-critical section.
1863    pub fn d2h_transfer_count(&self) -> u64 {
1864        self.d2h_transfer_count.load(Ordering::Relaxed)
1865    }
1866
1867    /// Reset the column-level D2H transfer counter to zero.
1868    pub fn reset_d2h_transfer_count(&self) {
1869        self.d2h_transfer_count.store(0, Ordering::Relaxed);
1870    }
1871
1872    /// Count of untracked control-plane metadata D2H reads
1873    /// (`dtoh_scalar_untracked` + `dtoh_small_metadata_untracked`).
1874    pub fn untracked_metadata_dtoh_count(&self) -> u64 {
1875        self.untracked_metadata_dtoh_count.load(Ordering::Relaxed)
1876    }
1877
1878    /// Reset the untracked metadata D2H read counter to zero.
1879    pub fn reset_untracked_metadata_dtoh_count(&self) {
1880        self.untracked_metadata_dtoh_count
1881            .store(0, Ordering::Relaxed);
1882    }
1883
1884    /// Enable the strict deterministic-Datalog D2H gate.
1885    ///
1886    /// While enabled, any data-plane device-to-host transfer (column downloads
1887    /// via `download_column` / `download_column_untracked`, and any internal
1888    /// transfer routed through `dtoh_sync_copy_into_tracked`) increments
1889    /// [`CudaKernelProvider::deterministic_d2h_violation_count`] and returns
1890    /// `XlogError::Execution` from the originating call.
1891    ///
1892    /// Metadata reads via [`CudaKernelProvider::dtoh_scalar_untracked`] are
1893    /// allowed and never trip the gate.
1894    ///
1895    /// Default is `false`; the runtime opts in via
1896    /// `RuntimeConfig::strict_deterministic_d2h`. v0.5.5 ships the gate
1897    /// opt-in only — known-violating relational paths (set difference,
1898    /// join count/materialize) are scheduled for replacement before the
1899    /// default flips.
1900    pub fn enable_strict_deterministic_d2h(&self) {
1901        self.strict_deterministic_d2h.store(true, Ordering::Relaxed);
1902    }
1903
1904    /// Disable the strict deterministic-Datalog D2H gate.
1905    pub fn disable_strict_deterministic_d2h(&self) {
1906        self.strict_deterministic_d2h
1907            .store(false, Ordering::Relaxed);
1908    }
1909
1910    /// Returns whether the strict deterministic-Datalog D2H gate is enabled.
1911    pub fn strict_deterministic_d2h_enabled(&self) -> bool {
1912        self.strict_deterministic_d2h.load(Ordering::Relaxed)
1913    }
1914
1915    /// Cumulative deterministic-D2H gate violations since the last reset.
1916    pub fn deterministic_d2h_violation_count(&self) -> u64 {
1917        self.deterministic_d2h_violations.load(Ordering::Relaxed)
1918    }
1919
1920    /// Reset the deterministic-D2H violation counter to zero.
1921    pub fn reset_deterministic_d2h_violations(&self) {
1922        self.deterministic_d2h_violations
1923            .store(0, Ordering::Relaxed);
1924    }
1925
1926    /// Chokepoint for the deterministic-D2H gate.
1927    ///
1928    /// If the gate is enabled, increments the violation counter and returns
1929    /// `XlogError::Execution` naming the offending operation and byte count.
1930    /// If the gate is disabled, returns `Ok(())` cheaply.
1931    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    /// Hard cap (in bytes) for [`Self::dtoh_small_metadata_untracked`].
1960    /// Set deliberately small (4 KB) so the helper cannot become a
1961    /// general-purpose vector D2H escape hatch — it's strictly for
1962    /// classifier histograms and similar small metadata round-trips.
1963    pub const DTOH_SMALL_METADATA_MAX_BYTES: usize = 4096;
1964
1965    /// Read a small metadata vector (≤ [`Self::DTOH_SMALL_METADATA_MAX_BYTES`])
1966    /// from device to host WITHOUT updating the D2H transfer tracker.
1967    ///
1968    /// Sibling of [`Self::dtoh_scalar_untracked`] for callers that need
1969    /// a few bucket counts (the WCOJ skew classifier reads a 3 × 64 ×
1970    /// `u32` = 768-byte histogram in one go) instead of `count` separate
1971    /// scalar reads. Like `dtoh_scalar_untracked`, this method is
1972    /// whitelisted by the strict deterministic-D2H gate
1973    /// ([`Self::enable_strict_deterministic_d2h`]) — it does NOT trip
1974    /// the gate, on purpose, because metadata reads are part of the
1975    /// determinism contract (just like a scalar `total` after a scan).
1976    ///
1977    /// # Hard contract — DO NOT WIDEN THE CAP
1978    /// The 4 KB cap is the contract. If a caller wants a larger D2H,
1979    /// it's a data-plane transfer and must go through the tracked
1980    /// `download_column*` path. Widening this cap turns the helper
1981    /// into a backdoor for tracked-bypass column reads, which would
1982    /// silently invalidate the strict deterministic-D2H gate.
1983    ///
1984    /// # Errors
1985    ///   * `XlogError::Kernel` if `count * size_of::<T>()` exceeds
1986    ///     `DTOH_SMALL_METADATA_MAX_BYTES`.
1987    ///   * `XlogError::Kernel` if `count` exceeds the device slice's
1988    ///     length, or if the inner sync copy fails.
1989    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    /// Read a single scalar from device to host WITHOUT updating the
2032    /// D2H transfer tracker. Use ONLY for metadata reads (e.g. total_nnz
2033    /// after an exclusive scan), never for data-plane transfers.
2034    ///
2035    /// This makes the "metadata != data-plane" contract explicit and
2036    /// auditable: callers that bypass tracking must call this method
2037    /// (which is grep-able) rather than reaching for device().inner().
2038    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    /// Upload host data to device while recording data-plane H2D transfer stats.
2067    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    /// Allocate a CUDA slice from host data while recording data-plane H2D
2083    /// transfer stats.
2084    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    /// Upload bounded launch metadata from host to device while recording it in
2099    /// the launch-metadata subcounter.
2100    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    /// Upload one launch-metadata scalar to device on a caller-owned stream
2122    /// while recording the transfer in the launch-metadata H2D counters.
2123    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    /// Compute exclusive prefix sum of u8 mask, returns (prefix_sum_vec, total_count)
2150    ///
2151    /// This is useful for compaction operations where we need to know:
2152    /// 1. The output position for each input element (prefix sum)
2153    /// 2. The total number of elements that pass the mask (count)
2154    ///
2155    /// # Arguments
2156    /// * `mask` - A slice of u8 values (0 or non-zero)
2157    ///
2158    /// # Returns
2159    /// A tuple of:
2160    /// - `Vec<u32>` containing the exclusive prefix sum
2161    /// - `u32` containing the total count of non-zero mask elements
2162    ///
2163    /// # Example
2164    /// ```ignore
2165    /// let mask = vec![1u8, 0, 1, 1, 0, 1];
2166    /// let (prefix_sum, count) = provider.prefix_sum_mask(&mask)?;
2167    /// // prefix_sum = [0, 1, 1, 2, 3, 3]
2168    /// // count = 4
2169    /// ```
2170    ///
2171    /// # Note
2172    /// For small inputs (<=256 elements), a CPU scan is used for efficiency.
2173    /// For larger inputs, a three-phase multi-block GPU scan is used.
2174    ///
2175    /// # Errors
2176    /// Returns `XlogError::Kernel` if kernel execution fails
2177    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            // SAFETY: multiblock_scan_phase2(uint32_t* block_sums, uint32_t num_blocks)
2212            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        // SAFETY: multiblock_scan_u32_phase1(uint32_t* data, uint32_t* block_sums, uint32_t n)
2238        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        // SAFETY: multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
2262        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    /// Stream-aware variant of [`Self::multiblock_scan_u32_inplace`].
2279    ///
2280    /// Runs every kernel of the recursive scan on `cu_stream`
2281    /// (no `device.synchronize()`), and records each intermediate
2282    /// `block_sums` allocation against the runtime so that when
2283    /// the helper returns and the local drops, the runtime's
2284    /// deallocate can queue `cuStreamWaitEvent(alloc_stream,
2285    /// recorded_event)` BEFORE `cuMemFreeAsync` — the same
2286    /// cross-stream lifetime safety the LaunchRecorder gives
2287    /// caller-provided buffers.
2288    ///
2289    /// `data` is not recorded here: the caller already records
2290    /// its own write of `data` against the same launch_stream
2291    /// (typically via `LaunchRecorder::write` BEFORE preflight).
2292    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            // SAFETY: kernel signature matches; data is mutated in place.
2313            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        // Fence alloc-ready → launch_stream for block_sums
2333        // before phase1 kernel writes it. The alloc was queued
2334        // on the manager's default stream; without this wait,
2335        // a launch_stream-queued kernel can begin before
2336        // cuMemAllocAsync completes and read pool-recycled
2337        // bytes when the streams differ.
2338        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        // SAFETY: kernel signature matches.
2357        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        // SAFETY: kernel signature matches.
2391        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        // Record `block_sums` use on `launch_stream` BEFORE it
2407        // drops at end-of-scope. Without this, the runtime's
2408        // deallocate would queue `cuMemFreeAsync` on alloc_stream
2409        // without waiting for the launch_stream chain that's
2410        // still reading/writing block_sums to complete.
2411        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    /// Allocate every recursive `block_sums` buffer needed by
2436    /// [`Self::multiblock_scan_u32_inplace_on_stream_with_scratch`].
2437    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    /// Stream-aware u32 scan with caller-owned scratch.
2452    ///
2453    /// This is the CUDA Graph compatible counterpart to
2454    /// [`Self::multiblock_scan_u32_inplace_on_stream`]: all scratch buffers are
2455    /// supplied by the caller, so graph capture sees a stable scan topology and
2456    /// stable intermediate addresses.
2457    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            // SAFETY: kernel signature matches; data is mutated in place.
2492            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        // SAFETY: kernel signature matches.
2534        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        // SAFETY: kernel signature matches.
2564        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    /// Stream-aware view-inplace variant of
2585    /// [`Self::multiblock_scan_u32_view_inplace`]. Same shape
2586    /// as [`Self::multiblock_scan_u32_inplace_on_stream`] but
2587    /// over a `CudaViewMut` (used by recorded radix sort
2588    /// digit loops that scan per-digit slices of the histogram
2589    /// in place). Records intermediate `block_sums` against
2590    /// the runtime before they drop at end-of-scope.
2591    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            // SAFETY: phase2 kernel signature.
2612            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        // Fence alloc-ready → launch_stream for block_sums
2635        // before phase1 kernel writes it. See the inplace
2636        // variant for the full rationale.
2637        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        // SAFETY: phase1 kernel signature.
2656        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        // SAFETY: phase3 kernel signature.
2690        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        // Record block_sums use before end-of-scope drop.
2709        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            // SAFETY: multiblock_scan_phase2(uint32_t* block_sums, uint32_t num_blocks)
2753            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        // SAFETY: multiblock_scan_u32_phase1(uint32_t* data, uint32_t* block_sums, uint32_t n)
2779        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        // SAFETY: multiblock_scan_phase3(uint32_t* prefix_sum, const uint32_t* block_offsets, uint32_t n)
2803        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    // ============== Internal Helper Methods ==============
2820
2821    /// Read a buffer's logical row count, using the host cache when available
2822    /// and falling back to a metadata-only device-to-host read when needed.
2823    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    /// Read and validate a buffer's logical row count for outward-facing APIs.
2833    ///
2834    /// This keeps exported/query-visible lengths tied to the device logical row
2835    /// count while still rejecting impossible metadata (`logical_rows > row_cap`).
2836    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    /// Reinterpret a `CudaBuffer` column as a `u32` slice for kernel access.
2923    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    /// Reinterpret a `CudaBuffer` column as an `f64` slice for kernel access.
2980    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    /// Create an empty buffer with the given schema (all columns are empty slices)
3009    ///
3010    /// # Arguments
3011    /// * `schema` - The schema for the empty buffer
3012    ///
3013    /// # Returns
3014    /// A new CudaBuffer with zero rows
3015    ///
3016    /// # Errors
3017    /// Returns `XlogError::Kernel` if allocation fails
3018    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            // Allocate zero-length column
3022            columns.push(self.memory.alloc::<u8>(0)?.into());
3023        }
3024        self.buffer_from_columns(columns, 0, schema)
3025    }
3026
3027    /// Create a zero-arity (nullary) relation buffer carrying `rows` unit tuples.
3028    ///
3029    /// A nullary relation holds exactly when it has at least one row; its single
3030    /// possible tuple is the empty tuple `()`. `create_buffer_from_slices` with no
3031    /// column slices routes to `create_empty_buffer` (0 rows), which represents the
3032    /// relation as *absent* — wrong for an asserted nullary fact. Nullary facts must
3033    /// use this path so presence is materialized as one row.
3034    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    /// Combine schemas from left and right buffers for join result
3063    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    /// Check if two schemas have compatible types (same arity and column types)
3074    ///
3075    /// This ignores column names, which is useful for Datalog operations where
3076    /// projected relations may have different column names but the same types.
3077    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        // Verify resolve_module_path finds portable PTX for all modules.
3165        // Uses a dummy cc (999) so cubin won't match — only portable PTX.
3166        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); // 1 GB
3243        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        // Verify all kernel functions can be retrieved
3272        let inner = device.inner();
3273
3274        // Join kernels
3275        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        // Dedup kernels
3288        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        // GroupBy kernels
3301        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        // Circuit kernels (XGCF forward/backward)
3333        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        // Neural fast-path kernels (AD chain weight fill + gradient scatter)
3360        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        // Ensure module names don't collide
3375        assert_ne!(JOIN_MODULE, DEDUP_MODULE);
3376        assert_ne!(JOIN_MODULE, GROUPBY_MODULE);
3377        assert_ne!(DEDUP_MODULE, GROUPBY_MODULE);
3378    }
3379
3380    // Helper function to create test provider
3381    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    // Helper function to create a CudaBuffer with U32 data
3461    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    // Helper function to create an empty buffer with correct column count
3482    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    // Helper function to read U32 data from CudaBuffer
3493    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    /// `clone_buffer` must propagate the host-side `cached_row_count` so
3600    /// downstream code can read the row count without a D2H round-trip.
3601    /// Without this propagation, buffers flowed through the relation store
3602    /// (`CompiledIlpProgram::put_relation` calls `clone_buffer` before
3603    /// storing) lose their host-visible count, forcing consumers to choose
3604    /// between an extra D2H (violating the native bounded exact-induction
3605    /// transfer-budget gates) and a hard error. This test pins the cache-propagation
3606    /// contract directly.
3607    #[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        // Source's cache is populated by the `create_buffer_from_*` path;
3623        // verify the precondition so a regression in that path shows up here
3624        // rather than silently passing the real assertion below.
3625        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    // ============== Hash Join Tests ==============
3740
3741    #[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        // Join empty with empty
3755        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        // Empty key columns
3774        let result = provider.hash_join(&left, &right, &[], &[0]);
3775        assert!(result.is_err());
3776
3777        // Mismatched key lengths
3778        let result = provider.hash_join(&left, &right, &[0], &[0, 0]);
3779        assert!(result.is_err());
3780    }
3781
3782    // ============== Dedup Tests ==============
3783
3784    #[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        // Empty key columns
3815        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        // Test dedup with duplicates: [3, 1, 2, 1, 3, 2]
3830        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        // Result should be sorted and deduped
3840        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        // Create input with duplicates: 0..500 ++ 250..750 = 1000 elements, 750 unique
3854        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        // Verify output is sorted
3867        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        // Verify expected values
3872        let expected: Vec<u32> = (0..750).collect();
3873        assert_eq!(result, expected);
3874    }
3875
3876    // ============== Union Tests ==============
3877
3878    #[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        // Empty union empty
3892        let result = provider.union(&empty, &empty);
3893        assert!(result.is_ok());
3894        assert!(result.unwrap().is_empty());
3895
3896        // Non-empty union empty
3897        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        // Different column names but same types should succeed (Datalog union semantics)
3919        let result = provider.union(&a, &b);
3920        assert!(result.is_ok());
3921
3922        // Different arity should fail - create a 2-column buffer
3923        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    // ============== Diff Tests ==============
3935
3936    #[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        // Empty diff empty
3950        let result = provider.diff(&empty, &empty);
3951        assert!(result.is_ok());
3952        assert!(result.unwrap().is_empty());
3953
3954        // Non-empty diff empty should return all of a
3955        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); // 1, 3, 5
3980
3981        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        // Different column names with same types should work (Datalog semantics)
4014        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        // Create buffers with different arities (this should fail)
4023        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    // ============== GroupBy Aggregation Tests ==============
4061
4062    #[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        // Empty key columns
4093        let result = provider.groupby_agg(&buffer, &[], AggOp::Count, 0);
4094        assert!(result.is_err());
4095
4096        // Value column out of bounds
4097        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        // Create buffer with U32 keys and F64 values
4112        // Group 0 (key=1): values 1.0, 2.0 -> logsumexp = log(e^1 + e^2) ≈ 2.31326
4113        // Group 1 (key=2): values 3.0, 4.0 -> logsumexp = log(e^3 + e^4) ≈ 4.31326
4114        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        // Create key column
4123        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        // Create value column
4135        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        // Run LogSumExp aggregation grouped by key column (0), aggregating value column (1)
4151        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        // Download results
4165        let result_values = provider
4166            .download_column::<f64>(&result, 1)
4167            .expect("download result");
4168
4169        // Expected values:
4170        // logsumexp(1.0, 2.0) = 2.0 + log(exp(1.0-2.0) + exp(2.0-2.0)) = 2.0 + log(e^-1 + 1) ≈ 2.31326
4171        // logsumexp(3.0, 4.0) = 4.0 + log(exp(3.0-4.0) + exp(4.0-4.0)) = 4.0 + log(e^-1 + 1) ≈ 4.31326
4172        let expected_0 = 2.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); // ≈ 2.31326
4173        let expected_1 = 4.0_f64 + ((-1.0_f64).exp() + 1.0_f64).ln(); // ≈ 4.31326
4174
4175        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    // ============== Schema Helper Tests ==============
4191
4192    #[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        // Count result schema (u64 to match predicate declarations)
4227        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        // Sum result schema
4233        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        // Min/Max result schema
4238        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        // Sum should return U64 to prevent overflow
4262        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        // Create buffers that produce more than 10 results when joined
4280        // Left: [1, 1, 1, 1, 2, 2, 2, 2] - 4 copies of 1, 4 copies of 2
4281        // Right: [1, 1, 1, 2, 2, 2] - 3 copies of 1, 3 copies of 2
4282        // Join produces: 4*3 + 4*3 = 24 results
4283        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        // Test with limit of 10 - should get at most 10
4287        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        // Test with None (default) - should get all 24 results
4297        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        // Test legacy API still works (backward compatibility)
4308        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    // ============== Arithmetic Operation Tests ==============
4320
4321    /// Helper to create a test provider for arithmetic tests
4322    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    /// Helper to create an i64 buffer for arithmetic tests
4333    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    /// Helper to create an f64 buffer for arithmetic tests
4341    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]); // Note: division by zero
4421
4422        let result = provider.div_columns(&a, &b).unwrap();
4423        let values = provider.download_column::<i64>(&result, 0).unwrap();
4424
4425        // Division by zero returns i64::MAX
4426        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]); // Note: mod by zero
4454
4455        let result = provider.mod_columns(&a, &b).unwrap();
4456        let values = provider.download_column::<i64>(&result, 0).unwrap();
4457
4458        // Mod by zero returns 0
4459        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        // IEEE 754: 1.0/0.0 = Inf, -1.0/0.0 = -Inf, 0.0/0.0 = NaN
4555        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        // pow always returns f64
4574        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        // sqrt(4) = 2, sqrt(9) = 3, cbrt(27) = 3
4591        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        // Truncation towards zero
4624        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]); // Different size
4651
4652        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        // Addition should wrap
4685        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); // MAX + 1 wraps to MIN
4688        assert_eq!(add_values[1], i64::MAX); // MIN - 1 wraps to MAX
4689    }
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}