1use crate::Tensor;
9use std::collections::HashMap;
10use std::sync::{Arc, RwLock};
11use torsh_core::sync::RwLockExt;
12use torsh_core::{device::DeviceType, dtype::TensorElement, error::Result};
13
14#[cfg(feature = "gpu")]
18pub struct GpuContext;
19
20#[cfg(feature = "gpu")]
21pub struct GpuKernel;
22
23#[cfg(feature = "gpu")]
24impl GpuContext {
25 pub fn new() -> Result<Self> {
26 Err(torsh_core::error::TorshError::InvalidArgument(
27 "GPU support temporarily unavailable".to_string(),
28 ))
29 }
30}
31
32#[cfg(feature = "gpu")]
33impl GpuKernel {
34 pub fn load(_context: &GpuContext, _name: &str) -> Result<Self> {
35 Err(torsh_core::error::TorshError::InvalidArgument(
36 "GPU support temporarily unavailable".to_string(),
37 ))
38 }
39
40 pub fn auto_tune(&mut self, _tuning_params: &[(String, f32)]) -> Result<()> {
41 Err(torsh_core::error::TorshError::InvalidArgument(
42 "GPU support temporarily unavailable".to_string(),
43 ))
44 }
45
46 pub fn enable_fusion(&mut self, _enable: bool) -> Result<()> {
47 Err(torsh_core::error::TorshError::InvalidArgument(
48 "GPU support temporarily unavailable".to_string(),
49 ))
50 }
51
52 pub fn enable_tensor_cores(&mut self, _enable: bool) -> Result<()> {
53 Err(torsh_core::error::TorshError::InvalidArgument(
54 "GPU support temporarily unavailable".to_string(),
55 ))
56 }
57
58 pub fn supports_tensor_cores(&self) -> bool {
59 false
60 }
61
62 pub fn execute<T>(&self, _input: &[T], _output: &mut [T]) -> Result<()> {
63 Err(torsh_core::error::TorshError::InvalidArgument(
64 "GPU support temporarily unavailable".to_string(),
65 ))
66 }
67}
68
69#[derive(Debug, Clone)]
71pub enum DeviceOptimization {
72 Cpu(CpuOptimization),
74 Gpu(GpuOptimization),
76 Metal(MetalOptimization),
78 WebGpu(WebGpuOptimization),
80}
81
82#[derive(Debug, Clone)]
84pub struct CpuOptimization {
85 pub use_simd: bool,
87 pub thread_count: Option<usize>,
89 pub cache_friendly: bool,
91 pub numa_aware: bool,
93}
94
95#[derive(Debug, Clone)]
97pub struct GpuOptimization {
98 pub use_pinned_memory: bool,
100 pub stream_count: u32,
102 pub mixed_precision: bool,
104 pub memory_pool_size: Option<usize>,
106
107 pub use_tensor_cores: bool,
110 pub auto_kernel_tuning: bool,
112 pub use_unified_memory: bool,
114 pub multi_gpu_strategy: MultiGpuStrategy,
116 pub backend_preference: Vec<GpuBackendType>,
118 pub memory_coalescing: bool,
120 pub kernel_fusion_level: u8,
122 pub dynamic_batching: bool,
124}
125
126#[derive(Debug, Clone)]
128pub enum MultiGpuStrategy {
129 Single,
131 DataParallel,
133 ModelParallel,
135 PipelineParallel,
137 Auto,
139}
140
141#[derive(Debug, Clone, PartialEq, Eq)]
143pub enum GpuBackendType {
144 Cuda,
146 Metal,
148 WebGpu,
150 Rocm,
152 OpenCl,
154}
155
156#[derive(Debug, Clone)]
158pub struct MetalOptimization {
159 pub use_mps: bool,
161 pub command_buffer_count: u32,
163 pub auto_memory_management: bool,
165}
166
167#[derive(Debug, Clone)]
169pub struct WebGpuOptimization {
170 pub use_compute_shaders: bool,
172 pub buffer_pool_size: Option<usize>,
174 pub pipeline_caching: bool,
176}
177
178#[derive(Debug)]
180pub struct OperationScheduler {
181 device_queues: HashMap<DeviceType, Vec<ScheduledOperation>>,
183 sync_state: HashMap<DeviceType, SyncState>,
185 operation_counter: Arc<RwLock<u64>>,
187}
188
189#[derive(Debug)]
191pub struct ScheduledOperation {
192 pub id: u64,
194 pub operation: OperationType,
196 pub priority: u8,
198 pub dependencies: Vec<DeviceType>,
200}
201
202#[derive(Debug)]
204pub enum OperationType {
205 Compute,
207 Transfer,
209 Synchronization,
211}
212
213#[derive(Debug)]
215pub struct SyncState {
216 pub last_operation: std::time::Instant,
218 pub pending_transfers: usize,
220 pub available: bool,
222}
223
224impl Default for CpuOptimization {
225 fn default() -> Self {
226 Self {
227 use_simd: true,
228 thread_count: None, cache_friendly: true,
230 numa_aware: true,
231 }
232 }
233}
234
235impl Default for GpuOptimization {
236 fn default() -> Self {
237 Self {
238 use_pinned_memory: true,
239 stream_count: 4,
240 mixed_precision: false,
241 memory_pool_size: Some(1024 * 1024 * 1024), use_tensor_cores: true, auto_kernel_tuning: true, use_unified_memory: true, multi_gpu_strategy: MultiGpuStrategy::Auto, backend_preference: vec![
249 GpuBackendType::Cuda, GpuBackendType::Metal, GpuBackendType::Rocm, GpuBackendType::WebGpu, GpuBackendType::OpenCl, ],
255 memory_coalescing: true, kernel_fusion_level: 2, dynamic_batching: true, }
259 }
260}
261
262impl Default for MetalOptimization {
263 fn default() -> Self {
264 Self {
265 use_mps: true,
266 command_buffer_count: 8,
267 auto_memory_management: true,
268 }
269 }
270}
271
272impl Default for WebGpuOptimization {
273 fn default() -> Self {
274 Self {
275 use_compute_shaders: true,
276 buffer_pool_size: Some(256 * 1024 * 1024), pipeline_caching: true,
278 }
279 }
280}
281
282impl<T: TensorElement + Copy> Tensor<T> {
283 pub fn to_device(&self, target_device: DeviceType) -> Result<Self> {
285 if self.device == target_device {
286 return Ok(self.clone());
287 }
288
289 let optimization = self.get_device_optimization(target_device);
291
292 match (self.device, target_device) {
294 (DeviceType::Cpu, DeviceType::Cuda(gpu_id)) => {
295 self.cpu_to_gpu_transfer(gpu_id as u32, optimization)
296 }
297 (DeviceType::Cuda(gpu_id), DeviceType::Cpu) => {
298 self.gpu_to_cpu_transfer(gpu_id as u32, optimization)
299 }
300 (DeviceType::Cpu, DeviceType::Metal(metal_id)) => {
301 self.cpu_to_metal_transfer(metal_id as u32, optimization)
302 }
303 (DeviceType::Metal(metal_id), DeviceType::Cpu) => {
304 self.metal_to_cpu_transfer(metal_id as u32, optimization)
305 }
306 _ => {
307 self.generic_device_transfer(target_device)
309 }
310 }
311 }
312
313 fn get_device_optimization(&self, device: DeviceType) -> DeviceOptimization {
315 match device {
316 DeviceType::Cpu => DeviceOptimization::Cpu(CpuOptimization::default()),
317 DeviceType::Cuda(_) => DeviceOptimization::Gpu(GpuOptimization::default()),
318 DeviceType::Metal(_) => DeviceOptimization::Metal(MetalOptimization::default()),
319 DeviceType::Wgpu(_) => DeviceOptimization::Gpu(GpuOptimization::default()),
320 }
321 }
322
323 fn cpu_to_gpu_transfer(&self, _gpu_id: u32, optimization: DeviceOptimization) -> Result<Self> {
330 #[cfg(feature = "gpu")]
331 if let Some(uploaded) =
332 crate::gpu_dispatch::try_upload_f32(self, DeviceType::Cuda(_gpu_id as usize))
333 {
334 return Ok(uploaded);
335 }
336
337 let data = self.to_vec()?;
338
339 if let DeviceOptimization::Gpu(gpu_opt) = optimization {
341 if gpu_opt.use_pinned_memory {
342 self.transfer_with_pinned_memory(data, DeviceType::Cuda(_gpu_id as usize))
344 } else {
345 Self::from_data(
347 data,
348 self.shape().dims().to_vec(),
349 DeviceType::Cuda(_gpu_id as usize),
350 )
351 }
352 } else {
353 Self::from_data(
354 data,
355 self.shape().dims().to_vec(),
356 DeviceType::Cuda(_gpu_id as usize),
357 )
358 }
359 }
360
361 fn gpu_to_cpu_transfer(&self, _gpu_id: u32, optimization: DeviceOptimization) -> Result<Self> {
363 let data = self.to_vec()?;
364
365 if let DeviceOptimization::Cpu(cpu_opt) = optimization {
367 if cpu_opt.numa_aware {
368 self.transfer_with_numa_awareness(data, DeviceType::Cpu)
370 } else {
371 Self::from_data(data, self.shape().dims().to_vec(), DeviceType::Cpu)
373 }
374 } else {
375 Self::from_data(data, self.shape().dims().to_vec(), DeviceType::Cpu)
376 }
377 }
378
379 fn cpu_to_metal_transfer(
381 &self,
382 _metal_id: u32,
383 optimization: DeviceOptimization,
384 ) -> Result<Self> {
385 let data = self.to_vec()?;
386
387 if let DeviceOptimization::Metal(metal_opt) = optimization {
389 if metal_opt.use_mps {
390 self.transfer_with_mps(data, DeviceType::Metal(_metal_id as usize))
392 } else {
393 Self::from_data(
395 data,
396 self.shape().dims().to_vec(),
397 DeviceType::Metal(_metal_id as usize),
398 )
399 }
400 } else {
401 Self::from_data(
402 data,
403 self.shape().dims().to_vec(),
404 DeviceType::Metal(_metal_id as usize),
405 )
406 }
407 }
408
409 fn metal_to_cpu_transfer(
411 &self,
412 _metal_id: u32,
413 optimization: DeviceOptimization,
414 ) -> Result<Self> {
415 let data = self.to_vec()?;
416
417 if let DeviceOptimization::Cpu(cpu_opt) = optimization {
419 if cpu_opt.cache_friendly {
420 self.transfer_with_cache_optimization(data, DeviceType::Cpu)
422 } else {
423 Self::from_data(data, self.shape().dims().to_vec(), DeviceType::Cpu)
425 }
426 } else {
427 Self::from_data(data, self.shape().dims().to_vec(), DeviceType::Cpu)
428 }
429 }
430
431 fn generic_device_transfer(&self, target_device: DeviceType) -> Result<Self> {
433 let data = self.to_vec()?;
434 Self::from_data(data, self.shape().dims().to_vec(), target_device)
435 }
436
437 fn transfer_with_pinned_memory(&self, data: Vec<T>, target_device: DeviceType) -> Result<Self> {
439 Self::from_data(data, self.shape().dims().to_vec(), target_device)
441 }
442
443 fn transfer_with_numa_awareness(
445 &self,
446 data: Vec<T>,
447 target_device: DeviceType,
448 ) -> Result<Self> {
449 Self::from_data(data, self.shape().dims().to_vec(), target_device)
451 }
452
453 fn transfer_with_mps(&self, data: Vec<T>, target_device: DeviceType) -> Result<Self> {
455 Self::from_data(data, self.shape().dims().to_vec(), target_device)
457 }
458
459 fn transfer_with_cache_optimization(
461 &self,
462 data: Vec<T>,
463 target_device: DeviceType,
464 ) -> Result<Self> {
465 let optimized_data = self.optimize_for_cache(data)?;
467 Self::from_data(optimized_data, self.shape().dims().to_vec(), target_device)
468 }
469
470 fn optimize_for_cache(&self, data: Vec<T>) -> Result<Vec<T>> {
472 Ok(data)
474 }
475
476 pub fn synchronize_devices(&self, devices: &[DeviceType]) -> Result<()> {
478 for device in devices {
480 self.synchronize_device(*device)?;
481 }
482 Ok(())
483 }
484
485 fn synchronize_device(&self, _device: DeviceType) -> Result<()> {
487 Ok(())
489 }
490
491 pub fn can_transfer_efficiently(&self, target_device: DeviceType) -> bool {
493 match (self.device, target_device) {
494 (a, b) if a == b => true,
496 (DeviceType::Cpu, DeviceType::Cuda(_)) | (DeviceType::Cuda(_), DeviceType::Cpu) => true,
498 (DeviceType::Cpu, DeviceType::Metal(_)) | (DeviceType::Metal(_), DeviceType::Cpu) => {
500 true
501 }
502 _ => false,
504 }
505 }
506
507 pub fn get_transfer_strategy(&self, target_device: DeviceType) -> TransferStrategy {
509 match (self.device, target_device) {
510 (a, b) if a == b => TransferStrategy::NoTransfer,
511 (DeviceType::Cpu, DeviceType::Cuda(_)) => TransferStrategy::DirectTransfer,
512 (DeviceType::Cuda(_), DeviceType::Cpu) => TransferStrategy::DirectTransfer,
513 (DeviceType::Cpu, DeviceType::Metal(_)) => TransferStrategy::DirectTransfer,
514 (DeviceType::Metal(_), DeviceType::Cpu) => TransferStrategy::DirectTransfer,
515 _ => TransferStrategy::ThroughCpu,
516 }
517 }
518}
519
520#[derive(Debug, Clone, PartialEq)]
522pub enum TransferStrategy {
523 NoTransfer,
525 DirectTransfer,
527 ThroughCpu,
529}
530
531impl OperationScheduler {
532 pub fn new() -> Self {
534 Self {
535 device_queues: HashMap::new(),
536 sync_state: HashMap::new(),
537 operation_counter: Arc::new(RwLock::new(0)),
538 }
539 }
540
541 pub fn schedule_operation(
543 &mut self,
544 device: DeviceType,
545 operation: OperationType,
546 priority: u8,
547 dependencies: Vec<DeviceType>,
548 ) -> Result<u64> {
549 let mut counter = self.operation_counter.write_or_recover();
551 *counter += 1;
552 let op_id = *counter;
553 drop(counter);
554
555 let scheduled_op = ScheduledOperation {
557 id: op_id,
558 operation,
559 priority,
560 dependencies,
561 };
562
563 self.device_queues
565 .entry(device)
566 .or_default()
567 .push(scheduled_op);
568
569 if let Some(queue) = self.device_queues.get_mut(&device) {
571 queue.sort_by(|a, b| b.priority.cmp(&a.priority));
572 }
573
574 self.sync_state.entry(device).or_insert_with(|| SyncState {
576 last_operation: std::time::Instant::now(),
577 pending_transfers: 0,
578 available: true,
579 });
580
581 Ok(op_id)
582 }
583
584 pub fn execute_next_operation(&mut self, device: DeviceType) -> Result<Option<u64>> {
586 let op = if let Some(queue) = self.device_queues.get_mut(&device) {
588 if queue.is_empty() {
589 None
590 } else {
591 Some(queue.remove(0)) }
593 } else {
594 None
595 };
596
597 if let Some(op) = op {
598 let dependencies_satisfied = self.check_dependencies(&op.dependencies)?;
600
601 if dependencies_satisfied {
602 self.execute_operation(&op)?;
604
605 if let Some(sync_state) = self.sync_state.get_mut(&device) {
607 sync_state.last_operation = std::time::Instant::now();
608 }
609
610 Ok(Some(op.id))
611 } else {
612 if let Some(queue) = self.device_queues.get_mut(&device) {
614 queue.insert(0, op);
615 }
616 Ok(None)
617 }
618 } else {
619 Ok(None)
620 }
621 }
622
623 fn check_dependencies(&self, dependencies: &[DeviceType]) -> Result<bool> {
625 for &dep_device in dependencies {
626 if let Some(sync_state) = self.sync_state.get(&dep_device) {
627 if !sync_state.available {
628 return Ok(false);
629 }
630 }
631 }
632 Ok(true)
633 }
634
635 fn execute_operation(&self, _operation: &ScheduledOperation) -> Result<()> {
637 std::thread::sleep(std::time::Duration::from_millis(1));
639 Ok(())
640 }
641
642 pub fn get_queue_length(&self, device: DeviceType) -> usize {
644 self.device_queues
645 .get(&device)
646 .map_or(0, |queue| queue.len())
647 }
648
649 pub fn clear_device_queue(&mut self, device: DeviceType) {
651 self.device_queues.remove(&device);
652 }
653}
654
655impl Default for OperationScheduler {
656 fn default() -> Self {
657 Self::new()
658 }
659}
660
661static GLOBAL_SCHEDULER: parking_lot::Mutex<Option<OperationScheduler>> =
663 parking_lot::Mutex::new(None);
664
665pub fn get_global_scheduler() -> parking_lot::MutexGuard<'static, Option<OperationScheduler>> {
667 let mut guard = GLOBAL_SCHEDULER.lock();
668 if guard.is_none() {
669 *guard = Some(OperationScheduler::new());
670 }
671 guard
672}
673
674pub fn initialize_global_scheduler() -> Result<()> {
676 let mut guard = GLOBAL_SCHEDULER.lock();
677 *guard = Some(OperationScheduler::new());
678 Ok(())
679}
680
681#[cfg(feature = "gpu")]
683impl<T: TensorElement + Copy + Default> Tensor<T> {
684 pub fn execute_gpu_kernel(&self, kernel_name: &str, _params: Vec<T>) -> Result<Self> {
686 let gpu_opt = match self.get_device_optimization(self.device) {
687 DeviceOptimization::Gpu(opt) => opt,
688 _ => {
689 return Err(torsh_core::error::TorshError::InvalidArgument(
690 "GPU kernel execution requires GPU device".to_string(),
691 ))
692 }
693 };
694
695 let gpu_context = self.create_optimal_gpu_context(&gpu_opt)?;
697
698 let input_buffer = self.create_gpu_buffer(&gpu_context, &gpu_opt)?;
700
701 let kernel = self.select_optimal_kernel(&gpu_context, kernel_name, &gpu_opt)?;
703
704 let mut output_buffer = vec![T::default(); input_buffer.len()];
706 kernel.execute(&input_buffer, &mut output_buffer)?;
707
708 self.gpu_buffer_to_tensor(output_buffer, &gpu_context, &gpu_opt)
710 }
711
712 #[allow(dead_code)]
715 fn create_optimal_gpu_context(&self, _gpu_opt: &GpuOptimization) -> Result<GpuContext> {
716 Err(torsh_core::error::TorshError::InvalidArgument(
748 "GPU backend creation temporarily disabled".to_string(),
749 ))
750 }
751
752 #[allow(dead_code)]
755 fn create_gpu_buffer(&self, _context: &GpuContext, _gpu_opt: &GpuOptimization) -> Result<Vec<T>>
756 where
757 T: Copy,
758 {
759 let data = self.to_vec()?;
760 Ok(data)
772 }
773
774 fn select_optimal_kernel(
776 &self,
777 context: &GpuContext,
778 kernel_name: &str,
779 gpu_opt: &GpuOptimization,
780 ) -> Result<GpuKernel> {
781 let mut kernel = GpuKernel::load(context, kernel_name).map_err(|e| {
782 torsh_core::error::TorshError::InvalidArgument(format!(
783 "Failed to load kernel '{}': {}",
784 kernel_name, e
785 ))
786 })?;
787
788 if gpu_opt.auto_kernel_tuning {
789 kernel.auto_tune(&[])?;
792 }
793
794 if gpu_opt.use_tensor_cores && kernel.supports_tensor_cores() {
795 kernel.enable_tensor_cores(true)?;
797 }
798
799 if gpu_opt.kernel_fusion_level > 0 {
800 kernel.enable_fusion(gpu_opt.kernel_fusion_level > 0)?;
802 }
803
804 Ok(kernel)
805 }
806
807 #[allow(dead_code)]
810 fn gpu_buffer_to_tensor(
811 &self,
812 buffer: Vec<T>, _context: &GpuContext,
814 _gpu_opt: &GpuOptimization,
815 ) -> Result<Self>
816 where
817 T: Copy,
818 {
819 Self::from_data(buffer, self.shape().dims().to_vec(), self.device)
829 }
830
831 pub fn distribute_multi_gpu(
833 &self,
834 gpu_count: usize,
835 strategy: Option<MultiGpuStrategy>,
836 ) -> Result<Vec<Self>> {
837 if gpu_count <= 1 {
838 return Ok(vec![self.clone()]);
839 }
840
841 let strategy = strategy.unwrap_or(MultiGpuStrategy::Auto);
842 let effective_strategy = match strategy {
843 MultiGpuStrategy::Auto => self.select_optimal_multi_gpu_strategy(gpu_count),
844 s => s,
845 };
846
847 match effective_strategy {
848 MultiGpuStrategy::DataParallel => self.data_parallel_distribution(gpu_count),
849 MultiGpuStrategy::ModelParallel => self.model_parallel_distribution(gpu_count),
850 MultiGpuStrategy::PipelineParallel => self.pipeline_parallel_distribution(gpu_count),
851 _ => Ok(vec![self.clone()]), }
853 }
854
855 fn select_optimal_multi_gpu_strategy(&self, gpu_count: usize) -> MultiGpuStrategy {
857 let _total_elements = self.numel();
858 let shape = self.shape();
859 let dims = shape.dims();
860
861 if dims.len() > 0 && dims[0] >= gpu_count * 4 {
863 return MultiGpuStrategy::DataParallel;
864 }
865
866 if dims.len() > 1 && dims.iter().skip(1).product::<usize>() > 1024 * 1024 {
868 return MultiGpuStrategy::ModelParallel;
869 }
870
871 if dims.len() > 3 {
873 return MultiGpuStrategy::PipelineParallel;
874 }
875
876 MultiGpuStrategy::DataParallel
878 }
879
880 fn data_parallel_distribution(&self, gpu_count: usize) -> Result<Vec<Self>> {
882 let shape = self.shape();
883 let dims = shape.dims();
884 if dims.is_empty() {
885 return Err(torsh_core::error::TorshError::InvalidArgument(
886 "Cannot distribute scalar tensor".to_string(),
887 ));
888 }
889
890 let batch_size = dims[0];
891 let chunk_size = (batch_size + gpu_count - 1) / gpu_count; let mut distributed_tensors = Vec::with_capacity(gpu_count);
894 let data = self.to_vec()?;
895 let elements_per_batch = dims.iter().skip(1).product::<usize>();
896
897 for gpu_id in 0..gpu_count {
898 let start_batch = gpu_id * chunk_size;
899 let end_batch = ((gpu_id + 1) * chunk_size).min(batch_size);
900
901 if start_batch >= batch_size {
902 break; }
904
905 let start_idx = start_batch * elements_per_batch;
906 let end_idx = end_batch * elements_per_batch;
907 let chunk_data = data[start_idx..end_idx].to_vec();
908
909 let mut chunk_dims = dims.to_vec();
910 chunk_dims[0] = end_batch - start_batch;
911
912 let chunk_tensor = Self::from_data(chunk_data, chunk_dims, DeviceType::Cuda(gpu_id))?;
913
914 distributed_tensors.push(chunk_tensor);
915 }
916
917 Ok(distributed_tensors)
918 }
919
920 fn model_parallel_distribution(&self, gpu_count: usize) -> Result<Vec<Self>> {
922 let shape = self.shape();
923 let dims = shape.dims();
924 if dims.len() < 2 {
925 return Err(torsh_core::error::TorshError::InvalidArgument(
926 "Model parallel requires at least 2D tensor".to_string(),
927 ));
928 }
929
930 let feature_dim = dims.len() - 1;
932 let feature_size = dims[feature_dim];
933 let chunk_size = (feature_size + gpu_count - 1) / gpu_count;
934
935 let mut distributed_tensors = Vec::with_capacity(gpu_count);
936 let _data = self.to_vec()?;
937
938 for gpu_id in 0..gpu_count {
939 let start_feature = gpu_id * chunk_size;
940 let end_feature = ((gpu_id + 1) * chunk_size).min(feature_size);
941
942 if start_feature >= feature_size {
943 break;
944 }
945
946 let mut chunk_dims = dims.to_vec();
949 chunk_dims[feature_dim] = end_feature - start_feature;
950
951 let chunk_size_total: usize = chunk_dims.iter().product();
953 let chunk_data = vec![T::default(); chunk_size_total];
954
955 let chunk_tensor = Self::from_data(chunk_data, chunk_dims, DeviceType::Cuda(gpu_id))?;
956
957 distributed_tensors.push(chunk_tensor);
958 }
959
960 Ok(distributed_tensors)
961 }
962
963 fn pipeline_parallel_distribution(&self, gpu_count: usize) -> Result<Vec<Self>> {
965 let mut distributed_tensors = Vec::with_capacity(gpu_count);
968
969 for gpu_id in 0..gpu_count {
970 let pipeline_tensor = Self::from_data(
971 self.to_vec()?,
972 self.shape().dims().to_vec(),
973 DeviceType::Cuda(gpu_id),
974 )?;
975 distributed_tensors.push(pipeline_tensor);
976 }
977
978 Ok(distributed_tensors)
979 }
980
981 #[allow(dead_code)]
984 pub fn enable_mixed_precision(
985 &mut self,
986 _precision: i32, ) -> Result<()> {
988 Err(torsh_core::error::TorshError::InvalidArgument(
998 "Mixed precision temporarily disabled".to_string(),
999 ))
1000 }
1001}
1002
1003#[cfg(test)]
1004mod tests {
1005 use super::*;
1006 use crate::Tensor;
1007
1008 #[test]
1009 fn test_device_transfer() {
1010 let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0], vec![2, 2], DeviceType::Cpu)
1011 .expect("tensor creation should succeed");
1012
1013 let same_device = tensor
1015 .to_device(DeviceType::Cpu)
1016 .expect("device transfer should succeed");
1017 assert_eq!(same_device.device(), DeviceType::Cpu);
1018
1019 assert_eq!(
1021 tensor.get_transfer_strategy(DeviceType::Cpu),
1022 TransferStrategy::NoTransfer
1023 );
1024 assert_eq!(
1025 tensor.get_transfer_strategy(DeviceType::Cuda(0)),
1026 TransferStrategy::DirectTransfer
1027 );
1028 }
1029
1030 #[test]
1031 fn test_operation_scheduler() {
1032 let mut scheduler = OperationScheduler::new();
1033
1034 let op1 = scheduler
1036 .schedule_operation(DeviceType::Cpu, OperationType::Compute, 5, vec![])
1037 .expect("operation should succeed");
1038
1039 let op2 = scheduler
1040 .schedule_operation(DeviceType::Cpu, OperationType::Compute, 10, vec![])
1041 .expect("operation should succeed");
1042
1043 assert_eq!(
1045 scheduler
1046 .execute_next_operation(DeviceType::Cpu)
1047 .expect("operation execution should succeed"),
1048 Some(op2)
1049 );
1050 assert_eq!(
1051 scheduler
1052 .execute_next_operation(DeviceType::Cpu)
1053 .expect("operation execution should succeed"),
1054 Some(op1)
1055 );
1056 }
1057
1058 #[test]
1059 fn test_transfer_efficiency() {
1060 let tensor = Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0], vec![2, 2], DeviceType::Cpu)
1061 .expect("tensor creation should succeed");
1062
1063 assert!(tensor.can_transfer_efficiently(DeviceType::Cpu));
1065
1066 assert!(tensor.can_transfer_efficiently(DeviceType::Cuda(0)));
1068
1069 assert!(tensor.can_transfer_efficiently(DeviceType::Metal(0)));
1071 }
1072
1073 #[test]
1074 fn test_device_optimization_defaults() {
1075 let cpu_opt = CpuOptimization::default();
1076 assert!(cpu_opt.use_simd);
1077 assert!(cpu_opt.cache_friendly);
1078 assert!(cpu_opt.numa_aware);
1079
1080 let gpu_opt = GpuOptimization::default();
1081 assert!(gpu_opt.use_pinned_memory);
1082 assert_eq!(gpu_opt.stream_count, 4);
1083 assert!(!gpu_opt.mixed_precision);
1084 }
1085
1086 #[test]
1087 fn test_global_scheduler() {
1088 initialize_global_scheduler().expect("scheduler initialization should succeed");
1089
1090 {
1091 let mut scheduler = get_global_scheduler();
1092 let scheduler = scheduler
1093 .as_mut()
1094 .expect("mutable reference should be available");
1095
1096 let op_id = scheduler
1097 .schedule_operation(DeviceType::Cpu, OperationType::Compute, 5, vec![])
1098 .expect("scheduler initialization should succeed");
1099
1100 assert_eq!(scheduler.get_queue_length(DeviceType::Cpu), 1);
1101 assert_eq!(
1102 scheduler
1103 .execute_next_operation(DeviceType::Cpu)
1104 .expect("operation execution should succeed"),
1105 Some(op_id)
1106 );
1107 }
1108 }
1109}