vyre_libs/nn/linear/inner/
builder.rs

1//! Linear builder + the canonical `linear()` Cat-A constructor.
2
3use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
4
5use crate::{
6    builder::{check_tensors, BuildOptions},
7    linear_algebra_substrate::MatmulBias,
8    region::{tag_program, wrap_anonymous},
9    tensor_ref::{TensorRef, TensorRefError},
10};
11
12use super::tiled::{linear_tiled, LINEAR_TILED_MIN_WORK, LINEAR_TILED_TILE};
13
14pub(super) const LINEAR_OP_ID: &str = "vyre-libs::nn::linear";
15/// Typed Cat-A builder for [`linear`].
16#[derive(Debug, Clone)]
17pub struct Linear {
18    x: TensorRef,
19    w: TensorRef,
20    b: TensorRef,
21    out: TensorRef,
22    options: BuildOptions,
23}
24
25impl Linear {
26    /// Create a builder for `out[i] = sum_k x[k] * w[k, i] + b[i]`.
27    #[must_use]
28    pub fn new(x: TensorRef, w: TensorRef, b: TensorRef, out: TensorRef) -> Self {
29        Self {
30            x,
31            w,
32            b,
33            out,
34            options: BuildOptions::default(),
35        }
36    }
37
38    /// Validate tensor metadata and materialize the linear Program.
39    ///
40    /// # Errors
41    ///
42    /// Returns [`TensorRefError`] when dtypes, names, shapes, or dimensions
43    /// violate the linear-layer contract.
44    pub fn build(self) -> Result<Program, TensorRefError> {
45        check_tensors(
46            LINEAR_OP_ID,
47            &[
48                (&self.x, DataType::U32),
49                (&self.w, DataType::U32),
50                (&self.b, DataType::U32),
51                (&self.out, DataType::U32),
52            ],
53        )?;
54        let x_shape = self.x.shape.as_ref();
55        let w_shape = self.w.shape.as_ref();
56        let b_shape = self.b.shape.as_ref();
57        let out_shape = self.out.shape.as_ref();
58        let expected_w = match x_shape {
59            [in_dim] => match out_shape {
60                [out_dim] => vec![*in_dim, *out_dim],
61                _ => vec![],
62            },
63            _ => vec![],
64        };
65        if w_shape != expected_w.as_slice() {
66            return Err(TensorRefError::ShapeMismatch {
67                name: self.w.name_str().to_string(),
68                found: self.w.shape.to_vec(),
69                expected: expected_w,
70                op: LINEAR_OP_ID,
71            });
72        }
73        if b_shape != out_shape {
74            return Err(TensorRefError::ShapeMismatch {
75                name: self.b.name_str().to_string(),
76                found: self.b.shape.to_vec(),
77                expected: self.out.shape.to_vec(),
78                op: LINEAR_OP_ID,
79            });
80        }
81        let &[in_dim] = x_shape else {
82            return Err(TensorRefError::ShapeMismatch {
83                name: self.x.name_str().to_string(),
84                found: self.x.shape.to_vec(),
85                expected: vec![1],
86                op: LINEAR_OP_ID,
87            });
88        };
89        let &[out_dim] = out_shape else {
90            return Err(TensorRefError::ShapeMismatch {
91                name: self.out.name_str().to_string(),
92                found: self.out.shape.to_vec(),
93                expected: vec![1],
94                op: LINEAR_OP_ID,
95            });
96        };
97        if in_dim == 0 {
98            return Err(TensorRefError::ShapeMismatch {
99                name: self.x.name_str().to_string(),
100                found: self.x.shape.to_vec(),
101                expected: vec![1],
102                op: LINEAR_OP_ID,
103            });
104        }
105        if out_dim == 0 {
106            return Err(TensorRefError::ShapeMismatch {
107                name: self.out.name_str().to_string(),
108                found: self.out.shape.to_vec(),
109                expected: vec![1],
110                op: LINEAR_OP_ID,
111            });
112        }
113        build_linear_program(
114            self.x.name_str(),
115            self.w.name_str(),
116            self.b.name_str(),
117            self.out.name_str(),
118            in_dim,
119            out_dim,
120            self.options,
121        )
122        .map_err(|_| TensorRefError::ElementCountOverflow {
123            name: self.w.name_str().to_string(),
124            shape: self.w.shape.to_vec(),
125        })
126    }
127}
128
129crate::builder::impl_cat_a_builder_options!(Linear);
130
131/// Build a Program that computes `out[i] = sum_k x[k] * w[k, i] + b[i]`.
132///
133/// Shapes: `x: [in_dim]`, `w: [in_dim, out_dim]`, `b: [out_dim]`,
134/// `out: [out_dim]`. Workgroup `[64, 1, 1]`  -  each invocation handles
135/// one output index.
136///
137/// # Errors
138/// Returns `Err` when `in_dim == 0` (FINDING-V7-TEST-010-LINEAR).
139pub fn linear(
140    x: &str,
141    w: &str,
142    b: &str,
143    out: &str,
144    in_dim: u32,
145    out_dim: u32,
146) -> Result<Program, String> {
147    if in_dim
148        .checked_mul(out_dim)
149        .is_some_and(|work| work >= LINEAR_TILED_MIN_WORK)
150    {
151        return linear_tiled(x, w, b, out, in_dim, out_dim, LINEAR_TILED_TILE);
152    }
153
154    Linear::new(
155        TensorRef::u32_1d(x, in_dim),
156        TensorRef::u32_2d(w, in_dim, out_dim),
157        TensorRef::u32_1d(b, out_dim),
158        TensorRef::u32_1d(out, out_dim),
159    )
160    .build()
161    .map_err(|error| format!("Fix: {LINEAR_OP_ID} build failed: {error}"))
162}
163
164/// Build a row-batched affine projection.
165///
166/// Shapes: `x: [rows, in_dim]`, `w: [in_dim, out_dim]`,
167/// `b: [out_dim]`, and `out: [rows, out_dim]`.
168///
169/// # Errors
170///
171/// Returns `Err` for zero dimensions or any flattened element-count overflow.
172pub fn linear_rows(
173    x: &str,
174    w: &str,
175    b: &str,
176    out: &str,
177    rows: u32,
178    in_dim: u32,
179    out_dim: u32,
180) -> Result<Program, String> {
181    linear_rows_impl(
182        x,
183        w,
184        Some(b),
185        out,
186        rows,
187        in_dim,
188        out_dim,
189        DataType::F32,
190        false,
191    )
192}
193
194/// Build a row-batched bias-free projection.
195///
196/// Shapes: `x: [rows, in_dim]`, `w: [in_dim, out_dim]`, and
197/// `out: [rows, out_dim]`.
198///
199/// # Errors
200///
201/// Returns `Err` for zero dimensions or any flattened element-count overflow.
202pub fn linear_rows_no_bias(
203    x: &str,
204    w: &str,
205    out: &str,
206    rows: u32,
207    in_dim: u32,
208    out_dim: u32,
209) -> Result<Program, String> {
210    linear_rows_impl(x, w, None, out, rows, in_dim, out_dim, DataType::F32, false)
211}
212
213/// Build a typed row-batched bias-free projection with F32 accumulation.
214///
215/// # Errors
216///
217/// Returns `Err` for unsupported dtypes, zero dimensions, or flattened
218/// element-count overflow.
219#[allow(clippy::too_many_arguments)]
220pub fn linear_rows_no_bias_typed(
221    x: &str,
222    w: &str,
223    out: &str,
224    rows: u32,
225    in_dim: u32,
226    out_dim: u32,
227    dtype: DataType,
228) -> Result<Program, String> {
229    linear_rows_impl(x, w, None, out, rows, in_dim, out_dim, dtype, false)
230}
231
232/// Build a typed bias-free projection from checkpoint-native `[out_dim, in_dim]` weights.
233///
234/// F32 accumulation is used for F16, BF16, and F32 source tensors.
235///
236/// # Errors
237///
238/// Returns `Err` for unsupported dtypes, zero dimensions, or flattened
239/// element-count overflow.
240#[allow(clippy::too_many_arguments)]
241pub fn linear_rows_no_bias_out_in_typed(
242    x: &str,
243    w: &str,
244    out: &str,
245    rows: u32,
246    in_dim: u32,
247    out_dim: u32,
248    dtype: DataType,
249) -> Result<Program, String> {
250    linear_rows_impl(x, w, None, out, rows, in_dim, out_dim, dtype, true)
251}
252
253fn linear_rows_impl(
254    x: &str,
255    w: &str,
256    bias: Option<&str>,
257    out: &str,
258    rows: u32,
259    in_dim: u32,
260    out_dim: u32,
261    dtype: DataType,
262    weight_out_in: bool,
263) -> Result<Program, String> {
264    if rows == 0 || in_dim == 0 || out_dim == 0 {
265        return Err(
266            "Fix: linear_rows requires nonzero rows, input dimension, and output dimension"
267                .to_string(),
268        );
269    }
270    if !matches!(dtype, DataType::F16 | DataType::BF16 | DataType::F32) {
271        return Err(format!(
272            "Fix: linear_rows supports F16, BF16, or F32 tensors; got {dtype:?}"
273        ));
274    }
275    rows.checked_mul(in_dim).ok_or_else(|| {
276        "Fix: linear_rows rows*in_dim overflows u32; split the row batch".to_string()
277    })?;
278    rows.checked_mul(out_dim).ok_or_else(|| {
279        "Fix: linear_rows rows*out_dim overflows u32; split the row batch".to_string()
280    })?;
281    in_dim.checked_mul(out_dim).ok_or_else(|| {
282        "Fix: linear_rows in_dim*out_dim overflows u32; shard the projection".to_string()
283    })?;
284    let input_count = rows * in_dim;
285    let output_count = rows * out_dim;
286    let weight_count = in_dim * out_dim;
287    let index = Expr::var("index");
288    let row = Expr::div(index.clone(), Expr::u32(out_dim));
289    let column = Expr::rem(index.clone(), Expr::u32(out_dim));
290    let accumulator = bias.map_or_else(
291        || Expr::f32(0.0),
292        |name| Expr::cast(DataType::F32, Expr::load(name, Expr::var("column"))),
293    );
294    let weight_index = if weight_out_in {
295        Expr::add(
296            Expr::mul(Expr::var("column"), Expr::u32(in_dim)),
297            Expr::var("inner"),
298        )
299    } else {
300        Expr::add(
301            Expr::mul(Expr::var("inner"), Expr::u32(out_dim)),
302            Expr::var("column"),
303        )
304    };
305    let body = vec![
306        Node::let_bind("index", Expr::InvocationId { axis: 0 }),
307        Node::if_then(
308            Expr::lt(index.clone(), Expr::u32(output_count)),
309            vec![
310                Node::let_bind("row", row),
311                Node::let_bind("column", column),
312                Node::let_bind("accumulator", accumulator),
313                Node::loop_for(
314                    "inner",
315                    Expr::u32(0),
316                    Expr::u32(in_dim),
317                    vec![Node::assign(
318                        "accumulator",
319                        Expr::add(
320                            Expr::var("accumulator"),
321                            Expr::mul(
322                                Expr::cast(
323                                    DataType::F32,
324                                    Expr::load(
325                                        x,
326                                        Expr::add(
327                                            Expr::mul(Expr::var("row"), Expr::u32(in_dim)),
328                                            Expr::var("inner"),
329                                        ),
330                                    ),
331                                ),
332                                Expr::cast(DataType::F32, Expr::load(w, weight_index.clone())),
333                            ),
334                        ),
335                    )],
336                ),
337                Node::Store {
338                    buffer: out.into(),
339                    index,
340                    value: Expr::cast(dtype.clone(), Expr::var("accumulator")),
341                },
342            ],
343        ),
344    ];
345    let mut buffers = vec![
346        BufferDecl::storage(x, 0, BufferAccess::ReadOnly, dtype.clone()).with_count(input_count),
347        BufferDecl::storage(w, 1, BufferAccess::ReadOnly, dtype.clone()).with_count(weight_count),
348    ];
349    if let Some(name) = bias {
350        buffers.push(
351            BufferDecl::storage(name, 2, BufferAccess::ReadOnly, dtype.clone()).with_count(out_dim),
352        );
353    }
354    let output_slot = if bias.is_some() { 3 } else { 2 };
355    buffers.push(BufferDecl::output(out, output_slot, dtype).with_count(output_count));
356    Ok(Program::wrapped(
357        buffers,
358        [64, 1, 1],
359        vec![wrap_anonymous("vyre-libs::nn::linear_rows", body)],
360    ))
361}
362
363fn build_linear_program(
364    x: &str,
365    w: &str,
366    b: &str,
367    out: &str,
368    in_dim: u32,
369    out_dim: u32,
370    options: BuildOptions,
371) -> Result<Program, String> {
372    in_dim.checked_mul(out_dim).ok_or_else(|| {
373        "Fix: linear in_dim*out_dim overflows u32; reduce dimensions.".to_string()
374    })?;
375    let mut builder = MatmulBias::new(
376        TensorRef::u32_2d(x, 1, in_dim),
377        TensorRef::u32_2d(w, in_dim, out_dim),
378        TensorRef::u32_1d(b, out_dim),
379        TensorRef::u32_2d(out, 1, out_dim),
380    );
381    if let Some(workgroup_size) = options.workgroup_size {
382        builder = builder.with_workgroup_size(workgroup_size);
383    }
384    let program = builder
385        .build()
386        .map_err(|error| format!("Fix: linear matmul_bias build failed: {error}"))?;
387    Ok(tag_program(
388        options.region_generator.unwrap_or(LINEAR_OP_ID),
389        program,
390    ))
391}
vyre_libs/nn/linear/inner/builder.rs

vyre_libs/nn/linear/inner/
builder.rs
