1use std::sync::Arc;
20use torsh_core::{
21 dtype::TensorElement,
22 error::{Result, TorshError},
23};
24
25use crate::memory_pool::global_acquire_uninit;
26
27#[cfg(feature = "simd")]
30mod simd_imports {
31 pub use scirs2_core::ndarray::Array1;
39}
40
41#[cfg(feature = "simd")]
42use simd_imports::*;
43
44#[cfg(feature = "parallel")]
46use scirs2_core::chunking::{
47 CacheAwareness, ChunkConfig, ChunkStrategy, ComputeIntensity, GpuChunkSettings, MemoryPattern,
48 NumaStrategy,
49};
50
51use crate::core_ops::{Operation, Tensor};
63
64#[cfg(feature = "simd")]
66pub(crate) mod adaptive_simd {
67 use super::*;
68 use scirs2_core::ndarray::ArrayView1;
69
70 pub fn adaptive_simd_relu_f32(input: &ArrayView1<f32>) -> Array1<f32> {
73 scirs2_core::simd::activation::simd_relu_f32(input)
74 }
75
76 pub fn adaptive_simd_sigmoid_f32(input: &ArrayView1<f32>) -> Array1<f32> {
79 scirs2_core::simd::transcendental::simd_sigmoid_f32(input)
80 }
81
82 pub fn adaptive_simd_gelu_f32(input: &ArrayView1<f32>) -> Array1<f32> {
85 scirs2_core::simd::transcendental::simd_gelu_f32(input)
86 }
87}
88
89#[cfg(feature = "simd")]
90#[cfg(feature = "parallel")]
92mod intelligent_chunking {
93 use super::*;
94
95 #[derive(Debug, Clone, Copy)]
97 pub enum TensorOpType {
98 ElementWise,
100 Activation,
102 }
103
104 pub fn create_optimal_chunk_config(
108 tensor_size: usize,
109 op_type: TensorOpType,
110 _device: torsh_core::device::DeviceType,
111 is_gpu_available: bool,
112 ) -> ChunkConfig {
113 match op_type {
114 TensorOpType::ElementWise => ChunkConfig {
115 strategy: if tensor_size > 100_000 {
116 ChunkStrategy::MemoryOptimized
117 } else {
118 ChunkStrategy::CacheOptimized
119 },
120 min_chunk_size: 64,
121 max_chunk_size: 8192,
122 prefer_work_stealing: true,
123 memory_pattern: MemoryPattern::Sequential,
124 compute_intensity: ComputeIntensity::MemoryBound,
125 enable_monitoring: false,
126 load_balance_factor: 0.1,
127 cache_awareness: CacheAwareness::L2,
128 numa_strategy: NumaStrategy::LocalPreferred,
129 gpu_settings: if is_gpu_available {
130 Some(GpuChunkSettings::default())
131 } else {
132 None
133 },
134 },
135
136 TensorOpType::Activation => ChunkConfig {
137 strategy: ChunkStrategy::CacheOptimized,
138 min_chunk_size: 64,
139 max_chunk_size: 4096,
140 prefer_work_stealing: true,
141 memory_pattern: MemoryPattern::Sequential,
142 compute_intensity: ComputeIntensity::ComputeIntensive,
143 enable_monitoring: false,
144 load_balance_factor: 0.1,
145 cache_awareness: CacheAwareness::L1,
146 numa_strategy: NumaStrategy::LocalPreferred,
147 gpu_settings: if is_gpu_available {
148 Some(GpuChunkSettings {
149 gpu_memory_ratio: 0.7,
150 gpu_min_chunk: 2048,
151 overlap_compute: true,
152 gpu_bandwidth: None, transfer_bandwidth: None, })
155 } else {
156 None
157 },
158 },
159 }
160 }
161
162 pub fn intelligent_parallel_process<T, F, R>(
164 data: Vec<T>,
165 op_type: TensorOpType,
166 device: torsh_core::device::DeviceType,
167 operation: F,
168 ) -> Vec<R>
169 where
170 T: Send + Sync,
171 R: Send + Sync,
172 F: Fn(T) -> R + Send + Sync,
173 {
174 let is_gpu_available = matches!(
175 device,
176 torsh_core::device::DeviceType::Cuda(_)
177 | torsh_core::device::DeviceType::Metal(_)
178 | torsh_core::device::DeviceType::Wgpu(_)
179 );
180
181 let _chunk_config =
182 create_optimal_chunk_config(data.len(), op_type, device, is_gpu_available);
183
184 #[cfg(feature = "parallel")]
186 {
187 use scirs2_core::parallel_ops::*;
188 data.into_par_iter().map(operation).collect()
189 }
190 #[cfg(not(feature = "parallel"))]
191 {
192 data.into_iter().map(operation).collect()
193 }
194 }
195}
196
197#[cfg(feature = "parallel")]
198use intelligent_chunking::*;
199
200fn can_broadcast(shape1: &[usize], shape2: &[usize]) -> bool {
202 let max_dims = shape1.len().max(shape2.len());
203
204 for i in 0..max_dims {
205 let dim1 = if i < shape1.len() {
206 shape1[shape1.len() - 1 - i]
207 } else {
208 1
209 };
210 let dim2 = if i < shape2.len() {
211 shape2[shape2.len() - 1 - i]
212 } else {
213 1
214 };
215
216 if dim1 != dim2 && dim1 != 1 && dim2 != 1 {
217 return false;
218 }
219 }
220 true
221}
222
223fn compute_broadcast_shape(shape1: &[usize], shape2: &[usize]) -> Result<Vec<usize>> {
225 let max_dims = shape1.len().max(shape2.len());
226 let mut result = Vec::with_capacity(max_dims);
227
228 for i in 0..max_dims {
229 let dim1 = if i < shape1.len() {
230 shape1[shape1.len() - 1 - i]
231 } else {
232 1
233 };
234 let dim2 = if i < shape2.len() {
235 shape2[shape2.len() - 1 - i]
236 } else {
237 1
238 };
239
240 if dim1 == dim2 {
241 result.push(dim1);
242 } else if dim1 == 1 {
243 result.push(dim2);
244 } else if dim2 == 1 {
245 result.push(dim1);
246 } else {
247 return Err(TorshError::ShapeMismatch {
248 expected: shape1.to_vec(),
249 got: shape2.to_vec(),
250 });
251 }
252 }
253
254 result.reverse();
255 Ok(result)
256}
257
258fn compute_broadcast_index(
260 flat_idx: usize,
261 broadcast_shape: &[usize],
262 original_shape: &[usize],
263) -> usize {
264 let mut result = 0;
265 let mut remaining = flat_idx;
266
267 let dims_diff = broadcast_shape.len() - original_shape.len();
268
269 for (i, &broadcast_dim) in broadcast_shape.iter().enumerate() {
270 let stride = broadcast_shape[i + 1..].iter().product::<usize>().max(1);
272 let coord = remaining / stride;
273 remaining %= stride;
274
275 debug_assert!(
277 coord < broadcast_dim,
278 "Coordinate {} out of bounds for dimension {} of size {}",
279 coord,
280 i,
281 broadcast_dim
282 );
283
284 if i >= dims_diff {
285 let original_dim = original_shape[i - dims_diff];
286 let adjusted_coord = if original_dim == 1 { 0 } else { coord };
287 result = result * original_dim + adjusted_coord;
288 }
289 }
290
291 result
292}
293
294impl<T: TensorElement + Copy> Tensor<T> {
295 pub fn add_scalar_(&mut self, scalar: T) -> Result<()>
297 where
298 T: Copy + std::ops::Add<Output = T>,
299 {
300 self.make_unique()?;
302 self.apply_(|x| x + scalar)
303 }
304
305 pub fn add_scalar(&self, scalar: T) -> Result<Self>
307 where
308 T: Copy + std::ops::Add<Output = T>,
309 {
310 self.map(|x| x + scalar)
311 }
312
313 pub fn sub_scalar_(&mut self, scalar: T) -> Result<()>
315 where
316 T: Copy + std::ops::Sub<Output = T>,
317 {
318 self.make_unique()?;
319 self.apply_(|x| x - scalar)
320 }
321
322 pub fn sub_scalar(&self, scalar: T) -> Result<Self>
324 where
325 T: Copy + std::ops::Sub<Output = T>,
326 {
327 self.map(|x| x - scalar)
328 }
329
330 pub fn mul_scalar_(&mut self, scalar: T) -> Result<()>
332 where
333 T: Copy + std::ops::Mul<Output = T>,
334 {
335 self.make_unique()?;
337 self.apply_(|x| x * scalar)
338 }
339
340 pub fn mul_scalar(&self, scalar: T) -> Result<Self>
342 where
343 T: Copy + std::ops::Mul<Output = T>,
344 {
345 self.map(|x| x * scalar)
346 }
347
348 pub fn div_scalar_(&mut self, scalar: T) -> Result<()>
350 where
351 T: Copy + std::ops::Div<Output = T>,
352 {
353 self.make_unique()?;
354 self.apply_(|x| x / scalar)
355 }
356
357 pub fn div_scalar(&self, scalar: T) -> Result<Self>
359 where
360 T: Copy + std::ops::Div<Output = T>,
361 {
362 self.map(|x| x / scalar)
363 }
364
365 pub fn add(&self, other: &Self) -> Result<Self>
367 where
368 T: std::ops::Add<Output = T>,
369 {
370 if self.shape() != other.shape() {
372 return self.broadcast_add(other);
373 }
374
375 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() {
377 let self_data = self.data()?;
378 if self_data.len() >= 1024 {
379 let other_data = other.data()?;
380 let a_f32: &[f32] = unsafe {
382 std::slice::from_raw_parts(self_data.as_ptr() as *const f32, self_data.len())
383 };
384 let b_f32: &[f32] = unsafe {
385 std::slice::from_raw_parts(other_data.as_ptr() as *const f32, other_data.len())
386 };
387 let n = self_data.len();
388 let mut buf = global_acquire_uninit::<f32>(n);
389 {
390 let uninit = buf.as_uninit_slice_mut();
391 for slot in uninit.iter_mut() {
392 slot.write(0.0);
393 }
394 }
395 let mut out = buf.into_vec(n);
396 crate::simd_ops_f32::add_into_f32(a_f32, b_f32, &mut out);
397 let result_data: Vec<T> = unsafe {
399 let mut v = std::mem::ManuallyDrop::new(out);
400 Vec::from_raw_parts(v.as_mut_ptr() as *mut T, v.len(), v.capacity())
401 };
402 let mut result =
403 Self::from_data(result_data, self.shape().dims().to_vec(), self.device)?;
404 if self.requires_grad || other.requires_grad {
405 result.requires_grad = true;
406 result.operation = Operation::Add {
407 lhs: Arc::new(self.clone()),
408 rhs: Arc::new(other.clone()),
409 };
410 }
411 return Ok(result);
412 }
413 }
414
415 let mut result = self.elementwise_operation(other, |a, b| a + b)?;
417
418 if self.requires_grad || other.requires_grad {
420 result.requires_grad = true;
421 result.operation = Operation::Add {
422 lhs: Arc::new(self.clone()),
423 rhs: Arc::new(other.clone()),
424 };
425 }
426
427 Ok(result)
428 }
429
430 fn broadcast_add(&self, other: &Self) -> Result<Self>
432 where
433 T: std::ops::Add<Output = T>,
434 {
435 let self_shape_binding = self.shape();
437 let other_shape_binding = other.shape();
438 let self_shape = self_shape_binding.dims();
439 let other_shape = other_shape_binding.dims();
440
441 if !can_broadcast(self_shape, other_shape) {
443 return Err(TorshError::ShapeMismatch {
444 expected: self_shape.to_vec(),
445 got: other_shape.to_vec(),
446 });
447 }
448
449 let broadcast_shape = compute_broadcast_shape(self_shape, other_shape)?;
451
452 let self_data = self.data()?;
454 let other_data = other.data()?;
455
456 let total_elems: usize = broadcast_shape.iter().product();
458 let mut buf = global_acquire_uninit::<T>(total_elems);
459 let uninit = buf.as_uninit_slice_mut();
460 let mut count = 0;
461
462 for i in 0..total_elems {
463 let self_idx = compute_broadcast_index(i, &broadcast_shape, self_shape);
464 let other_idx = compute_broadcast_index(i, &broadcast_shape, other_shape);
465
466 let self_val = *self_data
467 .get(self_idx)
468 .ok_or_else(|| TorshError::IndexError {
469 index: self_idx,
470 size: self_data.len(),
471 })?;
472 let other_val = *other_data
473 .get(other_idx)
474 .ok_or_else(|| TorshError::IndexError {
475 index: other_idx,
476 size: other_data.len(),
477 })?;
478 uninit[count].write(self_val + other_val);
479 count += 1;
480 }
481
482 let result_data = buf.into_vec(count);
483 let mut result = Self::from_data(result_data, broadcast_shape, self.device)?;
484
485 if self.requires_grad || other.requires_grad {
487 result.requires_grad = true;
488 result.operation = Operation::Add {
489 lhs: Arc::new(self.clone()),
490 rhs: Arc::new(other.clone()),
491 };
492 }
493
494 Ok(result)
495 }
496
497 pub fn sub(&self, other: &Self) -> Result<Self>
499 where
500 T: std::ops::Sub<Output = T>,
501 {
502 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
504 && self.shape() == other.shape()
505 {
506 let self_data = self.data()?;
507 if self_data.len() >= 1024 {
508 let other_data = other.data()?;
509 let a_f32: &[f32] = unsafe {
510 std::slice::from_raw_parts(self_data.as_ptr() as *const f32, self_data.len())
511 };
512 let b_f32: &[f32] = unsafe {
513 std::slice::from_raw_parts(other_data.as_ptr() as *const f32, other_data.len())
514 };
515 let n = self_data.len();
516 let mut buf = global_acquire_uninit::<f32>(n);
517 {
518 let uninit = buf.as_uninit_slice_mut();
519 for slot in uninit.iter_mut() {
520 slot.write(0.0);
521 }
522 }
523 let mut out = buf.into_vec(n);
524 crate::simd_ops_f32::sub_into_f32(a_f32, b_f32, &mut out);
525 let result_data: Vec<T> = unsafe {
526 let mut v = std::mem::ManuallyDrop::new(out);
527 Vec::from_raw_parts(v.as_mut_ptr() as *mut T, v.len(), v.capacity())
528 };
529 let mut result =
530 Self::from_data(result_data, self.shape().dims().to_vec(), self.device)?;
531 if self.requires_grad || other.requires_grad {
532 result.requires_grad = true;
533 result.operation = crate::Operation::Sub {
534 lhs: Arc::new(self.clone()),
535 rhs: Arc::new(other.clone()),
536 };
537 }
538 return Ok(result);
539 }
540 }
541
542 let mut result = self.elementwise_operation(other, |a, b| a - b)?;
543
544 if self.requires_grad || other.requires_grad {
546 result.requires_grad = true;
547 result.operation = crate::Operation::Sub {
548 lhs: Arc::new(self.clone()),
549 rhs: Arc::new(other.clone()),
550 };
551 }
552
553 Ok(result)
554 }
555
556 pub fn mul(&self, other: &Self) -> Result<Self>
558 where
559 T: std::ops::Mul<Output = T>,
560 {
561 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
563 && self.shape() == other.shape()
564 {
565 let self_data = self.data()?;
566 if self_data.len() >= 1024 {
567 let other_data = other.data()?;
568 let a_f32: &[f32] = unsafe {
569 std::slice::from_raw_parts(self_data.as_ptr() as *const f32, self_data.len())
570 };
571 let b_f32: &[f32] = unsafe {
572 std::slice::from_raw_parts(other_data.as_ptr() as *const f32, other_data.len())
573 };
574 let n = self_data.len();
575 let mut buf = global_acquire_uninit::<f32>(n);
576 {
577 let uninit = buf.as_uninit_slice_mut();
578 for slot in uninit.iter_mut() {
579 slot.write(0.0);
580 }
581 }
582 let mut out = buf.into_vec(n);
583 crate::simd_ops_f32::mul_into_f32(a_f32, b_f32, &mut out);
584 let result_data: Vec<T> = unsafe {
585 let mut v = std::mem::ManuallyDrop::new(out);
586 Vec::from_raw_parts(v.as_mut_ptr() as *mut T, v.len(), v.capacity())
587 };
588 return Self::from_data(result_data, self.shape().dims().to_vec(), self.device);
589 }
590 }
591 self.elementwise_operation(other, |a, b| a * b)
592 }
593
594 pub fn div(&self, other: &Self) -> Result<Self>
596 where
597 T: std::ops::Div<Output = T>,
598 {
599 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
601 && self.shape() == other.shape()
602 {
603 let self_data = self.data()?;
604 if self_data.len() >= 1024 {
605 let other_data = other.data()?;
606 let a_f32: &[f32] = unsafe {
607 std::slice::from_raw_parts(self_data.as_ptr() as *const f32, self_data.len())
608 };
609 let b_f32: &[f32] = unsafe {
610 std::slice::from_raw_parts(other_data.as_ptr() as *const f32, other_data.len())
611 };
612 let n = self_data.len();
613 let mut buf = global_acquire_uninit::<f32>(n);
614 {
615 let uninit = buf.as_uninit_slice_mut();
616 for slot in uninit.iter_mut() {
617 slot.write(0.0);
618 }
619 }
620 let mut out = buf.into_vec(n);
621 crate::simd_ops_f32::div_into_f32(a_f32, b_f32, &mut out);
622 let result_data: Vec<T> = unsafe {
623 let mut v = std::mem::ManuallyDrop::new(out);
624 Vec::from_raw_parts(v.as_mut_ptr() as *mut T, v.len(), v.capacity())
625 };
626 return Self::from_data(result_data, self.shape().dims().to_vec(), self.device);
627 }
628 }
629 self.elementwise_operation(other, |a, b| a / b)
630 }
631
632 fn broadcast_binary_op<F>(&self, other: &Self, op: F) -> Result<Self>
634 where
635 F: Fn(T, T) -> T + Send + Sync,
636 {
637 use crate::broadcast::BroadcastOps;
638
639 let self_shape_binding = self.shape();
640 let self_shape = self_shape_binding.dims();
641 let other_shape_binding = other.shape();
642 let other_shape = other_shape_binding.dims();
643
644 let broadcast_shape = BroadcastOps::compute_broadcast_shape(self_shape, other_shape)?;
646
647 let self_data = self.data()?;
648 let other_data = other.data()?;
649
650 let total_elements = broadcast_shape.iter().product::<usize>();
651 let mut buf = global_acquire_uninit::<T>(total_elements);
652 let uninit = buf.as_uninit_slice_mut();
653 let mut count = 0;
654
655 let mut indices = vec![0; broadcast_shape.len()];
657 for _ in 0..total_elements {
658 let self_idx = self.compute_broadcast_index(&indices, self_shape, &broadcast_shape)?;
660 let other_idx =
661 other.compute_broadcast_index(&indices, other_shape, &broadcast_shape)?;
662
663 let result = op(self_data[self_idx], other_data[other_idx]);
664 uninit[count].write(result);
665 count += 1;
666
667 Self::increment_indices(&mut indices, &broadcast_shape);
669 }
670
671 let result_data = buf.into_vec(count);
672 Self::from_data(result_data, broadcast_shape, self.device)
673 }
674
675 fn increment_indices(indices: &mut [usize], shape: &[usize]) {
677 for i in (0..indices.len()).rev() {
678 indices[i] += 1;
679 if indices[i] < shape[i] {
680 break;
681 }
682 indices[i] = 0;
683 }
684 }
685
686 fn compute_broadcast_index(
688 &self,
689 broadcast_indices: &[usize],
690 original_shape: &[usize],
691 broadcast_shape: &[usize],
692 ) -> Result<usize> {
693 let ndim_diff = broadcast_shape.len() - original_shape.len();
694 let mut flat_index = 0;
695 let mut stride = 1;
696
697 for i in (0..original_shape.len()).rev() {
698 let broadcast_idx = broadcast_indices[ndim_diff + i];
699 let original_size = original_shape[i];
700
701 let actual_idx = if original_size == 1 { 0 } else { broadcast_idx };
703
704 flat_index += actual_idx * stride;
705 stride *= original_size;
706 }
707
708 Ok(flat_index)
709 }
710
711 fn elementwise_operation<F>(&self, other: &Self, op: F) -> Result<Self>
713 where
714 F: Fn(T, T) -> T + Send + Sync,
715 {
716 if self.shape() != other.shape() {
718 return self.broadcast_binary_op(other, op);
719 }
720
721 let self_data = self.data()?;
722 let other_data = other.data()?;
723
724 #[cfg(feature = "simd")]
726 {
727 if self_data.len() > 1000 {
728 let result_data = self.simd_elementwise_operation(&self_data, &other_data, op)?;
730 return Self::from_data(result_data, self.shape().dims().to_vec(), self.device);
731 }
732 }
733
734 #[cfg(feature = "parallel")]
736 {
737 if self_data.len() > 100 {
738 let paired_data: Vec<(T, T)> = self_data
740 .iter()
741 .zip(other_data.iter())
742 .map(|(&a, &b)| (a, b))
743 .collect();
744 let result_data = intelligent_parallel_process(
745 paired_data,
746 TensorOpType::ElementWise, self.device.clone(),
748 |(a, b)| op(a, b),
749 );
750 return Self::from_data(result_data, self.shape().dims().to_vec(), self.device);
751 }
752 }
753
754 let result_data: Vec<T> = self_data
756 .iter()
757 .zip(other_data.iter())
758 .map(|(&a, &b)| op(a, b))
759 .collect();
760
761 Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
762 }
763
764 #[cfg(feature = "simd")]
766 #[allow(dead_code)]
767 fn simd_elementwise_operation<F>(&self, data_a: &[T], data_b: &[T], op: F) -> Result<Vec<T>>
768 where
769 F: Fn(T, T) -> T + Send + Sync,
770 T: TensorElement,
771 {
772 #[cfg(feature = "simd")]
774 {
775 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() {
777 let _a_f32 = unsafe { std::mem::transmute::<&[T], &[f32]>(data_a) };
778 let _b_f32 = unsafe { std::mem::transmute::<&[T], &[f32]>(data_b) };
779
780 #[cfg(feature = "parallel")]
784 {
785 use scirs2_core::parallel_ops::*;
786 return Ok(data_a
787 .par_iter()
788 .zip(data_b.par_iter())
789 .map(|(&a, &b)| op(a, b))
790 .collect());
791 }
792 #[cfg(not(feature = "parallel"))]
793 {
794 return Ok(data_a
795 .iter()
796 .zip(data_b.iter())
797 .map(|(&a, &b)| op(a, b))
798 .collect());
799 }
800 }
801 }
802
803 Ok(data_a
805 .iter()
806 .zip(data_b.iter())
807 .map(|(&a, &b)| op(a, b))
808 .collect())
809 }
810}
811
812impl<T: TensorElement + Copy> Tensor<T> {
817 pub fn add_(&mut self, other: &Self) -> Result<&mut Self>
831 where
832 T: std::ops::Add<Output = T>,
833 {
834 self.inplace_binary_op(
835 other,
836 "add_",
837 crate::simd_ops_f32::add_assign_f32,
838 |a, b| a + b,
839 )
840 }
841
842 pub fn sub_(&mut self, other: &Self) -> Result<&mut Self>
847 where
848 T: std::ops::Sub<Output = T>,
849 {
850 self.inplace_binary_op(
851 other,
852 "sub_",
853 crate::simd_ops_f32::sub_assign_f32,
854 |a, b| a - b,
855 )
856 }
857
858 pub fn mul_(&mut self, other: &Self) -> Result<&mut Self>
863 where
864 T: std::ops::Mul<Output = T>,
865 {
866 self.inplace_binary_op(
867 other,
868 "mul_",
869 crate::simd_ops_f32::mul_assign_f32,
870 |a, b| a * b,
871 )
872 }
873
874 pub fn div_(&mut self, other: &Self) -> Result<&mut Self>
879 where
880 T: std::ops::Div<Output = T>,
881 {
882 self.inplace_binary_op(
883 other,
884 "div_",
885 crate::simd_ops_f32::div_assign_f32,
886 |a, b| a / b,
887 )
888 }
889
890 fn inplace_binary_op<S, F>(
894 &mut self,
895 other: &Self,
896 op_name: &str,
897 simd_f32: S,
898 scalar_op: F,
899 ) -> Result<&mut Self>
900 where
901 S: Fn(&mut [f32], &[f32]),
902 F: Fn(T, T) -> T,
903 {
904 if self.requires_grad {
905 return Err(TorshError::InvalidArgument(format!(
906 "In-place operation `{op_name}` on tensor that requires grad is not allowed"
907 )));
908 }
909
910 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>()
912 && self.shape() == other.shape()
913 && self.numel() >= 1024
914 {
915 self.make_unique()?;
916 let other_data = other.data()?;
917 self.storage.with_slice_mut(|out_t: &mut [T]| {
918 let out_f32: &mut [f32] = unsafe {
921 std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
922 };
923 let rhs_f32: &[f32] = unsafe {
924 std::slice::from_raw_parts(other_data.as_ptr() as *const f32, other_data.len())
925 };
926 simd_f32(out_f32, rhs_f32);
927 Ok(())
928 })?;
929 return Ok(self);
930 }
931
932 let self_dims_vec = self.shape().dims().to_vec();
934 let other_dims_vec = other.shape().dims().to_vec();
935 let other_data = other.data()?;
936
937 if self_dims_vec == other_dims_vec {
938 self.make_unique()?;
940 self.storage.with_slice_mut(|slice: &mut [T]| {
941 debug_assert_eq!(slice.len(), other_data.len());
942 for (dst, src) in slice.iter_mut().zip(other_data.iter()) {
943 *dst = scalar_op(*dst, *src);
944 }
945 Ok(())
946 })?;
947 return Ok(self);
948 }
949
950 use crate::broadcast::BroadcastOps;
953 let broadcast_shape =
954 BroadcastOps::compute_broadcast_shape(&self_dims_vec, &other_dims_vec)?;
955 if broadcast_shape != self_dims_vec {
956 return Err(TorshError::ShapeMismatch {
957 expected: self_dims_vec,
958 got: broadcast_shape,
959 });
960 }
961
962 self.make_unique()?;
963 let total = self_dims_vec.iter().product::<usize>();
964 let self_dims = self_dims_vec.as_slice();
965 let other_dims = other_dims_vec.as_slice();
966 self.storage.with_slice_mut(|slice: &mut [T]| {
967 for flat_idx in 0..total {
968 let multi_index = BroadcastOps::flat_to_multi_index(flat_idx, self_dims);
969 let other_idx =
970 BroadcastOps::compute_broadcast_index(&multi_index, other_dims, self_dims)?;
971 slice[flat_idx] = scalar_op(slice[flat_idx], other_data[other_idx]);
972 }
973 Ok(())
974 })?;
975
976 Ok(self)
977 }
978}
979
980impl<T: TensorElement + Copy> Tensor<T> {
982 pub fn add_op(&self, other: &Self) -> Result<Self>
984 where
985 T: std::ops::Add<Output = T>,
986 {
987 self.add(other)
988 }
989
990 pub fn mul_op(&self, other: &Self) -> Result<Self>
992 where
993 T: std::ops::Mul<Output = T>,
994 {
995 self.mul(other)
996 }
997
998 pub fn sigmoid(&self) -> Result<Self>
1000 where
1001 T: torsh_core::dtype::FloatElement,
1002 {
1003 #[cfg(feature = "gpu")]
1006 if let Some(result) =
1007 crate::gpu_dispatch::try_unary_f32(self, crate::gpu_dispatch::UnaryOp::Sigmoid)
1008 {
1009 return Ok(result);
1010 }
1011
1012 #[cfg(feature = "simd")]
1014 {
1015 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() > 1000 {
1016 return self.simd_sigmoid_f32();
1017 }
1018 }
1019
1020 #[cfg(feature = "parallel")]
1022 {
1023 if self.numel() > 100 {
1024 let one = <T as scirs2_core::numeric::One>::one();
1025 return self.parallel_map(|x| {
1026 one / (one + (-x).exp())
1028 });
1029 }
1030 }
1031
1032 let one = <T as scirs2_core::numeric::One>::one();
1034 let neg_self = self.neg()?;
1035 let exp_neg = neg_self.exp()?;
1036 let one_plus_exp = exp_neg.add_scalar(one)?;
1037 let ones = Self::ones(self.shape().dims(), self.device)?;
1038 ones.div(&one_plus_exp)
1039 }
1040
1041 #[cfg(feature = "simd")]
1043 fn simd_sigmoid_f32(&self) -> Result<Self> {
1044 use scirs2_core::ndarray::ArrayView1;
1045
1046 let data = self.data()?;
1047
1048 let data_f32: &[f32] =
1050 unsafe { std::slice::from_raw_parts(data.as_ptr() as *const f32, data.len()) };
1051
1052 let data_view = ArrayView1::from(data_f32);
1054
1055 let result_array = adaptive_simd::adaptive_simd_sigmoid_f32(&data_view);
1057
1058 let result_vec: Vec<T> = result_array
1060 .to_vec()
1061 .into_iter()
1062 .map(|f| unsafe { std::mem::transmute_copy::<f32, T>(&f) })
1063 .collect();
1064
1065 Self::from_data(
1066 result_vec,
1067 self.shape().dims().to_vec(),
1068 self.device.clone(),
1069 )
1070 }
1071
1072 pub fn relu(&self) -> Result<Self>
1074 where
1075 T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
1076 {
1077 #[cfg(feature = "gpu")]
1080 if let Some(result) =
1081 crate::gpu_dispatch::try_unary_f32(self, crate::gpu_dispatch::UnaryOp::Relu)
1082 {
1083 return Ok(result);
1084 }
1085
1086 let zero = <T as scirs2_core::numeric::Zero>::zero();
1087
1088 #[cfg(feature = "simd")]
1090 {
1091 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() > 1000 {
1092 return self.simd_relu_f32();
1093 }
1094 }
1095
1096 #[cfg(feature = "parallel")]
1098 {
1099 if self.numel() > 100 {
1100 return self.parallel_map(|x| if x > zero { x } else { zero });
1101 }
1102 }
1103
1104 self.map(|x| if x > zero { x } else { zero })
1106 }
1107
1108 #[cfg(feature = "simd")]
1110 fn simd_relu_f32(&self) -> Result<Self> {
1111 use scirs2_core::ndarray::ArrayView1;
1112
1113 let data = self.data()?;
1114
1115 let data_f32: &[f32] =
1117 unsafe { std::slice::from_raw_parts(data.as_ptr() as *const f32, data.len()) };
1118
1119 let data_view = ArrayView1::from(data_f32);
1121
1122 let result_array = adaptive_simd::adaptive_simd_relu_f32(&data_view);
1124
1125 let result_vec: Vec<T> = result_array
1127 .to_vec()
1128 .into_iter()
1129 .map(|f| unsafe { std::mem::transmute_copy::<f32, T>(&f) })
1130 .collect();
1131
1132 Self::from_data(
1133 result_vec,
1134 self.shape().dims().to_vec(),
1135 self.device.clone(),
1136 )
1137 }
1138
1139 #[cfg(feature = "parallel")]
1141 fn parallel_map<F>(&self, op: F) -> Result<Self>
1142 where
1143 F: Fn(T) -> T + Send + Sync,
1144 {
1145 let data = self.data()?;
1146
1147 let result_data = intelligent_parallel_process(
1149 data.iter().copied().collect::<Vec<_>>(),
1150 TensorOpType::Activation, self.device.clone(),
1152 op,
1153 );
1154
1155 Self::from_data(result_data, self.shape().dims().to_vec(), self.device)
1156 }
1157
1158 pub fn minimum(&self, other: &Self) -> Result<Self>
1160 where
1161 T: std::cmp::PartialOrd,
1162 {
1163 self.elementwise_operation(other, |a, b| if a < b { a } else { b })
1164 }
1165
1166 pub fn maximum(&self, other: &Self) -> Result<Self>
1168 where
1169 T: std::cmp::PartialOrd,
1170 {
1171 self.elementwise_operation(other, |a, b| if a > b { a } else { b })
1172 }
1173
1174 pub fn clamp(&self, min: T, max: T) -> Result<Self>
1176 where
1177 T: std::cmp::PartialOrd + Copy,
1178 {
1179 let data = self.to_vec()?;
1180 let clamped_data: Vec<T> = data
1181 .iter()
1182 .map(|&x| {
1183 if x < min {
1184 min
1185 } else if x > max {
1186 max
1187 } else {
1188 x
1189 }
1190 })
1191 .collect();
1192
1193 Self::from_data(
1194 clamped_data,
1195 self.shape().dims().to_vec(),
1196 self.device.clone(),
1197 )
1198 }
1199
1200 pub fn dot(&self, other: &Self) -> Result<Self>
1204 where
1205 T: std::ops::Mul<Output = T> + std::ops::Add<Output = T> + scirs2_core::numeric::Zero,
1206 {
1207 let elementwise = self.mul(other)?;
1209 elementwise.sum()
1210 }
1211}
1212
1213impl<T: TensorElement + Copy + scirs2_core::numeric::FromPrimitive> Tensor<T> {
1215 pub fn add_scirs2(&self, other: &Self) -> Result<Self>
1217 where
1218 T: std::ops::Add<Output = T> + scirs2_core::numeric::Float,
1219 {
1220 self.add(other)
1223 }
1224
1225 pub fn mul_scirs2(&self, other: &Self) -> Result<Self>
1227 where
1228 T: std::ops::Mul<Output = T> + scirs2_core::numeric::Float,
1229 {
1230 self.mul(other)
1233 }
1234
1235 pub fn sub_scirs2(&self, other: &Self) -> Result<Self>
1237 where
1238 T: std::ops::Sub<Output = T> + scirs2_core::numeric::Float,
1239 {
1240 self.sub(other)
1243 }
1244
1245 pub fn div_scirs2(&self, other: &Self) -> Result<Self>
1247 where
1248 T: std::ops::Div<Output = T> + scirs2_core::numeric::Float,
1249 {
1250 self.div(other)
1253 }
1254}
1255
1256impl<T: TensorElement + Copy> std::ops::Add for &Tensor<T>
1258where
1259 T: std::ops::Add<Output = T>,
1260{
1261 type Output = Tensor<T>;
1262
1263 fn add(self, rhs: Self) -> Self::Output {
1264 self.add(rhs).expect("tensor addition should succeed")
1265 }
1266}
1267
1268impl<T: TensorElement + Copy> std::ops::Sub for &Tensor<T>
1269where
1270 T: std::ops::Sub<Output = T>,
1271{
1272 type Output = Tensor<T>;
1273
1274 fn sub(self, rhs: Self) -> Self::Output {
1275 self.sub(rhs).expect("tensor subtraction should succeed")
1276 }
1277}
1278
1279impl<T: TensorElement + Copy> std::ops::Mul for &Tensor<T>
1280where
1281 T: std::ops::Mul<Output = T>,
1282{
1283 type Output = Tensor<T>;
1284
1285 fn mul(self, rhs: Self) -> Self::Output {
1286 self.mul(rhs).expect("tensor multiplication should succeed")
1287 }
1288}
1289
1290impl<T: TensorElement + Copy> std::ops::Div for &Tensor<T>
1291where
1292 T: std::ops::Div<Output = T>,
1293{
1294 type Output = Tensor<T>;
1295
1296 fn div(self, rhs: Self) -> Self::Output {
1297 self.div(rhs).expect("tensor division should succeed")
1298 }
1299}
1300
1301impl<T: TensorElement + Copy> std::ops::Neg for &Tensor<T>
1303where
1304 T: std::ops::Neg<Output = T>,
1305{
1306 type Output = Tensor<T>;
1307
1308 fn neg(self) -> Self::Output {
1309 self.map(|x| -x).expect("negation map should succeed")
1310 }
1311}
1312
1313impl<T: TensorElement + Copy + scirs2_core::numeric::Float> Tensor<T> {
1317 #[cfg(feature = "simd")]
1321 pub fn add_simd(&self, other: &Self) -> Result<Self>
1322 where
1323 T: scirs2_core::simd_ops::SimdUnifiedOps,
1324 {
1325 use scirs2_core::ndarray::Array1;
1326
1327 if self.shape().dims() != other.shape().dims() {
1329 return Err(TorshError::ShapeMismatch {
1330 expected: self.shape().dims().to_vec(),
1331 got: other.shape().dims().to_vec(),
1332 });
1333 }
1334
1335 let data_a = self.to_vec()?;
1337 let data_b = other.to_vec()?;
1338
1339 let arr_a = Array1::from_vec(data_a);
1341 let arr_b = Array1::from_vec(data_b);
1342
1343 let result_arr = T::simd_add(&arr_a.view(), &arr_b.view());
1345
1346 Tensor::from_vec(result_arr.to_vec(), self.shape().dims())
1348 }
1349
1350 #[cfg(feature = "simd")]
1354 pub fn mul_simd(&self, other: &Self) -> Result<Self>
1355 where
1356 T: scirs2_core::simd_ops::SimdUnifiedOps,
1357 {
1358 use scirs2_core::ndarray::Array1;
1359
1360 if self.shape().dims() != other.shape().dims() {
1362 return Err(TorshError::ShapeMismatch {
1363 expected: self.shape().dims().to_vec(),
1364 got: other.shape().dims().to_vec(),
1365 });
1366 }
1367
1368 let data_a = self.to_vec()?;
1370 let data_b = other.to_vec()?;
1371
1372 let arr_a = Array1::from_vec(data_a);
1374 let arr_b = Array1::from_vec(data_b);
1375
1376 let result_arr = T::simd_mul(&arr_a.view(), &arr_b.view());
1378
1379 Tensor::from_vec(result_arr.to_vec(), self.shape().dims())
1381 }
1382
1383 #[cfg(feature = "simd")]
1389 pub fn dot_simd(&self, other: &Self) -> Result<T>
1390 where
1391 T: scirs2_core::simd_ops::SimdUnifiedOps,
1392 {
1393 use scirs2_core::ndarray::Array1;
1394
1395 if self.shape().dims() != other.shape().dims() {
1397 return Err(TorshError::ShapeMismatch {
1398 expected: self.shape().dims().to_vec(),
1399 got: other.shape().dims().to_vec(),
1400 });
1401 }
1402
1403 let data_a = self.to_vec()?;
1405 let data_b = other.to_vec()?;
1406
1407 let arr_a = Array1::from_vec(data_a);
1409 let arr_b = Array1::from_vec(data_b);
1410
1411 Ok(T::simd_dot(&arr_a.view(), &arr_b.view()))
1413 }
1414
1415 pub fn reduce_memory_efficient<F>(&self, func: F) -> Result<T>
1417 where
1418 F: Fn(T, T) -> T + Send + Sync,
1419 {
1420 #[cfg(feature = "profiling")]
1421 {
1422 }
1424
1425 let data = self.to_vec()?;
1428 Ok(data
1429 .into_iter()
1430 .reduce(func)
1431 .unwrap_or_else(|| <T as scirs2_core::numeric::Zero>::zero()))
1432 }
1433}
1434
1435impl<T: TensorElement + Copy + std::ops::Mul<Output = T>> Tensor<T> {
1437 pub fn relu_(&mut self) -> Result<&mut Self>
1445 where
1446 T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
1447 {
1448 if self.requires_grad {
1449 return Err(TorshError::InvalidArgument(
1450 "In-place operation on tensor that requires grad is not allowed".to_string(),
1451 ));
1452 }
1453
1454 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() >= 1024 {
1456 self.make_unique()?;
1457 self.storage.with_slice_mut(|out_t: &mut [T]| {
1458 let out_f32: &mut [f32] = unsafe {
1460 std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
1461 };
1462 crate::simd_ops_f32::relu_assign_f32(out_f32);
1463 Ok(())
1464 })?;
1465 return Ok(self);
1466 }
1467
1468 let zero = <T as scirs2_core::numeric::Zero>::zero();
1469 let len = self.storage.len();
1470
1471 for i in 0..len {
1472 let current = self.storage.get(i)?;
1473 if current < zero {
1474 self.storage.set(i, zero)?;
1475 }
1476 }
1477
1478 Ok(self)
1479 }
1480
1481 pub fn sigmoid_(&mut self) -> Result<&mut Self>
1486 where
1487 T: torsh_core::dtype::FloatElement,
1488 {
1489 if self.requires_grad {
1490 return Err(TorshError::InvalidArgument(
1491 "In-place operation on tensor that requires grad is not allowed".to_string(),
1492 ));
1493 }
1494
1495 let one = <T as scirs2_core::numeric::One>::one();
1496 let len = self.storage.len();
1497
1498 for i in 0..len {
1499 let x = self.storage.get(i)?;
1500 let sigmoid_val = one / (one + (-x).exp());
1501 self.storage.set(i, sigmoid_val)?;
1502 }
1503
1504 Ok(self)
1505 }
1506
1507 pub fn tanh_(&mut self) -> Result<&mut Self>
1512 where
1513 T: torsh_core::dtype::FloatElement,
1514 {
1515 if self.requires_grad {
1516 return Err(TorshError::InvalidArgument(
1517 "In-place operation on tensor that requires grad is not allowed".to_string(),
1518 ));
1519 }
1520
1521 let len = self.storage.len();
1522
1523 for i in 0..len {
1524 let x = self.storage.get(i)?;
1525 self.storage.set(i, x.tanh())?;
1526 }
1527
1528 Ok(self)
1529 }
1530
1531 pub fn gelu_(&mut self) -> Result<&mut Self>
1536 where
1537 T: torsh_core::dtype::FloatElement,
1538 {
1539 if self.requires_grad {
1540 return Err(TorshError::InvalidArgument(
1541 "In-place operation on tensor that requires grad is not allowed".to_string(),
1542 ));
1543 }
1544
1545 let len = self.storage.len();
1546 let pi = T::from(std::f64::consts::PI).expect("numeric conversion should succeed");
1547 let two = T::from(2.0).expect("numeric conversion should succeed");
1548 let sqrt_2_over_pi = (two / pi).sqrt();
1549 let point_044715 = T::from(0.044715).expect("numeric conversion should succeed");
1550 let one = <T as scirs2_core::numeric::One>::one();
1551 let half = T::from(0.5).expect("numeric conversion should succeed");
1552
1553 for i in 0..len {
1554 let x = self.storage.get(i)?;
1555 let x_cubed = x * x * x;
1556 let tanh_input = sqrt_2_over_pi * (x + point_044715 * x_cubed);
1557 let gelu_val = half * x * (one + tanh_input.tanh());
1558 self.storage.set(i, gelu_val)?;
1559 }
1560
1561 Ok(self)
1562 }
1563
1564 pub fn leaky_relu_(&mut self, negative_slope: T) -> Result<&mut Self>
1569 where
1570 T: std::cmp::PartialOrd + scirs2_core::numeric::Zero,
1571 {
1572 if self.requires_grad {
1573 return Err(TorshError::InvalidArgument(
1574 "In-place operation on tensor that requires grad is not allowed".to_string(),
1575 ));
1576 }
1577
1578 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() >= 1024 {
1580 let slope_f32: f32 = unsafe { std::mem::transmute_copy::<T, f32>(&negative_slope) };
1582 self.make_unique()?;
1583 self.storage.with_slice_mut(|out_t: &mut [T]| {
1584 let out_f32: &mut [f32] = unsafe {
1585 std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
1586 };
1587 crate::simd_ops_f32::leaky_relu_assign_f32(out_f32, slope_f32);
1588 Ok(())
1589 })?;
1590 return Ok(self);
1591 }
1592
1593 let zero = <T as scirs2_core::numeric::Zero>::zero();
1594 let len = self.storage.len();
1595
1596 for i in 0..len {
1597 let x = self.storage.get(i)?;
1598 if x < zero {
1599 self.storage.set(i, negative_slope * x)?;
1600 }
1601 }
1602
1603 Ok(self)
1604 }
1605
1606 pub fn clamp_(&mut self, min: T, max: T) -> Result<&mut Self>
1611 where
1612 T: std::cmp::PartialOrd,
1613 {
1614 if self.requires_grad {
1615 return Err(TorshError::InvalidArgument(
1616 "In-place operation on tensor that requires grad is not allowed".to_string(),
1617 ));
1618 }
1619
1620 if std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() && self.numel() >= 1024 {
1622 let min_f32: f32 = unsafe { std::mem::transmute_copy::<T, f32>(&min) };
1624 let max_f32: f32 = unsafe { std::mem::transmute_copy::<T, f32>(&max) };
1625 self.make_unique()?;
1626 self.storage.with_slice_mut(|out_t: &mut [T]| {
1627 let out_f32: &mut [f32] = unsafe {
1628 std::slice::from_raw_parts_mut(out_t.as_mut_ptr() as *mut f32, out_t.len())
1629 };
1630 crate::simd_ops_f32::clamp_assign_f32(out_f32, min_f32, max_f32);
1631 Ok(())
1632 })?;
1633 return Ok(self);
1634 }
1635
1636 let len = self.storage.len();
1637
1638 for i in 0..len {
1639 let x = self.storage.get(i)?;
1640 let clamped = if x < min {
1641 min
1642 } else if x > max {
1643 max
1644 } else {
1645 x
1646 };
1647 self.storage.set(i, clamped)?;
1648 }
1649
1650 Ok(self)
1651 }
1652}
1653
1654#[path = "math_ops_trig.rs"]
1655mod math_ops_trig;
1656
1657#[cfg(test)]
1658#[path = "math_ops_tests.rs"]
1659mod tests;