rutensor/
einsum.rs

1use std::collections::BTreeMap;
2use crate::{ComputeType, Error, Mode, OperandDescriptor, OperationDescriptor, Plan, Result, TensorDescriptor};
3use crate::plan::merge_extent;
4use ruda_core::tensor::DType;
5use ruda_core::device::Device;
6use ruda_kernel::{dsl::Runtime, tensor::RudaTensor};
7
8#[derive(Clone, Debug)]
9enum Token { Label(Mode), Ellipsis }
10
11fn tokens(text: &str) -> Result<Vec<Token>> {
12    let mut result = Vec::new();
13    let bytes = text.as_bytes();
14    let mut index = 0;
15    let mut ellipsis = false;
16    while index < bytes.len() {
17        if bytes[index].is_ascii_alphabetic() {
18            result.push(Token::Label(bytes[index] as Mode));
19            index += 1;
20        } else if bytes[index..].starts_with(b"...") && !ellipsis {
21            result.push(Token::Ellipsis);
22            ellipsis = true;
23            index += 3;
24        } else {
25            return Err(Error::InvalidExpression("use ASCII letters and at most one ellipsis per operand".into()));
26        }
27    }
28    Ok(result)
29}
30
31fn expand(tokens: &[Token], ellipsis_modes: &[Mode]) -> Vec<Mode> {
32    let mut modes = Vec::new();
33    for token in tokens {
34        match token {
35            Token::Label(mode) => modes.push(*mode),
36            Token::Ellipsis => modes.extend_from_slice(ellipsis_modes),
37        }
38    }
39    modes
40}
41
42pub(crate) fn output_dtype(inputs: &[OperandDescriptor]) -> DType {
43    let first = inputs[0].tensor().dtype();
44    if inputs.iter().all(|input| input.tensor().dtype() == first) { return first; }
45    if inputs.iter().any(|input| input.tensor().dtype() == DType::F64) { DType::F64 } else { DType::F32 }
46}
47
48pub(crate) fn infer_output(inputs: &[OperandDescriptor], modes: &[Mode], dtype: DType) -> Result<TensorDescriptor> {
49    let mut extents = BTreeMap::new();
50    for input in inputs {
51        for (&mode, &extent) in input.modes().iter().zip(input.tensor().extents()) {
52            merge_extent(&mut extents, mode, extent)?;
53        }
54    }
55    let shape = modes.iter().map(|mode| extents.get(mode).copied().ok_or_else(||
56        Error::InvalidOperation(format!("output mode {mode} is absent from all inputs"))))
57        .collect::<Result<Vec<_>>>()?;
58    TensorDescriptor::contiguous(&shape, dtype)
59}
60
61/// A reusable einsum expression with fixed input shapes, strides and storage types.
62#[derive(Clone, Debug)]
63pub struct EinsumPlan {
64    expression: String,
65    plan: Plan,
66    path: Option<crate::path::ContractionPath>,
67}
68
69impl EinsumPlan {
70    pub fn new(expression: &str, inputs: &[TensorDescriptor]) -> Result<Self> {
71        Self::build(expression, inputs, None)
72    }
73
74    pub fn with_options(
75        expression: &str, inputs: &[TensorDescriptor], output_dtype: DType, compute: ComputeType,
76    ) -> Result<Self> {
77        Self::build(expression, inputs, Some((output_dtype, compute)))
78    }
79
80    fn build(expression: &str, inputs: &[TensorDescriptor], options: Option<(DType, ComputeType)>) -> Result<Self> {
81        if inputs.is_empty() {
82            return Err(Error::InvalidExpression("at least one tensor is required".into()));
83        }
84        let expression: String = expression.chars().filter(|c| !c.is_ascii_whitespace()).collect();
85        let mut sides = expression.split("->");
86        let lhs = sides.next().unwrap_or_default();
87        let rhs = sides.next();
88        if sides.next().is_some() {
89            return Err(Error::InvalidExpression("more than one output arrow".into()));
90        }
91        let input_text: Vec<_> = lhs.split(',').collect();
92        if input_text.len() != inputs.len() {
93            return Err(Error::InputCount { expected: input_text.len(), actual: inputs.len() });
94        }
95        let token_lists: Vec<_> = input_text.iter().map(|text| tokens(text)).collect::<Result<_>>()?;
96        let mut widths = Vec::new();
97        let mut counts = BTreeMap::<Mode, usize>::new();
98        for (tokens, descriptor) in token_lists.iter().zip(inputs) {
99            let explicit = tokens.iter().filter(|token| matches!(token, Token::Label(_))).count();
100            let has_ellipsis = tokens.iter().any(|token| matches!(token, Token::Ellipsis));
101            if explicit > descriptor.rank() || (!has_ellipsis && explicit != descriptor.rank()) {
102                return Err(Error::InvalidExpression("operand labels do not match its tensor rank".into()));
103            }
104            widths.push(descriptor.rank() - explicit);
105            for token in tokens {
106                if let Token::Label(mode) = token { *counts.entry(*mode).or_default() += 1; }
107            }
108        }
109        let width = widths.iter().copied().max().unwrap_or(0);
110        let ellipsis_modes = (0..width).map(|axis| {
111            let axis = i32::try_from(axis).map_err(|_| Error::Overflow)?;
112            i32::MIN.checked_add(axis).ok_or(Error::Overflow)
113        }).collect::<Result<Vec<_>>>()?;
114        let mut operands = Vec::new();
115        for ((tokens, descriptor), local_width) in token_lists.iter().zip(inputs).zip(widths) {
116            let modes = expand(tokens, &ellipsis_modes[width - local_width..]);
117            operands.push(OperandDescriptor::new(descriptor.clone(), &modes)?);
118        }
119        let output_modes = if let Some(rhs) = rhs {
120            expand(&tokens(rhs)?, &ellipsis_modes)
121        } else {
122            let mut modes = ellipsis_modes;
123            modes.extend(counts.iter().filter_map(|(&mode, &count)| (count == 1).then_some(mode)));
124            modes
125        };
126        let (dtype, compute) = options.unwrap_or_else(||
127            (output_dtype(&operands), ComputeType::for_operands(&operands)));
128        let output = infer_output(&operands, &output_modes, dtype)?;
129        let operation = OperationDescriptor::sum_product(operands, None, output, &output_modes, compute)?;
130        Ok(Self { expression, plan: Plan::new(operation)?, path: None })
131    }
132
133    pub fn expression(&self) -> &str { &self.expression }
134    pub fn plan(&self) -> &Plan { &self.plan }
135    pub fn output(&self) -> &TensorDescriptor { self.plan.operation().output() }
136
137    /// Choose a minimum-product binary contraction tree by exhaustive subset dynamic programming.
138    /// Search is exponential in operand count. Reassociation changes floating-point operation order;
139    /// intermediates use the declared compute dtype and only the final output is cast to its dtype.
140    pub fn optimize(mut self) -> Result<Self> {
141        if self.plan.operation().inputs().len() > 1 {
142            self.path = Some(crate::path::ContractionPath::new(self.plan.operation())?);
143        }
144        Ok(self)
145    }
146
147    pub fn planned_scalar_products(&self) -> Option<u128> {
148        self.path.as_ref().map(|path| path.products)
149    }
150
151    pub fn execute<R: Runtime>(&self, inputs: &[&RudaTensor<R>]) -> Result<RudaTensor<R>> {
152        if let Some(path) = &self.path {
153            self.plan.validate_inputs(inputs, &[1.0])?;
154            path.execute(inputs)
155        } else {
156            self.plan.execute(inputs, &[1.0])
157        }
158    }
159
160    pub fn execute_into<R: Runtime>(&self, inputs: &[&RudaTensor<R>], output: RudaTensor<R>) -> Result<RudaTensor<R>> {
161        if let Some(path) = &self.path {
162            self.plan.validate_inputs(inputs, &[1.0])?;
163            if !self.output().matches(&output) { return Err(Error::OutputMismatch); }
164            self.output().check_buffer(&output)?;
165            if inputs[0].device.to_id() != output.device.to_id() { return Err(Error::DeviceMismatch); }
166            if !output.can_mut() { return Err(Error::SharedOutput); }
167            path.execute_into(inputs, output)
168        } else {
169            self.plan.execute_into(inputs, output, &[1.0])
170        }
171    }
172}
173
174/// Evaluate an einsum expression on device tensors without reading values back to the host.
175///
176/// Ellipses are right-aligned for broadcasting. Repeated input labels select diagonals.
177/// Explicit output labels select and order free axes; all other axes are summed.
178pub fn einsum<R: Runtime>(expression: &str, inputs: &[&RudaTensor<R>]) -> Result<RudaTensor<R>> {
179    let descriptors = inputs.iter().map(|input| TensorDescriptor::from_tensor(*input)).collect::<Result<Vec<_>>>()?;
180    EinsumPlan::new(expression, &descriptors)?.execute(inputs)
181}
182
183pub fn einsum_with_options<R: Runtime>(
184    expression: &str, inputs: &[&RudaTensor<R>], output_dtype: DType, compute: ComputeType,
185) -> Result<RudaTensor<R>> {
186    let descriptors = inputs.iter().map(|input| TensorDescriptor::from_tensor(*input)).collect::<Result<Vec<_>>>()?;
187    EinsumPlan::with_options(expression, &descriptors, output_dtype, compute)?.execute(inputs)
188}
rutensor/einsum.rs

rutensor/
einsum.rs
