1use std::collections::HashMap;
25use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
26use std::sync::{Arc, Mutex, OnceLock, RwLock};
27
28use fuel_ir::backend::{BackendCapabilities, SubstrateClass, TransferPath};
29use fuel_ir::dispatch::OpKind;
30use fuel_ir::probe::BackendId;
31use fuel_ir::{DType, DeviceLocation, Error, Layout, Result};
32
33#[cfg(feature = "cuda")]
34use crate::kernel::KernelCaps;
35use crate::kernel::{KernelBindingTable, MatmulM, OpParams};
36use fuel_memory::{BackendStorage, Storage};
37
38#[derive(Debug, Default)]
41pub struct CapabilityRegistry {
42 backends: Vec<BackendCapabilities>,
43}
44
45impl CapabilityRegistry {
46 pub fn new() -> Self {
48 Self {
49 backends: Vec::new(),
50 }
51 }
52
53 pub fn register(&mut self, caps: BackendCapabilities) {
58 self.backends.push(caps);
59 }
60
61 pub fn backends(&self) -> &[BackendCapabilities] {
63 &self.backends
64 }
65
66 pub fn find_backends(&self, op: OpKind, dtype: DType) -> Vec<&BackendCapabilities> {
68 self.backends
69 .iter()
70 .filter(|caps| caps.op_dtype_support.contains(&(op, dtype)))
71 .collect()
72 }
73
74 pub fn find_backend_for(&self, op: OpKind, dtype: DType) -> Result<&BackendCapabilities> {
84 for caps in &self.backends {
85 if caps.op_dtype_support.contains(&(op, dtype)) {
86 return Ok(caps);
87 }
88 }
89 Err(Error::NoBackendForOp {
94 op,
95 dtypes: vec![dtype],
96 available_backends: self.backends.iter().map(|c| c.backend_id).collect(),
97 supported_combinations: self
98 .backends
99 .iter()
100 .flat_map(|c| {
101 c.op_dtype_support
102 .iter()
103 .map(|&(o, d)| (c.backend_id, o, vec![d]))
104 })
105 .collect(),
106 }
107 .bt())
108 }
109
110 pub fn build_transfer_matrix(&self) -> TransferMatrix {
115 let mut entries = HashMap::new();
116 for caps in &self.backends {
117 for (dst, path) in &caps.transfer_paths {
118 entries.insert((caps.device_location, *dst), *path);
119 }
120 }
121 TransferMatrix { entries }
122 }
123}
124
125#[derive(Debug, Default)]
130pub struct TransferMatrix {
131 entries: HashMap<(DeviceLocation, DeviceLocation), TransferPath>,
132}
133
134impl TransferMatrix {
135 pub fn path(&self, src: DeviceLocation, dst: DeviceLocation) -> Option<TransferPath> {
140 if src == dst {
141 return Some(TransferPath::SameDevice);
142 }
143 self.entries.get(&(src, dst)).copied()
144 }
145
146 pub fn path_or_staging(&self, src: DeviceLocation, dst: DeviceLocation) -> TransferPath {
154 self.path(src, dst).unwrap_or(TransferPath::HostStaging)
155 }
156
157 pub fn entries(
159 &self,
160 ) -> impl Iterator<Item = (&(DeviceLocation, DeviceLocation), &TransferPath)> {
161 self.entries.iter()
162 }
163}
164
165pub fn resolve_target_backend(
178 registry: &CapabilityRegistry,
179 op: OpKind,
180 dtype: DType,
181) -> Result<BackendId> {
182 registry
183 .find_backend_for(op, dtype)
184 .map(|caps| caps.backend_id)
185}
186
187pub fn resolve_target_backend_residency_aware(
203 registry: &CapabilityRegistry,
204 op: OpKind,
205 dtype: DType,
206 input_locations: &[DeviceLocation],
207) -> Result<BackendId> {
208 let candidates = registry.find_backends(op, dtype);
209 if candidates.is_empty() {
210 return registry.find_backend_for(op, dtype).map(|c| c.backend_id);
212 }
213 for caps in &candidates {
216 if input_locations.contains(&caps.device_location) {
217 return Ok(caps.backend_id);
218 }
219 }
220 Ok(candidates[0].backend_id)
223}
224
225pub fn cpu_input(s: &Storage) -> Result<&fuel_cpu_backend::CpuStorageBytes> {
232 match &s.inner {
233 BackendStorage::Cpu(c) => Ok(c),
234 #[allow(unreachable_patterns)]
235 _ => Err(Error::Msg("cpu kernel wrapper called with non-CPU input".to_string()).bt()),
236 }
237}
238
239pub fn cpu_output(s: &mut Storage) -> Result<&mut fuel_cpu_backend::CpuStorageBytes> {
245 match &mut s.inner {
246 BackendStorage::Cpu(c) => Ok(c),
247 #[allow(unreachable_patterns)]
248 _ => Err(Error::Msg("cpu kernel wrapper called with non-CPU output".to_string()).bt()),
249 }
250}
251
252pub fn read_storage(arc: &Arc<RwLock<Storage>>) -> Result<std::sync::RwLockReadGuard<'_, Storage>> {
257 arc.read()
258 .map_err(|_| Error::Msg("kernel wrapper: storage RwLock poisoned (read)".to_string()).bt())
259}
260
261pub fn write_storage(
262 arc: &Arc<RwLock<Storage>>,
263) -> Result<std::sync::RwLockWriteGuard<'_, Storage>> {
264 arc.write()
265 .map_err(|_| Error::Msg("kernel wrapper: storage RwLock poisoned (write)".to_string()).bt())
266}
267
268macro_rules! cpu_binary_wrapper {
273 ($wrapper:ident, $kernel:path, $op_name:literal) => {
274 pub(crate) fn $wrapper(
277 inputs: &[Arc<RwLock<Storage>>],
278 outputs: &mut [Arc<RwLock<Storage>>],
279 _layouts: &[Layout],
280 _params: &OpParams,
281 ) -> Result<()> {
282 if inputs.len() != 2 {
283 return Err(Error::Msg(format!(
284 "{} wrapper expects 2 inputs, got {}",
285 $op_name,
286 inputs.len(),
287 ))
288 .bt());
289 }
290 if outputs.len() != 1 {
291 return Err(Error::Msg(format!(
292 "{} wrapper expects 1 output, got {}",
293 $op_name,
294 outputs.len(),
295 ))
296 .bt());
297 }
298 let lhs_guard = read_storage(&inputs[0])?;
299 let rhs_guard = read_storage(&inputs[1])?;
300 let mut out_guard = write_storage(&outputs[0])?;
301 let lhs_cpu = cpu_input(&lhs_guard)?;
302 let rhs_cpu = cpu_input(&rhs_guard)?;
303 let out_cpu = cpu_output(&mut out_guard)?;
304 $kernel(lhs_cpu, rhs_cpu, out_cpu)
305 }
306 };
307}
308
309macro_rules! cpu_unary_wrapper {
312 ($wrapper:ident, $kernel:path, $op_name:literal) => {
313 pub(crate) fn $wrapper(
316 inputs: &[Arc<RwLock<Storage>>],
317 outputs: &mut [Arc<RwLock<Storage>>],
318 _layouts: &[Layout],
319 _params: &OpParams,
320 ) -> Result<()> {
321 if inputs.len() != 1 {
322 return Err(Error::Msg(format!(
323 "{} wrapper expects 1 input, got {}",
324 $op_name,
325 inputs.len(),
326 ))
327 .bt());
328 }
329 if outputs.len() != 1 {
330 return Err(Error::Msg(format!(
331 "{} wrapper expects 1 output, got {}",
332 $op_name,
333 outputs.len(),
334 ))
335 .bt());
336 }
337 let in_guard = read_storage(&inputs[0])?;
338 let mut out_guard = write_storage(&outputs[0])?;
339 let in_cpu = cpu_input(&in_guard)?;
340 let out_cpu = cpu_output(&mut out_guard)?;
341 $kernel(in_cpu, out_cpu)
342 }
343 };
344}
345
346cpu_binary_wrapper!(
347 add_elementwise_f32_cpu_wrapper,
348 fuel_cpu_backend::byte_kernels::add_f32,
349 "add_elementwise"
350);
351cpu_binary_wrapper!(
352 sub_elementwise_f32_cpu_wrapper,
353 fuel_cpu_backend::byte_kernels::sub_f32,
354 "sub_elementwise"
355);
356cpu_binary_wrapper!(
357 mul_elementwise_f32_cpu_wrapper,
358 fuel_cpu_backend::byte_kernels::mul_f32,
359 "mul_elementwise"
360);
361cpu_binary_wrapper!(
362 div_elementwise_f32_cpu_wrapper,
363 fuel_cpu_backend::byte_kernels::div_f32,
364 "div_elementwise"
365);
366
367cpu_unary_wrapper!(
368 relu_elementwise_f32_cpu_wrapper,
369 fuel_cpu_backend::byte_kernels::relu_f32,
370 "relu_elementwise"
371);
372cpu_unary_wrapper!(
373 neg_elementwise_f32_cpu_wrapper,
374 fuel_cpu_backend::byte_kernels::neg_f32,
375 "neg_elementwise"
376);
377cpu_unary_wrapper!(
378 sqr_elementwise_f32_cpu_wrapper,
379 fuel_cpu_backend::byte_kernels::sqr_f32,
380 "sqr_elementwise"
381);
382cpu_unary_wrapper!(
383 sqrt_elementwise_f32_cpu_wrapper,
384 fuel_cpu_backend::byte_kernels::sqrt_f32,
385 "sqrt_elementwise"
386);
387cpu_unary_wrapper!(
388 recip_elementwise_f32_cpu_wrapper,
389 fuel_cpu_backend::byte_kernels::recip_f32,
390 "recip_elementwise"
391);
392cpu_unary_wrapper!(
393 abs_elementwise_f32_cpu_wrapper,
394 fuel_cpu_backend::byte_kernels::abs_f32,
395 "abs_elementwise"
396);
397cpu_unary_wrapper!(
398 tanh_elementwise_f32_cpu_wrapper,
399 fuel_cpu_backend::byte_kernels::tanh_f32,
400 "tanh_elementwise"
401);
402cpu_unary_wrapper!(
403 exp_elementwise_f32_cpu_wrapper,
404 fuel_cpu_backend::byte_kernels::exp_f32,
405 "exp_elementwise"
406);
407cpu_unary_wrapper!(
408 log_elementwise_f32_cpu_wrapper,
409 fuel_cpu_backend::byte_kernels::log_f32,
410 "log_elementwise"
411);
412cpu_unary_wrapper!(
413 sin_elementwise_f32_cpu_wrapper,
414 fuel_cpu_backend::byte_kernels::sin_f32,
415 "sin_elementwise"
416);
417cpu_unary_wrapper!(
418 cos_elementwise_f32_cpu_wrapper,
419 fuel_cpu_backend::byte_kernels::cos_f32,
420 "cos_elementwise"
421);
422cpu_unary_wrapper!(
423 sigmoid_elementwise_f32_cpu_wrapper,
424 fuel_cpu_backend::byte_kernels::sigmoid_f32,
425 "sigmoid_elementwise"
426);
427cpu_unary_wrapper!(
428 silu_elementwise_f32_cpu_wrapper,
429 fuel_cpu_backend::byte_kernels::silu_f32,
430 "silu_elementwise"
431);
432cpu_unary_wrapper!(
433 gelu_elementwise_f32_cpu_wrapper,
434 fuel_cpu_backend::byte_kernels::gelu_f32,
435 "gelu_elementwise"
436);
437cpu_unary_wrapper!(
438 step_elementwise_f32_cpu_wrapper,
439 fuel_cpu_backend::byte_kernels::step_f32,
440 "step_elementwise"
441);
442
443cpu_binary_wrapper!(
449 add_elementwise_f64_cpu_wrapper,
450 fuel_cpu_backend::byte_kernels::add_f64,
451 "add_elementwise"
452);
453cpu_binary_wrapper!(
454 sub_elementwise_f64_cpu_wrapper,
455 fuel_cpu_backend::byte_kernels::sub_f64,
456 "sub_elementwise"
457);
458cpu_binary_wrapper!(
459 mul_elementwise_f64_cpu_wrapper,
460 fuel_cpu_backend::byte_kernels::mul_f64,
461 "mul_elementwise"
462);
463cpu_binary_wrapper!(
464 div_elementwise_f64_cpu_wrapper,
465 fuel_cpu_backend::byte_kernels::div_f64,
466 "div_elementwise"
467);
468
469cpu_unary_wrapper!(
470 relu_elementwise_f64_cpu_wrapper,
471 fuel_cpu_backend::byte_kernels::relu_f64,
472 "relu_elementwise"
473);
474cpu_unary_wrapper!(
475 neg_elementwise_f64_cpu_wrapper,
476 fuel_cpu_backend::byte_kernels::neg_f64,
477 "neg_elementwise"
478);
479cpu_unary_wrapper!(
480 sqr_elementwise_f64_cpu_wrapper,
481 fuel_cpu_backend::byte_kernels::sqr_f64,
482 "sqr_elementwise"
483);
484cpu_unary_wrapper!(
485 sqrt_elementwise_f64_cpu_wrapper,
486 fuel_cpu_backend::byte_kernels::sqrt_f64,
487 "sqrt_elementwise"
488);
489cpu_unary_wrapper!(
490 recip_elementwise_f64_cpu_wrapper,
491 fuel_cpu_backend::byte_kernels::recip_f64,
492 "recip_elementwise"
493);
494cpu_unary_wrapper!(
495 abs_elementwise_f64_cpu_wrapper,
496 fuel_cpu_backend::byte_kernels::abs_f64,
497 "abs_elementwise"
498);
499cpu_unary_wrapper!(
500 tanh_elementwise_f64_cpu_wrapper,
501 fuel_cpu_backend::byte_kernels::tanh_f64,
502 "tanh_elementwise"
503);
504cpu_unary_wrapper!(
505 exp_elementwise_f64_cpu_wrapper,
506 fuel_cpu_backend::byte_kernels::exp_f64,
507 "exp_elementwise"
508);
509cpu_unary_wrapper!(
510 log_elementwise_f64_cpu_wrapper,
511 fuel_cpu_backend::byte_kernels::log_f64,
512 "log_elementwise"
513);
514cpu_unary_wrapper!(
515 sin_elementwise_f64_cpu_wrapper,
516 fuel_cpu_backend::byte_kernels::sin_f64,
517 "sin_elementwise"
518);
519cpu_unary_wrapper!(
520 cos_elementwise_f64_cpu_wrapper,
521 fuel_cpu_backend::byte_kernels::cos_f64,
522 "cos_elementwise"
523);
524cpu_unary_wrapper!(
525 sigmoid_elementwise_f64_cpu_wrapper,
526 fuel_cpu_backend::byte_kernels::sigmoid_f64,
527 "sigmoid_elementwise"
528);
529cpu_unary_wrapper!(
530 silu_elementwise_f64_cpu_wrapper,
531 fuel_cpu_backend::byte_kernels::silu_f64,
532 "silu_elementwise"
533);
534cpu_unary_wrapper!(
535 gelu_elementwise_f64_cpu_wrapper,
536 fuel_cpu_backend::byte_kernels::gelu_f64,
537 "gelu_elementwise"
538);
539cpu_unary_wrapper!(
540 step_elementwise_f64_cpu_wrapper,
541 fuel_cpu_backend::byte_kernels::step_f64,
542 "step_elementwise"
543);
544
545cpu_binary_wrapper!(
546 maximum_elementwise_f32_cpu_wrapper,
547 fuel_cpu_backend::byte_kernels::maximum_f32,
548 "maximum_elementwise"
549);
550cpu_binary_wrapper!(
551 minimum_elementwise_f32_cpu_wrapper,
552 fuel_cpu_backend::byte_kernels::minimum_f32,
553 "minimum_elementwise"
554);
555cpu_binary_wrapper!(
556 maximum_elementwise_f64_cpu_wrapper,
557 fuel_cpu_backend::byte_kernels::maximum_f64,
558 "maximum_elementwise"
559);
560cpu_binary_wrapper!(
561 minimum_elementwise_f64_cpu_wrapper,
562 fuel_cpu_backend::byte_kernels::minimum_f64,
563 "minimum_elementwise"
564);
565
566cpu_binary_wrapper!(
568 add_elementwise_bf16_cpu_wrapper,
569 fuel_cpu_backend::byte_kernels::add_bf16,
570 "add_elementwise"
571);
572cpu_binary_wrapper!(
573 sub_elementwise_bf16_cpu_wrapper,
574 fuel_cpu_backend::byte_kernels::sub_bf16,
575 "sub_elementwise"
576);
577cpu_binary_wrapper!(
578 mul_elementwise_bf16_cpu_wrapper,
579 fuel_cpu_backend::byte_kernels::mul_bf16,
580 "mul_elementwise"
581);
582cpu_binary_wrapper!(
583 div_elementwise_bf16_cpu_wrapper,
584 fuel_cpu_backend::byte_kernels::div_bf16,
585 "div_elementwise"
586);
587cpu_binary_wrapper!(
588 maximum_elementwise_bf16_cpu_wrapper,
589 fuel_cpu_backend::byte_kernels::maximum_bf16,
590 "maximum_elementwise"
591);
592cpu_binary_wrapper!(
593 minimum_elementwise_bf16_cpu_wrapper,
594 fuel_cpu_backend::byte_kernels::minimum_bf16,
595 "minimum_elementwise"
596);
597
598cpu_unary_wrapper!(
599 relu_elementwise_bf16_cpu_wrapper,
600 fuel_cpu_backend::byte_kernels::relu_bf16,
601 "relu_elementwise"
602);
603cpu_unary_wrapper!(
604 neg_elementwise_bf16_cpu_wrapper,
605 fuel_cpu_backend::byte_kernels::neg_bf16,
606 "neg_elementwise"
607);
608cpu_unary_wrapper!(
609 sqr_elementwise_bf16_cpu_wrapper,
610 fuel_cpu_backend::byte_kernels::sqr_bf16,
611 "sqr_elementwise"
612);
613cpu_unary_wrapper!(
614 sqrt_elementwise_bf16_cpu_wrapper,
615 fuel_cpu_backend::byte_kernels::sqrt_bf16,
616 "sqrt_elementwise"
617);
618cpu_unary_wrapper!(
619 recip_elementwise_bf16_cpu_wrapper,
620 fuel_cpu_backend::byte_kernels::recip_bf16,
621 "recip_elementwise"
622);
623cpu_unary_wrapper!(
624 abs_elementwise_bf16_cpu_wrapper,
625 fuel_cpu_backend::byte_kernels::abs_bf16,
626 "abs_elementwise"
627);
628cpu_unary_wrapper!(
629 tanh_elementwise_bf16_cpu_wrapper,
630 fuel_cpu_backend::byte_kernels::tanh_bf16,
631 "tanh_elementwise"
632);
633cpu_unary_wrapper!(
634 exp_elementwise_bf16_cpu_wrapper,
635 fuel_cpu_backend::byte_kernels::exp_bf16,
636 "exp_elementwise"
637);
638cpu_unary_wrapper!(
639 log_elementwise_bf16_cpu_wrapper,
640 fuel_cpu_backend::byte_kernels::log_bf16,
641 "log_elementwise"
642);
643cpu_unary_wrapper!(
644 sin_elementwise_bf16_cpu_wrapper,
645 fuel_cpu_backend::byte_kernels::sin_bf16,
646 "sin_elementwise"
647);
648cpu_unary_wrapper!(
649 cos_elementwise_bf16_cpu_wrapper,
650 fuel_cpu_backend::byte_kernels::cos_bf16,
651 "cos_elementwise"
652);
653cpu_unary_wrapper!(
654 sigmoid_elementwise_bf16_cpu_wrapper,
655 fuel_cpu_backend::byte_kernels::sigmoid_bf16,
656 "sigmoid_elementwise"
657);
658cpu_unary_wrapper!(
659 silu_elementwise_bf16_cpu_wrapper,
660 fuel_cpu_backend::byte_kernels::silu_bf16,
661 "silu_elementwise"
662);
663cpu_unary_wrapper!(
664 gelu_elementwise_bf16_cpu_wrapper,
665 fuel_cpu_backend::byte_kernels::gelu_bf16,
666 "gelu_elementwise"
667);
668cpu_unary_wrapper!(
669 step_elementwise_bf16_cpu_wrapper,
670 fuel_cpu_backend::byte_kernels::step_bf16,
671 "step_elementwise"
672);
673
674cpu_binary_wrapper!(
676 add_elementwise_f16_cpu_wrapper,
677 fuel_cpu_backend::byte_kernels::add_f16,
678 "add_elementwise"
679);
680cpu_binary_wrapper!(
681 sub_elementwise_f16_cpu_wrapper,
682 fuel_cpu_backend::byte_kernels::sub_f16,
683 "sub_elementwise"
684);
685cpu_binary_wrapper!(
686 mul_elementwise_f16_cpu_wrapper,
687 fuel_cpu_backend::byte_kernels::mul_f16,
688 "mul_elementwise"
689);
690cpu_binary_wrapper!(
691 div_elementwise_f16_cpu_wrapper,
692 fuel_cpu_backend::byte_kernels::div_f16,
693 "div_elementwise"
694);
695cpu_binary_wrapper!(
696 maximum_elementwise_f16_cpu_wrapper,
697 fuel_cpu_backend::byte_kernels::maximum_f16,
698 "maximum_elementwise"
699);
700cpu_binary_wrapper!(
701 minimum_elementwise_f16_cpu_wrapper,
702 fuel_cpu_backend::byte_kernels::minimum_f16,
703 "minimum_elementwise"
704);
705
706cpu_unary_wrapper!(
707 relu_elementwise_f16_cpu_wrapper,
708 fuel_cpu_backend::byte_kernels::relu_f16,
709 "relu_elementwise"
710);
711cpu_unary_wrapper!(
712 neg_elementwise_f16_cpu_wrapper,
713 fuel_cpu_backend::byte_kernels::neg_f16,
714 "neg_elementwise"
715);
716cpu_unary_wrapper!(
717 sqr_elementwise_f16_cpu_wrapper,
718 fuel_cpu_backend::byte_kernels::sqr_f16,
719 "sqr_elementwise"
720);
721cpu_unary_wrapper!(
722 sqrt_elementwise_f16_cpu_wrapper,
723 fuel_cpu_backend::byte_kernels::sqrt_f16,
724 "sqrt_elementwise"
725);
726cpu_unary_wrapper!(
727 recip_elementwise_f16_cpu_wrapper,
728 fuel_cpu_backend::byte_kernels::recip_f16,
729 "recip_elementwise"
730);
731cpu_unary_wrapper!(
732 abs_elementwise_f16_cpu_wrapper,
733 fuel_cpu_backend::byte_kernels::abs_f16,
734 "abs_elementwise"
735);
736cpu_unary_wrapper!(
737 tanh_elementwise_f16_cpu_wrapper,
738 fuel_cpu_backend::byte_kernels::tanh_f16,
739 "tanh_elementwise"
740);
741cpu_unary_wrapper!(
742 exp_elementwise_f16_cpu_wrapper,
743 fuel_cpu_backend::byte_kernels::exp_f16,
744 "exp_elementwise"
745);
746cpu_unary_wrapper!(
747 log_elementwise_f16_cpu_wrapper,
748 fuel_cpu_backend::byte_kernels::log_f16,
749 "log_elementwise"
750);
751cpu_unary_wrapper!(
752 sin_elementwise_f16_cpu_wrapper,
753 fuel_cpu_backend::byte_kernels::sin_f16,
754 "sin_elementwise"
755);
756cpu_unary_wrapper!(
757 cos_elementwise_f16_cpu_wrapper,
758 fuel_cpu_backend::byte_kernels::cos_f16,
759 "cos_elementwise"
760);
761cpu_unary_wrapper!(
762 sigmoid_elementwise_f16_cpu_wrapper,
763 fuel_cpu_backend::byte_kernels::sigmoid_f16,
764 "sigmoid_elementwise"
765);
766cpu_unary_wrapper!(
767 silu_elementwise_f16_cpu_wrapper,
768 fuel_cpu_backend::byte_kernels::silu_f16,
769 "silu_elementwise"
770);
771cpu_unary_wrapper!(
772 gelu_elementwise_f16_cpu_wrapper,
773 fuel_cpu_backend::byte_kernels::gelu_f16,
774 "gelu_elementwise"
775);
776cpu_unary_wrapper!(
777 step_elementwise_f16_cpu_wrapper,
778 fuel_cpu_backend::byte_kernels::step_f16,
779 "step_elementwise"
780);
781
782cpu_binary_wrapper!(
788 eq_elementwise_f32_cpu_wrapper,
789 fuel_cpu_backend::byte_kernels::eq_f32_u8,
790 "eq_elementwise"
791);
792cpu_binary_wrapper!(
793 eq_elementwise_f64_cpu_wrapper,
794 fuel_cpu_backend::byte_kernels::eq_f64_u8,
795 "eq_elementwise"
796);
797cpu_binary_wrapper!(
798 eq_elementwise_bf16_cpu_wrapper,
799 fuel_cpu_backend::byte_kernels::eq_bf16_u8,
800 "eq_elementwise"
801);
802cpu_binary_wrapper!(
803 eq_elementwise_f16_cpu_wrapper,
804 fuel_cpu_backend::byte_kernels::eq_f16_u8,
805 "eq_elementwise"
806);
807
808cpu_binary_wrapper!(
809 ne_elementwise_f32_cpu_wrapper,
810 fuel_cpu_backend::byte_kernels::ne_f32_u8,
811 "ne_elementwise"
812);
813cpu_binary_wrapper!(
814 ne_elementwise_f64_cpu_wrapper,
815 fuel_cpu_backend::byte_kernels::ne_f64_u8,
816 "ne_elementwise"
817);
818cpu_binary_wrapper!(
819 ne_elementwise_bf16_cpu_wrapper,
820 fuel_cpu_backend::byte_kernels::ne_bf16_u8,
821 "ne_elementwise"
822);
823cpu_binary_wrapper!(
824 ne_elementwise_f16_cpu_wrapper,
825 fuel_cpu_backend::byte_kernels::ne_f16_u8,
826 "ne_elementwise"
827);
828
829cpu_binary_wrapper!(
830 lt_elementwise_f32_cpu_wrapper,
831 fuel_cpu_backend::byte_kernels::lt_f32_u8,
832 "lt_elementwise"
833);
834cpu_binary_wrapper!(
835 lt_elementwise_f64_cpu_wrapper,
836 fuel_cpu_backend::byte_kernels::lt_f64_u8,
837 "lt_elementwise"
838);
839cpu_binary_wrapper!(
840 lt_elementwise_bf16_cpu_wrapper,
841 fuel_cpu_backend::byte_kernels::lt_bf16_u8,
842 "lt_elementwise"
843);
844cpu_binary_wrapper!(
845 lt_elementwise_f16_cpu_wrapper,
846 fuel_cpu_backend::byte_kernels::lt_f16_u8,
847 "lt_elementwise"
848);
849
850cpu_binary_wrapper!(
851 le_elementwise_f32_cpu_wrapper,
852 fuel_cpu_backend::byte_kernels::le_f32_u8,
853 "le_elementwise"
854);
855cpu_binary_wrapper!(
856 le_elementwise_f64_cpu_wrapper,
857 fuel_cpu_backend::byte_kernels::le_f64_u8,
858 "le_elementwise"
859);
860cpu_binary_wrapper!(
861 le_elementwise_bf16_cpu_wrapper,
862 fuel_cpu_backend::byte_kernels::le_bf16_u8,
863 "le_elementwise"
864);
865cpu_binary_wrapper!(
866 le_elementwise_f16_cpu_wrapper,
867 fuel_cpu_backend::byte_kernels::le_f16_u8,
868 "le_elementwise"
869);
870
871cpu_binary_wrapper!(
872 gt_elementwise_f32_cpu_wrapper,
873 fuel_cpu_backend::byte_kernels::gt_f32_u8,
874 "gt_elementwise"
875);
876cpu_binary_wrapper!(
877 gt_elementwise_f64_cpu_wrapper,
878 fuel_cpu_backend::byte_kernels::gt_f64_u8,
879 "gt_elementwise"
880);
881cpu_binary_wrapper!(
882 gt_elementwise_bf16_cpu_wrapper,
883 fuel_cpu_backend::byte_kernels::gt_bf16_u8,
884 "gt_elementwise"
885);
886cpu_binary_wrapper!(
887 gt_elementwise_f16_cpu_wrapper,
888 fuel_cpu_backend::byte_kernels::gt_f16_u8,
889 "gt_elementwise"
890);
891
892cpu_binary_wrapper!(
893 ge_elementwise_f32_cpu_wrapper,
894 fuel_cpu_backend::byte_kernels::ge_f32_u8,
895 "ge_elementwise"
896);
897cpu_binary_wrapper!(
898 ge_elementwise_f64_cpu_wrapper,
899 fuel_cpu_backend::byte_kernels::ge_f64_u8,
900 "ge_elementwise"
901);
902cpu_binary_wrapper!(
903 ge_elementwise_bf16_cpu_wrapper,
904 fuel_cpu_backend::byte_kernels::ge_bf16_u8,
905 "ge_elementwise"
906);
907cpu_binary_wrapper!(
908 ge_elementwise_f16_cpu_wrapper,
909 fuel_cpu_backend::byte_kernels::ge_f16_u8,
910 "ge_elementwise"
911);
912
913macro_rules! cpu_where_wrapper {
919 ($wrapper:ident, $kernel:path, $op_name:literal) => {
920 pub(crate) fn $wrapper(
923 inputs: &[Arc<RwLock<Storage>>],
924 outputs: &mut [Arc<RwLock<Storage>>],
925 _layouts: &[Layout],
926 _params: &OpParams,
927 ) -> Result<()> {
928 if inputs.len() != 3 {
929 return Err(Error::Msg(format!(
930 "{} wrapper expects 3 inputs (cond, a, b), got {}",
931 $op_name,
932 inputs.len(),
933 ))
934 .bt());
935 }
936 if outputs.len() != 1 {
937 return Err(Error::Msg(format!(
938 "{} wrapper expects 1 output, got {}",
939 $op_name,
940 outputs.len(),
941 ))
942 .bt());
943 }
944 let cond_guard = read_storage(&inputs[0])?;
945 if cond_guard.dtype != DType::Bool {
951 return Err(Error::Msg(format!(
952 "{}: cond must be Bool, got {:?} (GAP-168(c))",
953 $op_name, cond_guard.dtype,
954 ))
955 .bt());
956 }
957 let a_guard = read_storage(&inputs[1])?;
958 let b_guard = read_storage(&inputs[2])?;
959 let mut out_guard = write_storage(&outputs[0])?;
960 let cond_cpu = cpu_input(&cond_guard)?;
961 let a_cpu = cpu_input(&a_guard)?;
962 let b_cpu = cpu_input(&b_guard)?;
963 let out_cpu = cpu_output(&mut out_guard)?;
964 $kernel(cond_cpu, a_cpu, b_cpu, out_cpu)
965 }
966 };
967}
968
969cpu_where_wrapper!(
970 where_f32_cpu_wrapper,
971 fuel_cpu_backend::byte_kernels::where_f32,
972 "where"
973);
974cpu_where_wrapper!(
975 where_f64_cpu_wrapper,
976 fuel_cpu_backend::byte_kernels::where_f64,
977 "where"
978);
979cpu_where_wrapper!(
980 where_bf16_cpu_wrapper,
981 fuel_cpu_backend::byte_kernels::where_bf16,
982 "where"
983);
984cpu_where_wrapper!(
985 where_f16_cpu_wrapper,
986 fuel_cpu_backend::byte_kernels::where_f16,
987 "where"
988);
989
990cpu_unary_wrapper!(
993 floor_elementwise_f32_cpu_wrapper,
994 fuel_cpu_backend::byte_kernels::floor_f32,
995 "floor_elementwise"
996);
997cpu_unary_wrapper!(
998 floor_elementwise_f64_cpu_wrapper,
999 fuel_cpu_backend::byte_kernels::floor_f64,
1000 "floor_elementwise"
1001);
1002cpu_unary_wrapper!(
1003 floor_elementwise_bf16_cpu_wrapper,
1004 fuel_cpu_backend::byte_kernels::floor_bf16,
1005 "floor_elementwise"
1006);
1007cpu_unary_wrapper!(
1008 floor_elementwise_f16_cpu_wrapper,
1009 fuel_cpu_backend::byte_kernels::floor_f16,
1010 "floor_elementwise"
1011);
1012
1013cpu_unary_wrapper!(
1014 ceil_elementwise_f32_cpu_wrapper,
1015 fuel_cpu_backend::byte_kernels::ceil_f32,
1016 "ceil_elementwise"
1017);
1018cpu_unary_wrapper!(
1019 ceil_elementwise_f64_cpu_wrapper,
1020 fuel_cpu_backend::byte_kernels::ceil_f64,
1021 "ceil_elementwise"
1022);
1023cpu_unary_wrapper!(
1024 ceil_elementwise_bf16_cpu_wrapper,
1025 fuel_cpu_backend::byte_kernels::ceil_bf16,
1026 "ceil_elementwise"
1027);
1028cpu_unary_wrapper!(
1029 ceil_elementwise_f16_cpu_wrapper,
1030 fuel_cpu_backend::byte_kernels::ceil_f16,
1031 "ceil_elementwise"
1032);
1033
1034cpu_unary_wrapper!(
1035 round_elementwise_f32_cpu_wrapper,
1036 fuel_cpu_backend::byte_kernels::round_f32,
1037 "round_elementwise"
1038);
1039cpu_unary_wrapper!(
1040 round_elementwise_f64_cpu_wrapper,
1041 fuel_cpu_backend::byte_kernels::round_f64,
1042 "round_elementwise"
1043);
1044cpu_unary_wrapper!(
1045 round_elementwise_bf16_cpu_wrapper,
1046 fuel_cpu_backend::byte_kernels::round_bf16,
1047 "round_elementwise"
1048);
1049cpu_unary_wrapper!(
1050 round_elementwise_f16_cpu_wrapper,
1051 fuel_cpu_backend::byte_kernels::round_f16,
1052 "round_elementwise"
1053);
1054
1055cpu_unary_wrapper!(
1056 sign_elementwise_f32_cpu_wrapper,
1057 fuel_cpu_backend::byte_kernels::sign_f32,
1058 "sign_elementwise"
1059);
1060cpu_unary_wrapper!(
1061 sign_elementwise_f64_cpu_wrapper,
1062 fuel_cpu_backend::byte_kernels::sign_f64,
1063 "sign_elementwise"
1064);
1065cpu_unary_wrapper!(
1066 sign_elementwise_bf16_cpu_wrapper,
1067 fuel_cpu_backend::byte_kernels::sign_bf16,
1068 "sign_elementwise"
1069);
1070cpu_unary_wrapper!(
1071 sign_elementwise_f16_cpu_wrapper,
1072 fuel_cpu_backend::byte_kernels::sign_f16,
1073 "sign_elementwise"
1074);
1075
1076cpu_unary_wrapper!(
1077 erf_elementwise_f32_cpu_wrapper,
1078 fuel_cpu_backend::byte_kernels::erf_f32,
1079 "erf_elementwise"
1080);
1081cpu_unary_wrapper!(
1082 erf_elementwise_f64_cpu_wrapper,
1083 fuel_cpu_backend::byte_kernels::erf_f64,
1084 "erf_elementwise"
1085);
1086cpu_unary_wrapper!(
1087 erf_elementwise_bf16_cpu_wrapper,
1088 fuel_cpu_backend::byte_kernels::erf_bf16,
1089 "erf_elementwise"
1090);
1091cpu_unary_wrapper!(
1092 erf_elementwise_f16_cpu_wrapper,
1093 fuel_cpu_backend::byte_kernels::erf_f16,
1094 "erf_elementwise"
1095);
1096
1097cpu_unary_wrapper!(
1098 gelu_erf_elementwise_f32_cpu_wrapper,
1099 fuel_cpu_backend::byte_kernels::gelu_erf_f32,
1100 "gelu_erf_elementwise"
1101);
1102cpu_unary_wrapper!(
1103 gelu_erf_elementwise_f64_cpu_wrapper,
1104 fuel_cpu_backend::byte_kernels::gelu_erf_f64,
1105 "gelu_erf_elementwise"
1106);
1107cpu_unary_wrapper!(
1108 gelu_erf_elementwise_bf16_cpu_wrapper,
1109 fuel_cpu_backend::byte_kernels::gelu_erf_bf16,
1110 "gelu_erf_elementwise"
1111);
1112cpu_unary_wrapper!(
1113 gelu_erf_elementwise_f16_cpu_wrapper,
1114 fuel_cpu_backend::byte_kernels::gelu_erf_f16,
1115 "gelu_erf_elementwise"
1116);
1117
1118cpu_binary_wrapper!(
1119 pow_elementwise_f32_cpu_wrapper,
1120 fuel_cpu_backend::byte_kernels::pow_f32,
1121 "pow_elementwise"
1122);
1123cpu_binary_wrapper!(
1124 pow_elementwise_f64_cpu_wrapper,
1125 fuel_cpu_backend::byte_kernels::pow_f64,
1126 "pow_elementwise"
1127);
1128cpu_binary_wrapper!(
1129 pow_elementwise_bf16_cpu_wrapper,
1130 fuel_cpu_backend::byte_kernels::pow_bf16,
1131 "pow_elementwise"
1132);
1133cpu_binary_wrapper!(
1134 pow_elementwise_f16_cpu_wrapper,
1135 fuel_cpu_backend::byte_kernels::pow_f16,
1136 "pow_elementwise"
1137);
1138
1139cpu_unary_wrapper!(
1140 rsqrt_elementwise_f32_cpu_wrapper,
1141 fuel_cpu_backend::byte_kernels::rsqrt_f32,
1142 "rsqrt_elementwise"
1143);
1144cpu_unary_wrapper!(
1145 rsqrt_elementwise_f64_cpu_wrapper,
1146 fuel_cpu_backend::byte_kernels::rsqrt_f64,
1147 "rsqrt_elementwise"
1148);
1149cpu_unary_wrapper!(
1150 rsqrt_elementwise_bf16_cpu_wrapper,
1151 fuel_cpu_backend::byte_kernels::rsqrt_bf16,
1152 "rsqrt_elementwise"
1153);
1154cpu_unary_wrapper!(
1155 rsqrt_elementwise_f16_cpu_wrapper,
1156 fuel_cpu_backend::byte_kernels::rsqrt_f16,
1157 "rsqrt_elementwise"
1158);
1159
1160cpu_binary_wrapper!(
1161 rem_elementwise_f32_cpu_wrapper,
1162 fuel_cpu_backend::byte_kernels::rem_f32,
1163 "rem_elementwise"
1164);
1165cpu_binary_wrapper!(
1166 rem_elementwise_f64_cpu_wrapper,
1167 fuel_cpu_backend::byte_kernels::rem_f64,
1168 "rem_elementwise"
1169);
1170cpu_binary_wrapper!(
1171 rem_elementwise_bf16_cpu_wrapper,
1172 fuel_cpu_backend::byte_kernels::rem_bf16,
1173 "rem_elementwise"
1174);
1175cpu_binary_wrapper!(
1176 rem_elementwise_f16_cpu_wrapper,
1177 fuel_cpu_backend::byte_kernels::rem_f16,
1178 "rem_elementwise"
1179);
1180
1181pub(crate) fn flip_cpu_wrapper(
1186 inputs: &[Arc<RwLock<Storage>>],
1187 outputs: &mut [Arc<RwLock<Storage>>],
1188 _layouts: &[Layout],
1189 params: &OpParams,
1190) -> Result<()> {
1191 if inputs.len() != 1 || outputs.len() != 1 {
1192 return Err(Error::Msg(format!(
1193 "flip wrapper expects 1 input + 1 output, got {} + {}",
1194 inputs.len(),
1195 outputs.len(),
1196 ))
1197 .bt());
1198 }
1199 let (outer, dim_size, inner) = match params {
1200 OpParams::Flip {
1201 outer_count,
1202 dim_size,
1203 inner_count,
1204 ..
1205 } => (*outer_count, *dim_size, *inner_count),
1206 other => {
1207 return Err(Error::Msg(format!(
1208 "flip wrapper expects OpParams::Flip, got {other:?}",
1209 ))
1210 .bt());
1211 }
1212 };
1213 let in_guard = read_storage(&inputs[0])?;
1214 let mut out_guard = write_storage(&outputs[0])?;
1215 let dtype_size = out_guard.dtype.size_in_bytes();
1216 let in_cpu = cpu_input(&in_guard)?;
1217 let out_cpu = cpu_output(&mut out_guard)?;
1218 fuel_cpu_backend::byte_kernels::flip_cpu(in_cpu, out_cpu, outer, dim_size, inner, dtype_size)
1219}
1220
1221macro_rules! cpu_cumsum_wrapper {
1225 ($wrapper:ident, $kernel:path, $op_name:literal) => {
1226 pub(crate) fn $wrapper(
1227 inputs: &[Arc<RwLock<Storage>>],
1228 outputs: &mut [Arc<RwLock<Storage>>],
1229 _layouts: &[Layout],
1230 params: &OpParams,
1231 ) -> Result<()> {
1232 if inputs.len() != 1 || outputs.len() != 1 {
1233 return Err(Error::Msg(format!(
1234 "{} wrapper expects 1 input + 1 output, got {} + {}",
1235 $op_name,
1236 inputs.len(),
1237 outputs.len(),
1238 ))
1239 .bt());
1240 }
1241 let (outer, dim_size, inner) = match params {
1242 OpParams::CumSum {
1243 outer_count,
1244 dim_size,
1245 inner_count,
1246 ..
1247 } => (*outer_count, *dim_size, *inner_count),
1248 other => {
1249 return Err(Error::Msg(format!(
1250 "{} wrapper expects OpParams::CumSum, got {other:?}",
1251 $op_name,
1252 ))
1253 .bt())
1254 }
1255 };
1256 let in_guard = read_storage(&inputs[0])?;
1257 let mut out_guard = write_storage(&outputs[0])?;
1258 let in_cpu = cpu_input(&in_guard)?;
1259 let out_cpu = cpu_output(&mut out_guard)?;
1260 $kernel(in_cpu, out_cpu, outer, dim_size, inner)
1261 }
1262 };
1263}
1264
1265cpu_cumsum_wrapper!(
1266 cumsum_f32_cpu_wrapper,
1267 fuel_cpu_backend::byte_kernels::cumsum_f32,
1268 "cumsum_f32"
1269);
1270cpu_cumsum_wrapper!(
1271 cumsum_f64_cpu_wrapper,
1272 fuel_cpu_backend::byte_kernels::cumsum_f64,
1273 "cumsum_f64"
1274);
1275cpu_cumsum_wrapper!(
1276 cumsum_bf16_cpu_wrapper,
1277 fuel_cpu_backend::byte_kernels::cumsum_bf16,
1278 "cumsum_bf16"
1279);
1280cpu_cumsum_wrapper!(
1281 cumsum_f16_cpu_wrapper,
1282 fuel_cpu_backend::byte_kernels::cumsum_f16,
1283 "cumsum_f16"
1284);
1285
1286fn triu_cpu_wrapper(
1289 inputs: &[Arc<RwLock<Storage>>],
1290 outputs: &mut [Arc<RwLock<Storage>>],
1291 _layouts: &[Layout],
1292 params: &OpParams,
1293) -> Result<()> {
1294 triangular_wrapper_inner(
1295 inputs,
1296 outputs,
1297 params,
1298 true,
1299 "triu_cpu_wrapper",
1300 )
1301}
1302
1303fn tril_cpu_wrapper(
1304 inputs: &[Arc<RwLock<Storage>>],
1305 outputs: &mut [Arc<RwLock<Storage>>],
1306 _layouts: &[Layout],
1307 params: &OpParams,
1308) -> Result<()> {
1309 triangular_wrapper_inner(
1310 inputs,
1311 outputs,
1312 params,
1313 false,
1314 "tril_cpu_wrapper",
1315 )
1316}
1317
1318fn triangular_wrapper_inner(
1319 inputs: &[Arc<RwLock<Storage>>],
1320 outputs: &mut [Arc<RwLock<Storage>>],
1321 params: &OpParams,
1322 keep_upper: bool,
1323 op_name: &str,
1324) -> Result<()> {
1325 if inputs.len() != 1 || outputs.len() != 1 {
1326 return Err(Error::Msg(format!(
1327 "{op_name}: expects 1 input + 1 output, got {} + {}",
1328 inputs.len(),
1329 outputs.len(),
1330 ))
1331 .bt());
1332 }
1333 let (batch, rows, cols, diag) = match params {
1334 OpParams::Triangular {
1335 batch_count,
1336 rows,
1337 cols,
1338 diagonal,
1339 } => (*batch_count, *rows, *cols, *diagonal),
1340 other => {
1341 return Err(Error::Msg(format!(
1342 "{op_name}: expects OpParams::Triangular, got {other:?}",
1343 ))
1344 .bt());
1345 }
1346 };
1347 let in_guard = read_storage(&inputs[0])?;
1348 let mut out_guard = write_storage(&outputs[0])?;
1349 let dtype_size = out_guard.dtype.size_in_bytes();
1350 let in_cpu = cpu_input(&in_guard)?;
1351 let out_cpu = cpu_output(&mut out_guard)?;
1352 fuel_cpu_backend::byte_kernels::triangular_cpu(
1353 in_cpu, out_cpu, batch, rows, cols, diag, keep_upper, dtype_size,
1354 )
1355}
1356
1357macro_rules! cpu_log_softmax_wrapper {
1360 ($wrapper:ident, $kernel:path, $op_name:literal) => {
1361 pub(crate) fn $wrapper(
1362 inputs: &[Arc<RwLock<Storage>>],
1363 outputs: &mut [Arc<RwLock<Storage>>],
1364 _layouts: &[Layout],
1365 params: &OpParams,
1366 ) -> Result<()> {
1367 if inputs.len() != 1 || outputs.len() != 1 {
1368 return Err(Error::Msg(format!(
1369 "{} wrapper expects 1 input + 1 output, got {} + {}",
1370 $op_name,
1371 inputs.len(),
1372 outputs.len(),
1373 ))
1374 .bt());
1375 }
1376 let (outer, last_dim) = match params {
1377 OpParams::LogSoftmaxLastDim {
1378 outer_count,
1379 last_dim,
1380 } => (*outer_count, *last_dim),
1381 other => {
1382 return Err(Error::Msg(format!(
1383 "{} wrapper expects OpParams::LogSoftmaxLastDim, got {other:?}",
1384 $op_name,
1385 ))
1386 .bt());
1387 }
1388 };
1389 let in_guard = read_storage(&inputs[0])?;
1390 let mut out_guard = write_storage(&outputs[0])?;
1391 let in_cpu = cpu_input(&in_guard)?;
1392 let out_cpu = cpu_output(&mut out_guard)?;
1393 $kernel(in_cpu, out_cpu, outer, last_dim)
1394 }
1395 };
1396}
1397
1398cpu_log_softmax_wrapper!(
1399 log_softmax_f32_cpu_wrapper,
1400 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_f32,
1401 "log_softmax_f32"
1402);
1403cpu_log_softmax_wrapper!(
1404 log_softmax_f64_cpu_wrapper,
1405 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_f64,
1406 "log_softmax_f64"
1407);
1408cpu_log_softmax_wrapper!(
1409 log_softmax_bf16_cpu_wrapper,
1410 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_bf16,
1411 "log_softmax_bf16"
1412);
1413cpu_log_softmax_wrapper!(
1414 log_softmax_f16_cpu_wrapper,
1415 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_f16,
1416 "log_softmax_f16"
1417);
1418
1419macro_rules! cpu_log_softmax_backward_wrapper {
1422 ($wrapper:ident, $kernel:path, $op_name:literal) => {
1423 pub(crate) fn $wrapper(
1426 inputs: &[Arc<RwLock<Storage>>],
1427 outputs: &mut [Arc<RwLock<Storage>>],
1428 _layouts: &[Layout],
1429 params: &OpParams,
1430 ) -> Result<()> {
1431 if inputs.len() != 2 || outputs.len() != 1 {
1432 return Err(Error::Msg(format!(
1433 "{} wrapper expects 2 inputs (y, g) + 1 output, got {} + {}",
1434 $op_name,
1435 inputs.len(),
1436 outputs.len(),
1437 ))
1438 .bt());
1439 }
1440 let (outer, last_dim) = match params {
1441 OpParams::LogSoftmaxLastDim {
1442 outer_count,
1443 last_dim,
1444 } => (*outer_count, *last_dim),
1445 other => {
1446 return Err(Error::Msg(format!(
1447 "{} wrapper expects OpParams::LogSoftmaxLastDim, got {other:?}",
1448 $op_name,
1449 ))
1450 .bt());
1451 }
1452 };
1453 let y_guard = read_storage(&inputs[0])?;
1454 let g_guard = read_storage(&inputs[1])?;
1455 let mut out_guard = write_storage(&outputs[0])?;
1456 let y_cpu = cpu_input(&y_guard)?;
1457 let g_cpu = cpu_input(&g_guard)?;
1458 let out_cpu = cpu_output(&mut out_guard)?;
1459 $kernel(y_cpu, g_cpu, out_cpu, outer, last_dim)
1460 }
1461 };
1462}
1463
1464cpu_log_softmax_backward_wrapper!(
1465 log_softmax_backward_f32_cpu_wrapper,
1466 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_backward_f32,
1467 "log_softmax_backward_f32"
1468);
1469cpu_log_softmax_backward_wrapper!(
1470 log_softmax_backward_f64_cpu_wrapper,
1471 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_backward_f64,
1472 "log_softmax_backward_f64"
1473);
1474cpu_log_softmax_backward_wrapper!(
1475 log_softmax_backward_bf16_cpu_wrapper,
1476 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_backward_bf16,
1477 "log_softmax_backward_bf16"
1478);
1479cpu_log_softmax_backward_wrapper!(
1480 log_softmax_backward_f16_cpu_wrapper,
1481 fuel_cpu_backend::byte_kernels::log_softmax_last_dim_backward_f16,
1482 "log_softmax_backward_f16"
1483);
1484
1485macro_rules! cpu_softmax_last_dim_backward_wrapper {
1489 ($wrapper:ident, $kernel:path) => {
1490 pub(crate) fn $wrapper(
1493 inputs: &[Arc<RwLock<Storage>>],
1494 outputs: &mut [Arc<RwLock<Storage>>],
1495 _layouts: &[Layout],
1496 params: &OpParams,
1497 ) -> Result<()> {
1498 if inputs.len() != 2 || outputs.len() != 1 {
1499 return Err(Error::Msg(format!(
1500 "softmax_last_dim_backward wrapper expects 2 inputs (y, g) + 1 output, got {} + {}",
1501 inputs.len(), outputs.len(),
1502 )).bt());
1503 }
1504 let (outer, last_dim) = match params {
1505 OpParams::SoftmaxLastDim { outer_count, last_dim } => (*outer_count, *last_dim),
1506 other => {
1507 return Err(Error::Msg(format!(
1508 "softmax_last_dim_backward wrapper expects OpParams::SoftmaxLastDim, got {other:?}",
1509 )).bt());
1510 }
1511 };
1512 let y_guard = read_storage(&inputs[0])?;
1513 let g_guard = read_storage(&inputs[1])?;
1514 let mut out_guard = write_storage(&outputs[0])?;
1515 let y_cpu = cpu_input(&y_guard)?;
1516 let g_cpu = cpu_input(&g_guard)?;
1517 let out_cpu = cpu_output(&mut out_guard)?;
1518 $kernel(y_cpu, g_cpu, out_cpu, outer, last_dim)
1519 }
1520 };
1521}
1522
1523cpu_softmax_last_dim_backward_wrapper!(
1524 softmax_last_dim_backward_f32_cpu_wrapper,
1525 fuel_cpu_backend::byte_kernels::softmax_last_dim_backward_f32
1526);
1527cpu_softmax_last_dim_backward_wrapper!(
1528 softmax_last_dim_backward_f64_cpu_wrapper,
1529 fuel_cpu_backend::byte_kernels::softmax_last_dim_backward_f64
1530);
1531cpu_softmax_last_dim_backward_wrapper!(
1532 softmax_last_dim_backward_bf16_cpu_wrapper,
1533 fuel_cpu_backend::byte_kernels::softmax_last_dim_backward_bf16
1534);
1535cpu_softmax_last_dim_backward_wrapper!(
1536 softmax_last_dim_backward_f16_cpu_wrapper,
1537 fuel_cpu_backend::byte_kernels::softmax_last_dim_backward_f16
1538);
1539
1540macro_rules! cpu_norm_last_dim_backward_wrapper {
1543 ($wrapper:ident, $kernel:path, $op_name:literal) => {
1544 pub(crate) fn $wrapper(
1547 inputs: &[Arc<RwLock<Storage>>],
1548 outputs: &mut [Arc<RwLock<Storage>>],
1549 _layouts: &[Layout],
1550 params: &OpParams,
1551 ) -> Result<()> {
1552 if inputs.len() != 2 || outputs.len() != 1 {
1553 return Err(Error::Msg(format!(
1554 "{} wrapper expects 2 inputs (x, g) + 1 output, got {} + {}",
1555 $op_name,
1556 inputs.len(),
1557 outputs.len(),
1558 ))
1559 .bt());
1560 }
1561 let (outer, last_dim, eps) = match params {
1562 OpParams::NormLastDim {
1563 outer_count,
1564 last_dim,
1565 eps,
1566 } => (*outer_count, *last_dim, *eps),
1567 other => {
1568 return Err(Error::Msg(format!(
1569 "{} wrapper expects OpParams::NormLastDim, got {other:?}",
1570 $op_name,
1571 ))
1572 .bt());
1573 }
1574 };
1575 let x_guard = read_storage(&inputs[0])?;
1576 let g_guard = read_storage(&inputs[1])?;
1577 let mut out_guard = write_storage(&outputs[0])?;
1578 let x_cpu = cpu_input(&x_guard)?;
1579 let g_cpu = cpu_input(&g_guard)?;
1580 let out_cpu = cpu_output(&mut out_guard)?;
1581 $kernel(x_cpu, g_cpu, out_cpu, outer, last_dim, eps)
1582 }
1583 };
1584}
1585
1586cpu_norm_last_dim_backward_wrapper!(
1587 layer_norm_last_dim_backward_f32_cpu_wrapper,
1588 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_backward_f32,
1589 "layer_norm_backward_f32"
1590);
1591cpu_norm_last_dim_backward_wrapper!(
1592 layer_norm_last_dim_backward_f64_cpu_wrapper,
1593 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_backward_f64,
1594 "layer_norm_backward_f64"
1595);
1596cpu_norm_last_dim_backward_wrapper!(
1597 layer_norm_last_dim_backward_bf16_cpu_wrapper,
1598 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_backward_bf16,
1599 "layer_norm_backward_bf16"
1600);
1601cpu_norm_last_dim_backward_wrapper!(
1602 layer_norm_last_dim_backward_f16_cpu_wrapper,
1603 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_backward_f16,
1604 "layer_norm_backward_f16"
1605);
1606
1607cpu_norm_last_dim_backward_wrapper!(
1608 rms_norm_last_dim_backward_f32_cpu_wrapper,
1609 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_backward_f32,
1610 "rms_norm_backward_f32"
1611);
1612cpu_norm_last_dim_backward_wrapper!(
1613 rms_norm_last_dim_backward_f64_cpu_wrapper,
1614 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_backward_f64,
1615 "rms_norm_backward_f64"
1616);
1617cpu_norm_last_dim_backward_wrapper!(
1618 rms_norm_last_dim_backward_bf16_cpu_wrapper,
1619 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_backward_bf16,
1620 "rms_norm_backward_bf16"
1621);
1622cpu_norm_last_dim_backward_wrapper!(
1623 rms_norm_last_dim_backward_f16_cpu_wrapper,
1624 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_backward_f16,
1625 "rms_norm_backward_f16"
1626);
1627
1628macro_rules! cpu_reduce_max_to_backward_wrapper {
1633 ($wrapper:ident, $kernel:path) => {
1634 pub(crate) fn $wrapper(
1635 inputs: &[Arc<RwLock<Storage>>],
1636 outputs: &mut [Arc<RwLock<Storage>>],
1637 _layouts: &[Layout],
1638 params: &OpParams,
1639 ) -> Result<()> {
1640 if inputs.len() != 2 || outputs.len() != 1 {
1641 return Err(Error::Msg(format!(
1642 "reduce_max_to_backward wrapper expects 2 inputs (x, upstream) + 1 output, got {} + {}",
1643 inputs.len(), outputs.len(),
1644 )).bt());
1645 }
1646 let (input_shape, output_shape) = match params {
1647 OpParams::ReduceMaxToBackward { input_shape, output_shape } => {
1648 (input_shape.as_slice(), output_shape.as_slice())
1649 }
1650 other => {
1651 return Err(Error::Msg(format!(
1652 "reduce_max_to_backward wrapper expects OpParams::ReduceMaxToBackward, got {other:?}",
1653 )).bt());
1654 }
1655 };
1656 let x_guard = read_storage(&inputs[0])?;
1657 let up_guard = read_storage(&inputs[1])?;
1658 let mut out_guard = write_storage(&outputs[0])?;
1659 let x_cpu = cpu_input(&x_guard)?;
1660 let up_cpu = cpu_input(&up_guard)?;
1661 let out_cpu = cpu_output(&mut out_guard)?;
1662 $kernel(x_cpu, up_cpu, out_cpu, input_shape, output_shape)
1663 }
1664 };
1665}
1666
1667cpu_reduce_max_to_backward_wrapper!(
1668 reduce_max_to_backward_f32_cpu_wrapper,
1669 fuel_cpu_backend::byte_kernels::reduce_max_to_backward_f32
1670);
1671cpu_reduce_max_to_backward_wrapper!(
1672 reduce_max_to_backward_f64_cpu_wrapper,
1673 fuel_cpu_backend::byte_kernels::reduce_max_to_backward_f64
1674);
1675cpu_reduce_max_to_backward_wrapper!(
1676 reduce_max_to_backward_bf16_cpu_wrapper,
1677 fuel_cpu_backend::byte_kernels::reduce_max_to_backward_bf16
1678);
1679cpu_reduce_max_to_backward_wrapper!(
1680 reduce_max_to_backward_f16_cpu_wrapper,
1681 fuel_cpu_backend::byte_kernels::reduce_max_to_backward_f16
1682);
1683
1684pub(crate) fn masked_fill_cpu_wrapper(
1687 inputs: &[Arc<RwLock<Storage>>],
1688 outputs: &mut [Arc<RwLock<Storage>>],
1689 _layouts: &[Layout],
1690 params: &OpParams,
1691) -> Result<()> {
1692 if inputs.len() != 2 || outputs.len() != 1 {
1693 return Err(Error::Msg(format!(
1694 "masked_fill_cpu_wrapper: expects 2 inputs (x, mask) + 1 output, got {} + {}",
1695 inputs.len(),
1696 outputs.len(),
1697 ))
1698 .bt());
1699 }
1700 let fill_bytes = match params {
1701 OpParams::MaskedFill { fill_bytes } => fill_bytes.clone(),
1702 other => {
1703 return Err(Error::Msg(format!(
1704 "masked_fill_cpu_wrapper: expects OpParams::MaskedFill, got {other:?}",
1705 ))
1706 .bt());
1707 }
1708 };
1709 let in_guard = read_storage(&inputs[0])?;
1710 let mask_guard = read_storage(&inputs[1])?;
1711 if mask_guard.dtype != DType::Bool {
1718 return Err(Error::Msg(format!(
1719 "masked_fill_cpu_wrapper: mask must be Bool, got {:?} (GAP-168(c))",
1720 mask_guard.dtype,
1721 ))
1722 .bt());
1723 }
1724 let mut out_guard = write_storage(&outputs[0])?;
1725 let dtype_size = out_guard.dtype.size_in_bytes();
1726 let in_cpu = cpu_input(&in_guard)?;
1727 let mask_cpu = cpu_input(&mask_guard)?;
1728 let out_cpu = cpu_output(&mut out_guard)?;
1729 fuel_cpu_backend::byte_kernels::masked_fill_cpu(
1730 in_cpu,
1731 mask_cpu,
1732 out_cpu,
1733 &fill_bytes,
1734 dtype_size,
1735 )
1736}
1737
1738pub(crate) fn pad_cpu_wrapper(
1744 inputs: &[Arc<RwLock<Storage>>],
1745 outputs: &mut [Arc<RwLock<Storage>>],
1746 _layouts: &[Layout],
1747 params: &OpParams,
1748) -> Result<()> {
1749 if inputs.len() != 1 || outputs.len() != 1 {
1750 return Err(Error::Msg(format!(
1751 "pad wrapper expects 1 input + 1 output, got {} + {}",
1752 inputs.len(),
1753 outputs.len(),
1754 ))
1755 .bt());
1756 }
1757 let (in_shape, out_shape, padding, mode_tag, fill_bytes) = match params {
1758 OpParams::Pad {
1759 in_shape,
1760 out_shape,
1761 padding,
1762 mode_tag,
1763 fill_bytes,
1764 } => (in_shape, out_shape, padding, *mode_tag, fill_bytes),
1765 other => {
1766 return Err(
1767 Error::Msg(format!("pad wrapper expects OpParams::Pad, got {other:?}",)).bt(),
1768 );
1769 }
1770 };
1771 let in_guard = read_storage(&inputs[0])?;
1772 let mut out_guard = write_storage(&outputs[0])?;
1773 let dtype_size = out_guard.dtype.size_in_bytes();
1774 let in_cpu = cpu_input(&in_guard)?;
1775 let out_cpu = cpu_output(&mut out_guard)?;
1776 match mode_tag {
1778 0 => fuel_cpu_backend::byte_kernels::pad_const_cpu(
1779 in_cpu, out_cpu, in_shape, out_shape, padding, dtype_size, fill_bytes,
1780 ),
1781 1 => {
1782 for (k, (&n, &(b, a))) in in_shape.iter().zip(padding.iter()).enumerate() {
1785 if n > 0 && (b > n - 1 || a > n - 1) {
1786 return Err(Error::Msg(format!(
1787 "pad reflect: axis {k} has dim_size {n}; before ({b}) and \
1788 after ({a}) must each be <= dim_size - 1",
1789 ))
1790 .bt());
1791 }
1792 }
1793 fuel_cpu_backend::byte_kernels::pad_reflect_cpu(
1794 in_cpu, out_cpu, in_shape, out_shape, padding, dtype_size,
1795 )
1796 }
1797 2 => fuel_cpu_backend::byte_kernels::pad_replicate_cpu(
1798 in_cpu, out_cpu, in_shape, out_shape, padding, dtype_size,
1799 ),
1800 other => Err(Error::Msg(format!("pad: unknown mode_tag {other}",)).bt()),
1801 }
1802}
1803
1804macro_rules! cpu_pad_backward_wrapper {
1808 ($wrapper:ident, $kernel:path, $op_name:literal) => {
1809 pub(crate) fn $wrapper(
1810 inputs: &[Arc<RwLock<Storage>>],
1811 outputs: &mut [Arc<RwLock<Storage>>],
1812 _layouts: &[Layout],
1813 params: &OpParams,
1814 ) -> Result<()> {
1815 if inputs.len() != 1 || outputs.len() != 1 {
1816 return Err(Error::Msg(format!(
1817 "{} wrapper expects 1 input + 1 output, got {} + {}",
1818 $op_name,
1819 inputs.len(),
1820 outputs.len(),
1821 ))
1822 .bt());
1823 }
1824 let (in_shape, out_shape, padding, mode_tag) = match params {
1825 OpParams::PadBackward {
1826 in_shape,
1827 out_shape,
1828 padding,
1829 mode_tag,
1830 } => (in_shape, out_shape, padding, *mode_tag),
1831 other => {
1832 return Err(Error::Msg(format!(
1833 "{} wrapper expects OpParams::PadBackward, got {other:?}",
1834 $op_name,
1835 ))
1836 .bt())
1837 }
1838 };
1839 let in_guard = read_storage(&inputs[0])?;
1840 let mut out_guard = write_storage(&outputs[0])?;
1841 let in_cpu = cpu_input(&in_guard)?;
1842 let out_cpu = cpu_output(&mut out_guard)?;
1843 $kernel(in_cpu, out_cpu, in_shape, out_shape, padding, mode_tag)
1844 }
1845 };
1846}
1847
1848cpu_pad_backward_wrapper!(
1849 pad_backward_f32_cpu_wrapper,
1850 fuel_cpu_backend::byte_kernels::pad_backward_f32,
1851 "pad_backward_f32"
1852);
1853cpu_pad_backward_wrapper!(
1854 pad_backward_f64_cpu_wrapper,
1855 fuel_cpu_backend::byte_kernels::pad_backward_f64,
1856 "pad_backward_f64"
1857);
1858cpu_pad_backward_wrapper!(
1859 pad_backward_bf16_cpu_wrapper,
1860 fuel_cpu_backend::byte_kernels::pad_backward_bf16,
1861 "pad_backward_bf16"
1862);
1863cpu_pad_backward_wrapper!(
1864 pad_backward_f16_cpu_wrapper,
1865 fuel_cpu_backend::byte_kernels::pad_backward_f16,
1866 "pad_backward_f16"
1867);
1868
1869pub(crate) fn roll_cpu_wrapper(
1872 inputs: &[Arc<RwLock<Storage>>],
1873 outputs: &mut [Arc<RwLock<Storage>>],
1874 _layouts: &[Layout],
1875 params: &OpParams,
1876) -> Result<()> {
1877 if inputs.len() != 1 || outputs.len() != 1 {
1878 return Err(Error::Msg(format!(
1879 "roll wrapper expects 1 input + 1 output, got {} + {}",
1880 inputs.len(),
1881 outputs.len(),
1882 ))
1883 .bt());
1884 }
1885 let (outer, dim_size, inner, shift) = match params {
1886 OpParams::Roll {
1887 outer_count,
1888 dim_size,
1889 inner_count,
1890 shift,
1891 ..
1892 } => (*outer_count, *dim_size, *inner_count, *shift),
1893 other => {
1894 return Err(Error::Msg(format!(
1895 "roll wrapper expects OpParams::Roll, got {other:?}",
1896 ))
1897 .bt());
1898 }
1899 };
1900 let in_guard = read_storage(&inputs[0])?;
1901 let mut out_guard = write_storage(&outputs[0])?;
1902 let dtype_size = out_guard.dtype.size_in_bytes();
1903 let in_cpu = cpu_input(&in_guard)?;
1904 let out_cpu = cpu_output(&mut out_guard)?;
1905 fuel_cpu_backend::byte_kernels::roll_cpu(
1906 in_cpu, out_cpu, outer, dim_size, inner, shift, dtype_size,
1907 )
1908}
1909
1910macro_rules! cpu_arg_dim_wrapper {
1914 ($wrapper:ident, $kernel:path, $op_name:literal) => {
1915 fn $wrapper(
1916 inputs: &[Arc<RwLock<Storage>>],
1917 outputs: &mut [Arc<RwLock<Storage>>],
1918 layouts: &[Layout],
1919 params: &OpParams,
1920 ) -> Result<()> {
1921 if inputs.len() != 1 || outputs.len() != 1 {
1922 return Err(Error::Msg(format!(
1923 "{} wrapper expects 1 input + 1 output, got {} + {}",
1924 $op_name,
1925 inputs.len(),
1926 outputs.len(),
1927 ))
1928 .bt());
1929 }
1930 let dims = match params {
1931 OpParams::Reduce { dims, .. } => dims,
1932 other => {
1933 return Err(Error::Msg(format!(
1934 "{} wrapper expects OpParams::Reduce, got {other:?}",
1935 $op_name,
1936 ))
1937 .bt())
1938 }
1939 };
1940 if dims.len() != 1 {
1941 return Err(Error::Msg(format!(
1942 "{} wrapper expects exactly 1 reduce dim, got {dims:?}",
1943 $op_name,
1944 ))
1945 .bt());
1946 }
1947 let dim = dims[0];
1948 let input_layout = layouts.first().ok_or_else(|| {
1949 Error::Msg(format!(
1950 "{} wrapper: layouts empty (executor must pass input layout at [0])",
1951 $op_name,
1952 ))
1953 .bt()
1954 })?;
1955 let in_guard = read_storage(&inputs[0])?;
1956 if in_guard.dtype != DType::F32 {
1957 return Err(Error::Msg(format!(
1958 "{}: only F32 input is wired today, got {:?}",
1959 $op_name, in_guard.dtype,
1960 ))
1961 .bt());
1962 }
1963 let mut out_guard = write_storage(&outputs[0])?;
1964 let in_cpu = cpu_input(&in_guard)?;
1965 let out_cpu = cpu_output(&mut out_guard)?;
1966 $kernel(in_cpu, out_cpu, input_layout.shape().dims(), dim)
1967 }
1968 };
1969}
1970
1971cpu_arg_dim_wrapper!(
1972 argmax_dim_u32_cpu_wrapper,
1973 fuel_cpu_backend::byte_kernels::argmax_dim_f32,
1974 "argmax_dim"
1975);
1976cpu_arg_dim_wrapper!(
1977 argmin_dim_u32_cpu_wrapper,
1978 fuel_cpu_backend::byte_kernels::argmin_dim_f32,
1979 "argmin_dim"
1980);
1981
1982pub(crate) fn index_add_f32_cpu_wrapper(
1987 inputs: &[Arc<RwLock<Storage>>],
1988 outputs: &mut [Arc<RwLock<Storage>>],
1989 _layouts: &[Layout],
1990 params: &OpParams,
1991) -> Result<()> {
1992 if inputs.len() != 3 || outputs.len() != 1 {
1993 return Err(Error::Msg(format!(
1994 "index_add wrapper expects 3 inputs (base, indices, src) + 1 output, got {} + {}",
1995 inputs.len(),
1996 outputs.len(),
1997 ))
1998 .bt());
1999 }
2000 let (outer_count, base_dim_size, n_indices, inner_count) = match params {
2001 OpParams::IndexAdd {
2002 outer_count,
2003 base_dim_size,
2004 n_indices,
2005 inner_count,
2006 } => (*outer_count, *base_dim_size, *n_indices, *inner_count),
2007 other => {
2008 return Err(Error::Msg(format!(
2009 "index_add wrapper expects OpParams::IndexAdd, got {other:?}",
2010 ))
2011 .bt());
2012 }
2013 };
2014 let base_guard = read_storage(&inputs[0])?;
2015 let idx_guard = read_storage(&inputs[1])?;
2016 if idx_guard.dtype != DType::U32 {
2017 return Err(Error::Msg(format!(
2018 "index_add: indices must be U32, got {:?}",
2019 idx_guard.dtype,
2020 ))
2021 .bt());
2022 }
2023 let src_guard = read_storage(&inputs[2])?;
2024 let mut out_guard = write_storage(&outputs[0])?;
2025 let base_cpu = cpu_input(&base_guard)?;
2026 let idx_cpu = cpu_input(&idx_guard)?;
2027 let src_cpu = cpu_input(&src_guard)?;
2028 let out_cpu = cpu_output(&mut out_guard)?;
2029 fuel_cpu_backend::byte_kernels::index_add_f32(
2030 base_cpu,
2031 idx_cpu,
2032 src_cpu,
2033 out_cpu,
2034 outer_count,
2035 base_dim_size,
2036 n_indices,
2037 inner_count,
2038 )
2039}
2040
2041pub(crate) fn scatter_add_f32_cpu_wrapper(
2046 inputs: &[Arc<RwLock<Storage>>],
2047 outputs: &mut [Arc<RwLock<Storage>>],
2048 _layouts: &[Layout],
2049 params: &OpParams,
2050) -> Result<()> {
2051 if inputs.len() != 3 || outputs.len() != 1 {
2052 return Err(Error::Msg(format!(
2053 "scatter_add wrapper expects 3 inputs (base, indices, src) + 1 output, got {} + {}",
2054 inputs.len(),
2055 outputs.len(),
2056 ))
2057 .bt());
2058 }
2059 let (base_shape, src_shape, dim) = match params {
2060 OpParams::ScatterAdd {
2061 base_shape,
2062 src_shape,
2063 dim,
2064 } => (base_shape, src_shape, *dim),
2065 other => {
2066 return Err(Error::Msg(format!(
2067 "scatter_add wrapper expects OpParams::ScatterAdd, got {other:?}",
2068 ))
2069 .bt());
2070 }
2071 };
2072 let base_guard = read_storage(&inputs[0])?;
2073 let idx_guard = read_storage(&inputs[1])?;
2074 if idx_guard.dtype != DType::U32 {
2075 return Err(Error::Msg(format!(
2076 "scatter_add: indices must be U32, got {:?}",
2077 idx_guard.dtype,
2078 ))
2079 .bt());
2080 }
2081 let src_guard = read_storage(&inputs[2])?;
2082 let mut out_guard = write_storage(&outputs[0])?;
2083 let base_cpu = cpu_input(&base_guard)?;
2084 let idx_cpu = cpu_input(&idx_guard)?;
2085 let src_cpu = cpu_input(&src_guard)?;
2086 let out_cpu = cpu_output(&mut out_guard)?;
2087 fuel_cpu_backend::byte_kernels::scatter_add_f32(
2088 base_cpu, idx_cpu, src_cpu, out_cpu, base_shape, src_shape, dim,
2089 )
2090}
2091
2092pub(crate) fn rope_f32_cpu_wrapper(
2097 inputs: &[Arc<RwLock<Storage>>],
2098 outputs: &mut [Arc<RwLock<Storage>>],
2099 _layouts: &[Layout],
2100 params: &OpParams,
2101) -> Result<()> {
2102 if inputs.len() != 3 || outputs.len() != 1 {
2103 return Err(Error::Msg(format!(
2104 "rope wrapper expects 3 inputs (x, cos, sin) + 1 output, got {} + {}",
2105 inputs.len(),
2106 outputs.len(),
2107 ))
2108 .bt());
2109 }
2110 let (outer_count, seq, head_dim) = match params {
2111 OpParams::Rope {
2112 outer_count,
2113 seq,
2114 head_dim,
2115 } => (*outer_count, *seq, *head_dim),
2116 other => {
2117 return Err(Error::Msg(format!(
2118 "rope wrapper expects OpParams::Rope, got {other:?}",
2119 ))
2120 .bt());
2121 }
2122 };
2123 let x_guard = read_storage(&inputs[0])?;
2124 let cos_guard = read_storage(&inputs[1])?;
2125 let sin_guard = read_storage(&inputs[2])?;
2126 let mut out_guard = write_storage(&outputs[0])?;
2127 let x_cpu = cpu_input(&x_guard)?;
2128 let cos_cpu = cpu_input(&cos_guard)?;
2129 let sin_cpu = cpu_input(&sin_guard)?;
2130 let out_cpu = cpu_output(&mut out_guard)?;
2131 fuel_cpu_backend::byte_kernels::rope_f32(
2132 x_cpu,
2133 cos_cpu,
2134 sin_cpu,
2135 out_cpu,
2136 outer_count,
2137 seq,
2138 head_dim,
2139 )
2140}
2141
2142pub(crate) fn gather_cpu_wrapper(
2147 inputs: &[Arc<RwLock<Storage>>],
2148 outputs: &mut [Arc<RwLock<Storage>>],
2149 _layouts: &[Layout],
2150 params: &OpParams,
2151) -> Result<()> {
2152 if inputs.len() != 2 || outputs.len() != 1 {
2153 return Err(Error::Msg(format!(
2154 "gather wrapper expects 2 inputs + 1 output, got {} + {}",
2155 inputs.len(),
2156 outputs.len(),
2157 ))
2158 .bt());
2159 }
2160 let (source_shape, output_shape, dim) = match params {
2161 OpParams::Gather {
2162 source_shape,
2163 output_shape,
2164 dim,
2165 } => (source_shape, output_shape, *dim),
2166 other => {
2167 return Err(Error::Msg(format!(
2168 "gather wrapper expects OpParams::Gather, got {other:?}",
2169 ))
2170 .bt());
2171 }
2172 };
2173 let src_guard = read_storage(&inputs[0])?;
2174 let idx_guard = read_storage(&inputs[1])?;
2175 if idx_guard.dtype != DType::U32 {
2176 return Err(Error::Msg(format!(
2177 "gather: indices must be U32, got {:?}",
2178 idx_guard.dtype,
2179 ))
2180 .bt());
2181 }
2182 let mut out_guard = write_storage(&outputs[0])?;
2183 let dtype_size = out_guard.dtype.size_in_bytes();
2184 let src_cpu = cpu_input(&src_guard)?;
2185 let idx_cpu = cpu_input(&idx_guard)?;
2186 let out_cpu = cpu_output(&mut out_guard)?;
2187 fuel_cpu_backend::byte_kernels::gather_cpu(
2188 src_cpu,
2189 idx_cpu,
2190 out_cpu,
2191 source_shape,
2192 output_shape,
2193 dim,
2194 dtype_size,
2195 )
2196}
2197
2198pub(crate) fn index_select_cpu_wrapper(
2204 inputs: &[Arc<RwLock<Storage>>],
2205 outputs: &mut [Arc<RwLock<Storage>>],
2206 _layouts: &[Layout],
2207 params: &OpParams,
2208) -> Result<()> {
2209 if inputs.len() != 2 {
2210 return Err(Error::Msg(format!(
2211 "index_select wrapper expects 2 inputs (source, indices), got {}",
2212 inputs.len(),
2213 ))
2214 .bt());
2215 }
2216 if outputs.len() != 1 {
2217 return Err(Error::Msg(format!(
2218 "index_select wrapper expects 1 output, got {}",
2219 outputs.len(),
2220 ))
2221 .bt());
2222 }
2223 let (outer_count, source_dim_size, n_indices, inner_count) = match params {
2224 OpParams::IndexSelect {
2225 outer_count,
2226 source_dim_size,
2227 n_indices,
2228 inner_count,
2229 } => (*outer_count, *source_dim_size, *n_indices, *inner_count),
2230 other => {
2231 return Err(Error::Msg(format!(
2232 "index_select wrapper expects OpParams::IndexSelect, got {other:?}",
2233 ))
2234 .bt());
2235 }
2236 };
2237 let src_guard = read_storage(&inputs[0])?;
2238 let idx_guard = read_storage(&inputs[1])?;
2239 if idx_guard.dtype != DType::U32 {
2240 return Err(Error::Msg(format!(
2241 "index_select: indices must be U32, got {:?}",
2242 idx_guard.dtype,
2243 ))
2244 .bt());
2245 }
2246 let mut out_guard = write_storage(&outputs[0])?;
2247 let dtype_size = out_guard.dtype.size_in_bytes();
2248 let src_cpu = cpu_input(&src_guard)?;
2249 let idx_cpu = cpu_input(&idx_guard)?;
2250 let out_cpu = cpu_output(&mut out_guard)?;
2251 fuel_cpu_backend::byte_kernels::index_select_cpu(
2252 src_cpu,
2253 idx_cpu,
2254 out_cpu,
2255 outer_count,
2256 source_dim_size,
2257 n_indices,
2258 inner_count,
2259 dtype_size,
2260 )
2261}
2262
2263macro_rules! cpu_norm_last_dim_wrapper {
2267 ($wrapper:ident, $kernel:path, $op_name:literal) => {
2268 pub(crate) fn $wrapper(
2269 inputs: &[Arc<RwLock<Storage>>],
2270 outputs: &mut [Arc<RwLock<Storage>>],
2271 _layouts: &[Layout],
2272 params: &OpParams,
2273 ) -> Result<()> {
2274 if inputs.len() != 1 || outputs.len() != 1 {
2275 return Err(Error::Msg(format!(
2276 "{} wrapper expects 1 input + 1 output, got {} + {}",
2277 $op_name,
2278 inputs.len(),
2279 outputs.len(),
2280 ))
2281 .bt());
2282 }
2283 let (outer_count, last_dim, eps) = match params {
2284 OpParams::NormLastDim {
2285 outer_count,
2286 last_dim,
2287 eps,
2288 } => (*outer_count, *last_dim, *eps),
2289 other => {
2290 return Err(Error::Msg(format!(
2291 "{} wrapper expects OpParams::NormLastDim, got {other:?}",
2292 $op_name,
2293 ))
2294 .bt())
2295 }
2296 };
2297 let in_guard = read_storage(&inputs[0])?;
2298 let mut out_guard = write_storage(&outputs[0])?;
2299 let in_cpu = cpu_input(&in_guard)?;
2300 let out_cpu = cpu_output(&mut out_guard)?;
2301 $kernel(in_cpu, out_cpu, outer_count, last_dim, eps)
2302 }
2303 };
2304}
2305
2306cpu_norm_last_dim_wrapper!(
2307 rms_norm_last_dim_f32_cpu_wrapper,
2308 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_f32,
2309 "rms_norm_last_dim"
2310);
2311cpu_norm_last_dim_wrapper!(
2312 layer_norm_last_dim_f32_cpu_wrapper,
2313 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_f32,
2314 "layer_norm_last_dim"
2315);
2316cpu_norm_last_dim_wrapper!(
2317 rms_norm_last_dim_bf16_cpu_wrapper,
2318 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_bf16,
2319 "rms_norm_last_dim"
2320);
2321cpu_norm_last_dim_wrapper!(
2322 rms_norm_last_dim_f16_cpu_wrapper,
2323 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_f16,
2324 "rms_norm_last_dim"
2325);
2326cpu_norm_last_dim_wrapper!(
2327 layer_norm_last_dim_bf16_cpu_wrapper,
2328 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_bf16,
2329 "layer_norm_last_dim"
2330);
2331cpu_norm_last_dim_wrapper!(
2332 layer_norm_last_dim_f16_cpu_wrapper,
2333 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_f16,
2334 "layer_norm_last_dim"
2335);
2336cpu_norm_last_dim_wrapper!(
2337 rms_norm_last_dim_f64_cpu_wrapper,
2338 fuel_cpu_backend::byte_kernels::rms_norm_last_dim_f64,
2339 "rms_norm_last_dim"
2340);
2341cpu_norm_last_dim_wrapper!(
2342 layer_norm_last_dim_f64_cpu_wrapper,
2343 fuel_cpu_backend::byte_kernels::layer_norm_last_dim_f64,
2344 "layer_norm_last_dim"
2345);
2346
2347macro_rules! cpu_softmax_last_dim_wrapper {
2354 ($wrapper:ident, $kernel:path) => {
2355 pub(crate) fn $wrapper(
2356 inputs: &[Arc<RwLock<Storage>>],
2357 outputs: &mut [Arc<RwLock<Storage>>],
2358 _layouts: &[Layout],
2359 params: &OpParams,
2360 ) -> Result<()> {
2361 if inputs.len() != 1 || outputs.len() != 1 {
2362 return Err(Error::Msg(format!(
2363 "softmax_last_dim wrapper expects 1 input + 1 output, got {} + {}",
2364 inputs.len(),
2365 outputs.len(),
2366 ))
2367 .bt());
2368 }
2369 let (outer_count, last_dim) =
2370 match params {
2371 OpParams::SoftmaxLastDim {
2372 outer_count,
2373 last_dim,
2374 } => (*outer_count, *last_dim),
2375 other => return Err(Error::Msg(format!(
2376 "softmax_last_dim wrapper expects OpParams::SoftmaxLastDim, got {other:?}",
2377 ))
2378 .bt()),
2379 };
2380 let in_guard = read_storage(&inputs[0])?;
2381 let mut out_guard = write_storage(&outputs[0])?;
2382 let in_cpu = cpu_input(&in_guard)?;
2383 let out_cpu = cpu_output(&mut out_guard)?;
2384 $kernel(in_cpu, out_cpu, outer_count, last_dim)
2385 }
2386 };
2387}
2388
2389cpu_softmax_last_dim_wrapper!(
2390 softmax_last_dim_f32_cpu_wrapper,
2391 fuel_cpu_backend::byte_kernels::softmax_last_dim_f32
2392);
2393cpu_softmax_last_dim_wrapper!(
2394 softmax_last_dim_f64_cpu_wrapper,
2395 fuel_cpu_backend::byte_kernels::softmax_last_dim_f64
2396);
2397cpu_softmax_last_dim_wrapper!(
2398 softmax_last_dim_bf16_cpu_wrapper,
2399 fuel_cpu_backend::byte_kernels::softmax_last_dim_bf16
2400);
2401cpu_softmax_last_dim_wrapper!(
2402 softmax_last_dim_f16_cpu_wrapper,
2403 fuel_cpu_backend::byte_kernels::softmax_last_dim_f16
2404);
2405
2406macro_rules! cpu_rope_wrapper {
2412 ($wrapper:ident, $kernel:path) => {
2413 pub(crate) fn $wrapper(
2414 inputs: &[Arc<RwLock<Storage>>],
2415 outputs: &mut [Arc<RwLock<Storage>>],
2416 _layouts: &[Layout],
2417 params: &OpParams,
2418 ) -> Result<()> {
2419 if inputs.len() != 3 || outputs.len() != 1 {
2420 return Err(Error::Msg(format!(
2421 "rope wrapper expects 3 inputs + 1 output, got {} + {}",
2422 inputs.len(),
2423 outputs.len(),
2424 ))
2425 .bt());
2426 }
2427 let (outer_count, seq, head_dim) = match params {
2428 OpParams::Rope {
2429 outer_count,
2430 seq,
2431 head_dim,
2432 } => (*outer_count, *seq, *head_dim),
2433 other => {
2434 return Err(Error::Msg(format!(
2435 "rope wrapper expects OpParams::Rope, got {other:?}",
2436 ))
2437 .bt())
2438 }
2439 };
2440 let x_guard = read_storage(&inputs[0])?;
2441 let cos_guard = read_storage(&inputs[1])?;
2442 let sin_guard = read_storage(&inputs[2])?;
2443 let mut out_guard = write_storage(&outputs[0])?;
2444 let x_cpu = cpu_input(&x_guard)?;
2445 let cos_cpu = cpu_input(&cos_guard)?;
2446 let sin_cpu = cpu_input(&sin_guard)?;
2447 let out_cpu = cpu_output(&mut out_guard)?;
2448 $kernel(x_cpu, cos_cpu, sin_cpu, out_cpu, outer_count, seq, head_dim)
2449 }
2450 };
2451}
2452
2453cpu_rope_wrapper!(
2454 rope_bf16_cpu_wrapper,
2455 fuel_cpu_backend::byte_kernels::rope_bf16
2456);
2457cpu_rope_wrapper!(
2458 rope_f16_cpu_wrapper,
2459 fuel_cpu_backend::byte_kernels::rope_f16
2460);
2461cpu_rope_wrapper!(
2462 rope_f64_cpu_wrapper,
2463 fuel_cpu_backend::byte_kernels::rope_f64
2464);
2465
2466pub(crate) fn qmatmul_f32_cpu_wrapper(
2475 inputs: &[Arc<RwLock<Storage>>],
2476 outputs: &mut [Arc<RwLock<Storage>>],
2477 _layouts: &[Layout],
2478 params: &OpParams,
2479) -> Result<()> {
2480 if inputs.len() != 2 {
2481 return Err(Error::Msg(format!(
2482 "qmatmul wrapper expects 2 inputs (activations, weight_bytes), got {}",
2483 inputs.len(),
2484 ))
2485 .bt());
2486 }
2487 if outputs.len() != 1 {
2488 return Err(Error::Msg(format!(
2489 "qmatmul wrapper expects 1 output, got {}",
2490 outputs.len(),
2491 ))
2492 .bt());
2493 }
2494 let (quant_type, batch_count, m, n, k) = match params {
2495 OpParams::QMatMul {
2496 quant_type,
2497 batch_count,
2498 m,
2499 n,
2500 k,
2501 } => (*quant_type, *batch_count, *m, *n, *k),
2502 other => {
2503 return Err(Error::Msg(format!(
2504 "qmatmul wrapper expects OpParams::QMatMul, got {other:?}",
2505 ))
2506 .bt());
2507 }
2508 };
2509 let act_guard = read_storage(&inputs[0])?;
2510 let w_guard = read_storage(&inputs[1])?;
2511 if w_guard.dtype != DType::U32 {
2512 return Err(Error::Msg(format!(
2513 "qmatmul: weight bytes must be U32-typed (raw block stream), got {:?}",
2514 w_guard.dtype,
2515 ))
2516 .bt());
2517 }
2518 let mut out_guard = write_storage(&outputs[0])?;
2519 let act_cpu = cpu_input(&act_guard)?;
2520 let w_cpu = cpu_input(&w_guard)?;
2521 let out_cpu = cpu_output(&mut out_guard)?;
2522 use fuel_graph::QuantType;
2523 match quant_type {
2524 QuantType::Q4_0 => fuel_cpu_backend::byte_kernels::qmatmul_q4_0_f32(
2525 act_cpu,
2526 w_cpu,
2527 out_cpu,
2528 batch_count,
2529 m,
2530 n,
2531 k,
2532 ),
2533 QuantType::Q4_1 => fuel_cpu_backend::byte_kernels::qmatmul_q4_1_f32(
2534 act_cpu,
2535 w_cpu,
2536 out_cpu,
2537 batch_count,
2538 m,
2539 n,
2540 k,
2541 ),
2542 QuantType::Q5_0 => fuel_cpu_backend::byte_kernels::qmatmul_q5_0_f32(
2543 act_cpu,
2544 w_cpu,
2545 out_cpu,
2546 batch_count,
2547 m,
2548 n,
2549 k,
2550 ),
2551 QuantType::Q5_1 => fuel_cpu_backend::byte_kernels::qmatmul_q5_1_f32(
2552 act_cpu,
2553 w_cpu,
2554 out_cpu,
2555 batch_count,
2556 m,
2557 n,
2558 k,
2559 ),
2560 QuantType::Q8_0 => fuel_cpu_backend::byte_kernels::qmatmul_q8_0_f32(
2561 act_cpu,
2562 w_cpu,
2563 out_cpu,
2564 batch_count,
2565 m,
2566 n,
2567 k,
2568 ),
2569 QuantType::Q8_1 => fuel_cpu_backend::byte_kernels::qmatmul_q8_1_f32(
2570 act_cpu,
2571 w_cpu,
2572 out_cpu,
2573 batch_count,
2574 m,
2575 n,
2576 k,
2577 ),
2578 QuantType::Q2K => fuel_cpu_backend::byte_kernels::qmatmul_q2k_f32(
2579 act_cpu,
2580 w_cpu,
2581 out_cpu,
2582 batch_count,
2583 m,
2584 n,
2585 k,
2586 ),
2587 QuantType::Q3K => fuel_cpu_backend::byte_kernels::qmatmul_q3k_f32(
2588 act_cpu,
2589 w_cpu,
2590 out_cpu,
2591 batch_count,
2592 m,
2593 n,
2594 k,
2595 ),
2596 QuantType::Q4_K_M => fuel_cpu_backend::byte_kernels::qmatmul_q4_k_m_f32(
2597 act_cpu,
2598 w_cpu,
2599 out_cpu,
2600 batch_count,
2601 m,
2602 n,
2603 k,
2604 ),
2605 QuantType::Q5K => fuel_cpu_backend::byte_kernels::qmatmul_q5k_f32(
2606 act_cpu,
2607 w_cpu,
2608 out_cpu,
2609 batch_count,
2610 m,
2611 n,
2612 k,
2613 ),
2614 QuantType::Q6K => fuel_cpu_backend::byte_kernels::qmatmul_q6k_f32(
2615 act_cpu,
2616 w_cpu,
2617 out_cpu,
2618 batch_count,
2619 m,
2620 n,
2621 k,
2622 ),
2623 }
2624}
2625
2626pub(crate) fn write_slice_cpu_wrapper(
2634 inputs: &[Arc<RwLock<Storage>>],
2635 outputs: &mut [Arc<RwLock<Storage>>],
2636 _layouts: &[Layout],
2637 params: &OpParams,
2638) -> Result<()> {
2639 if inputs.len() != 1 || outputs.len() != 1 {
2640 return Err(Error::Msg(format!(
2641 "write_slice wrapper expects 1 input (source) + 1 output (dest), \
2642 got {} + {}",
2643 inputs.len(),
2644 outputs.len(),
2645 ))
2646 .bt());
2647 }
2648 let (dest_shape, ranges) = match params {
2649 OpParams::WriteSlice {
2650 dest_shape, ranges, ..
2651 } => (dest_shape, ranges),
2652 other => {
2653 return Err(Error::Msg(format!(
2654 "write_slice wrapper expects OpParams::WriteSlice, got {other:?}",
2655 ))
2656 .bt());
2657 }
2658 };
2659 let src_guard = read_storage(&inputs[0])?;
2660 let src_cpu = cpu_input(&src_guard)?;
2661 let mut dest_guard = write_storage(&outputs[0])?;
2662 let dtype_size = dest_guard.dtype.size_in_bytes();
2663 let dest_cpu = cpu_output(&mut dest_guard)?;
2664 fuel_cpu_backend::byte_kernels::write_slice_cpu(
2665 src_cpu, dest_cpu, dest_shape, ranges, dtype_size,
2666 )
2667}
2668
2669pub(crate) fn write_slice_rotating_cpu_wrapper(
2677 inputs: &[Arc<RwLock<Storage>>],
2678 outputs: &mut [Arc<RwLock<Storage>>],
2679 _layouts: &[Layout],
2680 params: &OpParams,
2681) -> Result<()> {
2682 if inputs.len() != 2 || outputs.len() != 1 {
2683 return Err(Error::Msg(format!(
2684 "write_slice_rotating wrapper expects 2 inputs (source, position) + 1 output (dest), \
2685 got {} + {}",
2686 inputs.len(),
2687 outputs.len(),
2688 ))
2689 .bt());
2690 }
2691 let (dest_shape, axis, modulus, ranges) =
2692 match params {
2693 OpParams::WriteSliceRotating {
2694 dest_shape,
2695 axis,
2696 modulus,
2697 ranges,
2698 } => (dest_shape, *axis, *modulus, ranges),
2699 other => return Err(Error::Msg(format!(
2700 "write_slice_rotating wrapper expects OpParams::WriteSliceRotating, got {other:?}",
2701 ))
2702 .bt()),
2703 };
2704 let src_guard = read_storage(&inputs[0])?;
2705 let src_cpu = cpu_input(&src_guard)?;
2706 let pos_guard = read_storage(&inputs[1])?;
2707 let pos_cpu = cpu_input(&pos_guard)?;
2708 let mut dest_guard = write_storage(&outputs[0])?;
2709 let dtype_size = dest_guard.dtype.size_in_bytes();
2710 let dest_cpu = cpu_output(&mut dest_guard)?;
2711 fuel_cpu_backend::byte_kernels::write_slice_rotating_cpu(
2712 src_cpu, pos_cpu, dest_cpu, dest_shape, axis, modulus, ranges, dtype_size,
2713 )
2714}
2715
2716pub(crate) fn write_slice_doff_cpu_wrapper(
2723 inputs: &[Arc<RwLock<Storage>>],
2724 outputs: &mut [Arc<RwLock<Storage>>],
2725 _layouts: &[Layout],
2726 params: &OpParams,
2727) -> Result<()> {
2728 if inputs.len() != 2 || outputs.len() != 1 {
2729 return Err(Error::Msg(format!(
2730 "write_slice_doff wrapper expects 2 inputs (source, offset) + 1 output (dest), \
2731 got {} + {}",
2732 inputs.len(),
2733 outputs.len(),
2734 ))
2735 .bt());
2736 }
2737 let (dest_shape, axis, ranges) = match params {
2738 OpParams::WriteSliceDoff {
2739 dest_shape,
2740 axis,
2741 ranges,
2742 } => (dest_shape, *axis, ranges),
2743 other => {
2744 return Err(Error::Msg(format!(
2745 "write_slice_doff wrapper expects OpParams::WriteSliceDoff, got {other:?}",
2746 ))
2747 .bt());
2748 }
2749 };
2750 let src_guard = read_storage(&inputs[0])?;
2751 let src_cpu = cpu_input(&src_guard)?;
2752 let off_guard = read_storage(&inputs[1])?;
2753 let off_cpu = cpu_input(&off_guard)?;
2754 let mut dest_guard = write_storage(&outputs[0])?;
2755 let dtype_size = dest_guard.dtype.size_in_bytes();
2756 let dest_cpu = cpu_output(&mut dest_guard)?;
2757 fuel_cpu_backend::byte_kernels::write_slice_doff_cpu(
2758 src_cpu, off_cpu, dest_cpu, dest_shape, axis, ranges, dtype_size,
2759 )
2760}
2761
2762pub(crate) fn concat_cpu_wrapper(
2763 inputs: &[Arc<RwLock<Storage>>],
2764 outputs: &mut [Arc<RwLock<Storage>>],
2765 _layouts: &[Layout],
2766 params: &OpParams,
2767) -> Result<()> {
2768 if inputs.is_empty() {
2769 return Err(Error::Msg("concat wrapper expects ≥ 1 input, got 0".to_string()).bt());
2770 }
2771 if outputs.len() != 1 {
2772 return Err(Error::Msg(format!(
2773 "concat wrapper expects 1 output, got {}",
2774 outputs.len(),
2775 ))
2776 .bt());
2777 }
2778 let (outer_count, input_dim_sizes, inner_count) = match params {
2779 OpParams::Concat {
2780 outer_count,
2781 input_dim_sizes,
2782 inner_count,
2783 ..
2784 } => (*outer_count, input_dim_sizes, *inner_count),
2785 other => {
2786 return Err(Error::Msg(format!(
2787 "concat wrapper expects OpParams::Concat, got {other:?}",
2788 ))
2789 .bt());
2790 }
2791 };
2792 if input_dim_sizes.len() != inputs.len() {
2793 return Err(Error::Msg(format!(
2794 "concat wrapper: OpParams declares {} inputs but the work item \
2795 carries {}",
2796 input_dim_sizes.len(),
2797 inputs.len(),
2798 ))
2799 .bt());
2800 }
2801 let in_guards: Vec<_> = inputs
2802 .iter()
2803 .map(read_storage)
2804 .collect::<Result<Vec<_>>>()?;
2805 let mut in_cpus: Vec<&fuel_cpu_backend::CpuStorageBytes> = Vec::with_capacity(in_guards.len());
2806 for g in &in_guards {
2807 in_cpus.push(cpu_input(g)?);
2808 }
2809 let mut out_guard = write_storage(&outputs[0])?;
2810 let dtype_size = out_guard.dtype.size_in_bytes();
2811 let out_cpu = cpu_output(&mut out_guard)?;
2812 fuel_cpu_backend::byte_kernels::concat_cpu(
2813 &in_cpus,
2814 out_cpu,
2815 outer_count,
2816 input_dim_sizes,
2817 inner_count,
2818 dtype_size,
2819 )
2820}
2821
2822macro_rules! cpu_index_add_wrapper {
2825 ($wrapper:ident, $kernel:path, $idx_ck:literal) => {
2826 pub(crate) fn $wrapper(
2829 inputs: &[Arc<RwLock<Storage>>],
2830 outputs: &mut [Arc<RwLock<Storage>>],
2831 _layouts: &[Layout],
2832 params: &OpParams,
2833 ) -> Result<()> {
2834 if inputs.len() != 3 || outputs.len() != 1 {
2835 return Err(Error::Msg(format!(
2836 "{}: expects 3 inputs + 1 output, got {} + {}",
2837 $idx_ck,
2838 inputs.len(),
2839 outputs.len(),
2840 ))
2841 .bt());
2842 }
2843 let (outer_count, base_dim_size, n_indices, inner_count) = match params {
2844 OpParams::IndexAdd {
2845 outer_count,
2846 base_dim_size,
2847 n_indices,
2848 inner_count,
2849 } => (*outer_count, *base_dim_size, *n_indices, *inner_count),
2850 other => {
2851 return Err(Error::Msg(format!(
2852 "{}: expects OpParams::IndexAdd, got {other:?}",
2853 $idx_ck,
2854 ))
2855 .bt())
2856 }
2857 };
2858 let base_guard = read_storage(&inputs[0])?;
2859 let idx_guard = read_storage(&inputs[1])?;
2860 if idx_guard.dtype != DType::U32 {
2861 return Err(Error::Msg(format!("{}: indices must be U32", $idx_ck,)).bt());
2862 }
2863 let src_guard = read_storage(&inputs[2])?;
2864 let mut out_guard = write_storage(&outputs[0])?;
2865 let base_cpu = cpu_input(&base_guard)?;
2866 let idx_cpu = cpu_input(&idx_guard)?;
2867 let src_cpu = cpu_input(&src_guard)?;
2868 let out_cpu = cpu_output(&mut out_guard)?;
2869 $kernel(
2870 base_cpu,
2871 idx_cpu,
2872 src_cpu,
2873 out_cpu,
2874 outer_count,
2875 base_dim_size,
2876 n_indices,
2877 inner_count,
2878 )
2879 }
2880 };
2881}
2882
2883cpu_index_add_wrapper!(
2884 index_add_f64_cpu_wrapper,
2885 fuel_cpu_backend::byte_kernels::index_add_f64,
2886 "index_add_f64"
2887);
2888cpu_index_add_wrapper!(
2889 index_add_bf16_cpu_wrapper,
2890 fuel_cpu_backend::byte_kernels::index_add_bf16,
2891 "index_add_bf16"
2892);
2893cpu_index_add_wrapper!(
2894 index_add_f16_cpu_wrapper,
2895 fuel_cpu_backend::byte_kernels::index_add_f16,
2896 "index_add_f16"
2897);
2898
2899macro_rules! cpu_scatter_add_wrapper {
2901 ($wrapper:ident, $kernel:path, $name_str:literal) => {
2902 pub(crate) fn $wrapper(
2904 inputs: &[Arc<RwLock<Storage>>],
2905 outputs: &mut [Arc<RwLock<Storage>>],
2906 _layouts: &[Layout],
2907 params: &OpParams,
2908 ) -> Result<()> {
2909 if inputs.len() != 3 || outputs.len() != 1 {
2910 return Err(Error::Msg(format!(
2911 "{}: expects 3 inputs + 1 output, got {} + {}",
2912 $name_str,
2913 inputs.len(),
2914 outputs.len(),
2915 ))
2916 .bt());
2917 }
2918 let (base_shape, src_shape, dim) = match params {
2919 OpParams::ScatterAdd {
2920 base_shape,
2921 src_shape,
2922 dim,
2923 } => (base_shape, src_shape, *dim),
2924 other => {
2925 return Err(Error::Msg(format!(
2926 "{}: expects OpParams::ScatterAdd, got {other:?}",
2927 $name_str,
2928 ))
2929 .bt())
2930 }
2931 };
2932 let base_guard = read_storage(&inputs[0])?;
2933 let idx_guard = read_storage(&inputs[1])?;
2934 if idx_guard.dtype != DType::U32 {
2935 return Err(Error::Msg(format!("{}: indices must be U32", $name_str,)).bt());
2936 }
2937 let src_guard = read_storage(&inputs[2])?;
2938 let mut out_guard = write_storage(&outputs[0])?;
2939 let base_cpu = cpu_input(&base_guard)?;
2940 let idx_cpu = cpu_input(&idx_guard)?;
2941 let src_cpu = cpu_input(&src_guard)?;
2942 let out_cpu = cpu_output(&mut out_guard)?;
2943 $kernel(
2944 base_cpu, idx_cpu, src_cpu, out_cpu, base_shape, src_shape, dim,
2945 )
2946 }
2947 };
2948}
2949
2950cpu_scatter_add_wrapper!(
2951 scatter_add_f64_cpu_wrapper,
2952 fuel_cpu_backend::byte_kernels::scatter_add_f64,
2953 "scatter_add_f64"
2954);
2955cpu_scatter_add_wrapper!(
2956 scatter_add_bf16_cpu_wrapper,
2957 fuel_cpu_backend::byte_kernels::scatter_add_bf16,
2958 "scatter_add_bf16"
2959);
2960cpu_scatter_add_wrapper!(
2961 scatter_add_f16_cpu_wrapper,
2962 fuel_cpu_backend::byte_kernels::scatter_add_f16,
2963 "scatter_add_f16"
2964);
2965
2966macro_rules! cpu_affine_wrapper_native {
2970 ($wrapper:ident, $kernel:path, $T:ty) => {
2971 pub(crate) fn $wrapper(
2974 inputs: &[Arc<RwLock<Storage>>],
2975 outputs: &mut [Arc<RwLock<Storage>>],
2976 _layouts: &[Layout],
2977 params: &OpParams,
2978 ) -> Result<()> {
2979 if inputs.len() != 1 || outputs.len() != 1 {
2980 return Err(Error::Msg("affine: 1 input + 1 output".to_string()).bt());
2981 }
2982 let (mul, add) = match params {
2983 OpParams::Affine { mul, add } => (*mul as $T, *add as $T),
2984 other => return Err(Error::Msg(format!("affine: bad params {other:?}")).bt()),
2985 };
2986 let in_guard = read_storage(&inputs[0])?;
2987 let mut out_guard = write_storage(&outputs[0])?;
2988 let in_cpu = cpu_input(&in_guard)?;
2989 let out_cpu = cpu_output(&mut out_guard)?;
2990 $kernel(in_cpu, out_cpu, mul, add)
2991 }
2992 };
2993}
2994
2995cpu_affine_wrapper_native!(
2996 affine_f64_cpu_wrapper,
2997 fuel_cpu_backend::byte_kernels::affine_f64,
2998 f64
2999);
3000
3001macro_rules! cpu_affine_wrapper_half {
3002 ($wrapper:ident, $kernel:path) => {
3003 pub(crate) fn $wrapper(
3006 inputs: &[Arc<RwLock<Storage>>],
3007 outputs: &mut [Arc<RwLock<Storage>>],
3008 _layouts: &[Layout],
3009 params: &OpParams,
3010 ) -> Result<()> {
3011 if inputs.len() != 1 || outputs.len() != 1 {
3012 return Err(Error::Msg("affine: 1 input + 1 output".to_string()).bt());
3013 }
3014 let (mul, add) = match params {
3015 OpParams::Affine { mul, add } => (*mul as f32, *add as f32),
3016 other => return Err(Error::Msg(format!("affine: bad params {other:?}")).bt()),
3017 };
3018 let in_guard = read_storage(&inputs[0])?;
3019 let mut out_guard = write_storage(&outputs[0])?;
3020 let in_cpu = cpu_input(&in_guard)?;
3021 let out_cpu = cpu_output(&mut out_guard)?;
3022 $kernel(in_cpu, out_cpu, mul, add)
3023 }
3024 };
3025}
3026
3027cpu_affine_wrapper_half!(
3028 affine_bf16_cpu_wrapper,
3029 fuel_cpu_backend::byte_kernels::affine_bf16
3030);
3031cpu_affine_wrapper_half!(
3032 affine_f16_cpu_wrapper,
3033 fuel_cpu_backend::byte_kernels::affine_f16
3034);
3035
3036macro_rules! cpu_clamp_wrapper {
3037 ($wrapper:ident, $kernel:path, $T:ty) => {
3038 pub(crate) fn $wrapper(
3041 inputs: &[Arc<RwLock<Storage>>],
3042 outputs: &mut [Arc<RwLock<Storage>>],
3043 _layouts: &[Layout],
3044 params: &OpParams,
3045 ) -> Result<()> {
3046 if inputs.len() != 1 || outputs.len() != 1 {
3047 return Err(Error::Msg("clamp: 1 input + 1 output".to_string()).bt());
3048 }
3049 let (min, max) = match params {
3050 OpParams::Clamp { min, max } => (*min as $T, *max as $T),
3051 other => return Err(Error::Msg(format!("clamp: bad params {other:?}")).bt()),
3052 };
3053 let in_guard = read_storage(&inputs[0])?;
3054 let mut out_guard = write_storage(&outputs[0])?;
3055 let in_cpu = cpu_input(&in_guard)?;
3056 let out_cpu = cpu_output(&mut out_guard)?;
3057 $kernel(in_cpu, out_cpu, min, max)
3058 }
3059 };
3060}
3061
3062cpu_clamp_wrapper!(
3063 clamp_f64_cpu_wrapper,
3064 fuel_cpu_backend::byte_kernels::clamp_f64,
3065 f64
3066);
3067cpu_clamp_wrapper!(
3068 clamp_bf16_cpu_wrapper,
3069 fuel_cpu_backend::byte_kernels::clamp_bf16,
3070 f32
3071);
3072cpu_clamp_wrapper!(
3073 clamp_f16_cpu_wrapper,
3074 fuel_cpu_backend::byte_kernels::clamp_f16,
3075 f32
3076);
3077
3078macro_rules! cpu_powi_wrapper {
3079 ($wrapper:ident, $kernel:path) => {
3080 pub(crate) fn $wrapper(
3083 inputs: &[Arc<RwLock<Storage>>],
3084 outputs: &mut [Arc<RwLock<Storage>>],
3085 _layouts: &[Layout],
3086 params: &OpParams,
3087 ) -> Result<()> {
3088 if inputs.len() != 1 || outputs.len() != 1 {
3089 return Err(Error::Msg("powi: 1 input + 1 output".to_string()).bt());
3090 }
3091 let exp = match params {
3092 OpParams::PowI { exp } => *exp,
3093 other => return Err(Error::Msg(format!("powi: bad params {other:?}")).bt()),
3094 };
3095 let in_guard = read_storage(&inputs[0])?;
3096 let mut out_guard = write_storage(&outputs[0])?;
3097 let in_cpu = cpu_input(&in_guard)?;
3098 let out_cpu = cpu_output(&mut out_guard)?;
3099 $kernel(in_cpu, out_cpu, exp)
3100 }
3101 };
3102}
3103
3104cpu_powi_wrapper!(
3105 powi_f64_cpu_wrapper,
3106 fuel_cpu_backend::byte_kernels::powi_f64
3107);
3108cpu_powi_wrapper!(
3109 powi_bf16_cpu_wrapper,
3110 fuel_cpu_backend::byte_kernels::powi_bf16
3111);
3112cpu_powi_wrapper!(
3113 powi_f16_cpu_wrapper,
3114 fuel_cpu_backend::byte_kernels::powi_f16
3115);
3116
3117macro_rules! cpu_arg_dim_wrapper_typed {
3121 ($wrapper:ident, $kernel:path, $expected_in_dtype:expr, $op_name:literal) => {
3122 fn $wrapper(
3123 inputs: &[Arc<RwLock<Storage>>],
3124 outputs: &mut [Arc<RwLock<Storage>>],
3125 layouts: &[Layout],
3126 params: &OpParams,
3127 ) -> Result<()> {
3128 if inputs.len() != 1 || outputs.len() != 1 {
3129 return Err(Error::Msg(format!("{}: 1 input + 1 output", $op_name,)).bt());
3130 }
3131 let dims = match params {
3132 OpParams::Reduce { dims, .. } => dims,
3133 other => {
3134 return Err(Error::Msg(format!(
3135 "{}: expects OpParams::Reduce, got {other:?}",
3136 $op_name,
3137 ))
3138 .bt())
3139 }
3140 };
3141 if dims.len() != 1 {
3142 return Err(Error::Msg(format!(
3143 "{}: expects exactly 1 reduce dim, got {dims:?}",
3144 $op_name,
3145 ))
3146 .bt());
3147 }
3148 let dim = dims[0];
3149 let input_layout = layouts.first().ok_or_else(|| {
3150 Error::Msg(format!(
3151 "{}: layouts empty (executor must pass input layout at [0])",
3152 $op_name,
3153 ))
3154 .bt()
3155 })?;
3156 let in_guard = read_storage(&inputs[0])?;
3157 if in_guard.dtype != $expected_in_dtype {
3158 return Err(Error::Msg(format!(
3159 "{}: expects input dtype {:?}, got {:?}",
3160 $op_name, $expected_in_dtype, in_guard.dtype,
3161 ))
3162 .bt());
3163 }
3164 let mut out_guard = write_storage(&outputs[0])?;
3165 let in_cpu = cpu_input(&in_guard)?;
3166 let out_cpu = cpu_output(&mut out_guard)?;
3167 $kernel(in_cpu, out_cpu, input_layout.shape().dims(), dim)
3168 }
3169 };
3170}
3171
3172fn argmax_dim_u32_cpu_dispatch(
3176 inputs: &[Arc<RwLock<Storage>>],
3177 outputs: &mut [Arc<RwLock<Storage>>],
3178 _layouts: &[Layout],
3179 params: &OpParams,
3180) -> Result<()> {
3181 let in_dtype = read_storage(&inputs[0])?.dtype;
3182 match in_dtype {
3183 DType::F32 => argmax_dim_u32_cpu_wrapper(inputs, outputs, _layouts, params),
3184 DType::F64 => argmax_dim_f64_only_wrapper(inputs, outputs, _layouts, params),
3185 DType::BF16 => argmax_dim_bf16_only_wrapper(inputs, outputs, _layouts, params),
3186 DType::F16 => argmax_dim_f16_only_wrapper(inputs, outputs, _layouts, params),
3187 other => Err(Error::Msg(format!("argmax_dim: unsupported input dtype {other:?}",)).bt()),
3188 }
3189}
3190
3191fn argmin_dim_u32_cpu_dispatch(
3192 inputs: &[Arc<RwLock<Storage>>],
3193 outputs: &mut [Arc<RwLock<Storage>>],
3194 _layouts: &[Layout],
3195 params: &OpParams,
3196) -> Result<()> {
3197 let in_dtype = read_storage(&inputs[0])?.dtype;
3198 match in_dtype {
3199 DType::F32 => argmin_dim_u32_cpu_wrapper(inputs, outputs, _layouts, params),
3200 DType::F64 => argmin_dim_f64_only_wrapper(inputs, outputs, _layouts, params),
3201 DType::BF16 => argmin_dim_bf16_only_wrapper(inputs, outputs, _layouts, params),
3202 DType::F16 => argmin_dim_f16_only_wrapper(inputs, outputs, _layouts, params),
3203 other => Err(Error::Msg(format!("argmin_dim: unsupported input dtype {other:?}",)).bt()),
3204 }
3205}
3206
3207cpu_arg_dim_wrapper_typed!(
3208 argmax_dim_f64_only_wrapper,
3209 fuel_cpu_backend::byte_kernels::argmax_dim_f64,
3210 DType::F64,
3211 "argmax_dim_f64"
3212);
3213cpu_arg_dim_wrapper_typed!(
3214 argmin_dim_f64_only_wrapper,
3215 fuel_cpu_backend::byte_kernels::argmin_dim_f64,
3216 DType::F64,
3217 "argmin_dim_f64"
3218);
3219cpu_arg_dim_wrapper_typed!(
3220 argmax_dim_bf16_only_wrapper,
3221 fuel_cpu_backend::byte_kernels::argmax_dim_bf16,
3222 DType::BF16,
3223 "argmax_dim_bf16"
3224);
3225cpu_arg_dim_wrapper_typed!(
3226 argmin_dim_bf16_only_wrapper,
3227 fuel_cpu_backend::byte_kernels::argmin_dim_bf16,
3228 DType::BF16,
3229 "argmin_dim_bf16"
3230);
3231cpu_arg_dim_wrapper_typed!(
3232 argmax_dim_f16_only_wrapper,
3233 fuel_cpu_backend::byte_kernels::argmax_dim_f16,
3234 DType::F16,
3235 "argmax_dim_f16"
3236);
3237cpu_arg_dim_wrapper_typed!(
3238 argmin_dim_f16_only_wrapper,
3239 fuel_cpu_backend::byte_kernels::argmin_dim_f16,
3240 DType::F16,
3241 "argmin_dim_f16"
3242);
3243
3244pub(crate) fn affine_f32_cpu_wrapper(
3247 inputs: &[Arc<RwLock<Storage>>],
3248 outputs: &mut [Arc<RwLock<Storage>>],
3249 _layouts: &[Layout],
3250 params: &OpParams,
3251) -> Result<()> {
3252 if inputs.len() != 1 || outputs.len() != 1 {
3253 return Err(Error::Msg(format!(
3254 "affine wrapper expects 1 input + 1 output, got {} + {}",
3255 inputs.len(),
3256 outputs.len(),
3257 ))
3258 .bt());
3259 }
3260 let (mul, add) = match params {
3261 OpParams::Affine { mul, add } => (*mul as f32, *add as f32),
3262 other => {
3263 return Err(Error::Msg(format!(
3264 "affine wrapper expects OpParams::Affine, got {other:?}",
3265 ))
3266 .bt());
3267 }
3268 };
3269 let in_guard = read_storage(&inputs[0])?;
3270 let mut out_guard = write_storage(&outputs[0])?;
3271 let in_cpu = cpu_input(&in_guard)?;
3272 let out_cpu = cpu_output(&mut out_guard)?;
3273 fuel_cpu_backend::byte_kernels::affine_f32(in_cpu, out_cpu, mul, add)
3274}
3275
3276macro_rules! cpu_affine_inplace_wrapper {
3281 ($name:ident, $kernel:ident, $T:ty) => {
3282 pub(crate) fn $name(
3283 inputs: &[Arc<RwLock<Storage>>],
3284 outputs: &mut [Arc<RwLock<Storage>>],
3285 _layouts: &[Layout],
3286 params: &OpParams,
3287 ) -> Result<()> {
3288 if !inputs.is_empty() || outputs.len() != 1 {
3289 return Err(Error::Msg(format!(
3290 concat!(
3291 stringify!($name),
3292 ": expected 0 inputs + 1 output (target adopted by executor), got {} + {}"
3293 ),
3294 inputs.len(),
3295 outputs.len(),
3296 ))
3297 .bt());
3298 }
3299 let (mul, add) = match params {
3300 OpParams::Affine { mul, add } => (*mul as $T, *add as $T),
3301 other => {
3302 return Err(Error::Msg(format!(
3303 concat!(stringify!($name), ": expected OpParams::Affine, got {:?}"),
3304 other,
3305 ))
3306 .bt())
3307 }
3308 };
3309 let mut out_guard = write_storage(&outputs[0])?;
3310 let out_cpu = cpu_output(&mut out_guard)?;
3311 fuel_cpu_backend::byte_kernels::$kernel(out_cpu, mul, add)
3312 }
3313 };
3314 ($name:ident, $kernel:ident, $T:ty, half) => {
3315 pub(crate) fn $name(
3316 inputs: &[Arc<RwLock<Storage>>],
3317 outputs: &mut [Arc<RwLock<Storage>>],
3318 _layouts: &[Layout],
3319 params: &OpParams,
3320 ) -> Result<()> {
3321 if !inputs.is_empty() || outputs.len() != 1 {
3322 return Err(Error::Msg(format!(
3323 concat!(
3324 stringify!($name),
3325 ": expected 0 inputs + 1 output (target adopted by executor), got {} + {}"
3326 ),
3327 inputs.len(),
3328 outputs.len(),
3329 ))
3330 .bt());
3331 }
3332 let (mul, add) = match params {
3333 OpParams::Affine { mul, add } => (*mul, *add),
3334 other => {
3335 return Err(Error::Msg(format!(
3336 concat!(stringify!($name), ": expected OpParams::Affine, got {:?}"),
3337 other,
3338 ))
3339 .bt())
3340 }
3341 };
3342 let mut out_guard = write_storage(&outputs[0])?;
3343 let out_cpu = cpu_output(&mut out_guard)?;
3344 fuel_cpu_backend::byte_kernels::$kernel(out_cpu, mul, add)
3345 }
3346 };
3347}
3348
3349cpu_affine_inplace_wrapper!(inplace_affine_f32_cpu_wrapper, affine_inplace_f32, f32);
3350cpu_affine_inplace_wrapper!(inplace_affine_f64_cpu_wrapper, affine_inplace_f64, f64);
3351cpu_affine_inplace_wrapper!(
3352 inplace_affine_bf16_cpu_wrapper,
3353 affine_inplace_bf16,
3354 f64,
3355 half
3356);
3357cpu_affine_inplace_wrapper!(
3358 inplace_affine_f16_cpu_wrapper,
3359 affine_inplace_f16,
3360 f64,
3361 half
3362);
3363
3364macro_rules! cpu_clamp_inplace_wrapper {
3368 ($name:ident, $kernel:ident, $T:ty) => {
3369 pub(crate) fn $name(
3370 inputs: &[Arc<RwLock<Storage>>],
3371 outputs: &mut [Arc<RwLock<Storage>>],
3372 _layouts: &[Layout],
3373 params: &OpParams,
3374 ) -> Result<()> {
3375 if !inputs.is_empty() || outputs.len() != 1 {
3376 return Err(Error::Msg(format!(
3377 concat!(
3378 stringify!($name),
3379 ": expected 0 inputs + 1 output (target adopted by executor), got {} + {}"
3380 ),
3381 inputs.len(),
3382 outputs.len(),
3383 ))
3384 .bt());
3385 }
3386 let (min, max) = match params {
3387 OpParams::Clamp { min, max } => (*min as $T, *max as $T),
3388 other => {
3389 return Err(Error::Msg(format!(
3390 concat!(stringify!($name), ": expected OpParams::Clamp, got {:?}"),
3391 other,
3392 ))
3393 .bt())
3394 }
3395 };
3396 let mut out_guard = write_storage(&outputs[0])?;
3397 let out_cpu = cpu_output(&mut out_guard)?;
3398 fuel_cpu_backend::byte_kernels::$kernel(out_cpu, min, max)
3399 }
3400 };
3401 ($name:ident, $kernel:ident, $T:ty, half) => {
3402 pub(crate) fn $name(
3403 inputs: &[Arc<RwLock<Storage>>],
3404 outputs: &mut [Arc<RwLock<Storage>>],
3405 _layouts: &[Layout],
3406 params: &OpParams,
3407 ) -> Result<()> {
3408 if !inputs.is_empty() || outputs.len() != 1 {
3409 return Err(Error::Msg(format!(
3410 concat!(
3411 stringify!($name),
3412 ": expected 0 inputs + 1 output (target adopted by executor), got {} + {}"
3413 ),
3414 inputs.len(),
3415 outputs.len(),
3416 ))
3417 .bt());
3418 }
3419 let (min, max) = match params {
3420 OpParams::Clamp { min, max } => (*min, *max),
3421 other => {
3422 return Err(Error::Msg(format!(
3423 concat!(stringify!($name), ": expected OpParams::Clamp, got {:?}"),
3424 other,
3425 ))
3426 .bt())
3427 }
3428 };
3429 let mut out_guard = write_storage(&outputs[0])?;
3430 let out_cpu = cpu_output(&mut out_guard)?;
3431 fuel_cpu_backend::byte_kernels::$kernel(out_cpu, min, max)
3432 }
3433 };
3434}
3435
3436cpu_clamp_inplace_wrapper!(clamp_inplace_f32_cpu_wrapper, clamp_inplace_f32, f32);
3437cpu_clamp_inplace_wrapper!(clamp_inplace_f64_cpu_wrapper, clamp_inplace_f64, f64);
3438cpu_clamp_inplace_wrapper!(
3439 clamp_inplace_bf16_cpu_wrapper,
3440 clamp_inplace_bf16,
3441 f64,
3442 half
3443);
3444cpu_clamp_inplace_wrapper!(clamp_inplace_f16_cpu_wrapper, clamp_inplace_f16, f64, half);
3445
3446macro_rules! cpu_powi_inplace_wrapper {
3449 ($name:ident, $kernel:ident) => {
3450 pub(crate) fn $name(
3451 inputs: &[Arc<RwLock<Storage>>],
3452 outputs: &mut [Arc<RwLock<Storage>>],
3453 _layouts: &[Layout],
3454 params: &OpParams,
3455 ) -> Result<()> {
3456 if !inputs.is_empty() || outputs.len() != 1 {
3457 return Err(Error::Msg(format!(
3458 concat!(
3459 stringify!($name),
3460 ": expected 0 inputs + 1 output (target adopted by executor), got {} + {}"
3461 ),
3462 inputs.len(),
3463 outputs.len(),
3464 ))
3465 .bt());
3466 }
3467 let exp = match params {
3468 OpParams::PowI { exp } => *exp,
3469 other => {
3470 return Err(Error::Msg(format!(
3471 concat!(stringify!($name), ": expected OpParams::PowI, got {:?}"),
3472 other,
3473 ))
3474 .bt())
3475 }
3476 };
3477 let mut out_guard = write_storage(&outputs[0])?;
3478 let out_cpu = cpu_output(&mut out_guard)?;
3479 fuel_cpu_backend::byte_kernels::$kernel(out_cpu, exp)
3480 }
3481 };
3482}
3483
3484cpu_powi_inplace_wrapper!(powi_inplace_f32_cpu_wrapper, powi_inplace_f32);
3485cpu_powi_inplace_wrapper!(powi_inplace_f64_cpu_wrapper, powi_inplace_f64);
3486cpu_powi_inplace_wrapper!(powi_inplace_bf16_cpu_wrapper, powi_inplace_bf16);
3487cpu_powi_inplace_wrapper!(powi_inplace_f16_cpu_wrapper, powi_inplace_f16);
3488
3489macro_rules! cpu_unary_inplace_wrapper {
3494 ($name:ident, $kernel:ident) => {
3495 pub(crate) fn $name(
3496 inputs: &[Arc<RwLock<Storage>>],
3497 outputs: &mut [Arc<RwLock<Storage>>],
3498 _layouts: &[Layout],
3499 _params: &OpParams,
3500 ) -> Result<()> {
3501 if !inputs.is_empty() || outputs.len() != 1 {
3502 return Err(Error::Msg(format!(
3503 concat!(
3504 stringify!($name),
3505 ": expected 0 inputs + 1 output (target adopted by executor), got {} + {}"
3506 ),
3507 inputs.len(),
3508 outputs.len(),
3509 ))
3510 .bt());
3511 }
3512 let mut out_guard = write_storage(&outputs[0])?;
3513 let out_cpu = cpu_output(&mut out_guard)?;
3514 fuel_cpu_backend::byte_kernels::$kernel(out_cpu)
3515 }
3516 };
3517}
3518
3519cpu_unary_inplace_wrapper!(relu_inplace_f32_cpu_wrapper, relu_inplace_f32);
3520cpu_unary_inplace_wrapper!(silu_inplace_f32_cpu_wrapper, silu_inplace_f32);
3521cpu_unary_inplace_wrapper!(gelu_inplace_f32_cpu_wrapper, gelu_inplace_f32);
3522cpu_unary_inplace_wrapper!(tanh_inplace_f32_cpu_wrapper, tanh_inplace_f32);
3523cpu_unary_inplace_wrapper!(sigmoid_inplace_f32_cpu_wrapper, sigmoid_inplace_f32);
3524
3525cpu_unary_inplace_wrapper!(relu_inplace_f64_cpu_wrapper, relu_inplace_f64);
3526cpu_unary_inplace_wrapper!(silu_inplace_f64_cpu_wrapper, silu_inplace_f64);
3527cpu_unary_inplace_wrapper!(gelu_inplace_f64_cpu_wrapper, gelu_inplace_f64);
3528cpu_unary_inplace_wrapper!(tanh_inplace_f64_cpu_wrapper, tanh_inplace_f64);
3529cpu_unary_inplace_wrapper!(sigmoid_inplace_f64_cpu_wrapper, sigmoid_inplace_f64);
3530
3531cpu_unary_inplace_wrapper!(relu_inplace_bf16_cpu_wrapper, relu_inplace_bf16);
3532cpu_unary_inplace_wrapper!(silu_inplace_bf16_cpu_wrapper, silu_inplace_bf16);
3533cpu_unary_inplace_wrapper!(gelu_inplace_bf16_cpu_wrapper, gelu_inplace_bf16);
3534cpu_unary_inplace_wrapper!(tanh_inplace_bf16_cpu_wrapper, tanh_inplace_bf16);
3535cpu_unary_inplace_wrapper!(sigmoid_inplace_bf16_cpu_wrapper, sigmoid_inplace_bf16);
3536
3537cpu_unary_inplace_wrapper!(relu_inplace_f16_cpu_wrapper, relu_inplace_f16);
3538cpu_unary_inplace_wrapper!(silu_inplace_f16_cpu_wrapper, silu_inplace_f16);
3539cpu_unary_inplace_wrapper!(gelu_inplace_f16_cpu_wrapper, gelu_inplace_f16);
3540cpu_unary_inplace_wrapper!(tanh_inplace_f16_cpu_wrapper, tanh_inplace_f16);
3541cpu_unary_inplace_wrapper!(sigmoid_inplace_f16_cpu_wrapper, sigmoid_inplace_f16);
3542
3543cpu_unary_inplace_wrapper!(neg_inplace_f32_cpu_wrapper, neg_inplace_f32);
3547cpu_unary_inplace_wrapper!(neg_inplace_f64_cpu_wrapper, neg_inplace_f64);
3548cpu_unary_inplace_wrapper!(neg_inplace_bf16_cpu_wrapper, neg_inplace_bf16);
3549cpu_unary_inplace_wrapper!(neg_inplace_f16_cpu_wrapper, neg_inplace_f16);
3550
3551cpu_unary_inplace_wrapper!(abs_inplace_f32_cpu_wrapper, abs_inplace_f32);
3552cpu_unary_inplace_wrapper!(abs_inplace_f64_cpu_wrapper, abs_inplace_f64);
3553cpu_unary_inplace_wrapper!(abs_inplace_bf16_cpu_wrapper, abs_inplace_bf16);
3554cpu_unary_inplace_wrapper!(abs_inplace_f16_cpu_wrapper, abs_inplace_f16);
3555
3556cpu_unary_inplace_wrapper!(sqr_inplace_f32_cpu_wrapper, sqr_inplace_f32);
3557cpu_unary_inplace_wrapper!(sqr_inplace_f64_cpu_wrapper, sqr_inplace_f64);
3558cpu_unary_inplace_wrapper!(sqr_inplace_bf16_cpu_wrapper, sqr_inplace_bf16);
3559cpu_unary_inplace_wrapper!(sqr_inplace_f16_cpu_wrapper, sqr_inplace_f16);
3560
3561cpu_unary_inplace_wrapper!(sqrt_inplace_f32_cpu_wrapper, sqrt_inplace_f32);
3562cpu_unary_inplace_wrapper!(sqrt_inplace_f64_cpu_wrapper, sqrt_inplace_f64);
3563cpu_unary_inplace_wrapper!(sqrt_inplace_bf16_cpu_wrapper, sqrt_inplace_bf16);
3564cpu_unary_inplace_wrapper!(sqrt_inplace_f16_cpu_wrapper, sqrt_inplace_f16);
3565
3566cpu_unary_inplace_wrapper!(rsqrt_inplace_f32_cpu_wrapper, rsqrt_inplace_f32);
3567cpu_unary_inplace_wrapper!(rsqrt_inplace_f64_cpu_wrapper, rsqrt_inplace_f64);
3568cpu_unary_inplace_wrapper!(rsqrt_inplace_bf16_cpu_wrapper, rsqrt_inplace_bf16);
3569cpu_unary_inplace_wrapper!(rsqrt_inplace_f16_cpu_wrapper, rsqrt_inplace_f16);
3570
3571cpu_unary_inplace_wrapper!(recip_inplace_f32_cpu_wrapper, recip_inplace_f32);
3572cpu_unary_inplace_wrapper!(recip_inplace_f64_cpu_wrapper, recip_inplace_f64);
3573cpu_unary_inplace_wrapper!(recip_inplace_bf16_cpu_wrapper, recip_inplace_bf16);
3574cpu_unary_inplace_wrapper!(recip_inplace_f16_cpu_wrapper, recip_inplace_f16);
3575
3576cpu_unary_inplace_wrapper!(exp_inplace_f32_cpu_wrapper, exp_inplace_f32);
3577cpu_unary_inplace_wrapper!(exp_inplace_f64_cpu_wrapper, exp_inplace_f64);
3578cpu_unary_inplace_wrapper!(exp_inplace_bf16_cpu_wrapper, exp_inplace_bf16);
3579cpu_unary_inplace_wrapper!(exp_inplace_f16_cpu_wrapper, exp_inplace_f16);
3580
3581cpu_unary_inplace_wrapper!(log_inplace_f32_cpu_wrapper, log_inplace_f32);
3582cpu_unary_inplace_wrapper!(log_inplace_f64_cpu_wrapper, log_inplace_f64);
3583cpu_unary_inplace_wrapper!(log_inplace_bf16_cpu_wrapper, log_inplace_bf16);
3584cpu_unary_inplace_wrapper!(log_inplace_f16_cpu_wrapper, log_inplace_f16);
3585
3586cpu_unary_inplace_wrapper!(sin_inplace_f32_cpu_wrapper, sin_inplace_f32);
3587cpu_unary_inplace_wrapper!(sin_inplace_f64_cpu_wrapper, sin_inplace_f64);
3588cpu_unary_inplace_wrapper!(sin_inplace_bf16_cpu_wrapper, sin_inplace_bf16);
3589cpu_unary_inplace_wrapper!(sin_inplace_f16_cpu_wrapper, sin_inplace_f16);
3590
3591cpu_unary_inplace_wrapper!(cos_inplace_f32_cpu_wrapper, cos_inplace_f32);
3592cpu_unary_inplace_wrapper!(cos_inplace_f64_cpu_wrapper, cos_inplace_f64);
3593cpu_unary_inplace_wrapper!(cos_inplace_bf16_cpu_wrapper, cos_inplace_bf16);
3594cpu_unary_inplace_wrapper!(cos_inplace_f16_cpu_wrapper, cos_inplace_f16);
3595
3596cpu_unary_inplace_wrapper!(sign_inplace_f32_cpu_wrapper, sign_inplace_f32);
3597cpu_unary_inplace_wrapper!(sign_inplace_f64_cpu_wrapper, sign_inplace_f64);
3598cpu_unary_inplace_wrapper!(sign_inplace_bf16_cpu_wrapper, sign_inplace_bf16);
3599cpu_unary_inplace_wrapper!(sign_inplace_f16_cpu_wrapper, sign_inplace_f16);
3600
3601cpu_unary_inplace_wrapper!(floor_inplace_f32_cpu_wrapper, floor_inplace_f32);
3602cpu_unary_inplace_wrapper!(floor_inplace_f64_cpu_wrapper, floor_inplace_f64);
3603cpu_unary_inplace_wrapper!(floor_inplace_bf16_cpu_wrapper, floor_inplace_bf16);
3604cpu_unary_inplace_wrapper!(floor_inplace_f16_cpu_wrapper, floor_inplace_f16);
3605
3606cpu_unary_inplace_wrapper!(ceil_inplace_f32_cpu_wrapper, ceil_inplace_f32);
3607cpu_unary_inplace_wrapper!(ceil_inplace_f64_cpu_wrapper, ceil_inplace_f64);
3608cpu_unary_inplace_wrapper!(ceil_inplace_bf16_cpu_wrapper, ceil_inplace_bf16);
3609cpu_unary_inplace_wrapper!(ceil_inplace_f16_cpu_wrapper, ceil_inplace_f16);
3610
3611cpu_unary_inplace_wrapper!(round_inplace_f32_cpu_wrapper, round_inplace_f32);
3612cpu_unary_inplace_wrapper!(round_inplace_f64_cpu_wrapper, round_inplace_f64);
3613cpu_unary_inplace_wrapper!(round_inplace_bf16_cpu_wrapper, round_inplace_bf16);
3614cpu_unary_inplace_wrapper!(round_inplace_f16_cpu_wrapper, round_inplace_f16);
3615
3616cpu_unary_inplace_wrapper!(erf_inplace_f32_cpu_wrapper, erf_inplace_f32);
3617cpu_unary_inplace_wrapper!(erf_inplace_f64_cpu_wrapper, erf_inplace_f64);
3618cpu_unary_inplace_wrapper!(erf_inplace_bf16_cpu_wrapper, erf_inplace_bf16);
3619cpu_unary_inplace_wrapper!(erf_inplace_f16_cpu_wrapper, erf_inplace_f16);
3620
3621cpu_unary_inplace_wrapper!(gelu_erf_inplace_f32_cpu_wrapper, gelu_erf_inplace_f32);
3622cpu_unary_inplace_wrapper!(gelu_erf_inplace_f64_cpu_wrapper, gelu_erf_inplace_f64);
3623cpu_unary_inplace_wrapper!(gelu_erf_inplace_bf16_cpu_wrapper, gelu_erf_inplace_bf16);
3624cpu_unary_inplace_wrapper!(gelu_erf_inplace_f16_cpu_wrapper, gelu_erf_inplace_f16);
3625
3626pub(crate) fn clamp_elementwise_f32_cpu_wrapper(
3628 inputs: &[Arc<RwLock<Storage>>],
3629 outputs: &mut [Arc<RwLock<Storage>>],
3630 _layouts: &[Layout],
3631 params: &OpParams,
3632) -> Result<()> {
3633 if inputs.len() != 1 || outputs.len() != 1 {
3634 return Err(Error::Msg(format!(
3635 "clamp wrapper expects 1 input + 1 output, got {} + {}",
3636 inputs.len(),
3637 outputs.len(),
3638 ))
3639 .bt());
3640 }
3641 let (min, max) = match params {
3642 OpParams::Clamp { min, max } => (*min as f32, *max as f32),
3643 other => {
3644 return Err(Error::Msg(format!(
3645 "clamp wrapper expects OpParams::Clamp, got {other:?}",
3646 ))
3647 .bt());
3648 }
3649 };
3650 let in_guard = read_storage(&inputs[0])?;
3651 let mut out_guard = write_storage(&outputs[0])?;
3652 let in_cpu = cpu_input(&in_guard)?;
3653 let out_cpu = cpu_output(&mut out_guard)?;
3654 fuel_cpu_backend::byte_kernels::clamp_f32(in_cpu, out_cpu, min, max)
3655}
3656
3657pub(crate) fn powi_elementwise_f32_cpu_wrapper(
3659 inputs: &[Arc<RwLock<Storage>>],
3660 outputs: &mut [Arc<RwLock<Storage>>],
3661 _layouts: &[Layout],
3662 params: &OpParams,
3663) -> Result<()> {
3664 if inputs.len() != 1 || outputs.len() != 1 {
3665 return Err(Error::Msg(format!(
3666 "powi wrapper expects 1 input + 1 output, got {} + {}",
3667 inputs.len(),
3668 outputs.len(),
3669 ))
3670 .bt());
3671 }
3672 let exp = match params {
3673 OpParams::PowI { exp } => *exp,
3674 other => {
3675 return Err(Error::Msg(format!(
3676 "powi wrapper expects OpParams::PowI, got {other:?}",
3677 ))
3678 .bt());
3679 }
3680 };
3681 let in_guard = read_storage(&inputs[0])?;
3682 let mut out_guard = write_storage(&outputs[0])?;
3683 let in_cpu = cpu_input(&in_guard)?;
3684 let out_cpu = cpu_output(&mut out_guard)?;
3685 fuel_cpu_backend::byte_kernels::powi_f32(in_cpu, out_cpu, exp)
3686}
3687
3688macro_rules! cpu_powi_backward_wrapper {
3691 ($wrapper:ident, $kernel:path, $op_name:literal) => {
3692 pub(crate) fn $wrapper(
3695 inputs: &[Arc<RwLock<Storage>>],
3696 outputs: &mut [Arc<RwLock<Storage>>],
3697 _layouts: &[Layout],
3698 params: &OpParams,
3699 ) -> Result<()> {
3700 if inputs.len() != 2 || outputs.len() != 1 {
3701 return Err(Error::Msg(format!(
3702 "{} wrapper expects 2 inputs + 1 output, got {} + {}",
3703 $op_name,
3704 inputs.len(),
3705 outputs.len(),
3706 ))
3707 .bt());
3708 }
3709 let exp = match params {
3710 OpParams::PowI { exp } => *exp,
3711 other => {
3712 return Err(Error::Msg(format!(
3713 "{} wrapper expects OpParams::PowI, got {other:?}",
3714 $op_name,
3715 ))
3716 .bt())
3717 }
3718 };
3719 let x_guard = read_storage(&inputs[0])?;
3720 let up_guard = read_storage(&inputs[1])?;
3721 let mut out_guard = write_storage(&outputs[0])?;
3722 let x_cpu = cpu_input(&x_guard)?;
3723 let up_cpu = cpu_input(&up_guard)?;
3724 let out_cpu = cpu_output(&mut out_guard)?;
3725 $kernel(x_cpu, up_cpu, out_cpu, exp)
3726 }
3727 };
3728}
3729
3730cpu_powi_backward_wrapper!(
3731 powi_backward_f32_cpu_wrapper,
3732 fuel_cpu_backend::byte_kernels::powi_backward_f32,
3733 "powi_backward_f32"
3734);
3735cpu_powi_backward_wrapper!(
3736 powi_backward_f64_cpu_wrapper,
3737 fuel_cpu_backend::byte_kernels::powi_backward_f64,
3738 "powi_backward_f64"
3739);
3740cpu_powi_backward_wrapper!(
3741 powi_backward_bf16_cpu_wrapper,
3742 fuel_cpu_backend::byte_kernels::powi_backward_bf16,
3743 "powi_backward_bf16"
3744);
3745cpu_powi_backward_wrapper!(
3746 powi_backward_f16_cpu_wrapper,
3747 fuel_cpu_backend::byte_kernels::powi_backward_f16,
3748 "powi_backward_f16"
3749);
3750
3751macro_rules! cpu_reduce_wrapper {
3756 ($wrapper:ident, $kernel:path, $op_name:literal) => {
3757 pub(crate) fn $wrapper(
3758 inputs: &[Arc<RwLock<Storage>>],
3759 outputs: &mut [Arc<RwLock<Storage>>],
3760 layouts: &[Layout],
3761 params: &OpParams,
3762 ) -> Result<()> {
3763 if inputs.len() != 1 {
3764 return Err(Error::Msg(format!(
3765 "{} wrapper expects 1 input, got {}",
3766 $op_name,
3767 inputs.len(),
3768 ))
3769 .bt());
3770 }
3771 if outputs.len() != 1 {
3772 return Err(Error::Msg(format!(
3773 "{} wrapper expects 1 output, got {}",
3774 $op_name,
3775 outputs.len(),
3776 ))
3777 .bt());
3778 }
3779 let dims = match params {
3780 OpParams::Reduce { dims, .. } => dims,
3781 other => {
3782 return Err(Error::Msg(format!(
3783 "{} wrapper expects OpParams::Reduce, got {:?}",
3784 $op_name, other,
3785 ))
3786 .bt())
3787 }
3788 };
3789 let input_layout = layouts.first().ok_or_else(|| {
3790 Error::Msg(format!(
3791 "{} wrapper: layouts empty (executor must pass input layout at [0])",
3792 $op_name,
3793 ))
3794 .bt()
3795 })?;
3796 let input_shape: &[usize] = input_layout.shape().dims();
3803 let in_guard = read_storage(&inputs[0])?;
3804 let mut out_guard = write_storage(&outputs[0])?;
3805 let in_cpu = cpu_input(&in_guard)?;
3806 let out_cpu = cpu_output(&mut out_guard)?;
3807 $kernel(in_cpu, out_cpu, input_shape, dims)
3808 }
3809 };
3810}
3811
3812cpu_reduce_wrapper!(
3813 sum_reduce_f32_cpu_wrapper,
3814 fuel_cpu_backend::byte_kernels::sum_reduce_f32,
3815 "sum_reduce"
3816);
3817cpu_reduce_wrapper!(
3818 max_reduce_f32_cpu_wrapper,
3819 fuel_cpu_backend::byte_kernels::max_reduce_f32,
3820 "max_reduce"
3821);
3822cpu_reduce_wrapper!(
3823 min_reduce_f32_cpu_wrapper,
3824 fuel_cpu_backend::byte_kernels::min_reduce_f32,
3825 "min_reduce"
3826);
3827cpu_reduce_wrapper!(
3828 mean_reduce_f32_cpu_wrapper,
3829 fuel_cpu_backend::byte_kernels::mean_reduce_f32,
3830 "mean_reduce"
3831);
3832cpu_reduce_wrapper!(
3833 sum_reduce_f64_cpu_wrapper,
3834 fuel_cpu_backend::byte_kernels::sum_reduce_f64,
3835 "sum_reduce"
3836);
3837cpu_reduce_wrapper!(
3838 max_reduce_f64_cpu_wrapper,
3839 fuel_cpu_backend::byte_kernels::max_reduce_f64,
3840 "max_reduce"
3841);
3842cpu_reduce_wrapper!(
3843 min_reduce_f64_cpu_wrapper,
3844 fuel_cpu_backend::byte_kernels::min_reduce_f64,
3845 "min_reduce"
3846);
3847cpu_reduce_wrapper!(
3848 mean_reduce_f64_cpu_wrapper,
3849 fuel_cpu_backend::byte_kernels::mean_reduce_f64,
3850 "mean_reduce"
3851);
3852
3853cpu_reduce_wrapper!(
3854 sum_reduce_bf16_cpu_wrapper,
3855 fuel_cpu_backend::byte_kernels::sum_reduce_bf16,
3856 "sum_reduce"
3857);
3858cpu_reduce_wrapper!(
3859 max_reduce_bf16_cpu_wrapper,
3860 fuel_cpu_backend::byte_kernels::max_reduce_bf16,
3861 "max_reduce"
3862);
3863cpu_reduce_wrapper!(
3864 min_reduce_bf16_cpu_wrapper,
3865 fuel_cpu_backend::byte_kernels::min_reduce_bf16,
3866 "min_reduce"
3867);
3868cpu_reduce_wrapper!(
3869 mean_reduce_bf16_cpu_wrapper,
3870 fuel_cpu_backend::byte_kernels::mean_reduce_bf16,
3871 "mean_reduce"
3872);
3873
3874cpu_reduce_wrapper!(
3875 sum_reduce_f16_cpu_wrapper,
3876 fuel_cpu_backend::byte_kernels::sum_reduce_f16,
3877 "sum_reduce"
3878);
3879cpu_reduce_wrapper!(
3880 max_reduce_f16_cpu_wrapper,
3881 fuel_cpu_backend::byte_kernels::max_reduce_f16,
3882 "max_reduce"
3883);
3884cpu_reduce_wrapper!(
3885 min_reduce_f16_cpu_wrapper,
3886 fuel_cpu_backend::byte_kernels::min_reduce_f16,
3887 "min_reduce"
3888);
3889cpu_reduce_wrapper!(
3890 mean_reduce_f16_cpu_wrapper,
3891 fuel_cpu_backend::byte_kernels::mean_reduce_f16,
3892 "mean_reduce"
3893);
3894
3895macro_rules! cpu_cast_wrapper {
3901 (
3902 $wrapper:ident,
3903 $target_dtype:expr,
3904 $target_name:literal,
3905 { $( $source:pat => $kernel:path ),+ $(,)? } $(,)?
3906 ) => {
3907 pub(crate) fn $wrapper(
3910 inputs: &[Arc<RwLock<Storage>>],
3911 outputs: &mut [Arc<RwLock<Storage>>],
3912 _layouts: &[Layout],
3913 _params: &OpParams,
3914 ) -> Result<()> {
3915 if inputs.len() != 1 {
3916 return Err(Error::Msg(format!(
3917 "cast→{} wrapper expects 1 input, got {}",
3918 $target_name,
3919 inputs.len(),
3920 ))
3921 .bt());
3922 }
3923 if outputs.len() != 1 {
3924 return Err(Error::Msg(format!(
3925 "cast→{} wrapper expects 1 output, got {}",
3926 $target_name,
3927 outputs.len(),
3928 ))
3929 .bt());
3930 }
3931 let in_guard = read_storage(&inputs[0])?;
3932 let mut out_guard = write_storage(&outputs[0])?;
3933 let source_dtype = in_guard.dtype;
3934 let in_cpu = cpu_input(&in_guard)?;
3935 let out_cpu = cpu_output(&mut out_guard)?;
3936 match source_dtype {
3937 $( $source => $kernel(in_cpu, out_cpu), )+
3938 other => Err(Error::Msg(format!(
3939 "cast→{}: source dtype {:?} not yet wired (Phase C \
3940 extends the cast matrix as needed)",
3941 $target_name, other,
3942 ))
3943 .bt()),
3944 }
3945 }
3946 };
3947}
3948
3949macro_rules! cpu_conv2d_wrapper {
3955 ($name:ident, $kernel:path) => {
3956 pub(crate) fn $name(
3957 inputs: &[Arc<RwLock<Storage>>],
3958 outputs: &mut [Arc<RwLock<Storage>>],
3959 _layouts: &[Layout],
3960 params: &OpParams,
3961 ) -> Result<()> {
3962 if inputs.len() != 2 && inputs.len() != 3 {
3963 return Err(Error::Msg(format!(
3964 "conv2d wrapper expects 2 or 3 inputs (x, weight, [bias]), got {}",
3965 inputs.len(),
3966 ))
3967 .bt());
3968 }
3969 if outputs.len() != 1 {
3970 return Err(Error::Msg(format!(
3971 "conv2d wrapper expects 1 output, got {}",
3972 outputs.len(),
3973 ))
3974 .bt());
3975 }
3976 let (x_shape, w_shape, out_shape, stride, padding, dilation, groups) = match params {
3977 OpParams::Conv2D {
3978 x_shape,
3979 w_shape,
3980 out_shape,
3981 stride,
3982 padding,
3983 dilation,
3984 groups,
3985 } => (
3986 *x_shape, *w_shape, *out_shape, *stride, *padding, *dilation, *groups,
3987 ),
3988 other => {
3989 return Err(Error::Msg(format!(
3990 "conv2d wrapper expects OpParams::Conv2D, got {other:?}",
3991 ))
3992 .bt())
3993 }
3994 };
3995 let x_guard = read_storage(&inputs[0])?;
3996 let w_guard = read_storage(&inputs[1])?;
3997 let bias_guard = match inputs.get(2) {
3998 Some(arc) => Some(read_storage(arc)?),
3999 None => None,
4000 };
4001 let mut out_guard = write_storage(&outputs[0])?;
4002 let x_cpu = cpu_input(&x_guard)?;
4003 let w_cpu = cpu_input(&w_guard)?;
4004 let bias_cpu = match &bias_guard {
4005 Some(g) => Some(cpu_input(g)?),
4006 None => None,
4007 };
4008 let out_cpu = cpu_output(&mut out_guard)?;
4009 $kernel(
4010 x_cpu, w_cpu, bias_cpu, out_cpu, x_shape, w_shape, out_shape, stride, padding,
4011 dilation, groups,
4012 )
4013 }
4014 };
4015}
4016
4017cpu_conv2d_wrapper!(
4018 conv2d_f32_cpu_wrapper,
4019 fuel_cpu_backend::byte_kernels::conv2d_f32
4020);
4021cpu_conv2d_wrapper!(
4022 conv2d_f64_cpu_wrapper,
4023 fuel_cpu_backend::byte_kernels::conv2d_f64
4024);
4025cpu_conv2d_wrapper!(
4026 conv2d_bf16_cpu_wrapper,
4027 fuel_cpu_backend::byte_kernels::conv2d_bf16
4028);
4029cpu_conv2d_wrapper!(
4030 conv2d_f16_cpu_wrapper,
4031 fuel_cpu_backend::byte_kernels::conv2d_f16
4032);
4033
4034macro_rules! cpu_conv_transpose2d_wrapper {
4040 ($name:ident, $kernel:path) => {
4041 pub(crate) fn $name(
4042 inputs: &[Arc<RwLock<Storage>>],
4043 outputs: &mut [Arc<RwLock<Storage>>],
4044 _layouts: &[Layout],
4045 params: &OpParams,
4046 ) -> Result<()> {
4047 if inputs.len() != 2 && inputs.len() != 3 {
4048 return Err(Error::Msg(format!(
4049 "conv_transpose2d wrapper expects 2 or 3 inputs (x, weight, [bias]), got {}",
4050 inputs.len(),
4051 ))
4052 .bt());
4053 }
4054 if outputs.len() != 1 {
4055 return Err(Error::Msg(format!(
4056 "conv_transpose2d wrapper expects 1 output, got {}",
4057 outputs.len(),
4058 ))
4059 .bt());
4060 }
4061 let (x_shape, w_shape, out_shape, stride, padding, dilation, groups) =
4062 match params {
4063 OpParams::ConvTranspose2D {
4064 x_shape,
4065 w_shape,
4066 out_shape,
4067 stride,
4068 padding,
4069 output_padding: _,
4070 dilation,
4071 groups,
4072 } => (
4073 *x_shape, *w_shape, *out_shape, *stride, *padding, *dilation, *groups,
4074 ),
4075 other => return Err(Error::Msg(format!(
4076 "conv_transpose2d wrapper expects OpParams::ConvTranspose2D, got {other:?}",
4077 ))
4078 .bt()),
4079 };
4080 let x_guard = read_storage(&inputs[0])?;
4081 let w_guard = read_storage(&inputs[1])?;
4082 let bias_guard = match inputs.get(2) {
4083 Some(arc) => Some(read_storage(arc)?),
4084 None => None,
4085 };
4086 let mut out_guard = write_storage(&outputs[0])?;
4087 let x_cpu = cpu_input(&x_guard)?;
4088 let w_cpu = cpu_input(&w_guard)?;
4089 let bias_cpu = match &bias_guard {
4090 Some(g) => Some(cpu_input(g)?),
4091 None => None,
4092 };
4093 let out_cpu = cpu_output(&mut out_guard)?;
4094 $kernel(
4095 x_cpu, w_cpu, bias_cpu, out_cpu, x_shape, w_shape, out_shape, stride, padding,
4096 dilation, groups,
4097 )
4098 }
4099 };
4100}
4101
4102cpu_conv_transpose2d_wrapper!(
4103 conv_transpose2d_f32_cpu_wrapper,
4104 fuel_cpu_backend::byte_kernels::conv_transpose2d_f32
4105);
4106cpu_conv_transpose2d_wrapper!(
4107 conv_transpose2d_f64_cpu_wrapper,
4108 fuel_cpu_backend::byte_kernels::conv_transpose2d_f64
4109);
4110cpu_conv_transpose2d_wrapper!(
4111 conv_transpose2d_bf16_cpu_wrapper,
4112 fuel_cpu_backend::byte_kernels::conv_transpose2d_bf16
4113);
4114cpu_conv_transpose2d_wrapper!(
4115 conv_transpose2d_f16_cpu_wrapper,
4116 fuel_cpu_backend::byte_kernels::conv_transpose2d_f16
4117);
4118
4119macro_rules! cpu_reduce_sum_to_wrapper {
4124 ($name:ident, $kernel:path) => {
4125 pub(crate) fn $name(
4126 inputs: &[Arc<RwLock<Storage>>],
4127 outputs: &mut [Arc<RwLock<Storage>>],
4128 _layouts: &[Layout],
4129 params: &OpParams,
4130 ) -> Result<()> {
4131 if inputs.len() != 1 {
4132 return Err(Error::Msg(format!(
4133 "reduce_sum_to wrapper expects 1 input, got {}",
4134 inputs.len(),
4135 ))
4136 .bt());
4137 }
4138 if outputs.len() != 1 {
4139 return Err(Error::Msg(format!(
4140 "reduce_sum_to wrapper expects 1 output, got {}",
4141 outputs.len(),
4142 ))
4143 .bt());
4144 }
4145 let (input_shape, output_shape) = match params {
4146 OpParams::ReduceSumTo {
4147 input_shape,
4148 output_shape,
4149 } => (input_shape.clone(), output_shape.clone()),
4150 other => {
4151 return Err(Error::Msg(format!(
4152 "reduce_sum_to wrapper expects OpParams::ReduceSumTo, got {other:?}",
4153 ))
4154 .bt())
4155 }
4156 };
4157 let in_guard = read_storage(&inputs[0])?;
4158 let mut out_guard = write_storage(&outputs[0])?;
4159 let in_cpu = cpu_input(&in_guard)?;
4160 let out_cpu = cpu_output(&mut out_guard)?;
4161 $kernel(in_cpu, out_cpu, &input_shape, &output_shape)
4162 }
4163 };
4164}
4165
4166cpu_reduce_sum_to_wrapper!(
4167 reduce_sum_to_f32_cpu_wrapper,
4168 fuel_cpu_backend::byte_kernels::reduce_sum_to_f32
4169);
4170cpu_reduce_sum_to_wrapper!(
4171 reduce_sum_to_f64_cpu_wrapper,
4172 fuel_cpu_backend::byte_kernels::reduce_sum_to_f64
4173);
4174cpu_reduce_sum_to_wrapper!(
4175 reduce_sum_to_bf16_cpu_wrapper,
4176 fuel_cpu_backend::byte_kernels::reduce_sum_to_bf16
4177);
4178cpu_reduce_sum_to_wrapper!(
4179 reduce_sum_to_f16_cpu_wrapper,
4180 fuel_cpu_backend::byte_kernels::reduce_sum_to_f16
4181);
4182
4183macro_rules! cpu_reduce_max_to_wrapper {
4188 ($name:ident, $kernel:path) => {
4189 pub(crate) fn $name(
4190 inputs: &[Arc<RwLock<Storage>>],
4191 outputs: &mut [Arc<RwLock<Storage>>],
4192 _layouts: &[Layout],
4193 params: &OpParams,
4194 ) -> Result<()> {
4195 if inputs.len() != 1 {
4196 return Err(Error::Msg(format!(
4197 "reduce_max_to wrapper expects 1 input, got {}",
4198 inputs.len(),
4199 ))
4200 .bt());
4201 }
4202 if outputs.len() != 1 {
4203 return Err(Error::Msg(format!(
4204 "reduce_max_to wrapper expects 1 output, got {}",
4205 outputs.len(),
4206 ))
4207 .bt());
4208 }
4209 let (input_shape, output_shape) = match params {
4210 OpParams::ReduceMaxTo {
4211 input_shape,
4212 output_shape,
4213 } => (input_shape.clone(), output_shape.clone()),
4214 other => {
4215 return Err(Error::Msg(format!(
4216 "reduce_max_to wrapper expects OpParams::ReduceMaxTo, got {other:?}",
4217 ))
4218 .bt())
4219 }
4220 };
4221 let in_guard = read_storage(&inputs[0])?;
4222 let mut out_guard = write_storage(&outputs[0])?;
4223 let in_cpu = cpu_input(&in_guard)?;
4224 let out_cpu = cpu_output(&mut out_guard)?;
4225 $kernel(in_cpu, out_cpu, &input_shape, &output_shape)
4226 }
4227 };
4228}
4229
4230cpu_reduce_max_to_wrapper!(
4231 reduce_max_to_f32_cpu_wrapper,
4232 fuel_cpu_backend::byte_kernels::reduce_max_to_f32
4233);
4234cpu_reduce_max_to_wrapper!(
4235 reduce_max_to_f64_cpu_wrapper,
4236 fuel_cpu_backend::byte_kernels::reduce_max_to_f64
4237);
4238cpu_reduce_max_to_wrapper!(
4239 reduce_max_to_bf16_cpu_wrapper,
4240 fuel_cpu_backend::byte_kernels::reduce_max_to_bf16
4241);
4242cpu_reduce_max_to_wrapper!(
4243 reduce_max_to_f16_cpu_wrapper,
4244 fuel_cpu_backend::byte_kernels::reduce_max_to_f16
4245);
4246
4247macro_rules! cpu_fused_linear_wrapper {
4250 ($wrapper:ident, $kernel:path) => {
4251 pub(crate) fn $wrapper(
4254 inputs: &[Arc<RwLock<Storage>>],
4255 outputs: &mut [Arc<RwLock<Storage>>],
4256 _layouts: &[Layout],
4257 params: &OpParams,
4258 ) -> Result<()> {
4259 if inputs.len() != 3 {
4260 return Err(Error::Msg(format!(
4261 "fused_linear wrapper expects 3 inputs (lhs, rhs, bias), got {}",
4262 inputs.len(),
4263 ))
4264 .bt());
4265 }
4266 if outputs.len() != 1 {
4267 return Err(Error::Msg(format!(
4268 "fused_linear wrapper expects 1 output, got {}",
4269 outputs.len(),
4270 ))
4271 .bt());
4272 }
4273 let (lhs_batch_dims, rhs_batch_dims, m, n, k) = match params {
4274 OpParams::Matmul {
4275 lhs_batch_dims,
4276 rhs_batch_dims,
4277 m,
4278 n,
4279 k,
4280 ..
4281 } => (lhs_batch_dims, rhs_batch_dims, *m, *n, *k),
4282 other => {
4283 return Err(Error::Msg(format!(
4284 "fused_linear wrapper expects OpParams::Matmul, got {other:?}",
4285 ))
4286 .bt())
4287 }
4288 };
4289 let lhs_guard = read_storage(&inputs[0])?;
4290 let rhs_guard = read_storage(&inputs[1])?;
4291 let bias_guard = read_storage(&inputs[2])?;
4292 let mut out_guard = write_storage(&outputs[0])?;
4293 let lhs_cpu = cpu_input(&lhs_guard)?;
4294 let rhs_cpu = cpu_input(&rhs_guard)?;
4295 let bias_cpu = cpu_input(&bias_guard)?;
4296 let out_cpu = cpu_output(&mut out_guard)?;
4297 $kernel(
4298 lhs_cpu,
4299 rhs_cpu,
4300 bias_cpu,
4301 out_cpu,
4302 lhs_batch_dims,
4303 rhs_batch_dims,
4304 m,
4305 n,
4306 k,
4307 )
4308 }
4309 };
4310}
4311
4312cpu_fused_linear_wrapper!(
4313 fused_linear_f32_cpu_wrapper,
4314 fuel_cpu_backend::byte_kernels::fused_linear_f32
4315);
4316cpu_fused_linear_wrapper!(
4317 fused_linear_f64_cpu_wrapper,
4318 fuel_cpu_backend::byte_kernels::fused_linear_f64
4319);
4320cpu_fused_linear_wrapper!(
4321 fused_linear_bf16_cpu_wrapper,
4322 fuel_cpu_backend::byte_kernels::fused_linear_bf16
4323);
4324cpu_fused_linear_wrapper!(
4325 fused_linear_f16_cpu_wrapper,
4326 fuel_cpu_backend::byte_kernels::fused_linear_f16
4327);
4328
4329macro_rules! cpu_fused_softmax_cross_entropy_wrapper {
4336 ($wrapper:ident, $kernel:path) => {
4337 pub(crate) fn $wrapper(
4338 inputs: &[Arc<RwLock<Storage>>],
4339 outputs: &mut [Arc<RwLock<Storage>>],
4340 _layouts: &[Layout],
4341 params: &OpParams,
4342 ) -> Result<()> {
4343 if inputs.len() != 2 {
4344 return Err(Error::Msg(format!(
4345 "fused_softmax_cross_entropy wrapper expects 2 inputs (logits, targets), got {}",
4346 inputs.len(),
4347 ))
4348 .bt());
4349 }
4350 if outputs.len() != 1 {
4351 return Err(Error::Msg(format!(
4352 "fused_softmax_cross_entropy wrapper expects 1 output, got {}",
4353 outputs.len(),
4354 ))
4355 .bt());
4356 }
4357 let (n_rows, vocab, reduction, ignore_index) = match params {
4358 OpParams::FusedSoftmaxCrossEntropy {
4359 n_rows, vocab, reduction, ignore_index,
4360 } => (*n_rows, *vocab, *reduction, *ignore_index),
4361 other => {
4362 return Err(Error::Msg(format!(
4363 "fused_softmax_cross_entropy wrapper expects \
4364 OpParams::FusedSoftmaxCrossEntropy, got {other:?}",
4365 ))
4366 .bt());
4367 }
4368 };
4369 let reduction_tag = match reduction {
4370 fuel_graph::registry::Reduction::Mean => fuel_cpu_backend::byte_kernels::REDUCTION_MEAN,
4371 fuel_graph::registry::Reduction::Sum => fuel_cpu_backend::byte_kernels::REDUCTION_SUM,
4372 fuel_graph::registry::Reduction::None => fuel_cpu_backend::byte_kernels::REDUCTION_NONE,
4373 };
4374 let logits_guard = read_storage(&inputs[0])?;
4375 let targets_guard = read_storage(&inputs[1])?;
4376 let mut out_guard = write_storage(&outputs[0])?;
4377 let logits_cpu = cpu_input(&logits_guard)?;
4378 let targets_cpu = cpu_input(&targets_guard)?;
4379 let out_cpu = cpu_output(&mut out_guard)?;
4380 $kernel(
4381 logits_cpu, targets_cpu, out_cpu,
4382 n_rows, vocab, reduction_tag, ignore_index,
4383 )
4384 }
4385 };
4386}
4387
4388cpu_fused_softmax_cross_entropy_wrapper!(
4389 fused_softmax_cross_entropy_f32_cpu_wrapper,
4390 fuel_cpu_backend::byte_kernels::fused_softmax_cross_entropy_f32
4391);
4392cpu_fused_softmax_cross_entropy_wrapper!(
4393 fused_softmax_cross_entropy_f64_cpu_wrapper,
4394 fuel_cpu_backend::byte_kernels::fused_softmax_cross_entropy_f64
4395);
4396cpu_fused_softmax_cross_entropy_wrapper!(
4397 fused_softmax_cross_entropy_bf16_cpu_wrapper,
4398 fuel_cpu_backend::byte_kernels::fused_softmax_cross_entropy_bf16
4399);
4400cpu_fused_softmax_cross_entropy_wrapper!(
4401 fused_softmax_cross_entropy_f16_cpu_wrapper,
4402 fuel_cpu_backend::byte_kernels::fused_softmax_cross_entropy_f16
4403);
4404
4405macro_rules! cpu_causal_conv1d_wrapper {
4410 ($wrapper:ident, $kernel:path) => {
4411 pub(crate) fn $wrapper(
4412 inputs: &[Arc<RwLock<Storage>>],
4413 outputs: &mut [Arc<RwLock<Storage>>],
4414 _layouts: &[Layout],
4415 params: &OpParams,
4416 ) -> Result<()> {
4417 if inputs.len() != 3 {
4418 return Err(Error::Msg(format!(
4419 "causal_conv1d wrapper expects 3 inputs (x, weight, bias), got {}",
4420 inputs.len(),
4421 ))
4422 .bt());
4423 }
4424 if outputs.len() != 1 {
4425 return Err(Error::Msg(format!(
4426 "causal_conv1d wrapper expects 1 output, got {}",
4427 outputs.len(),
4428 ))
4429 .bt());
4430 }
4431 let (batch, channels, seq_in, seq_out, kernel, use_silu) = match params {
4432 OpParams::CausalConv1d {
4433 batch,
4434 channels,
4435 seq_in,
4436 seq_out,
4437 kernel,
4438 use_silu,
4439 } => (*batch, *channels, *seq_in, *seq_out, *kernel, *use_silu),
4440 other => {
4441 return Err(Error::Msg(format!(
4442 "causal_conv1d wrapper expects OpParams::CausalConv1d, got {other:?}",
4443 ))
4444 .bt());
4445 }
4446 };
4447 let x_guard = read_storage(&inputs[0])?;
4448 let w_guard = read_storage(&inputs[1])?;
4449 let bias_guard = read_storage(&inputs[2])?;
4450 let mut out_guard = write_storage(&outputs[0])?;
4451 let x_cpu = cpu_input(&x_guard)?;
4452 let w_cpu = cpu_input(&w_guard)?;
4453 let bias_cpu = cpu_input(&bias_guard)?;
4454 let out_cpu = cpu_output(&mut out_guard)?;
4455 $kernel(
4456 x_cpu, w_cpu, bias_cpu, out_cpu, batch, channels, seq_in, seq_out, kernel, use_silu,
4457 )
4458 }
4459 };
4460}
4461
4462cpu_causal_conv1d_wrapper!(
4463 causal_conv1d_f32_cpu_wrapper,
4464 fuel_cpu_backend::byte_kernels::causal_conv1d_f32
4465);
4466cpu_causal_conv1d_wrapper!(
4467 causal_conv1d_f64_cpu_wrapper,
4468 fuel_cpu_backend::byte_kernels::causal_conv1d_f64
4469);
4470cpu_causal_conv1d_wrapper!(
4471 causal_conv1d_bf16_cpu_wrapper,
4472 fuel_cpu_backend::byte_kernels::causal_conv1d_bf16
4473);
4474cpu_causal_conv1d_wrapper!(
4475 causal_conv1d_f16_cpu_wrapper,
4476 fuel_cpu_backend::byte_kernels::causal_conv1d_f16
4477);
4478
4479macro_rules! cpu_selective_scan_wrapper {
4484 ($wrapper:ident, $kernel:path) => {
4485 pub(crate) fn $wrapper(
4486 inputs: &[Arc<RwLock<Storage>>],
4487 outputs: &mut [Arc<RwLock<Storage>>],
4488 _layouts: &[Layout],
4489 params: &OpParams,
4490 ) -> Result<()> {
4491 if inputs.len() != 5 {
4492 return Err(Error::Msg(format!(
4493 "selective_scan wrapper expects 5 inputs (u, delta, a, b, c), got {}",
4494 inputs.len(),
4495 ))
4496 .bt());
4497 }
4498 if outputs.len() != 1 {
4499 return Err(Error::Msg(format!(
4500 "selective_scan wrapper expects 1 output, got {}",
4501 outputs.len(),
4502 ))
4503 .bt());
4504 }
4505 let (batch, seqlen, dim, dstate, delta_softplus) = match params {
4506 OpParams::SelectiveScan {
4507 batch,
4508 seqlen,
4509 dim,
4510 dstate,
4511 delta_softplus,
4512 } => (*batch, *seqlen, *dim, *dstate, *delta_softplus),
4513 other => {
4514 return Err(Error::Msg(format!(
4515 "selective_scan wrapper expects OpParams::SelectiveScan, got {other:?}",
4516 ))
4517 .bt());
4518 }
4519 };
4520 let u_guard = read_storage(&inputs[0])?;
4521 let delta_guard = read_storage(&inputs[1])?;
4522 let a_guard = read_storage(&inputs[2])?;
4523 let b_guard = read_storage(&inputs[3])?;
4524 let c_guard = read_storage(&inputs[4])?;
4525 let mut out_guard = write_storage(&outputs[0])?;
4526 let u_cpu = cpu_input(&u_guard)?;
4527 let delta_cpu = cpu_input(&delta_guard)?;
4528 let a_cpu = cpu_input(&a_guard)?;
4529 let b_cpu = cpu_input(&b_guard)?;
4530 let c_cpu = cpu_input(&c_guard)?;
4531 let out_cpu = cpu_output(&mut out_guard)?;
4532 $kernel(
4533 u_cpu,
4534 delta_cpu,
4535 a_cpu,
4536 b_cpu,
4537 c_cpu,
4538 out_cpu,
4539 batch,
4540 seqlen,
4541 dim,
4542 dstate,
4543 delta_softplus,
4544 )
4545 }
4546 };
4547}
4548
4549cpu_selective_scan_wrapper!(
4550 selective_scan_f32_cpu_wrapper,
4551 fuel_cpu_backend::byte_kernels::selective_scan_f32
4552);
4553cpu_selective_scan_wrapper!(
4554 selective_scan_f64_cpu_wrapper,
4555 fuel_cpu_backend::byte_kernels::selective_scan_f64
4556);
4557cpu_selective_scan_wrapper!(
4558 selective_scan_bf16_cpu_wrapper,
4559 fuel_cpu_backend::byte_kernels::selective_scan_bf16
4560);
4561cpu_selective_scan_wrapper!(
4562 selective_scan_f16_cpu_wrapper,
4563 fuel_cpu_backend::byte_kernels::selective_scan_f16
4564);
4565
4566macro_rules! cpu_nonzero_indices_wrapper {
4573 ($wrapper:ident, $kernel:path) => {
4574 pub(crate) fn $wrapper(
4575 inputs: &[Arc<RwLock<Storage>>],
4576 outputs: &mut [Arc<RwLock<Storage>>],
4577 _layouts: &[Layout],
4578 params: &OpParams,
4579 ) -> Result<()> {
4580 if inputs.len() != 1 {
4581 return Err(Error::Msg(format!(
4582 "nonzero_indices wrapper expects 1 input, got {}",
4583 inputs.len(),
4584 ))
4585 .bt());
4586 }
4587 if outputs.len() != 1 {
4588 return Err(Error::Msg(format!(
4589 "nonzero_indices wrapper expects 1 output, got {}",
4590 outputs.len(),
4591 ))
4592 .bt());
4593 }
4594 let capacity = match params {
4595 OpParams::NonZeroIndices { capacity, .. } => *capacity,
4596 other => {
4597 return Err(Error::Msg(format!(
4598 "nonzero_indices wrapper expects OpParams::NonZeroIndices, got {other:?}",
4599 ))
4600 .bt());
4601 }
4602 };
4603 let x_guard = read_storage(&inputs[0])?;
4604 let mut out_guard = write_storage(&outputs[0])?;
4605 let x_cpu = cpu_input(&x_guard)?;
4606 let out_cpu = cpu_output(&mut out_guard)?;
4607 $kernel(x_cpu, out_cpu, capacity)
4608 }
4609 };
4610}
4611
4612cpu_nonzero_indices_wrapper!(
4613 nonzero_indices_f32_cpu_wrapper,
4614 fuel_cpu_backend::byte_kernels::nonzero_indices_f32
4615);
4616cpu_nonzero_indices_wrapper!(
4617 nonzero_indices_u32_cpu_wrapper,
4618 fuel_cpu_backend::byte_kernels::nonzero_indices_u32
4619);
4620
4621macro_rules! cpu_ssd_chunk_scan_wrapper {
4626 ($wrapper:ident, $kernel:path) => {
4627 pub(crate) fn $wrapper(
4628 inputs: &[Arc<RwLock<Storage>>],
4629 outputs: &mut [Arc<RwLock<Storage>>],
4630 _layouts: &[Layout],
4631 params: &OpParams,
4632 ) -> Result<()> {
4633 if inputs.len() != 5 {
4634 return Err(Error::Msg(format!(
4635 "ssd_chunk_scan wrapper expects 5 inputs (x, dt, a, b, c), got {}",
4636 inputs.len(),
4637 ))
4638 .bt());
4639 }
4640 if outputs.len() != 1 {
4641 return Err(Error::Msg(format!(
4642 "ssd_chunk_scan wrapper expects 1 output, got {}",
4643 outputs.len(),
4644 ))
4645 .bt());
4646 }
4647 let (batch, seqlen, heads, head_dim, state_dim, chunk_size) = match params {
4648 OpParams::SsdChunkScan {
4649 batch,
4650 seqlen,
4651 heads,
4652 head_dim,
4653 state_dim,
4654 chunk_size,
4655 } => (*batch, *seqlen, *heads, *head_dim, *state_dim, *chunk_size),
4656 other => {
4657 return Err(Error::Msg(format!(
4658 "ssd_chunk_scan wrapper expects OpParams::SsdChunkScan, got {other:?}",
4659 ))
4660 .bt());
4661 }
4662 };
4663 let x_guard = read_storage(&inputs[0])?;
4664 let dt_guard = read_storage(&inputs[1])?;
4665 let a_guard = read_storage(&inputs[2])?;
4666 let b_guard = read_storage(&inputs[3])?;
4667 let c_guard = read_storage(&inputs[4])?;
4668 let mut out_guard = write_storage(&outputs[0])?;
4669 let x_cpu = cpu_input(&x_guard)?;
4670 let dt_cpu = cpu_input(&dt_guard)?;
4671 let a_cpu = cpu_input(&a_guard)?;
4672 let b_cpu = cpu_input(&b_guard)?;
4673 let c_cpu = cpu_input(&c_guard)?;
4674 let out_cpu = cpu_output(&mut out_guard)?;
4675 $kernel(
4676 x_cpu, dt_cpu, a_cpu, b_cpu, c_cpu, out_cpu, batch, seqlen, heads, head_dim,
4677 state_dim, chunk_size,
4678 )
4679 }
4680 };
4681}
4682
4683cpu_ssd_chunk_scan_wrapper!(
4684 ssd_chunk_scan_f32_cpu_wrapper,
4685 fuel_cpu_backend::byte_kernels::ssd_chunk_scan_f32
4686);
4687cpu_ssd_chunk_scan_wrapper!(
4688 ssd_chunk_scan_f64_cpu_wrapper,
4689 fuel_cpu_backend::byte_kernels::ssd_chunk_scan_f64
4690);
4691cpu_ssd_chunk_scan_wrapper!(
4692 ssd_chunk_scan_bf16_cpu_wrapper,
4693 fuel_cpu_backend::byte_kernels::ssd_chunk_scan_bf16
4694);
4695cpu_ssd_chunk_scan_wrapper!(
4696 ssd_chunk_scan_f16_cpu_wrapper,
4697 fuel_cpu_backend::byte_kernels::ssd_chunk_scan_f16
4698);
4699
4700macro_rules! cpu_nf4_matmul_wrapper {
4704 ($wrapper:ident, $kernel:path) => {
4705 fn $wrapper(
4706 inputs: &[Arc<RwLock<Storage>>],
4707 outputs: &mut [Arc<RwLock<Storage>>],
4708 _layouts: &[Layout],
4709 params: &OpParams,
4710 ) -> Result<()> {
4711 if inputs.len() != 3 {
4712 return Err(Error::Msg(format!(
4713 "nf4_matmul wrapper expects 3 inputs (activations, w_packed, absmax), got {}",
4714 inputs.len(),
4715 ))
4716 .bt());
4717 }
4718 if outputs.len() != 1 {
4719 return Err(Error::Msg(format!(
4720 "nf4_matmul wrapper expects 1 output, got {}",
4721 outputs.len(),
4722 ))
4723 .bt());
4724 }
4725 let (batch, m, n, k, block_size) = match params {
4726 OpParams::Nf4Matmul {
4727 batch,
4728 m,
4729 n,
4730 k,
4731 block_size,
4732 } => (*batch, *m, *n, *k, *block_size),
4733 other => {
4734 return Err(Error::Msg(format!(
4735 "nf4_matmul wrapper expects OpParams::Nf4Matmul, got {other:?}",
4736 ))
4737 .bt());
4738 }
4739 };
4740 let a_guard = read_storage(&inputs[0])?;
4741 let w_guard = read_storage(&inputs[1])?;
4742 let abs_guard = read_storage(&inputs[2])?;
4743 let mut out_guard = write_storage(&outputs[0])?;
4744 let a_cpu = cpu_input(&a_guard)?;
4745 let w_cpu = cpu_input(&w_guard)?;
4746 let abs_cpu = cpu_input(&abs_guard)?;
4747 let out_cpu = cpu_output(&mut out_guard)?;
4748 $kernel(a_cpu, w_cpu, abs_cpu, out_cpu, batch, m, n, k, block_size)
4749 }
4750 };
4751}
4752
4753cpu_nf4_matmul_wrapper!(
4754 nf4_matmul_f32_cpu_wrapper,
4755 fuel_cpu_backend::byte_kernels::nf4_matmul_f32
4756);
4757cpu_nf4_matmul_wrapper!(
4758 nf4_matmul_f16_cpu_wrapper,
4759 fuel_cpu_backend::byte_kernels::nf4_matmul_f16
4760);
4761cpu_nf4_matmul_wrapper!(
4762 nf4_matmul_bf16_cpu_wrapper,
4763 fuel_cpu_backend::byte_kernels::nf4_matmul_bf16
4764);
4765
4766macro_rules! cpu_flash_attn_wrapper {
4770 ($wrapper:ident, $kernel:path) => {
4771 pub(crate) fn $wrapper(
4774 inputs: &[Arc<RwLock<Storage>>],
4775 outputs: &mut [Arc<RwLock<Storage>>],
4776 _layouts: &[Layout],
4777 params: &OpParams,
4778 ) -> Result<()> {
4779 if inputs.len() != 3 && inputs.len() != 4 {
4780 return Err(Error::Msg(format!(
4781 "flash_attn wrapper expects 3 or 4 inputs (q, k, v, [alibi]), got {}",
4782 inputs.len(),
4783 ))
4784 .bt());
4785 }
4786 if outputs.len() != 1 {
4787 return Err(Error::Msg(format!(
4788 "flash_attn wrapper expects 1 output, got {}",
4789 outputs.len(),
4790 ))
4791 .bt());
4792 }
4793 let (b, hq, hkv, sq, sk, d, k_len, scale, causal, wl, wr, softcap) = match params {
4794 OpParams::FlashAttn {
4795 b,
4796 hq,
4797 hkv,
4798 sq,
4799 sk,
4800 d,
4801 k_len,
4802 softmax_scale,
4803 causal,
4804 window_size_left,
4805 window_size_right,
4806 softcap,
4807 } => (
4808 *b,
4809 *hq,
4810 *hkv,
4811 *sq,
4812 *sk,
4813 *d,
4814 *k_len,
4815 *softmax_scale,
4816 *causal,
4817 *window_size_left,
4818 *window_size_right,
4819 *softcap,
4820 ),
4821 other => {
4822 return Err(Error::Msg(format!(
4823 "flash_attn wrapper expects OpParams::FlashAttn, got {other:?}",
4824 ))
4825 .bt())
4826 }
4827 };
4828 let q_guard = read_storage(&inputs[0])?;
4829 let k_guard = read_storage(&inputs[1])?;
4830 let v_guard = read_storage(&inputs[2])?;
4831 let alibi_guard = match inputs.get(3) {
4832 Some(arc) => Some(read_storage(arc)?),
4833 None => None,
4834 };
4835 let mut out_guard = write_storage(&outputs[0])?;
4836 let q_cpu = cpu_input(&q_guard)?;
4837 let k_cpu = cpu_input(&k_guard)?;
4838 let v_cpu = cpu_input(&v_guard)?;
4839 let alibi_cpu = match &alibi_guard {
4840 Some(g) => Some(cpu_input(g)?),
4841 None => None,
4842 };
4843 let out_cpu = cpu_output(&mut out_guard)?;
4844 $kernel(
4845 q_cpu, k_cpu, v_cpu, alibi_cpu, out_cpu, b, hq, hkv, sq, sk, d, k_len, scale,
4846 causal, wl, wr, softcap,
4847 )
4848 }
4849 };
4850}
4851
4852cpu_flash_attn_wrapper!(
4853 flash_attn_f32_cpu_wrapper,
4854 fuel_cpu_backend::byte_kernels::flash_attn_f32
4855);
4856cpu_flash_attn_wrapper!(
4857 flash_attn_f64_cpu_wrapper,
4858 fuel_cpu_backend::byte_kernels::flash_attn_f64
4859);
4860cpu_flash_attn_wrapper!(
4861 flash_attn_bf16_cpu_wrapper,
4862 fuel_cpu_backend::byte_kernels::flash_attn_bf16
4863);
4864cpu_flash_attn_wrapper!(
4865 flash_attn_f16_cpu_wrapper,
4866 fuel_cpu_backend::byte_kernels::flash_attn_f16
4867);
4868
4869macro_rules! cpu_flash_attn_backward_wrapper {
4874 ($wrapper:ident, $kernel:path, $which:expr) => {
4875 pub(crate) fn $wrapper(
4878 inputs: &[Arc<RwLock<Storage>>],
4879 outputs: &mut [Arc<RwLock<Storage>>],
4880 _layouts: &[Layout],
4881 params: &OpParams,
4882 ) -> Result<()> {
4883 if inputs.len() != 4 && inputs.len() != 5 {
4884 return Err(Error::Msg(format!(
4885 "flash_attn_backward wrapper expects 4 or 5 inputs (q, k, v, do, [alibi]), got {}",
4886 inputs.len(),
4887 ))
4888 .bt());
4889 }
4890 if outputs.len() != 1 {
4891 return Err(Error::Msg(format!(
4892 "flash_attn_backward wrapper expects 1 output, got {}",
4893 outputs.len(),
4894 ))
4895 .bt());
4896 }
4897 let (b, hq, hkv, sq, sk, d, scale, causal, wl, wr, softcap) = match params {
4898 OpParams::FlashAttn {
4899 b, hq, hkv, sq, sk, d,
4900 softmax_scale, causal,
4901 window_size_left, window_size_right, softcap,
4902 k_len: _,
4905 } => (
4906 *b, *hq, *hkv, *sq, *sk, *d,
4907 *softmax_scale, *causal,
4908 *window_size_left, *window_size_right, *softcap,
4909 ),
4910 other => {
4911 return Err(Error::Msg(format!(
4912 "flash_attn_backward wrapper expects OpParams::FlashAttn, got {other:?}",
4913 ))
4914 .bt())
4915 }
4916 };
4917 let q_guard = read_storage(&inputs[0])?;
4918 let k_guard = read_storage(&inputs[1])?;
4919 let v_guard = read_storage(&inputs[2])?;
4920 let do_guard = read_storage(&inputs[3])?;
4921 let alibi_guard = match inputs.get(4) {
4922 Some(arc) => Some(read_storage(arc)?),
4923 None => None,
4924 };
4925 let mut out_guard = write_storage(&outputs[0])?;
4926 let q_cpu = cpu_input(&q_guard)?;
4927 let k_cpu = cpu_input(&k_guard)?;
4928 let v_cpu = cpu_input(&v_guard)?;
4929 let do_cpu = cpu_input(&do_guard)?;
4930 let alibi_cpu = match &alibi_guard {
4931 Some(g) => Some(cpu_input(g)?),
4932 None => None,
4933 };
4934 let out_cpu = cpu_output(&mut out_guard)?;
4935 $kernel(
4936 q_cpu, k_cpu, v_cpu, do_cpu, alibi_cpu, out_cpu, $which,
4937 b, hq, hkv, sq, sk, d,
4938 scale, causal, wl, wr, softcap,
4939 )
4940 }
4941 };
4942}
4943
4944cpu_flash_attn_backward_wrapper!(
4945 flash_attn_backward_q_f32_cpu_wrapper,
4946 fuel_cpu_backend::byte_kernels::flash_attn_backward_f32,
4947 fuel_cpu_backend::byte_kernels::FaBackwardWhich::Q
4948);
4949cpu_flash_attn_backward_wrapper!(
4950 flash_attn_backward_k_f32_cpu_wrapper,
4951 fuel_cpu_backend::byte_kernels::flash_attn_backward_f32,
4952 fuel_cpu_backend::byte_kernels::FaBackwardWhich::K
4953);
4954cpu_flash_attn_backward_wrapper!(
4955 flash_attn_backward_v_f32_cpu_wrapper,
4956 fuel_cpu_backend::byte_kernels::flash_attn_backward_f32,
4957 fuel_cpu_backend::byte_kernels::FaBackwardWhich::V
4958);
4959cpu_flash_attn_backward_wrapper!(
4960 flash_attn_backward_q_f64_cpu_wrapper,
4961 fuel_cpu_backend::byte_kernels::flash_attn_backward_f64,
4962 fuel_cpu_backend::byte_kernels::FaBackwardWhich::Q
4963);
4964cpu_flash_attn_backward_wrapper!(
4965 flash_attn_backward_k_f64_cpu_wrapper,
4966 fuel_cpu_backend::byte_kernels::flash_attn_backward_f64,
4967 fuel_cpu_backend::byte_kernels::FaBackwardWhich::K
4968);
4969cpu_flash_attn_backward_wrapper!(
4970 flash_attn_backward_v_f64_cpu_wrapper,
4971 fuel_cpu_backend::byte_kernels::flash_attn_backward_f64,
4972 fuel_cpu_backend::byte_kernels::FaBackwardWhich::V
4973);
4974cpu_flash_attn_backward_wrapper!(
4975 flash_attn_backward_q_bf16_cpu_wrapper,
4976 fuel_cpu_backend::byte_kernels::flash_attn_backward_bf16,
4977 fuel_cpu_backend::byte_kernels::FaBackwardWhich::Q
4978);
4979cpu_flash_attn_backward_wrapper!(
4980 flash_attn_backward_k_bf16_cpu_wrapper,
4981 fuel_cpu_backend::byte_kernels::flash_attn_backward_bf16,
4982 fuel_cpu_backend::byte_kernels::FaBackwardWhich::K
4983);
4984cpu_flash_attn_backward_wrapper!(
4985 flash_attn_backward_v_bf16_cpu_wrapper,
4986 fuel_cpu_backend::byte_kernels::flash_attn_backward_bf16,
4987 fuel_cpu_backend::byte_kernels::FaBackwardWhich::V
4988);
4989cpu_flash_attn_backward_wrapper!(
4990 flash_attn_backward_q_f16_cpu_wrapper,
4991 fuel_cpu_backend::byte_kernels::flash_attn_backward_f16,
4992 fuel_cpu_backend::byte_kernels::FaBackwardWhich::Q
4993);
4994cpu_flash_attn_backward_wrapper!(
4995 flash_attn_backward_k_f16_cpu_wrapper,
4996 fuel_cpu_backend::byte_kernels::flash_attn_backward_f16,
4997 fuel_cpu_backend::byte_kernels::FaBackwardWhich::K
4998);
4999cpu_flash_attn_backward_wrapper!(
5000 flash_attn_backward_v_f16_cpu_wrapper,
5001 fuel_cpu_backend::byte_kernels::flash_attn_backward_f16,
5002 fuel_cpu_backend::byte_kernels::FaBackwardWhich::V
5003);
5004
5005macro_rules! cpu_paged_attn_wrapper {
5008 ($wrapper:ident, $kernel:path) => {
5009 pub(crate) fn $wrapper(
5013 inputs: &[Arc<RwLock<Storage>>],
5014 outputs: &mut [Arc<RwLock<Storage>>],
5015 _layouts: &[Layout],
5016 params: &OpParams,
5017 ) -> Result<()> {
5018 if inputs.len() != 5 && inputs.len() != 6 {
5019 return Err(Error::Msg(format!(
5020 "paged_attn wrapper expects 5 or 6 inputs, got {}",
5021 inputs.len(),
5022 ))
5023 .bt());
5024 }
5025 if outputs.len() != 1 {
5026 return Err(Error::Msg(format!(
5027 "paged_attn wrapper expects 1 output, got {}",
5028 outputs.len(),
5029 ))
5030 .bt());
5031 }
5032 let (b, hq, hkv, sq, d, block_size, max_blocks_per_seq, num_blocks, scale, softcap) =
5033 match params {
5034 OpParams::PagedAttn {
5035 b,
5036 hq,
5037 hkv,
5038 sq,
5039 d,
5040 block_size,
5041 max_blocks_per_seq,
5042 num_blocks,
5043 softmax_scale,
5044 softcap,
5045 } => (
5046 *b,
5047 *hq,
5048 *hkv,
5049 *sq,
5050 *d,
5051 *block_size,
5052 *max_blocks_per_seq,
5053 *num_blocks,
5054 *softmax_scale,
5055 *softcap,
5056 ),
5057 other => {
5058 return Err(Error::Msg(format!(
5059 "paged_attn wrapper expects OpParams::PagedAttn, got {other:?}",
5060 ))
5061 .bt())
5062 }
5063 };
5064 let q_g = read_storage(&inputs[0])?;
5065 let kc_g = read_storage(&inputs[1])?;
5066 let vc_g = read_storage(&inputs[2])?;
5067 let bt_g = read_storage(&inputs[3])?;
5068 let cl_g = read_storage(&inputs[4])?;
5069 let alibi_g = match inputs.get(5) {
5070 Some(arc) => Some(read_storage(arc)?),
5071 None => None,
5072 };
5073 if bt_g.dtype != DType::U32 || cl_g.dtype != DType::U32 {
5074 return Err(Error::Msg(format!(
5075 "paged_attn: block_table and context_lens must be U32, got {:?} / {:?}",
5076 bt_g.dtype, cl_g.dtype,
5077 ))
5078 .bt());
5079 }
5080 let mut out_guard = write_storage(&outputs[0])?;
5081 let q_cpu = cpu_input(&q_g)?;
5082 let kc_cpu = cpu_input(&kc_g)?;
5083 let vc_cpu = cpu_input(&vc_g)?;
5084 let bt_cpu = cpu_input(&bt_g)?;
5085 let cl_cpu = cpu_input(&cl_g)?;
5086 let alibi_cpu = match &alibi_g {
5087 Some(g) => Some(cpu_input(g)?),
5088 None => None,
5089 };
5090 let out_cpu = cpu_output(&mut out_guard)?;
5091 $kernel(
5092 q_cpu,
5093 kc_cpu,
5094 vc_cpu,
5095 bt_cpu,
5096 cl_cpu,
5097 alibi_cpu,
5098 out_cpu,
5099 b,
5100 hq,
5101 hkv,
5102 sq,
5103 d,
5104 block_size,
5105 max_blocks_per_seq,
5106 num_blocks,
5107 scale,
5108 softcap,
5109 )
5110 }
5111 };
5112}
5113
5114cpu_paged_attn_wrapper!(
5115 paged_attn_f32_cpu_wrapper,
5116 fuel_cpu_backend::byte_kernels::paged_attn_f32
5117);
5118cpu_paged_attn_wrapper!(
5119 paged_attn_f64_cpu_wrapper,
5120 fuel_cpu_backend::byte_kernels::paged_attn_f64
5121);
5122cpu_paged_attn_wrapper!(
5123 paged_attn_bf16_cpu_wrapper,
5124 fuel_cpu_backend::byte_kernels::paged_attn_bf16
5125);
5126cpu_paged_attn_wrapper!(
5127 paged_attn_f16_cpu_wrapper,
5128 fuel_cpu_backend::byte_kernels::paged_attn_f16
5129);
5130
5131cpu_cast_wrapper!(
5136 cast_to_f32_cpu_wrapper,
5137 DType::F32,
5138 "f32",
5139 {
5140 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_f32,
5141 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_f32,
5142 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_f32,
5143 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_f32,
5144 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_f32,
5145 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_f32,
5146 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_f32,
5147 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_f32,
5148 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_f32,
5149 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_f32,
5150 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_f32,
5151 },
5152);
5153cpu_cast_wrapper!(
5154 cast_to_f64_cpu_wrapper,
5155 DType::F64,
5156 "f64",
5157 {
5158 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_f64,
5159 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_f64,
5160 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_f64,
5161 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_f64,
5162 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_f64,
5163 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_f64,
5164 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_f64,
5165 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_f64,
5166 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_f64,
5167 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_f64,
5168 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_f64,
5169 },
5170);
5171cpu_cast_wrapper!(
5172 cast_to_bf16_cpu_wrapper,
5173 DType::BF16,
5174 "bf16",
5175 {
5176 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_bf16,
5177 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_bf16,
5178 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_bf16,
5179 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_bf16,
5180 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_bf16,
5181 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_bf16,
5182 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_bf16,
5183 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_bf16,
5184 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_bf16,
5185 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_bf16,
5186 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_bf16,
5187 },
5188);
5189cpu_cast_wrapper!(
5190 cast_to_f16_cpu_wrapper,
5191 DType::F16,
5192 "f16",
5193 {
5194 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_f16,
5195 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_f16,
5196 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_f16,
5197 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_f16,
5198 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_f16,
5199 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_f16,
5200 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_f16,
5201 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_f16,
5202 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_f16,
5203 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_f16,
5204 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_f16,
5205 },
5206);
5207cpu_cast_wrapper!(
5208 cast_to_f8e4m3_cpu_wrapper,
5209 DType::F8E4M3,
5210 "f8e4m3",
5211 {
5212 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_f8e4m3,
5213 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_f8e4m3,
5214 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_f8e4m3,
5215 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_f8e4m3,
5216 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_f8e4m3,
5217 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_f8e4m3,
5218 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_f8e4m3,
5219 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_f8e4m3,
5220 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_f8e4m3,
5221 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_f8e4m3,
5222 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_f8e4m3,
5223 },
5224);
5225cpu_cast_wrapper!(
5226 cast_to_u8_cpu_wrapper,
5227 DType::U8,
5228 "u8",
5229 {
5230 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_u8,
5231 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_u8,
5232 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_u8,
5233 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_u8,
5234 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_u8,
5235 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_u8,
5236 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_u8,
5237 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_u8,
5238 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_u8,
5239 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_u8,
5240 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_u8,
5241 },
5242);
5243cpu_cast_wrapper!(
5244 cast_to_i8_cpu_wrapper,
5245 DType::I8,
5246 "i8",
5247 {
5248 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_i8,
5249 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_i8,
5250 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_i8,
5251 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_i8,
5252 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_i8,
5253 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_i8,
5254 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_i8,
5255 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_i8,
5256 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_i8,
5257 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_i8,
5258 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_i8,
5259 },
5260);
5261cpu_cast_wrapper!(
5262 cast_to_u32_cpu_wrapper,
5263 DType::U32,
5264 "u32",
5265 {
5266 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_u32,
5267 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_u32,
5268 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_u32,
5269 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_u32,
5270 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_u32,
5271 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_u32,
5272 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_u32,
5273 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_u32,
5274 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_u32,
5275 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_u32,
5276 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_u32,
5277 },
5278);
5279cpu_cast_wrapper!(
5280 cast_to_i16_cpu_wrapper,
5281 DType::I16,
5282 "i16",
5283 {
5284 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_i16,
5285 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_i16,
5286 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_i16,
5287 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_i16,
5288 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_i16,
5289 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_i16,
5290 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_i16,
5291 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_i16,
5292 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_i16,
5293 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_i16,
5294 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_i16,
5295 },
5296);
5297cpu_cast_wrapper!(
5298 cast_to_i32_cpu_wrapper,
5299 DType::I32,
5300 "i32",
5301 {
5302 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_i32,
5303 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_i32,
5304 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_i32,
5305 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_i32,
5306 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_i32,
5307 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_i32,
5308 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_i32,
5309 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_i32,
5310 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_i32,
5311 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_i32,
5312 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_i32,
5313 },
5314);
5315cpu_cast_wrapper!(
5316 cast_to_i64_cpu_wrapper,
5317 DType::I64,
5318 "i64",
5319 {
5320 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_i64,
5321 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_i64,
5322 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_i64,
5323 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_i64,
5324 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_i64,
5325 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_i64,
5326 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_i64,
5327 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_i64,
5328 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_i64,
5329 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_i64,
5330 DType::Bool => fuel_cpu_backend::byte_kernels::cast_bool_to_i64,
5331 },
5332);
5333
5334cpu_cast_wrapper!(
5338 cast_to_bool_cpu_wrapper,
5339 DType::Bool,
5340 "bool",
5341 {
5342 DType::F32 => fuel_cpu_backend::byte_kernels::cast_f32_to_bool,
5343 DType::F64 => fuel_cpu_backend::byte_kernels::cast_f64_to_bool,
5344 DType::F16 => fuel_cpu_backend::byte_kernels::cast_f16_to_bool,
5345 DType::BF16 => fuel_cpu_backend::byte_kernels::cast_bf16_to_bool,
5346 DType::F8E4M3 => fuel_cpu_backend::byte_kernels::cast_f8e4m3_to_bool,
5347 DType::U8 => fuel_cpu_backend::byte_kernels::cast_u8_to_bool,
5348 DType::I8 => fuel_cpu_backend::byte_kernels::cast_i8_to_bool,
5349 DType::U32 => fuel_cpu_backend::byte_kernels::cast_u32_to_bool,
5350 DType::I16 => fuel_cpu_backend::byte_kernels::cast_i16_to_bool,
5351 DType::I32 => fuel_cpu_backend::byte_kernels::cast_i32_to_bool,
5352 DType::I64 => fuel_cpu_backend::byte_kernels::cast_i64_to_bool,
5353 },
5354);
5355
5356pub(crate) fn matmul_f32_cpu_wrapper(
5365 inputs: &[Arc<RwLock<Storage>>],
5366 outputs: &mut [Arc<RwLock<Storage>>],
5367 _layouts: &[Layout],
5368 params: &OpParams,
5369) -> Result<()> {
5370 if inputs.len() != 2 {
5371 return Err(Error::Msg(format!(
5372 "matmul wrapper expects 2 inputs, got {}",
5373 inputs.len(),
5374 ))
5375 .bt());
5376 }
5377 if outputs.len() != 1 {
5378 return Err(Error::Msg(format!(
5379 "matmul wrapper expects 1 output, got {}",
5380 outputs.len(),
5381 ))
5382 .bt());
5383 }
5384 let (lhs_batch_dims, rhs_batch_dims, m, n, k, m_compute) = match params {
5385 OpParams::Matmul {
5386 lhs_batch_dims,
5387 rhs_batch_dims,
5388 m,
5389 n,
5390 k,
5391 m_compute,
5392 } => (lhs_batch_dims, rhs_batch_dims, *m, *n, *k, m_compute),
5393 other => {
5394 return Err(Error::Msg(format!(
5395 "matmul wrapper expects OpParams::Matmul, got {other:?}",
5396 ))
5397 .bt());
5398 }
5399 };
5400 let rows = match m_compute {
5404 MatmulM::All => m,
5405 MatmulM::Rows(c) => *c,
5406 MatmulM::Deferred(d) => {
5407 return Err(Error::Msg(format!(
5408 "matmul_f32 wrapper: unresolved data-determined row count {d:?} \
5409 (must be resolved at execute before the kernel runs)",
5410 ))
5411 .bt());
5412 }
5413 };
5414 let lhs_guard = read_storage(&inputs[0])?;
5415 let rhs_guard = read_storage(&inputs[1])?;
5416 let mut out_guard = write_storage(&outputs[0])?;
5417 let lhs_cpu = cpu_input(&lhs_guard)?;
5418 let rhs_cpu = cpu_input(&rhs_guard)?;
5419 let out_cpu = cpu_output(&mut out_guard)?;
5420 fuel_cpu_backend::byte_kernels::matmul_f32_capacity(
5421 lhs_cpu,
5422 rhs_cpu,
5423 out_cpu,
5424 lhs_batch_dims,
5425 rhs_batch_dims,
5426 m,
5427 rows,
5428 n,
5429 k,
5430 )
5431}
5432
5433macro_rules! cpu_matmul_wrapper {
5438 ($wrapper:ident, $kernel:path, $type_name:literal) => {
5439 pub(crate) fn $wrapper(
5442 inputs: &[Arc<RwLock<Storage>>],
5443 outputs: &mut [Arc<RwLock<Storage>>],
5444 _layouts: &[Layout],
5445 params: &OpParams,
5446 ) -> Result<()> {
5447 if inputs.len() != 2 {
5448 return Err(Error::Msg(format!(
5449 "matmul wrapper expects 2 inputs, got {}",
5450 inputs.len(),
5451 ))
5452 .bt());
5453 }
5454 if outputs.len() != 1 {
5455 return Err(Error::Msg(format!(
5456 "matmul wrapper expects 1 output, got {}",
5457 outputs.len(),
5458 ))
5459 .bt());
5460 }
5461 let (lhs_batch_dims, rhs_batch_dims, m, n, k) = match params {
5462 OpParams::Matmul {
5463 lhs_batch_dims,
5464 rhs_batch_dims,
5465 m,
5466 n,
5467 k,
5468 ..
5469 } => (lhs_batch_dims, rhs_batch_dims, *m, *n, *k),
5470 other => {
5471 return Err(Error::Msg(format!(
5472 "matmul wrapper expects OpParams::Matmul, got {other:?}",
5473 ))
5474 .bt())
5475 }
5476 };
5477 let lhs_guard = read_storage(&inputs[0])?;
5478 let rhs_guard = read_storage(&inputs[1])?;
5479 let mut out_guard = write_storage(&outputs[0])?;
5480 let lhs_cpu = cpu_input(&lhs_guard)?;
5481 let rhs_cpu = cpu_input(&rhs_guard)?;
5482 let out_cpu = cpu_output(&mut out_guard)?;
5483 let _ = $type_name;
5484 $kernel(
5485 lhs_cpu,
5486 rhs_cpu,
5487 out_cpu,
5488 lhs_batch_dims,
5489 rhs_batch_dims,
5490 m,
5491 n,
5492 k,
5493 )
5494 }
5495 };
5496}
5497
5498cpu_matmul_wrapper!(
5499 matmul_bf16_cpu_wrapper,
5500 fuel_cpu_backend::byte_kernels::matmul_bf16,
5501 "matmul_bf16"
5502);
5503cpu_matmul_wrapper!(
5504 matmul_f16_cpu_wrapper,
5505 fuel_cpu_backend::byte_kernels::matmul_f16,
5506 "matmul_f16"
5507);
5508cpu_matmul_wrapper!(
5509 matmul_i8_cpu_wrapper,
5510 fuel_cpu_backend::byte_kernels::matmul_i8,
5511 "matmul_i8"
5512);
5513cpu_matmul_wrapper!(
5514 matmul_u8_cpu_wrapper,
5515 fuel_cpu_backend::byte_kernels::matmul_u8,
5516 "matmul_u8"
5517);
5518
5519pub(crate) fn matmul_f64_cpu_wrapper(
5526 inputs: &[Arc<RwLock<Storage>>],
5527 outputs: &mut [Arc<RwLock<Storage>>],
5528 _layouts: &[Layout],
5529 params: &OpParams,
5530) -> Result<()> {
5531 if inputs.len() != 2 {
5532 return Err(Error::Msg(format!(
5533 "matmul wrapper expects 2 inputs, got {}",
5534 inputs.len(),
5535 ))
5536 .bt());
5537 }
5538 if outputs.len() != 1 {
5539 return Err(Error::Msg(format!(
5540 "matmul wrapper expects 1 output, got {}",
5541 outputs.len(),
5542 ))
5543 .bt());
5544 }
5545 let (lhs_batch_dims, rhs_batch_dims, m, n, k) = match params {
5546 OpParams::Matmul {
5547 lhs_batch_dims,
5548 rhs_batch_dims,
5549 m,
5550 n,
5551 k,
5552 ..
5553 } => (lhs_batch_dims, rhs_batch_dims, *m, *n, *k),
5554 other => {
5555 return Err(Error::Msg(format!(
5556 "matmul wrapper expects OpParams::Matmul, got {other:?}",
5557 ))
5558 .bt());
5559 }
5560 };
5561 let lhs_guard = read_storage(&inputs[0])?;
5562 let rhs_guard = read_storage(&inputs[1])?;
5563 let mut out_guard = write_storage(&outputs[0])?;
5564 let lhs_cpu = cpu_input(&lhs_guard)?;
5565 let rhs_cpu = cpu_input(&rhs_guard)?;
5566 let out_cpu = cpu_output(&mut out_guard)?;
5567 fuel_cpu_backend::byte_kernels::matmul_f64(
5568 lhs_cpu,
5569 rhs_cpu,
5570 out_cpu,
5571 lhs_batch_dims,
5572 rhs_batch_dims,
5573 m,
5574 n,
5575 k,
5576 )
5577}
5578
5579fn copy_from_cpu_wrapper(
5605 inputs: &[Arc<RwLock<Storage>>],
5606 outputs: &mut [Arc<RwLock<Storage>>],
5607 _layouts: &[Layout],
5608 _params: &OpParams,
5609) -> Result<()> {
5610 if inputs.len() != 1 || outputs.len() != 1 {
5611 return Err(Error::Msg(format!(
5612 "copy_from_cpu_wrapper: expected 1 input + 1 output, got {} + {}",
5613 inputs.len(),
5614 outputs.len(),
5615 ))
5616 .bt());
5617 }
5618 let in_guard = read_storage(&inputs[0])?;
5619 let src = cpu_input(&in_guard)?;
5620 let mut out_guard = write_storage(&outputs[0])?;
5621 let n_out = out_guard.inner.len_bytes();
5628 let n_src = src.len_bytes();
5629 let n = n_src.min(n_out);
5630 let src_slice = &src.bytes()[..n];
5631 match &mut out_guard.inner {
5632 BackendStorage::Cpu(dst) => {
5633 dst.bytes_mut()[..n].copy_from_slice(src_slice);
5634 }
5635 #[cfg(feature = "cuda")]
5636 BackendStorage::Cuda(dst) => {
5637 if n == n_out {
5644 dst.write_from_host(src_slice)?;
5645 } else {
5646 let mut staged = vec![0_u8; n_out];
5649 staged[..n].copy_from_slice(src_slice);
5650 dst.write_from_host(&staged)?;
5651 }
5652 }
5653 #[cfg(feature = "vulkan")]
5654 BackendStorage::Vulkan(dst) => {
5655 let backend = dst
5656 .backend()
5657 .ok_or_else(|| {
5658 Error::Msg(
5659 "copy_from_cpu_wrapper: Vulkan output has no attached \
5660 backend handle. The executor's Op::Copy arm must allocate \
5661 via VulkanBackend::alloc_bytes_handle (which attaches \
5662 the handle)."
5663 .to_string(),
5664 )
5665 .bt()
5666 })?
5667 .clone();
5668 if n == n_out {
5669 backend.write_bytes(dst, src_slice)?;
5670 } else {
5671 let mut staged = vec![0_u8; n_out];
5672 staged[..n].copy_from_slice(src_slice);
5673 backend.write_bytes(dst, &staged)?;
5674 }
5675 }
5676 #[allow(unreachable_patterns)]
5677 other => {
5678 return Err(Error::Msg(format!(
5679 "copy_from_cpu_wrapper: output backend not wired ({other:?}); \
5680 CPU + CUDA + Vulkan covered, Metal extends when its \
5681 byte-storage substrate is ready.",
5682 ))
5683 .bt());
5684 }
5685 }
5686 Ok(())
5687}
5688
5689const CPU_ELEMENTWISE_BINARY_CONTRACT: &str =
5694 include_str!("../kernel-contracts/cpu/elementwise-binary.fkc.md");
5695
5696fn register_cpu_binary_from_contract(table: &mut KernelBindingTable) {
5728 let provider = crate::fkc::import_bundle_str(
5729 CPU_ELEMENTWISE_BINARY_CONTRACT,
5730 &crate::fkc::CpuLinkRegistry,
5731 )
5732 .expect(
5733 "authored CPU elementwise-binary contract must import \
5734 (embedded via include_str!, resolved through CpuLinkRegistry)",
5735 );
5736 debug_assert!(
5737 provider.fused.is_empty(),
5738 "elementwise-binary contract declares no fused ops",
5739 );
5740 let mut fused = crate::fused::FusedKernelRegistry::new();
5741 provider
5742 .register_into(table, &mut fused)
5743 .expect("CPU elementwise-binary contract must register into the binding table");
5744}
5745
5746const CPU_AFFINE_CLAMP_POWI_CONTRACT: &str =
5750 include_str!("../kernel-contracts/cpu/affine-clamp-powi.fkc.md");
5751
5752fn register_cpu_affine_clamp_powi_from_contract(table: &mut KernelBindingTable) {
5778 let provider =
5779 crate::fkc::import_bundle_str(CPU_AFFINE_CLAMP_POWI_CONTRACT, &crate::fkc::CpuLinkRegistry)
5780 .expect(
5781 "authored CPU affine/clamp/powi contract must import \
5782 (embedded via include_str!, resolved through CpuLinkRegistry)",
5783 );
5784 debug_assert!(
5785 provider.fused.is_empty(),
5786 "affine/clamp/powi contract declares no fused ops",
5787 );
5788 let mut fused = crate::fused::FusedKernelRegistry::new();
5789 provider
5790 .register_into(table, &mut fused)
5791 .expect("CPU affine/clamp/powi contract must register into the binding table");
5792}
5793
5794const CPU_ELEMENTWISE_UNARY_CONTRACT: &str =
5798 include_str!("../kernel-contracts/cpu/elementwise-unary.fkc.md");
5799
5800fn register_cpu_unary_from_contract(table: &mut KernelBindingTable) {
5827 let provider =
5828 crate::fkc::import_bundle_str(CPU_ELEMENTWISE_UNARY_CONTRACT, &crate::fkc::CpuLinkRegistry)
5829 .expect(
5830 "authored CPU elementwise-unary contract must import \
5831 (embedded via include_str!, resolved through CpuLinkRegistry)",
5832 );
5833 debug_assert!(
5834 provider.fused.is_empty(),
5835 "elementwise-unary contract declares no fused ops",
5836 );
5837 let mut fused = crate::fused::FusedKernelRegistry::new();
5838 provider
5839 .register_into(table, &mut fused)
5840 .expect("CPU elementwise-unary contract must register into the binding table");
5841}
5842
5843const CPU_COMPARE_WHERE_CONTRACT: &str =
5847 include_str!("../kernel-contracts/cpu/compare-where.fkc.md");
5848
5849fn register_cpu_compare_where_from_contract(table: &mut KernelBindingTable) {
5878 let provider =
5879 crate::fkc::import_bundle_str(CPU_COMPARE_WHERE_CONTRACT, &crate::fkc::CpuLinkRegistry)
5880 .expect(
5881 "authored CPU compare/where contract must import \
5882 (embedded via include_str!, resolved through CpuLinkRegistry)",
5883 );
5884 debug_assert!(
5885 provider.fused.is_empty(),
5886 "compare/where contract declares no fused ops",
5887 );
5888 let mut fused = crate::fused::FusedKernelRegistry::new();
5889 provider
5890 .register_into(table, &mut fused)
5891 .expect("CPU compare/where contract must register into the binding table");
5892}
5893
5894const CPU_REDUCE_CONTRACT: &str = include_str!("../kernel-contracts/cpu/reduce.fkc.md");
5898
5899fn register_cpu_reduce_from_contract(table: &mut KernelBindingTable) {
5929 let provider = crate::fkc::import_bundle_str(CPU_REDUCE_CONTRACT, &crate::fkc::CpuLinkRegistry)
5930 .expect(
5931 "authored CPU reduce contract must import \
5932 (embedded via include_str!, resolved through CpuLinkRegistry)",
5933 );
5934 debug_assert!(
5935 provider.fused.is_empty(),
5936 "reduce contract declares no fused ops",
5937 );
5938 let mut fused = crate::fused::FusedKernelRegistry::new();
5939 provider
5940 .register_into(table, &mut fused)
5941 .expect("CPU reduce contract must register into the binding table");
5942}
5943
5944const CPU_REDUCE_TO_CONTRACT: &str = include_str!("../kernel-contracts/cpu/reduce-to.fkc.md");
5948
5949fn register_cpu_reduce_to_from_contract(table: &mut KernelBindingTable) {
5981 let provider =
5982 crate::fkc::import_bundle_str(CPU_REDUCE_TO_CONTRACT, &crate::fkc::CpuLinkRegistry).expect(
5983 "authored CPU reduce-to contract must import \
5984 (embedded via include_str!, resolved through CpuLinkRegistry)",
5985 );
5986 debug_assert!(
5987 provider.fused.is_empty(),
5988 "reduce-to contract declares no fused ops",
5989 );
5990 let mut fused = crate::fused::FusedKernelRegistry::new();
5991 provider
5992 .register_into(table, &mut fused)
5993 .expect("CPU reduce-to contract must register into the binding table");
5994}
5995
5996const CPU_NORM_CONTRACT: &str = include_str!("../kernel-contracts/cpu/norm.fkc.md");
6000
6001fn register_cpu_norm_from_contract(table: &mut KernelBindingTable) {
6036 let provider = crate::fkc::import_bundle_str(CPU_NORM_CONTRACT, &crate::fkc::CpuLinkRegistry)
6037 .expect(
6038 "authored CPU norm contract must import \
6039 (embedded via include_str!, resolved through CpuLinkRegistry)",
6040 );
6041 debug_assert!(
6042 provider.fused.is_empty(),
6043 "norm forward contract declares no fused ops",
6044 );
6045 let mut fused = crate::fused::FusedKernelRegistry::new();
6046 provider
6047 .register_into(table, &mut fused)
6048 .expect("CPU norm contract must register into the binding table");
6049}
6050
6051const CPU_NORM_BACKWARD_CONTRACT: &str =
6055 include_str!("../kernel-contracts/cpu/norm-backward.fkc.md");
6056
6057fn register_cpu_norm_backward_from_contract(table: &mut KernelBindingTable) {
6091 let provider =
6092 crate::fkc::import_bundle_str(CPU_NORM_BACKWARD_CONTRACT, &crate::fkc::CpuLinkRegistry)
6093 .expect(
6094 "authored CPU norm-backward contract must import \
6095 (embedded via include_str!, resolved through CpuLinkRegistry)",
6096 );
6097 debug_assert!(
6098 provider.fused.is_empty(),
6099 "norm backward contract declares no fused ops",
6100 );
6101 let mut fused = crate::fused::FusedKernelRegistry::new();
6102 provider
6103 .register_into(table, &mut fused)
6104 .expect("CPU norm-backward contract must register into the binding table");
6105}
6106
6107const CPU_ROPE_CONTRACT: &str = include_str!("../kernel-contracts/cpu/rope.fkc.md");
6111
6112fn register_cpu_rope_from_contract(table: &mut KernelBindingTable) {
6146 let provider = crate::fkc::import_bundle_str(CPU_ROPE_CONTRACT, &crate::fkc::CpuLinkRegistry)
6147 .expect(
6148 "authored CPU rope contract must import \
6149 (embedded via include_str!, resolved through CpuLinkRegistry)",
6150 );
6151 debug_assert!(
6152 provider.fused.is_empty(),
6153 "rope contract declares no fused ops",
6154 );
6155 let mut fused = crate::fused::FusedKernelRegistry::new();
6156 provider
6157 .register_into(table, &mut fused)
6158 .expect("CPU rope contract must register into the binding table");
6159}
6160
6161const CPU_SSM_CONTRACT: &str = include_str!("../kernel-contracts/cpu/ssm.fkc.md");
6165
6166fn register_cpu_ssm_from_contract(table: &mut KernelBindingTable) {
6205 let provider = crate::fkc::import_bundle_str(CPU_SSM_CONTRACT, &crate::fkc::CpuLinkRegistry)
6206 .expect(
6207 "authored CPU ssm contract must import \
6208 (embedded via include_str!, resolved through CpuLinkRegistry)",
6209 );
6210 debug_assert!(
6211 provider.fused.is_empty(),
6212 "ssm contract declares no fused ops",
6213 );
6214 let mut fused = crate::fused::FusedKernelRegistry::new();
6215 provider
6216 .register_into(table, &mut fused)
6217 .expect("CPU ssm contract must register into the binding table");
6218}
6219
6220const CPU_CONV_CONTRACT: &str = include_str!("../kernel-contracts/cpu/conv.fkc.md");
6225
6226fn register_cpu_conv_from_contract(table: &mut KernelBindingTable) {
6266 let provider = crate::fkc::import_bundle_str(CPU_CONV_CONTRACT, &crate::fkc::CpuLinkRegistry)
6267 .expect(
6268 "authored CPU conv contract must import \
6269 (embedded via include_str!, resolved through CpuLinkRegistry)",
6270 );
6271 debug_assert!(
6272 provider.fused.is_empty(),
6273 "conv contract declares no fused ops",
6274 );
6275 let mut fused = crate::fused::FusedKernelRegistry::new();
6276 provider
6277 .register_into(table, &mut fused)
6278 .expect("CPU conv contract must register into the binding table");
6279}
6280
6281const CPU_PADDING_CONTRACT: &str = include_str!("../kernel-contracts/cpu/padding.fkc.md");
6286
6287fn register_cpu_padding_from_contract(table: &mut KernelBindingTable) {
6323 let provider =
6324 crate::fkc::import_bundle_str(CPU_PADDING_CONTRACT, &crate::fkc::CpuLinkRegistry).expect(
6325 "authored CPU padding contract must import \
6326 (embedded via include_str!, resolved through CpuLinkRegistry)",
6327 );
6328 debug_assert!(
6329 provider.fused.is_empty(),
6330 "padding contract declares no fused ops",
6331 );
6332 let mut fused = crate::fused::FusedKernelRegistry::new();
6333 provider
6334 .register_into(table, &mut fused)
6335 .expect("CPU padding contract must register into the binding table");
6336}
6337
6338const CPU_SHAPE_OPS_CONTRACT: &str = include_str!("../kernel-contracts/cpu/shape-ops.fkc.md");
6342
6343fn register_cpu_shape_ops_from_contract(table: &mut KernelBindingTable) {
6384 let provider =
6385 crate::fkc::import_bundle_str(CPU_SHAPE_OPS_CONTRACT, &crate::fkc::CpuLinkRegistry).expect(
6386 "authored CPU shape-ops contract must import \
6387 (embedded via include_str!, resolved through CpuLinkRegistry)",
6388 );
6389 debug_assert!(
6390 provider.fused.is_empty(),
6391 "shape-ops contract declares no fused ops",
6392 );
6393 let mut fused = crate::fused::FusedKernelRegistry::new();
6394 provider
6395 .register_into(table, &mut fused)
6396 .expect("CPU shape-ops contract must register into the binding table");
6397}
6398
6399const CPU_MATMUL_CONTRACT: &str = include_str!("../kernel-contracts/cpu/matmul.fkc.md");
6403
6404fn register_cpu_matmul_from_contract(table: &mut KernelBindingTable) {
6449 let provider = crate::fkc::import_bundle_str(CPU_MATMUL_CONTRACT, &crate::fkc::CpuLinkRegistry)
6450 .expect(
6451 "authored CPU matmul contract must import \
6452 (embedded via include_str!, resolved through CpuLinkRegistry)",
6453 );
6454 debug_assert!(
6455 provider.fused.is_empty(),
6456 "matmul contract declares no fused ops",
6457 );
6458 let mut fused = crate::fused::FusedKernelRegistry::new();
6459 provider
6460 .register_into(table, &mut fused)
6461 .expect("CPU matmul contract must register into the binding table");
6462}
6463
6464const CPU_ATTENTION_CONTRACT: &str = include_str!("../kernel-contracts/cpu/attention.fkc.md");
6469
6470fn register_cpu_attention_from_contract(table: &mut KernelBindingTable) {
6526 let provider =
6527 crate::fkc::import_bundle_str(CPU_ATTENTION_CONTRACT, &crate::fkc::CpuLinkRegistry).expect(
6528 "authored CPU attention contract must import \
6529 (embedded via include_str!, resolved through CpuLinkRegistry)",
6530 );
6531 debug_assert!(
6532 provider.fused.is_empty(),
6533 "attention contract declares no fused ops (all op_kind, incl. PagedAttn)",
6534 );
6535 let mut fused = crate::fused::FusedKernelRegistry::new();
6536 provider
6537 .register_into(table, &mut fused)
6538 .expect("CPU attention contract must register into the binding table");
6539}
6540
6541const CPU_INPLACE_CONTRACT: &str =
6545 include_str!("../kernel-contracts/cpu/inplace-unary-affine.fkc.md");
6546
6547fn register_cpu_inplace_from_contract(table: &mut KernelBindingTable) {
6587 let provider =
6588 crate::fkc::import_bundle_str(CPU_INPLACE_CONTRACT, &crate::fkc::CpuLinkRegistry).expect(
6589 "authored CPU in-place contract must import \
6590 (embedded via include_str!, resolved through CpuLinkRegistry)",
6591 );
6592 debug_assert!(
6593 provider.fused.is_empty(),
6594 "in-place contract declares no fused ops (all op_kind or describe-only umbrella)",
6595 );
6596 let mut fused = crate::fused::FusedKernelRegistry::new();
6597 provider
6598 .register_into(table, &mut fused)
6599 .expect("CPU in-place contract must register into the binding table");
6600}
6601
6602const CPU_CAST_CONTRACT: &str = include_str!("../kernel-contracts/cpu/cast.fkc.md");
6606
6607fn register_cpu_cast_from_contract(table: &mut KernelBindingTable) {
6637 let provider = crate::fkc::import_bundle_str(CPU_CAST_CONTRACT, &crate::fkc::CpuLinkRegistry)
6638 .expect(
6639 "authored CPU cast contract must import \
6640 (embedded via include_str!, resolved through CpuLinkRegistry)",
6641 );
6642 debug_assert!(
6643 provider.fused.is_empty(),
6644 "cast contract declares no fused ops",
6645 );
6646 let mut fused = crate::fused::FusedKernelRegistry::new();
6647 provider
6648 .register_into(table, &mut fused)
6649 .expect("CPU cast contract must register into the binding table");
6650}
6651
6652const CPU_INDEXING_CONTRACT: &str = include_str!("../kernel-contracts/cpu/indexing.fkc.md");
6657
6658fn register_cpu_indexing_from_contract(table: &mut KernelBindingTable) {
6697 let provider =
6698 crate::fkc::import_bundle_str(CPU_INDEXING_CONTRACT, &crate::fkc::CpuLinkRegistry).expect(
6699 "authored CPU indexing contract must import \
6700 (embedded via include_str!, resolved through CpuLinkRegistry)",
6701 );
6702 debug_assert!(
6703 provider.fused.is_empty(),
6704 "indexing contract declares no fused ops",
6705 );
6706 let mut fused = crate::fused::FusedKernelRegistry::new();
6707 provider
6708 .register_into(table, &mut fused)
6709 .expect("CPU indexing contract must register into the binding table");
6710}
6711
6712pub fn register_cpu_kernels(table: &mut KernelBindingTable) {
6718 use OpKind::*;
6719 let cpu = BackendId::Cpu;
6720 let f32_dt = DType::F32;
6721 let f64_dt = DType::F64;
6722 let bf16_dt = DType::BF16;
6723 let f16_dt = DType::F16;
6724 let u32_dt = DType::U32;
6725
6726 let unary = |t: DType| [t, t]; let u8_dt = DType::U8;
6736 register_cpu_binary_from_contract(table);
6774
6775 register_cpu_affine_clamp_powi_from_contract(table);
6785
6786 register_cpu_unary_from_contract(table);
6799
6800 register_cpu_compare_where_from_contract(table);
6811
6812 register_cpu_reduce_from_contract(table);
6824
6825 register_cpu_reduce_to_from_contract(table);
6837
6838 register_cpu_norm_from_contract(table);
6853
6854 register_cpu_norm_backward_from_contract(table);
6870
6871 register_cpu_rope_from_contract(table);
6884
6885 register_cpu_ssm_from_contract(table);
6900
6901 table.register(
6914 NonZeroIndices,
6915 &[f32_dt, u32_dt],
6916 cpu,
6917 nonzero_indices_f32_cpu_wrapper,
6918 );
6919 table.register(
6920 NonZeroIndices,
6921 &[u32_dt, u32_dt],
6922 cpu,
6923 nonzero_indices_u32_cpu_wrapper,
6924 );
6925
6926 register_cpu_conv_from_contract(table);
6944
6945 register_cpu_padding_from_contract(table);
6961
6962 register_cpu_shape_ops_from_contract(table);
6982
6983 register_cpu_matmul_from_contract(table);
7004
7005 register_cpu_attention_from_contract(table);
7022
7023 register_cpu_inplace_from_contract(table);
7036
7037 register_cpu_cast_from_contract(table);
7057
7058 register_cpu_indexing_from_contract(table);
7074
7075 table.register(
7132 Nf4Matmul,
7133 &[DType::F32, DType::U8, DType::F32, DType::F32],
7134 cpu,
7135 nf4_matmul_f32_cpu_wrapper,
7136 );
7137 table.register(
7138 Nf4Matmul,
7139 &[DType::F16, DType::U8, DType::F32, DType::F16],
7140 cpu,
7141 nf4_matmul_f16_cpu_wrapper,
7142 );
7143 table.register(
7144 Nf4Matmul,
7145 &[DType::BF16, DType::U8, DType::F32, DType::BF16],
7146 cpu,
7147 nf4_matmul_bf16_cpu_wrapper,
7148 );
7149
7150 table.register(Triu, &unary(f32_dt), cpu, triu_cpu_wrapper);
7191 table.register(Triu, &unary(f64_dt), cpu, triu_cpu_wrapper);
7192 table.register(Triu, &unary(bf16_dt), cpu, triu_cpu_wrapper);
7193 table.register(Triu, &unary(f16_dt), cpu, triu_cpu_wrapper);
7194 table.register(Triu, &unary(u32_dt), cpu, triu_cpu_wrapper);
7195 table.register(Triu, &unary(u8_dt), cpu, triu_cpu_wrapper);
7196 table.register(Tril, &unary(f32_dt), cpu, tril_cpu_wrapper);
7197 table.register(Tril, &unary(f64_dt), cpu, tril_cpu_wrapper);
7198 table.register(Tril, &unary(bf16_dt), cpu, tril_cpu_wrapper);
7199 table.register(Tril, &unary(f16_dt), cpu, tril_cpu_wrapper);
7200 table.register(Tril, &unary(u32_dt), cpu, tril_cpu_wrapper);
7201 table.register(Tril, &unary(u8_dt), cpu, tril_cpu_wrapper);
7202
7203 table.register(
7273 QMatMul,
7274 &[f32_dt, u32_dt, f32_dt],
7275 cpu,
7276 qmatmul_f32_cpu_wrapper,
7277 );
7278
7279 for dt in [f32_dt, f64_dt, bf16_dt, f16_dt] {
7290 table.register(ArgMaxDim, &[dt, u32_dt], cpu, argmax_dim_u32_cpu_dispatch);
7291 table.register(ArgMinDim, &[dt, u32_dt], cpu, argmin_dim_u32_cpu_dispatch);
7292 }
7293
7294 let copy_dtypes = [
7309 f32_dt,
7310 f64_dt,
7311 bf16_dt,
7312 f16_dt,
7313 u32_dt,
7314 u8_dt,
7315 DType::I16,
7316 DType::I32,
7317 DType::I64,
7318 DType::Bool,
7321 ];
7322 for dt in copy_dtypes {
7323 table.register(Copy, &[dt, dt], cpu, copy_from_cpu_wrapper);
7324 }
7325 table.fill_unset_cpu_precision(crate::fused::PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU);
7347
7348 table.fill_unset_cpu_cost(crate::cost::default_cost_for_op_kind);
7365}
7366
7367#[cfg(feature = "cuda")]
7381pub(crate) fn cuda_input(s: &Storage) -> Result<&fuel_cuda_backend::CudaStorageBytes> {
7382 match &s.inner {
7383 BackendStorage::Cuda(c) => Ok(c),
7384 #[allow(unreachable_patterns)]
7385 _ => Err(Error::Msg("cuda kernel wrapper called with non-CUDA input".to_string()).bt()),
7386 }
7387}
7388
7389#[cfg(feature = "cuda")]
7391pub(crate) fn cuda_output(s: &mut Storage) -> Result<&mut fuel_cuda_backend::CudaStorageBytes> {
7392 match &mut s.inner {
7393 BackendStorage::Cuda(c) => Ok(c),
7394 #[allow(unreachable_patterns)]
7395 _ => Err(Error::Msg("cuda kernel wrapper called with non-CUDA output".to_string()).bt()),
7396 }
7397}
7398
7399#[cfg(feature = "cuda")]
7422fn copy_from_cuda_wrapper(
7423 inputs: &[Arc<RwLock<Storage>>],
7424 outputs: &mut [Arc<RwLock<Storage>>],
7425 _layouts: &[Layout],
7426 _params: &OpParams,
7427) -> Result<()> {
7428 if inputs.len() != 1 || outputs.len() != 1 {
7429 return Err(Error::Msg(format!(
7430 "copy_from_cuda_wrapper: expected 1 input + 1 output, got {} + {}",
7431 inputs.len(),
7432 outputs.len(),
7433 ))
7434 .bt());
7435 }
7436 let in_guard = read_storage(&inputs[0])?;
7437 let cuda_src = cuda_input(&in_guard)?;
7438 let mut out_guard = write_storage(&outputs[0])?;
7439 match &mut out_guard.inner {
7440 BackendStorage::Cpu(_) => {
7441 let bytes = cuda_src.to_cpu_bytes_finer()?;
7453 let dst = cpu_output(&mut out_guard)?;
7454 let n = bytes.len().min(dst.len_bytes());
7455 dst.bytes_mut()[..n].copy_from_slice(&bytes[..n]);
7456 Ok(())
7457 }
7458 BackendStorage::Cuda(_) => {
7459 let dst = cuda_output(&mut out_guard)?;
7468 dst.copy_from_device(cuda_src)?;
7469 Ok(())
7470 }
7471 #[allow(unreachable_patterns)]
7472 other => Err(Error::Msg(format!(
7473 "copy_from_cuda_wrapper: unsupported output substrate {:?} \
7474 (CUDA sources copy to CPU or CUDA outputs only; cross-vendor \
7475 GPU transfer goes through host staging as two Copy hops)",
7476 std::mem::discriminant(other),
7477 ))
7478 .bt()),
7479 }
7480}
7481
7482#[cfg(feature = "cuda")]
7501pub fn register_cuda_kernels(table: &mut KernelBindingTable) {
7502 use OpKind::*;
7503 let cuda = BackendId::Cuda;
7504
7505 let copy_dtypes = [
7510 DType::F32,
7511 DType::BF16,
7512 DType::F16,
7513 DType::U32,
7514 DType::F64,
7515 DType::U8,
7516 DType::I16,
7517 DType::I32,
7518 DType::I64,
7519 DType::Bool,
7525 ];
7526 for dt in copy_dtypes {
7527 table.register(Copy, &[dt, dt], cuda, copy_from_cuda_wrapper);
7528 }
7529}
7530
7531static GLOBAL_REGISTRY: OnceLock<RwLock<CapabilityRegistry>> = OnceLock::new();
7544
7545static TOPOLOGY_GENERATION: AtomicU64 = AtomicU64::new(0);
7557
7558pub fn topology_generation() -> u64 {
7561 TOPOLOGY_GENERATION.load(Ordering::Acquire)
7562}
7563
7564pub fn bump_topology_generation() {
7571 TOPOLOGY_GENERATION.fetch_add(1, Ordering::Release);
7572}
7573
7574static DEVICE_INFLIGHT: OnceLock<RwLock<HashMap<DeviceLocation, AtomicU32>>> = OnceLock::new();
7625
7626fn device_inflight_table() -> &'static RwLock<HashMap<DeviceLocation, AtomicU32>> {
7627 DEVICE_INFLIGHT.get_or_init(|| RwLock::new(HashMap::new()))
7628}
7629
7630pub fn inflight_inc(loc: DeviceLocation) {
7638 let table = device_inflight_table();
7639 if let Ok(guard) = table.read()
7641 && let Some(slot) = guard.get(&loc)
7642 {
7643 slot.fetch_add(1, Ordering::Relaxed);
7644 return;
7645 }
7646 if let Ok(mut guard) = table.write() {
7650 guard
7651 .entry(loc)
7652 .or_insert_with(|| AtomicU32::new(0))
7653 .fetch_add(1, Ordering::Relaxed);
7654 }
7655}
7656
7657pub fn inflight_dec(loc: DeviceLocation) {
7665 let table = device_inflight_table();
7666 if let Ok(guard) = table.read()
7667 && let Some(slot) = guard.get(&loc)
7668 {
7669 let _ = slot.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
7671 Some(v.saturating_sub(1))
7672 });
7673 }
7674}
7675
7676pub fn inflight_count(loc: DeviceLocation) -> u32 {
7681 let table = device_inflight_table();
7682 table
7683 .read()
7684 .ok()
7685 .and_then(|guard| guard.get(&loc).map(|slot| slot.load(Ordering::Relaxed)))
7686 .unwrap_or(0)
7687}
7688
7689static GLOBAL_BINDINGS: OnceLock<RwLock<Arc<KernelBindingTable>>> = OnceLock::new();
7703
7704static BINDINGS_WRITER: Mutex<()> = Mutex::new(());
7728
7729fn bindings_cell() -> &'static RwLock<Arc<KernelBindingTable>> {
7733 GLOBAL_BINDINGS.get_or_init(|| {
7734 lock_trace(" BINDINGS init: begin");
7735 let mut t = KernelBindingTable::new();
7736 register_cpu_kernels(&mut t);
7737 lock_trace(" BINDINGS init: cpu done, entering optional backends");
7738 register_optional_backends(&mut t);
7739 lock_trace(" BINDINGS init: optional backends done");
7740 t.finalize()
7746 .expect("KernelBindingTable: duplicate kernel registration at init");
7747 RwLock::new(Arc::new(t))
7748 })
7749}
7750
7751fn default_cpu_caps() -> BackendCapabilities {
7752 use OpKind::*;
7753 use std::collections::HashSet;
7754 let mut op_dtype_support = HashSet::new();
7755 let f32_dt = DType::F32;
7759 for op in [
7760 AddElementwise,
7761 SubElementwise,
7762 MulElementwise,
7763 DivElementwise,
7764 ReluElementwise,
7765 NegElementwise,
7766 SqrElementwise,
7767 SqrtElementwise,
7768 RecipElementwise,
7769 AbsElementwise,
7770 TanhElementwise,
7771 ExpElementwise,
7772 LogElementwise,
7773 SinElementwise,
7774 CosElementwise,
7775 SigmoidElementwise,
7776 SiluElementwise,
7777 GeluElementwise,
7778 StepElementwise,
7779 SumReduce,
7780 MaxReduce,
7781 MinReduce,
7782 MeanReduce,
7783 MatMul,
7784 Conv2D,
7785 ConvTranspose2D,
7786 ReduceSumTo,
7787 ReduceMaxTo,
7788 FusedLinear,
7789 FlashAttn,
7790 PagedAttn,
7791 Affine,
7792 ClampElementwise,
7793 PowIElementwise,
7794 MaximumElementwise,
7795 MinimumElementwise,
7796 Concat,
7797 SoftmaxLastDim,
7798 RmsNormLastDim,
7799 LayerNormLastDim,
7800 IndexSelect,
7801 Gather,
7802 Rope,
7803 IndexAdd,
7804 ScatterAdd,
7805 QMatMul,
7806 Copy,
7807 ] {
7808 op_dtype_support.insert((op, f32_dt));
7809 }
7810 for op in [
7814 AddElementwise,
7815 SubElementwise,
7816 MulElementwise,
7817 DivElementwise,
7818 ReluElementwise,
7819 NegElementwise,
7820 SqrElementwise,
7821 SqrtElementwise,
7822 RecipElementwise,
7823 AbsElementwise,
7824 TanhElementwise,
7825 ExpElementwise,
7826 LogElementwise,
7827 SinElementwise,
7828 CosElementwise,
7829 SigmoidElementwise,
7830 SiluElementwise,
7831 GeluElementwise,
7832 StepElementwise,
7833 MaximumElementwise,
7834 MinimumElementwise,
7835 SumReduce,
7836 MaxReduce,
7837 MinReduce,
7838 MeanReduce,
7839 MatMul,
7840 Conv2D,
7841 ConvTranspose2D,
7842 ReduceSumTo,
7843 ReduceMaxTo,
7844 FusedLinear,
7845 FlashAttn,
7846 PagedAttn,
7847 ] {
7848 op_dtype_support.insert((op, DType::F64));
7849 }
7850 let half_elementwise_ops = [
7853 AddElementwise,
7854 SubElementwise,
7855 MulElementwise,
7856 DivElementwise,
7857 MaximumElementwise,
7858 MinimumElementwise,
7859 ReluElementwise,
7860 NegElementwise,
7861 SqrElementwise,
7862 SqrtElementwise,
7863 RecipElementwise,
7864 AbsElementwise,
7865 TanhElementwise,
7866 ExpElementwise,
7867 LogElementwise,
7868 SinElementwise,
7869 CosElementwise,
7870 SigmoidElementwise,
7871 SiluElementwise,
7872 GeluElementwise,
7873 StepElementwise,
7874 ];
7875 for op in half_elementwise_ops {
7876 op_dtype_support.insert((op, DType::BF16));
7877 op_dtype_support.insert((op, DType::F16));
7878 }
7879 for op in [SumReduce, MaxReduce, MinReduce, MeanReduce, MatMul] {
7881 op_dtype_support.insert((op, DType::BF16));
7882 op_dtype_support.insert((op, DType::F16));
7883 }
7884 for op in [
7889 SoftmaxLastDim,
7890 RmsNormLastDim,
7891 LayerNormLastDim,
7892 Rope,
7893 Conv2D,
7894 ConvTranspose2D,
7895 ReduceSumTo,
7896 ReduceMaxTo,
7897 FusedLinear,
7898 FlashAttn,
7899 PagedAttn,
7900 ] {
7901 op_dtype_support.insert((op, DType::BF16));
7902 op_dtype_support.insert((op, DType::F16));
7903 }
7904 for dt in [
7907 DType::F64,
7908 DType::BF16,
7909 DType::F16,
7910 DType::U32,
7911 DType::U8,
7912 DType::I16,
7913 DType::I32,
7914 DType::I64,
7915 ] {
7916 op_dtype_support.insert((Concat, dt));
7917 op_dtype_support.insert((IndexSelect, dt));
7918 op_dtype_support.insert((Gather, dt));
7919 }
7920 for dt in [DType::F64, DType::BF16, DType::F16] {
7922 op_dtype_support.insert((IndexAdd, dt));
7923 op_dtype_support.insert((ScatterAdd, dt));
7924 op_dtype_support.insert((Affine, dt));
7925 op_dtype_support.insert((ClampElementwise, dt));
7926 op_dtype_support.insert((PowIElementwise, dt));
7927 }
7928 for op in [SoftmaxLastDim, RmsNormLastDim, LayerNormLastDim, Rope] {
7930 op_dtype_support.insert((op, DType::F64));
7931 }
7932 for op in [ArgMaxDim, ArgMinDim] {
7934 op_dtype_support.insert((op, DType::U32));
7935 }
7936 for target in [DType::F32, DType::F64, DType::BF16, DType::F16] {
7941 op_dtype_support.insert((Cast, target));
7942 }
7943 let (compute_throughput_flops_per_ns, mem_bandwidth_bytes_per_ns) =
7944 crate::ranker::default_backend_rates(BackendId::Cpu);
7945 BackendCapabilities {
7946 backend_id: BackendId::Cpu,
7947 device_location: DeviceLocation::Cpu,
7948 op_dtype_support,
7949 required_alignment: 64,
7950 access_granularity_bits: 8,
7951 transfer_paths: vec![(DeviceLocation::Cpu, TransferPath::SameDevice)],
7952 storage_substrate: SubstrateClass::HostBytes,
7953 compute_throughput_flops_per_ns,
7954 mem_bandwidth_bytes_per_ns,
7955 }
7956}
7957
7958fn substrate_for_backend(backend: BackendId) -> SubstrateClass {
7962 match backend {
7963 BackendId::Cpu => SubstrateClass::HostBytes,
7964 BackendId::Cuda => SubstrateClass::CudaUntyped,
7965 BackendId::Vulkan => SubstrateClass::VulkanBuffer,
7966 BackendId::Metal => SubstrateClass::MetalBuffer,
7967 _ => SubstrateClass::HostBytes,
7970 }
7971}
7972
7973pub fn derive_backend_caps(
8009 backend: BackendId,
8010 device_location: DeviceLocation,
8011 table: &KernelBindingTable,
8012) -> BackendCapabilities {
8013 use std::collections::HashSet;
8014 let mut op_dtype_support: HashSet<(OpKind, DType)> = HashSet::new();
8015 for (op, dtypes, this_backend) in table.iter_keys() {
8016 if this_backend != backend {
8017 continue;
8018 }
8019 if let Some(&output_dt) = dtypes.last() {
8020 op_dtype_support.insert((op, output_dt));
8021 }
8022 }
8023 let transfer_paths = if device_location == DeviceLocation::Cpu {
8028 vec![(DeviceLocation::Cpu, TransferPath::SameDevice)]
8029 } else {
8030 vec![(DeviceLocation::Cpu, TransferPath::HostStaging)]
8031 };
8032 let (required_alignment, access_granularity_bits) = match backend {
8033 BackendId::Cpu => (64, 8),
8034 _ => (256, 8),
8038 };
8039 let (compute_throughput_flops_per_ns, mem_bandwidth_bytes_per_ns) =
8040 crate::ranker::default_backend_rates(backend);
8041 BackendCapabilities {
8042 backend_id: backend,
8043 device_location,
8044 op_dtype_support,
8045 required_alignment,
8046 access_granularity_bits,
8047 transfer_paths,
8048 storage_substrate: substrate_for_backend(backend),
8049 compute_throughput_flops_per_ns,
8050 mem_bandwidth_bytes_per_ns,
8051 }
8052}
8053
8054fn register_derived_gpu_caps(table: &KernelBindingTable) {
8074 use std::collections::HashSet;
8075 let gpu_backends: HashSet<BackendId> = table
8077 .iter_keys()
8078 .map(|(_, _, backend)| backend)
8079 .filter(|b| *b != BackendId::Cpu)
8080 .collect();
8081 for backend in gpu_backends {
8082 let device_location = match backend {
8083 BackendId::Cuda => DeviceLocation::Cuda { gpu_id: 0 },
8084 BackendId::Vulkan => DeviceLocation::Vulkan { gpu_id: 0 },
8085 BackendId::Metal => DeviceLocation::Metal { gpu_id: 0 },
8086 _ => DeviceLocation::Cpu,
8089 };
8090 register_backend_capabilities(derive_backend_caps(backend, device_location, table));
8091 }
8092}
8093
8094#[inline]
8104pub(crate) fn lock_trace(what: &str) {
8105 if std::env::var_os("FUEL_LOCK_TRACE").is_some() {
8106 eprintln!("[locktrace] t={:?} {}", std::thread::current().id(), what);
8107 }
8108}
8109
8110pub fn global_registry() -> std::sync::RwLockReadGuard<'static, CapabilityRegistry> {
8111 lock_trace("want REGISTRY.oncelock+read");
8112 let l = GLOBAL_REGISTRY.get_or_init(|| {
8113 lock_trace(" REGISTRY init: begin");
8114 let mut r = CapabilityRegistry::new();
8115 r.register(default_cpu_caps());
8116 lock_trace(" REGISTRY init: end");
8117 RwLock::new(r)
8118 });
8119 let g = l.read().unwrap();
8120 lock_trace("got REGISTRY.read");
8121 g
8122}
8123
8124pub fn register_backend_capabilities(caps: BackendCapabilities) {
8131 lock_trace("want REGISTRY.oncelock+WRITE (from register_backend_capabilities)");
8132 let lock = GLOBAL_REGISTRY.get_or_init(|| {
8133 lock_trace(" REGISTRY init: begin (via register_backend_capabilities)");
8134 let mut r = CapabilityRegistry::new();
8135 r.register(default_cpu_caps());
8136 lock_trace(" REGISTRY init: end");
8137 RwLock::new(r)
8138 });
8139 lock.write().unwrap().register(caps);
8140 lock_trace("got+released REGISTRY.WRITE");
8141 bump_topology_generation();
8142}
8143
8144pub fn global_bindings() -> Arc<KernelBindingTable> {
8166 lock_trace("want BINDINGS.oncelock+snapshot");
8167 let cell = bindings_cell();
8168 let snapshot = {
8169 let guard = cell.read().unwrap_or_else(|e| e.into_inner());
8174 Arc::clone(&guard)
8175 };
8177 lock_trace("got BINDINGS snapshot (guard already released)");
8178 snapshot
8179}
8180
8181fn register_optional_backends(table: &mut KernelBindingTable) {
8185 #[cfg(feature = "cuda")]
8186 {
8187 register_cuda_kernels(table);
8188 crate::baracuda_dispatch::register_baracuda_cuda_kernels(table);
8189 table.fill_unset_cost_for_backend(BackendId::Cuda, crate::cost::default_cost_for_op_kind);
8199 }
8200 #[cfg(feature = "vulkan")]
8201 {
8202 crate::vulkan_dispatch::register_vulkan_kernels(table);
8203 }
8204
8205 register_derived_gpu_caps(table);
8215
8216 let _ = table;
8217}
8218
8219pub fn extend_global_bindings(register: impl FnOnce(&mut KernelBindingTable)) {
8232 let cell = bindings_cell();
8233 {
8234 lock_trace("want BINDINGS.WRITER (extend_global_bindings)");
8239 let _writer = BINDINGS_WRITER.lock().unwrap_or_else(|e| e.into_inner());
8240 lock_trace("got BINDINGS.WRITER -- cloning table, extender runs OFF the RwLock");
8241
8242 let mut next = {
8243 let guard = cell.read().unwrap_or_else(|e| e.into_inner());
8244 (**guard).clone()
8245 };
8246 register(&mut next);
8247 lock_trace(" extender returned");
8248 next.finalize()
8254 .expect("KernelBindingTable: duplicate kernel registration after extend");
8255
8256 {
8258 let mut guard = cell.write().unwrap_or_else(|e| e.into_inner());
8259 *guard = Arc::new(next);
8260 }
8261 lock_trace("released BINDINGS.WRITE (swap only)");
8262 }
8263 bump_topology_generation();
8264}
8265
8266const FUSED_LINEAR_QUANT_CONTRACT: &str =
8270 include_str!("../kernel-contracts/fused/linear-quant.fkc.md");
8271
8272fn register_cpu_linear_quant_fused_from_contract(r: &mut crate::fused::FusedKernelRegistry) {
8304 let provider =
8305 crate::fkc::import_bundle_str(FUSED_LINEAR_QUANT_CONTRACT, &crate::fkc::CpuLinkRegistry)
8306 .expect(
8307 "authored linear-quant fused contract must import \
8308 (embedded via include_str!, resolved through CpuLinkRegistry)",
8309 );
8310 debug_assert!(
8311 provider.primitives.is_empty(),
8312 "linear-quant bundle declares only fused ops",
8313 );
8314 let mut table = KernelBindingTable::new();
8315 provider
8316 .register_into(&mut table, r)
8317 .expect("linear-quant fused contract must register into the fused registry");
8318}
8319
8320const FUSED_NORM_SOFTMAX_CONTRACT: &str =
8324 include_str!("../kernel-contracts/fused/norm-softmax.fkc.md");
8325
8326fn register_cpu_norm_softmax_fused_from_contract(r: &mut crate::fused::FusedKernelRegistry) {
8350 let provider =
8351 crate::fkc::import_bundle_str(FUSED_NORM_SOFTMAX_CONTRACT, &crate::fkc::CpuLinkRegistry)
8352 .expect(
8353 "authored norm-softmax fused contract must import \
8354 (embedded via include_str!, resolved through CpuLinkRegistry)",
8355 );
8356 debug_assert!(
8357 provider.primitives.is_empty(),
8358 "norm-softmax bundle declares only fused ops",
8359 );
8360 let mut table = KernelBindingTable::new();
8361 provider
8362 .register_into(&mut table, r)
8363 .expect("norm-softmax fused contract must register into the fused registry");
8364}
8365
8366const FUSED_CONV_ROPE_CONTRACT: &str = include_str!("../kernel-contracts/fused/conv-rope.fkc.md");
8370
8371fn register_cpu_conv_rope_fused_from_contract(r: &mut crate::fused::FusedKernelRegistry) {
8408 let provider =
8409 crate::fkc::import_bundle_str(FUSED_CONV_ROPE_CONTRACT, &crate::fkc::CpuLinkRegistry)
8410 .expect(
8411 "authored conv-rope fused contract must import \
8412 (embedded via include_str!, resolved through CpuLinkRegistry)",
8413 );
8414 debug_assert!(
8415 provider.primitives.is_empty(),
8416 "conv-rope bundle declares only fused ops",
8417 );
8418 let mut table = KernelBindingTable::new();
8419 provider
8420 .register_into(&mut table, r)
8421 .expect("conv-rope fused contract must register into the fused registry");
8422}
8423
8424pub fn register_default_fused_kernels(r: &mut crate::fused::FusedKernelRegistry) {
8458 use crate::fused::{
8464 ATTN_BACKWARD_CPU_PRECISION, ATTN_CPU_PRECISION, NF4_MATMUL_CPU_PRECISION,
8465 cost_attn_backward_cpu, cost_attn_cpu, cost_nf4_matmul_cpu,
8466 };
8467 use crate::register_fused;
8468 use fuel_graph::registry::FusedOps;
8469
8470 const FA_F32_NOA: &[DType] = &[DType::F32, DType::F32, DType::F32, DType::F32];
8482 const FA_F32_A: &[DType] = &[DType::F32, DType::F32, DType::F32, DType::F32, DType::F32];
8483 const FA_F64_NOA: &[DType] = &[DType::F64, DType::F64, DType::F64, DType::F64];
8484 const FA_F64_A: &[DType] = &[DType::F64, DType::F64, DType::F64, DType::F64, DType::F64];
8485 const FA_BF16_NOA: &[DType] = &[DType::BF16, DType::BF16, DType::BF16, DType::BF16];
8486 const FA_BF16_A: &[DType] = &[
8487 DType::BF16,
8488 DType::BF16,
8489 DType::BF16,
8490 DType::BF16,
8491 DType::BF16,
8492 ];
8493 const FA_F16_NOA: &[DType] = &[DType::F16, DType::F16, DType::F16, DType::F16];
8494 const FA_F16_A: &[DType] = &[DType::F16, DType::F16, DType::F16, DType::F16, DType::F16];
8495
8496 const FAB_F32_NOA: &[DType] = &[DType::F32, DType::F32, DType::F32, DType::F32, DType::F32];
8500 const FAB_F32_A: &[DType] = &[
8501 DType::F32,
8502 DType::F32,
8503 DType::F32,
8504 DType::F32,
8505 DType::F32,
8506 DType::F32,
8507 ];
8508 const FAB_F64_NOA: &[DType] = &[DType::F64, DType::F64, DType::F64, DType::F64, DType::F64];
8509 const FAB_F64_A: &[DType] = &[
8510 DType::F64,
8511 DType::F64,
8512 DType::F64,
8513 DType::F64,
8514 DType::F64,
8515 DType::F64,
8516 ];
8517 const FAB_BF16_NOA: &[DType] = &[
8518 DType::BF16,
8519 DType::BF16,
8520 DType::BF16,
8521 DType::BF16,
8522 DType::BF16,
8523 ];
8524 const FAB_BF16_A: &[DType] = &[
8525 DType::BF16,
8526 DType::BF16,
8527 DType::BF16,
8528 DType::BF16,
8529 DType::BF16,
8530 DType::BF16,
8531 ];
8532 const FAB_F16_NOA: &[DType] = &[DType::F16, DType::F16, DType::F16, DType::F16, DType::F16];
8533 const FAB_F16_A: &[DType] = &[
8534 DType::F16,
8535 DType::F16,
8536 DType::F16,
8537 DType::F16,
8538 DType::F16,
8539 DType::F16,
8540 ];
8541
8542 const PA_F32_NOA: &[DType] = &[
8544 DType::F32,
8545 DType::F32,
8546 DType::F32,
8547 DType::U32,
8548 DType::U32,
8549 DType::F32,
8550 ];
8551 const PA_F32_A: &[DType] = &[
8552 DType::F32,
8553 DType::F32,
8554 DType::F32,
8555 DType::U32,
8556 DType::U32,
8557 DType::F32,
8558 DType::F32,
8559 ];
8560 const PA_F64_NOA: &[DType] = &[
8561 DType::F64,
8562 DType::F64,
8563 DType::F64,
8564 DType::U32,
8565 DType::U32,
8566 DType::F64,
8567 ];
8568 const PA_F64_A: &[DType] = &[
8569 DType::F64,
8570 DType::F64,
8571 DType::F64,
8572 DType::U32,
8573 DType::U32,
8574 DType::F64,
8575 DType::F64,
8576 ];
8577 const PA_BF16_NOA: &[DType] = &[
8578 DType::BF16,
8579 DType::BF16,
8580 DType::BF16,
8581 DType::U32,
8582 DType::U32,
8583 DType::BF16,
8584 ];
8585 const PA_BF16_A: &[DType] = &[
8586 DType::BF16,
8587 DType::BF16,
8588 DType::BF16,
8589 DType::U32,
8590 DType::U32,
8591 DType::BF16,
8592 DType::BF16,
8593 ];
8594 const PA_F16_NOA: &[DType] = &[
8595 DType::F16,
8596 DType::F16,
8597 DType::F16,
8598 DType::U32,
8599 DType::U32,
8600 DType::F16,
8601 ];
8602 const PA_F16_A: &[DType] = &[
8603 DType::F16,
8604 DType::F16,
8605 DType::F16,
8606 DType::U32,
8607 DType::U32,
8608 DType::F16,
8609 DType::F16,
8610 ];
8611
8612 let cpu = BackendId::Cpu;
8619
8620 register_cpu_linear_quant_fused_from_contract(r);
8628
8629 register_cpu_norm_softmax_fused_from_contract(r);
8636
8637 register_cpu_conv_rope_fused_from_contract(r);
8646
8647 register_fused!(
8662 r,
8663 FusedOps::FLASH_ATTN,
8664 cpu,
8665 FA_F32_NOA,
8666 flash_attn_f32_cpu_wrapper,
8667 cost = cost_attn_cpu,
8668 precision = ATTN_CPU_PRECISION
8669 );
8670 register_fused!(
8671 r,
8672 FusedOps::FLASH_ATTN,
8673 cpu,
8674 FA_F32_A,
8675 flash_attn_f32_cpu_wrapper,
8676 cost = cost_attn_cpu,
8677 precision = ATTN_CPU_PRECISION
8678 );
8679 register_fused!(
8680 r,
8681 FusedOps::FLASH_ATTN,
8682 cpu,
8683 FA_F64_NOA,
8684 flash_attn_f64_cpu_wrapper,
8685 cost = cost_attn_cpu,
8686 precision = ATTN_CPU_PRECISION
8687 );
8688 register_fused!(
8689 r,
8690 FusedOps::FLASH_ATTN,
8691 cpu,
8692 FA_F64_A,
8693 flash_attn_f64_cpu_wrapper,
8694 cost = cost_attn_cpu,
8695 precision = ATTN_CPU_PRECISION
8696 );
8697 register_fused!(
8698 r,
8699 FusedOps::FLASH_ATTN,
8700 cpu,
8701 FA_BF16_NOA,
8702 flash_attn_bf16_cpu_wrapper,
8703 cost = cost_attn_cpu,
8704 precision = ATTN_CPU_PRECISION
8705 );
8706 register_fused!(
8707 r,
8708 FusedOps::FLASH_ATTN,
8709 cpu,
8710 FA_BF16_A,
8711 flash_attn_bf16_cpu_wrapper,
8712 cost = cost_attn_cpu,
8713 precision = ATTN_CPU_PRECISION
8714 );
8715 register_fused!(
8716 r,
8717 FusedOps::FLASH_ATTN,
8718 cpu,
8719 FA_F16_NOA,
8720 flash_attn_f16_cpu_wrapper,
8721 cost = cost_attn_cpu,
8722 precision = ATTN_CPU_PRECISION
8723 );
8724 register_fused!(
8725 r,
8726 FusedOps::FLASH_ATTN,
8727 cpu,
8728 FA_F16_A,
8729 flash_attn_f16_cpu_wrapper,
8730 cost = cost_attn_cpu,
8731 precision = ATTN_CPU_PRECISION
8732 );
8733
8734 register_fused!(
8739 r,
8740 FusedOps::FLASH_ATTN_BACKWARD_Q,
8741 cpu,
8742 FAB_F32_NOA,
8743 flash_attn_backward_q_f32_cpu_wrapper,
8744 cost = cost_attn_backward_cpu,
8745 precision = ATTN_BACKWARD_CPU_PRECISION
8746 );
8747 register_fused!(
8748 r,
8749 FusedOps::FLASH_ATTN_BACKWARD_Q,
8750 cpu,
8751 FAB_F32_A,
8752 flash_attn_backward_q_f32_cpu_wrapper,
8753 cost = cost_attn_backward_cpu,
8754 precision = ATTN_BACKWARD_CPU_PRECISION
8755 );
8756 register_fused!(
8757 r,
8758 FusedOps::FLASH_ATTN_BACKWARD_Q,
8759 cpu,
8760 FAB_F64_NOA,
8761 flash_attn_backward_q_f64_cpu_wrapper,
8762 cost = cost_attn_backward_cpu,
8763 precision = ATTN_BACKWARD_CPU_PRECISION
8764 );
8765 register_fused!(
8766 r,
8767 FusedOps::FLASH_ATTN_BACKWARD_Q,
8768 cpu,
8769 FAB_F64_A,
8770 flash_attn_backward_q_f64_cpu_wrapper,
8771 cost = cost_attn_backward_cpu,
8772 precision = ATTN_BACKWARD_CPU_PRECISION
8773 );
8774 register_fused!(
8775 r,
8776 FusedOps::FLASH_ATTN_BACKWARD_Q,
8777 cpu,
8778 FAB_BF16_NOA,
8779 flash_attn_backward_q_bf16_cpu_wrapper,
8780 cost = cost_attn_backward_cpu,
8781 precision = ATTN_BACKWARD_CPU_PRECISION
8782 );
8783 register_fused!(
8784 r,
8785 FusedOps::FLASH_ATTN_BACKWARD_Q,
8786 cpu,
8787 FAB_BF16_A,
8788 flash_attn_backward_q_bf16_cpu_wrapper,
8789 cost = cost_attn_backward_cpu,
8790 precision = ATTN_BACKWARD_CPU_PRECISION
8791 );
8792 register_fused!(
8793 r,
8794 FusedOps::FLASH_ATTN_BACKWARD_Q,
8795 cpu,
8796 FAB_F16_NOA,
8797 flash_attn_backward_q_f16_cpu_wrapper,
8798 cost = cost_attn_backward_cpu,
8799 precision = ATTN_BACKWARD_CPU_PRECISION
8800 );
8801 register_fused!(
8802 r,
8803 FusedOps::FLASH_ATTN_BACKWARD_Q,
8804 cpu,
8805 FAB_F16_A,
8806 flash_attn_backward_q_f16_cpu_wrapper,
8807 cost = cost_attn_backward_cpu,
8808 precision = ATTN_BACKWARD_CPU_PRECISION
8809 );
8810
8811 register_fused!(
8812 r,
8813 FusedOps::FLASH_ATTN_BACKWARD_K,
8814 cpu,
8815 FAB_F32_NOA,
8816 flash_attn_backward_k_f32_cpu_wrapper,
8817 cost = cost_attn_backward_cpu,
8818 precision = ATTN_BACKWARD_CPU_PRECISION
8819 );
8820 register_fused!(
8821 r,
8822 FusedOps::FLASH_ATTN_BACKWARD_K,
8823 cpu,
8824 FAB_F32_A,
8825 flash_attn_backward_k_f32_cpu_wrapper,
8826 cost = cost_attn_backward_cpu,
8827 precision = ATTN_BACKWARD_CPU_PRECISION
8828 );
8829 register_fused!(
8830 r,
8831 FusedOps::FLASH_ATTN_BACKWARD_K,
8832 cpu,
8833 FAB_F64_NOA,
8834 flash_attn_backward_k_f64_cpu_wrapper,
8835 cost = cost_attn_backward_cpu,
8836 precision = ATTN_BACKWARD_CPU_PRECISION
8837 );
8838 register_fused!(
8839 r,
8840 FusedOps::FLASH_ATTN_BACKWARD_K,
8841 cpu,
8842 FAB_F64_A,
8843 flash_attn_backward_k_f64_cpu_wrapper,
8844 cost = cost_attn_backward_cpu,
8845 precision = ATTN_BACKWARD_CPU_PRECISION
8846 );
8847 register_fused!(
8848 r,
8849 FusedOps::FLASH_ATTN_BACKWARD_K,
8850 cpu,
8851 FAB_BF16_NOA,
8852 flash_attn_backward_k_bf16_cpu_wrapper,
8853 cost = cost_attn_backward_cpu,
8854 precision = ATTN_BACKWARD_CPU_PRECISION
8855 );
8856 register_fused!(
8857 r,
8858 FusedOps::FLASH_ATTN_BACKWARD_K,
8859 cpu,
8860 FAB_BF16_A,
8861 flash_attn_backward_k_bf16_cpu_wrapper,
8862 cost = cost_attn_backward_cpu,
8863 precision = ATTN_BACKWARD_CPU_PRECISION
8864 );
8865 register_fused!(
8866 r,
8867 FusedOps::FLASH_ATTN_BACKWARD_K,
8868 cpu,
8869 FAB_F16_NOA,
8870 flash_attn_backward_k_f16_cpu_wrapper,
8871 cost = cost_attn_backward_cpu,
8872 precision = ATTN_BACKWARD_CPU_PRECISION
8873 );
8874 register_fused!(
8875 r,
8876 FusedOps::FLASH_ATTN_BACKWARD_K,
8877 cpu,
8878 FAB_F16_A,
8879 flash_attn_backward_k_f16_cpu_wrapper,
8880 cost = cost_attn_backward_cpu,
8881 precision = ATTN_BACKWARD_CPU_PRECISION
8882 );
8883
8884 register_fused!(
8885 r,
8886 FusedOps::FLASH_ATTN_BACKWARD_V,
8887 cpu,
8888 FAB_F32_NOA,
8889 flash_attn_backward_v_f32_cpu_wrapper,
8890 cost = cost_attn_backward_cpu,
8891 precision = ATTN_BACKWARD_CPU_PRECISION
8892 );
8893 register_fused!(
8894 r,
8895 FusedOps::FLASH_ATTN_BACKWARD_V,
8896 cpu,
8897 FAB_F32_A,
8898 flash_attn_backward_v_f32_cpu_wrapper,
8899 cost = cost_attn_backward_cpu,
8900 precision = ATTN_BACKWARD_CPU_PRECISION
8901 );
8902 register_fused!(
8903 r,
8904 FusedOps::FLASH_ATTN_BACKWARD_V,
8905 cpu,
8906 FAB_F64_NOA,
8907 flash_attn_backward_v_f64_cpu_wrapper,
8908 cost = cost_attn_backward_cpu,
8909 precision = ATTN_BACKWARD_CPU_PRECISION
8910 );
8911 register_fused!(
8912 r,
8913 FusedOps::FLASH_ATTN_BACKWARD_V,
8914 cpu,
8915 FAB_F64_A,
8916 flash_attn_backward_v_f64_cpu_wrapper,
8917 cost = cost_attn_backward_cpu,
8918 precision = ATTN_BACKWARD_CPU_PRECISION
8919 );
8920 register_fused!(
8921 r,
8922 FusedOps::FLASH_ATTN_BACKWARD_V,
8923 cpu,
8924 FAB_BF16_NOA,
8925 flash_attn_backward_v_bf16_cpu_wrapper,
8926 cost = cost_attn_backward_cpu,
8927 precision = ATTN_BACKWARD_CPU_PRECISION
8928 );
8929 register_fused!(
8930 r,
8931 FusedOps::FLASH_ATTN_BACKWARD_V,
8932 cpu,
8933 FAB_BF16_A,
8934 flash_attn_backward_v_bf16_cpu_wrapper,
8935 cost = cost_attn_backward_cpu,
8936 precision = ATTN_BACKWARD_CPU_PRECISION
8937 );
8938 register_fused!(
8939 r,
8940 FusedOps::FLASH_ATTN_BACKWARD_V,
8941 cpu,
8942 FAB_F16_NOA,
8943 flash_attn_backward_v_f16_cpu_wrapper,
8944 cost = cost_attn_backward_cpu,
8945 precision = ATTN_BACKWARD_CPU_PRECISION
8946 );
8947 register_fused!(
8948 r,
8949 FusedOps::FLASH_ATTN_BACKWARD_V,
8950 cpu,
8951 FAB_F16_A,
8952 flash_attn_backward_v_f16_cpu_wrapper,
8953 cost = cost_attn_backward_cpu,
8954 precision = ATTN_BACKWARD_CPU_PRECISION
8955 );
8956
8957 register_fused!(
8959 r,
8960 FusedOps::PAGED_ATTN,
8961 cpu,
8962 PA_F32_NOA,
8963 paged_attn_f32_cpu_wrapper,
8964 cost = cost_attn_cpu,
8965 precision = ATTN_CPU_PRECISION
8966 );
8967 register_fused!(
8968 r,
8969 FusedOps::PAGED_ATTN,
8970 cpu,
8971 PA_F32_A,
8972 paged_attn_f32_cpu_wrapper,
8973 cost = cost_attn_cpu,
8974 precision = ATTN_CPU_PRECISION
8975 );
8976 register_fused!(
8977 r,
8978 FusedOps::PAGED_ATTN,
8979 cpu,
8980 PA_F64_NOA,
8981 paged_attn_f64_cpu_wrapper,
8982 cost = cost_attn_cpu,
8983 precision = ATTN_CPU_PRECISION
8984 );
8985 register_fused!(
8986 r,
8987 FusedOps::PAGED_ATTN,
8988 cpu,
8989 PA_F64_A,
8990 paged_attn_f64_cpu_wrapper,
8991 cost = cost_attn_cpu,
8992 precision = ATTN_CPU_PRECISION
8993 );
8994 register_fused!(
8995 r,
8996 FusedOps::PAGED_ATTN,
8997 cpu,
8998 PA_BF16_NOA,
8999 paged_attn_bf16_cpu_wrapper,
9000 cost = cost_attn_cpu,
9001 precision = ATTN_CPU_PRECISION
9002 );
9003 register_fused!(
9004 r,
9005 FusedOps::PAGED_ATTN,
9006 cpu,
9007 PA_BF16_A,
9008 paged_attn_bf16_cpu_wrapper,
9009 cost = cost_attn_cpu,
9010 precision = ATTN_CPU_PRECISION
9011 );
9012 register_fused!(
9013 r,
9014 FusedOps::PAGED_ATTN,
9015 cpu,
9016 PA_F16_NOA,
9017 paged_attn_f16_cpu_wrapper,
9018 cost = cost_attn_cpu,
9019 precision = ATTN_CPU_PRECISION
9020 );
9021 register_fused!(
9022 r,
9023 FusedOps::PAGED_ATTN,
9024 cpu,
9025 PA_F16_A,
9026 paged_attn_f16_cpu_wrapper,
9027 cost = cost_attn_cpu,
9028 precision = ATTN_CPU_PRECISION
9029 );
9030
9031 const NF4_F32: &[DType] = &[DType::F32, DType::U8, DType::F32, DType::F32];
9042 const NF4_F16: &[DType] = &[DType::F16, DType::U8, DType::F32, DType::F16];
9043 const NF4_BF16: &[DType] = &[DType::BF16, DType::U8, DType::F32, DType::BF16];
9044 register_fused!(
9045 r,
9046 FusedOps::NF4_MATMUL,
9047 cpu,
9048 NF4_F32,
9049 nf4_matmul_f32_cpu_wrapper,
9050 cost = cost_nf4_matmul_cpu,
9051 precision = NF4_MATMUL_CPU_PRECISION
9052 );
9053 register_fused!(
9054 r,
9055 FusedOps::NF4_MATMUL,
9056 cpu,
9057 NF4_F16,
9058 nf4_matmul_f16_cpu_wrapper,
9059 cost = cost_nf4_matmul_cpu,
9060 precision = NF4_MATMUL_CPU_PRECISION
9061 );
9062 register_fused!(
9063 r,
9064 FusedOps::NF4_MATMUL,
9065 cpu,
9066 NF4_BF16,
9067 nf4_matmul_bf16_cpu_wrapper,
9068 cost = cost_nf4_matmul_cpu,
9069 precision = NF4_MATMUL_CPU_PRECISION
9070 );
9071}
9072
9073#[cfg(test)]
9074mod tests {
9075 use super::*;
9076
9077 #[test]
9105 fn linear_quant_fused_family_migrated_to_fkc_contract() {
9106 use crate::fused::{FusedKernelRegistry, KernelRevisionHash};
9107 use fuel_graph::registry::FusedOps;
9108
9109 let mut r = FusedKernelRegistry::new();
9110 register_default_fused_kernels(&mut r);
9111
9112 for (dt, expected) in [
9114 (
9115 DType::F32,
9116 fused_linear_f32_cpu_wrapper as *const () as usize,
9117 ),
9118 (
9119 DType::F64,
9120 fused_linear_f64_cpu_wrapper as *const () as usize,
9121 ),
9122 (
9123 DType::BF16,
9124 fused_linear_bf16_cpu_wrapper as *const () as usize,
9125 ),
9126 (
9127 DType::F16,
9128 fused_linear_f16_cpu_wrapper as *const () as usize,
9129 ),
9130 ] {
9131 let got = r
9132 .lookup_by_dtypes(FusedOps::FUSED_LINEAR, BackendId::Cpu, &[dt, dt, dt, dt])
9133 .unwrap_or_else(|| panic!("FUSED_LINEAR {dt:?} migrated impl present"));
9134 assert_eq!(
9135 got.kernel as usize, expected,
9136 "FUSED_LINEAR {dt:?} binds its exact per-dtype production wrapper",
9137 );
9138 assert_ne!(
9139 got.revision,
9140 KernelRevisionHash::UNTRACKED,
9141 "FKC-imported FUSED_LINEAR {dt:?} carries the contract's real revision \
9142 (hand-written register_fused! stamps UNTRACKED)",
9143 );
9144 }
9145
9146 let qm = r
9148 .lookup_by_dtypes(
9149 FusedOps::QMATMUL,
9150 BackendId::Cpu,
9151 &[DType::F32, DType::U32, DType::F32],
9152 )
9153 .expect("QMATMUL migrated impl present at [F32, U32, F32]");
9154 assert_eq!(
9155 qm.kernel as usize,
9156 qmatmul_f32_cpu_wrapper as *const () as usize
9157 );
9158 assert_ne!(qm.revision, KernelRevisionHash::UNTRACKED);
9159
9160 for (dt, expected) in [
9162 (
9163 DType::F32,
9164 inplace_affine_f32_cpu_wrapper as *const () as usize,
9165 ),
9166 (
9167 DType::F64,
9168 inplace_affine_f64_cpu_wrapper as *const () as usize,
9169 ),
9170 (
9171 DType::BF16,
9172 inplace_affine_bf16_cpu_wrapper as *const () as usize,
9173 ),
9174 (
9175 DType::F16,
9176 inplace_affine_f16_cpu_wrapper as *const () as usize,
9177 ),
9178 ] {
9179 let got = r
9180 .lookup_by_dtypes(FusedOps::INPLACE_AFFINE, BackendId::Cpu, &[dt, dt])
9181 .unwrap_or_else(|| panic!("INPLACE_AFFINE {dt:?} migrated impl present"));
9182 assert_eq!(got.kernel as usize, expected);
9183 assert_ne!(got.revision, KernelRevisionHash::UNTRACKED);
9184 }
9185
9186 for (dt, expected) in [
9188 (
9189 DType::F32,
9190 fused_softmax_cross_entropy_f32_cpu_wrapper as *const () as usize,
9191 ),
9192 (
9193 DType::F64,
9194 fused_softmax_cross_entropy_f64_cpu_wrapper as *const () as usize,
9195 ),
9196 (
9197 DType::BF16,
9198 fused_softmax_cross_entropy_bf16_cpu_wrapper as *const () as usize,
9199 ),
9200 (
9201 DType::F16,
9202 fused_softmax_cross_entropy_f16_cpu_wrapper as *const () as usize,
9203 ),
9204 ] {
9205 let got = r
9206 .lookup_by_dtypes(
9207 FusedOps::FUSED_SOFTMAX_CROSS_ENTROPY,
9208 BackendId::Cpu,
9209 &[dt, DType::I64, DType::F32],
9210 )
9211 .unwrap_or_else(|| panic!("FSCE {dt:?} migrated impl present"));
9212 assert_eq!(got.kernel as usize, expected);
9213 assert_ne!(got.revision, KernelRevisionHash::UNTRACKED);
9214 }
9215
9216 let fa = r
9218 .lookup_by_dtypes(
9219 FusedOps::FLASH_ATTN,
9220 BackendId::Cpu,
9221 &[DType::F32, DType::F32, DType::F32, DType::F32],
9222 )
9223 .expect("FLASH_ATTN hand-written impl present");
9224 assert_eq!(
9225 fa.kernel as usize,
9226 flash_attn_f32_cpu_wrapper as *const () as usize
9227 );
9228 assert_eq!(
9229 fa.revision,
9230 KernelRevisionHash::UNTRACKED,
9231 "unmigrated FLASH_ATTN keeps its hand-written UNTRACKED revision",
9232 );
9233
9234 let nf4 = r
9236 .lookup_by_dtypes(
9237 FusedOps::NF4_MATMUL,
9238 BackendId::Cpu,
9239 &[DType::F32, DType::U8, DType::F32, DType::F32],
9240 )
9241 .expect("NF4_MATMUL hand-written impl present (AFFINE_BLOCK deferred)");
9242 assert_eq!(
9243 nf4.kernel as usize,
9244 nf4_matmul_f32_cpu_wrapper as *const () as usize
9245 );
9246 assert_eq!(
9247 nf4.revision,
9248 KernelRevisionHash::UNTRACKED,
9249 "deferred NF4_MATMUL keeps its hand-written UNTRACKED revision",
9250 );
9251 }
9252
9253 #[test]
9278 fn norm_softmax_fused_family_migrated_to_fkc_contract() {
9279 use crate::fused::{FusedKernelRegistry, KernelRevisionHash};
9280 use fuel_graph::registry::FusedOps;
9281
9282 let mut r = FusedKernelRegistry::new();
9283 register_default_fused_kernels(&mut r);
9284
9285 let forward = [
9287 (
9288 FusedOps::SOFTMAX_LAST_DIM,
9289 [
9290 softmax_last_dim_f32_cpu_wrapper as *const () as usize,
9291 softmax_last_dim_f64_cpu_wrapper as *const () as usize,
9292 softmax_last_dim_bf16_cpu_wrapper as *const () as usize,
9293 softmax_last_dim_f16_cpu_wrapper as *const () as usize,
9294 ],
9295 ),
9296 (
9297 FusedOps::RMS_NORM_LAST_DIM,
9298 [
9299 rms_norm_last_dim_f32_cpu_wrapper as *const () as usize,
9300 rms_norm_last_dim_f64_cpu_wrapper as *const () as usize,
9301 rms_norm_last_dim_bf16_cpu_wrapper as *const () as usize,
9302 rms_norm_last_dim_f16_cpu_wrapper as *const () as usize,
9303 ],
9304 ),
9305 (
9306 FusedOps::LAYER_NORM_LAST_DIM,
9307 [
9308 layer_norm_last_dim_f32_cpu_wrapper as *const () as usize,
9309 layer_norm_last_dim_f64_cpu_wrapper as *const () as usize,
9310 layer_norm_last_dim_bf16_cpu_wrapper as *const () as usize,
9311 layer_norm_last_dim_f16_cpu_wrapper as *const () as usize,
9312 ],
9313 ),
9314 ];
9315 for (id, wrappers) in forward {
9316 for (dt, expected) in [DType::F32, DType::F64, DType::BF16, DType::F16]
9317 .iter()
9318 .zip(wrappers)
9319 {
9320 let got = r
9321 .lookup_by_dtypes(id, BackendId::Cpu, &[*dt, *dt])
9322 .unwrap_or_else(|| panic!("{id:?} {dt:?} migrated impl present"));
9323 assert_eq!(
9324 got.kernel as usize, expected,
9325 "{id:?} {dt:?} binds its exact per-dtype production wrapper",
9326 );
9327 assert_ne!(
9328 got.revision,
9329 KernelRevisionHash::UNTRACKED,
9330 "FKC-imported {id:?} {dt:?} carries the contract's real revision \
9331 (hand-written register_fused! stamps UNTRACKED)",
9332 );
9333 assert!(
9334 got.precision.bit_stable_on_same_hardware,
9335 "the 2026-07-03 audited:true flip rode through as bit-stable for {id:?} {dt:?} \
9336 (NOT the UNAUDITED an audited:false contract would have lowered to)",
9337 );
9338 }
9339 }
9340
9341 let backward = [
9344 (
9345 FusedOps::SOFTMAX_LAST_DIM_BACKWARD,
9346 [
9347 softmax_last_dim_backward_f32_cpu_wrapper as *const () as usize,
9348 softmax_last_dim_backward_f64_cpu_wrapper as *const () as usize,
9349 softmax_last_dim_backward_bf16_cpu_wrapper as *const () as usize,
9350 softmax_last_dim_backward_f16_cpu_wrapper as *const () as usize,
9351 ],
9352 ),
9353 (
9354 FusedOps::LAYER_NORM_LAST_DIM_BACKWARD,
9355 [
9356 layer_norm_last_dim_backward_f32_cpu_wrapper as *const () as usize,
9357 layer_norm_last_dim_backward_f64_cpu_wrapper as *const () as usize,
9358 layer_norm_last_dim_backward_bf16_cpu_wrapper as *const () as usize,
9359 layer_norm_last_dim_backward_f16_cpu_wrapper as *const () as usize,
9360 ],
9361 ),
9362 (
9363 FusedOps::RMS_NORM_LAST_DIM_BACKWARD,
9364 [
9365 rms_norm_last_dim_backward_f32_cpu_wrapper as *const () as usize,
9366 rms_norm_last_dim_backward_f64_cpu_wrapper as *const () as usize,
9367 rms_norm_last_dim_backward_bf16_cpu_wrapper as *const () as usize,
9368 rms_norm_last_dim_backward_f16_cpu_wrapper as *const () as usize,
9369 ],
9370 ),
9371 (
9372 FusedOps::REDUCE_MAX_TO_BACKWARD,
9373 [
9374 reduce_max_to_backward_f32_cpu_wrapper as *const () as usize,
9375 reduce_max_to_backward_f64_cpu_wrapper as *const () as usize,
9376 reduce_max_to_backward_bf16_cpu_wrapper as *const () as usize,
9377 reduce_max_to_backward_f16_cpu_wrapper as *const () as usize,
9378 ],
9379 ),
9380 (
9381 FusedOps::POWI_BACKWARD,
9382 [
9383 powi_backward_f32_cpu_wrapper as *const () as usize,
9384 powi_backward_f64_cpu_wrapper as *const () as usize,
9385 powi_backward_bf16_cpu_wrapper as *const () as usize,
9386 powi_backward_f16_cpu_wrapper as *const () as usize,
9387 ],
9388 ),
9389 ];
9390 for (id, wrappers) in backward {
9391 for (dt, expected) in [DType::F32, DType::F64, DType::BF16, DType::F16]
9392 .iter()
9393 .zip(wrappers)
9394 {
9395 let got = r
9396 .lookup_by_dtypes(id, BackendId::Cpu, &[*dt, *dt, *dt])
9397 .unwrap_or_else(|| panic!("{id:?} {dt:?} migrated impl present"));
9398 assert_eq!(got.kernel as usize, expected);
9399 assert_ne!(
9400 got.revision,
9401 KernelRevisionHash::UNTRACKED,
9402 "FKC-imported {id:?} {dt:?} carries the contract's real revision",
9403 );
9404 assert!(
9405 got.precision.bit_stable_on_same_hardware,
9406 "the audited:true flip rode through as bit-stable for {id:?} {dt:?}",
9407 );
9408 }
9409 }
9410
9411 let fa = r
9414 .lookup_by_dtypes(
9415 FusedOps::FLASH_ATTN,
9416 BackendId::Cpu,
9417 &[DType::F32, DType::F32, DType::F32, DType::F32],
9418 )
9419 .expect("FLASH_ATTN hand-written impl present");
9420 assert_eq!(
9421 fa.revision,
9422 KernelRevisionHash::UNTRACKED,
9423 "unmigrated FLASH_ATTN keeps its hand-written UNTRACKED revision",
9424 );
9425 }
9426
9427 #[test]
9446 fn conv_rope_fused_family_migrated_to_fkc_contract() {
9447 use crate::fused::{FusedKernelRegistry, KernelRevisionHash};
9448 use fuel_graph::registry::FusedOps;
9449
9450 let mut r = FusedKernelRegistry::new();
9451 register_default_fused_kernels(&mut r);
9452
9453 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
9454
9455 let rope_w = [
9457 rope_f32_cpu_wrapper as *const () as usize,
9458 rope_f64_cpu_wrapper as *const () as usize,
9459 rope_bf16_cpu_wrapper as *const () as usize,
9460 rope_f16_cpu_wrapper as *const () as usize,
9461 ];
9462 for (dt, expected) in dts.iter().zip(rope_w) {
9463 let got = r
9464 .lookup_by_dtypes(FusedOps::ROPE, BackendId::Cpu, &[*dt, *dt, *dt, *dt])
9465 .unwrap_or_else(|| panic!("ROPE {dt:?} migrated impl present"));
9466 assert_eq!(got.kernel as usize, expected);
9467 assert_ne!(
9468 got.revision,
9469 KernelRevisionHash::UNTRACKED,
9470 "ROPE {dt:?} real revision"
9471 );
9472 assert!(
9473 got.precision.bit_stable_on_same_hardware,
9474 "audited:true rode through bit-stable for ROPE {dt:?}",
9475 );
9476 }
9477
9478 let conv = [
9481 (
9482 FusedOps::CONV2D,
9483 [
9484 conv2d_f32_cpu_wrapper as *const () as usize,
9485 conv2d_f64_cpu_wrapper as *const () as usize,
9486 conv2d_bf16_cpu_wrapper as *const () as usize,
9487 conv2d_f16_cpu_wrapper as *const () as usize,
9488 ],
9489 ),
9490 (
9491 FusedOps::CONV_TRANSPOSE2D,
9492 [
9493 conv_transpose2d_f32_cpu_wrapper as *const () as usize,
9494 conv_transpose2d_f64_cpu_wrapper as *const () as usize,
9495 conv_transpose2d_bf16_cpu_wrapper as *const () as usize,
9496 conv_transpose2d_f16_cpu_wrapper as *const () as usize,
9497 ],
9498 ),
9499 ];
9500 for (id, wrappers) in conv {
9501 for (dt, expected) in dts.iter().zip(wrappers) {
9502 for key in [vec![*dt, *dt, *dt], vec![*dt, *dt, *dt, *dt]] {
9503 let with_bias = key.len() == 4;
9504 let got = r
9505 .lookup_by_dtypes(id, BackendId::Cpu, &key)
9506 .unwrap_or_else(|| {
9507 panic!("{id:?} {dt:?} (with_bias={with_bias}) migrated impl present")
9508 });
9509 assert_eq!(got.kernel as usize, expected);
9510 assert_ne!(
9511 got.revision,
9512 KernelRevisionHash::UNTRACKED,
9513 "{id:?} {dt:?} (with_bias={with_bias}) real revision",
9514 );
9515 assert!(
9516 got.precision.bit_stable_on_same_hardware,
9517 "audited:true rode through bit-stable for {id:?} {dt:?} (with_bias={with_bias})",
9518 );
9519 }
9520 }
9521 }
9522
9523 let cc1d_w = [
9525 causal_conv1d_f32_cpu_wrapper as *const () as usize,
9526 causal_conv1d_f64_cpu_wrapper as *const () as usize,
9527 causal_conv1d_bf16_cpu_wrapper as *const () as usize,
9528 causal_conv1d_f16_cpu_wrapper as *const () as usize,
9529 ];
9530 for (dt, expected) in dts.iter().zip(cc1d_w) {
9531 let got = r
9532 .lookup_by_dtypes(
9533 FusedOps::CAUSAL_CONV1D,
9534 BackendId::Cpu,
9535 &[*dt, *dt, *dt, *dt],
9536 )
9537 .unwrap_or_else(|| panic!("CAUSAL_CONV1D {dt:?} migrated impl present"));
9538 assert_eq!(got.kernel as usize, expected);
9539 assert_ne!(got.revision, KernelRevisionHash::UNTRACKED);
9540 assert!(got.precision.bit_stable_on_same_hardware);
9541 }
9542
9543 let scan = [
9546 (
9547 FusedOps::SELECTIVE_SCAN,
9548 [
9549 selective_scan_f32_cpu_wrapper as *const () as usize,
9550 selective_scan_f64_cpu_wrapper as *const () as usize,
9551 selective_scan_bf16_cpu_wrapper as *const () as usize,
9552 selective_scan_f16_cpu_wrapper as *const () as usize,
9553 ],
9554 ),
9555 (
9556 FusedOps::SSD_CHUNK_SCAN,
9557 [
9558 ssd_chunk_scan_f32_cpu_wrapper as *const () as usize,
9559 ssd_chunk_scan_f64_cpu_wrapper as *const () as usize,
9560 ssd_chunk_scan_bf16_cpu_wrapper as *const () as usize,
9561 ssd_chunk_scan_f16_cpu_wrapper as *const () as usize,
9562 ],
9563 ),
9564 ];
9565 for (id, wrappers) in scan {
9566 for (dt, expected) in dts.iter().zip(wrappers) {
9567 let got = r
9568 .lookup_by_dtypes(id, BackendId::Cpu, &[*dt; 6])
9569 .unwrap_or_else(|| panic!("{id:?} {dt:?} migrated impl present"));
9570 assert_eq!(got.kernel as usize, expected);
9571 assert_ne!(
9572 got.revision,
9573 KernelRevisionHash::UNTRACKED,
9574 "{id:?} {dt:?} real revision",
9575 );
9576 assert!(
9577 got.precision.bit_stable_on_same_hardware,
9578 "audited:true rode through bit-stable for {id:?} {dt:?}",
9579 );
9580 }
9581 }
9582
9583 let nf4 = r
9586 .lookup_by_dtypes(
9587 FusedOps::NF4_MATMUL,
9588 BackendId::Cpu,
9589 &[DType::F32, DType::U8, DType::F32, DType::F32],
9590 )
9591 .expect("NF4_MATMUL hand-written impl present");
9592 assert_eq!(
9593 nf4.revision,
9594 KernelRevisionHash::UNTRACKED,
9595 "deferred NF4_MATMUL keeps its hand-written UNTRACKED revision",
9596 );
9597 }
9598
9599 #[test]
9627 fn precision_guarantee_lint_bit_stable_cpu_coverage_primitives() {
9632 use fuel_ir::dispatch::OpKind;
9633 use fuel_ir::probe::BackendId;
9634
9635 const KNOWN_GAPS: &[(OpKind, &str)] = &[];
9640
9641 const ALL_OP_KINDS: &[OpKind] = &[
9647 OpKind::MatMul,
9648 OpKind::AddElementwise,
9649 OpKind::SubElementwise,
9650 OpKind::MulElementwise,
9651 OpKind::DivElementwise,
9652 OpKind::ReluElementwise,
9653 OpKind::NegElementwise,
9654 OpKind::SqrElementwise,
9655 OpKind::SqrtElementwise,
9656 OpKind::RecipElementwise,
9657 OpKind::AbsElementwise,
9658 OpKind::TanhElementwise,
9659 OpKind::ExpElementwise,
9660 OpKind::LogElementwise,
9661 OpKind::SinElementwise,
9662 OpKind::CosElementwise,
9663 OpKind::SigmoidElementwise,
9664 OpKind::SiluElementwise,
9665 OpKind::GeluTanhElementwise,
9666 OpKind::StepElementwise,
9667 OpKind::SumReduce,
9668 OpKind::MaxReduce,
9669 OpKind::MinReduce,
9670 OpKind::MeanReduce,
9671 OpKind::Cast,
9672 OpKind::Conv2D,
9673 OpKind::ConvTranspose2D,
9674 OpKind::ReduceSumTo,
9675 OpKind::ReduceMaxTo,
9676 OpKind::FusedLinear,
9677 OpKind::FlashAttn,
9678 OpKind::PagedAttn,
9679 OpKind::Affine,
9680 OpKind::ClampElementwise,
9681 OpKind::PowIElementwise,
9682 OpKind::MaximumElementwise,
9683 OpKind::MinimumElementwise,
9684 OpKind::EqualElementwise,
9685 OpKind::NotEqualElementwise,
9686 OpKind::LessElementwise,
9687 OpKind::LessEqualElementwise,
9688 OpKind::GreaterElementwise,
9689 OpKind::GreaterEqualElementwise,
9690 OpKind::Where,
9691 OpKind::FloorElementwise,
9692 OpKind::CeilElementwise,
9693 OpKind::RoundElementwise,
9694 OpKind::SignElementwise,
9695 OpKind::ErfElementwise,
9696 OpKind::GeluElementwise,
9697 OpKind::PowElementwise,
9698 OpKind::RsqrtElementwise,
9699 OpKind::RemElementwise,
9700 OpKind::Flip,
9701 OpKind::Roll,
9702 OpKind::CumSum,
9703 OpKind::Pad,
9704 OpKind::PadBackward,
9705 OpKind::Triu,
9706 OpKind::Tril,
9707 OpKind::LogSoftmaxLastDim,
9708 OpKind::LogSoftmaxLastDimBackward,
9709 OpKind::MaskedFill,
9710 OpKind::Concat,
9711 OpKind::SoftmaxLastDim,
9712 OpKind::SoftmaxLastDimBackward,
9713 OpKind::RmsNormLastDim,
9714 OpKind::RmsNormLastDimBackward,
9715 OpKind::LayerNormLastDim,
9716 OpKind::LayerNormLastDimBackward,
9717 OpKind::ReduceMaxToBackward,
9718 OpKind::IndexSelect,
9719 OpKind::Gather,
9720 OpKind::Rope,
9721 OpKind::IndexAdd,
9722 OpKind::ScatterAdd,
9723 OpKind::ArgMaxDim,
9724 OpKind::ArgMinDim,
9725 OpKind::QMatMul,
9726 OpKind::Copy,
9727 OpKind::FusedSoftmaxCrossEntropy,
9728 OpKind::CausalConv1d,
9729 OpKind::SelectiveScan,
9730 OpKind::SsdChunkScan,
9731 OpKind::Nf4Matmul,
9732 ];
9733
9734 let mut table = KernelBindingTable::new();
9738 register_cpu_kernels(&mut table);
9739
9740 let mut by_op_kind: std::collections::HashMap<
9743 OpKind,
9744 Vec<crate::fused::PrecisionGuarantee>,
9745 > = std::collections::HashMap::new();
9746 for (op, _dtypes, backend, precision) in table.iter_precision() {
9747 if backend == BackendId::Cpu {
9748 by_op_kind.entry(op).or_default().push(precision);
9749 }
9750 }
9751
9752 let mut failures: Vec<String> = Vec::new();
9753 let mut covered = 0usize;
9754 let mut allowlisted = 0usize;
9755
9756 for op in ALL_OP_KINDS.iter().copied() {
9757 if let Some((_, reason)) = KNOWN_GAPS.iter().find(|(g, _)| *g == op) {
9758 allowlisted += 1;
9759 if by_op_kind.contains_key(&op) {
9760 failures.push(format!(
9761 "OpKind::{op:?} is on the KNOWN_GAPS allowlist but DOES \
9762 have a CPU registration now. Reason given was: \
9763 {reason:?}. Remove the allowlist entry to enable the \
9764 bit-stable lint for this op.",
9765 ));
9766 }
9767 continue;
9768 }
9769 let precisions = by_op_kind.get(&op);
9770 let has_bit_stable_cpu =
9771 precisions.is_some_and(|ps| ps.iter().any(|p| p.bit_stable_on_same_hardware));
9772 if has_bit_stable_cpu {
9773 covered += 1;
9774 } else {
9775 failures.push(format!(
9776 "OpKind::{op:?} has no bit-stable CPU registration. \
9777 Architecture v1.0 §05 requires the always-built backend \
9778 (fuel-cpu-backend) to provide at least one \
9779 `bit_stable_on_same_hardware: true` kernel per primitive \
9780 op. Either add a CPU registration in \
9781 register_cpu_kernels (the fill_unset_cpu_precision pass \
9782 at the end will upgrade it to PRIMITIVE_DETERMINISTIC_CPU \
9783 automatically), or add a line to KNOWN_GAPS above with a \
9784 documented reason.",
9785 ));
9786 }
9787 }
9788
9789 assert!(
9790 failures.is_empty(),
9791 "Architecture v1.0 bit-stable CPU coverage lint (primitives) failed:\n{}",
9792 failures.join("\n"),
9793 );
9794 assert_eq!(
9795 allowlisted, 0,
9796 "KNOWN_GAPS allowlist is empty by design; saw {allowlisted} allowlisted",
9797 );
9798 assert!(covered > 0, "lint covered 0 OpKinds — table appears empty",);
9803 for op in ALL_OP_KINDS {
9807 assert!(
9808 by_op_kind.contains_key(op),
9809 "OpKind::{op:?} has no CPU registration at all — either add \
9810 one in register_cpu_kernels or document the gap in \
9811 KNOWN_GAPS.",
9812 );
9813 }
9814 }
9815
9816 #[test]
9835 fn cost_lint_per_op_kind_cpu_coverage() {
9836 use fuel_ir::dispatch::OpKind;
9837 use fuel_ir::probe::BackendId;
9838
9839 const KNOWN_GAPS: &[(OpKind, &str)] = &[];
9840
9841 const ALL_OP_KINDS: &[OpKind] = &[
9842 OpKind::MatMul,
9843 OpKind::AddElementwise,
9844 OpKind::SubElementwise,
9845 OpKind::MulElementwise,
9846 OpKind::DivElementwise,
9847 OpKind::ReluElementwise,
9848 OpKind::NegElementwise,
9849 OpKind::SqrElementwise,
9850 OpKind::SqrtElementwise,
9851 OpKind::RecipElementwise,
9852 OpKind::AbsElementwise,
9853 OpKind::TanhElementwise,
9854 OpKind::ExpElementwise,
9855 OpKind::LogElementwise,
9856 OpKind::SinElementwise,
9857 OpKind::CosElementwise,
9858 OpKind::SigmoidElementwise,
9859 OpKind::SiluElementwise,
9860 OpKind::GeluTanhElementwise,
9861 OpKind::StepElementwise,
9862 OpKind::SumReduce,
9863 OpKind::MaxReduce,
9864 OpKind::MinReduce,
9865 OpKind::MeanReduce,
9866 OpKind::Cast,
9867 OpKind::Conv2D,
9868 OpKind::ConvTranspose2D,
9869 OpKind::ReduceSumTo,
9870 OpKind::ReduceMaxTo,
9871 OpKind::FusedLinear,
9872 OpKind::FlashAttn,
9873 OpKind::PagedAttn,
9874 OpKind::Affine,
9875 OpKind::ClampElementwise,
9876 OpKind::PowIElementwise,
9877 OpKind::MaximumElementwise,
9878 OpKind::MinimumElementwise,
9879 OpKind::EqualElementwise,
9880 OpKind::NotEqualElementwise,
9881 OpKind::LessElementwise,
9882 OpKind::LessEqualElementwise,
9883 OpKind::GreaterElementwise,
9884 OpKind::GreaterEqualElementwise,
9885 OpKind::Where,
9886 OpKind::FloorElementwise,
9887 OpKind::CeilElementwise,
9888 OpKind::RoundElementwise,
9889 OpKind::SignElementwise,
9890 OpKind::ErfElementwise,
9891 OpKind::GeluElementwise,
9892 OpKind::PowElementwise,
9893 OpKind::RsqrtElementwise,
9894 OpKind::RemElementwise,
9895 OpKind::Flip,
9896 OpKind::Roll,
9897 OpKind::CumSum,
9898 OpKind::Pad,
9899 OpKind::PadBackward,
9900 OpKind::Triu,
9901 OpKind::Tril,
9902 OpKind::LogSoftmaxLastDim,
9903 OpKind::LogSoftmaxLastDimBackward,
9904 OpKind::MaskedFill,
9905 OpKind::Concat,
9906 OpKind::SoftmaxLastDim,
9907 OpKind::SoftmaxLastDimBackward,
9908 OpKind::RmsNormLastDim,
9909 OpKind::RmsNormLastDimBackward,
9910 OpKind::LayerNormLastDim,
9911 OpKind::LayerNormLastDimBackward,
9912 OpKind::ReduceMaxToBackward,
9913 OpKind::IndexSelect,
9914 OpKind::Gather,
9915 OpKind::Rope,
9916 OpKind::IndexAdd,
9917 OpKind::ScatterAdd,
9918 OpKind::ArgMaxDim,
9919 OpKind::ArgMinDim,
9920 OpKind::QMatMul,
9921 OpKind::Copy,
9922 OpKind::FusedSoftmaxCrossEntropy,
9923 OpKind::CausalConv1d,
9924 OpKind::SelectiveScan,
9925 OpKind::SsdChunkScan,
9926 OpKind::Nf4Matmul,
9927 ];
9928
9929 let mut table = KernelBindingTable::new();
9930 register_cpu_kernels(&mut table);
9931
9932 let mut by_op_kind: std::collections::HashMap<OpKind, Vec<crate::kernel::CostFn>> =
9934 std::collections::HashMap::new();
9935 for (op, _dtypes, backend, cost) in table.iter_cost() {
9936 if backend == BackendId::Cpu {
9937 by_op_kind.entry(op).or_default().push(cost);
9938 }
9939 }
9940
9941 let unknown_sentinel = crate::kernel::unknown_cost as *const () as usize;
9942
9943 let mut failures: Vec<String> = Vec::new();
9944 for op in ALL_OP_KINDS.iter().copied() {
9945 if let Some((_, reason)) = KNOWN_GAPS.iter().find(|(g, _)| *g == op) {
9946 if let Some(costs) = by_op_kind.get(&op) {
9947 let has_real = costs.iter().any(|c| (*c as usize) != unknown_sentinel);
9948 if has_real {
9949 failures.push(format!(
9950 "OpKind::{op:?} is on the cost-lint KNOWN_GAPS \
9951 allowlist but DOES have a non-default cost fn \
9952 now. Reason given: {reason:?}. Remove the \
9953 allowlist entry.",
9954 ));
9955 }
9956 }
9957 continue;
9958 }
9959 let costs = by_op_kind.get(&op);
9960 let has_real_cost =
9961 costs.is_some_and(|cs| cs.iter().any(|c| (*c as usize) != unknown_sentinel));
9962 if !has_real_cost {
9963 failures.push(format!(
9964 "OpKind::{op:?} has no non-default cost fn in any CPU \
9965 registration. Either add an arm in \
9966 `crate::cost::default_cost_for_op_kind` (the \
9967 fill_unset_cpu_cost pass at the end of \
9968 register_cpu_kernels will pick it up automatically), \
9969 register an explicit cost via `register_full(...)`, \
9970 or add a line to KNOWN_GAPS with a documented reason.",
9971 ));
9972 }
9973 }
9974
9975 assert!(
9976 failures.is_empty(),
9977 "Phase 7.6 step 8 cost-coverage lint failed:\n{}",
9978 failures.join("\n"),
9979 );
9980 for op in ALL_OP_KINDS {
9984 assert!(
9985 by_op_kind.contains_key(op),
9986 "OpKind::{op:?} has no CPU registration at all.",
9987 );
9988 }
9989 }
9990
9991 fn cpu_caps() -> BackendCapabilities {
9992 BackendCapabilities {
9993 backend_id: BackendId::Cpu,
9994 device_location: DeviceLocation::Cpu,
9995 op_dtype_support: [
9996 (OpKind::MatMul, DType::F32),
9997 (OpKind::MatMul, DType::F64),
9998 (OpKind::AddElementwise, DType::F32),
9999 (OpKind::AddElementwise, DType::F64),
10000 ]
10001 .into_iter()
10002 .collect(),
10003 required_alignment: 64,
10004 access_granularity_bits: 8,
10005 transfer_paths: vec![(DeviceLocation::Cpu, TransferPath::SameDevice)],
10006 storage_substrate: SubstrateClass::HostBytes,
10007 compute_throughput_flops_per_ns: 1.0,
10008 mem_bandwidth_bytes_per_ns: 4.0,
10009 }
10010 }
10011
10012 fn cuda_caps() -> BackendCapabilities {
10013 BackendCapabilities {
10014 backend_id: BackendId::Cuda,
10015 device_location: DeviceLocation::Cuda { gpu_id: 0 },
10016 op_dtype_support: [
10017 (OpKind::MatMul, DType::F32),
10018 (OpKind::MatMul, DType::F16),
10019 (OpKind::AddElementwise, DType::F32),
10020 ]
10021 .into_iter()
10022 .collect(),
10023 required_alignment: 256,
10024 access_granularity_bits: 8,
10025 transfer_paths: vec![
10026 (DeviceLocation::Cpu, TransferPath::DeviceCopy),
10027 (DeviceLocation::Cuda { gpu_id: 0 }, TransferPath::SameDevice),
10028 ],
10029 storage_substrate: SubstrateClass::CudaUntyped,
10030 compute_throughput_flops_per_ns: 30.0,
10031 mem_bandwidth_bytes_per_ns: 40.0,
10032 }
10033 }
10034
10035 #[test]
10037 fn empty_registry() {
10038 let r = CapabilityRegistry::new();
10039 assert!(r.backends().is_empty());
10040 let result = r.find_backend_for(OpKind::MatMul, DType::F32);
10041 assert!(result.is_err());
10042 }
10043
10044 #[test]
10046 fn find_backends_order_preserved() {
10047 let mut r = CapabilityRegistry::new();
10048 r.register(cuda_caps()); r.register(cpu_caps());
10050
10051 let backends = r.find_backends(OpKind::MatMul, DType::F32);
10052 assert_eq!(backends.len(), 2);
10053 assert_eq!(backends[0].backend_id, BackendId::Cuda);
10054 assert_eq!(backends[1].backend_id, BackendId::Cpu);
10055 }
10056
10057 #[test]
10059 fn find_backend_first_wins() {
10060 let mut r = CapabilityRegistry::new();
10061 r.register(cuda_caps());
10062 r.register(cpu_caps());
10063
10064 let chosen = r.find_backend_for(OpKind::MatMul, DType::F32).unwrap();
10065 assert_eq!(chosen.backend_id, BackendId::Cuda);
10066 }
10067
10068 #[test]
10070 fn cpu_handles_what_gpu_doesnt() {
10071 let mut r = CapabilityRegistry::new();
10072 r.register(cuda_caps());
10073 r.register(cpu_caps());
10074
10075 let chosen = r.find_backend_for(OpKind::MatMul, DType::F64).unwrap();
10077 assert_eq!(chosen.backend_id, BackendId::Cpu);
10078 }
10079
10080 #[test]
10082 fn unsupported_combo_errors_with_diagnostic() {
10083 let mut r = CapabilityRegistry::new();
10084 r.register(cpu_caps());
10085
10086 let err = match r.find_backend_for(OpKind::MatMul, DType::BF16) {
10087 Ok(_) => panic!("expected error"),
10088 Err(e) => e,
10089 };
10090 let msg = err.to_string();
10091 assert!(
10092 msg.contains("MatMul") || msg.contains("matmul"),
10093 "error names op: {msg}"
10094 );
10095 assert!(msg.contains("BF16"), "error names dtype: {msg}");
10096 assert!(msg.contains("Cpu"), "error names available backends: {msg}");
10097 }
10098
10099 #[test]
10101 fn transfer_matrix_same_device() {
10102 let r = CapabilityRegistry::new();
10103 let m = r.build_transfer_matrix();
10104 assert_eq!(
10105 m.path(DeviceLocation::Cpu, DeviceLocation::Cpu),
10106 Some(TransferPath::SameDevice)
10107 );
10108 }
10109
10110 #[test]
10112 fn transfer_matrix_built_from_caps() {
10113 let mut r = CapabilityRegistry::new();
10114 r.register(cuda_caps());
10115 let m = r.build_transfer_matrix();
10116
10117 assert_eq!(
10119 m.path(DeviceLocation::Cuda { gpu_id: 0 }, DeviceLocation::Cpu),
10120 Some(TransferPath::DeviceCopy)
10121 );
10122 }
10123
10124 #[test]
10127 fn host_staging_fallback() {
10128 let r = CapabilityRegistry::new();
10129 let m = r.build_transfer_matrix();
10130 assert_eq!(
10132 m.path_or_staging(DeviceLocation::Cpu, DeviceLocation::Cuda { gpu_id: 0 }),
10133 TransferPath::HostStaging
10134 );
10135 }
10136
10137 #[test]
10140 fn resolve_picks_first_match() {
10141 let mut r = CapabilityRegistry::new();
10142 r.register(cuda_caps());
10143 r.register(cpu_caps());
10144 assert_eq!(
10145 resolve_target_backend(&r, OpKind::MatMul, DType::F32).unwrap(),
10146 BackendId::Cuda
10147 );
10148 }
10149
10150 #[test]
10152 fn resolve_falls_back_to_cpu() {
10153 let mut r = CapabilityRegistry::new();
10154 r.register(cuda_caps());
10155 r.register(cpu_caps());
10156 assert_eq!(
10157 resolve_target_backend(&r, OpKind::MatMul, DType::F64).unwrap(),
10158 BackendId::Cpu
10159 );
10160 }
10161
10162 #[test]
10164 fn resolve_errors_on_unsupported() {
10165 let mut r = CapabilityRegistry::new();
10166 r.register(cpu_caps());
10167 let result = resolve_target_backend(&r, OpKind::MatMul, DType::BF16);
10168 assert!(result.is_err());
10169 }
10170
10171 #[test]
10174 fn resolve_residency_aware_prefers_local() {
10175 let mut r = CapabilityRegistry::new();
10176 r.register(cuda_caps());
10178 r.register(cpu_caps());
10179
10180 let chosen = resolve_target_backend_residency_aware(
10183 &r,
10184 OpKind::MatMul,
10185 DType::F32,
10186 &[DeviceLocation::Cpu],
10187 )
10188 .unwrap();
10189 assert_eq!(chosen, BackendId::Cpu);
10190 }
10191
10192 #[test]
10195 fn resolve_residency_aware_no_local_falls_back() {
10196 let mut r = CapabilityRegistry::new();
10197 r.register(cuda_caps());
10198 r.register(cpu_caps());
10199
10200 let chosen = resolve_target_backend_residency_aware(
10203 &r,
10204 OpKind::MatMul,
10205 DType::F32,
10206 &[DeviceLocation::Vulkan { gpu_id: 0 }],
10207 )
10208 .unwrap();
10209 assert_eq!(chosen, BackendId::Cuda);
10210 }
10211
10212 #[test]
10216 fn register_cpu_kernels_populates_table() {
10217 let mut table = KernelBindingTable::new();
10218 assert!(table.is_empty());
10219 register_cpu_kernels(&mut table);
10220 assert!(!table.is_empty());
10221 let _kernel = table
10223 .lookup(
10224 OpKind::AddElementwise,
10225 &[DType::F32, DType::F32, DType::F32],
10226 BackendId::Cpu,
10227 )
10228 .expect("registered");
10229 }
10230
10231 #[test]
10249 fn global_bindings_registers_binary_family_from_contract() {
10250 let table = global_bindings();
10251
10252 let families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
10254 (
10255 OpKind::AddElementwise,
10256 [
10257 add_elementwise_f32_cpu_wrapper,
10258 add_elementwise_f64_cpu_wrapper,
10259 add_elementwise_bf16_cpu_wrapper,
10260 add_elementwise_f16_cpu_wrapper,
10261 ],
10262 ),
10263 (
10264 OpKind::SubElementwise,
10265 [
10266 sub_elementwise_f32_cpu_wrapper,
10267 sub_elementwise_f64_cpu_wrapper,
10268 sub_elementwise_bf16_cpu_wrapper,
10269 sub_elementwise_f16_cpu_wrapper,
10270 ],
10271 ),
10272 (
10273 OpKind::MulElementwise,
10274 [
10275 mul_elementwise_f32_cpu_wrapper,
10276 mul_elementwise_f64_cpu_wrapper,
10277 mul_elementwise_bf16_cpu_wrapper,
10278 mul_elementwise_f16_cpu_wrapper,
10279 ],
10280 ),
10281 (
10282 OpKind::DivElementwise,
10283 [
10284 div_elementwise_f32_cpu_wrapper,
10285 div_elementwise_f64_cpu_wrapper,
10286 div_elementwise_bf16_cpu_wrapper,
10287 div_elementwise_f16_cpu_wrapper,
10288 ],
10289 ),
10290 (
10291 OpKind::MaximumElementwise,
10292 [
10293 maximum_elementwise_f32_cpu_wrapper,
10294 maximum_elementwise_f64_cpu_wrapper,
10295 maximum_elementwise_bf16_cpu_wrapper,
10296 maximum_elementwise_f16_cpu_wrapper,
10297 ],
10298 ),
10299 (
10300 OpKind::MinimumElementwise,
10301 [
10302 minimum_elementwise_f32_cpu_wrapper,
10303 minimum_elementwise_f64_cpu_wrapper,
10304 minimum_elementwise_bf16_cpu_wrapper,
10305 minimum_elementwise_f16_cpu_wrapper,
10306 ],
10307 ),
10308 (
10309 OpKind::PowElementwise,
10310 [
10311 pow_elementwise_f32_cpu_wrapper,
10312 pow_elementwise_f64_cpu_wrapper,
10313 pow_elementwise_bf16_cpu_wrapper,
10314 pow_elementwise_f16_cpu_wrapper,
10315 ],
10316 ),
10317 (
10318 OpKind::RemElementwise,
10319 [
10320 rem_elementwise_f32_cpu_wrapper,
10321 rem_elementwise_f64_cpu_wrapper,
10322 rem_elementwise_bf16_cpu_wrapper,
10323 rem_elementwise_f16_cpu_wrapper,
10324 ],
10325 ),
10326 ];
10327 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
10328
10329 let mut checked = 0usize;
10330 for (op, wrappers) in families {
10331 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
10332 let alts = table.lookup_alternatives(*op, &[*dt, *dt, *dt], BackendId::Cpu);
10333 let entry = alts
10334 .iter()
10335 .find(|e| e.kernel as usize == *expected as usize)
10336 .unwrap_or_else(|| {
10337 panic!(
10338 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
10339 FROM the elementwise-binary contract in global_bindings() \
10340 — found {} alternative(s) with sources {:?}",
10341 alts.len(),
10342 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10343 )
10344 });
10345 assert_eq!(
10346 entry.kernel_source, "portable-cpu",
10347 "{op:?}/{dt:?}: binary family must be contract-sourced \
10348 (kernel_source=\"portable-cpu\"); got {:?}",
10349 entry.kernel_source,
10350 );
10351 assert!(
10352 !entry.caps.strided_input,
10353 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
10354 );
10355 assert!(
10356 entry.precision.bit_stable_on_same_hardware,
10357 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
10358 );
10359 checked += 1;
10360 }
10361 }
10362 assert_eq!(checked, 32, "all 8 ops × 4 dtypes checked");
10363 }
10364
10365 #[test]
10375 fn global_bindings_registers_affine_clamp_powi_family_from_contract() {
10376 let table = global_bindings();
10377 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
10378
10379 let forward: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
10381 (
10382 OpKind::Affine,
10383 [
10384 affine_f32_cpu_wrapper,
10385 affine_f64_cpu_wrapper,
10386 affine_bf16_cpu_wrapper,
10387 affine_f16_cpu_wrapper,
10388 ],
10389 ),
10390 (
10391 OpKind::ClampElementwise,
10392 [
10393 clamp_elementwise_f32_cpu_wrapper,
10394 clamp_f64_cpu_wrapper,
10395 clamp_bf16_cpu_wrapper,
10396 clamp_f16_cpu_wrapper,
10397 ],
10398 ),
10399 (
10400 OpKind::PowIElementwise,
10401 [
10402 powi_elementwise_f32_cpu_wrapper,
10403 powi_f64_cpu_wrapper,
10404 powi_bf16_cpu_wrapper,
10405 powi_f16_cpu_wrapper,
10406 ],
10407 ),
10408 ];
10409
10410 let mut checked = 0usize;
10411 for (op, wrappers) in forward {
10412 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
10413 let alts = table.lookup_alternatives(*op, &[*dt, *dt], BackendId::Cpu);
10414 let entry = alts
10415 .iter()
10416 .find(|e| e.kernel as usize == *expected as usize)
10417 .unwrap_or_else(|| {
10418 panic!(
10419 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
10420 FROM the affine-clamp-powi contract in global_bindings() \
10421 — found {} alternative(s) with sources {:?}",
10422 alts.len(),
10423 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10424 )
10425 });
10426 assert_eq!(
10427 entry.kernel_source, "portable-cpu",
10428 "{op:?}/{dt:?}: family must be contract-sourced \
10429 (kernel_source=\"portable-cpu\"); got {:?}",
10430 entry.kernel_source,
10431 );
10432 assert!(
10433 !entry.caps.strided_input,
10434 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
10435 );
10436 assert!(
10437 entry.precision.bit_stable_on_same_hardware,
10438 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
10439 );
10440 checked += 1;
10441 }
10442 }
10443
10444 let bw: [crate::kernel::KernelRef; 4] = [
10446 powi_backward_f32_cpu_wrapper,
10447 powi_backward_f64_cpu_wrapper,
10448 powi_backward_bf16_cpu_wrapper,
10449 powi_backward_f16_cpu_wrapper,
10450 ];
10451 for (dt, expected) in dts.iter().zip(bw.iter()) {
10452 let alts = table.lookup_alternatives(
10453 OpKind::PowIElementwiseBackward,
10454 &[*dt, *dt, *dt],
10455 BackendId::Cpu,
10456 );
10457 let entry = alts
10458 .iter()
10459 .find(|e| e.kernel as usize == *expected as usize)
10460 .unwrap_or_else(|| {
10461 panic!(
10462 "PowIElementwiseBackward/{dt:?}/Cpu: the production wrapper must \
10463 be bound FROM the affine-clamp-powi contract in global_bindings() \
10464 — found {} alternative(s) with sources {:?}",
10465 alts.len(),
10466 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10467 )
10468 });
10469 assert_eq!(
10470 entry.kernel_source, "portable-cpu",
10471 "PowIElementwiseBackward/{dt:?}: must be contract-sourced; got {:?}",
10472 entry.kernel_source,
10473 );
10474 assert!(
10475 !entry.caps.strided_input,
10476 "PowIElementwiseBackward/{dt:?}: caps preserved (contiguous-only)",
10477 );
10478 assert!(
10479 entry.precision.bit_stable_on_same_hardware,
10480 "PowIElementwiseBackward/{dt:?}: the audited bit-stable claim rode through",
10481 );
10482 checked += 1;
10483 }
10484
10485 assert_eq!(checked, 16, "3 forward ops × 4 dtypes + powi_backward × 4");
10486 }
10487
10488 #[test]
10501 fn global_bindings_registers_unary_family_from_contract() {
10502 let table = global_bindings();
10503 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
10504
10505 let families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
10507 (
10508 OpKind::ReluElementwise,
10509 [
10510 relu_elementwise_f32_cpu_wrapper,
10511 relu_elementwise_f64_cpu_wrapper,
10512 relu_elementwise_bf16_cpu_wrapper,
10513 relu_elementwise_f16_cpu_wrapper,
10514 ],
10515 ),
10516 (
10517 OpKind::NegElementwise,
10518 [
10519 neg_elementwise_f32_cpu_wrapper,
10520 neg_elementwise_f64_cpu_wrapper,
10521 neg_elementwise_bf16_cpu_wrapper,
10522 neg_elementwise_f16_cpu_wrapper,
10523 ],
10524 ),
10525 (
10526 OpKind::SqrElementwise,
10527 [
10528 sqr_elementwise_f32_cpu_wrapper,
10529 sqr_elementwise_f64_cpu_wrapper,
10530 sqr_elementwise_bf16_cpu_wrapper,
10531 sqr_elementwise_f16_cpu_wrapper,
10532 ],
10533 ),
10534 (
10535 OpKind::SqrtElementwise,
10536 [
10537 sqrt_elementwise_f32_cpu_wrapper,
10538 sqrt_elementwise_f64_cpu_wrapper,
10539 sqrt_elementwise_bf16_cpu_wrapper,
10540 sqrt_elementwise_f16_cpu_wrapper,
10541 ],
10542 ),
10543 (
10544 OpKind::RecipElementwise,
10545 [
10546 recip_elementwise_f32_cpu_wrapper,
10547 recip_elementwise_f64_cpu_wrapper,
10548 recip_elementwise_bf16_cpu_wrapper,
10549 recip_elementwise_f16_cpu_wrapper,
10550 ],
10551 ),
10552 (
10553 OpKind::AbsElementwise,
10554 [
10555 abs_elementwise_f32_cpu_wrapper,
10556 abs_elementwise_f64_cpu_wrapper,
10557 abs_elementwise_bf16_cpu_wrapper,
10558 abs_elementwise_f16_cpu_wrapper,
10559 ],
10560 ),
10561 (
10562 OpKind::TanhElementwise,
10563 [
10564 tanh_elementwise_f32_cpu_wrapper,
10565 tanh_elementwise_f64_cpu_wrapper,
10566 tanh_elementwise_bf16_cpu_wrapper,
10567 tanh_elementwise_f16_cpu_wrapper,
10568 ],
10569 ),
10570 (
10571 OpKind::ExpElementwise,
10572 [
10573 exp_elementwise_f32_cpu_wrapper,
10574 exp_elementwise_f64_cpu_wrapper,
10575 exp_elementwise_bf16_cpu_wrapper,
10576 exp_elementwise_f16_cpu_wrapper,
10577 ],
10578 ),
10579 (
10580 OpKind::LogElementwise,
10581 [
10582 log_elementwise_f32_cpu_wrapper,
10583 log_elementwise_f64_cpu_wrapper,
10584 log_elementwise_bf16_cpu_wrapper,
10585 log_elementwise_f16_cpu_wrapper,
10586 ],
10587 ),
10588 (
10589 OpKind::SinElementwise,
10590 [
10591 sin_elementwise_f32_cpu_wrapper,
10592 sin_elementwise_f64_cpu_wrapper,
10593 sin_elementwise_bf16_cpu_wrapper,
10594 sin_elementwise_f16_cpu_wrapper,
10595 ],
10596 ),
10597 (
10598 OpKind::CosElementwise,
10599 [
10600 cos_elementwise_f32_cpu_wrapper,
10601 cos_elementwise_f64_cpu_wrapper,
10602 cos_elementwise_bf16_cpu_wrapper,
10603 cos_elementwise_f16_cpu_wrapper,
10604 ],
10605 ),
10606 (
10607 OpKind::SigmoidElementwise,
10608 [
10609 sigmoid_elementwise_f32_cpu_wrapper,
10610 sigmoid_elementwise_f64_cpu_wrapper,
10611 sigmoid_elementwise_bf16_cpu_wrapper,
10612 sigmoid_elementwise_f16_cpu_wrapper,
10613 ],
10614 ),
10615 (
10616 OpKind::SiluElementwise,
10617 [
10618 silu_elementwise_f32_cpu_wrapper,
10619 silu_elementwise_f64_cpu_wrapper,
10620 silu_elementwise_bf16_cpu_wrapper,
10621 silu_elementwise_f16_cpu_wrapper,
10622 ],
10623 ),
10624 (
10625 OpKind::StepElementwise,
10626 [
10627 step_elementwise_f32_cpu_wrapper,
10628 step_elementwise_f64_cpu_wrapper,
10629 step_elementwise_bf16_cpu_wrapper,
10630 step_elementwise_f16_cpu_wrapper,
10631 ],
10632 ),
10633 (
10634 OpKind::GeluTanhElementwise,
10635 [
10636 gelu_elementwise_f32_cpu_wrapper,
10637 gelu_elementwise_f64_cpu_wrapper,
10638 gelu_elementwise_bf16_cpu_wrapper,
10639 gelu_elementwise_f16_cpu_wrapper,
10640 ],
10641 ),
10642 (
10643 OpKind::FloorElementwise,
10644 [
10645 floor_elementwise_f32_cpu_wrapper,
10646 floor_elementwise_f64_cpu_wrapper,
10647 floor_elementwise_bf16_cpu_wrapper,
10648 floor_elementwise_f16_cpu_wrapper,
10649 ],
10650 ),
10651 (
10652 OpKind::CeilElementwise,
10653 [
10654 ceil_elementwise_f32_cpu_wrapper,
10655 ceil_elementwise_f64_cpu_wrapper,
10656 ceil_elementwise_bf16_cpu_wrapper,
10657 ceil_elementwise_f16_cpu_wrapper,
10658 ],
10659 ),
10660 (
10661 OpKind::RoundElementwise,
10662 [
10663 round_elementwise_f32_cpu_wrapper,
10664 round_elementwise_f64_cpu_wrapper,
10665 round_elementwise_bf16_cpu_wrapper,
10666 round_elementwise_f16_cpu_wrapper,
10667 ],
10668 ),
10669 (
10670 OpKind::SignElementwise,
10671 [
10672 sign_elementwise_f32_cpu_wrapper,
10673 sign_elementwise_f64_cpu_wrapper,
10674 sign_elementwise_bf16_cpu_wrapper,
10675 sign_elementwise_f16_cpu_wrapper,
10676 ],
10677 ),
10678 (
10679 OpKind::ErfElementwise,
10680 [
10681 erf_elementwise_f32_cpu_wrapper,
10682 erf_elementwise_f64_cpu_wrapper,
10683 erf_elementwise_bf16_cpu_wrapper,
10684 erf_elementwise_f16_cpu_wrapper,
10685 ],
10686 ),
10687 (
10688 OpKind::GeluElementwise,
10689 [
10690 gelu_erf_elementwise_f32_cpu_wrapper,
10691 gelu_erf_elementwise_f64_cpu_wrapper,
10692 gelu_erf_elementwise_bf16_cpu_wrapper,
10693 gelu_erf_elementwise_f16_cpu_wrapper,
10694 ],
10695 ),
10696 (
10697 OpKind::RsqrtElementwise,
10698 [
10699 rsqrt_elementwise_f32_cpu_wrapper,
10700 rsqrt_elementwise_f64_cpu_wrapper,
10701 rsqrt_elementwise_bf16_cpu_wrapper,
10702 rsqrt_elementwise_f16_cpu_wrapper,
10703 ],
10704 ),
10705 ];
10706
10707 let mut checked = 0usize;
10708 for (op, wrappers) in families {
10709 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
10710 let alts = table.lookup_alternatives(*op, &[*dt, *dt], BackendId::Cpu);
10711 let entry = alts
10712 .iter()
10713 .find(|e| e.kernel as usize == *expected as usize)
10714 .unwrap_or_else(|| {
10715 panic!(
10716 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
10717 FROM the elementwise-unary contract in global_bindings() \
10718 — found {} alternative(s) with sources {:?}",
10719 alts.len(),
10720 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10721 )
10722 });
10723 assert_eq!(
10724 entry.kernel_source, "portable-cpu",
10725 "{op:?}/{dt:?}: unary family must be contract-sourced \
10726 (kernel_source=\"portable-cpu\"); got {:?}",
10727 entry.kernel_source,
10728 );
10729 assert!(
10730 !entry.caps.strided_input,
10731 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
10732 );
10733 assert!(
10734 entry.precision.bit_stable_on_same_hardware,
10735 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
10736 );
10737 checked += 1;
10738 }
10739 }
10740 assert_eq!(checked, 88, "all 22 ops × 4 dtypes checked");
10741 }
10742
10743 #[test]
10763 fn global_bindings_registers_compare_where_family_from_contract() {
10764 let table = global_bindings();
10765 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
10766
10767 let compare_families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
10769 (
10770 OpKind::EqualElementwise,
10771 [
10772 eq_elementwise_f32_cpu_wrapper,
10773 eq_elementwise_f64_cpu_wrapper,
10774 eq_elementwise_bf16_cpu_wrapper,
10775 eq_elementwise_f16_cpu_wrapper,
10776 ],
10777 ),
10778 (
10779 OpKind::NotEqualElementwise,
10780 [
10781 ne_elementwise_f32_cpu_wrapper,
10782 ne_elementwise_f64_cpu_wrapper,
10783 ne_elementwise_bf16_cpu_wrapper,
10784 ne_elementwise_f16_cpu_wrapper,
10785 ],
10786 ),
10787 (
10788 OpKind::LessElementwise,
10789 [
10790 lt_elementwise_f32_cpu_wrapper,
10791 lt_elementwise_f64_cpu_wrapper,
10792 lt_elementwise_bf16_cpu_wrapper,
10793 lt_elementwise_f16_cpu_wrapper,
10794 ],
10795 ),
10796 (
10797 OpKind::LessEqualElementwise,
10798 [
10799 le_elementwise_f32_cpu_wrapper,
10800 le_elementwise_f64_cpu_wrapper,
10801 le_elementwise_bf16_cpu_wrapper,
10802 le_elementwise_f16_cpu_wrapper,
10803 ],
10804 ),
10805 (
10806 OpKind::GreaterElementwise,
10807 [
10808 gt_elementwise_f32_cpu_wrapper,
10809 gt_elementwise_f64_cpu_wrapper,
10810 gt_elementwise_bf16_cpu_wrapper,
10811 gt_elementwise_f16_cpu_wrapper,
10812 ],
10813 ),
10814 (
10815 OpKind::GreaterEqualElementwise,
10816 [
10817 ge_elementwise_f32_cpu_wrapper,
10818 ge_elementwise_f64_cpu_wrapper,
10819 ge_elementwise_bf16_cpu_wrapper,
10820 ge_elementwise_f16_cpu_wrapper,
10821 ],
10822 ),
10823 ];
10824
10825 let mut checked = 0usize;
10826 for (op, wrappers) in compare_families {
10827 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
10828 let alts = table.lookup_alternatives(*op, &[*dt, *dt, DType::Bool], BackendId::Cpu);
10831 let entry = alts
10832 .iter()
10833 .find(|e| e.kernel as usize == *expected as usize)
10834 .unwrap_or_else(|| {
10835 panic!(
10836 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
10837 FROM the compare-where contract in global_bindings() \
10838 — found {} alternative(s) with sources {:?}",
10839 alts.len(),
10840 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10841 )
10842 });
10843 assert_eq!(
10844 entry.kernel_source, "portable-cpu",
10845 "{op:?}/{dt:?}: compare family must be contract-sourced \
10846 (kernel_source=\"portable-cpu\"); got {:?}",
10847 entry.kernel_source,
10848 );
10849 assert!(
10850 !entry.caps.strided_input,
10851 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
10852 );
10853 assert!(
10854 entry.precision.bit_stable_on_same_hardware,
10855 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
10856 );
10857 checked += 1;
10858 }
10859 }
10860
10861 let where_wrappers: [crate::kernel::KernelRef; 4] = [
10863 where_f32_cpu_wrapper,
10864 where_f64_cpu_wrapper,
10865 where_bf16_cpu_wrapper,
10866 where_f16_cpu_wrapper,
10867 ];
10868 for (dt, expected) in dts.iter().zip(where_wrappers.iter()) {
10869 let alts = table.lookup_alternatives(
10871 OpKind::Where,
10872 &[DType::Bool, *dt, *dt, *dt],
10873 BackendId::Cpu,
10874 );
10875 let entry = alts
10876 .iter()
10877 .find(|e| e.kernel as usize == *expected as usize)
10878 .unwrap_or_else(|| {
10879 panic!(
10880 "Where/{dt:?}/Cpu: the production wrapper must be bound FROM \
10881 the compare-where contract (where fans to [U8,T,T,T]) in \
10882 global_bindings() — found {} alternative(s) with sources {:?}",
10883 alts.len(),
10884 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10885 )
10886 });
10887 assert_eq!(
10888 entry.kernel_source, "portable-cpu",
10889 "Where/{dt:?}: where must be contract-sourced \
10890 (kernel_source=\"portable-cpu\"); got {:?}",
10891 entry.kernel_source,
10892 );
10893 assert!(
10894 !entry.caps.strided_input,
10895 "Where/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
10896 );
10897 checked += 1;
10898 }
10899
10900 assert_eq!(checked, 28, "24 compare (6×4) + 4 where keys checked");
10901 }
10902
10903 #[test]
10920 fn global_bindings_registers_reduce_family_from_contract() {
10921 let table = global_bindings();
10922 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
10923
10924 let reduce_families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
10926 (
10927 OpKind::SumReduce,
10928 [
10929 sum_reduce_f32_cpu_wrapper,
10930 sum_reduce_f64_cpu_wrapper,
10931 sum_reduce_bf16_cpu_wrapper,
10932 sum_reduce_f16_cpu_wrapper,
10933 ],
10934 ),
10935 (
10936 OpKind::MeanReduce,
10937 [
10938 mean_reduce_f32_cpu_wrapper,
10939 mean_reduce_f64_cpu_wrapper,
10940 mean_reduce_bf16_cpu_wrapper,
10941 mean_reduce_f16_cpu_wrapper,
10942 ],
10943 ),
10944 (
10945 OpKind::MaxReduce,
10946 [
10947 max_reduce_f32_cpu_wrapper,
10948 max_reduce_f64_cpu_wrapper,
10949 max_reduce_bf16_cpu_wrapper,
10950 max_reduce_f16_cpu_wrapper,
10951 ],
10952 ),
10953 (
10954 OpKind::MinReduce,
10955 [
10956 min_reduce_f32_cpu_wrapper,
10957 min_reduce_f64_cpu_wrapper,
10958 min_reduce_bf16_cpu_wrapper,
10959 min_reduce_f16_cpu_wrapper,
10960 ],
10961 ),
10962 ];
10963
10964 let mut checked = 0usize;
10965 for (op, wrappers) in reduce_families {
10966 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
10967 let alts = table.lookup_alternatives(*op, &[*dt, *dt], BackendId::Cpu);
10969 let entry = alts
10970 .iter()
10971 .find(|e| e.kernel as usize == *expected as usize)
10972 .unwrap_or_else(|| {
10973 panic!(
10974 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
10975 FROM the reduce contract in global_bindings() \
10976 — found {} alternative(s) with sources {:?}",
10977 alts.len(),
10978 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
10979 )
10980 });
10981 assert_eq!(
10982 entry.kernel_source, "portable-cpu",
10983 "{op:?}/{dt:?}: reduce family must be contract-sourced \
10984 (kernel_source=\"portable-cpu\"); got {:?}",
10985 entry.kernel_source,
10986 );
10987 assert!(
10988 !entry.caps.strided_input,
10989 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
10990 );
10991 assert!(
10992 entry.precision.bit_stable_on_same_hardware,
10993 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
10994 );
10995 checked += 1;
10996 }
10997 }
10998 assert_eq!(checked, 16, "4 ops × 4 dtypes checked");
10999 }
11000
11001 #[test]
11013 fn global_bindings_registers_reduce_to_family_from_contract() {
11014 let table = global_bindings();
11015 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11016
11017 let unary_families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
11019 (
11020 OpKind::ReduceSumTo,
11021 [
11022 reduce_sum_to_f32_cpu_wrapper,
11023 reduce_sum_to_f64_cpu_wrapper,
11024 reduce_sum_to_bf16_cpu_wrapper,
11025 reduce_sum_to_f16_cpu_wrapper,
11026 ],
11027 ),
11028 (
11029 OpKind::ReduceMaxTo,
11030 [
11031 reduce_max_to_f32_cpu_wrapper,
11032 reduce_max_to_f64_cpu_wrapper,
11033 reduce_max_to_bf16_cpu_wrapper,
11034 reduce_max_to_f16_cpu_wrapper,
11035 ],
11036 ),
11037 ];
11038 let backward_family: (OpKind, [crate::kernel::KernelRef; 4]) = (
11041 OpKind::ReduceMaxToBackward,
11042 [
11043 reduce_max_to_backward_f32_cpu_wrapper,
11044 reduce_max_to_backward_f64_cpu_wrapper,
11045 reduce_max_to_backward_bf16_cpu_wrapper,
11046 reduce_max_to_backward_f16_cpu_wrapper,
11047 ],
11048 );
11049
11050 let assert_contract_bound =
11052 |op: OpKind, key: &[DType], expected: crate::kernel::KernelRef| {
11053 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
11054 let entry = alts
11055 .iter()
11056 .find(|e| e.kernel as usize == expected as usize)
11057 .unwrap_or_else(|| {
11058 panic!(
11059 "{op:?}/{key:?}/Cpu: the production wrapper must be bound \
11060 FROM the reduce-to contract in global_bindings() \
11061 — found {} alternative(s) with sources {:?}",
11062 alts.len(),
11063 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11064 )
11065 });
11066 assert_eq!(
11067 entry.kernel_source, "portable-cpu",
11068 "{op:?}/{key:?}: reduce-to family must be contract-sourced \
11069 (kernel_source=\"portable-cpu\"); got {:?}",
11070 entry.kernel_source,
11071 );
11072 assert!(
11073 !entry.caps.strided_input,
11074 "{op:?}/{key:?}: caps preserved (contiguous-only, strided_input=false)",
11075 );
11076 assert!(
11077 entry.precision.bit_stable_on_same_hardware,
11078 "{op:?}/{key:?}: the contract's audited bit-stable claim rode through",
11079 );
11080 };
11081
11082 let mut checked = 0usize;
11083 for (op, wrappers) in unary_families {
11084 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
11085 assert_contract_bound(*op, &[*dt, *dt], *expected);
11086 checked += 1;
11087 }
11088 }
11089 let (op, wrappers) = backward_family;
11090 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
11091 assert_contract_bound(op, &[*dt, *dt, *dt], *expected);
11092 checked += 1;
11093 }
11094 assert_eq!(checked, 12, "3 ops × 4 dtypes checked");
11095 }
11096
11097 #[test]
11110 fn global_bindings_registers_norm_family_from_contract() {
11111 let table = global_bindings();
11112 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11113
11114 let norm_families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
11117 (
11118 OpKind::SoftmaxLastDim,
11119 [
11120 softmax_last_dim_f32_cpu_wrapper,
11121 softmax_last_dim_f64_cpu_wrapper,
11122 softmax_last_dim_bf16_cpu_wrapper,
11123 softmax_last_dim_f16_cpu_wrapper,
11124 ],
11125 ),
11126 (
11127 OpKind::LogSoftmaxLastDim,
11128 [
11129 log_softmax_f32_cpu_wrapper,
11130 log_softmax_f64_cpu_wrapper,
11131 log_softmax_bf16_cpu_wrapper,
11132 log_softmax_f16_cpu_wrapper,
11133 ],
11134 ),
11135 (
11136 OpKind::RmsNormLastDim,
11137 [
11138 rms_norm_last_dim_f32_cpu_wrapper,
11139 rms_norm_last_dim_f64_cpu_wrapper,
11140 rms_norm_last_dim_bf16_cpu_wrapper,
11141 rms_norm_last_dim_f16_cpu_wrapper,
11142 ],
11143 ),
11144 (
11145 OpKind::LayerNormLastDim,
11146 [
11147 layer_norm_last_dim_f32_cpu_wrapper,
11148 layer_norm_last_dim_f64_cpu_wrapper,
11149 layer_norm_last_dim_bf16_cpu_wrapper,
11150 layer_norm_last_dim_f16_cpu_wrapper,
11151 ],
11152 ),
11153 ];
11154
11155 let mut checked = 0usize;
11156 for (op, wrappers) in norm_families {
11157 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
11158 let alts = table.lookup_alternatives(*op, &[*dt, *dt], BackendId::Cpu);
11160 let entry = alts
11161 .iter()
11162 .find(|e| e.kernel as usize == *expected as usize)
11163 .unwrap_or_else(|| {
11164 panic!(
11165 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
11166 FROM the norm contract in global_bindings() \
11167 — found {} alternative(s) with sources {:?}",
11168 alts.len(),
11169 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11170 )
11171 });
11172 assert_eq!(
11173 entry.kernel_source, "portable-cpu",
11174 "{op:?}/{dt:?}: norm family must be contract-sourced \
11175 (kernel_source=\"portable-cpu\"); got {:?}",
11176 entry.kernel_source,
11177 );
11178 assert!(
11179 !entry.caps.strided_input,
11180 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
11181 );
11182 assert!(
11183 entry.precision.bit_stable_on_same_hardware,
11184 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
11185 );
11186 checked += 1;
11187 }
11188 }
11189 assert_eq!(checked, 16, "4 ops × 4 dtypes checked");
11190 }
11191
11192 #[test]
11207 fn global_bindings_registers_norm_backward_family_from_contract() {
11208 let table = global_bindings();
11209 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11210
11211 let backward_families: &[(OpKind, [crate::kernel::KernelRef; 4])] = &[
11217 (
11218 OpKind::SoftmaxLastDimBackward,
11219 [
11220 softmax_last_dim_backward_f32_cpu_wrapper,
11221 softmax_last_dim_backward_f64_cpu_wrapper,
11222 softmax_last_dim_backward_bf16_cpu_wrapper,
11223 softmax_last_dim_backward_f16_cpu_wrapper,
11224 ],
11225 ),
11226 (
11227 OpKind::LogSoftmaxLastDimBackward,
11228 [
11229 log_softmax_backward_f32_cpu_wrapper,
11230 log_softmax_backward_f64_cpu_wrapper,
11231 log_softmax_backward_bf16_cpu_wrapper,
11232 log_softmax_backward_f16_cpu_wrapper,
11233 ],
11234 ),
11235 (
11236 OpKind::LayerNormLastDimBackward,
11237 [
11238 layer_norm_last_dim_backward_f32_cpu_wrapper,
11239 layer_norm_last_dim_backward_f64_cpu_wrapper,
11240 layer_norm_last_dim_backward_bf16_cpu_wrapper,
11241 layer_norm_last_dim_backward_f16_cpu_wrapper,
11242 ],
11243 ),
11244 (
11245 OpKind::RmsNormLastDimBackward,
11246 [
11247 rms_norm_last_dim_backward_f32_cpu_wrapper,
11248 rms_norm_last_dim_backward_f64_cpu_wrapper,
11249 rms_norm_last_dim_backward_bf16_cpu_wrapper,
11250 rms_norm_last_dim_backward_f16_cpu_wrapper,
11251 ],
11252 ),
11253 ];
11254
11255 let mut checked = 0usize;
11256 for (op, wrappers) in backward_families {
11257 for (dt, expected) in dts.iter().zip(wrappers.iter()) {
11258 let alts = table.lookup_alternatives(*op, &[*dt, *dt, *dt], BackendId::Cpu);
11260 let entry = alts
11261 .iter()
11262 .find(|e| e.kernel as usize == *expected as usize)
11263 .unwrap_or_else(|| {
11264 panic!(
11265 "{op:?}/{dt:?}/Cpu: the production wrapper must be bound \
11266 FROM the norm-backward contract in global_bindings() \
11267 — found {} alternative(s) with sources {:?}",
11268 alts.len(),
11269 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11270 )
11271 });
11272 assert_eq!(
11273 entry.kernel_source, "portable-cpu",
11274 "{op:?}/{dt:?}: norm-backward family must be contract-sourced \
11275 (kernel_source=\"portable-cpu\"); got {:?}",
11276 entry.kernel_source,
11277 );
11278 assert!(
11279 !entry.caps.strided_input,
11280 "{op:?}/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
11281 );
11282 assert!(
11283 entry.precision.bit_stable_on_same_hardware,
11284 "{op:?}/{dt:?}: the contract's audited bit-stable claim rode through",
11285 );
11286 checked += 1;
11287 }
11288 }
11289 assert_eq!(checked, 16, "4 ops × 4 dtypes checked");
11290 }
11291
11292 #[test]
11310 fn global_bindings_registers_rope_family_from_contract() {
11311 let table = global_bindings();
11312 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11313 let expected_wrappers: [crate::kernel::KernelRef; 4] = [
11315 rope_f32_cpu_wrapper,
11316 rope_f64_cpu_wrapper,
11317 rope_bf16_cpu_wrapper,
11318 rope_f16_cpu_wrapper,
11319 ];
11320
11321 let mut checked = 0usize;
11322 for (dt, expected) in dts.iter().zip(expected_wrappers.iter()) {
11323 let alts =
11325 table.lookup_alternatives(OpKind::Rope, &[*dt, *dt, *dt, *dt], BackendId::Cpu);
11326 let entry = alts
11327 .iter()
11328 .find(|e| e.kernel as usize == *expected as usize)
11329 .unwrap_or_else(|| {
11330 panic!(
11331 "Rope/{dt:?}/Cpu: the production wrapper must be bound \
11332 FROM the rope contract in global_bindings() \
11333 — found {} alternative(s) with sources {:?}",
11334 alts.len(),
11335 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11336 )
11337 });
11338 assert_eq!(
11339 entry.kernel_source, "portable-cpu",
11340 "Rope/{dt:?}: rope family must be contract-sourced \
11341 (kernel_source=\"portable-cpu\"); got {:?}",
11342 entry.kernel_source,
11343 );
11344 assert!(
11345 !entry.caps.strided_input,
11346 "Rope/{dt:?}: caps preserved (contiguous-only, strided_input=false)",
11347 );
11348 assert!(
11349 entry.precision.bit_stable_on_same_hardware,
11350 "Rope/{dt:?}: the contract's audited bit-stable claim rode through",
11351 );
11352 checked += 1;
11353 }
11354 assert_eq!(checked, 4, "1 op × 4 dtypes checked");
11355 }
11356
11357 #[test]
11376 fn global_bindings_registers_ssm_family_from_contract() {
11377 let table = global_bindings();
11378 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11379
11380 let check = |op: OpKind, key: &[DType], expected: crate::kernel::KernelRef, label: &str| {
11383 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
11384 let entry = alts
11385 .iter()
11386 .find(|e| e.kernel as usize == expected as usize)
11387 .unwrap_or_else(|| {
11388 panic!(
11389 "{label}/Cpu (key {key:?}): the production wrapper must be bound \
11390 FROM the ssm contract in global_bindings() — found {} \
11391 alternative(s) with sources {:?}",
11392 alts.len(),
11393 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11394 )
11395 });
11396 assert_eq!(
11397 entry.kernel_source, "portable-cpu",
11398 "{label}: family must be contract-sourced (kernel_source=\"portable-cpu\"); \
11399 got {:?}",
11400 entry.kernel_source,
11401 );
11402 assert!(
11403 !entry.caps.strided_input,
11404 "{label}: caps preserved (contiguous-only)"
11405 );
11406 assert!(
11407 entry.precision.bit_stable_on_same_hardware,
11408 "{label}: contract's bit-stable claim rode through",
11409 );
11410 };
11411
11412 let fsce: [crate::kernel::KernelRef; 4] = [
11414 fused_softmax_cross_entropy_f32_cpu_wrapper,
11415 fused_softmax_cross_entropy_f64_cpu_wrapper,
11416 fused_softmax_cross_entropy_bf16_cpu_wrapper,
11417 fused_softmax_cross_entropy_f16_cpu_wrapper,
11418 ];
11419 let conv: [crate::kernel::KernelRef; 4] = [
11420 causal_conv1d_f32_cpu_wrapper,
11421 causal_conv1d_f64_cpu_wrapper,
11422 causal_conv1d_bf16_cpu_wrapper,
11423 causal_conv1d_f16_cpu_wrapper,
11424 ];
11425 let sscan: [crate::kernel::KernelRef; 4] = [
11426 selective_scan_f32_cpu_wrapper,
11427 selective_scan_f64_cpu_wrapper,
11428 selective_scan_bf16_cpu_wrapper,
11429 selective_scan_f16_cpu_wrapper,
11430 ];
11431 let ssd: [crate::kernel::KernelRef; 4] = [
11432 ssd_chunk_scan_f32_cpu_wrapper,
11433 ssd_chunk_scan_f64_cpu_wrapper,
11434 ssd_chunk_scan_bf16_cpu_wrapper,
11435 ssd_chunk_scan_f16_cpu_wrapper,
11436 ];
11437
11438 let mut checked = 0usize;
11439 for (i, dt) in dts.iter().enumerate() {
11440 check(
11442 OpKind::FusedSoftmaxCrossEntropy,
11443 &[*dt, DType::I64, DType::F32],
11444 fsce[i],
11445 "FusedSoftmaxCrossEntropy",
11446 );
11447 check(OpKind::CausalConv1d, &[*dt; 4], conv[i], "CausalConv1d");
11449 check(OpKind::SelectiveScan, &[*dt; 6], sscan[i], "SelectiveScan");
11452 check(OpKind::SsdChunkScan, &[*dt; 6], ssd[i], "SsdChunkScan");
11455
11456 checked += 1;
11457 }
11458 assert_eq!(
11459 checked, 4,
11460 "4 ops × 4 dtypes checked (16 contract-sourced bindings)"
11461 );
11462 }
11463
11464 #[test]
11482 fn global_bindings_registers_conv_family_from_contract() {
11483 let table = global_bindings();
11484 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11485
11486 let conv2d_wrappers: [crate::kernel::KernelRef; 4] = [
11489 conv2d_f32_cpu_wrapper,
11490 conv2d_f64_cpu_wrapper,
11491 conv2d_bf16_cpu_wrapper,
11492 conv2d_f16_cpu_wrapper,
11493 ];
11494 let convt_wrappers: [crate::kernel::KernelRef; 4] = [
11496 conv_transpose2d_f32_cpu_wrapper,
11497 conv_transpose2d_f64_cpu_wrapper,
11498 conv_transpose2d_bf16_cpu_wrapper,
11499 conv_transpose2d_f16_cpu_wrapper,
11500 ];
11501
11502 let mut checked = 0usize;
11503 for (i, dt) in dts.iter().enumerate() {
11504 for (op, expected) in [
11505 (OpKind::Conv2D, conv2d_wrappers[i]),
11506 (OpKind::ConvTranspose2D, convt_wrappers[i]),
11507 ] {
11508 let no_bias: &[DType] = &[*dt, *dt, *dt];
11512 let with_bias: &[DType] = &[*dt, *dt, *dt, *dt];
11513 for (key, label) in [(no_bias, "no-bias"), (with_bias, "with-bias")] {
11514 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
11515 let entry = alts
11516 .iter()
11517 .find(|e| e.kernel as usize == expected as usize)
11518 .unwrap_or_else(|| {
11519 panic!(
11520 "{op:?}/{dt:?}/Cpu ({label}): the production wrapper must be \
11521 bound FROM the conv contract in global_bindings() — found {} \
11522 alternative(s) with sources {:?}",
11523 alts.len(),
11524 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11525 )
11526 });
11527 assert_eq!(
11528 entry.kernel_source, "portable-cpu",
11529 "{op:?}/{dt:?} ({label}): family must be contract-sourced \
11530 (kernel_source=\"portable-cpu\"); got {:?}",
11531 entry.kernel_source,
11532 );
11533 assert!(
11534 !entry.caps.strided_input,
11535 "{op:?}/{dt:?} ({label}): caps preserved (contiguous-only)",
11536 );
11537 assert!(
11538 entry.precision.bit_stable_on_same_hardware,
11539 "{op:?}/{dt:?} ({label}): contract's bit-stable claim rode through",
11540 );
11541 checked += 1;
11542 }
11543 }
11544 }
11545 assert_eq!(
11546 checked, 16,
11547 "2 ops × 4 dtypes × 2 operand-counts checked (16 contract-sourced bindings)"
11548 );
11549 }
11550
11551 #[test]
11575 fn global_bindings_registers_padding_family_from_contract() {
11576 let table = global_bindings();
11577 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11578 let pad_backward: [crate::kernel::KernelRef; 4] = [
11580 pad_backward_f32_cpu_wrapper,
11581 pad_backward_f64_cpu_wrapper,
11582 pad_backward_bf16_cpu_wrapper,
11583 pad_backward_f16_cpu_wrapper,
11584 ];
11585
11586 let mut checked = 0usize;
11587 for (i, dt) in dts.iter().enumerate() {
11588 let key: &[DType] = &[*dt, *dt];
11590 let expected = pad_backward[i];
11591 let alts = table.lookup_alternatives(OpKind::PadBackward, key, BackendId::Cpu);
11592 let entry = alts
11593 .iter()
11594 .find(|e| e.kernel as usize == expected as usize)
11595 .unwrap_or_else(|| {
11596 panic!(
11597 "PadBackward/{dt:?}/Cpu: the production wrapper must be bound \
11598 FROM the padding contract in global_bindings() — found {} \
11599 alternative(s) with sources {:?}",
11600 alts.len(),
11601 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11602 )
11603 });
11604 assert_eq!(
11605 entry.kernel_source, "portable-cpu",
11606 "PadBackward/{dt:?}: family must be contract-sourced \
11607 (kernel_source=\"portable-cpu\"); got {:?}",
11608 entry.kernel_source,
11609 );
11610 assert!(
11611 !entry.caps.strided_input,
11612 "PadBackward/{dt:?}: caps preserved (contiguous-only)",
11613 );
11614 assert!(
11615 entry.precision.bit_stable_on_same_hardware,
11616 "PadBackward/{dt:?}: contract's bit-stable claim rode through",
11617 );
11618 checked += 1;
11619 }
11620
11621 let pad_forward_dts = [
11629 DType::U8,
11630 DType::U32,
11631 DType::BF16,
11632 DType::F16,
11633 DType::F32,
11634 DType::F64,
11635 ];
11636 for dt in pad_forward_dts {
11637 let key: &[DType] = &[dt, dt];
11638 let alts = table.lookup_alternatives(OpKind::Pad, key, BackendId::Cpu);
11639 let entry = alts
11640 .iter()
11641 .find(|e| e.kernel as usize == pad_cpu_wrapper as *const () as usize)
11642 .unwrap_or_else(|| {
11643 panic!(
11644 "Pad/{dt:?}/Cpu: the production wrapper must be bound \
11645 FROM the padding contract in global_bindings() — found {} \
11646 alternative(s) with sources {:?}",
11647 alts.len(),
11648 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11649 )
11650 });
11651 assert_eq!(
11652 entry.kernel_source, "portable-cpu",
11653 "Pad/{dt:?}: forward family must be contract-sourced \
11654 (kernel_source=\"portable-cpu\"); got {:?}",
11655 entry.kernel_source,
11656 );
11657 assert!(
11658 !entry.caps.strided_input,
11659 "Pad/{dt:?}: caps preserved (contiguous-only)",
11660 );
11661 assert!(
11662 entry.precision.bit_stable_on_same_hardware,
11663 "Pad/{dt:?}: contract's bit-stable claim rode through",
11664 );
11665 checked += 1;
11666 }
11667
11668 assert_eq!(
11669 checked, 10,
11670 "4 PadBackward + 6 forward Pad dtypes checked (contract-sourced bindings)"
11671 );
11672 }
11673
11674 #[test]
11706 fn global_bindings_registers_shape_ops_family_from_contract() {
11707 let table = global_bindings();
11708 let mut checked = 0usize;
11709
11710 let mut check =
11712 |op: OpKind, key: &[DType], expected: crate::kernel::KernelRef, label: &str| {
11713 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
11714 let entry = alts
11715 .iter()
11716 .find(|e| e.kernel as usize == expected as usize)
11717 .unwrap_or_else(|| {
11718 panic!(
11719 "{op:?}/{label}/Cpu: the production wrapper must be bound \
11720 FROM the shape-ops contract in global_bindings() — found {} \
11721 alternative(s) with sources {:?}",
11722 alts.len(),
11723 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11724 )
11725 });
11726 assert_eq!(
11727 entry.kernel_source, "portable-cpu",
11728 "{op:?}/{label}: family must be contract-sourced \
11729 (kernel_source=\"portable-cpu\"); got {:?}",
11730 entry.kernel_source,
11731 );
11732 assert!(
11733 !entry.caps.strided_input,
11734 "{op:?}/{label}: caps preserved (contiguous-only)",
11735 );
11736 assert!(
11737 entry.precision.bit_stable_on_same_hardware,
11738 "{op:?}/{label}: contract's bit-stable claim rode through",
11739 );
11740 checked += 1;
11741 };
11742
11743 let byte_dts = [
11745 DType::F32,
11746 DType::F64,
11747 DType::BF16,
11748 DType::F16,
11749 DType::U32,
11750 DType::U8,
11751 ];
11752 for dt in byte_dts {
11753 check(OpKind::Flip, &[dt, dt], flip_cpu_wrapper, "flip");
11754 check(OpKind::Roll, &[dt, dt], roll_cpu_wrapper, "roll");
11755 check(
11757 OpKind::MaskedFill,
11758 &[dt, DType::Bool, dt], masked_fill_cpu_wrapper,
11760 "masked_fill",
11761 );
11762 }
11763
11764 let cumsum_dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11766 let cumsum_wrappers: [crate::kernel::KernelRef; 4] = [
11767 cumsum_f32_cpu_wrapper,
11768 cumsum_f64_cpu_wrapper,
11769 cumsum_bf16_cpu_wrapper,
11770 cumsum_f16_cpu_wrapper,
11771 ];
11772 for (i, dt) in cumsum_dts.iter().enumerate() {
11773 check(OpKind::CumSum, &[*dt, *dt], cumsum_wrappers[i], "cumsum");
11774 }
11775
11776 let concat_dts = [
11778 DType::F32,
11779 DType::F64,
11780 DType::BF16,
11781 DType::F16,
11782 DType::U32,
11783 DType::U8,
11784 DType::I16,
11785 DType::I32,
11786 DType::I64,
11787 ];
11788 for dt in concat_dts {
11789 check(OpKind::Concat, &[dt, dt], concat_cpu_wrapper, "concat");
11790 }
11791
11792 let scatter_dts = [
11799 DType::F32,
11800 DType::F64,
11801 DType::BF16,
11802 DType::F16,
11803 DType::U32,
11804 DType::U8,
11805 ];
11806 for dt in scatter_dts {
11807 check(
11808 OpKind::WriteSlice,
11809 &[dt, dt],
11810 write_slice_cpu_wrapper,
11811 "write_slice",
11812 );
11813 check(
11814 OpKind::WriteSliceRotating,
11815 &[dt, dt],
11816 write_slice_rotating_cpu_wrapper,
11817 "write_slice_rotating",
11818 );
11819 check(
11822 OpKind::WriteSliceDoff,
11823 &[dt, dt],
11824 write_slice_doff_cpu_wrapper,
11825 "write_slice_doff",
11826 );
11827 }
11828
11829 assert_eq!(
11830 checked,
11831 6 + 6 + 6 + 4 + 9 + 6 + 6 + 6,
11832 "flip(6) + roll(6) + masked_fill(6) + cumsum(4) + concat(9) + \
11833 write_slice(6) + write_slice_rotating(6) + write_slice_doff(6) contract-sourced bindings",
11834 );
11835 }
11836
11837 #[test]
11867 fn global_bindings_registers_indexing_family_from_contract() {
11868 let table = global_bindings();
11869 let mut checked = 0usize;
11870
11871 let mut check =
11873 |op: OpKind, key: &[DType], expected: crate::kernel::KernelRef, label: &str| {
11874 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
11875 let entry = alts
11876 .iter()
11877 .find(|e| e.kernel as usize == expected as usize)
11878 .unwrap_or_else(|| {
11879 panic!(
11880 "{op:?}/{label}/Cpu: the production wrapper must be bound \
11881 FROM the indexing contract in global_bindings() — found {} \
11882 alternative(s) with sources {:?}",
11883 alts.len(),
11884 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
11885 )
11886 });
11887 assert_eq!(
11888 entry.kernel_source, "portable-cpu",
11889 "{op:?}/{label}: family must be contract-sourced \
11890 (kernel_source=\"portable-cpu\"); got {:?}",
11891 entry.kernel_source,
11892 );
11893 assert!(
11894 !entry.caps.strided_input,
11895 "{op:?}/{label}: caps preserved (contiguous-only)",
11896 );
11897 assert!(
11898 entry.precision.bit_stable_on_same_hardware,
11899 "{op:?}/{label}: contract's bit-stable claim rode through",
11900 );
11901 checked += 1;
11902 };
11903
11904 let byte_dts = [
11907 DType::F32,
11908 DType::F64,
11909 DType::BF16,
11910 DType::F16,
11911 DType::U32,
11912 DType::U8,
11913 DType::I16,
11914 DType::I32,
11915 DType::I64,
11916 ];
11917 for dt in byte_dts {
11918 check(
11919 OpKind::IndexSelect,
11920 &[dt, DType::U32, dt],
11921 index_select_cpu_wrapper,
11922 "index_select",
11923 );
11924 check(
11925 OpKind::Gather,
11926 &[dt, DType::U32, dt],
11927 gather_cpu_wrapper,
11928 "gather",
11929 );
11930 }
11931
11932 let acc_dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
11935 let index_add_wrappers: [crate::kernel::KernelRef; 4] = [
11936 index_add_f32_cpu_wrapper,
11937 index_add_f64_cpu_wrapper,
11938 index_add_bf16_cpu_wrapper,
11939 index_add_f16_cpu_wrapper,
11940 ];
11941 let scatter_add_wrappers: [crate::kernel::KernelRef; 4] = [
11942 scatter_add_f32_cpu_wrapper,
11943 scatter_add_f64_cpu_wrapper,
11944 scatter_add_bf16_cpu_wrapper,
11945 scatter_add_f16_cpu_wrapper,
11946 ];
11947 for (i, dt) in acc_dts.iter().enumerate() {
11948 check(
11949 OpKind::IndexAdd,
11950 &[*dt, DType::U32, *dt, *dt],
11951 index_add_wrappers[i],
11952 "index_add",
11953 );
11954 check(
11955 OpKind::ScatterAdd,
11956 &[*dt, DType::U32, *dt, *dt],
11957 scatter_add_wrappers[i],
11958 "scatter_add",
11959 );
11960 }
11961
11962 assert_eq!(
11963 checked,
11964 9 + 9 + 4 + 4,
11965 "index_select(9) + gather(9) + index_add(4) + scatter_add(4) \
11966 contract-sourced bindings",
11967 );
11968 }
11969
11970 #[test]
12002 fn global_bindings_registers_matmul_family_from_contract() {
12003 let table = global_bindings();
12004 let mut checked = 0usize;
12005
12006 let mut check =
12010 |op: OpKind, key: &[DType], expected: crate::kernel::KernelRef, label: &str| {
12011 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
12012 let entry = alts
12013 .iter()
12014 .find(|e| e.kernel as usize == expected as usize)
12015 .unwrap_or_else(|| {
12016 panic!(
12017 "{op:?}/{label}/Cpu: the portable wrapper must be bound \
12018 FROM the matmul contract in global_bindings() — found {} \
12019 alternative(s) with sources {:?}",
12020 alts.len(),
12021 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
12022 )
12023 });
12024 assert_eq!(
12025 entry.kernel_source, "portable-cpu",
12026 "{op:?}/{label}: portable kernel must be contract-sourced \
12027 (kernel_source=\"portable-cpu\"); got {:?}",
12028 entry.kernel_source,
12029 );
12030 assert!(
12031 !entry.caps.strided_input,
12032 "{op:?}/{label}: caps preserved (contiguous-only)",
12033 );
12034 assert!(
12035 entry.precision.bit_stable_on_same_hardware,
12036 "{op:?}/{label}: contract's bit-stable claim rode through",
12037 );
12038 checked += 1;
12039 };
12040
12041 let mm_dts = [
12043 DType::F32,
12044 DType::F64,
12045 DType::BF16,
12046 DType::F16,
12047 DType::I8,
12048 DType::U8,
12049 ];
12050 let mm_wrappers: [crate::kernel::KernelRef; 6] = [
12051 matmul_f32_cpu_wrapper,
12052 matmul_f64_cpu_wrapper,
12053 matmul_bf16_cpu_wrapper,
12054 matmul_f16_cpu_wrapper,
12055 matmul_i8_cpu_wrapper,
12056 matmul_u8_cpu_wrapper,
12057 ];
12058 for (i, dt) in mm_dts.iter().enumerate() {
12059 check(OpKind::MatMul, &[*dt, *dt, *dt], mm_wrappers[i], "matmul");
12060 }
12061
12062 let fl_dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
12064 let fl_wrappers: [crate::kernel::KernelRef; 4] = [
12065 fused_linear_f32_cpu_wrapper,
12066 fused_linear_f64_cpu_wrapper,
12067 fused_linear_bf16_cpu_wrapper,
12068 fused_linear_f16_cpu_wrapper,
12069 ];
12070 for (i, dt) in fl_dts.iter().enumerate() {
12071 check(
12072 OpKind::FusedLinear,
12073 &[*dt, *dt, *dt, *dt],
12074 fl_wrappers[i],
12075 "fused_linear",
12076 );
12077 }
12078
12079 assert_eq!(
12080 checked,
12081 6 + 4,
12082 "matmul(6) + fused_linear(4) contract-sourced bindings",
12083 );
12084 }
12085
12086 #[test]
12116 fn global_bindings_registers_attention_family_from_contract() {
12117 let table = global_bindings();
12118 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
12119
12120 let fa_wrappers: [crate::kernel::KernelRef; 4] = [
12123 flash_attn_f32_cpu_wrapper,
12124 flash_attn_f64_cpu_wrapper,
12125 flash_attn_bf16_cpu_wrapper,
12126 flash_attn_f16_cpu_wrapper,
12127 ];
12128 let fabq: [crate::kernel::KernelRef; 4] = [
12130 flash_attn_backward_q_f32_cpu_wrapper,
12131 flash_attn_backward_q_f64_cpu_wrapper,
12132 flash_attn_backward_q_bf16_cpu_wrapper,
12133 flash_attn_backward_q_f16_cpu_wrapper,
12134 ];
12135 let fabk: [crate::kernel::KernelRef; 4] = [
12136 flash_attn_backward_k_f32_cpu_wrapper,
12137 flash_attn_backward_k_f64_cpu_wrapper,
12138 flash_attn_backward_k_bf16_cpu_wrapper,
12139 flash_attn_backward_k_f16_cpu_wrapper,
12140 ];
12141 let fabv: [crate::kernel::KernelRef; 4] = [
12142 flash_attn_backward_v_f32_cpu_wrapper,
12143 flash_attn_backward_v_f64_cpu_wrapper,
12144 flash_attn_backward_v_bf16_cpu_wrapper,
12145 flash_attn_backward_v_f16_cpu_wrapper,
12146 ];
12147 let pa_wrappers: [crate::kernel::KernelRef; 4] = [
12150 paged_attn_f32_cpu_wrapper,
12151 paged_attn_f64_cpu_wrapper,
12152 paged_attn_bf16_cpu_wrapper,
12153 paged_attn_f16_cpu_wrapper,
12154 ];
12155
12156 let check = |op: OpKind, key: &[DType], expected: crate::kernel::KernelRef, label: &str| {
12159 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
12160 let entry = alts
12161 .iter()
12162 .find(|e| e.kernel as usize == expected as usize)
12163 .unwrap_or_else(|| {
12164 panic!(
12165 "{op:?}/{label}/Cpu: the production wrapper must be bound \
12166 FROM the attention contract in global_bindings() — found {} \
12167 alternative(s) with sources {:?}",
12168 alts.len(),
12169 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
12170 )
12171 });
12172 assert_eq!(
12173 entry.kernel_source, "portable-cpu",
12174 "{op:?}/{label}: family must be contract-sourced \
12175 (kernel_source=\"portable-cpu\"); got {:?}",
12176 entry.kernel_source,
12177 );
12178 assert!(
12179 !entry.caps.strided_input,
12180 "{op:?}/{label}: caps preserved (contiguous-only)",
12181 );
12182 assert!(
12183 entry.precision.bit_stable_on_same_hardware,
12184 "{op:?}/{label}: contract's bit-stable claim rode through",
12185 );
12186 };
12187
12188 let mut checked = 0usize;
12189 for (i, dt) in dts.iter().enumerate() {
12190 let fa_noa: &[DType] = &[*dt, *dt, *dt, *dt];
12192 let fa_a: &[DType] = &[*dt, *dt, *dt, *dt, *dt];
12193 check(OpKind::FlashAttn, fa_noa, fa_wrappers[i], "flash no-alibi");
12194 check(OpKind::FlashAttn, fa_a, fa_wrappers[i], "flash with-alibi");
12195 checked += 2;
12196
12197 let bw_noa: &[DType] = &[*dt, *dt, *dt, *dt, *dt];
12199 let bw_a: &[DType] = &[*dt, *dt, *dt, *dt, *dt, *dt];
12200 for (op, ws) in [
12201 (OpKind::FlashAttnBackwardQ, &fabq),
12202 (OpKind::FlashAttnBackwardK, &fabk),
12203 (OpKind::FlashAttnBackwardV, &fabv),
12204 ] {
12205 check(op, bw_noa, ws[i], "backward no-alibi");
12206 check(op, bw_a, ws[i], "backward with-alibi");
12207 checked += 2;
12208 }
12209
12210 let pa_noa: &[DType] = &[*dt, *dt, *dt, DType::U32, DType::U32, *dt];
12219 let pa_a: &[DType] = &[*dt, *dt, *dt, DType::U32, DType::U32, *dt, *dt];
12220 check(OpKind::PagedAttn, pa_noa, pa_wrappers[i], "paged no-alibi");
12221 check(OpKind::PagedAttn, pa_a, pa_wrappers[i], "paged with-alibi");
12222 checked += 2;
12223 }
12224 assert_eq!(
12225 checked, 40,
12226 "FlashAttn(8) + FlashAttnBackward{{Q,K,V}}(24) + PagedAttn(8) = 40 \
12227 contract-sourced bindings",
12228 );
12229 }
12230
12231 #[test]
12246 fn global_bindings_registers_inplace_family_from_contract() {
12247 let table = global_bindings();
12248 let dts = [DType::F32, DType::F64, DType::BF16, DType::F16];
12249
12250 let check = |op: OpKind, dt: DType, expected: crate::kernel::KernelRef, label: &str| {
12253 let key: &[DType] = &[dt, dt];
12254 let alts = table.lookup_alternatives(op, key, BackendId::Cpu);
12255 let entry = alts
12256 .iter()
12257 .find(|e| e.kernel as usize == expected as usize)
12258 .unwrap_or_else(|| {
12259 panic!(
12260 "{op:?}/{label}/{dt:?}/Cpu: the production wrapper must be bound \
12261 FROM the in-place contract in global_bindings() — found {} \
12262 alternative(s) with sources {:?}",
12263 alts.len(),
12264 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
12265 )
12266 });
12267 assert_eq!(
12268 entry.kernel_source, "portable-cpu",
12269 "{op:?}/{label}/{dt:?}: family must be contract-sourced \
12270 (kernel_source=\"portable-cpu\"); got {:?}",
12271 entry.kernel_source,
12272 );
12273 assert!(
12274 !entry.caps.strided_input,
12275 "{op:?}/{label}/{dt:?}: caps preserved (contiguous-only)",
12276 );
12277 assert!(
12278 entry.precision.bit_stable_on_same_hardware,
12279 "{op:?}/{label}/{dt:?}: contract's bit-stable claim rode through",
12280 );
12281 };
12282
12283 let unary_family: &[(OpKind, &str, [crate::kernel::KernelRef; 4])] = &[
12286 (
12287 OpKind::ReluInplace,
12288 "relu",
12289 [
12290 relu_inplace_f32_cpu_wrapper,
12291 relu_inplace_f64_cpu_wrapper,
12292 relu_inplace_bf16_cpu_wrapper,
12293 relu_inplace_f16_cpu_wrapper,
12294 ],
12295 ),
12296 (
12297 OpKind::SiluInplace,
12298 "silu",
12299 [
12300 silu_inplace_f32_cpu_wrapper,
12301 silu_inplace_f64_cpu_wrapper,
12302 silu_inplace_bf16_cpu_wrapper,
12303 silu_inplace_f16_cpu_wrapper,
12304 ],
12305 ),
12306 (
12307 OpKind::GeluTanhInplace,
12308 "gelu_tanh",
12309 [
12310 gelu_inplace_f32_cpu_wrapper,
12311 gelu_inplace_f64_cpu_wrapper,
12312 gelu_inplace_bf16_cpu_wrapper,
12313 gelu_inplace_f16_cpu_wrapper,
12314 ],
12315 ),
12316 (
12317 OpKind::TanhInplace,
12318 "tanh",
12319 [
12320 tanh_inplace_f32_cpu_wrapper,
12321 tanh_inplace_f64_cpu_wrapper,
12322 tanh_inplace_bf16_cpu_wrapper,
12323 tanh_inplace_f16_cpu_wrapper,
12324 ],
12325 ),
12326 (
12327 OpKind::SigmoidInplace,
12328 "sigmoid",
12329 [
12330 sigmoid_inplace_f32_cpu_wrapper,
12331 sigmoid_inplace_f64_cpu_wrapper,
12332 sigmoid_inplace_bf16_cpu_wrapper,
12333 sigmoid_inplace_f16_cpu_wrapper,
12334 ],
12335 ),
12336 (
12337 OpKind::NegInplace,
12338 "neg",
12339 [
12340 neg_inplace_f32_cpu_wrapper,
12341 neg_inplace_f64_cpu_wrapper,
12342 neg_inplace_bf16_cpu_wrapper,
12343 neg_inplace_f16_cpu_wrapper,
12344 ],
12345 ),
12346 (
12347 OpKind::AbsInplace,
12348 "abs",
12349 [
12350 abs_inplace_f32_cpu_wrapper,
12351 abs_inplace_f64_cpu_wrapper,
12352 abs_inplace_bf16_cpu_wrapper,
12353 abs_inplace_f16_cpu_wrapper,
12354 ],
12355 ),
12356 (
12357 OpKind::SqrInplace,
12358 "sqr",
12359 [
12360 sqr_inplace_f32_cpu_wrapper,
12361 sqr_inplace_f64_cpu_wrapper,
12362 sqr_inplace_bf16_cpu_wrapper,
12363 sqr_inplace_f16_cpu_wrapper,
12364 ],
12365 ),
12366 (
12367 OpKind::SqrtInplace,
12368 "sqrt",
12369 [
12370 sqrt_inplace_f32_cpu_wrapper,
12371 sqrt_inplace_f64_cpu_wrapper,
12372 sqrt_inplace_bf16_cpu_wrapper,
12373 sqrt_inplace_f16_cpu_wrapper,
12374 ],
12375 ),
12376 (
12377 OpKind::RsqrtInplace,
12378 "rsqrt",
12379 [
12380 rsqrt_inplace_f32_cpu_wrapper,
12381 rsqrt_inplace_f64_cpu_wrapper,
12382 rsqrt_inplace_bf16_cpu_wrapper,
12383 rsqrt_inplace_f16_cpu_wrapper,
12384 ],
12385 ),
12386 (
12387 OpKind::RecipInplace,
12388 "recip",
12389 [
12390 recip_inplace_f32_cpu_wrapper,
12391 recip_inplace_f64_cpu_wrapper,
12392 recip_inplace_bf16_cpu_wrapper,
12393 recip_inplace_f16_cpu_wrapper,
12394 ],
12395 ),
12396 (
12397 OpKind::ExpInplace,
12398 "exp",
12399 [
12400 exp_inplace_f32_cpu_wrapper,
12401 exp_inplace_f64_cpu_wrapper,
12402 exp_inplace_bf16_cpu_wrapper,
12403 exp_inplace_f16_cpu_wrapper,
12404 ],
12405 ),
12406 (
12407 OpKind::LogInplace,
12408 "log",
12409 [
12410 log_inplace_f32_cpu_wrapper,
12411 log_inplace_f64_cpu_wrapper,
12412 log_inplace_bf16_cpu_wrapper,
12413 log_inplace_f16_cpu_wrapper,
12414 ],
12415 ),
12416 (
12417 OpKind::SinInplace,
12418 "sin",
12419 [
12420 sin_inplace_f32_cpu_wrapper,
12421 sin_inplace_f64_cpu_wrapper,
12422 sin_inplace_bf16_cpu_wrapper,
12423 sin_inplace_f16_cpu_wrapper,
12424 ],
12425 ),
12426 (
12427 OpKind::CosInplace,
12428 "cos",
12429 [
12430 cos_inplace_f32_cpu_wrapper,
12431 cos_inplace_f64_cpu_wrapper,
12432 cos_inplace_bf16_cpu_wrapper,
12433 cos_inplace_f16_cpu_wrapper,
12434 ],
12435 ),
12436 (
12437 OpKind::SignInplace,
12438 "sign",
12439 [
12440 sign_inplace_f32_cpu_wrapper,
12441 sign_inplace_f64_cpu_wrapper,
12442 sign_inplace_bf16_cpu_wrapper,
12443 sign_inplace_f16_cpu_wrapper,
12444 ],
12445 ),
12446 (
12447 OpKind::FloorInplace,
12448 "floor",
12449 [
12450 floor_inplace_f32_cpu_wrapper,
12451 floor_inplace_f64_cpu_wrapper,
12452 floor_inplace_bf16_cpu_wrapper,
12453 floor_inplace_f16_cpu_wrapper,
12454 ],
12455 ),
12456 (
12457 OpKind::CeilInplace,
12458 "ceil",
12459 [
12460 ceil_inplace_f32_cpu_wrapper,
12461 ceil_inplace_f64_cpu_wrapper,
12462 ceil_inplace_bf16_cpu_wrapper,
12463 ceil_inplace_f16_cpu_wrapper,
12464 ],
12465 ),
12466 (
12467 OpKind::RoundInplace,
12468 "round",
12469 [
12470 round_inplace_f32_cpu_wrapper,
12471 round_inplace_f64_cpu_wrapper,
12472 round_inplace_bf16_cpu_wrapper,
12473 round_inplace_f16_cpu_wrapper,
12474 ],
12475 ),
12476 (
12477 OpKind::ErfInplace,
12478 "erf",
12479 [
12480 erf_inplace_f32_cpu_wrapper,
12481 erf_inplace_f64_cpu_wrapper,
12482 erf_inplace_bf16_cpu_wrapper,
12483 erf_inplace_f16_cpu_wrapper,
12484 ],
12485 ),
12486 (
12487 OpKind::GeluInplace,
12488 "gelu",
12489 [
12490 gelu_erf_inplace_f32_cpu_wrapper,
12491 gelu_erf_inplace_f64_cpu_wrapper,
12492 gelu_erf_inplace_bf16_cpu_wrapper,
12493 gelu_erf_inplace_f16_cpu_wrapper,
12494 ],
12495 ),
12496 ];
12497
12498 let scalar_family: &[(OpKind, &str, [crate::kernel::KernelRef; 4])] = &[
12502 (
12503 OpKind::InplaceAffine,
12504 "affine",
12505 [
12506 inplace_affine_f32_cpu_wrapper,
12507 inplace_affine_f64_cpu_wrapper,
12508 inplace_affine_bf16_cpu_wrapper,
12509 inplace_affine_f16_cpu_wrapper,
12510 ],
12511 ),
12512 (
12513 OpKind::ClampInplace,
12514 "clamp",
12515 [
12516 clamp_inplace_f32_cpu_wrapper,
12517 clamp_inplace_f64_cpu_wrapper,
12518 clamp_inplace_bf16_cpu_wrapper,
12519 clamp_inplace_f16_cpu_wrapper,
12520 ],
12521 ),
12522 (
12523 OpKind::PowIInplace,
12524 "powi",
12525 [
12526 powi_inplace_f32_cpu_wrapper,
12527 powi_inplace_f64_cpu_wrapper,
12528 powi_inplace_bf16_cpu_wrapper,
12529 powi_inplace_f16_cpu_wrapper,
12530 ],
12531 ),
12532 ];
12533
12534 let mut checked = 0usize;
12535 for (op, label, ws) in unary_family.iter().chain(scalar_family.iter()) {
12536 for (i, dt) in dts.iter().enumerate() {
12537 check(*op, *dt, ws[i], label);
12538 checked += 1;
12539 }
12540 }
12541 assert_eq!(
12542 checked, 96,
12543 "21 in-place unary (×4) + InplaceAffine/ClampInplace/PowIInplace (×4) = 96 \
12544 contract-sourced bindings",
12545 );
12546 }
12547
12548 #[test]
12569 fn global_bindings_registers_cast_family_from_contract() {
12570 let table = global_bindings();
12571
12572 let targets: [(DType, crate::kernel::KernelRef); 11] = [
12575 (DType::F32, cast_to_f32_cpu_wrapper),
12576 (DType::F64, cast_to_f64_cpu_wrapper),
12577 (DType::F16, cast_to_f16_cpu_wrapper),
12578 (DType::BF16, cast_to_bf16_cpu_wrapper),
12579 (DType::F8E4M3, cast_to_f8e4m3_cpu_wrapper),
12580 (DType::U8, cast_to_u8_cpu_wrapper),
12581 (DType::I8, cast_to_i8_cpu_wrapper),
12582 (DType::U32, cast_to_u32_cpu_wrapper),
12583 (DType::I16, cast_to_i16_cpu_wrapper),
12584 (DType::I32, cast_to_i32_cpu_wrapper),
12585 (DType::I64, cast_to_i64_cpu_wrapper),
12586 ];
12587 let all_dts: [DType; 11] = [
12588 DType::F32,
12589 DType::F64,
12590 DType::F16,
12591 DType::BF16,
12592 DType::F8E4M3,
12593 DType::U8,
12594 DType::I8,
12595 DType::U32,
12596 DType::I16,
12597 DType::I32,
12598 DType::I64,
12599 ];
12600
12601 let mut checked = 0usize;
12602 for (dst, expected) in targets {
12603 for src in all_dts {
12604 if src == dst {
12605 continue; }
12607 let key: &[DType] = &[src, dst];
12608 let alts = table.lookup_alternatives(OpKind::Cast, key, BackendId::Cpu);
12609 let entry = alts
12610 .iter()
12611 .find(|e| e.kernel as usize == expected as usize)
12612 .unwrap_or_else(|| {
12613 panic!(
12614 "Cast {src:?}->{dst:?}/Cpu: the production per-target wrapper must be \
12615 bound FROM the cast contract in global_bindings() — found {} \
12616 alternative(s) with sources {:?}",
12617 alts.len(),
12618 alts.iter().map(|e| e.kernel_source).collect::<Vec<_>>(),
12619 )
12620 });
12621 assert_eq!(
12622 entry.kernel_source, "portable-cpu",
12623 "Cast {src:?}->{dst:?}: family must be contract-sourced \
12624 (kernel_source=\"portable-cpu\"); got {:?}",
12625 entry.kernel_source,
12626 );
12627 assert!(
12628 !entry.caps.strided_input,
12629 "Cast {src:?}->{dst:?}: caps preserved (contiguous-only)",
12630 );
12631 assert!(
12632 entry.precision.bit_stable_on_same_hardware,
12633 "Cast {src:?}->{dst:?}: contract's bit-stable claim rode through",
12634 );
12635 checked += 1;
12636 }
12637 }
12638 assert_eq!(
12639 checked, 110,
12640 "11 target dtypes × 10 sources each = 110 directed cast pairs (identity excluded)"
12641 );
12642 }
12643
12644 #[test]
12646 fn binding_table_lookup_miss_errors() {
12647 let mut table = KernelBindingTable::new();
12648 register_cpu_kernels(&mut table);
12649 let result = table.lookup(
12653 OpKind::AddElementwise,
12654 &[DType::I64, DType::I64, DType::I64],
12655 BackendId::Cpu,
12656 );
12657 assert!(result.is_err());
12658 }
12659
12660 #[test]
12663 fn cpu_f8e4m3_cast_pairs_registered() {
12664 let mut table = KernelBindingTable::new();
12665 register_cpu_kernels(&mut table);
12666 for other in [DType::F32, DType::F16, DType::BF16] {
12667 for (src, dst) in [(DType::F8E4M3, other), (other, DType::F8E4M3)] {
12668 table
12669 .lookup(OpKind::Cast, &[src, dst], BackendId::Cpu)
12670 .unwrap_or_else(|e| panic!("Cast {src:?} → {dst:?} (CPU) not registered: {e}"));
12671 }
12672 }
12673 }
12674
12675 #[test]
12678 fn cpu_int_matmul_registered() {
12679 let mut table = KernelBindingTable::new();
12680 register_cpu_kernels(&mut table);
12681 for t in [DType::I8, DType::U8] {
12682 table
12683 .lookup(OpKind::MatMul, &[t, t, t], BackendId::Cpu)
12684 .unwrap_or_else(|e| panic!("MatMul {t:?} (CPU) not registered: {e}"));
12685 }
12686 }
12687
12688 #[test]
12693 fn b5_end_to_end_add_elementwise_f32() {
12694 let mut registry = CapabilityRegistry::new();
12696 registry.register(cpu_caps());
12697
12698 let backend =
12700 resolve_target_backend(®istry, OpKind::AddElementwise, DType::F32).expect("resolve");
12701 assert_eq!(backend, BackendId::Cpu);
12702
12703 let mut bindings = KernelBindingTable::new();
12705 register_cpu_kernels(&mut bindings);
12706
12707 let kernel = bindings
12709 .lookup(
12710 OpKind::AddElementwise,
12711 &[DType::F32, DType::F32, DType::F32],
12712 backend,
12713 )
12714 .expect("lookup");
12715
12716 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0, 4.0]);
12718 let rhs = fuel_memory::from_slice_cpu(&[10.0_f32, 20.0, 30.0, 40.0]);
12719 let out = fuel_memory::alloc_cpu_zeroed(DType::F32, 4).expect("alloc");
12720
12721 let inputs = vec![Arc::new(RwLock::new(lhs)), Arc::new(RwLock::new(rhs))];
12722 let mut outputs = vec![Arc::new(RwLock::new(out))];
12723
12724 kernel(&inputs, &mut outputs, &[], &OpParams::None).expect("kernel");
12726
12727 let out_guard = outputs[0].read().unwrap();
12729 if let BackendStorage::Cpu(c) = &out_guard.inner {
12730 let typed: &[f32] = c.as_slice().expect("cast");
12731 assert_eq!(typed, &[11.0, 22.0, 33.0, 44.0]);
12732 } else {
12733 panic!("expected CPU output");
12734 }
12735 }
12736
12737 #[test]
12739 fn b5_wrapper_arity_check() {
12740 let mut bindings = KernelBindingTable::new();
12741 register_cpu_kernels(&mut bindings);
12742 let kernel = bindings
12743 .lookup(
12744 OpKind::AddElementwise,
12745 &[DType::F32, DType::F32, DType::F32],
12746 BackendId::Cpu,
12747 )
12748 .unwrap();
12749
12750 let lhs = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0]);
12752 let out = fuel_memory::alloc_cpu_zeroed(DType::F32, 2).unwrap();
12753 let inputs = vec![Arc::new(RwLock::new(lhs))];
12754 let mut outputs = vec![Arc::new(RwLock::new(out))];
12755
12756 let result = kernel(&inputs, &mut outputs, &[], &OpParams::None);
12757 assert!(result.is_err());
12758 }
12759
12760 #[test]
12777 fn held_snapshot_does_not_block_a_concurrent_write() {
12778 use std::sync::mpsc;
12779 use std::time::Duration;
12780
12781 let (tx, rx) = mpsc::channel();
12782 std::thread::spawn(move || {
12783 let snap = global_bindings();
12785 let before = snap
12786 .lookup(
12787 OpKind::AddElementwise,
12788 &[DType::F32, DType::F32, DType::F32],
12789 BackendId::Cpu,
12790 )
12791 .is_ok();
12792
12793 extend_global_bindings(|_t| {});
12797
12798 let _snap2 = global_bindings();
12802
12803 let after = snap
12806 .lookup(
12807 OpKind::AddElementwise,
12808 &[DType::F32, DType::F32, DType::F32],
12809 BackendId::Cpu,
12810 )
12811 .is_ok();
12812
12813 tx.send((before, after)).ok();
12814 });
12815
12816 match rx.recv_timeout(Duration::from_secs(30)) {
12817 Ok((before, after)) => {
12818 assert!(
12819 before,
12820 "CPU AddElementwise+F32 should resolve in the snapshot"
12821 );
12822 assert!(
12823 after,
12824 "snapshot must stay readable after a concurrent table swap"
12825 );
12826 }
12827 Err(_) => panic!(
12828 "GAP-015 regression: holding a binding-table snapshot across \
12829 extend_global_bindings() deadlocked (30s watchdog). The read \
12830 path must hand out an owned snapshot, not a lock guard."
12831 ),
12832 }
12833 }
12834
12835 #[test]
12837 fn global_registry_auto_registers_cpu() {
12838 let r = global_registry();
12840 assert!(!r.backends().is_empty());
12841 let cpu = r.backends().iter().find(|c| c.backend_id == BackendId::Cpu);
12843 assert!(cpu.is_some(), "CPU should be auto-registered");
12844 }
12845
12846 #[test]
12849 fn global_bindings_auto_registers_cpu_wrappers() {
12850 let b = global_bindings();
12851 let result = b.lookup(
12852 OpKind::AddElementwise,
12853 &[DType::F32, DType::F32, DType::F32],
12854 BackendId::Cpu,
12855 );
12856 assert!(
12857 result.is_ok(),
12858 "CPU AddElementwise+F32 should be registered"
12859 );
12860 }
12861
12862 #[cfg(feature = "cuda")]
12870 #[test]
12871 fn global_bindings_auto_registers_cuda_paths() {
12872 let b = global_bindings();
12873 let cublas = b.lookup(
12875 OpKind::MatMul,
12876 &[DType::F32, DType::F32, DType::F32],
12877 BackendId::Cuda,
12878 );
12879 assert!(
12880 cublas.is_ok(),
12881 "cuBLAS F32 MatMul should be auto-registered on Cuda"
12882 );
12883
12884 let baracuda_int8 = b.lookup(
12887 OpKind::MatMul,
12888 &[DType::I8, DType::I8, DType::I8],
12889 BackendId::Cuda,
12890 );
12891 assert!(
12892 baracuda_int8.is_ok(),
12893 "baracuda int8 MatMul should be auto-registered on Cuda",
12894 );
12895
12896 let alts = b.lookup_alternatives(
12901 OpKind::NegElementwise,
12902 &[DType::F32, DType::F32],
12903 BackendId::Cuda,
12904 );
12905 assert!(
12906 alts.len() == 1,
12907 "expected exactly 1 CUDA Neg F32 alternative (baracuda); got {}",
12908 alts.len(),
12909 );
12910 }
12911
12912 #[test]
12921 fn audit_multi_backend_coverage() {
12922 let b = global_bindings();
12923 assert!(
12924 !b.is_empty(),
12925 "the global binding table is empty, so this audit would report nothing"
12926 );
12927
12928 type CoverageGroup = ((OpKind, Vec<DType>), Vec<(BackendId, usize)>);
12934 let mut grouped: Vec<CoverageGroup> = Vec::new();
12935 for (op, dtypes, backend, _precision) in b.iter_precision() {
12936 let key = (op, dtypes.to_vec());
12937 if let Some((_, entry)) = grouped.iter_mut().find(|(k, _)| *k == key) {
12938 if let Some(slot) = entry.iter_mut().find(|(be, _)| *be == backend) {
12939 slot.1 += 1;
12940 } else {
12941 entry.push((backend, 1));
12942 }
12943 } else {
12944 grouped.push((key, vec![(backend, 1)]));
12945 }
12946 }
12947 grouped.sort_by(|a, b| {
12949 let oa = format!("{:?}", a.0.0);
12950 let ob = format!("{:?}", b.0.0);
12951 oa.cmp(&ob).then_with(|| {
12952 let da = format!("{:?}", a.0.1);
12953 let db = format!("{:?}", b.0.1);
12954 da.cmp(&db)
12955 })
12956 });
12957
12958 let total_keys = grouped.len();
12959 let mut multi_backend = 0;
12960 let mut multi_impl_single_backend = 0;
12961 let mut single_alt = 0;
12962
12963 println!("\n=== Picker-alternatives audit ===");
12964 println!("Total unique (op, dtypes) keys: {total_keys}");
12965 println!("Total alternatives (sum): {}", b.len());
12966 println!();
12967 println!("Keys with >1 alternative (per Judge-picking working set):");
12968 for ((op, dtypes), backends) in &grouped {
12969 let total_alts: usize = backends.iter().map(|(_, n)| *n).sum();
12970 if total_alts <= 1 {
12971 single_alt += 1;
12972 continue;
12973 }
12974 if backends.len() > 1 {
12975 multi_backend += 1;
12976 } else {
12977 multi_impl_single_backend += 1;
12978 }
12979 let backend_summary: Vec<String> = backends
12980 .iter()
12981 .map(|(be, n)| {
12982 if *n == 1 {
12983 format!("{be:?}")
12984 } else {
12985 format!("{be:?}×{n}")
12986 }
12987 })
12988 .collect();
12989 println!(
12990 " {op:?} {dtypes:?} → {} alts: [{}]",
12991 total_alts,
12992 backend_summary.join(", "),
12993 );
12994 }
12995 println!();
12996 println!("--- summary ---");
12997 println!("Single-alternative keys (no pick to make): {single_alt}");
12998 println!("Multi-backend keys (cross-backend pick): {multi_backend}");
12999 println!("Multi-impl single-backend keys: {multi_impl_single_backend}");
13000 }
13001
13002 #[test]
13016 fn inflight_inc_dec_balances_to_baseline() {
13017 let loc = DeviceLocation::Cuda { gpu_id: 99_001 };
13018 let base = inflight_count(loc);
13019
13020 const N: u32 = 7;
13022 for i in 0..N {
13023 inflight_inc(loc);
13024 assert_eq!(
13025 inflight_count(loc),
13026 base + i + 1,
13027 "inc must raise the count"
13028 );
13029 }
13030 assert_eq!(inflight_count(loc), base + N);
13031
13032 for i in 0..N {
13034 inflight_dec(loc);
13035 assert_eq!(
13036 inflight_count(loc),
13037 base + N - i - 1,
13038 "dec must lower the count",
13039 );
13040 }
13041 assert_eq!(
13042 inflight_count(loc),
13043 base,
13044 "balanced inc/dec must return the count to its pre-sequence baseline",
13045 );
13046 }
13047
13048 #[test]
13052 fn inflight_dec_is_underflow_safe() {
13053 let untouched = DeviceLocation::Vulkan { gpu_id: 99_002 };
13054 inflight_dec(untouched);
13056 assert_eq!(
13057 inflight_count(untouched),
13058 0,
13059 "dec on an untouched device stays 0"
13060 );
13061
13062 let loc = DeviceLocation::Vulkan { gpu_id: 99_003 };
13063 inflight_inc(loc);
13064 inflight_dec(loc);
13065 assert_eq!(inflight_count(loc), 0);
13066 inflight_dec(loc);
13068 assert_eq!(
13069 inflight_count(loc),
13070 0,
13071 "extra dec past zero must saturate, not wrap"
13072 );
13073 }
13074
13075 #[test]
13079 fn inflight_counts_are_per_device_independent() {
13080 let cuda = DeviceLocation::Cuda { gpu_id: 99_004 };
13081 let vulkan = DeviceLocation::Vulkan { gpu_id: 99_004 };
13082 let base_c = inflight_count(cuda);
13083 let base_v = inflight_count(vulkan);
13084
13085 inflight_inc(cuda);
13086 inflight_inc(cuda);
13087 inflight_inc(vulkan);
13088
13089 assert_eq!(inflight_count(cuda), base_c + 2);
13090 assert_eq!(inflight_count(vulkan), base_v + 1);
13091
13092 inflight_dec(cuda);
13093 inflight_dec(cuda);
13094 inflight_dec(vulkan);
13095 assert_eq!(inflight_count(cuda), base_c);
13096 assert_eq!(inflight_count(vulkan), base_v);
13097 }
13098
13099 #[test]
13108 fn derive_backend_caps_mirrors_binding_table_output_dtypes() {
13109 fn stub(
13110 _i: &[Arc<RwLock<Storage>>],
13111 _o: &mut [Arc<RwLock<Storage>>],
13112 _l: &[Layout],
13113 _p: &OpParams,
13114 ) -> Result<()> {
13115 Ok(())
13116 }
13117 fn stub_b(
13118 _i: &[Arc<RwLock<Storage>>],
13119 _o: &mut [Arc<RwLock<Storage>>],
13120 _l: &[Layout],
13121 _p: &OpParams,
13122 ) -> Result<()> {
13123 Ok(())
13124 }
13125 fn stub_c(
13126 _i: &[Arc<RwLock<Storage>>],
13127 _o: &mut [Arc<RwLock<Storage>>],
13128 _l: &[Layout],
13129 _p: &OpParams,
13130 ) -> Result<()> {
13131 Ok(())
13132 }
13133
13134 let mut table = KernelBindingTable::new();
13135 table.register(
13137 OpKind::AddElementwise,
13138 &[DType::F32, DType::F32, DType::F32],
13139 BackendId::Cpu,
13140 stub,
13141 );
13142 table.register(
13144 OpKind::MulElementwise,
13145 &[DType::F32, DType::F32, DType::F32],
13146 BackendId::Vulkan,
13147 stub_b,
13148 );
13149 table.register(
13152 OpKind::MatMul,
13153 &[DType::F16, DType::F16, DType::F32],
13154 BackendId::Vulkan,
13155 stub_c,
13156 );
13157
13158 let vk0 = DeviceLocation::Vulkan { gpu_id: 0 };
13159 let caps = derive_backend_caps(BackendId::Vulkan, vk0, &table);
13160
13161 assert_eq!(caps.backend_id, BackendId::Vulkan);
13162 assert_eq!(caps.device_location, vk0);
13163 assert_eq!(caps.storage_substrate, SubstrateClass::VulkanBuffer);
13164 assert!(
13166 caps.op_dtype_support
13167 .contains(&(OpKind::MulElementwise, DType::F32))
13168 );
13169 assert!(
13171 caps.op_dtype_support
13172 .contains(&(OpKind::MatMul, DType::F32))
13173 );
13174 assert!(
13175 !caps
13176 .op_dtype_support
13177 .contains(&(OpKind::MatMul, DType::F16))
13178 );
13179 assert!(
13181 !caps
13182 .op_dtype_support
13183 .contains(&(OpKind::AddElementwise, DType::F32))
13184 );
13185 assert!(caps
13188 .transfer_paths
13189 .iter()
13190 .any(|&(dst, path)| dst == DeviceLocation::Cpu && path == TransferPath::HostStaging));
13191 }
13192
13193 #[test]
13204 #[ignore = "requires a live CUDA device"]
13205 #[cfg(feature = "cuda")]
13206 fn copy_from_cuda_wrapper_cuda_output_write_into_matches_old_behavior() {
13207 use fuel_cuda_backend::{CudaDevice, CudaStorageBytes};
13208
13209 let Ok(dev) = CudaDevice::new(0) else {
13210 return fuel_test_support::hardware::skip(
13211 fuel_test_support::hardware::Hardware::Cuda,
13212 fuel_test_support::hardware::Missing::device("CudaDevice::new(0) failed"),
13213 );
13214 };
13215
13216 let src_data: Vec<f32> = (0..17).map(|i| i as f32 * 1.5 - 3.0).collect();
13217 let src_bytes: Vec<u8> = src_data.iter().flat_map(|f| f.to_le_bytes()).collect();
13218 let src_cb = CudaStorageBytes::from_cpu_bytes(&dev, &src_bytes).expect("h2d src");
13219 let src_arc: Arc<RwLock<Storage>> = Arc::new(RwLock::new(Storage::new(
13220 BackendStorage::Cuda(src_cb),
13221 DType::F32,
13222 )));
13223
13224 let garbage: Vec<u8> = vec![0xAAu8; src_bytes.len()];
13228 let dst_cb = CudaStorageBytes::from_cpu_bytes(&dev, &garbage).expect("h2d garbage");
13229 let dst_arc: Arc<RwLock<Storage>> = Arc::new(RwLock::new(Storage::new(
13230 BackendStorage::Cuda(dst_cb),
13231 DType::F32,
13232 )));
13233
13234 let before_ptr = {
13236 let g = dst_arc.read().unwrap();
13237 let BackendStorage::Cuda(c) = &g.inner else {
13238 panic!("not cuda")
13239 };
13240 c.buffer() as *const _ as usize
13241 };
13242
13243 let inputs = vec![Arc::clone(&src_arc)];
13244 let mut outputs = vec![Arc::clone(&dst_arc)];
13245 copy_from_cuda_wrapper(&inputs, &mut outputs, &[], &OpParams::None)
13246 .expect("copy_from_cuda_wrapper (write-into CUDA branch)");
13247
13248 let after_ptr = {
13249 let g = dst_arc.read().unwrap();
13250 let BackendStorage::Cuda(c) = &g.inner else {
13251 panic!("not cuda")
13252 };
13253 c.buffer() as *const _ as usize
13254 };
13255 assert_eq!(
13256 before_ptr, after_ptr,
13257 "write-into must preserve the output's buffer identity (device address)"
13258 );
13259
13260 let new_bytes = {
13261 let g = dst_arc.read().unwrap();
13262 let BackendStorage::Cuda(c) = &g.inner else {
13263 panic!("not cuda")
13264 };
13265 c.to_cpu_bytes().expect("d2h new")
13266 };
13267
13268 let old_bytes = {
13271 let g = src_arc.read().unwrap();
13272 let BackendStorage::Cuda(c) = &g.inner else {
13273 panic!("not cuda")
13274 };
13275 let copied = c
13276 .slot_copy_to_new(0, c.len_bytes())
13277 .expect("slot_copy_to_new (old allocate-and-replace path)");
13278 copied.to_cpu_bytes().expect("d2h old")
13279 };
13280
13281 assert_eq!(
13282 new_bytes, old_bytes,
13283 "write-into result must byte-match the old slot_copy_to_new result"
13284 );
13285 assert_eq!(new_bytes, src_bytes, "and both must match the source bytes");
13286 }
13287}