vyre_libs/nn/conv/
mod.rs

1//! Floating depthwise causal convolution for sequence models.
2
3use thiserror::Error;
4use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program, UnOp};
5
6use crate::region::wrap_anonymous;
7
8const OP_ID: &str = "vyre-libs::nn::depthwise_causal_conv1d";
9
10/// Optional post-convolution activation.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum CausalConvActivation {
13    /// Return the affine convolution result.
14    None,
15    /// Apply SiLU in F32 before source-dtype conversion.
16    Silu,
17}
18
19/// Invalid depthwise causal convolution construction.
20#[derive(Debug, Clone, PartialEq, Eq, Error)]
21pub enum DepthwiseCausalConv1dError {
22    /// One required tensor dimension is zero.
23    #[error("depthwise causal convolution requires nonzero batch, channels, sequence, and kernel dimensions")]
24    EmptyShape,
25    /// One flattened tensor size exceeds u32 indexing.
26    #[error("depthwise causal convolution tensor element count overflows u32; split the tensor")]
27    ElementCountOverflow,
28    /// Streaming state is unnecessary for a pointwise kernel.
29    #[error("causal convolution state update requires kernel >= 2")]
30    StateKernelTooShort,
31    /// Source dtype lacks the required floating conversion contract.
32    #[error("depthwise causal convolution supports F16, BF16, or F32 tensors; got {dtype:?}")]
33    UnsupportedDtype {
34        /// Rejected dtype.
35        dtype: DataType,
36    },
37}
38
39/// Build channel-major left-padded depthwise causal convolution.
40///
41/// `input` and `output` use `[batch, channels, sequence]`; `weight` uses
42/// `[channels, kernel]`; optional `bias` uses `[channels]`; optional `mask`
43/// uses U32 `[batch, sequence]`, where zero excludes an input position.
44/// Accumulation and optional SiLU execute in F32 before one output conversion.
45///
46/// # Errors
47///
48/// Returns [`DepthwiseCausalConv1dError`] for empty/overflowing shapes or an
49/// unsupported source dtype.
50#[allow(clippy::too_many_arguments)]
51pub fn depthwise_causal_conv1d(
52    input: &str,
53    weight: &str,
54    bias: Option<&str>,
55    mask: Option<&str>,
56    output: &str,
57    batch: u32,
58    channels: u32,
59    sequence: u32,
60    kernel: u32,
61    activation: CausalConvActivation,
62    dtype: DataType,
63) -> Result<Program, DepthwiseCausalConv1dError> {
64    if batch == 0 || channels == 0 || sequence == 0 || kernel == 0 {
65        return Err(DepthwiseCausalConv1dError::EmptyShape);
66    }
67    if !matches!(dtype, DataType::F16 | DataType::BF16 | DataType::F32) {
68        return Err(DepthwiseCausalConv1dError::UnsupportedDtype { dtype });
69    }
70    let channel_sequence = channels
71        .checked_mul(sequence)
72        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
73    let total = batch
74        .checked_mul(channel_sequence)
75        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
76    let weight_count = channels
77        .checked_mul(kernel)
78        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
79    let mask_count = batch
80        .checked_mul(sequence)
81        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
82
83    let index = Expr::var("index");
84    let batch_index = Expr::div(index.clone(), Expr::u32(channel_sequence));
85    let within_batch = Expr::sub(
86        index.clone(),
87        Expr::mul(batch_index.clone(), Expr::u32(channel_sequence)),
88    );
89    let channel = Expr::div(within_batch.clone(), Expr::u32(sequence));
90    let time = Expr::sub(
91        within_batch.clone(),
92        Expr::mul(channel.clone(), Expr::u32(sequence)),
93    );
94    let initial = bias.map_or_else(
95        || Expr::f32(0.0),
96        |name| Expr::cast(DataType::F32, Expr::load(name, channel.clone())),
97    );
98    let lag = Expr::sub(Expr::u32(kernel - 1), Expr::var("kernel_index"));
99    let position = Expr::sub(Expr::var("time"), lag.clone());
100    let input_index = Expr::add(
101        Expr::mul(Expr::var("batch"), Expr::u32(channel_sequence)),
102        Expr::add(
103            Expr::mul(Expr::var("channel"), Expr::u32(sequence)),
104            Expr::var("position"),
105        ),
106    );
107    let product = Expr::mul(
108        Expr::cast(DataType::F32, Expr::load(input, input_index)),
109        Expr::cast(
110            DataType::F32,
111            Expr::load(
112                weight,
113                Expr::add(
114                    Expr::mul(Expr::var("channel"), Expr::u32(kernel)),
115                    Expr::var("kernel_index"),
116                ),
117            ),
118        ),
119    );
120    let accumulate = Node::assign("accumulator", Expr::add(Expr::var("accumulator"), product));
121    let valid_body = if let Some(mask_name) = mask {
122        vec![
123            Node::let_bind("position", position),
124            Node::if_then(
125                Expr::ne(
126                    Expr::load(
127                        mask_name,
128                        Expr::add(
129                            Expr::mul(Expr::var("batch"), Expr::u32(sequence)),
130                            Expr::var("position"),
131                        ),
132                    ),
133                    Expr::u32(0),
134                ),
135                vec![accumulate],
136            ),
137        ]
138    } else {
139        vec![Node::let_bind("position", position), accumulate]
140    };
141    let activated = match activation {
142        CausalConvActivation::None => Expr::var("accumulator"),
143        CausalConvActivation::Silu => {
144            let value = Expr::var("accumulator");
145            Expr::div(
146                value.clone(),
147                Expr::add(
148                    Expr::f32(1.0),
149                    Expr::UnOp {
150                        op: UnOp::Exp,
151                        operand: Box::new(Expr::UnOp {
152                            op: UnOp::Negate,
153                            operand: Box::new(value),
154                        }),
155                    },
156                ),
157            )
158        }
159    };
160    let body = vec![
161        Node::let_bind("index", Expr::InvocationId { axis: 0 }),
162        Node::if_then(
163            Expr::lt(index.clone(), Expr::u32(total)),
164            vec![
165                Node::let_bind("batch", batch_index),
166                Node::let_bind("within_batch", within_batch),
167                Node::let_bind("channel", channel),
168                Node::let_bind("time", time),
169                Node::let_bind("accumulator", initial),
170                Node::loop_for(
171                    "kernel_index",
172                    Expr::u32(0),
173                    Expr::u32(kernel),
174                    vec![Node::if_then(Expr::ge(Expr::var("time"), lag), valid_body)],
175                ),
176                Node::Store {
177                    buffer: output.into(),
178                    index,
179                    value: Expr::cast(dtype.clone(), activated),
180                },
181            ],
182        ),
183    ];
184
185    let mut buffers = Vec::with_capacity(5);
186    buffers.push(
187        BufferDecl::storage(input, 0, BufferAccess::ReadOnly, dtype.clone()).with_count(total),
188    );
189    buffers.push(
190        BufferDecl::storage(weight, 1, BufferAccess::ReadOnly, dtype.clone())
191            .with_count(weight_count),
192    );
193    let mut binding = 2;
194    if let Some(name) = bias {
195        buffers.push(
196            BufferDecl::storage(name, binding, BufferAccess::ReadOnly, dtype.clone())
197                .with_count(channels),
198        );
199        binding += 1;
200    }
201    if let Some(name) = mask {
202        buffers.push(
203            BufferDecl::storage(name, binding, BufferAccess::ReadOnly, DataType::U32)
204                .with_count(mask_count),
205        );
206        binding += 1;
207    }
208    buffers.push(BufferDecl::output(output, binding, dtype).with_count(total));
209    Ok(Program::wrapped(
210        buffers,
211        [64, 1, 1],
212        vec![wrap_anonymous(OP_ID, body)],
213    ))
214}
215
216/// Build short-chunk convolution and its next explicit loop-carried state.
217///
218/// State tensors use `[batch, channels, kernel - 1]`. The next state is a
219/// separate output, so runtimes may recycle prior-state storage only after all
220/// reads finish.
221#[allow(clippy::too_many_arguments)]
222pub fn depthwise_causal_conv1d_update(
223    input: &str,
224    weight: &str,
225    bias: Option<&str>,
226    state_input: &str,
227    output: &str,
228    state_output: &str,
229    batch: u32,
230    channels: u32,
231    chunk: u32,
232    kernel: u32,
233    activation: CausalConvActivation,
234    dtype: DataType,
235) -> Result<Program, DepthwiseCausalConv1dError> {
236    if kernel < 2 {
237        return Err(DepthwiseCausalConv1dError::StateKernelTooShort);
238    }
239    if batch == 0 || channels == 0 || chunk == 0 {
240        return Err(DepthwiseCausalConv1dError::EmptyShape);
241    }
242    if !matches!(dtype, DataType::F16 | DataType::BF16 | DataType::F32) {
243        return Err(DepthwiseCausalConv1dError::UnsupportedDtype { dtype });
244    }
245    let state_len = kernel - 1;
246    let channel_chunk = channels
247        .checked_mul(chunk)
248        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
249    let output_count = batch
250        .checked_mul(channel_chunk)
251        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
252    let channel_state = channels
253        .checked_mul(state_len)
254        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
255    let state_count = batch
256        .checked_mul(channel_state)
257        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
258    let weight_count = channels
259        .checked_mul(kernel)
260        .ok_or(DepthwiseCausalConv1dError::ElementCountOverflow)?;
261
262    let initial = bias.map_or_else(
263        || Expr::f32(0.0),
264        |name| Expr::cast(DataType::F32, Expr::load(name, Expr::var("output_channel"))),
265    );
266    let state_sample = Expr::load(
267        state_input,
268        Expr::add(
269            Expr::mul(Expr::var("output_batch"), Expr::u32(channel_state)),
270            Expr::add(
271                Expr::mul(Expr::var("output_channel"), Expr::u32(state_len)),
272                Expr::var("combined"),
273            ),
274        ),
275    );
276    let input_sample = Expr::load(
277        input,
278        Expr::add(
279            Expr::mul(Expr::var("output_batch"), Expr::u32(channel_chunk)),
280            Expr::add(
281                Expr::mul(Expr::var("output_channel"), Expr::u32(chunk)),
282                Expr::sub(Expr::var("combined"), Expr::u32(state_len)),
283            ),
284        ),
285    );
286    let weight_sample = Expr::load(
287        weight,
288        Expr::add(
289            Expr::mul(Expr::var("output_channel"), Expr::u32(kernel)),
290            Expr::var("kernel_index"),
291        ),
292    );
293    let activated = match activation {
294        CausalConvActivation::None => Expr::var("accumulator"),
295        CausalConvActivation::Silu => {
296            let value = Expr::var("accumulator");
297            Expr::div(
298                value.clone(),
299                Expr::add(
300                    Expr::f32(1.0),
301                    Expr::UnOp {
302                        op: UnOp::Exp,
303                        operand: Box::new(Expr::UnOp {
304                            op: UnOp::Negate,
305                            operand: Box::new(value),
306                        }),
307                    },
308                ),
309            )
310        }
311    };
312    let output_body = vec![
313        Node::let_bind(
314            "output_batch",
315            Expr::div(Expr::var("dispatch_index"), Expr::u32(channel_chunk)),
316        ),
317        Node::let_bind(
318            "output_remainder",
319            Expr::sub(
320                Expr::var("dispatch_index"),
321                Expr::mul(Expr::var("output_batch"), Expr::u32(channel_chunk)),
322            ),
323        ),
324        Node::let_bind(
325            "output_channel",
326            Expr::div(Expr::var("output_remainder"), Expr::u32(chunk)),
327        ),
328        Node::let_bind(
329            "output_time",
330            Expr::sub(
331                Expr::var("output_remainder"),
332                Expr::mul(Expr::var("output_channel"), Expr::u32(chunk)),
333            ),
334        ),
335        Node::let_bind("accumulator", initial),
336        Node::loop_for(
337            "kernel_index",
338            Expr::u32(0),
339            Expr::u32(kernel),
340            vec![
341                Node::let_bind(
342                    "combined",
343                    Expr::add(Expr::var("output_time"), Expr::var("kernel_index")),
344                ),
345                Node::let_bind("sample", Expr::f32(0.0)),
346                Node::if_then(
347                    Expr::lt(Expr::var("combined"), Expr::u32(state_len)),
348                    vec![Node::assign(
349                        "sample",
350                        Expr::cast(DataType::F32, state_sample),
351                    )],
352                ),
353                Node::if_then(
354                    Expr::ge(Expr::var("combined"), Expr::u32(state_len)),
355                    vec![Node::assign(
356                        "sample",
357                        Expr::cast(DataType::F32, input_sample),
358                    )],
359                ),
360                Node::assign(
361                    "accumulator",
362                    Expr::add(
363                        Expr::var("accumulator"),
364                        Expr::mul(
365                            Expr::var("sample"),
366                            Expr::cast(DataType::F32, weight_sample),
367                        ),
368                    ),
369                ),
370            ],
371        ),
372        Node::Store {
373            buffer: output.into(),
374            index: Expr::var("dispatch_index"),
375            value: Expr::cast(dtype.clone(), activated),
376        },
377    ];
378
379    let prior_state_tail = Expr::load(
380        state_input,
381        Expr::add(
382            Expr::mul(Expr::var("state_batch"), Expr::u32(channel_state)),
383            Expr::add(
384                Expr::mul(Expr::var("state_channel"), Expr::u32(state_len)),
385                Expr::var("state_combined"),
386            ),
387        ),
388    );
389    let input_tail = Expr::load(
390        input,
391        Expr::add(
392            Expr::mul(Expr::var("state_batch"), Expr::u32(channel_chunk)),
393            Expr::add(
394                Expr::mul(Expr::var("state_channel"), Expr::u32(chunk)),
395                Expr::sub(Expr::var("state_combined"), Expr::u32(state_len)),
396            ),
397        ),
398    );
399    let state_body = vec![
400        Node::let_bind(
401            "state_batch",
402            Expr::div(Expr::var("dispatch_index"), Expr::u32(channel_state)),
403        ),
404        Node::let_bind(
405            "state_remainder",
406            Expr::sub(
407                Expr::var("dispatch_index"),
408                Expr::mul(Expr::var("state_batch"), Expr::u32(channel_state)),
409            ),
410        ),
411        Node::let_bind(
412            "state_channel",
413            Expr::div(Expr::var("state_remainder"), Expr::u32(state_len)),
414        ),
415        Node::let_bind(
416            "state_offset",
417            Expr::sub(
418                Expr::var("state_remainder"),
419                Expr::mul(Expr::var("state_channel"), Expr::u32(state_len)),
420            ),
421        ),
422        Node::let_bind(
423            "state_combined",
424            Expr::add(Expr::u32(chunk), Expr::var("state_offset")),
425        ),
426        Node::let_bind("next_state_value", Expr::f32(0.0)),
427        Node::if_then(
428            Expr::lt(Expr::var("state_combined"), Expr::u32(state_len)),
429            vec![Node::assign(
430                "next_state_value",
431                Expr::cast(DataType::F32, prior_state_tail),
432            )],
433        ),
434        Node::if_then(
435            Expr::ge(Expr::var("state_combined"), Expr::u32(state_len)),
436            vec![Node::assign(
437                "next_state_value",
438                Expr::cast(DataType::F32, input_tail),
439            )],
440        ),
441        Node::Store {
442            buffer: state_output.into(),
443            index: Expr::var("dispatch_index"),
444            value: Expr::cast(dtype.clone(), Expr::var("next_state_value")),
445        },
446    ];
447    let body = vec![
448        Node::let_bind("dispatch_index", Expr::InvocationId { axis: 0 }),
449        Node::if_then(
450            Expr::lt(Expr::var("dispatch_index"), Expr::u32(output_count)),
451            output_body,
452        ),
453        Node::if_then(
454            Expr::lt(Expr::var("dispatch_index"), Expr::u32(state_count)),
455            state_body,
456        ),
457    ];
458
459    let mut buffers = vec![
460        BufferDecl::storage(input, 0, BufferAccess::ReadOnly, dtype.clone())
461            .with_count(output_count),
462        BufferDecl::storage(weight, 1, BufferAccess::ReadOnly, dtype.clone())
463            .with_count(weight_count),
464        BufferDecl::storage(state_input, 2, BufferAccess::ReadWrite, dtype.clone())
465            .with_count(state_count),
466    ];
467    let mut binding = 3;
468    if let Some(name) = bias {
469        buffers.push(
470            BufferDecl::storage(name, binding, BufferAccess::ReadOnly, dtype.clone())
471                .with_count(channels),
472        );
473        binding += 1;
474    }
475    buffers.push(BufferDecl::output(output, binding, dtype.clone()).with_count(output_count));
476    buffers.push(
477        BufferDecl::storage(state_output, binding + 1, BufferAccess::ReadWrite, dtype)
478            .with_count(state_count),
479    );
480    Ok(Program::wrapped(
481        buffers,
482        [64, 1, 1],
483        vec![wrap_anonymous(
484            "vyre-libs::nn::depthwise_causal_conv1d_update",
485            body,
486        )],
487    ))
488}
vyre_libs/nn/conv/mod.rs

vyre_libs/nn/conv/
mod.rs
