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 crate::memory::DeviceMemoryView;
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, 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().is_multiple_of(head_arity) {
264        return Err(nary_err("pos_values length not a multiple of head_arity"));
265    }
266    if !request.neg_values.len().is_multiple_of(head_arity) {
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(slice: &TrackedCudaSlice<u8>, elements: usize) -> DeviceMemoryView<u64> {
283    // SAFETY: the arena stores u64 values and is checked before this view.
284    unsafe { slice.view().cast::<u64>() }
285        .and_then(|view| view.try_slice(..elements))
286        .expect("validated u64 arena extent")
287}
288
289/// One columnar u64 buffer requirement, validated fail-closed.
290fn require_u64_columns(buf: &CudaBuffer, label: &str) -> Result<(u32, u32)> {
291    let arity = u32::try_from(buf.arity()).map_err(|_| nary_err(format!("{label}: arity")))?;
292    if arity == 0 {
293        return Err(nary_err(format!("{label}: buffer has arity 0")));
294    }
295    for column in 0..buf.arity() {
296        match buf.schema().column_type(column) {
297            Some(ScalarType::U64) => {}
298            other => {
299                return Err(nary_err(format!(
300                    "{label}: column {column} is {other:?}, expected U64"
301                )));
302            }
303        }
304    }
305    let rows = buf
306        .cached_row_count()
307        .ok_or_else(|| nary_err(format!("{label}: cached_row_count absent")))?;
308    Ok((arity, rows))
309}
310
311impl super::CudaKernelProvider {
312    /// Score every flattened pattern against the example tuples on GPU,
313    /// from HOST slices.
314    ///
315    /// Returns `(pos_covered, neg_covered)`, one slot per pattern in
316    /// batch order. D2H budget: **2** counter-tracked transfers (one per
317    /// count array); all uploads are setup-phase H2D.
318    pub fn ilp_exact_nary_score(
319        &self,
320        request: &IlpExactNaryRequest<'_>,
321    ) -> Result<(Vec<u32>, Vec<u32>)> {
322        let (num_patterns, num_pos, num_neg) = validate_request(request)?;
323
324        macro_rules! upload_u64_bytes {
325            ($host:expr) => {{
326                let host: &[u64] = $host;
327                let bytes: Vec<u8> = host.iter().flat_map(|v| v.to_le_bytes()).collect();
328                let mut buf = self.memory.alloc::<u8>(bytes.len().max(U64_SIZE))?;
329                if !bytes.is_empty() {
330                    self.htod_sync_copy_into_tracked(&bytes, &mut buf)
331                        .map_err(|e| nary_err(format!("h2d u64 values: {e}")))?;
332                }
333                buf
334            }};
335        }
336
337        let cand_values_buf = upload_u64_bytes!(request.cand_values);
338        let pos_values_buf = upload_u64_bytes!(request.pos_values);
339        let neg_values_buf = upload_u64_bytes!(request.neg_values);
340
341        self.launch_nary(
342            &request.patterns(),
343            num_patterns,
344            num_pos,
345            num_neg,
346            request.cand_value_offset,
347            request.cand_rows,
348            u64_view(&cand_values_buf, request.cand_values.len()),
349            u64_view(&pos_values_buf, request.pos_values.len()),
350            u64_view(&neg_values_buf, request.neg_values.len()),
351        )
352    }
353
354    /// Score every flattened pattern with DEVICE-RESIDENT relations and
355    /// examples — the production ingest.
356    ///
357    /// `candidates[slot]` and the example buffers are columnar
358    /// [`CudaBuffer`]s with all-U64 columns; ingestion is one
359    /// device-to-device copy per column into the concatenated columnar
360    /// value buffers. No relation or example value crosses the host
361    /// boundary; the only D2H is the two count arrays (counter-tracked).
362    /// Example buffers must have arity == `patterns.head_arity`.
363    pub fn ilp_exact_nary_score_device(
364        &self,
365        patterns: &IlpExactNaryPatterns<'_>,
366        candidates: &[&CudaBuffer],
367        positives: &CudaBuffer,
368        negatives: &CudaBuffer,
369    ) -> Result<(Vec<u32>, Vec<u32>)> {
370        // Slot table from the buffers themselves.
371        let mut cand_value_offset: Vec<u32> = Vec::with_capacity(candidates.len());
372        let mut cand_rows: Vec<u32> = Vec::with_capacity(candidates.len());
373        let mut arities: Vec<u32> = Vec::with_capacity(candidates.len());
374        let mut total_elems: u32 = 0;
375        for (slot, buf) in candidates.iter().enumerate() {
376            let (arity, rows) = require_u64_columns(buf, &format!("candidate[{slot}]"))?;
377            cand_value_offset.push(total_elems);
378            cand_rows.push(rows);
379            arities.push(arity);
380            let elems = arity
381                .checked_mul(rows)
382                .and_then(|e| total_elems.checked_add(e))
383                .ok_or_else(|| nary_err("candidate value count exceeds u32"))?;
384            total_elems = elems;
385        }
386        let num_patterns =
387            validate_batch(patterns, &cand_value_offset, &cand_rows, total_elems as u64)?;
388
389        // The batch declares each atom's arity; the buffers know their own
390        // column count. A claimed arity wider than the relation reads
391        // straight past the slot into the NEXT relation's column and scores
392        // real-looking garbage, so the two must agree exactly.
393        for (atom, (&slot, &claimed)) in patterns
394            .atom_candidate_slot
395            .iter()
396            .zip(patterns.atom_arity.iter())
397            .enumerate()
398        {
399            let actual = arities.get(slot as usize).copied().ok_or_else(|| {
400                nary_err(format!("atom {atom}: candidate slot {slot} out of range"))
401            })?;
402            if claimed != actual {
403                return Err(nary_err(format!(
404                    "atom {atom} claims arity {claimed} for candidate slot \
405                     {slot}, but that relation has {actual} columns; the \
406                     kernel would read past the slot into the next relation",
407                )));
408            }
409        }
410
411        let (pos_arity, num_pos) = require_u64_columns(positives, "positives")?;
412        let (neg_arity, num_neg) = require_u64_columns(negatives, "negatives")?;
413        if pos_arity != patterns.head_arity {
414            return Err(nary_err(format!(
415                "positives arity {pos_arity} != head arity {}",
416                patterns.head_arity
417            )));
418        }
419        // Empty negatives are exempt from the arity match on purpose: a
420        // zero-row buffer is never dereferenced by the kernel, and callers
421        // legitimately pass a schema-arbitrary empty placeholder. Positives
422        // are checked unconditionally because an empty positive set is a
423        // real request shape whose head arity still defines the layout.
424        if neg_arity != patterns.head_arity && num_neg != 0 {
425            return Err(nary_err(format!(
426                "negatives arity {neg_arity} != head arity {}",
427                patterns.head_arity
428            )));
429        }
430
431        // ── D2D columnar concatenation (setup-phase, never host) ──────
432        let concat =
433            |bufs: &[(&CudaBuffer, u32, u32)], total: usize| -> Result<TrackedCudaSlice<u8>> {
434                let mut out = self.memory.alloc::<u8>((total * U64_SIZE).max(U64_SIZE))?;
435                let device = self.device.inner();
436                let mut element_offset: usize = 0;
437                for (buf, arity, rows) in bufs {
438                    let rows = *rows as usize;
439                    for column in 0..*arity as usize {
440                        if rows == 0 {
441                            continue;
442                        }
443                        let bytes = rows * U64_SIZE;
444                        let col = buf
445                            .column(column)
446                            .ok_or_else(|| nary_err(format!("missing column {column}")))?;
447                        let src = self.column_bytes_view(col, bytes)?;
448                        let byte_offset = element_offset * U64_SIZE;
449                        let mut dst = out.slice_mut(byte_offset..byte_offset + bytes);
450                        device
451                            .dtod_copy(&src, &mut dst)
452                            .map_err(|e| nary_err(format!("d2d column concat: {e}")))?;
453                        element_offset += rows;
454                    }
455                }
456                Ok(out)
457            };
458
459        let cand_triples: Vec<(&CudaBuffer, u32, u32)> = candidates
460            .iter()
461            .zip(arities.iter().zip(cand_rows.iter()))
462            .map(|(buf, (&a, &r))| (*buf, a, r))
463            .collect();
464        let cand_values_buf = concat(&cand_triples, total_elems as usize)?;
465        let pos_elems = (pos_arity as usize) * (num_pos as usize);
466        let pos_values_buf = concat(&[(positives, pos_arity, num_pos)], pos_elems)?;
467        let neg_elems = (neg_arity as usize) * (num_neg as usize);
468        let neg_values_buf = concat(&[(negatives, neg_arity, num_neg)], neg_elems)?;
469
470        self.launch_nary(
471            patterns,
472            num_patterns,
473            num_pos,
474            num_neg,
475            &cand_value_offset,
476            &cand_rows,
477            u64_view(&cand_values_buf, total_elems as usize),
478            u64_view(&pos_values_buf, pos_elems),
479            u64_view(&neg_values_buf, neg_elems),
480        )
481    }
482
483    /// Shared launch core: pack + upload the pattern batch and slot
484    /// table, launch, and read back the two count arrays.
485    #[expect(
486        clippy::too_many_arguments,
487        reason = "the launch core mirrors the fixed n-ary kernel inputs and its explicit output counters"
488    )]
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: DeviceMemoryView<u64>,
498        pos_values: DeviceMemoryView<u64>,
499        neg_values: DeviceMemoryView<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 xlog_core::{MemoryBudget, ScalarType, Schema};
600
601    use super::{IlpExactNaryPatterns, IlpExactNaryRequest};
602    use crate::CudaKernelProvider;
603
604    /// Provider fixture for the launcher tests.
605    ///
606    /// Returns `None` when CUDA is unavailable so a CPU-only host can still
607    /// run the suite — except under `XLOG_REQUIRE_CUDA=1`, where a missing
608    /// device is a hard failure instead of a skip. Without that arm, a test
609    /// that returns early is indistinguishable from a passing one in cargo's
610    /// output, so `cargo test -p xlog-cuda --lib nary` reports five green
611    /// tests on a machine that never touched a GPU — which is how this
612    /// filter gets run when certifying the launcher on a rented device.
613    ///
614    /// The manual release gate exports `XLOG_REQUIRE_CUDA=1` and invokes this
615    /// crate's library tests alongside its `xlog-cli`, `xlog-prob`, pyxlog,
616    /// and `xlog-cuda-tests` legs, so a device-less release run fails here
617    /// instead of reporting a green skip.
618    ///
619    /// Mirrors `xlog_cuda_tests::harness::enforce_cuda_required`, which
620    /// already covers the integration-level `TestContext`; that crate is a
621    /// dev-dependency of the integration tests only and is not reachable
622    /// from this in-crate unit-test module.
623    fn make_provider() -> Option<CudaKernelProvider> {
624        fn skip_unless_required(
625            context: &str,
626            error: impl std::fmt::Display,
627        ) -> Option<CudaKernelProvider> {
628            if std::env::var("XLOG_REQUIRE_CUDA").as_deref() == Ok("1") {
629                panic!("XLOG_REQUIRE_CUDA=1 but CUDA is unavailable ({context}): {error}");
630            }
631            eprintln!("Skipping n-ary launcher test: CUDA unavailable ({context}): {error}");
632            None
633        }
634
635        let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
636        match crate::CudaProviderBuilder::new(0, budget).build() {
637            Ok(provider) => Some(provider),
638            Err(error) => skip_unless_required("CudaProviderBuilder::build", error),
639        }
640    }
641
642    const JOIN: u32 = 0x8000_0000;
643
644    /// Build a columnar u64 buffer from per-column host arrays (mirrors
645    /// the binary launcher's test helper, generalized to N columns).
646    fn columnar_buffer(provider: &CudaKernelProvider, columns: &[&[u64]]) -> crate::CudaBuffer {
647        let rows = columns[0].len();
648        let schema = Schema::new(
649            (0..columns.len())
650                .map(|i| (format!("arg{i}"), ScalarType::U64))
651                .collect(),
652        );
653        if rows == 0 {
654            return provider.create_empty_buffer(schema).expect("empty buffer");
655        }
656        let device = provider.device().inner();
657        let mut device_columns = Vec::with_capacity(columns.len());
658        for column in columns {
659            assert_eq!(column.len(), rows, "ragged columns");
660            let bytes: Vec<u8> = column.iter().flat_map(|v| v.to_le_bytes()).collect();
661            let mut buf = provider.memory().alloc::<u8>(bytes.len()).expect("alloc");
662            device
663                .htod_sync_copy_into(&bytes, &mut buf)
664                .expect("h2d column");
665            device_columns.push(buf.into());
666        }
667        provider
668            .buffer_from_columns(device_columns, rows as u64, schema)
669            .expect("buffer_from_columns")
670    }
671
672    /// chain(L=0, R=1) over the shipped binary-kernel fixture:
673    /// p_B={(1,2),(2,3)}, p_C={(2,4),(3,5),(4,6)}, positives {(1,4),(2,5)},
674    /// negatives {(7,8)} — covers both positives, no negative.
675    #[test]
676    fn nary_kernel_matches_binary_chain_fixture() {
677        let Some(provider) = make_provider() else {
678            return;
679        };
680        let request = IlpExactNaryRequest {
681            body_offset: &[0],
682            body_len: &[2],
683            atom_candidate_slot: &[0, 1],
684            atom_arity: &[2, 2],
685            atom_binding_offset: &[0, 2],
686            // L(Head0, Join0), R(Join0, Head1)
687            binding_codes: &[0, JOIN, JOIN, 1],
688            // Columnar: p_B cols [1,2],[2,3]; p_C cols [2,3,4],[4,5,6].
689            cand_values: &[1, 2, 2, 3, 2, 3, 4, 4, 5, 6],
690            cand_value_offset: &[0, 4],
691            cand_rows: &[2, 3],
692            // Columnar examples: positives cols [1,2],[4,5].
693            pos_values: &[1, 2, 4, 5],
694            neg_values: &[7, 8],
695            head_arity: 2,
696        };
697        let (pos, neg) = provider.ilp_exact_nary_score(&request).unwrap();
698        assert_eq!(pos, vec![2]);
699        assert_eq!(neg, vec![0]);
700    }
701
702    /// The same chain fixture through the PRODUCTION device-buffer path:
703    /// candidates and examples live as columnar CudaBuffers and are
704    /// ingested with D2D column copies only.
705    #[test]
706    fn nary_device_ingest_matches_host_path() {
707        let Some(provider) = make_provider() else {
708            return;
709        };
710        let p_b = columnar_buffer(&provider, &[&[1, 2], &[2, 3]]);
711        let p_c = columnar_buffer(&provider, &[&[2, 3, 4], &[4, 5, 6]]);
712        let positives = columnar_buffer(&provider, &[&[1, 2], &[4, 5]]);
713        let negatives = columnar_buffer(&provider, &[&[7], &[8]]);
714        let patterns = IlpExactNaryPatterns {
715            body_offset: &[0],
716            body_len: &[2],
717            atom_candidate_slot: &[0, 1],
718            atom_arity: &[2, 2],
719            atom_binding_offset: &[0, 2],
720            binding_codes: &[0, JOIN, JOIN, 1],
721            head_arity: 2,
722        };
723        let (pos, neg) = provider
724            .ilp_exact_nary_score_device(&patterns, &[&p_b, &p_c], &positives, &negatives)
725            .unwrap();
726        assert_eq!(pos, vec![2]);
727        assert_eq!(neg, vec![0]);
728    }
729
730    /// Ternary fixture from the reference suite: H(x0,x1,x2) :-
731    /// T(x0,x1,z0), P(z0,x2) with T={(1,2,9),(4,5,8)}, P={(9,3),(8,7)}.
732    /// Of positives {(1,2,3),(4,5,6),(1,2,7)} exactly one is covered.
733    #[test]
734    fn nary_kernel_matches_ternary_reference_fixture() {
735        let Some(provider) = make_provider() else {
736            return;
737        };
738        let request = IlpExactNaryRequest {
739            body_offset: &[0],
740            body_len: &[2],
741            atom_candidate_slot: &[0, 1],
742            atom_arity: &[3, 2],
743            atom_binding_offset: &[0, 3],
744            binding_codes: &[0, 1, JOIN, JOIN, 2],
745            // Columnar: T cols [1,4],[2,5],[9,8]; P cols [9,8],[3,7].
746            cand_values: &[1, 4, 2, 5, 9, 8, 9, 8, 3, 7],
747            cand_value_offset: &[0, 6],
748            cand_rows: &[2, 2],
749            // Columnar examples: cols [1,4,1],[2,5,2],[3,6,7].
750            pos_values: &[1, 4, 1, 2, 5, 2, 3, 6, 7],
751            neg_values: &[],
752            head_arity: 3,
753        };
754        let (pos, neg) = provider.ilp_exact_nary_score(&request).unwrap();
755        assert_eq!(pos, vec![1]);
756        assert_eq!(neg, vec![0]);
757    }
758
759    /// Backtracking fixture: T={(1,8),(1,9)}, P={(9,2)} — the first T row
760    /// dead-ends and the kernel must revisit T at row 2 to find the cover.
761    #[test]
762    fn nary_kernel_backtracks_across_atoms() {
763        let Some(provider) = make_provider() else {
764            return;
765        };
766        let request = IlpExactNaryRequest {
767            body_offset: &[0],
768            body_len: &[2],
769            atom_candidate_slot: &[0, 1],
770            atom_arity: &[2, 2],
771            atom_binding_offset: &[0, 2],
772            binding_codes: &[0, JOIN, JOIN, 1],
773            // Columnar: T cols [1,1],[8,9]; P single row [9,2].
774            cand_values: &[1, 1, 8, 9, 9, 2],
775            cand_value_offset: &[0, 4],
776            cand_rows: &[2, 1],
777            // Columnar examples: positives {(1,2),(1,3)} -> cols [1,1],[2,3].
778            pos_values: &[1, 1, 2, 3],
779            neg_values: &[],
780            head_arity: 2,
781        };
782        let (pos, neg) = provider.ilp_exact_nary_score(&request).unwrap();
783        assert_eq!(pos, vec![1]);
784        assert_eq!(neg, vec![0]);
785    }
786
787    #[test]
788    fn validation_refuses_malformed_batches_without_a_device() {
789        // Validation is host-side and must refuse BEFORE any CUDA work,
790        // so these run everywhere.
791        use super::validate_request;
792
793        let base = IlpExactNaryRequest {
794            body_offset: &[0],
795            body_len: &[1],
796            atom_candidate_slot: &[0],
797            atom_arity: &[2],
798            atom_binding_offset: &[0],
799            binding_codes: &[0, 1],
800            cand_values: &[1, 2],
801            cand_value_offset: &[0],
802            cand_rows: &[1],
803            pos_values: &[1, 2],
804            neg_values: &[],
805            head_arity: 2,
806        };
807        assert!(validate_request(&base).is_ok());
808
809        let empty = IlpExactNaryRequest {
810            body_offset: &[],
811            body_len: &[],
812            ..base
813        };
814        assert!(validate_request(&empty).is_err());
815
816        let bad_slot = IlpExactNaryRequest {
817            atom_candidate_slot: &[5],
818            ..base
819        };
820        assert!(validate_request(&bad_slot).is_err());
821
822        let short_values = IlpExactNaryRequest {
823            cand_rows: &[9],
824            ..base
825        };
826        assert!(validate_request(&short_values).is_err());
827
828        let head_out_of_range = IlpExactNaryRequest {
829            binding_codes: &[0, 7],
830            ..base
831        };
832        assert!(validate_request(&head_out_of_range).is_err());
833
834        let ragged_tuples = IlpExactNaryRequest {
835            pos_values: &[1, 2, 3],
836            ..base
837        };
838        assert!(validate_request(&ragged_tuples).is_err());
839
840        let wide_head = IlpExactNaryRequest {
841            head_arity: 9,
842            pos_values: &[1, 2, 3, 4, 5, 6, 7, 8, 9],
843            binding_codes: &[0, 1],
844            ..base
845        };
846        assert!(validate_request(&wide_head).is_err());
847    }
848}
xlog_cuda/provider/ilp_exact_nary.rs

xlog_cuda/provider/
ilp_exact_nary.rs
