ruPRIM-host 0.1.21

ruPRIM CPU tensor primitives, reductions, indexing and SIMD kernels.
Documentation 
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

use ruda_core::tensor::{DType, host::HostTensor};

pub fn int_gather(
    dim: usize,
    tensor: HostTensor,
    indices: HostTensor,
) -> HostTensor {
    match tensor.dtype() {
        DType::I64 => crate::gather_scatter::gather::<i64>(tensor, dim, indices),
        DType::I32 => crate::gather_scatter::gather::<i32>(tensor, dim, indices),
        DType::I16 => crate::gather_scatter::gather::<i16>(tensor, dim, indices),
        DType::I8 => crate::gather_scatter::gather::<i8>(tensor, dim, indices),
        DType::U64 => crate::gather_scatter::gather::<u64>(tensor, dim, indices),
        DType::U32 => crate::gather_scatter::gather::<u32>(tensor, dim, indices),
        DType::U16 => crate::gather_scatter::gather::<u16>(tensor, dim, indices),
        DType::U8 => crate::gather_scatter::gather::<u8>(tensor, dim, indices),
        dt => panic!("int_gather: unsupported dtype {:?}", dt),
    }
}

pub fn int_scatter_add(
    dim: usize,
    tensor: HostTensor,
    indices: HostTensor,
    value: HostTensor,
) -> HostTensor {
    debug_assert_eq!(
        tensor.dtype(),
        value.dtype(),
        "int_scatter_add: dtype mismatch"
    );
    match tensor.dtype() {
        DType::I64 => {
            crate::gather_scatter::scatter_add::<i64>(tensor, dim, indices, value)
        }
        DType::I32 => {
            crate::gather_scatter::scatter_add::<i32>(tensor, dim, indices, value)
        }
        DType::I16 => {
            crate::gather_scatter::scatter_add::<i16>(tensor, dim, indices, value)
        }
        DType::I8 => crate::gather_scatter::scatter_add::<i8>(tensor, dim, indices, value),
        DType::U64 => {
            crate::gather_scatter::scatter_add::<u64>(tensor, dim, indices, value)
        }
        DType::U32 => {
            crate::gather_scatter::scatter_add::<u32>(tensor, dim, indices, value)
        }
        DType::U16 => {
            crate::gather_scatter::scatter_add::<u16>(tensor, dim, indices, value)
        }
        DType::U8 => crate::gather_scatter::scatter_add::<u8>(tensor, dim, indices, value),
        dt => panic!("int_scatter_add: unsupported dtype {:?}", dt),
    }
}

pub fn int_scatter_nd(
    data: HostTensor,
    indices: HostTensor,
    values: HostTensor,
    reduction: ruda_core::tensor::indexing::IndexingUpdateOp,
) -> HostTensor {
    match data.dtype() {
        DType::I64 => {
            crate::gather_scatter::scatter_nd::<i64>(data, indices, values, reduction)
        }
        DType::I32 => {
            crate::gather_scatter::scatter_nd::<i32>(data, indices, values, reduction)
        }
        DType::I16 => {
            crate::gather_scatter::scatter_nd::<i16>(data, indices, values, reduction)
        }
        DType::I8 => {
            crate::gather_scatter::scatter_nd::<i8>(data, indices, values, reduction)
        }
        DType::U64 => {
            crate::gather_scatter::scatter_nd::<u64>(data, indices, values, reduction)
        }
        DType::U32 => {
            crate::gather_scatter::scatter_nd::<u32>(data, indices, values, reduction)
        }
        DType::U16 => {
            crate::gather_scatter::scatter_nd::<u16>(data, indices, values, reduction)
        }
        DType::U8 => {
            crate::gather_scatter::scatter_nd::<u8>(data, indices, values, reduction)
        }
        dt => panic!("int_scatter_nd: unsupported dtype {:?}", dt),
    }
}

pub fn int_gather_nd(data: HostTensor, indices: HostTensor) -> HostTensor {
    match data.dtype() {
        DType::I64 => crate::gather_scatter::gather_nd::<i64>(data, indices),
        DType::I32 => crate::gather_scatter::gather_nd::<i32>(data, indices),
        DType::I16 => crate::gather_scatter::gather_nd::<i16>(data, indices),
        DType::I8 => crate::gather_scatter::gather_nd::<i8>(data, indices),
        DType::U64 => crate::gather_scatter::gather_nd::<u64>(data, indices),
        DType::U32 => crate::gather_scatter::gather_nd::<u32>(data, indices),
        DType::U16 => crate::gather_scatter::gather_nd::<u16>(data, indices),
        DType::U8 => crate::gather_scatter::gather_nd::<u8>(data, indices),
        dt => panic!("int_gather_nd: unsupported dtype {:?}", dt),
    }
}

pub fn int_select(
    tensor: HostTensor,
    dim: usize,
    indices: HostTensor,
) -> HostTensor {
    match tensor.dtype() {
        DType::I64 => crate::gather_scatter::select::<i64>(tensor, dim, indices),
        DType::I32 => crate::gather_scatter::select::<i32>(tensor, dim, indices),
        DType::I16 => crate::gather_scatter::select::<i16>(tensor, dim, indices),
        DType::I8 => crate::gather_scatter::select::<i8>(tensor, dim, indices),
        DType::U64 => crate::gather_scatter::select::<u64>(tensor, dim, indices),
        DType::U32 => crate::gather_scatter::select::<u32>(tensor, dim, indices),
        DType::U16 => crate::gather_scatter::select::<u16>(tensor, dim, indices),
        DType::U8 => crate::gather_scatter::select::<u8>(tensor, dim, indices),
        dt => panic!("int_select: unsupported dtype {:?}", dt),
    }
}

pub fn int_select_add(
    tensor: HostTensor,
    dim: usize,
    indices: HostTensor,
    value: HostTensor,
) -> HostTensor {
    debug_assert_eq!(
        tensor.dtype(),
        value.dtype(),
        "int_select_add: dtype mismatch"
    );
    match tensor.dtype() {
        DType::I64 => {
            crate::gather_scatter::select_add::<i64>(tensor, dim, indices, value)
        }
        DType::I32 => {
            crate::gather_scatter::select_add::<i32>(tensor, dim, indices, value)
        }
        DType::I16 => {
            crate::gather_scatter::select_add::<i16>(tensor, dim, indices, value)
        }
        DType::I8 => crate::gather_scatter::select_add::<i8>(tensor, dim, indices, value),
        DType::U64 => {
            crate::gather_scatter::select_add::<u64>(tensor, dim, indices, value)
        }
        DType::U32 => {
            crate::gather_scatter::select_add::<u32>(tensor, dim, indices, value)
        }
        DType::U16 => {
            crate::gather_scatter::select_add::<u16>(tensor, dim, indices, value)
        }
        DType::U8 => crate::gather_scatter::select_add::<u8>(tensor, dim, indices, value),
        dt => panic!("int_select_add: unsupported dtype {:?}", dt),
    }
}