teeny_kernels/nn/mlp/
flatten.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 teeny_core::dtype::Num;
18use teeny_macros::kernel;
19use teeny_triton::triton::{Axis, PaddingOption, Triton};
20
21/// Copy a [B, N] tensor with arbitrary input strides to a contiguous row-major [B, N] output.
22///
23/// This is the forward pass of the flatten operation. In a neural network pipeline it is
24/// used to materialise a contiguous copy of a potentially non-contiguous activation tensor
25/// (e.g. when transitioning from convolutional to fully-connected layers). When the input is
26/// already row-major contiguous (stride_ib = N, stride_in = 1) this is a simple memcpy; when
27/// the input uses a different layout (e.g. column-major: stride_ib = 1, stride_in = B) the
28/// kernel performs the necessary reordering so that downstream kernels can assume unit strides.
29///
30/// Grid: one flat 1D pid that encodes (pid_b, pid_n) = (pid / num_pid_n, pid % num_pid_n).
31#[kernel]
32pub fn flatten_forward<T: Triton, D: Num, const BLOCK_B: i32, const BLOCK_N: i32>(
33    input_ptr: T::Pointer<D>,
34    output_ptr: T::Pointer<D>,
35    B: i32,
36    N: i32,
37    stride_ib: i32,
38    stride_in: i32,
39) {
40    let pid = T::program_id(Axis::X);
41    let num_pid_n = T::cdiv(N, BLOCK_N);
42    let pid_b = pid / num_pid_n;
43    let pid_n = pid % num_pid_n;
44
45    let input_desc = T::make_tensor_descriptor(
46        input_ptr,
47        &[B, N],
48        &[stride_ib, stride_in],
49        &[BLOCK_B, BLOCK_N],
50        Some(PaddingOption::Zero),
51    );
52    let output_desc = T::make_tensor_descriptor(
53        output_ptr,
54        &[B, N],
55        &[N, 1],
56        &[BLOCK_B, BLOCK_N],
57        Some(PaddingOption::Zero),
58    );
59
60    let b_off = pid_b * BLOCK_B;
61    let n_off = pid_n * BLOCK_N;
62    let tile = T::load_tensor_descriptor(input_desc, &[b_off, n_off]);
63    T::store_tensor_descriptor(output_desc, &[b_off, n_off], tile);
64}
65
66/// Copy a contiguous row-major [B, N] gradient back to an output buffer with arbitrary strides.
67///
68/// This is the backward pass of the flatten operation. The gradient dy arrives as a contiguous
69/// row-major tensor (the upstream gradient matches the contiguous forward output); this kernel
70/// writes it back to dx using the original input strides so that the gradient is in the same
71/// memory layout as the forward input. When the forward input was already row-major contiguous
72/// this is again a simple memcpy; when the forward input used a different layout the kernel
73/// performs the inverse reordering.
74#[kernel]
75pub fn flatten_backward<T: Triton, D: Num, const BLOCK_B: i32, const BLOCK_N: i32>(
76    dy_ptr: T::Pointer<D>,
77    dx_ptr: T::Pointer<D>,
78    B: i32,
79    N: i32,
80    stride_dxb: i32,
81    stride_dxn: i32,
82) {
83    let pid = T::program_id(Axis::X);
84    let num_pid_n = T::cdiv(N, BLOCK_N);
85    let pid_b = pid / num_pid_n;
86    let pid_n = pid % num_pid_n;
87
88    let dy_desc = T::make_tensor_descriptor(
89        dy_ptr,
90        &[B, N],
91        &[N, 1],
92        &[BLOCK_B, BLOCK_N],
93        Some(PaddingOption::Zero),
94    );
95    let dx_desc = T::make_tensor_descriptor(
96        dx_ptr,
97        &[B, N],
98        &[stride_dxb, stride_dxn],
99        &[BLOCK_B, BLOCK_N],
100        Some(PaddingOption::Zero),
101    );
102
103    let b_off = pid_b * BLOCK_B;
104    let n_off = pid_n * BLOCK_N;
105    let tile = T::load_tensor_descriptor(dy_desc, &[b_off, n_off]);
106    T::store_tensor_descriptor(dx_desc, &[b_off, n_off], tile);
107}
108
109impl<D: Num + Send + Sync + 'static> teeny_core::model::RuntimeOp for FlattenForward<D> {
110    fn n_activation_inputs(&self) -> usize {
111        1
112    }
113
114    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
115        Vec::new()
116    }
117
118    fn pack_args(
119        &self,
120        inputs: &[(teeny_core::model::RawPtr, &[usize])],
121        _params: &[teeny_core::model::RawPtr],
122        output: teeny_core::model::RawPtr,
123        output_shape: &[usize],
124        _output_row_stride: i32,
125        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
126    ) {
127        // kernel args: input_ptr, output_ptr, B, N, stride_ib, stride_in
128        // output_shape = [B, N] where N = product of all non-batch input dims
129        let b = output_shape[0] as i32;
130        let n = output_shape[1] as i32;
131        // Input is row-major contiguous: stride_ib = N, stride_in = 1
132        visitor.visit_ptr(inputs[0].0);
133        visitor.visit_ptr(output);
134        visitor.visit_i32(b);
135        visitor.visit_i32(n);
136        visitor.visit_i32(n); // stride_ib = N
137        visitor.visit_i32(1); // stride_in = 1
138    }
139
140    fn block(&self) -> [u32; 3] {
141        [128, 1, 1]
142    }
143
144    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
145        // pid encodes (pid_b, pid_n) = (pid / num_pid_n, pid % num_pid_n)
146        let pb = output_shape[0].div_ceil(self.block_b as usize);
147        let pn = output_shape[1].div_ceil(self.block_n as usize);
148        [(pb * pn) as u32, 1, 1]
149    }
150}
151
152pub struct FlattenOp<'a, T: Num> {
153    pub forward: FlattenForward<T>,
154    pub backward: FlattenBackward<T>,
155    _marker: core::marker::PhantomData<&'a ()>,
156}
teeny_kernels/nn/mlp/flatten.rs

teeny_kernels/nn/mlp/
flatten.rs
