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
6#[duplicate_item(
7    func_name
8        TypeC TypeA Types
9        func_clone
10    ;
11    [assign_arbitary_cpu_serial]
12        [T] [T] [T: Clone]
13        [*ci = ai.clone()]
14    ;
15    [assign_arbitary_uninit_cpu_serial]
16        [MaybeUninit<T>] [T] [T: Clone]
17        [ci.write(ai.clone())]
18    ;
19    [assign_arbitary_promote_cpu_serial]
20        [TC] [TA] [TC: Clone, TA: Clone + DTypeCastAPI<TC>]
21        [*ci = ai.clone().into_cast()]
22    ;
23    [assign_arbitary_uninit_promote_cpu_serial]
24        [MaybeUninit<TC>] [TA] [TC: Clone, TA: Clone + DTypeCastAPI<TC>]
25        [ci.write(ai.clone().into_cast())]
26    ;
27)]
28pub fn func_name<Types, DC, DA>(
29    c: &mut [TypeC],
30    lc: &Layout<DC>,
31    a: &[TypeA],
32    la: &Layout<DA>,
33    order: FlagOrder,
34) -> Result<()>
35where
36    DC: DimAPI,
37    DA: DimAPI,
38{
39    let contig = match order {
40        RowMajor => lc.c_contig() && la.c_contig(),
41        ColMajor => lc.f_contig() && la.f_contig(),
42    };
43    if contig {
44        // contiguous case
45        let offset_c = lc.offset();
46        let offset_a = la.offset();
47        let size = lc.size();
48        c[offset_c..(offset_c + size)].iter_mut().zip(a[offset_a..(offset_a + size)].iter()).for_each(|(ci, ai)| {
49            func_clone;
50        });
51    } else {
52        // determine order by layout preference
53        let order = match order {
54            RowMajor => TensorIterOrder::C,
55            ColMajor => TensorIterOrder::F,
56        };
57        // generate col-major iterator
58        let lc = translate_to_col_major_unary(lc, order)?;
59        let la = translate_to_col_major_unary(la, order)?;
60        layout_col_major_dim_dispatch_2diff(&lc, &la, |(idx_c, idx_a)| {
61            let ci = &mut c[idx_c];
62            let ai = &a[idx_a];
63            func_clone;
64        })?;
65    }
66    Ok(())
67}
68
69#[duplicate_item(
70    func_name
71        TypeC TypeA Types
72        func_clone
73    ;
74    [assign_cpu_serial]
75        [T] [T] [T: Clone]
76        [*ci = ai.clone()]
77    ;
78    [assign_uninit_cpu_serial]
79        [MaybeUninit<T>] [T] [T: Clone]
80        [ci.write(ai.clone())]
81    ;
82    [assign_promote_cpu_serial]
83        [TC] [TA] [TC: Clone, TA: Clone + DTypeCastAPI<TC>]
84        [*ci = ai.clone().into_cast()]
85    ;
86    [assign_uninit_promote_cpu_serial]
87        [MaybeUninit<TC>] [TA] [TC: Clone, TA: Clone + DTypeCastAPI<TC>]
88        [ci.write(ai.clone().into_cast())]
89    ;
90)]
91pub fn func_name<Types, D>(c: &mut [TypeC], lc: &Layout<D>, a: &[TypeA], la: &Layout<D>) -> Result<()>
92where
93    D: DimAPI,
94{
95    let layouts_full = translate_to_col_major(&[lc, la], TensorIterOrder::K)?;
96    let layouts_full_ref = layouts_full.iter().collect_vec();
97    let (layouts_contig, size_contig) = translate_to_col_major_with_contig(&layouts_full_ref);
98
99    if size_contig >= CONTIG_SWITCH {
100        let lc = &layouts_contig[0];
101        let la = &layouts_contig[1];
102        layout_col_major_dim_dispatch_2(lc, la, |(idx_c, idx_a)| {
103            c[idx_c..(idx_c + size_contig)].iter_mut().zip(a[idx_a..(idx_a + size_contig)].iter()).for_each(
104                |(ci, ai)| {
105                    func_clone;
106                },
107            );
108        })?;
109    } else {
110        let lc = &layouts_full[0];
111        let la = &layouts_full[1];
112        layout_col_major_dim_dispatch_2(lc, la, |(idx_c, idx_a)| {
113            let ci = &mut c[idx_c];
114            let ai = &a[idx_a];
115            func_clone;
116        })?;
117    }
118    Ok(())
119}
120
121pub fn fill_cpu_serial<T, D>(c: &mut [T], lc: &Layout<D>, fill: T) -> Result<()>
122where
123    T: Clone,
124    D: DimAPI,
125{
126    fill_promote_cpu_serial(c, lc, fill)
127}
128
129pub fn fill_promote_cpu_serial<TC, TA, D>(c: &mut [TC], lc: &Layout<D>, fill: TA) -> Result<()>
130where
131    TA: Clone + DTypeCastAPI<TC>,
132    TC: Clone,
133    D: DimAPI,
134{
135    let fill = fill.clone().into_cast();
136
137    let layouts_full = [translate_to_col_major_unary(lc, TensorIterOrder::G)?];
138    let layouts_full_ref = layouts_full.iter().collect_vec();
139    let (layouts_contig, size_contig) = translate_to_col_major_with_contig(&layouts_full_ref);
140
141    if size_contig > CONTIG_SWITCH {
142        layout_col_major_dim_dispatch_1(&layouts_contig[0], |idx_c| {
143            for i in 0..size_contig {
144                c[idx_c + i] = fill.clone();
145            }
146        })?;
147    } else {
148        layout_col_major_dim_dispatch_1(&layouts_full[0], |idx_c| {
149            c[idx_c] = fill.clone();
150        })?;
151    }
152    Ok(())
153}
rstsr_native_impl/cpu_serial/assignment.rs

rstsr_native_impl/cpu_serial/
assignment.rs
