Skip to main content

ferrum_interfaces/vnext/operation/weight_contract/
hadamard.rs

1use std::num::NonZeroU32;
2
3use serde::{Deserialize, Serialize};
4
5use super::{PhysicalWeightComponentBinding, ResolvedWeightBinding, WeightComponentRole};
6use crate::vnext::{ResolvedStorageComponent, VNextError};
7
8/// Transpose the two outer feature axes before applying input signs.
9/// Dimensions are fastest-axis-first: `[inner, first, second]` becomes
10/// `[inner, second, first]`. The complete last tensor axis is permuted;
11/// token and batch coordinates are unchanged.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13#[serde(deny_unknown_fields)]
14pub struct GroupedFeatureTranspose {
15    pub inner_extent: u64,
16    pub first_outer_extent: u64,
17    pub second_outer_extent: u64,
18}
19
20impl GroupedFeatureTranspose {
21    pub fn width(&self) -> Option<u64> {
22        if self.inner_extent == 0 || self.first_outer_extent == 0 || self.second_outer_extent == 0 {
23            return None;
24        }
25        self.inner_extent
26            .checked_mul(self.first_outer_extent)?
27            .checked_mul(self.second_outer_extent)
28    }
29}
30
31/// Signs span the full last-axis width, not one repeated Hadamard block.
32/// Explicit signs bind immutable, exact-contiguous `TransformSigns` F32
33/// components whose source has validated every value as either -1 or +1.
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum HadamardSigns {
37    Identity,
38    Explicit(PhysicalWeightComponentBinding),
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
42#[serde(rename_all = "snake_case", deny_unknown_fields)]
43pub enum HadamardApplication {
44    /// Consume `H(signs * permutation(input))` before the packed matrix.
45    BeforeMatmul {
46        input_permutation: Option<GroupedFeatureTranspose>,
47    },
48    /// Restore a looked-up latent row with `signs * H(row)`.
49    /// A projection-input permutation has no meaning in this direction.
50    AfterEmbeddingLookup,
51}
52
53/// Normalized, blockwise Sylvester Walsh-Hadamard on the complete last axis.
54///
55/// For block width B, H[i,j] = (-1)^popcount(i & j) / sqrt(B).
56/// Providers perform butterflies, normalization and intermediate storage in
57/// F32; no F16 narrowing is allowed between this transform and its matrix
58/// projection. Embedding output is converted only after H and output signs.
59/// The logical operation's output dtype remains its declared dtype.
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61#[serde(deny_unknown_fields)]
62pub struct HadamardTransformSpec {
63    pub block_size: NonZeroU32,
64    pub signs: HadamardSigns,
65    pub application: HadamardApplication,
66}
67
68impl HadamardTransformSpec {
69    pub fn validate(&self, width: u64) -> Result<(), VNextError> {
70        let block = self.block_size.get();
71        if !block.is_power_of_two() || width == 0 || !width.is_multiple_of(u64::from(block)) {
72            return Err(invalid(
73                "block size must be a power of two dividing the nonzero last axis",
74            ));
75        }
76        if let HadamardApplication::BeforeMatmul {
77            input_permutation: Some(permutation),
78        } = &self.application
79        {
80            if permutation.width() != Some(width) {
81                return Err(invalid(
82                    "input permutation shape differs from the full last axis",
83                ));
84            }
85        }
86        Ok(())
87    }
88}
89
90fn invalid(reason: &str) -> VNextError {
91    VNextError::InvalidExecutionPlan {
92        reason: format!("invalid Hadamard transform: {reason}"),
93    }
94}
95
96/// Only an entire immutable auxiliary component can alias another weight.
97/// Callers separately require read-only weight bindings. Logical validation
98/// proves that this role is bound by Hadamard, with exact F32 vector storage.
99pub(crate) fn same_shared_transform_sign_component(
100    left: &ResolvedWeightBinding,
101    left_storage: &ResolvedStorageComponent,
102    right: &ResolvedWeightBinding,
103    right_storage: &ResolvedStorageComponent,
104) -> bool {
105    if left_storage != right_storage {
106        return false;
107    }
108    left.components().iter().any(|component| {
109        Some(component.component_id()) == left_storage.component_id()
110            && component.role() == WeightComponentRole::TransformSigns
111            && right.components().iter().any(|other| other == component)
112    })
113}