xlog_cuda/provider/
ilp_exact_nary.rs

1//! Launcher for the n-ary bounded exact-induction scoring kernel.
2//!
3//! Drives `kernels/ilp_exact_nary.cu`'s `ilp_exact_nary_score`: one block
4//! per flattened pattern, threads striding the example tuples, block-reduced
5//! positive/negative coverage counts per pattern.
6//!
7//! Two entry points share one launch core:
8//!
9//! * [`CudaKernelProvider::ilp_exact_nary_score`] — host-slice inputs.
10//!   The CPU-side leg of the parity chain (reference == flat == device)
11//!   and the harness the pod fixtures drive.
12//! * [`CudaKernelProvider::ilp_exact_nary_score_device`] — the PRODUCTION
13//!   ingest: candidate relations and example tuples arrive as
14//!   device-resident columnar [`CudaBuffer`]s and are concatenated with
15//!   device-to-device column copies. No relation or example value ever
16//!   visits the host; the only D2H is the two per-pattern count arrays.
17//!
18//! Everything is COLUMN-MAJOR in element units: relation cell
19//! (row, position) sits at `cand_value_offset[slot] + position *
20//! cand_rows[slot] + row`, and example position `i` of tuple `q` sits at
21//! `i * count + q` — the same formulas the flat host interpreter and the
22//! kernel share.
23//!
24//! Example tuples are a BAG, not a set: a duplicated positive counts once
25//! per occurrence, so a caller that passes duplicates inflates coverage.
26//! This is consistent across the reference, flat and device layers and is
27//! the caller's contract to uphold.
28//!
29//! Per-launch runtime is NOT bounded by the pattern cap: one block walks a
30//! backtracking search whose worst case is `rows^body_len`. On a display
31//! GPU a wide relation set can therefore trip the driver watchdog; bound
32//! the inputs, not just the pattern count.
33//!
34//! Layering: pattern arrays are PRIMITIVE flat host slices (the encoding
35//! produced by `xlog_induce::nary_layout::flatten_patterns`), so
36//! xlog-cuda keeps no dependency on xlog-induce. Every structural
37//! invariant the kernel relies on (offsets in bounds, per-slot arity
38//! consistency, binding indexes inside the fixed device state) is
39//! re-validated here fail-closed before any upload; the kernel never sees
40//! an out-of-bounds batch.
41
42use std::marker::PhantomData;
43use std::sync::atomic::Ordering;
44
45use crate::memory::{CudaBuffer, TrackedCudaSlice};
46use crate::{LaunchAsync, LaunchConfig};
47use xlog_core::{Result, ScalarType, XlogError};
48
49use super::{ilp_exact_nary_kernels, RawCudaView, ILP_EXACT_NARY_MODULE};
50
51/// MUST equal `ILP_EXACT_NARY_BLOCK_SIZE` in `kernels/ilp_exact_nary.cu`:
52/// the kernel sizes its static `__shared__` scratch from that macro, so a
53/// larger launch dimension would let threads write past it. The block
54/// reduction additionally assumes a power of two.
55const ILP_EXACT_NARY_BLOCK_SIZE: u32 = 256;
56const _: () = assert!(
57    ILP_EXACT_NARY_BLOCK_SIZE.is_power_of_two() && ILP_EXACT_NARY_BLOCK_SIZE <= 256,
58    "block size must stay a power of two within the kernel's shared scratch",
59);
60
61// Device-contract bounds; must equal both the kernel's fixed state and
62// xlog_induce::nary_layout's flattener bounds. Head arity is bounded by
63// the kernel's per-thread example gather array.
64const NARY_MAX_BODY_ATOMS: u32 = 8;
65const NARY_MAX_JOIN_VARS: u32 = 8;
66const NARY_MAX_HEAD_ARITY: u32 = 8;
67// The kernel is safe with a wider atom (its position loop is bounded by
68// the binding-code and relation extents), but the contract publishes
69// arity <= 8 and this launcher is exported raw — without this the bound
70// would hold only for callers who came through the flattener.
71const NARY_MAX_ATOM_ARITY: u32 = 8;
72
73const NARY_JOIN_FLAG: u32 = 0x8000_0000;
74const NARY_INDEX_MASK: u32 = 0xFF;
75
76const U64_SIZE: usize = std::mem::size_of::<u64>();
77
78/// Flattened pattern batch (host-born: patterns are enumerated on host).
79pub struct IlpExactNaryPatterns<'a> {
80    pub body_offset: &'a [u32],
81    pub body_len: &'a [u32],
82    pub atom_candidate_slot: &'a [u32],
83    pub atom_arity: &'a [u32],
84    pub atom_binding_offset: &'a [u32],
85    pub binding_codes: &'a [u32],
86    pub head_arity: u32,
87}
88
89/// Host-slice scoring request: the pattern batch plus candidate relations
90/// and example tuples as flat host arrays, exactly as the kernel consumes
91/// them (column-major, element units).
92pub struct IlpExactNaryRequest<'a> {
93    pub body_offset: &'a [u32],
94    pub body_len: &'a [u32],
95    pub atom_candidate_slot: &'a [u32],
96    pub atom_arity: &'a [u32],
97    pub atom_binding_offset: &'a [u32],
98    pub binding_codes: &'a [u32],
99    /// Concatenated COLUMN-MAJOR u64 relation values: within a
100    /// relation, cell (row, position) sits at `cand_value_offset[slot]
101    /// + position * cand_rows[slot] + row`. All offsets are u32
102    /// ELEMENT indexes, never bytes.
103    pub cand_values: &'a [u64],
104    /// Element offset of each relation's first value in `cand_values`.
105    pub cand_value_offset: &'a [u32],
106    pub cand_rows: &'a [u32],
107    /// COLUMNAR example tuples: position `i` of example `q` sits at
108    /// `i * example_count + q`; length is a multiple of `head_arity`.
109    pub pos_values: &'a [u64],
110    pub neg_values: &'a [u64],
111    pub head_arity: u32,
112}
113
114impl<'a> IlpExactNaryRequest<'a> {
115    fn patterns(&self) -> IlpExactNaryPatterns<'a> {
116        IlpExactNaryPatterns {
117            body_offset: self.body_offset,
118            body_len: self.body_len,
119            atom_candidate_slot: self.atom_candidate_slot,
120            atom_arity: self.atom_arity,
121            atom_binding_offset: self.atom_binding_offset,
122            binding_codes: self.binding_codes,
123            head_arity: self.head_arity,
124        }
125    }
126}
127
128fn nary_err(detail: impl std::fmt::Display) -> XlogError {
129    XlogError::Kernel(format!("ilp_exact_nary_score: {detail}"))
130}
131
132/// Validate the pattern batch against the candidate slot table. Shared by
133/// the host-slice and device-buffer paths; returns the pattern count.
134fn validate_batch(
135    patterns: &IlpExactNaryPatterns<'_>,
136    cand_value_offset: &[u32],
137    cand_rows: &[u32],
138    values_len: u64,
139) -> Result<u32> {
140    let pattern_count = patterns.body_offset.len();
141    if pattern_count == 0 {
142        return Err(nary_err("empty pattern batch"));
143    }
144    if patterns.body_len.len() != pattern_count {
145        return Err(nary_err("body_len length != pattern count"));
146    }
147    let atoms = patterns.atom_candidate_slot.len();
148    if patterns.atom_arity.len() != atoms || patterns.atom_binding_offset.len() != atoms {
149        return Err(nary_err("atom array lengths disagree"));
150    }
151    let bindings = patterns.binding_codes.len();
152    let slots = cand_value_offset.len();
153    if cand_rows.len() != slots {
154        return Err(nary_err("cand_rows length != cand_value_offset length"));
155    }
156    if patterns.head_arity == 0 {
157        return Err(nary_err("head_arity must be >= 1"));
158    }
159    if patterns.head_arity > NARY_MAX_HEAD_ARITY {
160        return Err(nary_err(format!(
161            "head_arity {} exceeds device bound {NARY_MAX_HEAD_ARITY}",
162            patterns.head_arity
163        )));
164    }
165
166    // Per-slot arity is implied by the atoms that read the slot; it must
167    // be consistent and its rows must fit the concatenated value buffer.
168    let mut slot_arity: Vec<Option<u32>> = vec![None; slots];
169    for (pattern, (&offset, &len)) in patterns
170        .body_offset
171        .iter()
172        .zip(patterns.body_len)
173        .enumerate()
174    {
175        if len == 0 {
176            return Err(nary_err(format!("pattern {pattern} has an empty body")));
177        }
178        if len > NARY_MAX_BODY_ATOMS {
179            return Err(nary_err(format!(
180                "pattern {pattern} has {len} body atoms; device bound is \
181                 {NARY_MAX_BODY_ATOMS}"
182            )));
183        }
184        let end = offset
185            .checked_add(len)
186            .ok_or_else(|| nary_err("body offset overflow"))?;
187        if end as usize > atoms {
188            return Err(nary_err(format!(
189                "pattern {pattern} body [{offset}, {end}) exceeds {atoms} atoms"
190            )));
191        }
192        for atom in offset as usize..end as usize {
193            let slot = patterns.atom_candidate_slot[atom] as usize;
194            if slot >= slots {
195                return Err(nary_err(format!(
196                    "atom {atom} references candidate slot {slot} of {slots}"
197                )));
198            }
199            let arity = patterns.atom_arity[atom];
200            if arity == 0 {
201                return Err(nary_err(format!("atom {atom} has arity 0")));
202            }
203            if arity > NARY_MAX_ATOM_ARITY {
204                return Err(nary_err(format!(
205                    "atom {atom} has arity {arity} > the device contract's \
206                     {NARY_MAX_ATOM_ARITY}"
207                )));
208            }
209            match slot_arity[slot] {
210                None => slot_arity[slot] = Some(arity),
211                Some(existing) if existing == arity => {}
212                Some(existing) => {
213                    return Err(nary_err(format!(
214                        "candidate slot {slot} read at arity {arity} and \
215                         arity {existing}; relation arity must be consistent"
216                    )));
217                }
218            }
219            let rows = cand_rows[slot] as u64;
220            let value_end = cand_value_offset[slot] as u64 + rows * arity as u64;
221            if value_end > values_len {
222                return Err(nary_err(format!(
223                    "candidate slot {slot} needs values up to {value_end}, \
224                     buffer holds {values_len}"
225                )));
226            }
227            let binding_offset = patterns.atom_binding_offset[atom];
228            let binding_end = binding_offset
229                .checked_add(arity)
230                .ok_or_else(|| nary_err("binding offset overflow"))?;
231            if binding_end as usize > bindings {
232                return Err(nary_err(format!(
233                    "atom {atom} bindings [{binding_offset}, {binding_end}) \
234                     exceed {bindings} codes"
235                )));
236            }
237            for position in binding_offset as usize..binding_end as usize {
238                let code = patterns.binding_codes[position];
239                let index = code & NARY_INDEX_MASK;
240                if code & NARY_JOIN_FLAG != 0 {
241                    if index >= NARY_MAX_JOIN_VARS {
242                        return Err(nary_err(format!(
243                            "binding {position} join index {index} >= device \
244                             bound {NARY_MAX_JOIN_VARS}"
245                        )));
246                    }
247                } else if index >= patterns.head_arity {
248                    return Err(nary_err(format!(
249                        "binding {position} head index {index} >= head arity \
250                         {}",
251                        patterns.head_arity
252                    )));
253                }
254            }
255        }
256    }
257
258    u32::try_from(pattern_count).map_err(|_| nary_err("pattern count exceeds u32"))
259}
260
261fn validate_request(request: &IlpExactNaryRequest<'_>) -> Result<(u32, u32, u32)> {
262    let head_arity = request.head_arity.max(1) as usize;
263    if request.pos_values.len() % head_arity != 0 {
264        return Err(nary_err("pos_values length not a multiple of head_arity"));
265    }
266    if request.neg_values.len() % head_arity != 0 {
267        return Err(nary_err("neg_values length not a multiple of head_arity"));
268    }
269    let num_patterns = validate_batch(
270        &request.patterns(),
271        request.cand_value_offset,
272        request.cand_rows,
273        request.cand_values.len() as u64,
274    )?;
275    let num_pos = u32::try_from(request.pos_values.len() / head_arity)
276        .map_err(|_| nary_err("positive tuple count exceeds u32"))?;
277    let num_neg = u32::try_from(request.neg_values.len() / head_arity)
278        .map_err(|_| nary_err("negative tuple count exceeds u32"))?;
279    Ok((num_patterns, num_pos, num_neg))
280}
281
282fn u64_view<'a>(slice: &'a TrackedCudaSlice<u8>, elements: usize) -> RawCudaView<'a, u64> {
283    RawCudaView {
284        ptr: *slice.device_ptr(),
285        len: elements,
286        stream: slice.stream().clone(),
287        _marker: PhantomData,
288    }
289}
290
291/// One columnar u64 buffer requirement, validated fail-closed.
292fn require_u64_columns(buf: &CudaBuffer, label: &str) -> Result<(u32, u32)> {
293    let arity = u32::try_from(buf.arity()).map_err(|_| nary_err(format!("{label}: arity")))?;
294    if arity == 0 {
295        return Err(nary_err(format!("{label}: buffer has arity 0")));
296    }
297    for column in 0..buf.arity() {
298        match buf.schema().column_type(column) {
299            Some(ScalarType::U64) => {}
300            other => {
301                return Err(nary_err(format!(
302                    "{label}: column {column} is {other:?}, expected U64"
303                )));
304            }
305        }
306    }
307    let rows = buf
308        .cached_row_count()
309        .ok_or_else(|| nary_err(format!("{label}: cached_row_count absent")))?;
310    let rows = u32::try_from(rows).map_err(|_| nary_err(format!("{label}: row count")))?;
311    Ok((arity, rows))
312}
313
314impl super::CudaKernelProvider {
315    /// Score every flattened pattern against the example tuples on GPU,
316    /// from HOST slices.
317    ///
318    /// Returns `(pos_covered, neg_covered)`, one slot per pattern in
319    /// batch order. D2H budget: **2** counter-tracked transfers (one per
320    /// count array); all uploads are setup-phase H2D.
321    pub fn ilp_exact_nary_score(
322        &self,
323        request: &IlpExactNaryRequest<'_>,
324    ) -> Result<(Vec<u32>, Vec<u32>)> {
325        let (num_patterns, num_pos, num_neg) = validate_request(request)?;
326
327        macro_rules! upload_u64_bytes {
328            ($host:expr) => {{
329                let host: &[u64] = $host;
330                let bytes: Vec<u8> = host.iter().flat_map(|v| v.to_le_bytes()).collect();
331                let mut buf = self.memory.alloc::<u8>(bytes.len().max(1))?;
332                if !bytes.is_empty() {
333                    self.htod_sync_copy_into_tracked(&bytes, &mut buf)
334                        .map_err(|e| nary_err(format!("h2d u64 values: {e}")))?;
335                }
336                buf
337            }};
338        }
339
340        let cand_values_buf = upload_u64_bytes!(request.cand_values);
341        let pos_values_buf = upload_u64_bytes!(request.pos_values);
342        let neg_values_buf = upload_u64_bytes!(request.neg_values);
343
344        self.launch_nary(
345            &request.patterns(),
346            num_patterns,
347            num_pos,
348            num_neg,
349            request.cand_value_offset,
350            request.cand_rows,
351            u64_view(&cand_values_buf, request.cand_values.len()),
352            u64_view(&pos_values_buf, request.pos_values.len()),
353            u64_view(&neg_values_buf, request.neg_values.len()),
354        )
355    }
356
357    /// Score every flattened pattern with DEVICE-RESIDENT relations and
358    /// examples — the production ingest.
359    ///
360    /// `candidates[slot]` and the example buffers are columnar
361    /// [`CudaBuffer`]s with all-U64 columns; ingestion is one
362    /// device-to-device copy per column into the concatenated columnar
363    /// value buffers. No relation or example value crosses the host
364    /// boundary; the only D2H is the two count arrays (counter-tracked).
365    /// Example buffers must have arity == `patterns.head_arity`.
366    pub fn ilp_exact_nary_score_device(
367        &self,
368        patterns: &IlpExactNaryPatterns<'_>,
369        candidates: &[&CudaBuffer],
370        positives: &CudaBuffer,
371        negatives: &CudaBuffer,
372    ) -> Result<(Vec<u32>, Vec<u32>)> {
373        // Slot table from the buffers themselves.
374        let mut cand_value_offset: Vec<u32> = Vec::with_capacity(candidates.len());
375        let mut cand_rows: Vec<u32> = Vec::with_capacity(candidates.len());
376        let mut arities: Vec<u32> = Vec::with_capacity(candidates.len());
377        let mut total_elems: u32 = 0;
378        for (slot, buf) in candidates.iter().enumerate() {
379            let (arity, rows) = require_u64_columns(buf, &format!("candidate[{slot}]"))?;
380            cand_value_offset.push(total_elems);
381            cand_rows.push(rows);
382            arities.push(arity);
383            let elems = arity
384                .checked_mul(rows)
385                .and_then(|e| total_elems.checked_add(e))
386                .ok_or_else(|| nary_err("candidate value count exceeds u32"))?;
387            total_elems = elems;
388        }
389        let num_patterns =
390            validate_batch(patterns, &cand_value_offset, &cand_rows, total_elems as u64)?;
391
392        // The batch declares each atom's arity; the buffers know their own
393        // column count. A claimed arity wider than the relation reads
394        // straight past the slot into the NEXT relation's column and scores
395        // real-looking garbage, so the two must agree exactly.
396        for (atom, (&slot, &claimed)) in patterns
397            .atom_candidate_slot
398            .iter()
399            .zip(patterns.atom_arity.iter())
400            .enumerate()
401        {
402            let actual = arities.get(slot as usize).copied().ok_or_else(|| {
403                nary_err(format!("atom {atom}: candidate slot {slot} out of range"))
404            })?;
405            if claimed != actual {
406                return Err(nary_err(format!(
407                    "atom {atom} claims arity {claimed} for candidate slot \
408                     {slot}, but that relation has {actual} columns; the \
409                     kernel would read past the slot into the next relation",
410                )));
411            }
412        }
413
414        let (pos_arity, num_pos) = require_u64_columns(positives, "positives")?;
415        let (neg_arity, num_neg) = require_u64_columns(negatives, "negatives")?;
416        if pos_arity != patterns.head_arity {
417            return Err(nary_err(format!(
418                "positives arity {pos_arity} != head arity {}",
419                patterns.head_arity
420            )));
421        }
422        // Empty negatives are exempt from the arity match on purpose: a
423        // zero-row buffer is never dereferenced by the kernel, and callers
424        // legitimately pass a schema-arbitrary empty placeholder. Positives
425        // are checked unconditionally because an empty positive set is a
426        // real request shape whose head arity still defines the layout.
427        if neg_arity != patterns.head_arity && num_neg != 0 {
428            return Err(nary_err(format!(
429                "negatives arity {neg_arity} != head arity {}",
430                patterns.head_arity
431            )));
432        }
433
434        // ── D2D columnar concatenation (setup-phase, never host) ──────
435        let concat =
436            |bufs: &[(&CudaBuffer, u32, u32)], total: usize| -> Result<TrackedCudaSlice<u8>> {
437                let mut out = self.memory.alloc::<u8>((total * U64_SIZE).max(1))?;
438                let device = self.device.inner();
439                let mut element_offset: usize = 0;
440                for (buf, arity, rows) in bufs {
441                    let rows = *rows as usize;
442                    for column in 0..*arity as usize {
443                        if rows == 0 {
444                            continue;
445                        }
446                        let bytes = rows * U64_SIZE;
447                        let col = buf
448                            .column(column)
449                            .ok_or_else(|| nary_err(format!("missing column {column}")))?;
450                        let src = self.column_bytes_view(col, bytes)?;
451                        let byte_offset = element_offset * U64_SIZE;
452                        let mut dst = out.slice_mut(byte_offset..byte_offset + bytes);
453                        device
454                            .dtod_copy(&src, &mut dst)
455                            .map_err(|e| nary_err(format!("d2d column concat: {e}")))?;
456                        element_offset += rows;
457                    }
458                }
459                Ok(out)
460            };
461
462        let cand_triples: Vec<(&CudaBuffer, u32, u32)> = candidates
463            .iter()
464            .zip(arities.iter().zip(cand_rows.iter()))
465            .map(|(buf, (&a, &r))| (*buf, a, r))
466            .collect();
467        let cand_values_buf = concat(&cand_triples, total_elems as usize)?;
468        let pos_elems = (pos_arity as usize) * (num_pos as usize);
469        let pos_values_buf = concat(&[(positives, pos_arity, num_pos)], pos_elems)?;
470        let neg_elems = (neg_arity as usize) * (num_neg as usize);
471        let neg_values_buf = concat(&[(negatives, neg_arity, num_neg)], neg_elems)?;
472
473        self.launch_nary(
474            patterns,
475            num_patterns,
476            num_pos,
477            num_neg,
478            &cand_value_offset,
479            &cand_rows,
480            u64_view(&cand_values_buf, total_elems as usize),
481            u64_view(&pos_values_buf, pos_elems),
482            u64_view(&neg_values_buf, neg_elems),
483        )
484    }
485
486    /// Shared launch core: pack + upload the pattern batch and slot
487    /// table, launch, and read back the two count arrays.
488    #[allow(clippy::too_many_arguments)]
489    fn launch_nary(
490        &self,
491        patterns: &IlpExactNaryPatterns<'_>,
492        num_patterns: u32,
493        num_pos: u32,
494        num_neg: u32,
495        cand_value_offset: &[u32],
496        cand_rows: &[u32],
497        cand_values: RawCudaView<'_, u64>,
498        pos_values: RawCudaView<'_, u64>,
499        neg_values: RawCudaView<'_, u64>,
500    ) -> Result<(Vec<u32>, Vec<u32>)> {
501        let device = self.device.inner();
502
503        // Pack the six pattern arrays into ONE u32 buffer in the exact
504        // section order the kernel unpacks (see ilp_exact_nary.cu): the
505        // launch ABI caps the argument tuple, so the batch rides packed.
506        let atoms = patterns.atom_candidate_slot.len();
507        let mut batch_host: Vec<u32> = Vec::with_capacity(
508            2 * num_patterns as usize + 3 * atoms + patterns.binding_codes.len(),
509        );
510        batch_host.extend_from_slice(patterns.body_offset);
511        batch_host.extend_from_slice(patterns.body_len);
512        batch_host.extend_from_slice(patterns.atom_candidate_slot);
513        batch_host.extend_from_slice(patterns.atom_arity);
514        batch_host.extend_from_slice(patterns.atom_binding_offset);
515        batch_host.extend_from_slice(patterns.binding_codes);
516        let params_host: Vec<u32> = vec![
517            num_patterns,
518            u32::try_from(atoms).map_err(|_| nary_err("atom count exceeds u32"))?,
519            num_pos,
520            num_neg,
521            patterns.head_arity,
522        ];
523
524        macro_rules! upload_u32 {
525            ($name:ident, $host:expr) => {{
526                let host: &[u32] = $host;
527                let mut buf = self.memory.alloc::<u32>(host.len().max(1))?;
528                if !host.is_empty() {
529                    self.htod_sync_copy_into_tracked(host, &mut buf)
530                        .map_err(|e| {
531                            nary_err(format!(concat!("h2d ", stringify!($name), ": {}"), e))
532                        })?;
533                }
534                buf
535            }};
536        }
537
538        let batch_buf = upload_u32!(batch, &batch_host);
539        let params_buf = upload_u32!(params, &params_host);
540        let cand_value_offset_buf = upload_u32!(cand_value_offset, cand_value_offset);
541        let cand_rows_buf = upload_u32!(cand_rows, cand_rows);
542
543        let mut pos_covered_buf = self.memory.alloc::<u32>(num_patterns as usize)?;
544        let mut neg_covered_buf = self.memory.alloc::<u32>(num_patterns as usize)?;
545        // The kernel writes every pattern slot exactly once — no zero-init.
546
547        let func = device
548            .get_func(
549                ILP_EXACT_NARY_MODULE,
550                ilp_exact_nary_kernels::ILP_EXACT_NARY_SCORE,
551            )
552            .ok_or_else(|| nary_err("kernel not loaded"))?;
553        unsafe {
554            func.launch(
555                LaunchConfig {
556                    grid_dim: (num_patterns, 1, 1),
557                    block_dim: (ILP_EXACT_NARY_BLOCK_SIZE, 1, 1),
558                    shared_mem_bytes: 0,
559                },
560                (
561                    &batch_buf,
562                    &params_buf,
563                    &cand_values,
564                    &cand_value_offset_buf,
565                    &cand_rows_buf,
566                    &pos_values,
567                    &neg_values,
568                    &mut pos_covered_buf,
569                    &mut neg_covered_buf,
570                ),
571            )
572            .map_err(|e| nary_err(format!("launch: {e}")))?;
573        }
574        self.device.synchronize()?;
575
576        let mut pos_covered = vec![0u32; num_patterns as usize];
577        self.d2h_transfer_count.fetch_add(1, Ordering::Relaxed);
578        device
579            .dtoh_sync_copy_into(&pos_covered_buf, &mut pos_covered)
580            .map_err(|e| nary_err(format!("dtoh pos_covered: {e}")))?;
581        let mut neg_covered = vec![0u32; num_patterns as usize];
582        self.d2h_transfer_count.fetch_add(1, Ordering::Relaxed);
583        device
584            .dtoh_sync_copy_into(&neg_covered_buf, &mut neg_covered)
585            .map_err(|e| nary_err(format!("dtoh neg_covered: {e}")))?;
586        Ok((pos_covered, neg_covered))
587    }
588}
589
590#[cfg(test)]
591mod tests {
592    //! CUDA-gated correctness tests for the n-ary launcher, pinned to the
593    //! same hand-computed fixtures the host reference scorer freezes
594    //! (`xlog_induce::nary_reference`). Skipped without a GPU so the suite
595    //! still runs on a CPU-only host, but never under `XLOG_REQUIRE_CUDA=1`
596    //! (see [`make_provider`]); the pod leg runs them for real. All values
597    //! are COLUMNAR.
598
599    use std::sync::Arc;
600
601    use xlog_core::{MemoryBudget, ScalarType, Schema};
602
603    use super::{IlpExactNaryPatterns, IlpExactNaryRequest};
604    use crate::{CudaDevice, CudaKernelProvider, GpuMemoryManager};
605
606    /// Provider fixture for the launcher tests.
607    ///
608    /// Returns `None` when CUDA is unavailable so a CPU-only host can still
609    /// run the suite — except under `XLOG_REQUIRE_CUDA=1`, where a missing
610    /// device is a hard failure instead of a skip. Without that arm, a test
611    /// that returns early is indistinguishable from a passing one in cargo's
612    /// output, so `cargo test -p xlog-cuda --lib nary` reports five green
613    /// tests on a machine that never touched a GPU — which is how this
614    /// filter gets run when certifying the launcher on a rented device.
615    ///
616    /// The manual release gate exports `XLOG_REQUIRE_CUDA=1` and invokes this
617    /// crate's library tests alongside its `xlog-cli`, `xlog-prob`, pyxlog,
618    /// and `xlog-cuda-tests` legs, so a device-less release run fails here
619    /// instead of reporting a green skip.
620    ///
621    /// Mirrors `xlog_cuda_tests::harness::enforce_cuda_required`, which
622    /// already covers the integration-level `TestContext`; that crate is a
623    /// dev-dependency of the integration tests only and is not reachable
624    /// from this in-crate unit-test module.
625    fn make_provider() -> Option<CudaKernelProvider> {
626        fn skip_unless_required(
627            context: &str,
628            error: impl std::fmt::Display,
629        ) -> Option<CudaKernelProvider> {
630            if std::env::var("XLOG_REQUIRE_CUDA").as_deref() == Ok("1") {
631                panic!("XLOG_REQUIRE_CUDA=1 but CUDA is unavailable ({context}): {error}");
632            }
633            eprintln!("Skipping n-ary launcher test: CUDA unavailable ({context}): {error}");
634            None
635        }
636
637        let device = match CudaDevice::new(0) {
638            Ok(device) => Arc::new(device),
639            Err(error) => return skip_unless_required("CudaDevice::new", error),
640        };
641        let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
642        let memory = Arc::new(GpuMemoryManager::new(device.clone(), budget));
643        match CudaKernelProvider::new(device, memory) {
644            Ok(provider) => Some(provider),
645            Err(error) => skip_unless_required("CudaKernelProvider::new", error),
646        }
647    }
648
649    const JOIN: u32 = 0x8000_0000;
650
651    /// Build a columnar u64 buffer from per-column host arrays (mirrors
652    /// the binary launcher's test helper, generalized to N columns).
653    fn columnar_buffer(provider: &CudaKernelProvider, columns: &[&[u64]]) -> crate::CudaBuffer {
654        let rows = columns[0].len();
655        let schema = Schema::new(
656            (0..columns.len())
657                .map(|i| (format!("arg{i}"), ScalarType::U64))
658                .collect(),
659        );
660        if rows == 0 {
661            return provider.create_empty_buffer(schema).expect("empty buffer");
662        }
663        let device = provider.device().inner();
664        let mut device_columns = Vec::with_capacity(columns.len());
665        for column in columns {
666            assert_eq!(column.len(), rows, "ragged columns");
667            let bytes: Vec<u8> = column.iter().flat_map(|v| v.to_le_bytes()).collect();
668            let mut buf = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
669            device
670                .htod_sync_copy_into(&bytes, &mut buf)
671                .expect("h2d column");
672            device_columns.push(buf.into());
673        }
674        provider
675            .buffer_from_columns(device_columns, rows as u64, schema)
676            .expect("buffer_from_columns")
677    }
678
679    /// chain(L=0, R=1) over the shipped binary-kernel fixture:
680    /// p_B={(1,2),(2,3)}, p_C={(2,4),(3,5),(4,6)}, positives {(1,4),(2,5)},
681    /// negatives {(7,8)} — covers both positives, no negative.
682    #[test]
683    fn nary_kernel_matches_binary_chain_fixture() {
684        let Some(provider) = make_provider() else {
685            return;
686        };
687        let request = IlpExactNaryRequest {
688            body_offset: &[0],
689            body_len: &[2],
690            atom_candidate_slot: &[0, 1],
691            atom_arity: &[2, 2],
692            atom_binding_offset: &[0, 2],
693            // L(Head0, Join0), R(Join0, Head1)
694            binding_codes: &[0, JOIN, JOIN, 1],
695            // Columnar: p_B cols [1,2],[2,3]; p_C cols [2,3,4],[4,5,6].
696            cand_values: &[1, 2, 2, 3, 2, 3, 4, 4, 5, 6],
697            cand_value_offset: &[0, 4],
698            cand_rows: &[2, 3],
699            // Columnar examples: positives cols [1,2],[4,5].
700            pos_values: &[1, 2, 4, 5],
701            neg_values: &[7, 8],
702            head_arity: 2,
703        };
704        let (pos, neg) = provider.ilp_exact_nary_score(&request).unwrap();
705        assert_eq!(pos, vec![2]);
706        assert_eq!(neg, vec![0]);
707    }
708
709    /// The same chain fixture through the PRODUCTION device-buffer path:
710    /// candidates and examples live as columnar CudaBuffers and are
711    /// ingested with D2D column copies only.
712    #[test]
713    fn nary_device_ingest_matches_host_path() {
714        let Some(provider) = make_provider() else {
715            return;
716        };
717        let p_b = columnar_buffer(&provider, &[&[1, 2], &[2, 3]]);
718        let p_c = columnar_buffer(&provider, &[&[2, 3, 4], &[4, 5, 6]]);
719        let positives = columnar_buffer(&provider, &[&[1, 2], &[4, 5]]);
720        let negatives = columnar_buffer(&provider, &[&[7], &[8]]);
721        let patterns = IlpExactNaryPatterns {
722            body_offset: &[0],
723            body_len: &[2],
724            atom_candidate_slot: &[0, 1],
725            atom_arity: &[2, 2],
726            atom_binding_offset: &[0, 2],
727            binding_codes: &[0, JOIN, JOIN, 1],
728            head_arity: 2,
729        };
730        let (pos, neg) = provider
731            .ilp_exact_nary_score_device(&patterns, &[&p_b, &p_c], &positives, &negatives)
732            .unwrap();
733        assert_eq!(pos, vec![2]);
734        assert_eq!(neg, vec![0]);
735    }
736
737    /// Ternary fixture from the reference suite: H(x0,x1,x2) :-
738    /// T(x0,x1,z0), P(z0,x2) with T={(1,2,9),(4,5,8)}, P={(9,3),(8,7)}.
739    /// Of positives {(1,2,3),(4,5,6),(1,2,7)} exactly one is covered.
740    #[test]
741    fn nary_kernel_matches_ternary_reference_fixture() {
742        let Some(provider) = make_provider() else {
743            return;
744        };
745        let request = IlpExactNaryRequest {
746            body_offset: &[0],
747            body_len: &[2],
748            atom_candidate_slot: &[0, 1],
749            atom_arity: &[3, 2],
750            atom_binding_offset: &[0, 3],
751            binding_codes: &[0, 1, JOIN, JOIN, 2],
752            // Columnar: T cols [1,4],[2,5],[9,8]; P cols [9,8],[3,7].
753            cand_values: &[1, 4, 2, 5, 9, 8, 9, 8, 3, 7],
754            cand_value_offset: &[0, 6],
755            cand_rows: &[2, 2],
756            // Columnar examples: cols [1,4,1],[2,5,2],[3,6,7].
757            pos_values: &[1, 4, 1, 2, 5, 2, 3, 6, 7],
758            neg_values: &[],
759            head_arity: 3,
760        };
761        let (pos, neg) = provider.ilp_exact_nary_score(&request).unwrap();
762        assert_eq!(pos, vec![1]);
763        assert_eq!(neg, vec![0]);
764    }
765
766    /// Backtracking fixture: T={(1,8),(1,9)}, P={(9,2)} — the first T row
767    /// dead-ends and the kernel must revisit T at row 2 to find the cover.
768    #[test]
769    fn nary_kernel_backtracks_across_atoms() {
770        let Some(provider) = make_provider() else {
771            return;
772        };
773        let request = IlpExactNaryRequest {
774            body_offset: &[0],
775            body_len: &[2],
776            atom_candidate_slot: &[0, 1],
777            atom_arity: &[2, 2],
778            atom_binding_offset: &[0, 2],
779            binding_codes: &[0, JOIN, JOIN, 1],
780            // Columnar: T cols [1,1],[8,9]; P single row [9,2].
781            cand_values: &[1, 1, 8, 9, 9, 2],
782            cand_value_offset: &[0, 4],
783            cand_rows: &[2, 1],
784            // Columnar examples: positives {(1,2),(1,3)} -> cols [1,1],[2,3].
785            pos_values: &[1, 1, 2, 3],
786            neg_values: &[],
787            head_arity: 2,
788        };
789        let (pos, neg) = provider.ilp_exact_nary_score(&request).unwrap();
790        assert_eq!(pos, vec![1]);
791        assert_eq!(neg, vec![0]);
792    }
793
794    #[test]
795    fn validation_refuses_malformed_batches_without_a_device() {
796        // Validation is host-side and must refuse BEFORE any CUDA work,
797        // so these run everywhere.
798        use super::validate_request;
799
800        let base = IlpExactNaryRequest {
801            body_offset: &[0],
802            body_len: &[1],
803            atom_candidate_slot: &[0],
804            atom_arity: &[2],
805            atom_binding_offset: &[0],
806            binding_codes: &[0, 1],
807            cand_values: &[1, 2],
808            cand_value_offset: &[0],
809            cand_rows: &[1],
810            pos_values: &[1, 2],
811            neg_values: &[],
812            head_arity: 2,
813        };
814        assert!(validate_request(&base).is_ok());
815
816        let empty = IlpExactNaryRequest {
817            body_offset: &[],
818            body_len: &[],
819            ..base
820        };
821        assert!(validate_request(&empty).is_err());
822
823        let bad_slot = IlpExactNaryRequest {
824            atom_candidate_slot: &[5],
825            ..base
826        };
827        assert!(validate_request(&bad_slot).is_err());
828
829        let short_values = IlpExactNaryRequest {
830            cand_rows: &[9],
831            ..base
832        };
833        assert!(validate_request(&short_values).is_err());
834
835        let head_out_of_range = IlpExactNaryRequest {
836            binding_codes: &[0, 7],
837            ..base
838        };
839        assert!(validate_request(&head_out_of_range).is_err());
840
841        let ragged_tuples = IlpExactNaryRequest {
842            pos_values: &[1, 2, 3],
843            ..base
844        };
845        assert!(validate_request(&ragged_tuples).is_err());
846
847        let wide_head = IlpExactNaryRequest {
848            head_arity: 9,
849            pos_values: &[1, 2, 3, 4, 5, 6, 7, 8, 9],
850            binding_codes: &[0, 1],
851            ..base
852        };
853        assert!(validate_request(&wide_head).is_err());
854    }
855}
xlog_cuda/provider/ilp_exact_nary.rs

xlog_cuda/provider/
ilp_exact_nary.rs
