luma_tensor/device/cpu/ops/
int_ops.rs

1//! `impl IntOps for Cpu`: dispatch `CpuIntStorage` variants to generic kernels.
2
3use std::borrow::Cow;
4
5use super::kernels::{elementwise as ew, indexing, matmul, reduce};
6use super::{Cpu, CpuBoolStorage, CpuFloatStorage, CpuIntStorage, int_ids_as_usize};
7use crate::dtype::{BoolDType, FloatDType, IntDType};
8use crate::{
9    BinaryOp, CmpOp, DType, Device, Error, IntOps, Layout, Result, Shape, dispatch_int, dispatch_int_raw, dispatch_int2, dispatch_int2_raw,
10};
11
12fn build_int(n: usize, dtype: IntDType, v: i64) -> CpuIntStorage {
13    match dtype {
14        IntDType::I32 => CpuIntStorage::I32(vec![v as i32; n]),
15        IntDType::U32 => CpuIntStorage::U32(vec![v as u32; n]),
16        IntDType::U8 => CpuIntStorage::U8(vec![v as u8; n]),
17    }
18}
19
20impl IntOps<Cpu> for Cpu {
21    fn i_zeros(shape: &Shape, _device: &Cpu, dtype: IntDType) -> Result<<Cpu as Device>::IntStorage> {
22        Ok(build_int(shape.element_count(), dtype, 0))
23    }
24
25    fn i_ones(shape: &Shape, _device: &Cpu, dtype: IntDType) -> Result<<Cpu as Device>::IntStorage> {
26        Ok(build_int(shape.element_count(), dtype, 1))
27    }
28
29    fn i_full(shape: &Shape, value: i64, _device: &Cpu, dtype: IntDType) -> Result<<Cpu as Device>::IntStorage> {
30        Ok(build_int(shape.element_count(), dtype, value))
31    }
32
33    fn i_from_i64<'a>(data: impl Into<Cow<'a, [i64]>>, _device: &Cpu) -> Result<<Cpu as Device>::IntStorage> {
34        let data = data.into();
35        Ok(match data {
36            Cow::Owned(v) => CpuIntStorage::I32(v.iter().map(|&x| x as i32).collect()),
37            Cow::Borrowed(s) => CpuIntStorage::I32(s.iter().map(|&x| x as i32).collect()),
38        })
39    }
40
41    fn i_from_i32<'a>(data: impl Into<Cow<'a, [i32]>>, _device: &Cpu) -> Result<<Cpu as Device>::IntStorage> {
42        let data = data.into();
43        Ok(match data {
44            Cow::Owned(v) => CpuIntStorage::I32(v),
45            Cow::Borrowed(s) => CpuIntStorage::I32(s.to_vec()),
46        })
47    }
48
49    fn i_from_u32<'a>(data: impl Into<Cow<'a, [u32]>>, _device: &Cpu) -> Result<<Cpu as Device>::IntStorage> {
50        let data = data.into();
51        Ok(match data {
52            Cow::Owned(v) => CpuIntStorage::U32(v),
53            Cow::Borrowed(s) => CpuIntStorage::U32(s.to_vec()),
54        })
55    }
56
57    fn i_from_u8<'a>(data: impl Into<Cow<'a, [u8]>>, _device: &Cpu) -> Result<<Cpu as Device>::IntStorage> {
58        let data = data.into();
59        Ok(match data {
60            Cow::Owned(v) => CpuIntStorage::U8(v),
61            Cow::Borrowed(s) => CpuIntStorage::U8(s.to_vec()),
62        })
63    }
64
65    fn i_from_bytes<'a>(
66        bytes: impl Into<Cow<'a, [u8]>>,
67        _shape: &Shape,
68        _device: &Cpu,
69        dtype: IntDType,
70    ) -> Result<<Cpu as Device>::IntStorage> {
71        let bytes = bytes.into();
72        Ok(match dtype {
73            IntDType::I32 => {
74                let v: Vec<i32> = bytes.chunks_exact(4).map(|c| i32::from_le_bytes(c.try_into().unwrap())).collect();
75                CpuIntStorage::I32(v)
76            }
77            IntDType::U32 => {
78                let v: Vec<u32> = bytes.chunks_exact(4).map(|c| u32::from_le_bytes(c.try_into().unwrap())).collect();
79                CpuIntStorage::U32(v)
80            }
81            IntDType::U8 => match bytes {
82                Cow::Owned(b) => CpuIntStorage::U8(b),
83                Cow::Borrowed(b) => CpuIntStorage::U8(b.to_vec()),
84            },
85        })
86    }
87
88    fn i_arange(start: i64, end: i64, step: i64, _device: &Cpu, dtype: IntDType) -> Result<(<Cpu as Device>::IntStorage, usize)> {
89        if step == 0 {
90            return Err(Error::Msg("arange step cannot be 0".into()));
91        }
92        let mut data = Vec::new();
93        let mut v = start;
94        if step > 0 {
95            while v < end {
96                data.push(v);
97                v += step;
98            }
99        } else {
100            while v > end {
101                data.push(v);
102                v += step;
103            }
104        }
105        let n = data.len();
106        let storage = match dtype {
107            IntDType::I32 => {
108                let v: Vec<i32> = data.iter().map(|&x| x as i32).collect();
109                Cpu::i_from_i32(v, &Cpu)?
110            }
111            IntDType::U32 => {
112                let v: Vec<u32> = data.iter().map(|&x| x as u32).collect();
113                Cpu::i_from_u32(v, &Cpu)?
114            }
115            IntDType::U8 => {
116                let v: Vec<u8> = data.iter().map(|&x| x as u8).collect();
117                Cpu::i_from_u8(v, &Cpu)?
118            }
119        };
120        Ok((storage, n))
121    }
122
123    fn i_contiguous(x: &<Cpu as Device>::IntStorage, l: &Layout) -> Result<<Cpu as Device>::IntStorage> {
124        Ok(dispatch_int!(x, |d| super::kernels::iter::gather(d, l)))
125    }
126
127    fn i_cast_float(x: &CpuIntStorage, layout: &Layout, to: FloatDType) -> Result<CpuFloatStorage> {
128        let s = match to {
129            FloatDType::F32 => CpuFloatStorage::F32(dispatch_int_raw!(x, |d| layout.storage_indices().map(|i| d[i] as f32).collect())),
130            FloatDType::F64 => CpuFloatStorage::F64(dispatch_int_raw!(x, |d| layout.storage_indices().map(|i| d[i] as f64).collect())),
131        };
132        Ok(s)
133    }
134
135    fn i_cast_int(x: &CpuIntStorage, layout: &Layout, to: IntDType) -> Result<CpuIntStorage> {
136        let s = match to {
137            IntDType::I32 => CpuIntStorage::I32(dispatch_int_raw!(x, |d| layout.storage_indices().map(|i| d[i] as i32).collect())),
138            IntDType::U32 => CpuIntStorage::U32(dispatch_int_raw!(x, |d| layout.storage_indices().map(|i| d[i] as u32).collect())),
139            IntDType::U8 => CpuIntStorage::U8(dispatch_int_raw!(x, |d| layout.storage_indices().map(|i| d[i] as u8).collect())),
140        };
141        Ok(s)
142    }
143
144    fn i_cast_bool(x: &CpuIntStorage, layout: &Layout, _to: BoolDType) -> Result<CpuBoolStorage> {
145        Ok(CpuBoolStorage(dispatch_int_raw!(x, |d| layout.storage_indices().map(|i| d[i] != 0).collect())))
146    }
147
148    fn i_to_vec(x: &<Cpu as Device>::IntStorage, layout: &Layout) -> Result<Vec<i64>> {
149        Ok(match x {
150            CpuIntStorage::I32(d) => layout.storage_indices().map(|i| d[i] as i64).collect(),
151            CpuIntStorage::U32(d) => layout.storage_indices().map(|i| d[i] as i64).collect(),
152            CpuIntStorage::U8(d) => layout.storage_indices().map(|i| d[i] as i64).collect(),
153        })
154    }
155
156    fn i_to_bytes<'a>(x: &'a <Cpu as Device>::IntStorage, layout: &Layout) -> Result<Cow<'a, [u8]>> {
157        if layout.is_contiguous() {
158            Ok(match x {
159                CpuIntStorage::I32(d) => Cow::Borrowed(bytemuck::cast_slice(d)),
160                CpuIntStorage::U32(d) => Cow::Borrowed(bytemuck::cast_slice(d)),
161                CpuIntStorage::U8(d) => Cow::Borrowed(d),
162            })
163        } else {
164            let contig = Self::i_contiguous(x, layout)?;
165            Ok(match contig {
166                CpuIntStorage::I32(d) => Cow::Owned(bytemuck::cast_slice(&d).to_vec()),
167                CpuIntStorage::U32(d) => Cow::Owned(bytemuck::cast_slice(&d).to_vec()),
168                CpuIntStorage::U8(d) => Cow::Owned(d),
169            })
170        }
171    }
172
173    fn i_binary(
174        lhs: &<Cpu as Device>::IntStorage,
175        lhs_l: &Layout,
176        rhs: &<Cpu as Device>::IntStorage,
177        rhs_l: &Layout,
178        op: BinaryOp,
179    ) -> Result<<Cpu as Device>::IntStorage> {
180        dispatch_int2!(lhs, rhs, "int binary", |a, b| ew::num_binary(a, lhs_l, b, rhs_l, op))
181    }
182
183    fn i_binary_scalar(lhs: &<Cpu as Device>::IntStorage, lhs_l: &Layout, rhs: i64, op: BinaryOp) -> Result<<Cpu as Device>::IntStorage> {
184        Ok(match lhs {
185            CpuIntStorage::I32(d) => CpuIntStorage::I32(ew::num_binary_scalar(d, lhs_l, rhs as i32, op)),
186            CpuIntStorage::U32(d) => CpuIntStorage::U32(ew::num_binary_scalar(d, lhs_l, rhs as u32, op)),
187            CpuIntStorage::U8(d) => CpuIntStorage::U8(ew::num_binary_scalar(d, lhs_l, rhs as u8, op)),
188        })
189    }
190
191    fn i_binary_(
192        dst: &mut <Cpu as Device>::IntStorage,
193        dst_l: &Layout,
194        src: &<Cpu as Device>::IntStorage,
195        src_l: &Layout,
196        op: BinaryOp,
197    ) -> Result<()> {
198        match (dst, src) {
199            (CpuIntStorage::I32(d), CpuIntStorage::I32(s)) => {
200                ew::binary_(d, dst_l, s, src_l, ew::num_binary_fn::<i32>(op));
201                Ok(())
202            }
203            (CpuIntStorage::U32(d), CpuIntStorage::U32(s)) => {
204                ew::binary_(d, dst_l, s, src_l, ew::num_binary_fn::<u32>(op));
205                Ok(())
206            }
207            (CpuIntStorage::U8(d), CpuIntStorage::U8(s)) => {
208                ew::binary_(d, dst_l, s, src_l, ew::num_binary_fn::<u8>(op));
209                Ok(())
210            }
211            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "int binary_" }),
212        }
213    }
214
215    fn i_binary_scalar_(dst: &mut <Cpu as Device>::IntStorage, dst_l: &Layout, rhs: i64, op: BinaryOp) -> Result<()> {
216        match dst {
217            CpuIntStorage::I32(d) => {
218                ew::binary_scalar_(d, dst_l, rhs as i32, ew::num_binary_fn::<i32>(op));
219                Ok(())
220            }
221            CpuIntStorage::U32(d) => {
222                ew::binary_scalar_(d, dst_l, rhs as u32, ew::num_binary_fn::<u32>(op));
223                Ok(())
224            }
225            CpuIntStorage::U8(d) => {
226                ew::binary_scalar_(d, dst_l, rhs as u8, ew::num_binary_fn::<u8>(op));
227                Ok(())
228            }
229        }
230    }
231
232    fn i_binary_scalar_lhs(
233        scalar: i64,
234        rhs: &<Cpu as Device>::IntStorage,
235        rhs_l: &Layout,
236        op: BinaryOp,
237    ) -> Result<<Cpu as Device>::IntStorage> {
238        Ok(match rhs {
239            CpuIntStorage::I32(d) => CpuIntStorage::I32(ew::num_scalar_binary(scalar as i32, d, rhs_l, op)),
240            CpuIntStorage::U32(d) => CpuIntStorage::U32(ew::num_scalar_binary(scalar as u32, d, rhs_l, op)),
241            CpuIntStorage::U8(d) => CpuIntStorage::U8(ew::num_scalar_binary(scalar as u8, d, rhs_l, op)),
242        })
243    }
244
245    fn i_unary(x: &<Cpu as Device>::IntStorage, l: &Layout, op: crate::UnaryOp<i64>) -> Result<<Cpu as Device>::IntStorage> {
246        use super::kernels::element::CpuNum;
247        match x {
248            CpuIntStorage::I32(d) => Ok(CpuIntStorage::I32(match op {
249                crate::UnaryOp::Neg => ew::unary(d, l, |v: i32| -v),
250                crate::UnaryOp::Abs => ew::unary(d, l, |v: i32| CpuNum::abs(v)),
251                crate::UnaryOp::Sign => ew::unary(d, l, |v: i32| CpuNum::signum(v)),
252                crate::UnaryOp::Affine(mul, add) => ew::unary(d, l, |v: i32| (v as i64 * mul + add) as i32),
253                crate::UnaryOp::Pow(exp) => ew::unary(d, l, |v: i32| (v as i64).pow(exp as u32) as i32),
254                crate::UnaryOp::Clamp(min, max) => {
255                    let lo = min.map(|v| v as i32);
256                    let hi = max.map(|v| v as i32);
257                    ew::unary(d, l, |v: i32| {
258                        let mut val = v;
259                        if let Some(lo) = lo {
260                            val = val.max(lo);
261                        }
262                        if let Some(hi) = hi {
263                            val = val.min(hi);
264                        }
265                        val
266                    })
267                }
268            })),
269            CpuIntStorage::U32(d) => Ok(CpuIntStorage::U32(match op {
270                crate::UnaryOp::Neg => ew::unary(d, l, |v: u32| -(v as i32) as u32),
271                crate::UnaryOp::Abs => ew::unary(d, l, |v: u32| CpuNum::abs(v)),
272                crate::UnaryOp::Sign => ew::unary(d, l, |v: u32| CpuNum::signum(v)),
273                crate::UnaryOp::Affine(mul, add) => ew::unary(d, l, |v: u32| (v as u64 * mul as u64 + add as u64) as u32),
274                crate::UnaryOp::Pow(exp) => ew::unary(d, l, |v: u32| (v as u64).pow(exp as u32) as u32),
275                crate::UnaryOp::Clamp(min, max) => {
276                    let lo = min.map(|v| v.max(0) as u32);
277                    let hi = max.map(|v| v as u32);
278                    ew::unary(d, l, |v: u32| {
279                        let mut val = v;
280                        if let Some(lo) = lo {
281                            val = val.max(lo);
282                        }
283                        if let Some(hi) = hi {
284                            val = val.min(hi);
285                        }
286                        val
287                    })
288                }
289            })),
290            CpuIntStorage::U8(d) => Ok(CpuIntStorage::U8(match op {
291                crate::UnaryOp::Neg => ew::unary(d, l, |v: u8| -(v as i32) as u8),
292                crate::UnaryOp::Abs => ew::unary(d, l, |v: u8| CpuNum::abs(v)),
293                crate::UnaryOp::Sign => ew::unary(d, l, |v: u8| CpuNum::signum(v)),
294                crate::UnaryOp::Affine(mul, add) => ew::unary(d, l, |v: u8| (v as i64 * mul + add) as u8),
295                crate::UnaryOp::Pow(exp) => ew::unary(d, l, |v: u8| (v as u64).pow(exp as u32) as u8),
296                crate::UnaryOp::Clamp(min, max) => {
297                    let lo = min.map(|v| v.max(0) as u8);
298                    let hi = max.map(|v| v as u8);
299                    ew::unary(d, l, |v: u8| {
300                        let mut val = v;
301                        if let Some(lo) = lo {
302                            val = val.max(lo);
303                        }
304                        if let Some(hi) = hi {
305                            val = val.min(hi);
306                        }
307                        val
308                    })
309                }
310            })),
311        }
312    }
313
314    fn i_unary_(dst: &mut <Cpu as Device>::IntStorage, dst_l: &Layout, op: crate::UnaryOp<i64>) -> Result<()> {
315        use super::kernels::element::CpuNum;
316        match dst {
317            CpuIntStorage::I32(d) => match op {
318                crate::UnaryOp::Neg => {
319                    ew::unary_(d, dst_l, |v: i32| -v);
320                    Ok(())
321                }
322                crate::UnaryOp::Abs => {
323                    ew::unary_(d, dst_l, |v: i32| CpuNum::abs(v));
324                    Ok(())
325                }
326                crate::UnaryOp::Sign => {
327                    ew::unary_(d, dst_l, |v: i32| CpuNum::signum(v));
328                    Ok(())
329                }
330                crate::UnaryOp::Affine(mul, add) => {
331                    ew::unary_(d, dst_l, |v: i32| (v as i64 * mul + add) as i32);
332                    Ok(())
333                }
334                crate::UnaryOp::Pow(exp) => {
335                    ew::unary_(d, dst_l, |v: i32| (v as i64).pow(exp as u32) as i32);
336                    Ok(())
337                }
338                crate::UnaryOp::Clamp(min, max) => {
339                    let lo = min.map(|v| v as i32);
340                    let hi = max.map(|v| v as i32);
341                    ew::unary_(d, dst_l, |v: i32| {
342                        let mut val = v;
343                        if let Some(lo) = lo {
344                            val = val.max(lo);
345                        }
346                        if let Some(hi) = hi {
347                            val = val.min(hi);
348                        }
349                        val
350                    });
351                    Ok(())
352                }
353            },
354            CpuIntStorage::U32(d) => match op {
355                crate::UnaryOp::Neg => {
356                    ew::unary_(d, dst_l, |v: u32| -(v as i32) as u32);
357                    Ok(())
358                }
359                crate::UnaryOp::Abs => {
360                    ew::unary_(d, dst_l, |v: u32| CpuNum::abs(v));
361                    Ok(())
362                }
363                crate::UnaryOp::Sign => {
364                    ew::unary_(d, dst_l, |v: u32| CpuNum::signum(v));
365                    Ok(())
366                }
367                crate::UnaryOp::Affine(mul, add) => {
368                    ew::unary_(d, dst_l, |v: u32| (v as u64 * mul as u64 + add as u64) as u32);
369                    Ok(())
370                }
371                crate::UnaryOp::Pow(exp) => {
372                    ew::unary_(d, dst_l, |v: u32| (v as u64).pow(exp as u32) as u32);
373                    Ok(())
374                }
375                crate::UnaryOp::Clamp(min, max) => {
376                    let lo = min.map(|v| v.max(0) as u32);
377                    let hi = max.map(|v| v as u32);
378                    ew::unary_(d, dst_l, |v: u32| {
379                        let mut val = v;
380                        if let Some(lo) = lo {
381                            val = val.max(lo);
382                        }
383                        if let Some(hi) = hi {
384                            val = val.min(hi);
385                        }
386                        val
387                    });
388                    Ok(())
389                }
390            },
391            CpuIntStorage::U8(d) => match op {
392                crate::UnaryOp::Neg => {
393                    ew::unary_(d, dst_l, |v: u8| -(v as i32) as u8);
394                    Ok(())
395                }
396                crate::UnaryOp::Abs => {
397                    ew::unary_(d, dst_l, |v: u8| CpuNum::abs(v));
398                    Ok(())
399                }
400                crate::UnaryOp::Sign => {
401                    ew::unary_(d, dst_l, |v: u8| CpuNum::signum(v));
402                    Ok(())
403                }
404                crate::UnaryOp::Affine(mul, add) => {
405                    ew::unary_(d, dst_l, |v: u8| (v as i64 * mul + add) as u8);
406                    Ok(())
407                }
408                crate::UnaryOp::Pow(exp) => {
409                    ew::unary_(d, dst_l, |v: u8| (v as u64).pow(exp as u32) as u8);
410                    Ok(())
411                }
412                crate::UnaryOp::Clamp(min, max) => {
413                    let lo = min.map(|v| v.max(0) as u8);
414                    let hi = max.map(|v| v as u8);
415                    ew::unary_(d, dst_l, |v: u8| {
416                        let mut val = v;
417                        if let Some(lo) = lo {
418                            val = val.max(lo);
419                        }
420                        if let Some(hi) = hi {
421                            val = val.min(hi);
422                        }
423                        val
424                    });
425                    Ok(())
426                }
427            },
428        }
429    }
430
431    fn i_matmul(
432        lhs: &<Cpu as Device>::IntStorage,
433        lhs_l: &Layout,
434        rhs: &<Cpu as Device>::IntStorage,
435        rhs_l: &Layout,
436    ) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
437        let result = match (lhs, rhs) {
438            (CpuIntStorage::I32(a), CpuIntStorage::I32(b)) => {
439                let (vec, shape) = matmul::matmul(a, lhs_l, b, rhs_l)?;
440                (CpuIntStorage::I32(vec), shape)
441            }
442            (CpuIntStorage::U32(a), CpuIntStorage::U32(b)) => {
443                let (vec, shape) = matmul::matmul(a, lhs_l, b, rhs_l)?;
444                (CpuIntStorage::U32(vec), shape)
445            }
446            (CpuIntStorage::U8(a), CpuIntStorage::U8(b)) => {
447                let (vec, shape) = matmul::matmul(a, lhs_l, b, rhs_l)?;
448                (CpuIntStorage::U8(vec), shape)
449            }
450            (l, r) => return Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "int matmul" }),
451        };
452        Ok(result)
453    }
454
455    fn i_cmp(
456        lhs: &<Cpu as Device>::IntStorage,
457        lhs_l: &Layout,
458        rhs: &<Cpu as Device>::IntStorage,
459        rhs_l: &Layout,
460        op: CmpOp,
461    ) -> Result<<Cpu as Device>::BoolStorage> {
462        let v = dispatch_int2_raw!(lhs, rhs, "int cmp", |a, b| ew::num_cmp(a, lhs_l, b, rhs_l, op))?;
463        Ok(CpuBoolStorage(v))
464    }
465
466    fn i_cmp_scalar(lhs: &<Cpu as Device>::IntStorage, lhs_l: &Layout, rhs: i64, op: CmpOp) -> Result<<Cpu as Device>::BoolStorage> {
467        Ok(match lhs {
468            CpuIntStorage::I32(d) => CpuBoolStorage(ew::cmp_scalar(d, lhs_l, rhs as i32, op)),
469            CpuIntStorage::U32(d) => CpuBoolStorage(ew::cmp_scalar(d, lhs_l, rhs as u32, op)),
470            CpuIntStorage::U8(d) => CpuBoolStorage(ew::cmp_scalar(d, lhs_l, rhs as u8, op)),
471        })
472    }
473
474    fn i_reduce(
475        x: &<Cpu as Device>::IntStorage,
476        l: &Layout,
477        dims: &[usize],
478        keepdim: bool,
479        op: crate::ReduceOp,
480    ) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
481        let reducer = reduce::Reducer::from(op);
482        match x {
483            CpuIntStorage::I32(d) => {
484                let (v, s) = reduce::reduce_dims(d, l, dims, keepdim, reducer)?;
485                Ok((CpuIntStorage::I32(v), s))
486            }
487            CpuIntStorage::U32(d) => {
488                let (v, s) = reduce::reduce_dims(d, l, dims, keepdim, reducer)?;
489                Ok((CpuIntStorage::U32(v), s))
490            }
491            CpuIntStorage::U8(d) => {
492                let (v, s) = reduce::reduce_dims(d, l, dims, keepdim, reducer)?;
493                Ok((CpuIntStorage::U8(v), s))
494            }
495        }
496    }
497
498    fn i_index_select(
499        x: &<Cpu as Device>::IntStorage,
500        x_l: &Layout,
501        idx: &<Cpu as Device>::IntStorage,
502        idx_l: &Layout,
503        dim: usize,
504    ) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
505        let ids = int_ids_as_usize(idx, idx_l);
506        match x {
507            CpuIntStorage::I32(d) => {
508                let (v, dims) = indexing::index_select(d, x_l, &ids, idx_l, dim)?;
509                Ok((CpuIntStorage::I32(v), Shape::from(dims)))
510            }
511            CpuIntStorage::U32(d) => {
512                let (v, dims) = indexing::index_select(d, x_l, &ids, idx_l, dim)?;
513                Ok((CpuIntStorage::U32(v), Shape::from(dims)))
514            }
515            CpuIntStorage::U8(d) => {
516                let (v, dims) = indexing::index_select(d, x_l, &ids, idx_l, dim)?;
517                Ok((CpuIntStorage::U8(v), Shape::from(dims)))
518            }
519        }
520    }
521
522    fn i_gather(
523        x: &<Cpu as Device>::IntStorage,
524        x_l: &Layout,
525        idx: &<Cpu as Device>::IntStorage,
526        idx_l: &Layout,
527        dim: usize,
528    ) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
529        let ids = int_ids_as_usize(idx, idx_l);
530        match x {
531            CpuIntStorage::I32(d) => {
532                let (v, dims) = indexing::gather(d, x_l, &ids, idx_l, dim)?;
533                Ok((CpuIntStorage::I32(v), Shape::from(dims)))
534            }
535            CpuIntStorage::U32(d) => {
536                let (v, dims) = indexing::gather(d, x_l, &ids, idx_l, dim)?;
537                Ok((CpuIntStorage::U32(v), Shape::from(dims)))
538            }
539            CpuIntStorage::U8(d) => {
540                let (v, dims) = indexing::gather(d, x_l, &ids, idx_l, dim)?;
541                Ok((CpuIntStorage::U8(v), Shape::from(dims)))
542            }
543        }
544    }
545
546    fn i_cat(srcs: &[(&<Cpu as Device>::IntStorage, &Layout)], dim: usize) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
547        if srcs.is_empty() {
548            return Err(Error::OpRequiresAtLeastOneTensor { op: "cat" });
549        }
550        let dt = srcs[0].0.dtype();
551        for (s, _) in srcs {
552            if s.dtype() != dt {
553                return Err(Error::DTypeMismatch { lhs: dt, rhs: s.dtype(), op: "cat" });
554            }
555        }
556        macro_rules! cat_variant {
557            ($variant:path, $getter:ident) => {{
558                let views: Vec<(&[_], &Layout)> = srcs.iter().map(|(s, l)| ($getter(s), *l)).collect();
559                let (v, shape) = super::kernels::shape::cat(&views, dim)?;
560                Ok(($variant(v), shape))
561            }};
562        }
563        match dt {
564            DType::I32 => cat_variant!(CpuIntStorage::I32, as_i32),
565            DType::U32 => cat_variant!(CpuIntStorage::U32, as_u32),
566            _ => cat_variant!(CpuIntStorage::U8, as_u8),
567        }
568    }
569
570    fn i_arg_reduce(
571        x: &<Cpu as Device>::IntStorage,
572        layout: &Layout,
573        dim: usize,
574        keepdim: bool,
575        take_max: bool,
576    ) -> Result<(<Cpu as Device>::IntStorage, Shape)> {
577        let (indices, shape) = dispatch_int_raw!(x, |d| reduce::arg_reduce(d, layout, dim, keepdim, take_max))?;
578        let indices_u32: Vec<u32> = indices.into_iter().map(|i| i as u32).collect();
579        Ok((CpuIntStorage::U32(indices_u32), shape))
580    }
581
582    fn i_index_add(
583        init: &<Cpu as Device>::IntStorage,
584        init_l: &Layout,
585        idx: &<Cpu as Device>::IntStorage,
586        idx_l: &Layout,
587        src: &<Cpu as Device>::IntStorage,
588        _src_l: &Layout,
589        dim: usize,
590    ) -> Result<<Cpu as Device>::IntStorage> {
591        let idx_u = int_ids_as_usize(idx, idx_l);
592        match (init, src) {
593            (CpuIntStorage::I32(init_d), CpuIntStorage::I32(src_d)) => {
594                let result = indexing::index_add(init_d, init_l, &idx_u, idx_l, src_d, dim)?;
595                Ok(CpuIntStorage::I32(result))
596            }
597            (CpuIntStorage::U32(init_d), CpuIntStorage::U32(src_d)) => {
598                let result = indexing::index_add(init_d, init_l, &idx_u, idx_l, src_d, dim)?;
599                Ok(CpuIntStorage::U32(result))
600            }
601            (CpuIntStorage::U8(init_d), CpuIntStorage::U8(src_d)) => {
602                let result = indexing::index_add(init_d, init_l, &idx_u, idx_l, src_d, dim)?;
603                Ok(CpuIntStorage::U8(result))
604            }
605            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "int index_add" }),
606        }
607    }
608
609    fn i_scatter_add(
610        init: &<Cpu as Device>::IntStorage,
611        init_l: &Layout,
612        idx: &<Cpu as Device>::IntStorage,
613        idx_l: &Layout,
614        src: &<Cpu as Device>::IntStorage,
615        _src_l: &Layout,
616        dim: usize,
617    ) -> Result<<Cpu as Device>::IntStorage> {
618        let idx_u = int_ids_as_usize(idx, idx_l);
619        match (init, src) {
620            (CpuIntStorage::I32(init_d), CpuIntStorage::I32(src_d)) => {
621                let result = indexing::scatter_add(init_d, init_l, &idx_u, idx_l, src_d, dim)?;
622                Ok(CpuIntStorage::I32(result))
623            }
624            (CpuIntStorage::U32(init_d), CpuIntStorage::U32(src_d)) => {
625                let result = indexing::scatter_add(init_d, init_l, &idx_u, idx_l, src_d, dim)?;
626                Ok(CpuIntStorage::U32(result))
627            }
628            (CpuIntStorage::U8(init_d), CpuIntStorage::U8(src_d)) => {
629                let result = indexing::scatter_add(init_d, init_l, &idx_u, idx_l, src_d, dim)?;
630                Ok(CpuIntStorage::U8(result))
631            }
632            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "int scatter_add" }),
633        }
634    }
635
636    fn i_pick(
637        mask: &<Cpu as Device>::BoolStorage,
638        mask_l: &Layout,
639        on_true: &<Cpu as Device>::IntStorage,
640        true_l: &Layout,
641        on_false: &<Cpu as Device>::IntStorage,
642        false_l: &Layout,
643    ) -> Result<<Cpu as Device>::IntStorage> {
644        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
645        match (on_true, on_false) {
646            (CpuIntStorage::I32(t), CpuIntStorage::I32(f)) => {
647                let tv: Vec<i32> = true_l.storage_indices().map(|i| t[i]).collect();
648                let fv: Vec<i32> = false_l.storage_indices().map(|i| f[i]).collect();
649                Ok(CpuIntStorage::I32(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { fv[i] }).collect()))
650            }
651            (CpuIntStorage::U32(t), CpuIntStorage::U32(f)) => {
652                let tv: Vec<u32> = true_l.storage_indices().map(|i| t[i]).collect();
653                let fv: Vec<u32> = false_l.storage_indices().map(|i| f[i]).collect();
654                Ok(CpuIntStorage::U32(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { fv[i] }).collect()))
655            }
656            (CpuIntStorage::U8(t), CpuIntStorage::U8(f)) => {
657                let tv: Vec<u8> = true_l.storage_indices().map(|i| t[i]).collect();
658                let fv: Vec<u8> = false_l.storage_indices().map(|i| f[i]).collect();
659                Ok(CpuIntStorage::U8(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { fv[i] }).collect()))
660            }
661            (l, r) => Err(Error::DTypeMismatch { lhs: l.dtype(), rhs: r.dtype(), op: "pick" }),
662        }
663    }
664
665    fn i_pick_true(
666        mask: &<Cpu as Device>::BoolStorage,
667        mask_l: &Layout,
668        value: i64,
669        on_false: &<Cpu as Device>::IntStorage,
670        false_l: &Layout,
671    ) -> Result<<Cpu as Device>::IntStorage> {
672        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
673        match on_false {
674            CpuIntStorage::I32(f) => {
675                let fv: Vec<i32> = false_l.storage_indices().map(|i| f[i]).collect();
676                let val = value as i32;
677                Ok(CpuIntStorage::I32(m.iter().enumerate().map(|(i, &c)| if c { val } else { fv[i] }).collect()))
678            }
679            CpuIntStorage::U32(f) => {
680                let fv: Vec<u32> = false_l.storage_indices().map(|i| f[i]).collect();
681                let val = value as u32;
682                Ok(CpuIntStorage::U32(m.iter().enumerate().map(|(i, &c)| if c { val } else { fv[i] }).collect()))
683            }
684            CpuIntStorage::U8(f) => {
685                let fv: Vec<u8> = false_l.storage_indices().map(|i| f[i]).collect();
686                let val = value as u8;
687                Ok(CpuIntStorage::U8(m.iter().enumerate().map(|(i, &c)| if c { val } else { fv[i] }).collect()))
688            }
689        }
690    }
691
692    fn i_pick_false(
693        mask: &<Cpu as Device>::BoolStorage,
694        mask_l: &Layout,
695        on_true: &<Cpu as Device>::IntStorage,
696        true_l: &Layout,
697        value: i64,
698    ) -> Result<<Cpu as Device>::IntStorage> {
699        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
700        match on_true {
701            CpuIntStorage::I32(t) => {
702                let tv: Vec<i32> = true_l.storage_indices().map(|i| t[i]).collect();
703                let val = value as i32;
704                Ok(CpuIntStorage::I32(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { val }).collect()))
705            }
706            CpuIntStorage::U32(t) => {
707                let tv: Vec<u32> = true_l.storage_indices().map(|i| t[i]).collect();
708                let val = value as u32;
709                Ok(CpuIntStorage::U32(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { val }).collect()))
710            }
711            CpuIntStorage::U8(t) => {
712                let tv: Vec<u8> = true_l.storage_indices().map(|i| t[i]).collect();
713                let val = value as u8;
714                Ok(CpuIntStorage::U8(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { val }).collect()))
715            }
716        }
717    }
718
719    fn i_allclose(a: &CpuIntStorage, a_l: &Layout, b: &CpuIntStorage, b_l: &Layout) -> Result<bool> {
720        match (a, b) {
721            (CpuIntStorage::I32(av), CpuIntStorage::I32(bv)) => {
722                Ok(a_l.storage_indices().zip(b_l.storage_indices()).all(|(ai, bi)| av[ai] == bv[bi]))
723            }
724            (CpuIntStorage::U32(av), CpuIntStorage::U32(bv)) => {
725                Ok(a_l.storage_indices().zip(b_l.storage_indices()).all(|(ai, bi)| av[ai] == bv[bi]))
726            }
727            (CpuIntStorage::U8(av), CpuIntStorage::U8(bv)) => {
728                Ok(a_l.storage_indices().zip(b_l.storage_indices()).all(|(ai, bi)| av[ai] == bv[bi]))
729            }
730            _ => Err(crate::Error::DTypeMismatch { lhs: a.dtype(), rhs: b.dtype(), op: "allclose" }),
731        }
732    }
733}
734
735fn as_i32(s: &CpuIntStorage) -> &[i32] {
736    match s {
737        CpuIntStorage::I32(d) => d,
738        _ => unreachable!(),
739    }
740}
741fn as_u32(s: &CpuIntStorage) -> &[u32] {
742    match s {
743        CpuIntStorage::U32(d) => d,
744        _ => unreachable!(),
745    }
746}
747fn as_u8(s: &CpuIntStorage) -> &[u8] {
748    match s {
749        CpuIntStorage::U8(d) => d,
750        _ => unreachable!(),
751    }
752}
luma_tensor/device/cpu/ops/int_ops.rs

luma_tensor/device/cpu/ops/
int_ops.rs
