nuvec_core/
primitive.rs

1//! What a graph node computes.
2//!
3//! Primitives form a closed enum (D-015): equality and hashing — the basis for
4//! common-subexpression elimination — are derived, and every per-primitive rule (output
5//! inference here; gradient, batching and fusion rules later) is an exhaustive `match`.
6//!
7//! Primitives require operands of equal shape and dtype. Promotion and broadcasting happen
8//! in the ops layer, which inserts `AsType` and `Broadcast` nodes (D-019).
9
10use smallvec::{SmallVec, smallvec};
11
12use crate::shape::{Shape, Strides};
13use crate::view::{self, DimSlice};
14use crate::{Array, DType, Error, Result};
15
16/// An element-wise function of one operand.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18#[non_exhaustive]
19pub enum UnaryOp {
20    /// `-x`.
21    Negative,
22    /// `|x|`; the magnitude for complex numbers.
23    Abs,
24    /// `e^x`.
25    Exp,
26    /// Natural logarithm.
27    Log,
28    /// Square root.
29    Sqrt,
30    /// Sine.
31    Sin,
32    /// Cosine.
33    Cos,
34    /// Hyperbolic tangent.
35    Tanh,
36    /// `1 / (1 + e^-x)`.
37    Sigmoid,
38}
39
40impl UnaryOp {
41    /// Whether the function computes in floating point (integer inputs are promoted).
42    pub(crate) fn is_transcendental(self) -> bool {
43        !matches!(self, Self::Negative | Self::Abs)
44    }
45
46    pub(crate) fn name(self) -> &'static str {
47        match self {
48            Self::Negative => "negative",
49            Self::Abs => "abs",
50            Self::Exp => "exp",
51            Self::Log => "log",
52            Self::Sqrt => "sqrt",
53            Self::Sin => "sin",
54            Self::Cos => "cos",
55            Self::Tanh => "tanh",
56            Self::Sigmoid => "sigmoid",
57        }
58    }
59}
60
61/// An element-wise function of two operands.
62#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
63#[non_exhaustive]
64pub enum BinaryOp {
65    /// `a + b`.
66    Add,
67    /// `a - b`.
68    Subtract,
69    /// `a * b`.
70    Multiply,
71    /// `a / b` (true division).
72    Divide,
73    /// The larger operand; NaN propagates.
74    Maximum,
75    /// The smaller operand; NaN propagates.
76    Minimum,
77    /// `a == b`.
78    Equal,
79    /// `a != b`.
80    NotEqual,
81    /// `a < b`.
82    Less,
83    /// `a <= b`.
84    LessEqual,
85    /// `a > b`.
86    Greater,
87    /// `a >= b`.
88    GreaterEqual,
89}
90
91impl BinaryOp {
92    /// Whether the result is boolean.
93    pub(crate) fn is_comparison(self) -> bool {
94        matches!(
95            self,
96            Self::Equal
97                | Self::NotEqual
98                | Self::Less
99                | Self::LessEqual
100                | Self::Greater
101                | Self::GreaterEqual
102        )
103    }
104
105    pub(crate) fn name(self) -> &'static str {
106        match self {
107            Self::Add => "add",
108            Self::Subtract => "subtract",
109            Self::Multiply => "multiply",
110            Self::Divide => "divide",
111            Self::Maximum => "maximum",
112            Self::Minimum => "minimum",
113            Self::Equal => "equal",
114            Self::NotEqual => "not_equal",
115            Self::Less => "less",
116            Self::LessEqual => "less_equal",
117            Self::Greater => "greater",
118            Self::GreaterEqual => "greater_equal",
119        }
120    }
121}
122
123/// A reduction over a set of axes.
124#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
125#[non_exhaustive]
126pub enum ReduceOp {
127    /// Sum (identity 0).
128    Sum,
129    /// Product (identity 1).
130    Prod,
131    /// Maximum; NaN propagates. No identity.
132    Max,
133    /// Minimum; NaN propagates. No identity.
134    Min,
135    /// Logical AND of bools (identity true).
136    And,
137    /// Logical OR of bools (identity false).
138    Or,
139}
140
141impl ReduceOp {
142    pub(crate) fn name(self) -> &'static str {
143        match self {
144            Self::Sum => "sum",
145            Self::Prod => "prod",
146            Self::Max => "max",
147            Self::Min => "min",
148            Self::And => "all",
149            Self::Or => "any",
150        }
151    }
152}
153
154/// The position of an extreme value along one axis.
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
156#[non_exhaustive]
157pub enum ArgReduceOp {
158    /// Index of the maximum (first occurrence; first NaN wins).
159    Max,
160    /// Index of the minimum (first occurrence; first NaN wins).
161    Min,
162}
163
164impl ArgReduceOp {
165    pub(crate) fn name(self) -> &'static str {
166        match self {
167            Self::Max => "argmax",
168            Self::Min => "argmin",
169        }
170    }
171}
172
173/// An operation recorded in the graph, with its parameters.
174#[derive(Debug, Clone, PartialEq, Eq, Hash)]
175#[non_exhaustive]
176pub enum Primitive {
177    /// Element-wise function of one operand.
178    Unary(UnaryOp),
179    /// Element-wise function of two operands of equal shape and dtype.
180    Binary(BinaryOp),
181    /// `where(condition, a, b)`: a boolean condition selects from `a` or `b`.
182    Select,
183    /// Conversion to another element type.
184    AsType(DType),
185    /// The same elements in another shape. A view when the layout permits, else a copy.
186    Reshape(Vec<usize>),
187    /// Axes reordered: output axis `i` is input axis `axes[i]`. A view.
188    Transpose(Vec<usize>),
189    /// A normalised slice per dimension. A view.
190    Slice {
191        /// First index per dimension.
192        start: Vec<usize>,
193        /// Step per dimension (non-zero, may be negative).
194        step: Vec<isize>,
195        /// Number of selected indices per dimension (the output shape).
196        len: Vec<usize>,
197    },
198    /// Presented as a larger shape with zero strides. A view.
199    Broadcast(Vec<usize>),
200    /// Reduction over sorted, unique `axes`; reduced dimensions become size 1.
201    Reduce {
202        /// The reduction.
203        op: ReduceOp,
204        /// Reduced axes, sorted and unique.
205        axes: Vec<usize>,
206    },
207    /// Index of an extreme value along `axis` (uint32); the axis becomes size 1.
208    ArgReduce {
209        /// Maximum or minimum.
210        op: ArgReduceOp,
211        /// The reduced axis.
212        axis: usize,
213    },
214    /// Batched matrix product of `[B..., M, K]` and `[B..., K, N]` with equal batch
215    /// dimensions.
216    Matmul,
217    /// Test-only two-output primitive: `(a + b, a * b)`.
218    #[cfg(test)]
219    AddMultiply,
220}
221
222/// The dtype, shape and (for views) strides of one output of a primitive.
223#[derive(Debug, Clone, PartialEq, Eq)]
224pub(crate) struct OutputSpec {
225    pub(crate) dtype: DType,
226    pub(crate) shape: Shape,
227    /// `None` means row-contiguous.
228    pub(crate) strides: Option<Strides>,
229}
230
231impl Primitive {
232    /// The operation's name, as used in error messages.
233    pub fn name(&self) -> &'static str {
234        match self {
235            Self::Unary(op) => op.name(),
236            Self::Binary(op) => op.name(),
237            Self::Select => "where",
238            Self::AsType(_) => "astype",
239            Self::Reshape(_) => "reshape",
240            Self::Transpose(_) => "transpose",
241            Self::Slice { .. } => "slice",
242            Self::Broadcast(_) => "broadcast_to",
243            Self::Reduce { op, .. } => op.name(),
244            Self::ArgReduce { op, .. } => op.name(),
245            Self::Matmul => "matmul",
246            #[cfg(test)]
247            Self::AddMultiply => "add_multiply",
248        }
249    }
250
251    /// Validates the inputs and infers the outputs, without computing anything.
252    pub(crate) fn output_specs(&self, inputs: &[Array]) -> Result<SmallVec<[OutputSpec; 1]>> {
253        let computed = |dtype: DType, shape: &[usize]| OutputSpec {
254            dtype,
255            shape: shape.into(),
256            strides: None,
257        };
258        let spec = match self {
259            Self::Unary(op) => {
260                let [a] = self.arity::<1>(inputs)?;
261                let dtype = match (op, a.dtype()) {
262                    (UnaryOp::Abs, DType::Complex64) => DType::F32,
263                    (_, d) => d,
264                };
265                computed(dtype, a.shape())
266            }
267            Self::Binary(op) => {
268                let [a, b] = self.arity::<2>(inputs)?;
269                self.same_shape_and_dtype(a, b)?;
270                computed(
271                    if op.is_comparison() {
272                        DType::Bool
273                    } else {
274                        a.dtype()
275                    },
276                    a.shape(),
277                )
278            }
279            Self::Select => {
280                let [c, a, b] = self.arity::<3>(inputs)?;
281                self.same_shape_and_dtype(a, b)?;
282                if c.dtype() != DType::Bool || c.shape() != a.shape() {
283                    return Err(Error::InvalidArgument(
284                        "where: condition must be a bool array of the same shape".into(),
285                    ));
286                }
287                computed(a.dtype(), a.shape())
288            }
289            Self::AsType(dtype) => {
290                let [a] = self.arity::<1>(inputs)?;
291                computed(*dtype, a.shape())
292            }
293            Self::Reshape(shape) => {
294                let [a] = self.arity::<1>(inputs)?;
295                OutputSpec {
296                    dtype: a.dtype(),
297                    shape: shape.as_slice().into(),
298                    strides: view::nocopy_reshape(a.shape(), a.strides(), shape),
299                }
300            }
301            Self::Transpose(axes) => {
302                let [a] = self.arity::<1>(inputs)?;
303                let (shape, strides) = view::permute(a.shape(), a.strides(), axes)?;
304                OutputSpec {
305                    dtype: a.dtype(),
306                    shape,
307                    strides: Some(strides),
308                }
309            }
310            Self::Slice { len, .. } => {
311                let [a] = self.arity::<1>(inputs)?;
312                let (strides, _) = view::slice_geometry(a.strides(), &self.dim_slices(a.ndim())?);
313                OutputSpec {
314                    dtype: a.dtype(),
315                    shape: len.as_slice().into(),
316                    strides: Some(strides),
317                }
318            }
319            Self::Broadcast(shape) => {
320                let [a] = self.arity::<1>(inputs)?;
321                let strides = view::broadcast_strides(a.shape(), a.strides(), shape)?;
322                OutputSpec {
323                    dtype: a.dtype(),
324                    shape: shape.as_slice().into(),
325                    strides: Some(strides),
326                }
327            }
328            Self::Reduce { axes, .. } => {
329                let [a] = self.arity::<1>(inputs)?;
330                computed(a.dtype(), &self.reduced_shape(a.shape(), axes)?)
331            }
332            Self::ArgReduce { axis, .. } => {
333                let [a] = self.arity::<1>(inputs)?;
334                computed(DType::U32, &self.reduced_shape(a.shape(), &[*axis])?)
335            }
336            Self::Matmul => {
337                let [a, b] = self.arity::<2>(inputs)?;
338                let (sa, sb) = (a.shape(), b.shape());
339                let n = sa.len();
340                let valid =
341                    n >= 2 && sb.len() == n && sa[..n - 2] == sb[..n - 2] && sa[n - 1] == sb[n - 2];
342                if !valid {
343                    return Err(Error::ShapeMismatch {
344                        op: "matmul",
345                        lhs: sa.to_vec(),
346                        rhs: sb.to_vec(),
347                    });
348                }
349                if a.dtype() != b.dtype() {
350                    return Err(Error::DTypeMismatch {
351                        op: "matmul",
352                        lhs: a.dtype(),
353                        rhs: b.dtype(),
354                    });
355                }
356                let mut shape: Shape = sa.into();
357                shape[n - 1] = sb[n - 1];
358                computed(a.dtype(), &shape)
359            }
360            #[cfg(test)]
361            Self::AddMultiply => {
362                let [a, b] = self.arity::<2>(inputs)?;
363                self.same_shape_and_dtype(a, b)?;
364                let spec = computed(a.dtype(), a.shape());
365                return Ok(smallvec![spec.clone(), spec]);
366            }
367        };
368        Ok(smallvec![spec])
369    }
370
371    /// The per-dimension slices of a `Slice` primitive over `ndim` dimensions.
372    pub(crate) fn dim_slices(&self, ndim: usize) -> Result<Vec<DimSlice>> {
373        match self {
374            Self::Slice { start, step, len }
375                if start.len() == ndim && step.len() == ndim && len.len() == ndim =>
376            {
377                Ok((0..ndim)
378                    .map(|i| DimSlice {
379                        start: start[i],
380                        step: step[i],
381                        len: len[i],
382                    })
383                    .collect())
384            }
385            _ => Err(Error::InvalidArgument(
386                "slice: one start, step and length per dimension".into(),
387            )),
388        }
389    }
390
391    /// `shape` with `axes` set to 1, after validating the axes.
392    fn reduced_shape(&self, shape: &[usize], axes: &[usize]) -> Result<Shape> {
393        let mut out: Shape = shape.into();
394        for &axis in axes {
395            *out.get_mut(axis).ok_or_else(|| {
396                Error::InvalidArgument(format!(
397                    "{}: axis {axis} out of range for shape {shape:?}",
398                    self.name()
399                ))
400            })? = 1;
401        }
402        Ok(out)
403    }
404
405    fn arity<'a, const N: usize>(&self, inputs: &'a [Array]) -> Result<&'a [Array; N]> {
406        inputs.try_into().map_err(|_| {
407            Error::InvalidArgument(format!(
408                "{} takes {N} inputs, got {}",
409                self.name(),
410                inputs.len()
411            ))
412        })
413    }
414
415    fn same_shape_and_dtype(&self, a: &Array, b: &Array) -> Result<()> {
416        if a.shape() != b.shape() {
417            return Err(Error::ShapeMismatch {
418                op: self.name(),
419                lhs: a.shape().to_vec(),
420                rhs: b.shape().to_vec(),
421            });
422        }
423        if a.dtype() != b.dtype() {
424            return Err(Error::DTypeMismatch {
425                op: self.name(),
426                lhs: a.dtype(),
427                rhs: b.dtype(),
428            });
429        }
430        Ok(())
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use std::collections::hash_map::DefaultHasher;
437    use std::hash::{Hash, Hasher};
438
439    use super::*;
440
441    fn hash(p: &Primitive) -> u64 {
442        let mut h = DefaultHasher::new();
443        p.hash(&mut h);
444        h.finish()
445    }
446
447    #[test]
448    fn equal_primitives_compare_and_hash_equal() {
449        let add = Primitive::Binary(BinaryOp::Add);
450        assert_eq!(add, Primitive::Binary(BinaryOp::Add));
451        assert_eq!(hash(&add), hash(&Primitive::Binary(BinaryOp::Add)));
452        assert_ne!(add, Primitive::Binary(BinaryOp::Multiply));
453        let s = |start: Vec<usize>| Primitive::Slice {
454            start,
455            step: vec![1],
456            len: vec![2],
457        };
458        assert_eq!(s(vec![1]), s(vec![1]));
459        assert_ne!(s(vec![1]), s(vec![2]));
460    }
461}
nuvec_core/primitive.rs

nuvec_core/
primitive.rs
