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_KERNEL: &str = "joint_component_enumerate";
37const MEMOIZED_KERNEL: &str = "joint_label_memoized";
39
40const JOINT_SOLVE_KERNELS: &[&str] = &[
41 FEASIBILITY_KERNEL,
42 TOP2_KERNEL,
43 COMPONENT_KERNEL,
44 MEMOIZED_KERNEL,
45];
46
47const CARRIER_BUDGET_BYTES: u64 = 64 * 1024 * 1024;
51
52#[derive(Debug)]
55pub enum CarrierError {
56 SchemaAlreadyRegistered {
60 catalog_sha: String,
62 solver_identity: String,
64 },
65 Allocation(xlog_core::XlogError),
67 ZeroCapacity {
71 dimension: &'static str,
73 },
74 SchemaNotRegistered,
78 SignaturesAlreadyBound,
81 SignatureShapeMismatch {
83 side: &'static str,
85 expected_words: usize,
87 got_words: usize,
89 },
90 SignaturesUnbound,
92 FeasibilityNotSolved,
95 InvalidComponentPlan {
99 detail: String,
101 },
102 AbstainOutOfRange {
104 abstain_label: u32,
106 labels: usize,
108 },
109 KernelUnavailable {
112 detail: String,
114 },
115 Launch(xlog_core::XlogError),
117 Solver(SolverError),
120}
121
122impl std::fmt::Display for CarrierError {
123 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 match self {
125 CarrierError::SchemaAlreadyRegistered {
126 catalog_sha,
127 solver_identity,
128 } => write!(
129 f,
130 "carrier schema already registered (catalog {catalog_sha}, \
131 solver {solver_identity}); registration is once-per-session"
132 ),
133 CarrierError::Allocation(err) => write!(f, "carrier allocation failed: {err}"),
134 CarrierError::ZeroCapacity { dimension } => write!(
135 f,
136 "carrier capacity dimension {dimension} is zero; refusing \
137 instead of silently clamping"
138 ),
139 CarrierError::SchemaNotRegistered => write!(
140 f,
141 "carrier schema is not registered; signature masks are \
142 catalog-bound and require the catalog anchor first"
143 ),
144 CarrierError::SignaturesAlreadyBound => {
145 write!(f, "signature masks already bound for this session")
146 }
147 CarrierError::SignatureShapeMismatch {
148 side,
149 expected_words,
150 got_words,
151 } => write!(
152 f,
153 "{side} signature mask has {got_words} u64 words, expected \
154 {expected_words} (labels x lanes)"
155 ),
156 CarrierError::SignaturesUnbound => write!(
157 f,
158 "solve refused: signature masks are not bound for this session"
159 ),
160 CarrierError::FeasibilityNotSolved => write!(
161 f,
162 "top-two stage refused: the feasibility stage has not \
163 populated the feasible sets this session"
164 ),
165 CarrierError::InvalidComponentPlan { detail } => {
166 write!(f, "invalid component plan: {detail}")
167 }
168 CarrierError::AbstainOutOfRange {
169 abstain_label,
170 labels,
171 } => write!(
172 f,
173 "abstain label {abstain_label} is outside the label universe \
174 of width {labels}"
175 ),
176 CarrierError::KernelUnavailable { detail } => {
177 write!(f, "joint-solve kernel unavailable: {detail}")
178 }
179 CarrierError::Launch(err) => write!(f, "carrier solve launch failed: {err}"),
180 CarrierError::Solver(err) => write!(f, "carrier solve refused: {err}"),
181 }
182 }
183}
184
185impl std::error::Error for CarrierError {}
186
187struct SilentSink;
189
190impl LoggingSink for SilentSink {
191 fn emit(&self, _record: LogRecord) -> Result<(), SinkError> {
192 Ok(())
193 }
194}
195
196#[derive(Debug, Clone, Copy, PartialEq, Eq)]
200pub enum CarrierBufferId {
201 Domains,
203 Scores,
205 Constraints,
207 Outputs,
209 FeasibleSets,
212 LogicalCounts,
217 MapResults,
223 SolveStatus,
230}
231
232pub struct CarrierExport {
238 pub slice: Arc<crate::memory::TrackedCudaSlice<u8>>,
240 pub stream: Arc<crate::CudaStream>,
242 pub elem_bytes: usize,
244 pub rows: usize,
246 pub cols: usize,
248}
249
250pub struct JointConstraintCarrier {
256 buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 8],
257 columns: [CudaColumn; 7],
258 signatures: Option<[CudaColumn; 2]>,
259 registered_schema: Option<(String, String)>,
260 feasibility_solved: bool,
261 pending_producer_events: Vec<cudarc::driver::sys::CUevent>,
266 pending_consumer_streams: Vec<cudarc::driver::sys::CUstream>,
271 entities: usize,
272 domain_lanes: usize,
273 candidates: usize,
274 labels: usize,
275 device: Arc<CudaDevice>,
276 pool: Arc<StreamPool>,
277 memory: Arc<GpuMemoryManager>,
278 runtime: Arc<XlogDeviceRuntime>,
279}
280
281fn label_words(labels: usize) -> usize {
283 labels.div_ceil(64)
284}
285
286fn shared_column(
293 slice: &Arc<crate::memory::TrackedCudaSlice<u8>>,
294 stream: &Arc<crate::CudaStream>,
295) -> CudaColumn {
296 let tensor = unsafe { crate::DlpackManagedTensor::from_raw(std::ptr::null_mut()) };
297 CudaColumn::dlpack_xlog_owned(Arc::clone(slice), Arc::clone(stream), tensor)
298}
299
300fn ensure_joint_solve_module(device: &Arc<CudaDevice>) -> Result<(), CarrierError> {
304 if JOINT_SOLVE_KERNELS
305 .iter()
306 .all(|k| device.inner().get_func(JOINT_SOLVE_MODULE, k).is_some())
307 {
308 return Ok(());
309 }
310 let cc = crate::provider::detect_compute_capability(device).map_err(|e| {
311 CarrierError::KernelUnavailable {
312 detail: e.to_string(),
313 }
314 })?;
315 let sources = crate::provider::load_module_sources("joint_solve", cc).map_err(|e| {
316 CarrierError::KernelUnavailable {
317 detail: e.to_string(),
318 }
319 })?;
320 let mut load_errors = Vec::new();
321 for source in sources {
322 let attempt = match source {
323 crate::provider::KernelModuleSource::File { path, .. } => device
324 .inner()
325 .load_file(&path, JOINT_SOLVE_MODULE, JOINT_SOLVE_KERNELS)
326 .map_err(|e| format!("{}: {e}", path.display())),
327 crate::provider::KernelModuleSource::EmbeddedPortablePtx { ptx } => device
328 .inner()
329 .load_ptx(
330 cudarc::nvrtc::Ptx::from_src(ptx),
331 JOINT_SOLVE_MODULE,
332 JOINT_SOLVE_KERNELS,
333 )
334 .map_err(|e| format!("embedded portable PTX: {e}")),
335 };
336 match attempt {
337 Ok(()) => return Ok(()),
338 Err(detail) => load_errors.push(detail),
339 }
340 }
341 Err(CarrierError::KernelUnavailable {
342 detail: if load_errors.is_empty() {
343 "no kernel artifact source available".to_string()
344 } else {
345 load_errors.join("; ")
346 },
347 })
348}
349
350impl Drop for JointConstraintCarrier {
351 fn drop(&mut self) {
352 for event in self.pending_producer_events.drain(..) {
355 unsafe {
358 let _ = cudarc::driver::result::event::destroy(event);
359 }
360 }
361 self.pending_consumer_streams.clear();
362 }
363}
364
365impl JointConstraintCarrier {
366 pub fn allocate(
370 device: Arc<CudaDevice>,
371 entities: usize,
372 domain_lanes: usize,
373 candidates: usize,
374 labels: usize,
375 ) -> Result<Self, CarrierError> {
376 for (dimension, value) in [
377 ("entities", entities),
378 ("domain_lanes", domain_lanes),
379 ("candidates", candidates),
380 ("labels", labels),
381 ] {
382 if value == 0 {
383 return Err(CarrierError::ZeroCapacity { dimension });
384 }
385 }
386
387 let pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
388 let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
389 AsyncCudaResource::new(Arc::clone(&device), 0, Arc::clone(&pool)),
390 );
391 let logging: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(LoggingResource::new(
392 async_resource,
393 Arc::new(SilentSink) as Arc<dyn LoggingSink>,
394 ));
395 let budget: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
396 GlobalDeviceBudget::new(logging, CARRIER_BUDGET_BYTES as usize),
397 );
398 let runtime = Arc::new(XlogDeviceRuntime::with_resource(
399 Arc::clone(&device),
400 0,
401 Arc::clone(&pool),
402 budget,
403 ));
404 let memory = Arc::new(GpuMemoryManager::with_runtime(
405 Arc::clone(&device),
406 MemoryBudget::with_limit(CARRIER_BUDGET_BYTES),
407 Arc::clone(&runtime),
408 ));
409
410 ensure_joint_solve_module(&device)?;
411
412 let domains = memory
413 .alloc::<u64>(entities * domain_lanes)
414 .map_err(CarrierError::Allocation)?;
415 let scores = memory
416 .alloc::<f32>(candidates * labels)
417 .map_err(CarrierError::Allocation)?;
418 let constraints = memory
419 .alloc::<u32>(candidates * 2)
420 .map_err(CarrierError::Allocation)?;
421 let outputs = memory
422 .alloc::<u32>(candidates)
423 .map_err(CarrierError::Allocation)?;
424 let feasible_sets = memory
425 .alloc::<u64>(candidates * label_words(labels))
426 .map_err(CarrierError::Allocation)?;
427 let logical_counts = memory.alloc::<u32>(4).map_err(CarrierError::Allocation)?;
428 let map_results = memory
429 .alloc::<u32>(candidates * 4)
430 .map_err(CarrierError::Allocation)?;
431 let solve_status = memory
432 .alloc::<u32>(candidates)
433 .map_err(CarrierError::Allocation)?;
434
435 let buffers: [Arc<crate::memory::TrackedCudaSlice<u8>>; 8] = [
439 Arc::new(domains.into_bytes()),
440 Arc::new(scores.into_bytes()),
441 Arc::new(constraints.into_bytes()),
442 Arc::new(outputs.into_bytes()),
443 Arc::new(feasible_sets.into_bytes()),
444 Arc::new(logical_counts.into_bytes()),
445 Arc::new(map_results.into_bytes()),
446 Arc::new(solve_status.into_bytes()),
447 ];
448 let stream = device.inner().stream().clone();
453 for buffer in &buffers {
454 unsafe {
457 cudarc::driver::result::memset_d8_async(
458 *buffer.device_ptr(),
459 0,
460 buffer.len(),
461 stream.cu_stream(),
462 )
463 .map_err(|e| {
464 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
465 "carrier zero-init failed: {e}"
466 )))
467 })?;
468 }
469 }
470 device.inner().synchronize().map_err(|e| {
471 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
472 "carrier zero-init sync failed: {e}"
473 )))
474 })?;
475
476 let columns = [
477 shared_column(&buffers[0], &stream),
478 shared_column(&buffers[1], &stream),
479 shared_column(&buffers[2], &stream),
480 shared_column(&buffers[3], &stream),
481 shared_column(&buffers[4], &stream),
482 shared_column(&buffers[6], &stream),
483 shared_column(&buffers[7], &stream),
484 ];
485
486 Ok(Self {
487 buffers,
488 columns,
489 signatures: None,
490 registered_schema: None,
491 feasibility_solved: false,
492 pending_producer_events: Vec::new(),
493 pending_consumer_streams: Vec::new(),
494 entities,
495 domain_lanes,
496 candidates,
497 labels,
498 device,
499 pool,
500 memory,
501 runtime,
502 })
503 }
504
505 pub fn note_producer_stream(&mut self, external_stream: u64) -> Result<(), CarrierError> {
513 if external_stream == 0 {
514 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(
515 "null producer stream handle".to_string(),
516 )));
517 }
518 unsafe {
523 let event = cudarc::driver::result::event::create(
524 cudarc::driver::sys::CUevent_flags::CU_EVENT_DISABLE_TIMING,
525 )
526 .map_err(|e| {
527 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
528 "producer event create failed: {e}"
529 )))
530 })?;
531 if let Err(e) = cudarc::driver::result::event::record(
532 event,
533 external_stream as cudarc::driver::sys::CUstream,
534 ) {
535 let _ = cudarc::driver::result::event::destroy(event);
536 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
537 "producer event record failed: {e}"
538 ))));
539 }
540 self.pending_producer_events.push(event);
541 }
542 Ok(())
543 }
544
545 fn drain_producer_waits(&mut self, cu_stream: &crate::CudaStream) -> Result<(), CarrierError> {
549 for event in self.pending_producer_events.drain(..) {
550 unsafe {
553 let wait = cudarc::driver::result::stream::wait_event(
554 cu_stream.cu_stream(),
555 event,
556 cudarc::driver::sys::CUevent_wait_flags::CU_EVENT_WAIT_DEFAULT,
557 );
558 let _ = cudarc::driver::result::event::destroy(event);
559 wait.map_err(|e| {
560 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
561 "producer event wait failed: {e}"
562 )))
563 })?;
564 }
565 }
566 Ok(())
567 }
568
569 pub fn note_consumer_stream(&mut self, external_stream: u64) -> Result<(), CarrierError> {
574 if external_stream == 0 {
575 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(
576 "null consumer stream handle".to_string(),
577 )));
578 }
579 self.pending_consumer_streams
580 .push(external_stream as cudarc::driver::sys::CUstream);
581 Ok(())
582 }
583
584 fn handoff_consumers(&mut self, cu_stream: &crate::CudaStream) -> Result<(), CarrierError> {
587 let consumer_streams = std::mem::take(&mut self.pending_consumer_streams);
588 if consumer_streams.is_empty() {
589 return Ok(());
590 }
591
592 unsafe {
597 let event = cudarc::driver::result::event::create(
598 cudarc::driver::sys::CUevent_flags::CU_EVENT_DISABLE_TIMING,
599 )
600 .map_err(|e| {
601 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
602 "consumer event create failed: {e}"
603 )))
604 })?;
605 if let Err(e) = cudarc::driver::result::event::record(event, cu_stream.cu_stream()) {
606 let _ = cudarc::driver::result::event::destroy(event);
607 return Err(CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
608 "consumer event record failed: {e}"
609 ))));
610 }
611
612 let wait_result = consumer_streams
613 .into_iter()
614 .try_for_each(|consumer_stream| {
615 cudarc::driver::result::stream::wait_event(
616 consumer_stream,
617 event,
618 cudarc::driver::sys::CUevent_wait_flags::CU_EVENT_WAIT_DEFAULT,
619 )
620 .map_err(|e| {
621 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
622 "consumer event wait failed: {e}"
623 )))
624 })
625 });
626 let destroy_result = cudarc::driver::result::event::destroy(event).map_err(|e| {
627 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
628 "consumer event destroy failed: {e}"
629 )))
630 });
631 wait_result?;
632 destroy_result?;
633 }
634 Ok(())
635 }
636
637 pub fn export_buffer(&self, id: CarrierBufferId) -> CarrierExport {
643 let (index, elem_bytes, rows, cols) = match id {
644 CarrierBufferId::Domains => (0, 8, self.entities, self.domain_lanes),
645 CarrierBufferId::Scores => (1, 4, self.candidates, self.labels),
646 CarrierBufferId::Constraints => (2, 4, self.candidates, 2),
647 CarrierBufferId::Outputs => (3, 4, self.candidates, 1),
648 CarrierBufferId::FeasibleSets => (4, 8, self.candidates, label_words(self.labels)),
649 CarrierBufferId::LogicalCounts => (5, 4, 1, 4),
650 CarrierBufferId::MapResults => (6, 4, self.candidates, 4),
651 CarrierBufferId::SolveStatus => (7, 4, self.candidates, 1),
652 };
653 CarrierExport {
654 slice: Arc::clone(&self.buffers[index]),
655 stream: self.device.inner().stream().clone(),
656 elem_bytes,
657 rows,
658 cols,
659 }
660 }
661
662 pub fn bind_signatures(
666 &mut self,
667 head_masks: &[u64],
668 tail_masks: &[u64],
669 ) -> Result<(), CarrierError> {
670 if self.registered_schema.is_none() {
671 return Err(CarrierError::SchemaNotRegistered);
672 }
673 if self.signatures.is_some() {
674 return Err(CarrierError::SignaturesAlreadyBound);
675 }
676 let expected_words = self.labels * self.domain_lanes;
677 for (side, masks) in [("head", head_masks), ("tail", tail_masks)] {
678 if masks.len() != expected_words {
679 return Err(CarrierError::SignatureShapeMismatch {
680 side,
681 expected_words,
682 got_words: masks.len(),
683 });
684 }
685 }
686
687 let head = self.upload_mask(head_masks)?;
688 let tail = self.upload_mask(tail_masks)?;
689 self.signatures = Some([head, tail]);
690 Ok(())
691 }
692
693 fn upload_mask(&self, masks: &[u64]) -> Result<CudaColumn, CarrierError> {
696 let mut slice = self
697 .memory
698 .alloc::<u64>(masks.len())
699 .map_err(CarrierError::Allocation)?;
700 self.device
701 .inner()
702 .htod_sync_copy_into(masks, &mut slice)
703 .map_err(|e| {
704 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
705 "signature mask upload failed: {e}"
706 )))
707 })?;
708 Ok(CudaColumn::owned(slice.into_bytes()))
709 }
710
711 pub fn solve_label_feasibility(
718 &mut self,
719 abstain_label: u32,
720 fuel: &mut FuelMeter,
721 ) -> Result<(), CarrierError> {
722 let result = self.solve_label_feasibility_inner(abstain_label, fuel);
723 if result.is_err() {
724 self.pending_consumer_streams.clear();
725 }
726 result
727 }
728
729 fn solve_label_feasibility_inner(
730 &mut self,
731 abstain_label: u32,
732 fuel: &mut FuelMeter,
733 ) -> Result<(), CarrierError> {
734 if self.registered_schema.is_none() {
735 return Err(CarrierError::SchemaNotRegistered);
736 }
737 if self.signatures.is_none() {
738 return Err(CarrierError::SignaturesUnbound);
739 }
740 if abstain_label as usize >= self.labels {
741 return Err(CarrierError::AbstainOutOfRange {
742 abstain_label,
743 labels: self.labels,
744 });
745 }
746 fuel.charge((self.candidates as u64) * (self.labels as u64))
747 .map_err(CarrierError::Solver)?;
748
749 let stream_id = self.pool.acquire().map_err(|e| {
750 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
751 "no launch stream available: {e:?}"
752 )))
753 })?;
754 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
755 CarrierError::Launch(xlog_core::XlogError::Kernel(
756 "launch stream did not resolve".to_string(),
757 ))
758 })?;
759
760 self.drain_producer_waits(&cu_stream)?;
761
762 let Some(signatures) = &self.signatures else {
763 return Err(CarrierError::SignaturesUnbound);
764 };
765 let [domains, _scores, constraints, outputs, feasible_sets, _map_results, _solve_status] =
766 &self.columns;
767 let [head_masks, tail_masks] = signatures;
768
769 let mut rec = LaunchRecorder::new_strict(stream_id);
770 rec.read_column(domains);
771 rec.read_column(constraints);
772 rec.read_column(head_masks);
773 rec.read_column(tail_masks);
774 rec.write_column(outputs);
775 rec.write_column(feasible_sets);
776 rec.preflight(&self.runtime).map_err(|e| {
777 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
778 "solve launch preflight failed: {e}"
779 )))
780 })?;
781
782 let kernel = self
783 .device
784 .inner()
785 .get_func(JOINT_SOLVE_MODULE, FEASIBILITY_KERNEL)
786 .ok_or_else(|| CarrierError::KernelUnavailable {
787 detail: format!("{FEASIBILITY_KERNEL} not resolvable after module load"),
788 })?;
789 let block = 256u32;
790 let grid = (self.candidates as u32).div_ceil(block);
791 unsafe {
798 kernel
799 .launch_on_stream(
800 &cu_stream,
801 LaunchConfig {
802 grid_dim: (grid, 1, 1),
803 block_dim: (block, 1, 1),
804 shared_mem_bytes: 0,
805 },
806 (
807 *domains.device_ptr(),
808 *constraints.device_ptr(),
809 *head_masks.device_ptr(),
810 *tail_masks.device_ptr(),
811 self.entities as u32,
812 self.candidates as u32,
813 self.labels as u32,
814 self.domain_lanes as u32,
815 abstain_label,
816 *outputs.device_ptr(),
817 *feasible_sets.device_ptr(),
818 ),
819 )
820 .map_err(|e| {
821 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
822 "solve launch failed: {e}"
823 )))
824 })?;
825 }
826 rec.commit(&self.runtime).map_err(|e| {
827 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
828 "solve launch commit failed: {e}"
829 )))
830 })?;
831 self.handoff_consumers(&cu_stream)?;
832 self.feasibility_solved = true;
833 Ok(())
834 }
835
836 pub fn solve_label_map_top2(&mut self, fuel: &mut FuelMeter) -> Result<(), CarrierError> {
849 let result = self.solve_label_map_top2_inner(fuel);
850 if result.is_err() {
851 self.pending_consumer_streams.clear();
852 }
853 result
854 }
855
856 fn solve_label_map_top2_inner(&mut self, fuel: &mut FuelMeter) -> Result<(), CarrierError> {
857 if self.registered_schema.is_none() {
858 return Err(CarrierError::SchemaNotRegistered);
859 }
860 if !self.feasibility_solved {
861 return Err(CarrierError::FeasibilityNotSolved);
862 }
863 fuel.charge((self.candidates as u64) * (self.labels as u64))
864 .map_err(CarrierError::Solver)?;
865
866 let stream_id = self.pool.acquire().map_err(|e| {
867 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
868 "no launch stream available: {e:?}"
869 )))
870 })?;
871 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
872 CarrierError::Launch(xlog_core::XlogError::Kernel(
873 "launch stream did not resolve".to_string(),
874 ))
875 })?;
876
877 self.drain_producer_waits(&cu_stream)?;
878
879 let [_domains, scores, _constraints, _outputs, feasible_sets, map_results, _solve_status] =
880 &self.columns;
881
882 let mut rec = LaunchRecorder::new_strict(stream_id);
883 rec.read_column(scores);
884 rec.read_column(feasible_sets);
885 rec.write_column(map_results);
886 rec.preflight(&self.runtime).map_err(|e| {
887 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
888 "top-two launch preflight failed: {e}"
889 )))
890 })?;
891
892 let kernel = self
893 .device
894 .inner()
895 .get_func(JOINT_SOLVE_MODULE, TOP2_KERNEL)
896 .ok_or_else(|| CarrierError::KernelUnavailable {
897 detail: format!("{TOP2_KERNEL} not resolvable after module load"),
898 })?;
899 let block = 256u32;
900 let grid = (self.candidates as u32).div_ceil(block);
901 unsafe {
905 kernel
906 .launch_on_stream(
907 &cu_stream,
908 LaunchConfig {
909 grid_dim: (grid, 1, 1),
910 block_dim: (block, 1, 1),
911 shared_mem_bytes: 0,
912 },
913 (
914 *scores.device_ptr(),
915 *feasible_sets.device_ptr(),
916 self.candidates as u32,
917 self.labels as u32,
918 *map_results.device_ptr(),
919 ),
920 )
921 .map_err(|e| {
922 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
923 "top-two launch failed: {e}"
924 )))
925 })?;
926 }
927 rec.commit(&self.runtime).map_err(|e| {
928 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
929 "top-two launch commit failed: {e}"
930 )))
931 })?;
932 self.handoff_consumers(&cu_stream)?;
933 Ok(())
934 }
935
936 pub fn solve_components_exact(
951 &mut self,
952 comp_offsets: &[u32],
953 comp_indices: &[u32],
954 fuel: &mut FuelMeter,
955 ) -> Result<(), CarrierError> {
956 let result = self.solve_components_exact_inner(comp_offsets, comp_indices, fuel);
957 if result.is_err() {
958 self.pending_consumer_streams.clear();
959 }
960 result
961 }
962
963 fn solve_components_exact_inner(
964 &mut self,
965 comp_offsets: &[u32],
966 comp_indices: &[u32],
967 fuel: &mut FuelMeter,
968 ) -> Result<(), CarrierError> {
969 if self.registered_schema.is_none() {
970 return Err(CarrierError::SchemaNotRegistered);
971 }
972 if !self.feasibility_solved {
973 return Err(CarrierError::FeasibilityNotSolved);
974 }
975 let invalid = |detail: String| CarrierError::InvalidComponentPlan { detail };
976 if comp_offsets.first() != Some(&0)
977 || comp_offsets.last().copied() != Some(comp_indices.len() as u32)
978 {
979 return Err(invalid(format!(
980 "offsets must run 0..={}, got first {:?} last {:?}",
981 comp_indices.len(),
982 comp_offsets.first(),
983 comp_offsets.last()
984 )));
985 }
986 if comp_offsets.windows(2).any(|w| w[0] > w[1]) {
987 return Err(invalid("offsets are not monotone".to_string()));
988 }
989 let mut seen = vec![false; self.candidates];
990 for &cand in comp_indices {
991 let slot = seen
992 .get_mut(cand as usize)
993 .ok_or_else(|| invalid(format!("candidate {cand} outside capacity")))?;
994 if *slot {
995 return Err(invalid(format!("candidate {cand} listed twice")));
996 }
997 *slot = true;
998 }
999 let num_components = comp_offsets.len() - 1;
1000 if num_components == 0 {
1001 return Ok(());
1002 }
1003
1004 let fuel_per_component = fuel.remaining() / num_components as u64;
1010 let authorized = fuel_per_component * num_components as u64;
1011 fuel.charge(authorized).map_err(CarrierError::Solver)?;
1012
1013 let stream_id = self.pool.acquire().map_err(|e| {
1014 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1015 "no launch stream available: {e:?}"
1016 )))
1017 })?;
1018 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
1019 CarrierError::Launch(xlog_core::XlogError::Kernel(
1020 "launch stream did not resolve".to_string(),
1021 ))
1022 })?;
1023 self.drain_producer_waits(&cu_stream)?;
1024
1025 let offsets_col = self.upload_plan(comp_offsets)?;
1030 let indices_col = self.upload_plan(comp_indices)?;
1031 let fuel_words = [0u32, 0u32];
1032 let fuel_col = self.upload_plan(&fuel_words)?;
1033
1034 let Some(signatures) = &self.signatures else {
1035 return Err(CarrierError::SignaturesUnbound);
1036 };
1037 let [domains, scores, constraints, _outputs, feasible_sets, map_results, solve_status] =
1038 &self.columns;
1039 let [head_masks, tail_masks] = signatures;
1040
1041 let mut rec = LaunchRecorder::new_strict(stream_id);
1042 rec.read_column(scores);
1043 rec.read_column(feasible_sets);
1044 rec.read_column(constraints);
1045 rec.read_column(domains);
1046 rec.read_column(head_masks);
1047 rec.read_column(tail_masks);
1048 rec.read_column(&offsets_col);
1049 rec.read_column(&indices_col);
1050 rec.write_column(map_results);
1051 rec.write_column(solve_status);
1052 rec.write_column(&fuel_col);
1053 rec.preflight(&self.runtime).map_err(|e| {
1054 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1055 "component solve preflight failed: {e}"
1056 )))
1057 })?;
1058
1059 let kernel = self
1060 .device
1061 .inner()
1062 .get_func(JOINT_SOLVE_MODULE, COMPONENT_KERNEL)
1063 .ok_or_else(|| CarrierError::KernelUnavailable {
1064 detail: format!("{COMPONENT_KERNEL} not resolvable after module load"),
1065 })?;
1066 unsafe {
1075 use std::ffi::c_void;
1076 let scores_p = *scores.device_ptr();
1077 let feasible_p = *feasible_sets.device_ptr();
1078 let pairs_p = *constraints.device_ptr();
1079 let domains_p = *domains.device_ptr();
1080 let head_p = *head_masks.device_ptr();
1081 let tail_p = *tail_masks.device_ptr();
1082 let offsets_p = *offsets_col.device_ptr();
1083 let indices_p = *indices_col.device_ptr();
1084 let num_components_v = num_components as u32;
1085 let num_labels_v = self.labels as u32;
1086 let lanes_v = self.domain_lanes as u32;
1087 let map_p = *map_results.device_ptr();
1088 let status_p = *solve_status.device_ptr();
1089 let fuel_p = *fuel_col.device_ptr();
1090 let mut params: [*mut c_void; 15] = [
1091 &scores_p as *const _ as *mut c_void,
1092 &feasible_p as *const _ as *mut c_void,
1093 &pairs_p as *const _ as *mut c_void,
1094 &domains_p as *const _ as *mut c_void,
1095 &head_p as *const _ as *mut c_void,
1096 &tail_p as *const _ as *mut c_void,
1097 &offsets_p as *const _ as *mut c_void,
1098 &indices_p as *const _ as *mut c_void,
1099 &num_components_v as *const _ as *mut c_void,
1100 &num_labels_v as *const _ as *mut c_void,
1101 &lanes_v as *const _ as *mut c_void,
1102 &fuel_per_component as *const _ as *mut c_void,
1103 &map_p as *const _ as *mut c_void,
1104 &status_p as *const _ as *mut c_void,
1105 &fuel_p as *const _ as *mut c_void,
1106 ];
1107 kernel
1108 .launch_on_stream(
1109 &cu_stream,
1110 LaunchConfig {
1111 grid_dim: (num_components as u32, 1, 1),
1112 block_dim: (32, 1, 1),
1113 shared_mem_bytes: 0,
1114 },
1115 &mut params[..],
1116 )
1117 .map_err(|e| {
1118 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1119 "component solve launch failed: {e}"
1120 )))
1121 })?;
1122 }
1123 rec.commit(&self.runtime).map_err(|e| {
1124 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1125 "component solve commit failed: {e}"
1126 )))
1127 })?;
1128
1129 let mut measured = [0u64; 1];
1133 unsafe {
1134 cudarc::driver::result::stream::synchronize(cu_stream.cu_stream()).map_err(|e| {
1135 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1136 "component solve completion wait failed: {e}"
1137 )))
1138 })?;
1139 cudarc::driver::result::memcpy_dtoh_sync(&mut measured, *fuel_col.device_ptr())
1140 .map_err(|e| {
1141 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1142 "fuel counter readback failed: {e}"
1143 )))
1144 })?;
1145 }
1146 fuel.refund(authorized.saturating_sub(measured[0]));
1147 self.handoff_consumers(&cu_stream)?;
1148 Ok(())
1149 }
1150
1151 pub fn solve_components_memoized(
1159 &mut self,
1160 comp_offsets: &[u32],
1161 comp_indices: &[u32],
1162 pinned_width: u32,
1163 fuel: &mut FuelMeter,
1164 ) -> Result<(), CarrierError> {
1165 let result =
1166 self.solve_components_memoized_inner(comp_offsets, comp_indices, pinned_width, fuel);
1167 if result.is_err() {
1168 self.pending_consumer_streams.clear();
1169 }
1170 result
1171 }
1172
1173 fn solve_components_memoized_inner(
1174 &mut self,
1175 comp_offsets: &[u32],
1176 comp_indices: &[u32],
1177 pinned_width: u32,
1178 fuel: &mut FuelMeter,
1179 ) -> Result<(), CarrierError> {
1180 if self.registered_schema.is_none() {
1181 return Err(CarrierError::SchemaNotRegistered);
1182 }
1183 if !self.feasibility_solved {
1184 return Err(CarrierError::FeasibilityNotSolved);
1185 }
1186 let invalid = |detail: String| CarrierError::InvalidComponentPlan { detail };
1187 if comp_offsets.first() != Some(&0)
1188 || comp_offsets.last().copied() != Some(comp_indices.len() as u32)
1189 {
1190 return Err(invalid(format!(
1191 "offsets must run 0..={}, got first {:?} last {:?}",
1192 comp_indices.len(),
1193 comp_offsets.first(),
1194 comp_offsets.last()
1195 )));
1196 }
1197 if comp_offsets.windows(2).any(|w| w[0] > w[1]) {
1198 return Err(invalid("offsets are not monotone".to_string()));
1199 }
1200 let mut seen = vec![false; self.candidates];
1201 for &cand in comp_indices {
1202 let slot = seen
1203 .get_mut(cand as usize)
1204 .ok_or_else(|| invalid(format!("candidate {cand} outside capacity")))?;
1205 if *slot {
1206 return Err(invalid(format!("candidate {cand} listed twice")));
1207 }
1208 *slot = true;
1209 }
1210 let num_components = comp_offsets.len() - 1;
1211 if num_components == 0 {
1212 return Ok(());
1213 }
1214
1215 let fuel_per_component = fuel.remaining() / num_components as u64;
1216 let authorized = fuel_per_component * num_components as u64;
1217 fuel.charge(authorized).map_err(CarrierError::Solver)?;
1218
1219 let stream_id = self.pool.acquire().map_err(|e| {
1220 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1221 "no launch stream available: {e:?}"
1222 )))
1223 })?;
1224 let cu_stream = self.pool.resolve(stream_id).ok_or_else(|| {
1225 CarrierError::Launch(xlog_core::XlogError::Kernel(
1226 "launch stream did not resolve".to_string(),
1227 ))
1228 })?;
1229 self.drain_producer_waits(&cu_stream)?;
1230
1231 let offsets_col = self.upload_plan(comp_offsets)?;
1232 let indices_col = self.upload_plan(comp_indices)?;
1233 let fuel_words = [0u32, 0u32];
1234 let fuel_col = self.upload_plan(&fuel_words)?;
1235
1236 let Some(signatures) = &self.signatures else {
1237 return Err(CarrierError::SignaturesUnbound);
1238 };
1239 let [domains, scores, constraints, _outputs, feasible_sets, map_results, solve_status] =
1240 &self.columns;
1241 let [head_masks, tail_masks] = signatures;
1242
1243 let mut rec = LaunchRecorder::new_strict(stream_id);
1244 rec.read_column(scores);
1245 rec.read_column(feasible_sets);
1246 rec.read_column(constraints);
1247 rec.read_column(domains);
1248 rec.read_column(head_masks);
1249 rec.read_column(tail_masks);
1250 rec.read_column(&offsets_col);
1251 rec.read_column(&indices_col);
1252 rec.write_column(map_results);
1253 rec.write_column(solve_status);
1254 rec.write_column(&fuel_col);
1255 rec.preflight(&self.runtime).map_err(|e| {
1256 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1257 "memoized solve preflight failed: {e}"
1258 )))
1259 })?;
1260
1261 let kernel = self
1262 .device
1263 .inner()
1264 .get_func(JOINT_SOLVE_MODULE, MEMOIZED_KERNEL)
1265 .ok_or_else(|| CarrierError::KernelUnavailable {
1266 detail: format!("{MEMOIZED_KERNEL} not resolvable after module load"),
1267 })?;
1268 unsafe {
1277 use std::ffi::c_void;
1278 let scores_p = *scores.device_ptr();
1279 let feasible_p = *feasible_sets.device_ptr();
1280 let pairs_p = *constraints.device_ptr();
1281 let domains_p = *domains.device_ptr();
1282 let head_p = *head_masks.device_ptr();
1283 let tail_p = *tail_masks.device_ptr();
1284 let offsets_p = *offsets_col.device_ptr();
1285 let indices_p = *indices_col.device_ptr();
1286 let num_components_v = num_components as u32;
1287 let num_labels_v = self.labels as u32;
1288 let lanes_v = self.domain_lanes as u32;
1289 let map_p = *map_results.device_ptr();
1290 let status_p = *solve_status.device_ptr();
1291 let fuel_p = *fuel_col.device_ptr();
1292 let mut params: [*mut c_void; 16] = [
1293 &scores_p as *const _ as *mut c_void,
1294 &feasible_p as *const _ as *mut c_void,
1295 &pairs_p as *const _ as *mut c_void,
1296 &domains_p as *const _ as *mut c_void,
1297 &head_p as *const _ as *mut c_void,
1298 &tail_p as *const _ as *mut c_void,
1299 &offsets_p as *const _ as *mut c_void,
1300 &indices_p as *const _ as *mut c_void,
1301 &num_components_v as *const _ as *mut c_void,
1302 &num_labels_v as *const _ as *mut c_void,
1303 &lanes_v as *const _ as *mut c_void,
1304 &pinned_width as *const _ as *mut c_void,
1305 &fuel_per_component as *const _ as *mut c_void,
1306 &map_p as *const _ as *mut c_void,
1307 &status_p as *const _ as *mut c_void,
1308 &fuel_p as *const _ as *mut c_void,
1309 ];
1310 kernel
1311 .launch_on_stream(
1312 &cu_stream,
1313 LaunchConfig {
1314 grid_dim: (num_components as u32, 1, 1),
1315 block_dim: (32, 1, 1),
1316 shared_mem_bytes: 0,
1317 },
1318 &mut params[..],
1319 )
1320 .map_err(|e| {
1321 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1322 "memoized solve launch failed: {e}"
1323 )))
1324 })?;
1325 }
1326 rec.commit(&self.runtime).map_err(|e| {
1327 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1328 "memoized solve commit failed: {e}"
1329 )))
1330 })?;
1331 let mut measured = [0u64; 1];
1335 unsafe {
1336 cudarc::driver::result::stream::synchronize(cu_stream.cu_stream()).map_err(|e| {
1337 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1338 "memoized solve completion wait failed: {e}"
1339 )))
1340 })?;
1341 cudarc::driver::result::memcpy_dtoh_sync(&mut measured, *fuel_col.device_ptr())
1342 .map_err(|e| {
1343 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1344 "fuel counter readback failed: {e}"
1345 )))
1346 })?;
1347 }
1348 fuel.refund(authorized.saturating_sub(measured[0]));
1349 self.handoff_consumers(&cu_stream)?;
1350 Ok(())
1351 }
1352
1353 fn upload_plan(&self, words: &[u32]) -> Result<CudaColumn, CarrierError> {
1356 let mut slice = self
1357 .memory
1358 .alloc::<u32>(words.len())
1359 .map_err(CarrierError::Allocation)?;
1360 self.device
1361 .inner()
1362 .htod_sync_copy_into(words, &mut slice)
1363 .map_err(|e| {
1364 CarrierError::Launch(xlog_core::XlogError::Kernel(format!(
1365 "component plan upload failed: {e}"
1366 )))
1367 })?;
1368 Ok(CudaColumn::owned(slice.into_bytes()))
1369 }
1370
1371 pub fn columns(&self) -> impl Iterator<Item = &CudaColumn> {
1375 self.columns.iter()
1376 }
1377
1378 pub fn register_schema(
1384 &mut self,
1385 catalog_sha: &str,
1386 solver_identity: &str,
1387 ) -> Result<(), CarrierError> {
1388 if let Some((catalog, solver)) = &self.registered_schema {
1389 return Err(CarrierError::SchemaAlreadyRegistered {
1390 catalog_sha: catalog.clone(),
1391 solver_identity: solver.clone(),
1392 });
1393 }
1394 self.registered_schema = Some((catalog_sha.to_string(), solver_identity.to_string()));
1395 Ok(())
1396 }
1397}