ferrum_interfaces/vnext/operation/weight_contract/
hadamard.rs1use std::num::NonZeroU32;
2
3use serde::{Deserialize, Serialize};
4
5use super::{PhysicalWeightComponentBinding, ResolvedWeightBinding, WeightComponentRole};
6use crate::vnext::{ResolvedStorageComponent, VNextError};
7
8#[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#[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 BeforeMatmul {
46 input_permutation: Option<GroupedFeatureTranspose>,
47 },
48 AfterEmbeddingLookup,
51}
52
53#[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
96pub(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}