vyre_libs/nn/linear/inner/
builder.rs

1//! Linear builder + the canonical `linear()` Cat-A constructor.
2
3use vyre::ir::{DataType, Program};
4
5use crate::{
6    builder::{check_tensors, BuildOptions},
7    region::tag_program,
8    tensor_ref::{TensorRef, TensorRefError},
9    MatmulBias,
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
164fn build_linear_program(
165    x: &str,
166    w: &str,
167    b: &str,
168    out: &str,
169    in_dim: u32,
170    out_dim: u32,
171    options: BuildOptions,
172) -> Result<Program, String> {
173    in_dim.checked_mul(out_dim).ok_or_else(|| {
174        "Fix: linear in_dim*out_dim overflows u32; reduce dimensions.".to_string()
175    })?;
176    let mut builder = MatmulBias::new(
177        TensorRef::u32_2d(x, 1, in_dim),
178        TensorRef::u32_2d(w, in_dim, out_dim),
179        TensorRef::u32_1d(b, out_dim),
180        TensorRef::u32_2d(out, 1, out_dim),
181    );
182    if let Some(workgroup_size) = options.workgroup_size {
183        builder = builder.with_workgroup_size(workgroup_size);
184    }
185    let program = builder
186        .build()
187        .map_err(|error| format!("Fix: linear matmul_bias build failed: {error}"))?;
188    Ok(tag_program(
189        options.region_generator.unwrap_or(LINEAR_OP_ID),
190        program,
191    ))
192}
vyre_libs/nn/linear/inner/builder.rs

vyre_libs/nn/linear/inner/
builder.rs
