teeny_core/device/
program.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 alloc::{format, string::String};
18use sha2::{Digest, Sha256};
19
20/// Receives each argument of a kernel in order. Implement this on a backend-
21/// specific "packer" struct (e.g. a CUDA arg-pointer array builder).
22pub trait ArgVisitor {
23    /// Visits a raw pointer argument.
24    fn visit_ptr(&mut self, ptr: *mut core::ffi::c_void);
25    /// Visits a `bool` argument.
26    fn visit_bool(&mut self, val: bool);
27    /// Visits an `i8` argument.
28    fn visit_i8(&mut self, val: i8);
29    /// Visits an `i16` argument.
30    fn visit_i16(&mut self, val: i16);
31    /// Visits an `i32` argument.
32    fn visit_i32(&mut self, val: i32);
33    /// Visits an `i64` argument.
34    fn visit_i64(&mut self, val: i64);
35    /// Visits a `u8` argument.
36    fn visit_u8(&mut self, val: u8);
37    /// Visits a `u16` argument.
38    fn visit_u16(&mut self, val: u16);
39    /// Visits a `u32` argument.
40    fn visit_u32(&mut self, val: u32);
41    /// Visits a `u64` argument.
42    fn visit_u64(&mut self, val: u64);
43    /// Visits an `f32` argument.
44    fn visit_f32(&mut self, val: f32);
45    /// Visits an `f64` argument.
46    fn visit_f64(&mut self, val: f64);
47}
48
49/// Implemented by each concrete argument type; dispatches to the right
50/// `ArgVisitor` method.
51pub trait KernelArg {
52    /// Visits `self` on `visitor`.
53    fn visit<V: ArgVisitor>(&self, visitor: &mut V);
54}
55
56impl<T> KernelArg for *mut T {
57    #[inline]
58    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
59        visitor.visit_ptr(*self as *mut core::ffi::c_void);
60    }
61}
62
63impl KernelArg for bool {
64    #[inline]
65    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
66        visitor.visit_bool(*self);
67    }
68}
69
70impl KernelArg for i8 {
71    #[inline]
72    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
73        visitor.visit_i8(*self);
74    }
75}
76
77impl KernelArg for i16 {
78    #[inline]
79    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
80        visitor.visit_i16(*self);
81    }
82}
83
84impl KernelArg for i32 {
85    #[inline]
86    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
87        visitor.visit_i32(*self);
88    }
89}
90
91impl KernelArg for i64 {
92    #[inline]
93    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
94        visitor.visit_i64(*self);
95    }
96}
97
98impl KernelArg for u8 {
99    #[inline]
100    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
101        visitor.visit_u8(*self);
102    }
103}
104
105impl KernelArg for u16 {
106    #[inline]
107    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
108        visitor.visit_u16(*self);
109    }
110}
111
112impl KernelArg for u32 {
113    #[inline]
114    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
115        visitor.visit_u32(*self);
116    }
117}
118
119impl KernelArg for u64 {
120    #[inline]
121    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
122        visitor.visit_u64(*self);
123    }
124}
125
126impl KernelArg for f32 {
127    #[inline]
128    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
129        visitor.visit_f32(*self);
130    }
131}
132
133impl KernelArg for f64 {
134    #[inline]
135    fn visit<V: ArgVisitor>(&self, visitor: &mut V) {
136        visitor.visit_f64(*self);
137    }
138}
139
140/// Implemented for tuples of `KernelArg`s; visits each element in order.
141/// The proc macro generates `type Args<'a> = (A, B, C, ...)` and the blanket
142/// tuple impls below make that satisfy this bound automatically.
143pub trait KernelArgs {
144    /// Visits each element of `self` in order.
145    fn visit_args<V: ArgVisitor>(&self, visitor: &mut V);
146}
147
148impl KernelArgs for () {
149    #[inline]
150    fn visit_args<V: ArgVisitor>(&self, _visitor: &mut V) {}
151}
152
153macro_rules! impl_kernel_args {
154    ($( $T:ident : $idx:tt ),+) => {
155        impl<$($T: KernelArg),+> KernelArgs for ($($T,)+) {
156            #[inline]
157            fn visit_args<V: ArgVisitor>(&self, visitor: &mut V) {
158                $( self.$idx.visit(visitor); )+
159            }
160        }
161    };
162}
163
164impl_kernel_args!(A:0);
165impl_kernel_args!(A:0, B:1);
166impl_kernel_args!(A:0, B:1, C:2);
167impl_kernel_args!(A:0, B:1, C:2, D:3);
168impl_kernel_args!(A:0, B:1, C:2, D:3, E:4);
169impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5);
170impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6);
171impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7);
172impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8);
173impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9);
174impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10);
175impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11);
176impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12);
177impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13);
178impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13, O:14);
179impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13, O:14, P:15);
180impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13, O:14, P:15, Q:16);
181impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13, O:14, P:15, Q:16, R:17);
182impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13, O:14, P:15, Q:16, R:17, S:18);
183impl_kernel_args!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11, M:12, N:13, O:14, P:15, Q:16, R:17, S:18, Ty:19);
184
185/// A kernel definition: its source, name, and identity.
186pub trait Kernel {
187    /// This kernel's launch-argument tuple type.
188    type Args<'a>: KernelArgs;
189
190    /// A content hash of this kernel's source, used as a cache key.
191    fn id(&self) -> String {
192        let mut hasher = Sha256::new();
193        hasher.update(self.source().as_bytes());
194        hasher
195            .finalize()
196            .iter()
197            .map(|b| format!("{:02x}", b))
198            .collect()
199    }
200
201    /// This kernel's name.
202    fn name(&self) -> &str;
203
204    /// This kernel's full source (DSL + kernel body + entry point wrapper).
205    fn source(&self) -> &str;
206
207    /// This kernel's body source, without the entry-point wrapper.
208    fn kernel_source(&self) -> &str;
209
210    /// Returns the Rust source of the generated C-ABI entry-point wrapper function.
211    fn entry_point_source(&self) -> &str;
212
213    /// Returns the PTX symbol name for this kernel: `"{name}_entry_point"`.
214    fn entry_point_name(&self) -> String {
215        format!("{}_entry_point", self.name())
216    }
217}
218
219/// A kernel `K` compiled and loaded onto a device, ready to launch.
220pub trait Program<'a, K: Kernel>: Sized {}
teeny_core/device/program.rs

teeny_core/device/
program.rs
