teeny_kernels/nn/fused/
conv2d_bn_silu_gemm.rs

1/*
2 * Copyright (c) 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#![allow(non_snake_case)]
18
19use core::ops::BitAnd;
20
21use teeny_macros::kernel;
22use teeny_triton::triton::{
23    types::{AddOffsets, Comparison, Tensor},
24    *,
25};
26
27/// Fused Conv2d(1×1, stride=1, pad=0, groups=1) + BatchNorm2d + SiLU using GEMM.
28///
29/// A 1×1 stride=1 no-padding convolution is mathematically equivalent to:
30///   Y[N, M] = W[N, K] @ X[K, M]
31/// where:
32///   N = C_OUT, K = C_IN, M = OH * OW  (batch handled in grid)
33///
34/// Input X is stored NCHW = [B, K, OH, OW], viewed as 2-D [K, OH*OW] per batch
35/// (column-major spatial: stride_K = OH*OW, stride_spatial = 1).
36///
37/// Weight W is [C_OUT, C_IN] row-major.
38///
39/// T::dot uses TF32 Tensor Cores on sm_87+ (Jetson Orin) for ~8× throughput
40/// vs direct scalar accumulation.
41///
42/// Restrictions (enforced by dispatch in graph/mod.rs):
43///   - kernel_h == 1, kernel_w == 1
44///   - stride_h == 1, stride_w == 1
45///   - padding_h == 0, padding_w == 0
46///   - groups == 1
47///
48/// BN parameters must be precomputed (same convention as conv2d_bn_silu_forward).
49///
50/// Grid: pid = b * num_pid_per_batch + group_id * (GROUP_M * num_pid_n)
51///             + pid_in_group  (same L2-locality grouping as Triton's matmul tutorial)
52///
53/// Inference-only; no backward pass.
54#[kernel]
55pub fn conv2d_bn_silu_gemm_forward<
56    T: Triton,
57    const BLOCK_M: i32,
58    const BLOCK_N: i32,
59    const BLOCK_K: i32,
60    const GROUP_M: i32,
61>(
62    x_ptr: T::Pointer<f32>,
63    w_ptr: T::Pointer<f32>,
64    bn_scale_ptr: T::Pointer<f32>,
65    bn_shift_ptr: T::Pointer<f32>,
66    y_ptr: T::Pointer<f32>,
67    B: i32,
68    C_IN: i32,
69    C_OUT: i32,
70    M: i32, // OH * OW per batch
71) where
72    T::I32Tensor: Tensor<i32, 1>,
73    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
74    T::BoolTensor: BitAnd<Output = T::BoolTensor>,
75    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
76{
77    let pid = T::program_id(Axis::X);
78
79    let num_pid_m = T::cdiv(M, BLOCK_M);
80    let num_pid_n = T::cdiv(C_OUT, BLOCK_N);
81    let pids_per_batch = num_pid_m * num_pid_n;
82
83    // Decode batch dimension from pid
84    let b = pid / pids_per_batch;
85    let pid_local = pid % pids_per_batch;
86
87    // L2-locality grouping (same as Triton matmul tutorial)
88    let num_pid_in_group = GROUP_M * num_pid_n;
89    let group_id = pid_local / num_pid_in_group;
90    let first_pid_m = group_id * GROUP_M;
91    let remaining_m = num_pid_m - first_pid_m;
92    let group_size_m = if remaining_m < GROUP_M {
93        remaining_m
94    } else {
95        GROUP_M
96    };
97    let pid_in_group = pid_local % num_pid_in_group;
98    let pid_m = first_pid_m + (pid_in_group % group_size_m);
99    let pid_n = pid_in_group / group_size_m;
100
101    // ── Input descriptor: X viewed as [B*C_IN, M] ────────────────────────────
102    // NCHW layout: x[b, c_in, oh, ow] = x_flat[b*(C_IN*M) + c_in*M + oh*OW + ow]
103    // For batch b: the C_IN rows for that batch start at row b*C_IN.
104    // stride along K (row) = M (spatial extent per channel)
105    // stride along M (col) = 1
106    let x_desc = T::make_tensor_descriptor(
107        x_ptr,
108        &[B * C_IN, M],
109        &[M, 1],
110        &[BLOCK_K, BLOCK_M],
111        Some(PaddingOption::Zero),
112    );
113
114    // ── Weight descriptor: W [C_OUT, C_IN] row-major ─────────────────────────
115    let w_desc = T::make_tensor_descriptor(
116        w_ptr,
117        &[C_OUT, C_IN],
118        &[C_IN, 1],
119        &[BLOCK_N, BLOCK_K],
120        Some(PaddingOption::Zero),
121    );
122
123    // ── GEMM: acc [BLOCK_N, BLOCK_M] = sum_k W_tile @ X_tile ─────────────────
124    let mut acc = T::zeros::<f32>(&[BLOCK_N, BLOCK_M]);
125    let k_tiles = T::cdiv(C_IN, BLOCK_K);
126    for k in 0..k_tiles {
127        // x_tile: [BLOCK_K, BLOCK_M] — rows are the C_IN channels for batch b
128        let x_tile = T::load_tensor_descriptor(x_desc, &[b * C_IN + k * BLOCK_K, pid_m * BLOCK_M]);
129
130        // w_tile: [BLOCK_N, BLOCK_K]
131        let w_tile = T::load_tensor_descriptor(w_desc, &[pid_n * BLOCK_N, k * BLOCK_K]);
132
133        // [BLOCK_N, BLOCK_K] @ [BLOCK_K, BLOCK_M] → [BLOCK_N, BLOCK_M]
134        // TF32 precision is what actually routes this dot to the tensor-core MMA
135        // path (see getMmaTypeDot in Triton's MMAv2.cpp) — IEEE forces the
136        // software FMA fallback and silently disables tensor cores entirely.
137        acc = T::dot::<f32, f32>(w_tile, x_tile, Some(acc), InputPrecision::TF32, None);
138    }
139
140    // ── BatchNorm epilog ──────────────────────────────────────────────────────
141    let bn_off = T::arange(0, BLOCK_N) + pid_n * BLOCK_N;
142    let bn_n_mask = bn_off.lt(C_OUT);
143    let bn_scale = T::load(
144        bn_scale_ptr.add_offsets(bn_off),
145        Some(bn_n_mask),
146        Some(T::zeros::<f32>(&[BLOCK_N])),
147        &[],
148        None,
149        None,
150        None,
151        false,
152    );
153    let bn_shift = T::load(
154        bn_shift_ptr.add_offsets(bn_off),
155        Some(bn_n_mask),
156        Some(T::zeros::<f32>(&[BLOCK_N])),
157        &[],
158        None,
159        None,
160        None,
161        false,
162    );
163    let scale_2d = T::broadcast_to(T::expand_dims(bn_scale, 1), &[BLOCK_N, BLOCK_M]);
164    let shift_2d = T::broadcast_to(T::expand_dims(bn_shift, 1), &[BLOCK_N, BLOCK_M]);
165    let bn_out = scale_2d * acc + shift_2d;
166
167    // ── SiLU epilog: y = x * sigmoid(x) ─────────────────────────────────────
168    let y = bn_out * T::sigmoid(bn_out);
169
170    // ── Store: y NCHW [B, C_OUT, OH, OW] viewed as [B*C_OUT, M] ─────────────
171    // y[b, c_out, oh, ow] = y_flat[(b*C_OUT + c_out)*M + oh*OW + ow]
172    let y_desc = T::make_tensor_descriptor(
173        y_ptr,
174        &[B * C_OUT, M],
175        &[M, 1],
176        &[BLOCK_N, BLOCK_M],
177        Some(PaddingOption::Zero),
178    );
179    T::store_tensor_descriptor(y_desc, &[b * C_OUT + pid_n * BLOCK_N, pid_m * BLOCK_M], y);
180}
181
182// ── RuntimeOp ────────────────────────────────────────────────────────────────
183//
184// Params layout: [weight [C_OUT, C_IN], bn_scale [C_OUT], bn_shift [C_OUT]]
185// pack_args order: x_ptr, w_ptr, bn_scale_ptr, bn_shift_ptr, y_ptr,
186//                  B, C_IN, C_OUT, M (= OH * OW)
187
188impl teeny_core::model::RuntimeOp for Conv2dBnSiluGemmForward {
189    fn n_activation_inputs(&self) -> usize {
190        1
191    }
192
193    fn param_shapes(&self, input_shapes: &[&[usize]], output_shape: &[usize]) -> Vec<Vec<usize>> {
194        let c_in = input_shapes[0][1];
195        let c_out = output_shape[1];
196        // 1×1 weight: [C_OUT, C_IN] (no KH/KW dims needed — both are 1)
197        vec![vec![c_out, c_in], vec![c_out], vec![c_out]]
198    }
199
200    fn param_names(&self) -> &'static [&'static str] {
201        &["weight", "bn_scale", "bn_shift"]
202    }
203
204    fn pack_args(
205        &self,
206        inputs: &[(teeny_core::model::RawPtr, &[usize])],
207        params: &[teeny_core::model::RawPtr],
208        output: teeny_core::model::RawPtr,
209        output_shape: &[usize],
210        _output_row_stride: i32,
211        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
212    ) {
213        let input_shape = inputs[0].1;
214        let b = input_shape[0] as i32;
215        let c_in = input_shape[1] as i32;
216        let c_out = output_shape[1] as i32;
217        let m = (output_shape[2] * output_shape[3]) as i32; // OH * OW
218        visitor.visit_ptr(inputs[0].0); // x_ptr
219        visitor.visit_ptr(params[0]); // w_ptr
220        visitor.visit_ptr(params[1]); // bn_scale_ptr
221        visitor.visit_ptr(params[2]); // bn_shift_ptr
222        visitor.visit_ptr(output); // y_ptr
223        visitor.visit_i32(b);
224        visitor.visit_i32(c_in);
225        visitor.visit_i32(c_out);
226        visitor.visit_i32(m);
227    }
228
229    fn block(&self) -> [u32; 3] {
230        [128, 1, 1]
231    }
232
233    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
234        let b = output_shape[0];
235        let m = output_shape[2] * output_shape[3]; // OH * OW
236        let c_out = output_shape[1];
237        let pm = m.div_ceil(self.block_m as usize);
238        let pn = c_out.div_ceil(self.block_n as usize);
239        [(b * pm * pn) as u32, 1, 1]
240    }
241}
teeny_kernels/nn/fused/conv2d_bn_silu_gemm.rs

teeny_kernels/nn/fused/
conv2d_bn_silu_gemm.rs
