vyre_primitives/math/
dot_partial.rs

1//! Shared partial dot-product accumulator.
2//!
3//! This is the reusable inner kernel extracted from attention-style
4//! score passes: walk `dk` from `0..d`, load `q[q_base + dk]` and
5//! `k[k_base + dk]`, and accumulate the product into `accum_var`.
6
7use std::sync::Arc;
8
9use vyre_foundation::ir::model::expr::Ident;
10use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
11
12/// Stable Tier 2.5 op id for the attention dot-product child region.
13pub const OP_ID: &str = "vyre-primitives::math::dot_partial";
14
15/// Emit the `dk` loop that accumulates a partial dot product into `accum_var`.
16#[must_use]
17pub fn dot_partial(
18    q_buffer: &str,
19    k_buffer: &str,
20    accum_var: &str,
21    q_base: Expr,
22    k_base: Expr,
23    d: u32,
24) -> Node {
25    if d <= 8 {
26        return Node::Block(
27            (0..d)
28                .map(|lane| {
29                    Node::assign(
30                        accum_var,
31                        Expr::add(
32                            Expr::var(accum_var),
33                            Expr::mul(
34                                Expr::cast(
35                                    DataType::F32,
36                                    Expr::load(
37                                        q_buffer,
38                                        Expr::add(q_base.clone(), Expr::u32(lane)),
39                                    ),
40                                ),
41                                Expr::cast(
42                                    DataType::F32,
43                                    Expr::load(
44                                        k_buffer,
45                                        Expr::add(k_base.clone(), Expr::u32(lane)),
46                                    ),
47                                ),
48                            ),
49                        ),
50                    )
51                })
52                .collect(),
53        );
54    }
55
56    Node::loop_for(
57        "dk",
58        Expr::u32(0),
59        Expr::u32(d),
60        vec![Node::assign(
61            accum_var,
62            Expr::add(
63                Expr::var(accum_var),
64                Expr::mul(
65                    Expr::cast(
66                        DataType::F32,
67                        Expr::load(q_buffer, Expr::add(q_base, Expr::var("dk"))),
68                    ),
69                    Expr::cast(
70                        DataType::F32,
71                        Expr::load(k_buffer, Expr::add(k_base, Expr::var("dk"))),
72                    ),
73                ),
74            ),
75        )],
76    )
77}
78
79/// Standalone dot-partial Program.
80#[must_use]
81pub fn dot_partial_program(q_buffer: &str, k_buffer: &str, out: &str, d: u32) -> Program {
82    Program::wrapped(
83        vec![
84            BufferDecl::storage(q_buffer, 0, BufferAccess::ReadOnly, DataType::F32).with_count(d),
85            BufferDecl::storage(k_buffer, 1, BufferAccess::ReadOnly, DataType::F32).with_count(d),
86            BufferDecl::storage(out, 2, BufferAccess::ReadWrite, DataType::F32).with_count(1),
87        ],
88        [1, 1, 1],
89        vec![Node::Region {
90            generator: Ident::from(OP_ID),
91            source_region: None,
92            body: Arc::new(vec![
93                Node::let_bind("accum", Expr::f32(0.0)),
94                dot_partial(q_buffer, k_buffer, "accum", Expr::u32(0), Expr::u32(0), d),
95                Node::store(out, Expr::u32(0), Expr::var("accum")),
96            ]),
97        }],
98    )
99}
100
101#[cfg(feature = "inventory-registry")]
102inventory::submit! {
103    vyre_foundation::operation::OperationRegistration::primitive(
104        OP_ID,
105        || dot_partial_program("q", "k", "out", 2),
106        Some(|| {
107            let to_f32_bytes = |w: &[f32]| crate::wire::pack_f32_slice(w);
108            vec![vec![
109                to_f32_bytes(&[2.0, 3.0]),
110                to_f32_bytes(&[4.0, 5.0]),
111                vec![0u8; 4],
112            ]]
113        }),
114        Some(|| {
115            let to_f32_bytes = |w: &[f32]| crate::wire::pack_f32_slice(w);
116            vec![vec![to_f32_bytes(&[23.0])]]
117        }),
118    )
119}
vyre_primitives/math/dot_partial.rs

vyre_primitives/math/
dot_partial.rs
