rstsr_native_impl/cpu_serial/
assignment.rs

1use crate::prelude_dev::*;
2
3// this value is used to determine whether to use contiguous inner iteration
4const CONTIG_SWITCH: usize = 16;
5
6pub fn assign_arbitary_cpu_serial<T, DC, DA>(
7    c: &mut [T],
8    lc: &Layout<DC>,
9    a: &[T],
10    la: &Layout<DA>,
11    order: FlagOrder,
12) -> Result<()>
13where
14    T: Clone,
15    DC: DimAPI,
16    DA: DimAPI,
17{
18    let contig = match order {
19        RowMajor => lc.c_contig() && la.c_contig(),
20        ColMajor => lc.f_contig() && la.f_contig(),
21    };
22    if contig {
23        // contiguous case
24        let offset_c = lc.offset();
25        let offset_a = la.offset();
26        let size = lc.size();
27        c[offset_c..(offset_c + size)]
28            .iter_mut()
29            .zip(a[offset_a..(offset_a + size)].iter())
30            .for_each(|(ci, ai)| *ci = ai.clone());
31    } else {
32        // determine order by layout preference
33        let order = match order {
34            RowMajor => TensorIterOrder::C,
35            ColMajor => TensorIterOrder::F,
36        };
37        // generate col-major iterator
38        let lc = translate_to_col_major_unary(lc, order)?;
39        let la = translate_to_col_major_unary(la, order)?;
40        layout_col_major_dim_dispatch_2diff(&lc, &la, |(idx_c, idx_a)| c[idx_c] = a[idx_a].clone())?;
41    }
42    Ok(())
43}
44
45pub fn assign_cpu_serial<T, D>(c: &mut [T], lc: &Layout<D>, a: &[T], la: &Layout<D>) -> Result<()>
46where
47    T: Clone,
48    D: DimAPI,
49{
50    let layouts_full = translate_to_col_major(&[lc, la], TensorIterOrder::K)?;
51    let layouts_full_ref = layouts_full.iter().collect_vec();
52    let (layouts_contig, size_contig) = translate_to_col_major_with_contig(&layouts_full_ref);
53
54    if size_contig >= CONTIG_SWITCH {
55        let lc = &layouts_contig[0];
56        let la = &layouts_contig[1];
57        layout_col_major_dim_dispatch_2(lc, la, |(idx_c, idx_a)| {
58            c[idx_c..(idx_c + size_contig)]
59                .iter_mut()
60                .zip(a[idx_a..(idx_a + size_contig)].iter())
61                .for_each(|(ci, ai)| *ci = ai.clone());
62        })?;
63    } else {
64        let lc = &layouts_full[0];
65        let la = &layouts_full[1];
66        layout_col_major_dim_dispatch_2(lc, la, |(idx_c, idx_a)| c[idx_c] = a[idx_a].clone())?;
67    }
68    Ok(())
69}
70
71pub fn fill_cpu_serial<T, D>(c: &mut [T], lc: &Layout<D>, fill: T) -> Result<()>
72where
73    T: Clone,
74    D: DimAPI,
75{
76    let layouts_full = [translate_to_col_major_unary(lc, TensorIterOrder::G)?];
77    let layouts_full_ref = layouts_full.iter().collect_vec();
78    let (layouts_contig, size_contig) = translate_to_col_major_with_contig(&layouts_full_ref);
79
80    if size_contig > CONTIG_SWITCH {
81        layout_col_major_dim_dispatch_1(&layouts_contig[0], |idx_c| {
82            for i in 0..size_contig {
83                c[idx_c + i] = fill.clone();
84            }
85        })?;
86    } else {
87        layout_col_major_dim_dispatch_1(&layouts_full[0], |idx_c| {
88            c[idx_c] = fill.clone();
89        })?;
90    }
91    Ok(())
92}
rstsr_native_impl/cpu_serial/assignment.rs

rstsr_native_impl/cpu_serial/
assignment.rs
