1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
use crate::prelude::*;
use ndarray::Dimension;
pub(crate) unsafe fn scatter_contig_data<T: Datum>(
mut src: *const T,
dst: *mut T,
dst_len_and_strides: &[(usize, usize)],
) {
match &*dst_len_and_strides {
&[(len_a, stride_a)] => {
for a in 0..len_a {
*dst.offset((a * stride_a) as isize) = (*src).clone();
src = src.offset(1);
}
}
&[(len_a, stride_a), (len_b, stride_b)] => {
for a in 0..len_a {
for b in 0..len_b {
*dst.offset((a * stride_a + b * stride_b) as isize) = (&*src).clone();
src = src.offset(1);
}
}
}
&[(len_a, stride_a), (len_b, stride_b), (len_c, stride_c)] => {
for a in 0..len_a {
for b in 0..len_b {
for c in 0..len_c {
*dst.offset((a * stride_a + b * stride_b + c * stride_c) as isize) =
(&*src).clone();
src = src.offset(1);
}
}
}
}
&[(len_a, stride_a), (len_b, stride_b), (len_c, stride_c), (len_d, stride_d)] => {
for a in 0..len_a {
for b in 0..len_b {
for c in 0..len_c {
for d in 0..len_d {
*dst.offset(
(a * stride_a + b * stride_b + c * stride_c + d * stride_d)
as isize,
) = (&*src).clone();
src = src.offset(1);
}
}
}
}
}
&[(len_a, stride_a), (len_b, stride_b), (len_c, stride_c), (len_d, stride_d), (len_e, stride_e)] => {
for a in 0..len_a {
for b in 0..len_b {
for c in 0..len_c {
for d in 0..len_d {
for e in 0..len_e {
*dst.offset(
(a * stride_a
+ b * stride_b
+ c * stride_c
+ d * stride_d
+ e * stride_e)
as isize,
) = (&*src).clone();
src = src.offset(1);
}
}
}
}
}
}
_ => {
let shape: TVec<usize> = dst_len_and_strides.iter().map(|pair| pair.0).collect();
for coords in ndarray::indices(&*shape) {
let offset = coords
.slice()
.iter()
.zip(dst_len_and_strides.iter())
.map(|(x, (_len, stride))| x * stride)
.sum::<usize>();
*dst.offset(offset as isize) = (&*src).clone();
src = src.offset(1);
}
}
}
}
