rustymath_kernels/operation/
at.rs

1use super::*;
2
3/// Writes to the addresses of `ids`, replacing or combining with the current items. Invalid positions, including fills, write nothing; repeated positions combine with the earlier result.
4pub fn at(
5    combine: Option<&Kernel>, root: WriteRoot<'_>, values: crate::agreement::Read<'_>, ids: Writes<'_>, checked: bool, settings: Settings, output: &mut Buffer
6) -> Result<(), Error> {
7    let (wanted, n) = (output.dtype(), ids.len());
8    if values.plan.len() != 1 && values.plan.len() != n { return Err(Error::new(Kind::Length, "write values do not match positions")); }
9    let projection = combine.map_or(Projection::AsIs, |k| k.projection(settings));
10    match combine {
11        None => {
12            root.initialize(wanted, checked, Projection::AsIs, output)?;
13            with_dtype!(wanted, T => at_typed::<T>(values, ids, checked, None, output))?;
14        }
15        Some(kernel) => kernel
16            .projected(&[values], None, Seed::Identity, settings, |inputs, _| kernel.write(root, inputs[0], ids, checked, projection, settings, output))?,
17    }
18    finish(output, wanted, crate::domain::Variant { checked, ..crate::domain::Variant::checked(wanted) })
19}
20
21/// Writes `values` into `output` at `ids`, a chunk at a time: over the items there, or merged with them by `combine`, which gives any failure it finds. `combine` is one `dyn` call for each chunk, so this compiles once for each item type.
22pub(crate) fn at_typed<T: crate::buffer::Cast>(
23    values: crate::agreement::Read<'_>, ids: Writes<'_>, checked: bool,
24    combine: Option<&dyn Fn(&mut [T], crate::gather::Writes<'_>, Arg<'_, T>) -> Option<Kind>>, output: &mut Buffer) -> Result<(), Error> {
25    let out = T::vec_mut(output).expect("write output has its working dtype");
26    let mut scratch = Vec::new();
27    let mut result = Ok(());
28    ids.runs_while(&mut |start, positions| {
29        result = (|| {
30            for (k, positions) in positions.chunks(crate::CHUNK).enumerate() {
31                let start = start + k * crate::CHUNK;
32                let source = values.part::<T>(if values.plan.len() == 1 { 0..1 } else { start..start + positions.len() }, checked, &mut scratch)?;
33                let writes = Writes::Offsets(positions);
34                match combine {
35                    None => match source {
36                        Arg::One(value) => writes.write(out, &[value], crate::gather::Same), Arg::Many(values) => writes.write(out, values, crate::gather::Same)
37                    },
38                    Some(combine) => {
39                        if let Some(kind) = combine(&mut out[..], writes, source) {
40                            return Err(Error::new(kind, "modified write is outside the function's domain"));
41                        }
42                    }
43                }
44            }
45            Ok(())
46        })();
47        result.is_ok()
48    });
49    result
50}
51
52impl Kernel {
53    /// `values` written into `root`'s items at `ids`, combined by this kernel: by its fold's reduction where that runs at the variant's type or the type the formula computes in (`with_binary_body!`), and otherwise item by item (`Route::At`).
54    fn write(
55        &self, root: WriteRoot<'_>, values: crate::agreement::Read<'_>, ids: Writes<'_>, checked: bool, projection: Projection, settings: Settings,
56        output: &mut Buffer) -> Result<(), Error> {
57        let variant = self.variant;
58        if let Some(fold) = self.fold() {
59            let working = match variant.dtype {
60                DType::Boolean => DType::Integer(Width::U8), DType::Float(FloatWidth::F16) => DType::Float(FloatWidth::F32), dtype => dtype
61            };
62            native_reduction!(fold, variant.dtype, variant.fails_on_overflow(), |R, T| {
63                type S = <R as crate::reduce::Reduction<T>>::Total;
64                if S::DTYPE == variant.dtype || S::DTYPE == working {
65                    root.initialize(S::DTYPE, checked, projection, output)?;
66                    return at_typed::<S>(
67                        values, ids, checked, Some(&|out, writes, items| {
68                            writes.combine(out, items, |x, y| {
69                                let (value, ok) = <R as crate::reduce::Reduction<T>>::chain(x, y);
70                                (value, (!ok && variant.fails_on_overflow()).then_some(Kind::Overflow))
71                            })
72                        }), output,
73                    );
74                }
75                None::<()>
76            });
77        }
78        self.definition.binary(variant, settings, Route::At { root, values, ids, checked, projection }, output)
79    }
80}
rustymath_kernels/operation/at.rs

rustymath_kernels/operation/
at.rs
