ferrum_quantization/
dense.rs

1//! Dense linear projection — the baseline, uses `B::gemm` directly.
2//!
3//! Supports an optional learnable bias (Bert / Clip / many encoder models).
4//! When `bias` is set, `forward` lowers to `gemm + add_bias` (one extra
5//! dispatch on GPU backends, still part of the current command buffer).
6
7use ferrum_kernels::backend::Backend;
8use ferrum_kernels::LinearMetadata;
9
10use crate::traits::Linear;
11
12/// Dense linear projection.
13///
14/// Holds a single weight matrix laid out row-major as `[out_features, in_features]`.
15/// `forward` delegates to `B::gemm` plus (optional) `B::add_bias`.
16pub struct DenseLinear<B: Backend> {
17    weight: B::Buffer,
18    bias: Option<B::Buffer>,
19    in_features: usize,
20    out_features: usize,
21    metadata: LinearMetadata,
22}
23
24impl<B: Backend> DenseLinear<B> {
25    /// Build a weight-only dense projection (no bias).
26    pub fn from_rows(weight_row_major: &[f32], out_features: usize, in_features: usize) -> Self {
27        debug_assert_eq!(
28            weight_row_major.len(),
29            out_features * in_features,
30            "DenseLinear weight length mismatch"
31        );
32        let weight = B::from_slice(weight_row_major);
33        Self {
34            weight,
35            bias: None,
36            in_features,
37            out_features,
38            metadata: LinearMetadata::default(),
39        }
40    }
41
42    /// Build a dense projection with a bias vector of length `out_features`.
43    pub fn from_rows_with_bias(
44        weight_row_major: &[f32],
45        bias: &[f32],
46        out_features: usize,
47        in_features: usize,
48    ) -> Self {
49        debug_assert_eq!(bias.len(), out_features, "DenseLinear bias length mismatch");
50        Self {
51            weight: B::from_slice(weight_row_major),
52            bias: Some(B::from_slice(bias)),
53            in_features,
54            out_features,
55            metadata: LinearMetadata::default(),
56        }
57    }
58
59    /// Construct by moving already-allocated `Backend` buffers.
60    pub fn from_buffer(weight: B::Buffer, out_features: usize, in_features: usize) -> Self {
61        Self {
62            weight,
63            bias: None,
64            in_features,
65            out_features,
66            metadata: LinearMetadata::default(),
67        }
68    }
69
70    pub fn with_bias(mut self, bias: B::Buffer) -> Self {
71        self.bias = Some(bias);
72        self
73    }
74
75    pub fn with_metadata(mut self, metadata: LinearMetadata) -> Self {
76        self.metadata = metadata;
77        self
78    }
79
80    pub fn weight(&self) -> &B::Buffer {
81        &self.weight
82    }
83
84    pub fn bias(&self) -> Option<&B::Buffer> {
85        self.bias.as_ref()
86    }
87}
88
89impl<B: Backend> Linear<B> for DenseLinear<B> {
90    fn in_features(&self) -> usize {
91        self.in_features
92    }
93
94    fn out_features(&self) -> usize {
95        self.out_features
96    }
97
98    fn metadata(&self) -> LinearMetadata {
99        self.metadata
100    }
101
102    fn forward(&self, ctx: &mut B::Context, input: &B::Buffer, out: &mut B::Buffer, m: usize) {
103        B::gemm(
104            ctx,
105            input,
106            &self.weight,
107            out,
108            m,
109            self.out_features,
110            self.in_features,
111        );
112        if let Some(bias) = &self.bias {
113            B::add_bias(ctx, out, bias, m, self.out_features);
114        }
115    }
116}
ferrum_quantization/dense.rs

ferrum_quantization/
dense.rs
