tenferro_bf16_proof/
einsum.rs

1//! A bfloat16 contraction through the runtime's extension module.
2//!
3//! #1793's precision table lists a bf16 CPU einsum "with documented f32 internal
4//! calculation/accumulation", and asks for a contraction whose result distinguishes f32
5//! accumulation from repeated bfloat16 rounding, compared against an independent reference with the
6//! specified output rounding. This module supplies that: the operands are widened to `f32`, the
7//! contraction accumulates there, and the result is rounded to bfloat16 once.
8//!
9//! What it does not do is inherit a wider pattern surface. The operation is the pairwise
10//! contraction the table's example needs, so a trace, a label repeated in the output, or more than
11//! two operands is refused with a typed error rather than folded, which keeps this module's body to
12//! the accumulation contract the row is about.
13
14use std::any::Any;
15use std::hash::Hasher;
16use std::marker::PhantomData;
17use std::sync::Arc;
18
19use tenferro_ops::ext_op::{ExtensionAliasDeclaration, ExtensionEffectDeclaration};
20use tenferro_runtime::extension::{ExtensionOp, ExtensionShapeContext, SymDim};
21use tenferro_runtime::{
22    CoreCapabilityKind, EngineId, ErasedExecutionContext, ErrorPhase, ExecutionContextIdentity,
23    ExtensionCacheStore, ExtensionEngine, ExtensionModule, ExtensionModuleId,
24    ExtensionModuleRegistrar, ExtensionPlanningConfig, ExtensionPrepareRequest, PrepareCapability,
25    PrepareError, PreparedOperation, PreparedOperationBinding, PreparedOperationExecutor,
26    PreparedOperationPlan, ProviderContractError, RuntimeConfigError, SpecializationProjection,
27};
28use tenferro_tensor::{DType, DynRank, Engine, Host, Tensor, TensorRead, TypedTensor};
29
30use crate::Bf16;
31
32/// The family the bfloat16 contraction belongs to.
33pub const BF16_EINSUM_FAMILY: &str = "tenferro-bf16-proof.einsum.v1";
34
35/// The canonical identity of the bfloat16 scalar a program declares.
36pub const BF16_SCALAR_IDENTITY: &str = "tenferro-bf16-proof.bf16.v1";
37
38/// A pairwise contraction in bfloat16.
39///
40/// The labels are the ones an ordinary einsum subscript string names, so `"ik,kj->ij"` is
41/// `(&[0, 1], &[1, 2], &[0, 2])`. The operands are widened to `f32`, the contraction accumulates
42/// there, and the result is rounded to bfloat16 once.
43///
44/// # Examples
45///
46/// ```rust
47/// use tenferro_bf16_proof::einsum::Bf16Einsum;
48///
49/// let op = Bf16Einsum::new(&[0, 1], &[1, 2], &[0, 2]).expect("a contraction");
50/// assert_eq!(op.labels(), (&[0, 1][..], &[1, 2][..], &[0, 2][..]));
51/// ```
52#[derive(Clone, Debug, PartialEq, Eq)]
53pub struct Bf16Einsum {
54    lhs: Vec<u32>,
55    rhs: Vec<u32>,
56    out: Vec<u32>,
57}
58
59impl Bf16Einsum {
60    /// Build the pairwise contraction `lhs,rhs->out`.
61    ///
62    /// # Errors
63    ///
64    /// Returns a validation error (an invalid `pattern` argument, built with
65    /// [`tenferro_tensor::Error::invalid_argument`]) when the pattern is not a pairwise contraction: an operand
66    /// carries no label, a label repeats inside one operand, an output label repeats, or an output
67    /// label appears in no operand.
68    ///
69    /// # Examples
70    ///
71    /// ```rust
72    /// use tenferro_bf16_proof::einsum::Bf16Einsum;
73    ///
74    /// assert!(Bf16Einsum::new(&[0, 1], &[1, 2], &[0, 2]).is_ok());
75    /// assert!(Bf16Einsum::new(&[], &[1, 2], &[0, 2]).is_err());
76    /// // A label repeated in the output is refused, exactly as one repeated inside an operand is.
77    /// assert!(Bf16Einsum::new(&[0, 1], &[1, 2], &[0, 0]).is_err());
78    /// ```
79    pub fn new(lhs: &[u32], rhs: &[u32], out: &[u32]) -> tenferro_runtime::Result<Self> {
80        let invalid = |message: &str| {
81            tenferro_runtime::Error::from(tenferro_tensor::Error::invalid_argument(
82                "bf16_einsum",
83                "pattern",
84                message,
85            ))
86        };
87        if lhs.is_empty() || rhs.is_empty() {
88            return Err(invalid("an operand must carry at least one label"));
89        }
90        // A label repeated inside an operand is a trace, and one repeated in the output would
91        // build a duplicate output axis; neither is a pairwise contraction.
92        for labels in [lhs, rhs, out] {
93            let mut seen = labels.to_vec();
94            seen.sort_unstable();
95            seen.dedup();
96            if seen.len() != labels.len() {
97                return Err(invalid(
98                    "a label repeats within one operand or the output, which this pairwise contraction does not support",
99                ));
100            }
101        }
102        for label in out {
103            if !lhs.contains(label) && !rhs.contains(label) {
104                return Err(invalid(
105                    "an output label must appear in at least one operand",
106                ));
107            }
108        }
109        Ok(Self {
110            lhs: lhs.to_vec(),
111            rhs: rhs.to_vec(),
112            out: out.to_vec(),
113        })
114    }
115
116    /// The pattern's three label lists.
117    ///
118    /// # Examples
119    ///
120    /// ```rust
121    /// use tenferro_bf16_proof::einsum::Bf16Einsum;
122    ///
123    /// let op = Bf16Einsum::new(&[0, 1], &[1, 2], &[0, 2]).expect("a contraction");
124    /// assert_eq!(op.labels(), (&[0, 1][..], &[1, 2][..], &[0, 2][..]));
125    /// ```
126    #[must_use]
127    pub fn labels(&self) -> (&[u32], &[u32], &[u32]) {
128        (&self.lhs, &self.rhs, &self.out)
129    }
130}
131
132impl ExtensionOp for Bf16Einsum {
133    fn family_id(&self) -> &'static str {
134        BF16_EINSUM_FAMILY
135    }
136
137    fn payload_hash(&self, hasher: &mut dyn Hasher) {
138        for labels in [&self.lhs, &self.rhs, &self.out] {
139            hasher.write_usize(labels.len());
140            for label in labels {
141                hasher.write_u32(*label);
142            }
143        }
144    }
145
146    fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
147        other
148            .as_any()
149            .downcast_ref::<Self>()
150            .is_some_and(|other| other == self)
151    }
152
153    fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
154        Arc::new(self.clone())
155    }
156
157    fn as_any(&self) -> &dyn Any {
158        self
159    }
160
161    fn input_count(&self) -> usize {
162        2
163    }
164
165    fn output_count(&self) -> usize {
166        1
167    }
168
169    fn semantic_effects(&self) -> ExtensionEffectDeclaration<'_> {
170        ExtensionEffectDeclaration::Declared(&[])
171    }
172
173    fn semantic_aliases(&self) -> ExtensionAliasDeclaration<'_> {
174        ExtensionAliasDeclaration::AllFresh
175    }
176
177    fn scalar_identity(&self) -> Option<&'static str> {
178        Some(BF16_SCALAR_IDENTITY)
179    }
180
181    fn infer_output_meta(
182        &self,
183        ctx: &mut ExtensionShapeContext<'_>,
184    ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
185        let dtype = ctx.input_dtype(0)?;
186        if !matches!(dtype, DType::External(_)) {
187            return Err(tenferro_tensor::Error::unsupported_dtype(
188                "bf16_einsum",
189                dtype,
190                "bf16_einsum takes an externally defined scalar",
191            ));
192        }
193        if ctx.input_dtype(1)? != dtype {
194            return Err(tenferro_tensor::Error::invalid_argument(
195                "bf16_einsum",
196                "inputs",
197                "both operands must carry the same scalar",
198            ));
199        }
200        let lhs = ctx.input_shape(0)?;
201        let rhs = ctx.input_shape(1)?;
202        if lhs.len() != self.lhs.len() || rhs.len() != self.rhs.len() {
203            return Err(tenferro_tensor::Error::rank_mismatch(
204                "bf16_einsum",
205                self.lhs.len().max(self.rhs.len()),
206                lhs.len().min(rhs.len()),
207            ));
208        }
209        let mut out_shape = Vec::with_capacity(self.out.len());
210        for label in &self.out {
211            let extent = self
212                .lhs
213                .iter()
214                .position(|candidate| candidate == label)
215                .map(|axis| lhs[axis].clone())
216                .or_else(|| {
217                    self.rhs
218                        .iter()
219                        .position(|candidate| candidate == label)
220                        .map(|axis| rhs[axis].clone())
221                })
222                .ok_or_else(|| {
223                    tenferro_tensor::Error::invalid_argument(
224                        "bf16_einsum",
225                        "pattern",
226                        "an output label must appear in at least one operand",
227                    )
228                })?;
229            out_shape.push(extent);
230        }
231        Ok(vec![(dtype, out_shape)])
232    }
233}
234
235/// The number of elements a shape describes.
236fn element_count(shape: &[usize]) -> usize {
237    shape.iter().product()
238}
239
240/// Advance a column-major odometer by one.
241fn advance(index: &mut [usize], shape: &[usize]) {
242    for axis in 0..shape.len() {
243        index[axis] += 1;
244        if index[axis] < shape[axis] {
245            return;
246        }
247        index[axis] = 0;
248    }
249}
250
251/// The column-major offset an operand's labels select at one output and contracted index.
252fn offset_for(
253    input_labels: &[u32],
254    input_shape: &[usize],
255    out_labels: &[u32],
256    out_index: &[usize],
257    summed_labels: &[u32],
258    summed_index: &[usize],
259) -> usize {
260    let mut offset = 0usize;
261    let mut stride = 1usize;
262    for (axis, label) in input_labels.iter().enumerate() {
263        let position = out_labels
264            .iter()
265            .position(|candidate| candidate == label)
266            .map(|index| out_index[index])
267            .or_else(|| {
268                summed_labels
269                    .iter()
270                    .position(|candidate| candidate == label)
271                    .map(|index| summed_index[index])
272            })
273            .unwrap_or(0);
274        offset += position * stride;
275        stride *= input_shape[axis];
276    }
277    offset
278}
279
280/// The stored values of a bfloat16 tensor, widened to `f32`.
281fn values_of(op: &'static str, tensor: &Tensor) -> tenferro_runtime::Result<Vec<f32>> {
282    match tensor.external_payload() {
283        Some(payload) => payload
284            .downcast_ref::<Bf16>()
285            .map(|stored| {
286                stored
287                    .as_slice()
288                    .iter()
289                    .map(|value| value.to_f32())
290                    .collect()
291            })
292            .ok_or_else(|| {
293                tenferro_runtime::Error::from(tenferro_tensor::Error::invalid_argument(
294                    op,
295                    "input",
296                    "the operand does not carry a bfloat16 payload",
297                ))
298            }),
299        None => Err(tenferro_runtime::Error::from(
300            tenferro_tensor::Error::unsupported_dtype(
301                op,
302                tensor.dtype(),
303                "bf16_einsum takes an externally defined bfloat16 scalar",
304            ),
305        )),
306    }
307}
308
309/// Contract two operands, accumulating in `f32` and rounding once.
310fn contract(op: &dyn ExtensionOp, inputs: &[&Tensor]) -> tenferro_runtime::Result<Vec<Tensor>> {
311    let name = "bf16_einsum";
312    let Some(contraction) = op.as_any().downcast_ref::<Bf16Einsum>() else {
313        return Err(tenferro_runtime::Error::from(
314            tenferro_tensor::Error::invalid_argument(
315                name,
316                "payload",
317                "the operation is not a bfloat16 contraction",
318            ),
319        ));
320    };
321    if inputs.len() != 2 {
322        return Err(tenferro_runtime::Error::from(
323            tenferro_tensor::Error::invalid_argument(
324                name,
325                "input",
326                "a bf16 contraction takes two operands",
327            ),
328        ));
329    }
330    let (lhs_labels, rhs_labels, out_labels) = contraction.labels();
331    let lhs_values = values_of(name, inputs[0])?;
332    let rhs_values = values_of(name, inputs[1])?;
333    let lhs_shape = inputs[0].shape().to_vec();
334    let rhs_shape = inputs[1].shape().to_vec();
335
336    let mut extents: Vec<(u32, usize)> = Vec::new();
337    for (labels, shape) in [(lhs_labels, &lhs_shape), (rhs_labels, &rhs_shape)] {
338        for (axis, label) in labels.iter().enumerate() {
339            match extents.iter().find(|(existing, _)| existing == label) {
340                Some((_, existing)) if *existing != shape[axis] => {
341                    return Err(tenferro_runtime::Error::from(
342                        tenferro_tensor::Error::invalid_argument(
343                            name,
344                            "inputs",
345                            "the operands disagree on the extent of a shared label",
346                        ),
347                    ))
348                }
349                Some(_) => {}
350                None => extents.push((*label, shape[axis])),
351            }
352        }
353    }
354    let extent_of = |label: u32| {
355        extents
356            .iter()
357            .find(|(existing, _)| *existing == label)
358            .map(|(_, extent)| *extent)
359            .unwrap_or(1)
360    };
361    let out_shape: Vec<usize> = out_labels.iter().map(|label| extent_of(*label)).collect();
362    let mut summed_labels: Vec<u32> = Vec::new();
363    for label in lhs_labels.iter().chain(rhs_labels.iter()) {
364        if !out_labels.contains(label) && !summed_labels.contains(label) {
365            summed_labels.push(*label);
366        }
367    }
368    let summed_shape: Vec<usize> = summed_labels
369        .iter()
370        .map(|label| extent_of(*label))
371        .collect();
372
373    let out_count = element_count(&out_shape);
374    let summed_count = element_count(&summed_shape);
375    let mut accumulated = vec![0.0_f32; out_count];
376    let mut out_index = vec![0usize; out_shape.len()];
377    let mut summed_index = vec![0usize; summed_shape.len()];
378    for slot in accumulated.iter_mut() {
379        for value in summed_index.iter_mut() {
380            *value = 0;
381        }
382        let mut total = 0.0_f32;
383        for _ in 0..summed_count {
384            let lhs_offset = offset_for(
385                lhs_labels,
386                &lhs_shape,
387                out_labels,
388                &out_index,
389                &summed_labels,
390                &summed_index,
391            );
392            let rhs_offset = offset_for(
393                rhs_labels,
394                &rhs_shape,
395                out_labels,
396                &out_index,
397                &summed_labels,
398                &summed_index,
399            );
400            total += lhs_values[lhs_offset] * rhs_values[rhs_offset];
401            advance(&mut summed_index, &summed_shape);
402        }
403        *slot = total;
404        advance(&mut out_index, &out_shape);
405    }
406
407    // One rounding, after the accumulation rather than at every step.
408    let rounded: Vec<Bf16> = accumulated.into_iter().map(Bf16::from_f32).collect();
409    let tensor = TypedTensor::<_, DynRank, Host>::from_host_vec_col_major(out_shape, rounded)
410        .map_err(tenferro_runtime::Error::from)?;
411    Ok(vec![Tensor::external(
412        tenferro_tensor::ErasedHostTensor::new(tensor),
413    )])
414}
415
416/// The engine the bfloat16 contraction is registered under.
417#[derive(Debug)]
418struct Bf16EinsumEngine<B: std::fmt::Debug + Send + Sync> {
419    family_id: &'static str,
420    engine_id: EngineId,
421    _backend: PhantomData<B>,
422}
423
424impl<B: Engine + std::fmt::Debug + Send + Sync + 'static> ExtensionEngine for Bf16EinsumEngine<B> {
425    fn family_id(&self) -> &'static str {
426        self.family_id
427    }
428
429    fn engine_id(&self) -> &EngineId {
430        &self.engine_id
431    }
432
433    fn context_identity(&self) -> ExecutionContextIdentity {
434        ExecutionContextIdentity::of::<B>()
435    }
436
437    fn prepare(
438        &self,
439        request: ExtensionPrepareRequest<'_>,
440    ) -> Result<PrepareCapability, PrepareError> {
441        if request.operation().family_id() != self.family_id {
442            return Err(PrepareError::ProviderContract {
443                source: ProviderContractError::WrongOperationFamily {
444                    expected: CoreCapabilityKind::Elementwise,
445                    operation: self.family_id,
446                },
447            });
448        }
449        let prepared = Arc::new(Bf16EinsumPrepared::<B> {
450            binding: request.binding().clone(),
451            specialization: request.specialization().clone(),
452            op: request.operation().clone_arc(),
453            _backend: PhantomData,
454        });
455        Ok(PrepareCapability::Prepared(
456            PreparedOperationPlan::executable(prepared.clone(), prepared),
457        ))
458    }
459}
460
461/// The planning config the runtime keys the bfloat16 family under.
462///
463/// The runtime keeps one planning config per engine identity, so the family states its identity
464/// here rather than leaving the engine without one.
465#[derive(Debug)]
466struct Bf16EinsumPlanning {
467    family_id: &'static str,
468}
469
470impl ExtensionPlanningConfig for Bf16EinsumPlanning {
471    fn family_id(&self) -> &'static str {
472        self.family_id
473    }
474
475    fn as_any(&self) -> &dyn Any {
476        self
477    }
478
479    fn payload_hash(&self, _state: &mut dyn Hasher) {}
480
481    fn payload_eq(&self, other: &dyn ExtensionPlanningConfig) -> bool {
482        other.family_id() == self.family_id
483    }
484
485    fn retained_bytes(&self) -> usize {
486        0
487    }
488}
489
490/// The prepared bfloat16 contraction.
491#[derive(Debug)]
492struct Bf16EinsumPrepared<B: std::fmt::Debug + Send + Sync> {
493    binding: PreparedOperationBinding,
494    specialization: SpecializationProjection,
495    op: Arc<dyn ExtensionOp>,
496    _backend: PhantomData<B>,
497}
498
499impl<B: Engine + std::fmt::Debug + Send + Sync + 'static> PreparedOperation
500    for Bf16EinsumPrepared<B>
501{
502    fn binding(&self) -> &PreparedOperationBinding {
503        &self.binding
504    }
505
506    fn specialization(&self) -> &SpecializationProjection {
507        &self.specialization
508    }
509
510    fn retained_bytes(&self) -> usize {
511        0
512    }
513}
514
515impl<B: Engine + std::fmt::Debug + Send + Sync + 'static> PreparedOperationExecutor
516    for Bf16EinsumPrepared<B>
517{
518    fn execute(
519        &self,
520        context: &mut ErasedExecutionContext<'_>,
521        extension_caches: &mut ExtensionCacheStore,
522        inputs: &[TensorRead<'_>],
523    ) -> tenferro_runtime::Result<Vec<Tensor>> {
524        // The runtime passes the caller's engine; the body needs contiguous inputs.
525        let engine = context
526            .downcast_mut::<B>(self.binding.context_identity())
527            .map_err(|source| {
528                tenferro_runtime::Error::runtime_state_source(
529                    "extension",
530                    ErrorPhase::Execution,
531                    source,
532                )
533            })?;
534        let _ = extension_caches;
535        let materialized = inputs
536            .iter()
537            .cloned()
538            .map(|input| engine.to_contiguous_read(input))
539            .collect::<tenferro_tensor::Result<Vec<Tensor>>>()
540            .map_err(tenferro_runtime::Error::from)?;
541        let borrowed: Vec<&Tensor> = materialized.iter().collect();
542        contract(self.op.as_ref(), &borrowed)
543    }
544}
545
546/// The module an application installs to run the bfloat16 contraction.
547#[derive(Debug)]
548struct Bf16EinsumModule<B: std::fmt::Debug + Send + Sync> {
549    module_id: ExtensionModuleId,
550    engine_id: EngineId,
551    _backend: PhantomData<B>,
552}
553
554impl<B: Engine + std::fmt::Debug + Send + Sync + 'static> ExtensionModule for Bf16EinsumModule<B> {
555    fn module_id(&self) -> &ExtensionModuleId {
556        &self.module_id
557    }
558
559    fn configure(
560        &self,
561        registrar: &mut ExtensionModuleRegistrar<'_>,
562    ) -> Result<(), tenferro_runtime::ExtensionModuleError> {
563        registrar.register_engine(Arc::new(Bf16EinsumEngine::<B> {
564            family_id: BF16_EINSUM_FAMILY,
565            engine_id: self.engine_id.clone(),
566            _backend: PhantomData,
567        }))?;
568        registrar.register_planning_config(
569            self.engine_id.clone(),
570            Arc::new(Bf16EinsumPlanning {
571                family_id: BF16_EINSUM_FAMILY,
572            }),
573        )?;
574        Ok(())
575    }
576}
577
578/// The module for the CPU backend's engine identity.
579///
580/// # Errors
581///
582/// Returns [`RuntimeConfigError`] when the module identity cannot be built or the backend has no
583/// engine identity.
584///
585/// # Examples
586///
587/// ```rust
588/// use tenferro_bf16_proof::einsum::module;
589///
590/// assert!(module().is_ok());
591/// ```
592pub fn module() -> Result<Arc<dyn ExtensionModule>, RuntimeConfigError> {
593    Ok(Arc::new(Bf16EinsumModule::<tenferro_cpu::CpuEngine> {
594        module_id: ExtensionModuleId::new("tenferro-bf16-proof.module")?,
595        engine_id: tenferro_cpu::runtime_engine_id()?,
596        _backend: PhantomData,
597    }))
598}
tenferro_bf16_proof/einsum.rs

tenferro_bf16_proof/
einsum.rs
