ruprim_host/gather_scatter/
dispatch_float.rs

1use ruda_core::tensor::{DType, host::HostTensor};
2use half::{bf16, f16};
3
4pub fn float_gather(
5    dim: usize,
6    tensor: HostTensor,
7    indices: HostTensor,
8) -> HostTensor {
9    match tensor.dtype() {
10        DType::F32 => crate::gather_scatter::gather::<f32>(tensor, dim, indices),
11        DType::F64 => crate::gather_scatter::gather::<f64>(tensor, dim, indices),
12        DType::F16 => crate::gather_scatter::gather::<f16>(tensor, dim, indices),
13        DType::BF16 => crate::gather_scatter::gather::<bf16>(tensor, dim, indices),
14        _ => panic!("float_gather: unsupported dtype {:?}", tensor.dtype()),
15    }
16}
17
18pub fn float_scatter_add(
19    dim: usize,
20    tensor: HostTensor,
21    indices: HostTensor,
22    value: HostTensor,
23) -> HostTensor {
24    match tensor.dtype() {
25        DType::F32 => {
26            crate::gather_scatter::scatter_add::<f32>(tensor, dim, indices, value)
27        }
28        DType::F64 => {
29            crate::gather_scatter::scatter_add::<f64>(tensor, dim, indices, value)
30        }
31        DType::F16 => {
32            crate::gather_scatter::scatter_add::<f16>(tensor, dim, indices, value)
33        }
34        DType::BF16 => {
35            crate::gather_scatter::scatter_add::<bf16>(tensor, dim, indices, value)
36        }
37        _ => panic!("float_scatter_add: unsupported dtype {:?}", tensor.dtype()),
38    }
39}
40
41pub fn float_scatter_nd(
42    data: HostTensor,
43    indices: HostTensor,
44    values: HostTensor,
45    reduction: ruda_core::tensor::indexing::IndexingUpdateOp,
46) -> HostTensor {
47    match data.dtype() {
48        DType::F32 => {
49            crate::gather_scatter::scatter_nd::<f32>(data, indices, values, reduction)
50        }
51        DType::F64 => {
52            crate::gather_scatter::scatter_nd::<f64>(data, indices, values, reduction)
53        }
54        DType::F16 => {
55            crate::gather_scatter::scatter_nd::<f16>(data, indices, values, reduction)
56        }
57        DType::BF16 => {
58            crate::gather_scatter::scatter_nd::<bf16>(data, indices, values, reduction)
59        }
60        _ => panic!("float_scatter_nd: unsupported dtype {:?}", data.dtype()),
61    }
62}
63
64pub fn float_gather_nd(data: HostTensor, indices: HostTensor) -> HostTensor {
65    match data.dtype() {
66        DType::F32 => crate::gather_scatter::gather_nd::<f32>(data, indices),
67        DType::F64 => crate::gather_scatter::gather_nd::<f64>(data, indices),
68        DType::F16 => crate::gather_scatter::gather_nd::<f16>(data, indices),
69        DType::BF16 => crate::gather_scatter::gather_nd::<bf16>(data, indices),
70        _ => panic!("float_gather_nd: unsupported dtype {:?}", data.dtype()),
71    }
72}
73
74pub fn float_select(
75    tensor: HostTensor,
76    dim: usize,
77    indices: HostTensor,
78) -> HostTensor {
79    match tensor.dtype() {
80        DType::F32 => crate::gather_scatter::select::<f32>(tensor, dim, indices),
81        DType::F64 => crate::gather_scatter::select::<f64>(tensor, dim, indices),
82        DType::F16 => crate::gather_scatter::select::<f16>(tensor, dim, indices),
83        DType::BF16 => crate::gather_scatter::select::<bf16>(tensor, dim, indices),
84        _ => panic!("float_select: unsupported dtype {:?}", tensor.dtype()),
85    }
86}
87
88pub fn float_select_add(
89    tensor: HostTensor,
90    dim: usize,
91    indices: HostTensor,
92    value: HostTensor,
93) -> HostTensor {
94    match tensor.dtype() {
95        DType::F32 => {
96            crate::gather_scatter::select_add::<f32>(tensor, dim, indices, value)
97        }
98        DType::F64 => {
99            crate::gather_scatter::select_add::<f64>(tensor, dim, indices, value)
100        }
101        DType::F16 => {
102            crate::gather_scatter::select_add::<f16>(tensor, dim, indices, value)
103        }
104        DType::BF16 => {
105            crate::gather_scatter::select_add::<bf16>(tensor, dim, indices, value)
106        }
107        _ => panic!("float_select_add: unsupported dtype {:?}", tensor.dtype()),
108    }
109}
110
ruprim_host/gather_scatter/dispatch_float.rs

ruprim_host/gather_scatter/
dispatch_float.rs
