1use std::sync::Arc;
18
19use xlog_core::MemoryBudget;
20
21use crate::device_runtime::{
22 AsyncCudaResource, DeviceMemoryResource, GlobalDeviceBudget, LogRecord, LoggingResource,
23 LoggingSink, SinkError, StreamPool, XlogDeviceRuntime,
24};
25use crate::joint_solver::{FuelMeter, SolverError};
26use crate::launch::LaunchRecorder;
27use crate::memory::{CudaColumn, GpuMemoryManager};
28use crate::provider::JOINT_SOLVE_MODULE;
29use crate::{CudaDevice, LaunchAsync, LaunchConfig};
30
31const FEASIBILITY_KERNEL: &str = "joint_label_feasibility";
33const TOP2_KERNEL: &str = "joint_label_top2";
35const COMPONENT_PLAN_INIT_KERNEL: &str = "joint_component_plan_init";
37const COMPONENT_ENTITY_OWNERS_KERNEL: &str = "joint_component_entity_owners";
38const COMPONENT_UNION_KERNEL: &str = "joint_component_union";
39const COMPONENT_COMPRESS_KERNEL: &str = "joint_component_compress";
40const COMPONENT_KERNEL: &str = "joint_component_enumerate";
42const CHAIN_DP_KERNEL: &str = "joint_component_chain_dp";
44const BRANCH_AND_BOUND_KERNEL: &str = "joint_component_branch_and_bound";
45
46const JOINT_SOLVE_KERNELS: &[&str] = &[
47 FEASIBILITY_KERNEL,
48 TOP2_KERNEL,
49 COMPONENT_PLAN_INIT_KERNEL,
50 COMPONENT_ENTITY_OWNERS_KERNEL,
51 COMPONENT_UNION_KERNEL,
52 COMPONENT_COMPRESS_KERNEL,
53 COMPONENT_KERNEL,
54 CHAIN_DP_KERNEL,
55 BRANCH_AND_BOUND_KERNEL,
56];
57
58const CARRIER_BUDGET_BYTES: u64 = 64 * 1024 * 1024;
62
63#[derive(Debug)]
66pub enum CarrierError {
67 SchemaAlreadyRegistered {
71 catalog_sha: String,
73 solver_identity: String,
75 },
76 Allocation(xlog_core::XlogError),
78 ZeroCapacity {
82 dimension: &'static str,
84 },
85 SchemaNotRegistered,
89 SignaturesAlreadyBound,
92 SignatureShapeMismatch {
94 side: &'static str,
96 expected_words: usize,
98 got_words: usize,
100 },
101 SignaturesUnbound,
103 FeasibilityNotSolved,
106 AbstainOutOfRange {
108 abstain_label: u32,
110 labels: usize,
112 },
113 KernelUnavailable {
116 detail: String,
118 },
119 Launch(xlog_core::XlogError),
121 Solver(SolverError),
124}
125
126impl std::fmt::Display for CarrierError {
127 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128 match self {
129 CarrierError::SchemaAlreadyRegistered {
130 catalog_sha,
131 solver_identity,
132 } => write!(
133 f,
134 "carrier schema already registered (catalog {catalog_sha}, \
135 solver {solver_identity}); registration is once-per-session"
136 ),
137 CarrierError::Allocation(err) => write!(f, "carrier allocation failed: {err}"),
138 CarrierError::ZeroCapacity { dimension } => write!(
139 f,
140 "carrier capacity dimension {dimension} is zero; refusing \
141 instead of silently clamping"
142 ),
143 CarrierError::SchemaNotRegistered => write!(
144 f,
145 "carrier schema is not registered; signature masks are \
146 catalog-bound and require the catalog anchor first"
147 ),
148 CarrierError::SignaturesAlreadyBound => {
149 write!(f, "signature masks already bound for this session")
150 }
151 CarrierError::SignatureShapeMismatch {
152 side,
153 expected_words,
154 got_words,
155 } => write!(
156 f,
157 "{side} signature mask has {got_words} u64 words, expected \
158 {expected_words} (labels x lanes)"
159 ),
160 CarrierError::SignaturesUnbound => write!(
161 f,
162 "solve refused: signature masks are not bound for this session"
163 ),
164 CarrierError::FeasibilityNotSolved => write!(
165 f,
166 "top-two stage refused: the feasibility stage has not \
167 populated the feasible sets this session"
168 ),
169 CarrierError::AbstainOutOfRange {
170 abstain_label,
171 labels,
172 } => write!(
173 f,
174 "abstain label {abstain_label} is outside the label universe \
175 of width {labels}"
176 ),
177 CarrierError::KernelUnavailable { detail } => {
178 write!(f, "joint-solve kernel unavailable: {detail}")
179 }
180 CarrierError::Launch(err) => write!(f, "carrier solve launch failed: {err}"),
181 CarrierError::Solver(err) => write!(f, "carrier solve refused: {err}"),
182 }
183 }
184}
185
186impl std::error::Error for CarrierError {}
187
188struct SilentSink;
190
191impl LoggingSink for SilentSink {
192 fn emit(&self, _record: LogRecord) -> Result<(), SinkError> {
193 Ok(())
194 }
195}
196
197#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub enum CarrierBufferId {
202 Domains,
204 Scores,
206 Constraints,
208 ArgumentArities,
210 Outputs,
212 FeasibleSets,
215 LogicalCounts,
220 MapResults,
226 SolveStatus,
233}
234
235pub struct CarrierExport {
241 pub slice: Arc<crate::memory::TrackedCudaSlice<u8>>,
243 pub stream: Arc<crate::CudaStream>,
245 pub elem_bytes: usize,
247 pub rows: usize,
249 pub cols: usize,
251}
252
253pub struct JointConstraintCarrier {
259 buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 9],
260 columns: [CudaColumn; 8],
261 component_buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 4],
262 component_columns: [CudaColumn; 4],
263 search_columns: [CudaColumn; 4],
264 signatures: Option<CudaColumn>,
265 registered_schema: Option<(String, String)>,
266 feasibility_solved: bool,
267 pending_producer_events: Vec<cudarc::driver::sys::CUevent>,
272 pending_consumer_streams: Vec<cudarc::driver::sys::CUstream>,
277 entities: usize,
278 domain_lanes: usize,
279 candidates: usize,
280 labels: usize,
281 max_arity: usize,
282 device: Arc<CudaDevice>,
283 pool: Arc<StreamPool>,
284 memory: Arc<GpuMemoryManager>,
285 runtime: Arc<XlogDeviceRuntime>,
286}
287
288fn label_words(labels: usize) -> usize {
290 labels.div_ceil(64)
291}
292
293fn shared_column(
300 slice: &Arc<crate::memory::TrackedCudaSlice<u8>>,
301 stream: &Arc<crate::CudaStream>,
302) -> CudaColumn {
303 let tensor = unsafe { crate::DlpackManagedTensor::from_raw(std::ptr::null_mut()) };
304 CudaColumn::dlpack_xlog_owned(Arc::clone(slice), Arc::clone(stream), tensor)
305}
306
307fn ensure_joint_solve_module(device: &Arc<CudaDevice>) -> Result<(), CarrierError> {
311 if JOINT_SOLVE_KERNELS
312 .iter()
313 .all(|k| device.inner().get_func(JOINT_SOLVE_MODULE, k).is_some())
314 {
315 return Ok(());
316 }
317 let cc = crate::provider::detect_compute_capability(device).map_err(|e| {
318 CarrierError::KernelUnavailable {
319 detail: e.to_string(),
320 }
321 })?;
322 let sources = crate::provider::load_module_sources("joint_solve", cc).map_err(|e| {
323 CarrierError::KernelUnavailable {
324 detail: e.to_string(),
325 }
326 })?;
327 let mut load_errors = Vec::new();
328 for source in sources {
329 let attempt = match source {
330 crate::provider::KernelModuleSource::File { path, .. } => device
331 .inner()
332 .load_file(&path, JOINT_SOLVE_MODULE, JOINT_SOLVE_KERNELS)
333 .map_err(|e| format!("{}: {e}", path.display())),
334 crate::provider::KernelModuleSource::EmbeddedPortablePtx { ptx } => device
335 .inner()
336 .load_ptx(
337 cudarc::nvrtc::Ptx::from_src(ptx),
338 JOINT_SOLVE_MODULE,
339 JOINT_SOLVE_KERNELS,
340 )
341 .map_err(|e| format!("embedded portable PTX: {e}")),
342 };
343 match attempt {
344 Ok(()) => return Ok(()),
345 Err(detail) => load_errors.push(detail),
346 }
347 }
348 Err(CarrierError::KernelUnavailable {
349 detail: if load_errors.is_empty() {
350 "no kernel artifact source available".to_string()
351 } else {
352 load_errors.join("; ")
353 },
354 })
355}
356
357impl Drop for JointConstraintCarrier {
358 fn drop(&mut self) {
359 for event in self.pending_producer_events.drain(..) {
362 unsafe {
365 let _ = cudarc::driver::result::event::destroy(event);
366 }
367 }
368 self.pending_consumer_streams.clear();
369 }
370}
371
372impl JointConstraintCarrier {
373 pub fn allocate(
377 device: Arc<CudaDevice>,
378 entities: usize,
379 domain_lanes: usize,
380 candidates: usize,
381 labels: usize,
382 ) -> Result<Self, CarrierError> {
383 let carrier =
384 Self::allocate_with_max_arity(device, entities, domain_lanes, candidates, labels, 2)?;
385 let binary_arities = vec![2u32; candidates];
386 unsafe {
387 cudarc::driver::result::memcpy_htod_sync(
388 *carrier.buffers[3].device_ptr(),
389 &binary_arities,
390 )
391 .map_err(|error| {
392 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
393 "legacy binary arity initialization failed: {error}"
394 )))
395 })?;
396 }
397 Ok(carrier)
398 }
399
400 pub fn allocate_with_max_arity(
404 device: Arc<CudaDevice>,
405 entities: usize,
406 domain_lanes: usize,
407 candidates: usize,
408 labels: usize,
409 max_arity: usize,
410 ) -> Result<Self, CarrierError> {
411 for (dimension, value) in [
412 ("entities", entities),
413 ("domain_lanes", domain_lanes),
414 ("candidates", candidates),
415 ("labels", labels),
416 ("max_arity", max_arity),
417 ] {
418 if value == 0 {
419 return Err(CarrierError::ZeroCapacity { dimension });
420 }
421 }
422
423 let pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
424 let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
425 AsyncCudaResource::new(Arc::clone(&device), 0, Arc::clone(&pool)),
426 );
427 let logging: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(LoggingResource::new(
428 async_resource,
429 Arc::new(SilentSink) as Arc<dyn LoggingSink>,
430 ));
431 let budget: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
432 GlobalDeviceBudget::new(logging, CARRIER_BUDGET_BYTES as usize),
433 );
434 let runtime = Arc::new(XlogDeviceRuntime::with_resource(
435 Arc::clone(&device),
436 0,
437 Arc::clone(&pool),
438 budget,
439 ));
440 let memory = Arc::new(GpuMemoryManager::with_runtime(
441 Arc::clone(&device),
442 MemoryBudget::with_limit(CARRIER_BUDGET_BYTES),
443 Arc::clone(&runtime),
444 ));
445
446 ensure_joint_solve_module(&device)?;
447
448 let domains = memory
449 .alloc::<u64>(entities * domain_lanes)
450 .map_err(CarrierError::Allocation)?;
451 let scores = memory
452 .alloc::<f32>(candidates * labels)
453 .map_err(CarrierError::Allocation)?;
454 let constraints = memory
455 .alloc::<u32>(candidates * max_arity)
456 .map_err(CarrierError::Allocation)?;
457 let argument_arities = memory
458 .alloc::<u32>(candidates)
459 .map_err(CarrierError::Allocation)?;
460 let outputs = memory
461 .alloc::<u32>(candidates)
462 .map_err(CarrierError::Allocation)?;
463 let feasible_sets = memory
464 .alloc::<u64>(candidates * label_words(labels))
465 .map_err(CarrierError::Allocation)?;
466 let logical_counts = memory.alloc::<u32>(4).map_err(CarrierError::Allocation)?;
467 let map_results = memory
468 .alloc::<u32>(candidates * 4)
469 .map_err(CarrierError::Allocation)?;
470 let solve_status = memory
471 .alloc::<u32>(candidates)
472 .map_err(CarrierError::Allocation)?;
473 let component_parents = memory
474 .alloc::<u32>(candidates)
475 .map_err(CarrierError::Allocation)?;
476 let component_entity_owners = memory
477 .alloc::<u32>(entities)
478 .map_err(CarrierError::Allocation)?;
479 let component_count = memory.alloc::<u32>(1).map_err(CarrierError::Allocation)?;
480 let component_fuel = memory.alloc::<u64>(1).map_err(CarrierError::Allocation)?;
481 let search_assignment = memory
482 .alloc::<u32>(candidates)
483 .map_err(CarrierError::Allocation)?;
484 let search_best_label = memory
485 .alloc::<u32>(candidates)
486 .map_err(CarrierError::Allocation)?;
487 let search_best_total = memory
488 .alloc::<f32>(candidates)
489 .map_err(CarrierError::Allocation)?;
490 let search_alt_total = memory
491 .alloc::<f32>(candidates)
492 .map_err(CarrierError::Allocation)?;
493
494 let buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 9] = [
498 Arc::new(domains.into_bytes()),
499 Arc::new(scores.into_bytes()),
500 Arc::new(constraints.into_bytes()),
501 Arc::new(argument_arities.into_bytes()),
502 Arc::new(outputs.into_bytes()),
503 Arc::new(feasible_sets.into_bytes()),
504 Arc::new(logical_counts.into_bytes()),
505 Arc::new(map_results.into_bytes()),
506 Arc::new(solve_status.into_bytes()),
507 ];
508 let component_buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 4] = [
509 Arc::new(component_parents.into_bytes()),
510 Arc::new(component_entity_owners.into_bytes()),
511 Arc::new(component_count.into_bytes()),
512 Arc::new(component_fuel.into_bytes()),
513 ];
514 let search_buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 4] = [
515 Arc::new(search_assignment.into_bytes()),
516 Arc::new(search_best_label.into_bytes()),
517 Arc::new(search_best_total.into_bytes()),
518 Arc::new(search_alt_total.into_bytes()),
519 ];
520 let stream = device.inner().stream().clone();
525 for buffer in buffers
526 .iter()
527 .chain(component_buffers.iter())
528 .chain(search_buffers.iter())
529 {
530 unsafe {
533 cudarc::driver::result::memset_d8_async(
534 *buffer.device_ptr(),
535 0,
536 buffer.len(),
537 stream.cu_stream(),
538 )
539 .map_err(|e| {
540 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
541 "carrier zero-init failed: {e}"
542 )))
543 })?;
544 }
545 }
546 device.inner().synchronize().map_err(|e| {
547 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
548 "carrier zero-init sync failed: {e}"
549 )))
550 })?;
551
552 let columns = [
553 shared_column(&buffers[0], &stream),
554 shared_column(&buffers[1], &stream),
555 shared_column(&buffers[2], &stream),
556 shared_column(&buffers[3], &stream),
557 shared_column(&buffers[4], &stream),
558 shared_column(&buffers[5], &stream),
559 shared_column(&buffers[7], &stream),
560 shared_column(&buffers[8], &stream),
561 ];
562 let component_columns = [
563 shared_column(&component_buffers[0], &stream),
564 shared_column(&component_buffers[1], &stream),
565 shared_column(&component_buffers[2], &stream),
566 shared_column(&component_buffers[3], &stream),
567 ];
568 let search_columns = [
569 shared_column(&search_buffers[0], &stream),
570 shared_column(&search_buffers[1], &stream),
571 shared_column(&search_buffers[2], &stream),
572 shared_column(&search_buffers[3], &stream),
573 ];
574
575 Ok(Self {
576 buffers,
577 columns,
578 component_buffers,
579 component_columns,
580 search_columns,
581 signatures: None,
582 registered_schema: None,
583 feasibility_solved: false,
584 pending_producer_events: Vec::new(),
585 pending_consumer_streams: Vec::new(),
586 entities,
587 domain_lanes,
588 candidates,
589 labels,
590 max_arity,
591 device,
592 pool,
593 memory,
594 runtime,
595 })
596 }
597
598 pub fn note_producer_stream(&mut self, external_stream: u64) -> Result<(), CarrierError> {
606 if external_stream == 0 {
607 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(
608 "null producer stream handle".to_string(),
609 )));
610 }
611 unsafe {
616 let event = cudarc::driver::result::event::create(
617 cudarc::driver::sys::CUevent_flags::CU_EVENT_DISABLE_TIMING,
618 )
619 .map_err(|e| {
620 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
621 "producer event create failed: {e}"
622 )))
623 })?;
624 if let Err(e) = cudarc::driver::result::event::record(
625 event,
626 external_stream as cudarc::driver::sys::CUstream,
627 ) {
628 let _ = cudarc::driver::result::event::destroy(event);
629 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
630 "producer event record failed: {e}"
631 ))));
632 }
633 self.pending_producer_events.push(event);
634 }
635 Ok(())
636 }
637
638 fn drain_producer_waits(&mut self, cu_stream: &crate::CudaStream) -> Result<(), CarrierError> {
642 for event in self.pending_producer_events.drain(..) {
643 unsafe {
646 let wait = cudarc::driver::result::stream::wait_event(
647 cu_stream.cu_stream(),
648 event,
649 cudarc::driver::sys::CUevent_wait_flags::CU_EVENT_WAIT_DEFAULT,
650 );
651 let _ = cudarc::driver::result::event::destroy(event);
652 wait.map_err(|e| {
653 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
654 "producer event wait failed: {e}"
655 )))
656 })?;
657 }
658 }
659 Ok(())
660 }
661
662 pub fn note_consumer_stream(&mut self, external_stream: u64) -> Result<(), CarrierError> {
667 if external_stream == 0 {
668 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(
669 "null consumer stream handle".to_string(),
670 )));
671 }
672 self.pending_consumer_streams
673 .push(external_stream as cudarc::driver::sys::CUstream);
674 Ok(())
675 }
676
677 fn handoff_consumers(&mut self, cu_stream: &crate::CudaStream) -> Result<(), CarrierError> {
680 let consumer_streams = std::mem::take(&mut self.pending_consumer_streams);
681 if consumer_streams.is_empty() {
682 return Ok(());
683 }
684
685 unsafe {
690 let event = cudarc::driver::result::event::create(
691 cudarc::driver::sys::CUevent_flags::CU_EVENT_DISABLE_TIMING,
692 )
693 .map_err(|e| {
694 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
695 "consumer event create failed: {e}"
696 )))
697 })?;
698 if let Err(e) = cudarc::driver::result::event::record(event, cu_stream.cu_stream()) {
699 let _ = cudarc::driver::result::event::destroy(event);
700 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
701 "consumer event record failed: {e}"
702 ))));
703 }
704
705 let wait_result = consumer_streams
706 .into_iter()
707 .try_for_each(|consumer_stream| {
708 cudarc::driver::result::stream::wait_event(
709 consumer_stream,
710 event,
711 cudarc::driver::sys::CUevent_wait_flags::CU_EVENT_WAIT_DEFAULT,
712 )
713 .map_err(|e| {
714 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
715 "consumer event wait failed: {e}"
716 )))
717 })
718 });
719 let destroy_result = cudarc::driver::result::event::destroy(event).map_err(|e| {
720 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
721 "consumer event destroy failed: {e}"
722 )))
723 });
724 wait_result?;
725 destroy_result?;
726 }
727 Ok(())
728 }
729
730 pub fn export_buffer(&self, id: CarrierBufferId) -> CarrierExport {
736 let (index, elem_bytes, rows, cols) = match id {
737 CarrierBufferId::Domains => (0, 8, self.entities, self.domain_lanes),
738 CarrierBufferId::Scores => (1, 4, self.candidates, self.labels),
739 CarrierBufferId::Constraints => (2, 4, self.candidates, self.max_arity),
740 CarrierBufferId::ArgumentArities => (3, 4, self.candidates, 1),
741 CarrierBufferId::Outputs => (4, 4, self.candidates, 1),
742 CarrierBufferId::FeasibleSets => (5, 8, self.candidates, label_words(self.labels)),
743 CarrierBufferId::LogicalCounts => (6, 4, 1, 4),
744 CarrierBufferId::MapResults => (7, 4, self.candidates, 4),
745 CarrierBufferId::SolveStatus => (8, 4, self.candidates, 1),
746 };
747 CarrierExport {
748 slice: Arc::clone(&self.buffers[index]),
749 stream: self.device.inner().stream().clone(),
750 elem_bytes,
751 rows,
752 cols,
753 }
754 }
755
756 pub fn bind_signatures(
760 &mut self,
761 head_masks: &[u64],
762 tail_masks: &[u64],
763 ) -> Result<(), CarrierError> {
764 if self.registered_schema.is_none() {
765 return Err(CarrierError::SchemaNotRegistered);
766 }
767 if self.signatures.is_some() {
768 return Err(CarrierError::SignaturesAlreadyBound);
769 }
770 let expected_words = self.labels * self.domain_lanes;
771 for (side, masks) in [("head", head_masks), ("tail", tail_masks)] {
772 if masks.len() != expected_words {
773 return Err(CarrierError::SignatureShapeMismatch {
774 side,
775 expected_words,
776 got_words: masks.len(),
777 });
778 }
779 }
780 if self.max_arity != 2 {
781 return Err(CarrierError::SignatureShapeMismatch {
782 side: "roles",
783 expected_words: self.max_arity * expected_words,
784 got_words: 2 * expected_words,
785 });
786 }
787 let mut role_masks = Vec::with_capacity(2 * expected_words);
788 role_masks.extend_from_slice(head_masks);
789 role_masks.extend_from_slice(tail_masks);
790 self.signatures = Some(self.upload_mask(&role_masks)?);
791 Ok(())
792 }
793
794 pub fn bind_role_signatures(&mut self, role_masks: &[u64]) -> Result<(), CarrierError> {
797 if self.registered_schema.is_none() {
798 return Err(CarrierError::SchemaNotRegistered);
799 }
800 if self.signatures.is_some() {
801 return Err(CarrierError::SignaturesAlreadyBound);
802 }
803 let expected_words = self.max_arity * self.labels * self.domain_lanes;
804 if role_masks.len() != expected_words {
805 return Err(CarrierError::SignatureShapeMismatch {
806 side: "roles",
807 expected_words,
808 got_words: role_masks.len(),
809 });
810 }
811 self.signatures = Some(self.upload_mask(role_masks)?);
812 Ok(())
813 }
814
815 fn upload_mask(&self, masks: &[u64]) -> Result<CudaColumn, CarrierError> {
818 let mut slice = self
819 .memory
820 .alloc::<u64>(masks.len())
821 .map_err(CarrierError::Allocation)?;
822 self.device
823 .inner()
824 .htod_sync_copy_into(masks, &mut slice)
825 .map_err(|e| {
826 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
827 "signature mask upload failed: {e}"
828 )))
829 })?;
830 Ok(CudaColumn::owned(slice.into_bytes()))
831 }
832
833 pub fn solve_label_feasibility(
840 &mut self,
841 abstain_label: u32,
842 fuel: &mut FuelMeter,
843 ) -> Result<(), CarrierError> {
844 let result = self.solve_label_feasibility_inner(abstain_label, fuel);
845 if result.is_err() {
846 self.pending_consumer_streams.clear();
847 }
848 result
849 }
850
851 fn solve_label_feasibility_inner(
852 &mut self,
853 abstain_label: u32,
854 fuel: &mut FuelMeter,
855 ) -> Result<(), CarrierError> {
856 if self.registered_schema.is_none() {
857 return Err(CarrierError::SchemaNotRegistered);
858 }
859 if self.signatures.is_none() {
860 return Err(CarrierError::SignaturesUnbound);
861 }
862 if abstain_label as usize >= self.labels {
863 return Err(CarrierError::AbstainOutOfRange {
864 abstain_label,
865 labels: self.labels,
866 });
867 }
868 fuel.charge((self.candidates as u64) * (self.labels as u64))
869 .map_err(CarrierError::Solver)?;
870
871 let stream_id = self.pool.acquire().map_err(|e| {
872 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
873 "no launch stream available: {e:?}"
874 )))
875 })?;
876 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
877 CarrierError::Launch(xlog_core::XlogError::Kernel(
878 "launch stream did not resolve".to_string(),
879 ))
880 })?;
881
882 self.drain_producer_waits(&cu_stream)?;
883
884 let Some(signatures) = &self.signatures else {
885 return Err(CarrierError::SignaturesUnbound);
886 };
887 let [domains, _scores, arguments, argument_arities, outputs, feasible_sets, _map_results, _solve_status] =
888 &self.columns;
889 let role_masks = signatures;
890
891 let mut rec = LaunchRecorder::new_strict(stream_id);
892 rec.read_column(domains);
893 rec.read_column(arguments);
894 rec.read_column(argument_arities);
895 rec.read_column(role_masks);
896 rec.write_column(outputs);
897 rec.write_column(feasible_sets);
898 rec.preflight(&self.runtime).map_err(|e| {
899 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
900 "solve launch preflight failed: {e}"
901 )))
902 })?;
903
904 let kernel = self
905 .device
906 .inner()
907 .get_func(JOINT_SOLVE_MODULE, FEASIBILITY_KERNEL)
908 .ok_or_else(|| CarrierError::KernelUnavailable {
909 detail: format!("{FEASIBILITY_KERNEL} not resolvable after module load"),
910 })?;
911 let block = 256u32;
912 let grid = (self.candidates as u32).div_ceil(block);
913 unsafe {
920 kernel
921 .launch_on_stream(
922 &cu_stream,
923 LaunchConfig {
924 grid_dim: (grid, 1, 1),
925 block_dim: (block, 1, 1),
926 shared_mem_bytes: 0,
927 },
928 (
929 *domains.device_ptr(),
930 *arguments.device_ptr(),
931 *argument_arities.device_ptr(),
932 *role_masks.device_ptr(),
933 self.max_arity as u32,
934 self.entities as u32,
935 self.candidates as u32,
936 self.labels as u32,
937 self.domain_lanes as u32,
938 abstain_label,
939 *outputs.device_ptr(),
940 *feasible_sets.device_ptr(),
941 ),
942 )
943 .map_err(|e| {
944 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
945 "solve launch failed: {e}"
946 )))
947 })?;
948 }
949 rec.commit(&self.runtime).map_err(|e| {
950 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
951 "solve launch commit failed: {e}"
952 )))
953 })?;
954 self.handoff_consumers(&cu_stream)?;
955 self.feasibility_solved = true;
956 Ok(())
957 }
958
959 pub fn solve_label_map_top2(&mut self, fuel: &mut FuelMeter) -> Result<(), CarrierError> {
972 let result = self.solve_label_map_top2_inner(fuel);
973 if result.is_err() {
974 self.pending_consumer_streams.clear();
975 }
976 result
977 }
978
979 fn solve_label_map_top2_inner(&mut self, fuel: &mut FuelMeter) -> Result<(), CarrierError> {
980 if self.registered_schema.is_none() {
981 return Err(CarrierError::SchemaNotRegistered);
982 }
983 if !self.feasibility_solved {
984 return Err(CarrierError::FeasibilityNotSolved);
985 }
986 fuel.charge((self.candidates as u64) * (self.labels as u64))
987 .map_err(CarrierError::Solver)?;
988
989 let stream_id = self.pool.acquire().map_err(|e| {
990 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
991 "no launch stream available: {e:?}"
992 )))
993 })?;
994 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
995 CarrierError::Launch(xlog_core::XlogError::Kernel(
996 "launch stream did not resolve".to_string(),
997 ))
998 })?;
999
1000 self.drain_producer_waits(&cu_stream)?;
1001
1002 let [_domains, scores, _arguments, _argument_arities, _outputs, feasible_sets, map_results, _solve_status] =
1003 &self.columns;
1004
1005 let mut rec = LaunchRecorder::new_strict(stream_id);
1006 rec.read_column(scores);
1007 rec.read_column(feasible_sets);
1008 rec.write_column(map_results);
1009 rec.preflight(&self.runtime).map_err(|e| {
1010 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1011 "top-two launch preflight failed: {e}"
1012 )))
1013 })?;
1014
1015 let kernel = self
1016 .device
1017 .inner()
1018 .get_func(JOINT_SOLVE_MODULE, TOP2_KERNEL)
1019 .ok_or_else(|| CarrierError::KernelUnavailable {
1020 detail: format!("{TOP2_KERNEL} not resolvable after module load"),
1021 })?;
1022 let block = 256u32;
1023 let grid = (self.candidates as u32).div_ceil(block);
1024 unsafe {
1028 kernel
1029 .launch_on_stream(
1030 &cu_stream,
1031 LaunchConfig {
1032 grid_dim: (grid, 1, 1),
1033 block_dim: (block, 1, 1),
1034 shared_mem_bytes: 0,
1035 },
1036 (
1037 *scores.device_ptr(),
1038 *feasible_sets.device_ptr(),
1039 self.candidates as u32,
1040 self.labels as u32,
1041 *map_results.device_ptr(),
1042 ),
1043 )
1044 .map_err(|e| {
1045 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1046 "top-two launch failed: {e}"
1047 )))
1048 })?;
1049 }
1050 rec.commit(&self.runtime).map_err(|e| {
1051 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1052 "top-two launch commit failed: {e}"
1053 )))
1054 })?;
1055 self.handoff_consumers(&cu_stream)?;
1056 Ok(())
1057 }
1058
1059 pub fn solve_components_exact(&mut self, fuel: &mut FuelMeter) -> Result<(), CarrierError> {
1067 let result = self.solve_components_exact_inner(fuel);
1068 if result.is_err() {
1069 self.pending_consumer_streams.clear();
1070 }
1071 result
1072 }
1073
1074 fn solve_components_exact_inner(&mut self, fuel: &mut FuelMeter) -> Result<(), CarrierError> {
1075 if self.registered_schema.is_none() {
1076 return Err(CarrierError::SchemaNotRegistered);
1077 }
1078 if !self.feasibility_solved {
1079 return Err(CarrierError::FeasibilityNotSolved);
1080 }
1081
1082 let authorized = fuel.remaining();
1087 fuel.charge(authorized).map_err(CarrierError::Solver)?;
1088
1089 let stream_id = self.pool.acquire().map_err(|e| {
1090 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1091 "no launch stream available: {e:?}"
1092 )))
1093 })?;
1094 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
1095 CarrierError::Launch(xlog_core::XlogError::Kernel(
1096 "launch stream did not resolve".to_string(),
1097 ))
1098 })?;
1099 self.drain_producer_waits(&cu_stream)?;
1100
1101 let Some(signatures) = &self.signatures else {
1102 return Err(CarrierError::SignaturesUnbound);
1103 };
1104 let [domains, scores, arguments, argument_arities, _outputs, feasible_sets, map_results, solve_status] =
1105 &self.columns;
1106 let [component_parents, component_entity_owners, component_count, component_fuel] =
1107 &self.component_columns;
1108 let [search_assignment, search_best_label, search_best_total, search_alt_total] =
1109 &self.search_columns;
1110 let role_masks = signatures;
1111
1112 let mut rec = LaunchRecorder::new_strict(stream_id);
1113 rec.read_column(scores);
1114 rec.read_column(feasible_sets);
1115 rec.read_column(arguments);
1116 rec.read_column(argument_arities);
1117 rec.read_column(domains);
1118 rec.read_column(role_masks);
1119 rec.read_column(component_parents);
1120 rec.write_column(component_parents);
1121 rec.read_column(component_entity_owners);
1122 rec.write_column(component_entity_owners);
1123 rec.read_column(component_count);
1124 rec.write_column(component_count);
1125 rec.write_column(component_fuel);
1126 for search_column in [
1127 search_assignment,
1128 search_best_label,
1129 search_best_total,
1130 search_alt_total,
1131 ] {
1132 rec.read_column(search_column);
1133 rec.write_column(search_column);
1134 }
1135 rec.write_column(map_results);
1136 rec.write_column(solve_status);
1137 rec.preflight(&self.runtime).map_err(|e| {
1138 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1139 "component solve preflight failed: {e}"
1140 )))
1141 })?;
1142
1143 let plan_init_kernel = self
1144 .device
1145 .inner()
1146 .get_func(JOINT_SOLVE_MODULE, COMPONENT_PLAN_INIT_KERNEL)
1147 .ok_or_else(|| CarrierError::KernelUnavailable {
1148 detail: format!("{COMPONENT_PLAN_INIT_KERNEL} not resolvable after module load"),
1149 })?;
1150 let entity_owners_kernel = self
1151 .device
1152 .inner()
1153 .get_func(JOINT_SOLVE_MODULE, COMPONENT_ENTITY_OWNERS_KERNEL)
1154 .ok_or_else(|| CarrierError::KernelUnavailable {
1155 detail: format!(
1156 "{COMPONENT_ENTITY_OWNERS_KERNEL} not resolvable after module load"
1157 ),
1158 })?;
1159 let union_kernel = self
1160 .device
1161 .inner()
1162 .get_func(JOINT_SOLVE_MODULE, COMPONENT_UNION_KERNEL)
1163 .ok_or_else(|| CarrierError::KernelUnavailable {
1164 detail: format!("{COMPONENT_UNION_KERNEL} not resolvable after module load"),
1165 })?;
1166 let compress_kernel = self
1167 .device
1168 .inner()
1169 .get_func(JOINT_SOLVE_MODULE, COMPONENT_COMPRESS_KERNEL)
1170 .ok_or_else(|| CarrierError::KernelUnavailable {
1171 detail: format!("{COMPONENT_COMPRESS_KERNEL} not resolvable after module load"),
1172 })?;
1173 let exact_kernel = self
1174 .device
1175 .inner()
1176 .get_func(JOINT_SOLVE_MODULE, COMPONENT_KERNEL)
1177 .ok_or_else(|| CarrierError::KernelUnavailable {
1178 detail: format!("{COMPONENT_KERNEL} not resolvable after module load"),
1179 })?;
1180 let chain_dp_kernel = self
1181 .device
1182 .inner()
1183 .get_func(JOINT_SOLVE_MODULE, CHAIN_DP_KERNEL)
1184 .ok_or_else(|| CarrierError::KernelUnavailable {
1185 detail: format!("{CHAIN_DP_KERNEL} not resolvable after module load"),
1186 })?;
1187 let branch_and_bound_kernel = self
1188 .device
1189 .inner()
1190 .get_func(JOINT_SOLVE_MODULE, BRANCH_AND_BOUND_KERNEL)
1191 .ok_or_else(|| CarrierError::KernelUnavailable {
1192 detail: format!("{BRANCH_AND_BOUND_KERNEL} not resolvable after module load"),
1193 })?;
1194 unsafe {
1199 use std::ffi::c_void;
1200 let scores_p = *scores.device_ptr();
1201 let feasible_p = *feasible_sets.device_ptr();
1202 let arguments_p = *arguments.device_ptr();
1203 let arities_p = *argument_arities.device_ptr();
1204 let domains_p = *domains.device_ptr();
1205 let role_masks_p = *role_masks.device_ptr();
1206 let parents_p = *component_parents.device_ptr();
1207 let entity_owners_p = *component_entity_owners.device_ptr();
1208 let component_count_p = *component_count.device_ptr();
1209 let component_fuel_p = *component_fuel.device_ptr();
1210 let search_assignment_p = *search_assignment.device_ptr();
1211 let search_best_label_p = *search_best_label.device_ptr();
1212 let search_best_total_p = *search_best_total.device_ptr();
1213 let search_alt_total_p = *search_alt_total.device_ptr();
1214 let num_candidates_v = self.candidates as u32;
1215 let num_entities_v = self.entities as u32;
1216 let num_labels_v = self.labels as u32;
1217 let lanes_v = self.domain_lanes as u32;
1218 let max_arity_v = self.max_arity as u32;
1219 let map_p = *map_results.device_ptr();
1220 let status_p = *solve_status.device_ptr();
1221 let planner_threads = 256u32;
1222 let planner_rows = num_candidates_v.max(num_entities_v);
1223 let planner_grid = planner_rows.div_ceil(planner_threads);
1224 let candidate_grid = num_candidates_v.div_ceil(planner_threads);
1225
1226 let mut init_params: [*mut c_void; 6] = [
1227 &parents_p as *const _ as *mut c_void,
1228 &num_candidates_v as *const _ as *mut c_void,
1229 &entity_owners_p as *const _ as *mut c_void,
1230 &num_entities_v as *const _ as *mut c_void,
1231 &component_count_p as *const _ as *mut c_void,
1232 &component_fuel_p as *const _ as *mut c_void,
1233 ];
1234 plan_init_kernel
1235 .launch_on_stream(
1236 &cu_stream,
1237 LaunchConfig {
1238 grid_dim: (planner_grid, 1, 1),
1239 block_dim: (planner_threads, 1, 1),
1240 shared_mem_bytes: 0,
1241 },
1242 &mut init_params[..],
1243 )
1244 .map_err(|e| {
1245 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1246 "component plan initialization failed: {e}"
1247 )))
1248 })?;
1249
1250 let mut owner_params: [*mut c_void; 6] = [
1251 &arguments_p as *const _ as *mut c_void,
1252 &arities_p as *const _ as *mut c_void,
1253 &num_candidates_v as *const _ as *mut c_void,
1254 &max_arity_v as *const _ as *mut c_void,
1255 &num_entities_v as *const _ as *mut c_void,
1256 &entity_owners_p as *const _ as *mut c_void,
1257 ];
1258 entity_owners_kernel
1259 .launch_on_stream(
1260 &cu_stream,
1261 LaunchConfig {
1262 grid_dim: (candidate_grid, 1, 1),
1263 block_dim: (planner_threads, 1, 1),
1264 shared_mem_bytes: 0,
1265 },
1266 &mut owner_params[..],
1267 )
1268 .map_err(|e| {
1269 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1270 "component entity ownership failed: {e}"
1271 )))
1272 })?;
1273
1274 let mut union_params: [*mut c_void; 7] = [
1275 &arguments_p as *const _ as *mut c_void,
1276 &arities_p as *const _ as *mut c_void,
1277 &num_candidates_v as *const _ as *mut c_void,
1278 &max_arity_v as *const _ as *mut c_void,
1279 &num_entities_v as *const _ as *mut c_void,
1280 &entity_owners_p as *const _ as *mut c_void,
1281 &parents_p as *const _ as *mut c_void,
1282 ];
1283 union_kernel
1284 .launch_on_stream(
1285 &cu_stream,
1286 LaunchConfig {
1287 grid_dim: (candidate_grid, 1, 1),
1288 block_dim: (planner_threads, 1, 1),
1289 shared_mem_bytes: 0,
1290 },
1291 &mut union_params[..],
1292 )
1293 .map_err(|e| {
1294 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1295 "component union failed: {e}"
1296 )))
1297 })?;
1298
1299 let mut compress_params: [*mut c_void; 3] = [
1300 &parents_p as *const _ as *mut c_void,
1301 &num_candidates_v as *const _ as *mut c_void,
1302 &component_count_p as *const _ as *mut c_void,
1303 ];
1304 compress_kernel
1305 .launch_on_stream(
1306 &cu_stream,
1307 LaunchConfig {
1308 grid_dim: (candidate_grid, 1, 1),
1309 block_dim: (planner_threads, 1, 1),
1310 shared_mem_bytes: 0,
1311 },
1312 &mut compress_params[..],
1313 )
1314 .map_err(|e| {
1315 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1316 "component compression failed: {e}"
1317 )))
1318 })?;
1319
1320 let mut exact_params: [*mut c_void; 16] = [
1321 &scores_p as *const _ as *mut c_void,
1322 &feasible_p as *const _ as *mut c_void,
1323 &arguments_p as *const _ as *mut c_void,
1324 &arities_p as *const _ as *mut c_void,
1325 &domains_p as *const _ as *mut c_void,
1326 &role_masks_p as *const _ as *mut c_void,
1327 &parents_p as *const _ as *mut c_void,
1328 &component_count_p as *const _ as *mut c_void,
1329 &num_candidates_v as *const _ as *mut c_void,
1330 &num_labels_v as *const _ as *mut c_void,
1331 &lanes_v as *const _ as *mut c_void,
1332 &max_arity_v as *const _ as *mut c_void,
1333 &authorized as *const _ as *mut c_void,
1334 &map_p as *const _ as *mut c_void,
1335 &status_p as *const _ as *mut c_void,
1336 &component_fuel_p as *const _ as *mut c_void,
1337 ];
1338 exact_kernel
1339 .launch_on_stream(
1340 &cu_stream,
1341 LaunchConfig {
1342 grid_dim: (num_candidates_v, 1, 1),
1343 block_dim: (32, 1, 1),
1344 shared_mem_bytes: 0,
1345 },
1346 &mut exact_params[..],
1347 )
1348 .map_err(|e| {
1349 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1350 "component solve launch failed: {e}"
1351 )))
1352 })?;
1353 chain_dp_kernel
1354 .launch_on_stream(
1355 &cu_stream,
1356 LaunchConfig {
1357 grid_dim: (num_candidates_v, 1, 1),
1358 block_dim: (32, 1, 1),
1359 shared_mem_bytes: 0,
1360 },
1361 &mut exact_params[..],
1362 )
1363 .map_err(|e| {
1364 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1365 "device-discovered chain DP launch failed: {e}"
1366 )))
1367 })?;
1368 let mut branch_params: [*mut c_void; 20] = [
1369 &scores_p as *const _ as *mut c_void,
1370 &feasible_p as *const _ as *mut c_void,
1371 &arguments_p as *const _ as *mut c_void,
1372 &arities_p as *const _ as *mut c_void,
1373 &domains_p as *const _ as *mut c_void,
1374 &role_masks_p as *const _ as *mut c_void,
1375 &parents_p as *const _ as *mut c_void,
1376 &component_count_p as *const _ as *mut c_void,
1377 &num_candidates_v as *const _ as *mut c_void,
1378 &num_labels_v as *const _ as *mut c_void,
1379 &lanes_v as *const _ as *mut c_void,
1380 &max_arity_v as *const _ as *mut c_void,
1381 &authorized as *const _ as *mut c_void,
1382 &map_p as *const _ as *mut c_void,
1383 &status_p as *const _ as *mut c_void,
1384 &component_fuel_p as *const _ as *mut c_void,
1385 &search_assignment_p as *const _ as *mut c_void,
1386 &search_best_label_p as *const _ as *mut c_void,
1387 &search_best_total_p as *const _ as *mut c_void,
1388 &search_alt_total_p as *const _ as *mut c_void,
1389 ];
1390 branch_and_bound_kernel
1391 .launch_on_stream(
1392 &cu_stream,
1393 LaunchConfig {
1394 grid_dim: (num_candidates_v, 1, 1),
1395 block_dim: (32, 1, 1),
1396 shared_mem_bytes: 0,
1397 },
1398 &mut branch_params[..],
1399 )
1400 .map_err(|e| {
1401 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1402 "general exact branch-and-bound launch failed: {e}"
1403 )))
1404 })?;
1405 }
1406 rec.commit(&self.runtime).map_err(|e| {
1407 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1408 "component solve commit failed: {e}"
1409 )))
1410 })?;
1411
1412 let mut measured = [0u64; 1];
1416 unsafe {
1417 cudarc::driver::result::stream::synchronize(cu_stream.cu_stream()).map_err(|e| {
1418 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1419 "component solve completion wait failed: {e}"
1420 )))
1421 })?;
1422 cudarc::driver::result::memcpy_dtoh_sync(
1423 &mut measured,
1424 *self.component_buffers[3].device_ptr(),
1425 )
1426 .map_err(|e| {
1427 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1428 "fuel counter readback failed: {e}"
1429 )))
1430 })?;
1431 }
1432 fuel.refund(authorized.saturating_sub(measured[0]));
1433 self.handoff_consumers(&cu_stream)?;
1434 Ok(())
1435 }
1436
1437 pub fn columns(&self) -> impl Iterator<Item = &CudaColumn> {
1441 self.columns.iter()
1442 }
1443
1444 pub fn register_schema(
1450 &mut self,
1451 catalog_sha: &str,
1452 solver_identity: &str,
1453 ) -> Result<(), CarrierError> {
1454 if let Some((catalog, solver)) = &self.registered_schema {
1455 return Err(CarrierError::SchemaAlreadyRegistered {
1456 catalog_sha: catalog.clone(),
1457 solver_identity: solver.clone(),
1458 });
1459 }
1460 self.registered_schema = Some((catalog_sha.to_string(), solver_identity.to_string()));
1461 Ok(())
1462 }
1463}