vision_rs/models/yolo/yolo26/blocks/
c2psa.rs

1/*
2 * Copyright 2026 Teenygrad
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17
18use teeny_core::{dtype::Float, graph::{CustomData, Op, SymTensor}, name_scope::name_scope};
19
20use crate::models::yolo::kernels::attention::psa::{
21    FlashAttn2PsaOp, PsaExtractVOp, PsaMergeAttnOp, PsaPackQkvOp,
22};
23use super::conv::{conv, conv_bn};
24
25// ── Graph helpers ─────────────────────────────────────────────────────────────
26
27fn channel_chunk(x: SymTensor, c_total: usize, chunk_c: usize, chunk_offset: usize) -> SymTensor {
28    let op = Op::ChannelChunk { c_total, chunk_c, chunk_offset };
29    let shape = vec![x.shape[0], Some(chunk_c), x.shape[2], x.shape[3]];
30    let node_id = x.graph.borrow_mut().add_node(op, vec![x.node_id], x.dtype, shape.clone());
31    SymTensor { node_id, graph: x.graph.clone(), dtype: x.dtype, shape }
32}
33
34fn channel_cat(tensors: Vec<SymTensor>, c_total: usize) -> SymTensor {
35    let first = &tensors[0];
36    let shape = vec![first.shape[0], Some(c_total), first.shape[2], first.shape[3]];
37    let inputs: Vec<usize> = tensors.iter().map(|t| t.node_id).collect();
38    let node_id = first.graph.borrow_mut().add_node(
39        Op::ChannelCat { c_total }, inputs, first.dtype, shape.clone(),
40    );
41    SymTensor { node_id, graph: first.graph.clone(), dtype: first.dtype, shape }
42}
43
44fn elem_add(a: SymTensor, b: SymTensor) -> SymTensor {
45    let shape = a.shape.clone();
46    let node_id = a.graph.borrow_mut().add_node(
47        Op::Add, vec![a.node_id, b.node_id], a.dtype, shape.clone(),
48    );
49    SymTensor { node_id, graph: a.graph.clone(), dtype: a.dtype, shape }
50}
51
52// ── PSA Attention ─────────────────────────────────────────────────────────────
53
54/// Multi-head self-attention with Flash Attention 2 + position encoding.
55///
56/// Matches `ultralytics.nn.modules.block.Attention.forward(x)`:
57///   qkv conv+BN → pack_qkv → FA2×2 → merge_attn
58///   ↕ + extract_v → pe_dw_conv+BN → (merge + pe) → proj_conv+BN
59///
60/// Input/output: `[B, c, H, W]`  (residual add is applied by the caller).
61fn psa_attention<D: Float + Send + Sync + 'static>(c: usize, num_heads: usize, key_dim: usize)
62    -> impl Fn(SymTensor) -> SymTensor
63{
64    let qkv_h    = num_heads * 4 * key_dim;
65    let qkv_conv = conv_bn::<D>(c, qkv_h, 1, 1, 1);
66    let pe_dw    = conv_bn::<D>(c, c, 3, 1, c);       // depthwise position encoding
67    let proj     = conv_bn::<D>(c, c, 1, 1, 1);
68
69    move |x: SymTensor| {
70        let h = x.shape[2].unwrap_or(1);
71        let w = x.shape[3].unwrap_or(1);
72
73        // [B, c, H, W] → [B, qkv_h, H, W]
74        let qkv = { let _g = name_scope("qkv"); qkv_conv(x) };
75
76        // [B, qkv_h, H, W] → [4, BH, N, KEY_DIM]
77        let packed = qkv.record_custom(
78            CustomData::new(PsaPackQkvOp::<D>::new(key_dim as i32, num_heads)),
79            &[],
80            None,
81        );
82
83        // [4, BH, N, KEY_DIM] → [BH, N, KEY_DIM]  (V_lo, V_hi separately)
84        let lo = packed.record_custom(
85            CustomData::new(FlashAttn2PsaOp::<D>::new_lo(key_dim as i32)),
86            &[],
87            None,
88        );
89        let hi = packed.record_custom(
90            CustomData::new(FlashAttn2PsaOp::<D>::new_hi(key_dim as i32)),
91            &[],
92            None,
93        );
94
95        // (lo, hi) [BH, N, KEY_DIM] → [B, c, H, W]
96        let merged = lo.record_custom(
97            CustomData::new(PsaMergeAttnOp::<D>::new(key_dim as i32, num_heads, h, w)),
98            &[&hi],
99            None,
100        );
101
102        // [B, qkv_h, H, W] → [B, c, H, W]  (V channels in NCHW for PE)
103        let v_nchw = qkv.record_custom(
104            CustomData::new(PsaExtractVOp::<D>::new(key_dim as i32, num_heads)),
105            &[],
106            None,
107        );
108
109        // PE depthwise conv, then add to merged attention
110        let pe      = { let _g = name_scope("pe"); pe_dw(v_nchw) };
111        let attn_pe = elem_add(merged, pe);
112
113        // Final projection
114        { let _g = name_scope("proj"); proj(attn_pe) }
115    }
116}
117
118// ── PSABlock ──────────────────────────────────────────────────────────────────
119
120/// Single PSABlock iteration — attention residual + FFN residual.
121///
122/// Matches `ultralytics.nn.modules.block.PSABlock(c, attn_ratio=0.5, num_heads, shortcut=True)`.
123pub(super) fn psa_block<D: Float + Send + Sync + 'static>(c: usize, num_heads: usize, key_dim: usize)
124    -> impl Fn(SymTensor) -> SymTensor
125{
126    let attn = psa_attention::<D>(c, num_heads, key_dim);
127    let ffn0 = conv::<D>(c, 2 * c, 1, 1);
128    let ffn1 = conv_bn::<D>(2 * c, c, 1, 1, 1);
129    move |b: SymTensor| {
130        let b = elem_add(b.clone(), { let _g = name_scope("attn"); attn(b) });
131        let ffn_out = {
132            let tmp = { let _g = name_scope("ffn.0"); ffn0(b.clone()) };
133            let _g = name_scope("ffn.1"); ffn1(tmp)
134        };
135        elem_add(b, ffn_out)
136    }
137}
138
139// ── C2PSA ─────────────────────────────────────────────────────────────────────
140
141/// C2PSA: cross-stage partial network with PSA attention blocks.
142///
143/// Matches `ultralytics.nn.modules.block.C2PSA(c1, c2, n, e)`.
144///
145/// Forward:
146///   h        = cv1(x)           // [B, 2*c, H, W]
147///   [a, b]   = split(h, c)      // each [B, c, H, W]
148///   b        = PSABlock(b) × n
149///   output   = cv2(cat(a, b))   // [B, c2, H, W]
150pub fn c2psa<D: Float + Send + Sync + 'static>(
151    c_in: usize,
152    c_out: usize,
153    n: usize,
154    e: f32,
155) -> impl Fn(SymTensor) -> SymTensor {
156    assert_eq!(c_in, c_out, "C2PSA requires c_in == c_out");
157    let c = (c_out as f32 * e) as usize;
158    let num_heads = c / 64;
159    let key_dim = 32;
160
161    let cv1 = conv::<D>(c_in, 2 * c, 1, 1);
162    let cv2 = conv::<D>(2 * c, c_out, 1, 1);
163    let blocks: Vec<Box<dyn Fn(SymTensor) -> SymTensor>> = (0..n)
164        .map(|_| -> Box<dyn Fn(SymTensor) -> SymTensor> {
165            Box::new(psa_block::<D>(c, num_heads, key_dim))
166        })
167        .collect();
168
169    move |x: SymTensor| {
170        let h  = { let _g = name_scope("cv1"); cv1(x) };
171        let a  = channel_chunk(h.clone(), 2 * c, c, 0);
172        let mut b = channel_chunk(h, 2 * c, c, c);
173        for (i, blk) in blocks.iter().enumerate() {
174            let _g = name_scope(format!("m.{i}"));
175            b = blk(b);
176        }
177        { let _g = name_scope("cv2"); cv2(channel_cat(vec![a, b], 2 * c)) }
178    }
179}
vision_rs/models/yolo/yolo26/blocks/c2psa.rs

vision_rs/models/yolo/yolo26/blocks/
c2psa.rs
