teeny_kernels/nn/tensor/
channel_chunk.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
17use core::marker::PhantomData;
18use teeny_core::dtype::Num;
19use teeny_macros::kernel;
20use teeny_triton::triton::{
21    types::{AddOffsets, Comparison},
22    *,
23};
24
25/// Channel-wise chunk (split) forward — NC layout.
26///
27/// Extracts one contiguous channel slice `[chunk_offset, chunk_offset + chunk_c)`
28/// from a wide NC tensor and writes it into a narrow NC output tensor.
29///
30/// Index mapping (NC layout, index = `n * C + c`):
31///   `y[n * chunk_c + ci] = x[n * c_total + chunk_offset + ci]`
32///
33/// This is the structural inverse of `channel_cat_forward`. The backward
34/// of this op is `channel_cat_forward` with the same parameters.
35///
36/// Grid: `n_spatial * cdiv(chunk_c, BLOCK_SIZE)` CTAs.
37/// `pid` is decoded into `(pid_n, ci_tile)` via scalar integer division —
38/// no tensor-level division or modulo is required.
39#[kernel]
40pub fn channel_chunk_forward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
41    x_ptr: T::Pointer<D>, // input:  n_spatial * c_total  (wide NC)
42    y_ptr: T::Pointer<D>, // output: n_spatial * chunk_c  (narrow NC)
43    c_total: i32,         // total input channels
44    chunk_c: i32,         // output channels per chunk
45    chunk_offset: i32,    // first input channel index: k * chunk_c
46) where
47    T::I32Tensor: types::Tensor<i32, 1>,
48    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
49    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
50{
51    let pid = T::program_id(Axis::X);
52    let num_c_tiles = T::cdiv(chunk_c, BLOCK_SIZE);
53    let pid_n = pid / num_c_tiles;
54    let ci_tile = pid % num_c_tiles;
55    let ci_start = ci_tile * BLOCK_SIZE;
56
57    let ci_offsets = T::arange(0, BLOCK_SIZE) + ci_start;
58    let in_bounds = ci_offsets.lt(chunk_c);
59
60    let in_offsets = ci_offsets + (pid_n * c_total + chunk_offset);
61    let out_offsets = ci_offsets + (pid_n * chunk_c);
62
63    let x = T::load(
64        x_ptr.add_offsets(in_offsets),
65        Some(in_bounds),
66        None,
67        &[],
68        None,
69        None,
70        None,
71        false,
72    );
73    T::store(
74        y_ptr.add_offsets(out_offsets),
75        x,
76        Some(in_bounds),
77        &[],
78        None,
79        None,
80    );
81}
82
83/// Channel-wise chunk backward — propagates the gradient of one chunk
84/// output back into the full-width input gradient tensor.
85///
86/// Index mapping:
87///   `dx[n * c_total + chunk_offset + ci] = dy[n * chunk_c + ci]`
88///
89/// No atomic operations are required: each chunk's backward writes to a
90/// disjoint channel range `[chunk_offset, chunk_offset + chunk_c)` of `dx`.
91///
92/// Grid: same as `channel_chunk_forward`.
93#[kernel]
94pub fn channel_chunk_backward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
95    dy_ptr: T::Pointer<D>, // upstream grad: n_spatial * chunk_c  (narrow NC)
96    dx_ptr: T::Pointer<D>, // input grad:    n_spatial * c_total  (wide NC)
97    c_total: i32,
98    chunk_c: i32,
99    chunk_offset: i32,
100) where
101    T::I32Tensor: types::Tensor<i32, 1>,
102    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
103    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
104{
105    let pid = T::program_id(Axis::X);
106    let num_c_tiles = T::cdiv(chunk_c, BLOCK_SIZE);
107    let pid_n = pid / num_c_tiles;
108    let ci_tile = pid % num_c_tiles;
109    let ci_start = ci_tile * BLOCK_SIZE;
110
111    let ci_offsets = T::arange(0, BLOCK_SIZE) + ci_start;
112    let in_bounds = ci_offsets.lt(chunk_c);
113
114    let dy_offsets = ci_offsets + (pid_n * chunk_c);
115    let dx_offsets = ci_offsets + (pid_n * c_total + chunk_offset);
116
117    let grad = T::load(
118        dy_ptr.add_offsets(dy_offsets),
119        Some(in_bounds),
120        None,
121        &[],
122        None,
123        None,
124        None,
125        false,
126    );
127    T::store(
128        dx_ptr.add_offsets(dx_offsets),
129        grad,
130        Some(in_bounds),
131        &[],
132        None,
133        None,
134    );
135}
136
137pub struct ChannelChunkOp<'a, D: Num> {
138    pub forward: ChannelChunkForward<D>,
139    pub backward: ChannelChunkBackward<D>,
140    _marker: PhantomData<&'a ()>,
141}
142
143/// Runtime op for channel_chunk forward + backward.
144///
145/// Single-launch: extracts one channel slice from a wide NCHW input.
146/// `chunk_c` and `chunk_offset` are fixed at graph construction time
147/// from `Op::ChannelChunk`.
148pub struct ChannelChunkRuntimeOp<D: Num + Send + Sync + 'static> {
149    fwd: ChannelChunkForward<D>,
150    bwd: ChannelChunkBackward<D>,
151    chunk_c: usize,
152    chunk_offset: usize,
153}
154
155impl<D: Num + Send + Sync + 'static> ChannelChunkRuntimeOp<D> {
156    pub fn new(block_size: i32, chunk_c: usize, chunk_offset: usize) -> Self {
157        Self {
158            fwd: ChannelChunkForward::<D>::new(block_size),
159            bwd: ChannelChunkBackward::<D>::new(block_size),
160            chunk_c,
161            chunk_offset,
162        }
163    }
164
165    pub fn forward_source(&self) -> &str {
166        &self.fwd.source
167    }
168    pub fn backward_source(&self) -> &str {
169        &self.bwd.source
170    }
171    pub fn kernel_name(&self) -> &str {
172        self.fwd.name
173    }
174}
175
176impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for ChannelChunkRuntimeOp<D> {
177    fn n_activation_inputs(&self) -> usize {
178        1
179    }
180
181    fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
182        Vec::new()
183    }
184
185    fn pack_args(
186        &self,
187        inputs: &[(teeny_core::model::RawPtr, &[usize])],
188        _params: &[teeny_core::model::RawPtr],
189        output: teeny_core::model::RawPtr,
190        _output_shape: &[usize],
191        _output_row_stride: i32,
192        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
193    ) {
194        // Input is NCHW [B, C_total, H, W]. Treat as NC: N=B, C=C_total*H*W.
195        let input_shape = inputs[0].1;
196        let c_total = (input_shape[1] * input_shape[2] * input_shape[3]) as i32;
197        let chunk_c = self.chunk_c as i32;
198        // Scale chunk_offset from channel index to NC-layout offset (×H×W).
199        let h = input_shape[2];
200        let w = input_shape[3];
201        let chunk_offset = (self.chunk_offset * h * w) as i32;
202
203        visitor.visit_ptr(inputs[0].0);
204        visitor.visit_ptr(output);
205        visitor.visit_i32(c_total);
206        visitor.visit_i32(chunk_c * (h as i32) * (w as i32)); // chunk_c in NC units
207        visitor.visit_i32(chunk_offset);
208    }
209
210    fn block(&self) -> [u32; 3] {
211        [self.fwd.block_size as u32, 1, 1]
212    }
213
214    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
215        // Grid over n_spatial * ceil(chunk_c_nc / block_size)
216        // output_shape = [B, chunk_c, H, W]
217        let n_spatial = output_shape[0];
218        let chunk_c_nc = output_shape[1] * output_shape[2] * output_shape[3];
219        let num_tiles = chunk_c_nc.div_ceil(self.fwd.block_size as usize);
220        [(n_spatial * num_tiles) as u32, 1, 1]
221    }
222
223    #[cfg(feature = "training")]
224    fn has_backward(&self) -> bool {
225        true
226    }
227
228    #[cfg(feature = "training")]
229    fn pack_backward_args(
230        &self,
231        inputs: &[(teeny_core::model::RawPtr, &[usize])],
232        _params: &[teeny_core::model::RawPtr],
233        _output: teeny_core::model::RawPtr,
234        _output_shape: &[usize],
235        grad_output: teeny_core::model::RawPtr,
236        _grad_output_row_stride: i32,
237        grad_inputs: &[teeny_core::model::RawPtr],
238        _grad_params: &[teeny_core::model::RawPtr],
239        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
240    ) {
241        // channel_chunk_backward: (dy_ptr, dx_ptr, c_total, chunk_c, chunk_offset)
242        // dy = grad_output (narrow [B, chunk_c, H, W])
243        // dx = grad_inputs[0] (wide [B, C_total, H, W])
244        let input_shape = inputs[0].1; // original input shape [B, C_total, H, W]
245        let h = input_shape[2];
246        let w = input_shape[3];
247        let c_total = (input_shape[1] * h * w) as i32;
248        let chunk_c_nc = (self.chunk_c * h * w) as i32;
249        let chunk_offset = (self.chunk_offset * h * w) as i32;
250
251        visitor.visit_ptr(grad_output);
252        visitor.visit_ptr(grad_inputs[0]);
253        visitor.visit_i32(c_total);
254        visitor.visit_i32(chunk_c_nc);
255        visitor.visit_i32(chunk_offset);
256    }
257
258    #[cfg(feature = "training")]
259    fn backward_grid(&self, input_shapes: &[&[usize]], _output_shape: &[usize]) -> [u32; 3] {
260        // Grid over input grad: n_spatial * ceil(chunk_c_nc / block_size)
261        let s = input_shapes[0];
262        let n_spatial = s[0];
263        let chunk_c_nc = self.chunk_c * s[2] * s[3];
264        let num_tiles = chunk_c_nc.div_ceil(self.fwd.block_size as usize);
265        [(n_spatial * num_tiles) as u32, 1, 1]
266    }
267}
teeny_kernels/nn/tensor/channel_chunk.rs

teeny_kernels/nn/tensor/
channel_chunk.rs
