1use crate::memory::DeviceRead;
10use std::sync::atomic::Ordering;
11
12use crate::{LaunchAsync, LaunchConfig};
13use xlog_core::{resolve_bool, Result, ScalarType, XlogError};
14
15use super::{ilp_exact_kernels, ILP_EXACT_MODULE};
16use crate::memory::{CudaBuffer, TrackedCudaSlice};
17
18const ILP_EXACT_BLOCK_SIZE: u32 = 256;
19const ILP_EXACT_TOPK_FIELDS: usize = 9;
20const ENV_ILP_EXACT_CHAIN_SMEM: &str = "XLOG_ILP_EXACT_CHAIN_SMEM";
21const ENV_ILP_EXACT_CHAIN_SMEM_MIN_ROWS: &str = "XLOG_ILP_EXACT_CHAIN_SMEM_MIN_ROWS";
22const DEFAULT_ILP_EXACT_CHAIN_SMEM_MIN_ROWS: u32 = 256;
23
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
25pub struct IlpExactTopkCandidate {
26 pub topology_idx: u32,
27 pub left_idx: u32,
28 pub right_idx: u32,
29 pub positives_covered: u32,
30 pub negatives_covered: u32,
31 pub local_rank: u32,
32 pub next_positives_covered: u32,
33 pub next_negatives_covered: u32,
34 pub tie_class_size: u32,
35}
36
37struct IlpExactDeviceScores {
38 candidate_count: usize,
39 #[cfg(test)]
40 slot_count: usize,
41 pos_covered: TrackedCudaSlice<u32>,
42 neg_covered: TrackedCudaSlice<u32>,
43}
44
45#[derive(Clone, Copy, Debug, Eq, PartialEq)]
46enum ExactPairLayout {
47 U64,
48 U32,
49 Symbol,
50}
51
52impl ExactPairLayout {
53 fn elem_size(self) -> usize {
54 match self {
55 Self::U64 => std::mem::size_of::<u64>(),
56 Self::U32 | Self::Symbol => std::mem::size_of::<u32>(),
57 }
58 }
59}
60
61fn ilp_exact_chain_smem_enabled() -> Result<bool> {
62 resolve_bool(None, ENV_ILP_EXACT_CHAIN_SMEM, true)
63}
64
65fn chain_smem_shared_bytes(layout: ExactPairLayout) -> u32 {
66 let block = ILP_EXACT_BLOCK_SIZE as usize;
67 let bytes = (2usize * block * layout.elem_size()) + (block * std::mem::size_of::<u32>());
68 u32::try_from(bytes).expect("chain smem byte count fits in u32")
69}
70
71fn ilp_exact_chain_smem_min_rows() -> u32 {
72 std::env::var(ENV_ILP_EXACT_CHAIN_SMEM_MIN_ROWS)
73 .ok()
74 .and_then(|value| value.trim().parse::<u32>().ok())
75 .unwrap_or(DEFAULT_ILP_EXACT_CHAIN_SMEM_MIN_ROWS)
76}
77
78impl super::CudaKernelProvider {
79 #[cfg(test)]
98 fn ilp_exact_score(
99 &self,
100 candidate_buffers: &[&CudaBuffer],
101 positives: &CudaBuffer,
102 negatives: &CudaBuffer,
103 ) -> Result<(Vec<u32>, Vec<u32>)> {
104 let scores = self.ilp_exact_score_device(candidate_buffers, positives, negatives)?;
105 let device = self.device.inner();
106 self.device.synchronize()?;
107
108 let mut pos_covered = vec![0u32; scores.slot_count];
109 self.d2h_transfer_count.fetch_add(1, Ordering::Relaxed);
110 device
111 .dtoh_sync_copy_into(&scores.pos_covered, &mut pos_covered)
112 .map_err(|e| XlogError::Kernel(format!("ilp_exact_score: dtoh pos_covered: {}", e)))?;
113
114 let mut neg_covered = vec![0u32; scores.slot_count];
115 self.d2h_transfer_count.fetch_add(1, Ordering::Relaxed);
116 device
117 .dtoh_sync_copy_into(&scores.neg_covered, &mut neg_covered)
118 .map_err(|e| XlogError::Kernel(format!("ilp_exact_score: dtoh neg_covered: {}", e)))?;
119
120 Ok((pos_covered, neg_covered))
121 }
122
123 pub fn ilp_exact_score_topk(
126 &self,
127 candidate_buffers: &[&CudaBuffer],
128 positives: &CudaBuffer,
129 negatives: &CudaBuffer,
130 k_per_topology: u32,
131 ) -> Result<Vec<IlpExactTopkCandidate>> {
132 if k_per_topology == 0 {
133 return Ok(Vec::new());
134 }
135
136 let scores = self.ilp_exact_score_device(candidate_buffers, positives, negatives)?;
137 let out_rows = 4usize
138 .checked_mul(k_per_topology as usize)
139 .ok_or_else(|| XlogError::Kernel("ilp_exact_score_topk: output row overflow".into()))?;
140 let out_words = out_rows.checked_mul(ILP_EXACT_TOPK_FIELDS).ok_or_else(|| {
141 XlogError::Kernel("ilp_exact_score_topk: output word overflow".into())
142 })?;
143 let mut selected_buf = self.memory.alloc::<u32>(out_words)?;
144 let device = self.device.inner();
145 let func = device
146 .get_func(ILP_EXACT_MODULE, ilp_exact_kernels::ILP_EXACT_SELECT_TOPK)
147 .ok_or_else(|| {
148 XlogError::Kernel(format!(
149 "{} kernel not loaded",
150 ilp_exact_kernels::ILP_EXACT_SELECT_TOPK
151 ))
152 })?;
153
154 unsafe {
155 func.clone().launch(
156 LaunchConfig {
157 grid_dim: (4, 1, 1),
158 block_dim: (1, 1, 1),
159 shared_mem_bytes: 0,
160 },
161 (
162 &scores.pos_covered,
163 &scores.neg_covered,
164 scores.candidate_count as u32,
165 k_per_topology,
166 &mut selected_buf,
167 ),
168 )
169 }
170 .map_err(|e| XlogError::Kernel(format!("ilp_exact_select_topk launch: {}", e)))?;
171
172 self.device.synchronize()?;
173 let mut words = vec![0u32; out_words];
174 self.d2h_transfer_count.fetch_add(1, Ordering::Relaxed);
175 device
176 .dtoh_sync_copy_into(&selected_buf, &mut words)
177 .map_err(|e| {
178 XlogError::Kernel(format!("ilp_exact_score_topk: dtoh selected: {}", e))
179 })?;
180
181 let mut selected = Vec::new();
182 for chunk in words.as_chunks::<ILP_EXACT_TOPK_FIELDS>().0 {
183 if chunk[3] == 0 {
184 continue;
185 }
186 selected.push(IlpExactTopkCandidate {
187 topology_idx: chunk[0],
188 left_idx: chunk[1],
189 right_idx: chunk[2],
190 positives_covered: chunk[3],
191 negatives_covered: chunk[4],
192 local_rank: chunk[5],
193 next_positives_covered: chunk[6],
194 next_negatives_covered: chunk[7],
195 tie_class_size: chunk[8],
196 });
197 }
198 Ok(selected)
199 }
200
201 fn ilp_exact_score_device(
202 &self,
203 candidate_buffers: &[&CudaBuffer],
204 positives: &CudaBuffer,
205 negatives: &CudaBuffer,
206 ) -> Result<IlpExactDeviceScores> {
207 let c = candidate_buffers.len();
208 if c == 0 {
209 return Err(XlogError::Kernel(
210 "ilp_exact_score: candidate list is empty (filter at the engine)".to_string(),
211 ));
212 }
213 let c_u32 = u32::try_from(c).map_err(|_| {
214 XlogError::Kernel(format!(
215 "ilp_exact_score: candidate count {} exceeds u32::MAX",
216 c
217 ))
218 })?;
219
220 let layout = validate_exact_pair_buffer(positives, "positives")?;
222 require_exact_pair_layout(negatives, "negatives", layout)?;
223 let pos_rows = cached_rows(positives, "positives")?;
224 let neg_rows = cached_rows(negatives, "negatives")?;
225
226 let mut cand_rows: Vec<u32> = Vec::with_capacity(c);
227 for (i, buf) in candidate_buffers.iter().enumerate() {
228 let label = format!("candidate[{}]", i);
229 require_exact_pair_layout(buf, &label, layout)?;
230 cand_rows.push(cached_rows(buf, &label)?);
231 }
232
233 let mut cand_offsets_host: Vec<u32> = Vec::with_capacity(c + 1);
235 let mut running: u32 = 0;
236 cand_offsets_host.push(0);
237 for &r in &cand_rows {
238 running = running.checked_add(r).ok_or_else(|| {
239 XlogError::Kernel("ilp_exact_score: candidate row count overflow u32".to_string())
240 })?;
241 cand_offsets_host.push(running);
242 }
243 let total_rows = running as usize;
244 let elem_size = layout.elem_size();
245 let total_bytes = total_rows * elem_size;
246
247 let device = self.device.inner();
248
249 let mut cand_arg0_buf = self.memory.alloc::<u8>(total_bytes)?;
253 let mut cand_arg1_buf = self.memory.alloc::<u8>(total_bytes)?;
254 if total_bytes > 0 {
255 let mut byte_offset: usize = 0;
256 for (i, buf) in candidate_buffers.iter().enumerate() {
257 let rows = cand_rows[i] as usize;
258 if rows == 0 {
259 continue;
260 }
261 let bytes = rows * elem_size;
262
263 let src0 = buf.column(0).ok_or_else(|| {
264 XlogError::Kernel(format!("candidate[{}] missing column 0", i))
265 })?;
266 let src1 = buf.column(1).ok_or_else(|| {
267 XlogError::Kernel(format!("candidate[{}] missing column 1", i))
268 })?;
269 let src_view0 = self.column_bytes_view(src0, bytes)?;
270 let src_view1 = self.column_bytes_view(src1, bytes)?;
271 let mut dst0 = cand_arg0_buf.slice_mut(byte_offset..byte_offset + bytes);
272 let mut dst1 = cand_arg1_buf.slice_mut(byte_offset..byte_offset + bytes);
273 device.dtod_copy(&src_view0, &mut dst0).map_err(|e| {
274 XlogError::Kernel(format!(
275 "ilp_exact_score: d2d concat arg0 (candidate {}): {}",
276 i, e
277 ))
278 })?;
279 device.dtod_copy(&src_view1, &mut dst1).map_err(|e| {
280 XlogError::Kernel(format!(
281 "ilp_exact_score: d2d concat arg1 (candidate {}): {}",
282 i, e
283 ))
284 })?;
285 byte_offset += bytes;
286 }
287 }
288
289 let mut cand_offsets_buf = self.memory.alloc::<u32>(c + 1)?;
291 self.htod_sync_copy_into_tracked(&cand_offsets_host, &mut cand_offsets_buf)
292 .map_err(|e| XlogError::Kernel(format!("ilp_exact_score: h2d cand_offsets: {}", e)))?;
293
294 let n_slots = 4usize
296 .checked_mul(c)
297 .and_then(|v| v.checked_mul(c))
298 .ok_or_else(|| {
299 XlogError::Kernel("ilp_exact_score: n_slots = 4 * C * C overflow".to_string())
300 })?;
301 let mut pos_covered_buf = self.memory.alloc::<u32>(n_slots)?;
302 let mut neg_covered_buf = self.memory.alloc::<u32>(n_slots)?;
303 let pos_col0 = positives
306 .column(0)
307 .ok_or_else(|| XlogError::Kernel("positives: missing column 0".to_string()))?;
308 let pos_col1 = positives
309 .column(1)
310 .ok_or_else(|| XlogError::Kernel("positives: missing column 1".to_string()))?;
311 let neg_col0 = negatives
312 .column(0)
313 .ok_or_else(|| XlogError::Kernel("negatives: missing column 0".to_string()))?;
314 let neg_col1 = negatives
315 .column(1)
316 .ok_or_else(|| XlogError::Kernel("negatives: missing column 1".to_string()))?;
317
318 let max_candidate_rows = cand_rows.iter().copied().max().unwrap_or(0);
320 let chain_smem_enabled = ilp_exact_chain_smem_enabled()?
321 && max_candidate_rows >= ilp_exact_chain_smem_min_rows();
322 let shared_mem_bytes = if chain_smem_enabled {
323 chain_smem_shared_bytes(layout)
324 } else {
325 0
326 };
327 match layout {
328 ExactPairLayout::U64 => {
329 let cand_arg0_view = unsafe { cand_arg0_buf.device_view().cast::<u64>() }
330 .and_then(|view| view.try_slice(..total_rows))
331 .ok_or_else(|| XlogError::Kernel("device scalar view is invalid".into()))?;
332 let cand_arg1_view = unsafe { cand_arg1_buf.device_view().cast::<u64>() }
333 .and_then(|view| view.try_slice(..total_rows))
334 .ok_or_else(|| XlogError::Kernel("device scalar view is invalid".into()))?;
335 let pos_arg0_view = self.column_as_u64_view(pos_col0, pos_rows as usize)?;
336 let pos_arg1_view = self.column_as_u64_view(pos_col1, pos_rows as usize)?;
337 let neg_arg0_view = self.column_as_u64_view(neg_col0, neg_rows as usize)?;
338 let neg_arg1_view = self.column_as_u64_view(neg_col1, neg_rows as usize)?;
339 let kernel_name = if chain_smem_enabled {
340 ilp_exact_kernels::ILP_EXACT_SCORE_CHAIN_SMEM
341 } else {
342 ilp_exact_kernels::ILP_EXACT_SCORE
343 };
344 let func = device
345 .get_func(ILP_EXACT_MODULE, kernel_name)
346 .ok_or_else(|| {
347 XlogError::Kernel(format!("{} kernel not loaded", kernel_name))
348 })?;
349 unsafe {
350 func.clone().launch(
351 LaunchConfig {
352 grid_dim: (c_u32, c_u32, 4),
353 block_dim: (ILP_EXACT_BLOCK_SIZE, 1, 1),
354 shared_mem_bytes,
355 },
356 (
357 &cand_arg0_view,
358 &cand_arg1_view,
359 &cand_offsets_buf,
360 c_u32,
361 &pos_arg0_view,
362 &pos_arg1_view,
363 pos_rows,
364 &neg_arg0_view,
365 &neg_arg1_view,
366 neg_rows,
367 &mut pos_covered_buf,
368 &mut neg_covered_buf,
369 ),
370 )
371 }
372 .map_err(|e| XlogError::Kernel(format!("ilp_exact_score launch: {}", e)))?;
373 }
374 ExactPairLayout::U32 | ExactPairLayout::Symbol => {
375 let cand_arg0_view = unsafe { cand_arg0_buf.device_view().cast::<u32>() }
376 .and_then(|view| view.try_slice(..total_rows))
377 .ok_or_else(|| XlogError::Kernel("device scalar view is invalid".into()))?;
378 let cand_arg1_view = unsafe { cand_arg1_buf.device_view().cast::<u32>() }
379 .and_then(|view| view.try_slice(..total_rows))
380 .ok_or_else(|| XlogError::Kernel("device scalar view is invalid".into()))?;
381 let pos_arg0_view = self.column_as_u32_view(pos_col0, pos_rows as usize)?;
382 let pos_arg1_view = self.column_as_u32_view(pos_col1, pos_rows as usize)?;
383 let neg_arg0_view = self.column_as_u32_view(neg_col0, neg_rows as usize)?;
384 let neg_arg1_view = self.column_as_u32_view(neg_col1, neg_rows as usize)?;
385 let kernel_name = if chain_smem_enabled {
386 ilp_exact_kernels::ILP_EXACT_SCORE_CHAIN_SMEM_U32
387 } else {
388 ilp_exact_kernels::ILP_EXACT_SCORE_U32
389 };
390 let func = device
391 .get_func(ILP_EXACT_MODULE, kernel_name)
392 .ok_or_else(|| {
393 XlogError::Kernel(format!("{} kernel not loaded", kernel_name))
394 })?;
395 unsafe {
396 func.clone().launch(
397 LaunchConfig {
398 grid_dim: (c_u32, c_u32, 4),
399 block_dim: (ILP_EXACT_BLOCK_SIZE, 1, 1),
400 shared_mem_bytes,
401 },
402 (
403 &cand_arg0_view,
404 &cand_arg1_view,
405 &cand_offsets_buf,
406 c_u32,
407 &pos_arg0_view,
408 &pos_arg1_view,
409 pos_rows,
410 &neg_arg0_view,
411 &neg_arg1_view,
412 neg_rows,
413 &mut pos_covered_buf,
414 &mut neg_covered_buf,
415 ),
416 )
417 }
418 .map_err(|e| XlogError::Kernel(format!("ilp_exact_score_u32 launch: {}", e)))?;
419 }
420 }
421
422 Ok(IlpExactDeviceScores {
423 candidate_count: c,
424 #[cfg(test)]
425 slot_count: n_slots,
426 pos_covered: pos_covered_buf,
427 neg_covered: neg_covered_buf,
428 })
429 }
430}
431
432fn validate_exact_pair_buffer(buf: &CudaBuffer, label: &str) -> Result<ExactPairLayout> {
433 if buf.arity() != 2 {
434 return Err(XlogError::Kernel(format!(
435 "ilp_exact_score: {} buffer arity = {}, expected 2",
436 label,
437 buf.arity(),
438 )));
439 }
440 let mut layout: Option<ExactPairLayout> = None;
441 for col_idx in 0..2 {
442 let t = buf.schema().column_type(col_idx).ok_or_else(|| {
443 XlogError::Kernel(format!(
444 "ilp_exact_score: {} buffer missing column {} type",
445 label, col_idx,
446 ))
447 })?;
448 let col_layout = match t {
449 ScalarType::U64 => ExactPairLayout::U64,
450 ScalarType::U32 => ExactPairLayout::U32,
451 ScalarType::Symbol => ExactPairLayout::Symbol,
452 _ => {
453 return Err(XlogError::Kernel(format!(
454 "ilp_exact_score: {} buffer column {} type = {:?}, expected U64, U32, or Symbol",
455 label, col_idx, t,
456 )));
457 }
458 };
459 if let Some(expected) = layout {
460 if expected != col_layout {
461 return Err(XlogError::Kernel(format!(
462 "ilp_exact_score: {} buffer column {} type mismatch: {:?} vs {:?}",
463 label, col_idx, expected, col_layout,
464 )));
465 }
466 } else {
467 layout = Some(col_layout);
468 }
469 }
470 Ok(layout.expect("arity 2 loop sets layout"))
471}
472
473fn require_exact_pair_layout(
474 buf: &CudaBuffer,
475 label: &str,
476 expected: ExactPairLayout,
477) -> Result<()> {
478 let actual = validate_exact_pair_buffer(buf, label)?;
479 if actual != expected {
480 return Err(XlogError::Kernel(format!(
481 "ilp_exact_score: {} buffer type mismatch: expected {:?}, got {:?}",
482 label, expected, actual,
483 )));
484 }
485 Ok(())
486}
487
488fn cached_rows(buf: &CudaBuffer, label: &str) -> Result<u32> {
489 buf.cached_row_count().ok_or_else(|| {
490 XlogError::Kernel(format!(
491 "ilp_exact_score: {} buffer has no cached row count \
492 (DLPack ingest and create_empty_buffer both populate it)",
493 label
494 ))
495 })
496}
497
498#[cfg(test)]
499mod tests {
500 use xlog_core::{MemoryBudget, ScalarType, Schema};
508
509 use crate::{CudaKernelProvider, CudaProviderBuilder};
510
511 fn make_provider() -> Option<CudaKernelProvider> {
512 let budget = MemoryBudget::with_limit(1024 * 1024 * 1024);
513 CudaProviderBuilder::new(0, budget).build().ok()
514 }
515
516 fn pair_buffer(provider: &CudaKernelProvider, arg0: &[u64], arg1: &[u64]) -> crate::CudaBuffer {
520 assert_eq!(arg0.len(), arg1.len());
521 let schema = Schema::new(vec![
522 ("arg0".to_string(), ScalarType::U64),
523 ("arg1".to_string(), ScalarType::U64),
524 ]);
525 if arg0.is_empty() {
526 return provider
527 .create_empty_buffer(schema)
528 .expect("empty pair buffer");
529 }
530 let device = provider.device().inner();
534 let arg0_bytes: Vec<u8> = arg0.iter().flat_map(|v| v.to_le_bytes()).collect();
535 let arg1_bytes: Vec<u8> = arg1.iter().flat_map(|v| v.to_le_bytes()).collect();
536 let mut col0 = provider
537 .memory()
538 .alloc::<u8>(arg0_bytes.len())
539 .expect("alloc");
540 let mut col1 = provider
541 .memory()
542 .alloc::<u8>(arg1_bytes.len())
543 .expect("alloc");
544 device
545 .htod_sync_copy_into(&arg0_bytes, &mut col0)
546 .expect("h2d arg0");
547 device
548 .htod_sync_copy_into(&arg1_bytes, &mut col1)
549 .expect("h2d arg1");
550 provider
551 .buffer_from_columns(vec![col0.into(), col1.into()], arg0.len() as u64, schema)
552 .expect("buffer_from_columns")
553 }
554
555 fn pair_buffer_u32(
556 provider: &CudaKernelProvider,
557 arg0: &[u32],
558 arg1: &[u32],
559 typ: ScalarType,
560 ) -> crate::CudaBuffer {
561 assert_eq!(arg0.len(), arg1.len());
562 assert!(matches!(typ, ScalarType::U32 | ScalarType::Symbol));
563 let schema = Schema::new(vec![("arg0".to_string(), typ), ("arg1".to_string(), typ)]);
564 if arg0.is_empty() {
565 return provider
566 .create_empty_buffer(schema)
567 .expect("empty pair buffer");
568 }
569 let device = provider.device().inner();
570 let arg0_bytes: Vec<u8> = arg0.iter().flat_map(|v| v.to_le_bytes()).collect();
571 let arg1_bytes: Vec<u8> = arg1.iter().flat_map(|v| v.to_le_bytes()).collect();
572 let mut col0 = provider
573 .memory()
574 .alloc::<u8>(arg0_bytes.len())
575 .expect("alloc");
576 let mut col1 = provider
577 .memory()
578 .alloc::<u8>(arg1_bytes.len())
579 .expect("alloc");
580 device
581 .htod_sync_copy_into(&arg0_bytes, &mut col0)
582 .expect("h2d arg0");
583 device
584 .htod_sync_copy_into(&arg1_bytes, &mut col1)
585 .expect("h2d arg1");
586 provider
587 .buffer_from_columns(vec![col0.into(), col1.into()], arg0.len() as u64, schema)
588 .expect("buffer_from_columns")
589 }
590
591 fn pair_buffer_i32(
592 provider: &CudaKernelProvider,
593 arg0: &[i32],
594 arg1: &[i32],
595 ) -> crate::CudaBuffer {
596 assert_eq!(arg0.len(), arg1.len());
597 let schema = Schema::new(vec![
598 ("arg0".to_string(), ScalarType::I32),
599 ("arg1".to_string(), ScalarType::I32),
600 ]);
601 if arg0.is_empty() {
602 return provider
603 .create_empty_buffer(schema)
604 .expect("empty pair buffer");
605 }
606 let device = provider.device().inner();
607 let arg0_bytes: Vec<u8> = arg0.iter().flat_map(|v| v.to_le_bytes()).collect();
608 let arg1_bytes: Vec<u8> = arg1.iter().flat_map(|v| v.to_le_bytes()).collect();
609 let mut col0 = provider
610 .memory()
611 .alloc::<u8>(arg0_bytes.len())
612 .expect("alloc");
613 let mut col1 = provider
614 .memory()
615 .alloc::<u8>(arg1_bytes.len())
616 .expect("alloc");
617 device
618 .htod_sync_copy_into(&arg0_bytes, &mut col0)
619 .expect("h2d arg0");
620 device
621 .htod_sync_copy_into(&arg1_bytes, &mut col1)
622 .expect("h2d arg1");
623 provider
624 .buffer_from_columns(vec![col0.into(), col1.into()], arg0.len() as u64, schema)
625 .expect("buffer_from_columns")
626 }
627
628 #[test]
637 fn ilp_exact_score_matches_hand_computed_fixture() {
638 let provider = match make_provider() {
639 Some(p) => p,
640 None => {
641 eprintln!("Skipping test: no CUDA device available");
642 return;
643 }
644 };
645
646 let p_b = pair_buffer(&provider, &[1, 2], &[2, 3]);
648 let p_c = pair_buffer(&provider, &[2, 3, 4], &[4, 5, 6]);
649
650 let positives = pair_buffer(&provider, &[1, 2], &[4, 5]);
652 let negatives = pair_buffer(&provider, &[7], &[8]);
653
654 let (pos, neg) = provider
655 .ilp_exact_score(&[&p_b, &p_c], &positives, &negatives)
656 .expect("ilp_exact_score launch");
657
658 let mut expected_pos = vec![0u32; 16];
663 expected_pos[1] = 2;
664 assert_eq!(
665 pos, expected_pos,
666 "positives coverage mismatch: expected {:?}, got {:?}",
667 expected_pos, pos,
668 );
669
670 let expected_neg = vec![0u32; 16];
672 assert_eq!(
673 neg, expected_neg,
674 "negatives coverage mismatch: expected {:?}, got {:?}",
675 expected_neg, neg,
676 );
677 }
678
679 #[test]
680 fn ilp_exact_score_topk_reduces_on_device_to_compact_result() {
681 let provider = match make_provider() {
682 Some(p) => p,
683 None => {
684 eprintln!("Skipping test: no CUDA device available");
685 return;
686 }
687 };
688
689 let p_b = pair_buffer(&provider, &[1, 2], &[2, 3]);
690 let p_c = pair_buffer(&provider, &[2, 3, 4], &[4, 5, 6]);
691 let positives = pair_buffer(&provider, &[1, 2], &[4, 5]);
692 let negatives = pair_buffer(&provider, &[7], &[8]);
693
694 provider.reset_d2h_transfer_count();
695 let selected = provider
696 .ilp_exact_score_topk(&[&p_b, &p_c], &positives, &negatives, 2)
697 .expect("ilp_exact_score_topk launch");
698
699 assert_eq!(provider.d2h_transfer_count(), 1);
700 assert_eq!(selected.len(), 1);
701 let winner = selected[0];
702 assert_eq!(winner.topology_idx, 0);
703 assert_eq!(winner.left_idx, 0);
704 assert_eq!(winner.right_idx, 1);
705 assert_eq!(winner.positives_covered, 2);
706 assert_eq!(winner.negatives_covered, 0);
707 assert_eq!(winner.local_rank, 0);
708 assert_eq!(winner.next_positives_covered, 0);
709 assert_eq!(winner.next_negatives_covered, 0);
710 assert_eq!(winner.tie_class_size, 1);
711 }
712
713 #[test]
714 fn ilp_exact_score_topk_preserves_rank_next_and_tie_diagnostics() {
715 let provider = match make_provider() {
716 Some(p) => p,
717 None => {
718 eprintln!("Skipping test: no CUDA device available");
719 return;
720 }
721 };
722
723 let p_all = pair_buffer(&provider, &[1, 2], &[1, 2]);
724 let p_one = pair_buffer(&provider, &[1], &[1]);
725 let p_two = pair_buffer(&provider, &[2], &[2]);
726 let positives = pair_buffer(&provider, &[1, 2], &[1, 2]);
727 let negatives = pair_buffer(&provider, &[9], &[9]);
728
729 let selected = provider
730 .ilp_exact_score_topk(&[&p_all, &p_one, &p_two], &positives, &negatives, 2)
731 .expect("ilp_exact_score_topk launch");
732
733 let star_rank0 = selected
734 .iter()
735 .find(|row| row.topology_idx == 1 && row.local_rank == 0)
736 .expect("star rank 0");
737 assert_eq!(star_rank0.left_idx, 0);
738 assert_eq!(star_rank0.right_idx, 0);
739 assert_eq!(star_rank0.positives_covered, 2);
740 assert_eq!(star_rank0.negatives_covered, 0);
741 assert_eq!(star_rank0.next_positives_covered, 1);
742 assert_eq!(star_rank0.next_negatives_covered, 0);
743 assert_eq!(star_rank0.tie_class_size, 1);
744
745 let star_rank1 = selected
746 .iter()
747 .find(|row| row.topology_idx == 1 && row.local_rank == 1)
748 .expect("star rank 1");
749 assert_eq!(star_rank1.left_idx, 0);
750 assert_eq!(star_rank1.right_idx, 1);
751 assert_eq!(star_rank1.positives_covered, 1);
752 assert_eq!(star_rank1.negatives_covered, 0);
753 assert_eq!(star_rank1.next_positives_covered, 1);
754 assert_eq!(star_rank1.next_negatives_covered, 0);
755 assert_eq!(star_rank1.tie_class_size, 6);
756 }
757
758 #[test]
764 fn ilp_exact_score_is_deterministic_across_runs() {
765 let provider = match make_provider() {
766 Some(p) => p,
767 None => {
768 eprintln!("Skipping test: no CUDA device available");
769 return;
770 }
771 };
772
773 let p_b = pair_buffer(&provider, &[1, 2], &[2, 3]);
774 let p_c = pair_buffer(&provider, &[2, 3, 4], &[4, 5, 6]);
775 let positives = pair_buffer(&provider, &[1, 2], &[4, 5]);
776 let negatives = pair_buffer(&provider, &[7], &[8]);
777
778 let run_a = provider
779 .ilp_exact_score(&[&p_b, &p_c], &positives, &negatives)
780 .unwrap();
781 let run_b = provider
782 .ilp_exact_score(&[&p_b, &p_c], &positives, &negatives)
783 .unwrap();
784 assert_eq!(run_a.0, run_b.0, "pos coverage drifted across runs");
785 assert_eq!(run_a.1, run_b.1, "neg coverage drifted across runs");
786 }
787
788 #[test]
793 fn ilp_exact_score_handles_empty_negatives() {
794 let provider = match make_provider() {
795 Some(p) => p,
796 None => {
797 eprintln!("Skipping test: no CUDA device available");
798 return;
799 }
800 };
801
802 let p_b = pair_buffer(&provider, &[1, 2], &[2, 3]);
803 let p_c = pair_buffer(&provider, &[2, 3, 4], &[4, 5, 6]);
804 let positives = pair_buffer(&provider, &[1, 2], &[4, 5]);
805 let negatives = pair_buffer(&provider, &[], &[]);
806
807 let (pos, neg) = provider
808 .ilp_exact_score(&[&p_b, &p_c], &positives, &negatives)
809 .unwrap();
810
811 let mut expected_pos = vec![0u32; 16];
812 expected_pos[1] = 2;
813 assert_eq!(pos, expected_pos);
814 assert_eq!(neg, vec![0u32; 16]);
815 }
816
817 #[test]
818 fn ilp_exact_score_accepts_u32_pair_buffers() {
819 let provider = match make_provider() {
820 Some(p) => p,
821 None => {
822 eprintln!("Skipping test: no CUDA device available");
823 return;
824 }
825 };
826
827 let p_b = pair_buffer_u32(&provider, &[1, 2], &[2, 3], ScalarType::U32);
828 let p_c = pair_buffer_u32(&provider, &[2, 3, 4], &[4, 5, 6], ScalarType::U32);
829 let positives = pair_buffer_u32(&provider, &[1, 2], &[4, 5], ScalarType::U32);
830 let negatives = pair_buffer_u32(&provider, &[7], &[8], ScalarType::U32);
831
832 let (pos, neg) = provider
833 .ilp_exact_score(&[&p_b, &p_c], &positives, &negatives)
834 .expect("U32 ilp_exact_score launch");
835
836 let mut expected_pos = vec![0u32; 16];
837 expected_pos[1] = 2;
838 assert_eq!(pos, expected_pos);
839 assert_eq!(neg, vec![0u32; 16]);
840 }
841
842 #[test]
843 fn ilp_exact_score_accepts_symbol_pair_buffers() {
844 let provider = match make_provider() {
845 Some(p) => p,
846 None => {
847 eprintln!("Skipping test: no CUDA device available");
848 return;
849 }
850 };
851
852 let p_b = pair_buffer_u32(&provider, &[1, 2], &[2, 3], ScalarType::Symbol);
853 let p_c = pair_buffer_u32(&provider, &[2, 3, 4], &[4, 5, 6], ScalarType::Symbol);
854 let positives = pair_buffer_u32(&provider, &[1, 2], &[4, 5], ScalarType::Symbol);
855 let negatives = pair_buffer_u32(&provider, &[7], &[8], ScalarType::Symbol);
856
857 let (pos, neg) = provider
858 .ilp_exact_score(&[&p_b, &p_c], &positives, &negatives)
859 .expect("Symbol ilp_exact_score launch");
860
861 let mut expected_pos = vec![0u32; 16];
862 expected_pos[1] = 2;
863 assert_eq!(pos, expected_pos);
864 assert_eq!(neg, vec![0u32; 16]);
865 }
866
867 #[test]
868 fn ilp_exact_score_rejects_mixed_pair_types() {
869 let provider = match make_provider() {
870 Some(p) => p,
871 None => {
872 eprintln!("Skipping test: no CUDA device available");
873 return;
874 }
875 };
876
877 let p_b = pair_buffer_u32(&provider, &[1, 2], &[2, 3], ScalarType::U32);
878 let positives = pair_buffer(&provider, &[1, 2], &[4, 5]);
879 let negatives = pair_buffer(&provider, &[7], &[8]);
880
881 let err = provider
882 .ilp_exact_score(&[&p_b], &positives, &negatives)
883 .expect_err("mixed U64/U32 buffers must be rejected");
884 assert!(
885 err.to_string().contains("expected U64") || err.to_string().contains("type mismatch"),
886 "unexpected error: {err}"
887 );
888 }
889
890 #[test]
891 fn ilp_exact_score_rejects_unsupported_pair_types() {
892 let provider = match make_provider() {
893 Some(p) => p,
894 None => {
895 eprintln!("Skipping test: no CUDA device available");
896 return;
897 }
898 };
899
900 let p_b = pair_buffer_i32(&provider, &[1, 2], &[2, 3]);
901 let positives = pair_buffer_i32(&provider, &[1, 2], &[4, 5]);
902 let negatives = pair_buffer_i32(&provider, &[7], &[8]);
903
904 let err = provider
905 .ilp_exact_score(&[&p_b], &positives, &negatives)
906 .expect_err("I32 pair buffers must be rejected");
907 assert!(
908 err.to_string().contains("expected U64, U32, or Symbol"),
909 "unexpected error: {err}"
910 );
911 }
912}