xlog_runtime/
ilp_registry.rs

1//! ILP (Inductive Logic Programming) registry for tensor mask management.
2
3use std::collections::HashMap;
4use xlog_core::XlogError;
5use xlog_cuda::{CudaBuffer, CudaKernelProvider};
6
7/// Reads the device-side row count through the provider's public metadata API.
8pub fn read_device_row_count(
9    provider: &CudaKernelProvider,
10    buffer: &CudaBuffer,
11) -> Result<usize, XlogError> {
12    provider.device_row_count(buffer)
13}
14
15/// Registry for ILP tensor masks.
16pub struct IlpRegistry {
17    masks: HashMap<String, IlpMask>,
18}
19
20/// A registered ILP mask — Dense (imported via DLPack) or Sparse (candidate entries only).
21#[allow(clippy::large_enum_variant)]
22pub enum IlpMask {
23    /// Dense mask with hard and soft weight buffers.
24    Dense {
25        /// Hard mask (binary on/off).
26        hard: CudaBuffer,
27        /// Soft mask (continuous weights).
28        soft: CudaBuffer,
29        /// Number of body relations in the schema.
30        schema_size: usize,
31    },
32    /// Sparse mask listing active (i,j,k) entries on host.
33    Sparse {
34        /// Active rule entries.
35        active_entries: Vec<(u32, u32, u32)>,
36        /// Number of body relations in the schema.
37        schema_size: usize,
38    },
39    /// Sparse mask with device-resident active flags.
40    SparseDevice {
41        /// Candidate (i,j,k) entries in evaluation order.
42        candidate_order: Vec<(u32, u32, u32)>,
43        /// Device-resident boolean flags per candidate.
44        active_flags: CudaBuffer,
45        /// Number of currently selected candidates.
46        selected_count: usize,
47        /// Number of body relations in the schema.
48        schema_size: usize,
49    },
50}
51
52impl IlpMask {
53    /// Return the schema size (number of body relations).
54    pub fn schema_size(&self) -> usize {
55        match self {
56            IlpMask::Dense { schema_size, .. } => *schema_size,
57            IlpMask::Sparse { schema_size, .. } => *schema_size,
58            IlpMask::SparseDevice { schema_size, .. } => *schema_size,
59        }
60    }
61}
62
63/// Tag metadata from TensorMaskedJoin execution.
64/// Retains per-entry projected join buffers for batch credit queries.
65pub struct IlpTaggedResult {
66    /// Per-entry metadata and result buffers.
67    pub entries: Vec<IlpTagEntry>,
68}
69
70/// Metadata for a single active rule (i,j,k), its result cardinality,
71/// and the projected join result buffer (retained for batch credit queries).
72pub struct IlpTagEntry {
73    /// First body relation index.
74    pub i: u32,
75    /// Second body relation index.
76    pub j: u32,
77    /// Head relation index.
78    pub k: u32,
79    /// Number of result rows for this entry.
80    pub num_rows: u32,
81    /// The projected join result buffer, retained for batch credit queries.
82    pub buffer: Option<CudaBuffer>,
83}
84
85impl IlpRegistry {
86    /// Create an empty ILP registry.
87    pub fn new() -> Self {
88        Self {
89            masks: HashMap::new(),
90        }
91    }
92
93    /// Clear all registered masks, releasing GPU buffers.
94    pub fn clear(&mut self) {
95        self.masks.clear();
96    }
97
98    /// Register a dense ILP mask (hard + soft weight buffers).
99    pub fn insert_mask(
100        &mut self,
101        name: String,
102        hard: CudaBuffer,
103        soft: CudaBuffer,
104        schema_size: usize,
105    ) {
106        self.masks.insert(
107            name,
108            IlpMask::Dense {
109                hard,
110                soft,
111                schema_size,
112            },
113        );
114    }
115
116    /// Insert a mask built from sparse candidate data.
117    ///
118    /// Performs deterministic top-k ranking (desc soft value, then lower index)
119    /// and stores the selected (i,j,k) entries directly — no dense buffer.
120    pub fn insert_mask_from_sparse(
121        &mut self,
122        name: String,
123        schema_size: usize,
124        active_ijk: &[(u32, u32, u32)],
125        active_soft: &[f32],
126        budget: usize,
127    ) -> Result<(), XlogError> {
128        if active_ijk.len() != active_soft.len() {
129            return Err(XlogError::Execution(format!(
130                "active_ijk length {} != active_soft length {}",
131                active_ijk.len(),
132                active_soft.len()
133            )));
134        }
135
136        // Deterministic top-k over positive-probability candidates: descending
137        // soft value, then ascending index for ties.
138        let mut ranked: Vec<(usize, f32)> = active_soft.iter().copied().enumerate().collect();
139        ranked.retain(|(_, soft)| *soft > 0.0);
140        ranked.sort_by(|a, b| {
141            b.1.partial_cmp(&a.1)
142                .unwrap_or(std::cmp::Ordering::Equal)
143                .then(a.0.cmp(&b.0))
144        });
145        ranked.truncate(budget.min(ranked.len()));
146
147        let entries: Vec<(u32, u32, u32)> =
148            ranked.iter().map(|&(idx, _)| active_ijk[idx]).collect();
149
150        self.masks.insert(
151            name,
152            IlpMask::Sparse {
153                active_entries: entries,
154                schema_size,
155            },
156        );
157        Ok(())
158    }
159
160    /// Insert an already-selected sparse mask, preserving caller order exactly.
161    pub fn insert_selected_mask(
162        &mut self,
163        name: String,
164        schema_size: usize,
165        active_entries: &[(u32, u32, u32)],
166    ) {
167        self.masks.insert(
168            name,
169            IlpMask::Sparse {
170                active_entries: active_entries.to_vec(),
171                schema_size,
172            },
173        );
174    }
175
176    /// Insert a device-resident sparse mask with active flags on GPU.
177    pub fn insert_selected_mask_device(
178        &mut self,
179        name: String,
180        schema_size: usize,
181        candidate_order: Vec<(u32, u32, u32)>,
182        active_flags: CudaBuffer,
183        selected_count: usize,
184    ) {
185        self.masks.insert(
186            name,
187            IlpMask::SparseDevice {
188                candidate_order,
189                active_flags,
190                selected_count,
191                schema_size,
192            },
193        );
194    }
195
196    /// Look up a registered mask by name.
197    pub fn get_mask(&self, name: &str) -> Option<&IlpMask> {
198        self.masks.get(name)
199    }
200
201    /// Returns true if any registered mask uses the sparse-device representation.
202    pub fn has_sparse_device_mask(&self) -> bool {
203        self.masks
204            .values()
205            .any(|mask| matches!(mask, IlpMask::SparseDevice { .. }))
206    }
207}
208
209impl Default for IlpRegistry {
210    fn default() -> Self {
211        Self::new()
212    }
213}
xlog_runtime/ilp_registry.rs

xlog_runtime/
ilp_registry.rs
