teeny_kernels/nn/tensor/
channel_cat.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 concatenation forward — NC layout.
26///
27/// Copies one input chunk into the channel region `[chunk_offset, chunk_offset + chunk_c)`
28/// of a wide output NC tensor. Call once per input tensor to build the full concat.
29///
30/// Index mapping:
31///   `y[n * c_total + chunk_offset + ci] = x[n * chunk_c + ci]`
32///
33/// This is the structural inverse of `channel_chunk_forward`. The backward
34/// of this op is `channel_chunk_forward` with the same parameters.
35///
36/// Grid: `n_spatial * cdiv(chunk_c, BLOCK_SIZE)` CTAs.
37#[kernel]
38pub fn channel_cat_forward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
39    x_ptr: T::Pointer<D>, // one input: n_spatial * chunk_c  (narrow NC)
40    y_ptr: T::Pointer<D>, // output:    n_spatial * c_total  (wide NC)
41    chunk_c: i32,         // input channels for this chunk
42    c_total: i32,         // total output channels
43    chunk_offset: i32,    // first output channel index: k * chunk_c
44) where
45    T::I32Tensor: types::Tensor<i32, 1>,
46    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
47    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
48{
49    let pid = T::program_id(Axis::X);
50    let num_c_tiles = T::cdiv(chunk_c, BLOCK_SIZE);
51    let pid_n = pid / num_c_tiles;
52    let ci_tile = pid % num_c_tiles;
53    let ci_start = ci_tile * BLOCK_SIZE;
54
55    let ci_offsets = T::arange(0, BLOCK_SIZE) + ci_start;
56    let in_bounds = ci_offsets.lt(chunk_c);
57
58    let in_offsets = ci_offsets + (pid_n * chunk_c);
59    let out_offsets = ci_offsets + (pid_n * c_total + chunk_offset);
60
61    let x = T::load(
62        x_ptr.add_offsets(in_offsets),
63        Some(in_bounds),
64        None,
65        &[],
66        None,
67        None,
68        None,
69        false,
70    );
71    T::store(
72        y_ptr.add_offsets(out_offsets),
73        x,
74        Some(in_bounds),
75        &[],
76        None,
77        None,
78    );
79}
80
81/// Channel-wise cat backward — extracts the gradient slice for one input
82/// from the combined upstream gradient tensor.
83///
84/// Index mapping:
85///   `dx[n * chunk_c + ci] = dy[n * c_total + chunk_offset + ci]`
86///
87/// No atomic operations are required: this kernel reads from a specific
88/// disjoint channel range of `dy` and writes to its own output buffer.
89///
90/// Grid: same as `channel_cat_forward`.
91#[kernel]
92pub fn channel_cat_backward<T: Triton, D: Num, const BLOCK_SIZE: i32>(
93    dy_ptr: T::Pointer<D>, // upstream grad: n_spatial * c_total  (wide NC)
94    dx_ptr: T::Pointer<D>, // input grad:    n_spatial * chunk_c  (narrow NC)
95    chunk_c: i32,
96    c_total: i32,
97    chunk_offset: i32,
98) where
99    T::I32Tensor: types::Tensor<i32, 1>,
100    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
101    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
102{
103    let pid = T::program_id(Axis::X);
104    let num_c_tiles = T::cdiv(chunk_c, BLOCK_SIZE);
105    let pid_n = pid / num_c_tiles;
106    let ci_tile = pid % num_c_tiles;
107    let ci_start = ci_tile * BLOCK_SIZE;
108
109    let ci_offsets = T::arange(0, BLOCK_SIZE) + ci_start;
110    let in_bounds = ci_offsets.lt(chunk_c);
111
112    let dy_offsets = ci_offsets + (pid_n * c_total + chunk_offset);
113    let dx_offsets = ci_offsets + (pid_n * chunk_c);
114
115    let grad = T::load(
116        dy_ptr.add_offsets(dy_offsets),
117        Some(in_bounds),
118        None,
119        &[],
120        None,
121        None,
122        None,
123        false,
124    );
125    T::store(
126        dx_ptr.add_offsets(dx_offsets),
127        grad,
128        Some(in_bounds),
129        &[],
130        None,
131        None,
132    );
133}
134
135pub struct ChannelCatOp<'a, D: Num> {
136    pub forward: ChannelCatForward<D>,
137    pub backward: ChannelCatBackward<D>,
138    _marker: PhantomData<&'a ()>,
139}
140
141/// Combined runtime op for channel-cat forward + backward.
142///
143/// Forward: launches one `channel_cat_forward` call per input chunk, all
144/// writing to disjoint channel ranges of the shared output buffer.
145///
146/// Backward: launches one `channel_cat_backward` call per input chunk,
147/// each extracting the gradient slice for its chunk from the full concat
148/// upstream gradient.
149pub struct ChannelCatRuntimeOp<D: Num + Send + Sync + 'static> {
150    fwd: ChannelCatForward<D>,
151    bwd: ChannelCatBackward<D>,
152    n_inputs: usize,
153}
154
155impl<D: Num + Send + Sync + 'static> ChannelCatRuntimeOp<D> {
156    pub fn new(block_size: i32, n_inputs: usize) -> Self {
157        Self {
158            fwd: ChannelCatForward::<D>::new(block_size),
159            bwd: ChannelCatBackward::<D>::new(block_size),
160            n_inputs,
161        }
162    }
163
164    /// Returns the compiled forward kernel source for embedding in KernelExecutable.
165    pub fn forward_source(&self) -> &str {
166        &self.fwd.source
167    }
168    /// Returns the compiled backward kernel source.
169    pub fn backward_source(&self) -> &str {
170        &self.bwd.source
171    }
172    /// Returns the kernel name (for the forward kernel).
173    pub fn kernel_name(&self) -> &str {
174        self.fwd.name
175    }
176}
177
178impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for ChannelCatRuntimeOp<D> {
179    fn n_activation_inputs(&self) -> usize {
180        self.n_inputs
181    }
182
183    fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
184        Vec::new()
185    }
186
187    // pack_args is used only as a fallback (n_launches==1 never occurs for ChannelCat).
188    fn pack_args(
189        &self,
190        inputs: &[(teeny_core::model::RawPtr, &[usize])],
191        params: &[teeny_core::model::RawPtr],
192        output: teeny_core::model::RawPtr,
193        output_shape: &[usize],
194        output_row_stride: i32,
195        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
196    ) {
197        self.pack_args_for_launch(
198            0,
199            inputs,
200            params,
201            output,
202            output_shape,
203            output_row_stride,
204            visitor,
205        );
206    }
207
208    fn n_launches(&self) -> usize {
209        self.n_inputs
210    }
211
212    fn pack_args_for_launch(
213        &self,
214        launch_idx: usize,
215        inputs: &[(teeny_core::model::RawPtr, &[usize])],
216        _params: &[teeny_core::model::RawPtr],
217        output: teeny_core::model::RawPtr,
218        _output_shape: &[usize],
219        _output_row_stride: i32,
220        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
221    ) {
222        // Each input is NCHW [B, C_i, H, W]. Treat as NC where N=B, C=C_i*H*W.
223        let chunk_offset: i32 = inputs[..launch_idx]
224            .iter()
225            .map(|(_, s)| (s[1] * s[2] * s[3]) as i32)
226            .sum();
227        let x_ptr = inputs[launch_idx].0;
228        let input_shape = inputs[launch_idx].1;
229        let chunk_c = (input_shape[1] * input_shape[2] * input_shape[3]) as i32;
230        let c_total: i32 = inputs
231            .iter()
232            .map(|(_, s)| (s[1] * s[2] * s[3]) as i32)
233            .sum();
234
235        visitor.visit_ptr(x_ptr);
236        visitor.visit_ptr(output);
237        visitor.visit_i32(chunk_c);
238        visitor.visit_i32(c_total);
239        visitor.visit_i32(chunk_offset);
240    }
241
242    fn grid_for_launch(
243        &self,
244        launch_idx: usize,
245        input_shapes: &[&[usize]],
246        _output_shape: &[usize],
247    ) -> [u32; 3] {
248        let s = input_shapes[launch_idx];
249        let n_spatial = s[0];
250        let chunk_c = s[1] * s[2] * s[3];
251        let num_tiles = chunk_c.div_ceil(self.fwd.block_size as usize);
252        [(n_spatial * num_tiles) as u32, 1, 1]
253    }
254
255    fn block(&self) -> [u32; 3] {
256        [self.fwd.block_size as u32, 1, 1]
257    }
258
259    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
260        let n_spatial = output_shape[0];
261        let c = output_shape[1] * output_shape[2] * output_shape[3];
262        let num_tiles = c.div_ceil(self.fwd.block_size as usize);
263        [(n_spatial * num_tiles) as u32, 1, 1]
264    }
265
266    #[cfg(feature = "training")]
267    fn has_backward(&self) -> bool {
268        true
269    }
270
271    #[cfg(feature = "training")]
272    fn n_backward_launches(&self) -> usize {
273        self.n_inputs
274    }
275
276    #[cfg(feature = "training")]
277    fn pack_backward_args(
278        &self,
279        inputs: &[(teeny_core::model::RawPtr, &[usize])],
280        params: &[teeny_core::model::RawPtr],
281        output: teeny_core::model::RawPtr,
282        output_shape: &[usize],
283        grad_output: teeny_core::model::RawPtr,
284        grad_output_row_stride: i32,
285        grad_inputs: &[teeny_core::model::RawPtr],
286        grad_params: &[teeny_core::model::RawPtr],
287        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
288    ) {
289        self.pack_backward_args_for_launch(
290            0,
291            inputs,
292            params,
293            output,
294            output_shape,
295            grad_output,
296            grad_output_row_stride,
297            grad_inputs,
298            grad_params,
299            visitor,
300        );
301    }
302
303    #[cfg(feature = "training")]
304    #[allow(clippy::too_many_arguments)]
305    fn pack_backward_args_for_launch(
306        &self,
307        launch_idx: usize,
308        inputs: &[(teeny_core::model::RawPtr, &[usize])],
309        _params: &[teeny_core::model::RawPtr],
310        _output: teeny_core::model::RawPtr,
311        _output_shape: &[usize],
312        grad_output: teeny_core::model::RawPtr,
313        _grad_output_row_stride: i32,
314        grad_inputs: &[teeny_core::model::RawPtr],
315        _grad_params: &[teeny_core::model::RawPtr],
316        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
317    ) {
318        let chunk_offset: i32 = inputs[..launch_idx]
319            .iter()
320            .map(|(_, s)| (s[1] * s[2] * s[3]) as i32)
321            .sum();
322        let input_shape = inputs[launch_idx].1;
323        let chunk_c = (input_shape[1] * input_shape[2] * input_shape[3]) as i32;
324        let c_total: i32 = inputs
325            .iter()
326            .map(|(_, s)| (s[1] * s[2] * s[3]) as i32)
327            .sum();
328
329        visitor.visit_ptr(grad_output);
330        visitor.visit_ptr(grad_inputs[launch_idx]);
331        visitor.visit_i32(chunk_c);
332        visitor.visit_i32(c_total);
333        visitor.visit_i32(chunk_offset);
334    }
335
336    #[cfg(feature = "training")]
337    fn backward_grid(&self, input_shapes: &[&[usize]], output_shape: &[usize]) -> [u32; 3] {
338        self.backward_grid_for_launch(0, input_shapes, output_shape)
339    }
340
341    #[cfg(feature = "training")]
342    fn backward_grid_for_launch(
343        &self,
344        launch_idx: usize,
345        input_shapes: &[&[usize]],
346        output_shape: &[usize],
347    ) -> [u32; 3] {
348        self.grid_for_launch(launch_idx, input_shapes, output_shape)
349    }
350}
teeny_kernels/nn/tensor/channel_cat.rs

teeny_kernels/nn/tensor/
channel_cat.rs
