1use std::sync::Arc;
18use teeny_core::{
19 graph::{DtypeRepr, Graph, Op, Shape},
20 model::{ExecutableOp, Lowering, LoweringMode, RuntimeOp},
21 utils::dag::Dag,
22};
23
24use crate::nn::{
25 activation::extra::{
26 LogSoftmaxBackward, LogSoftmaxForward, PreluForward, ShrinkRuntimeOp, SwishBackward,
27 SwishForward, ThresholdedReluRuntimeOp,
28 },
29 activation::{
30 elu::{
31 CeluForward, CeluForwardDispatch, EluForward, EluForwardDispatch, SeluForward,
32 SeluForwardDispatch,
33 },
34 gelu::{GeluForwardDispatch, MishForward, MishForwardDispatch},
35 hard::{
36 HardshrinkForward, HardshrinkForwardDispatch, HardsigmoidForward,
37 HardsigmoidForwardDispatch, HardswishForward, HardswishForwardDispatch,
38 HardtanhForward, HardtanhForwardDispatch, Relu6Forward, Relu6ForwardDispatch,
39 },
40 misc::{
41 LeakyReluForward, LeakyReluForwardDispatch, SoftplusForward, SoftplusForwardDispatch,
42 SoftshrinkForward, SoftshrinkForwardDispatch, SoftsignForward, SoftsignForwardDispatch,
43 ThresholdForward, ThresholdForwardDispatch,
44 },
45 relu::{ReluBackward, ReluForward},
46 sigmoid::{
47 LogsigmoidForward, LogsigmoidForwardDispatch, SigmoidForwardDispatch,
48 SiluForwardDispatch,
49 },
50 softmax::SoftmaxForward,
51 tanh::{TanhForward, TanhForwardDispatch, TanhshrinkForward, TanhshrinkForwardDispatch},
52 },
53 conv::{
54 conv1d::Conv1dForward,
55 conv2d::{Conv2dBackward, Conv2dBiasForward, Conv2dForward},
56 conv3d::Conv3dForward,
57 },
58 fused::{
59 conv2d_bn_silu::Conv2dBnSiluForward, conv2d_bn_silu_gemm::Conv2dBnSiluGemmForward,
60 conv2d_bn_silu_tiled::Conv2dBnSiluTiledForward,
61 },
62 mlp::{
63 flatten::FlattenForward,
64 linear::{LinearBackward, LinearForward},
65 },
66 norm::{
67 batchnorm::{BatchNorm2dNchwInferenceRuntimeOp, BatchNormForwardInference},
68 groupnorm::GroupNormForwardInference,
69 instancenorm::InstanceNormForwardInference,
70 layernorm::{LayerNormForwardInference, LayerNormForwardInferenceRuntimeOp},
71 rmsnorm::RmsNormForward,
72 },
73 pad::{
74 circular_pad1d::CircularPad1dForward, circular_pad2d::CircularPad2dForward,
75 circular_pad3d::CircularPad3dForward, constant_pad1d::ConstantPad1dForward,
76 constant_pad2d::ConstantPad2dForward, constant_pad3d::ConstantPad3dForward,
77 reflection_pad1d::ReflectionPad1dForward, reflection_pad2d::ReflectionPad2dForward,
78 reflection_pad3d::ReflectionPad3dForward, replication_pad1d::ReplicationPad1dForward,
79 replication_pad2d::ReplicationPad2dForward, replication_pad3d::ReplicationPad3dForward,
80 },
81 pool::{
82 avgpool1d::Avgpool1dForward,
83 avgpool2d::Avgpool2dForward,
84 avgpool3d::Avgpool3dForward,
85 lppool1d::Lppool1dForward,
86 lppool2d::Lppool2dForward,
87 lppool3d::Lppool3dForward,
88 maxpool1d::Maxpool1dForward,
89 maxpool2d::{Maxpool2dBackward, Maxpool2dForward},
90 maxpool3d::Maxpool3dForward,
91 },
92 tensor::{
93 channel_bias_add::{ChannelBiasAddRuntimeOp, NchwBiasAddRuntimeOp},
94 channel_cat::ChannelCatRuntimeOp,
95 channel_chunk::ChannelChunkRuntimeOp,
96 elemwise_add::{ElemwiseAddBackward, ElemwiseAddForward},
97 elemwise_binary::{
98 ClipRuntimeOp, ElemwiseDivBackward, ElemwiseDivForward, ElemwiseEqualForward,
99 ElemwiseFmodForward, ElemwiseGreaterEqualForward, ElemwiseGreaterForward,
100 ElemwiseLessEqualForward, ElemwiseLessForward, ElemwiseMaxBackward, ElemwiseMaxForward,
101 ElemwiseMeanBackward, ElemwiseMeanForward, ElemwiseMinBackward, ElemwiseMinForward,
102 ElemwiseMulBackward, ElemwiseMulForward, ElemwisePowBackward, ElemwisePowForward,
103 ElemwiseSubBackward, ElemwiseSubForward, ElemwiseSumBackward, ElemwiseSumForward,
104 ElemwiseWhereBackward, ElemwiseWhereForward,
105 },
106 elemwise_unary::{
107 ElemwiseAbsBackward, ElemwiseAbsForward, ElemwiseAcosBackward, ElemwiseAcosForward,
108 ElemwiseAcoshBackward, ElemwiseAcoshForward, ElemwiseAsinBackward, ElemwiseAsinForward,
109 ElemwiseAsinhBackward, ElemwiseAsinhForward, ElemwiseAtanBackward, ElemwiseAtanForward,
110 ElemwiseAtanhBackward, ElemwiseAtanhForward, ElemwiseCeilForward, ElemwiseCosBackward,
111 ElemwiseCosForward, ElemwiseCoshBackward, ElemwiseCoshForward, ElemwiseErfBackward,
112 ElemwiseErfForward, ElemwiseExpBackward, ElemwiseExpForward, ElemwiseFloorForward,
113 ElemwiseIsnanForward, ElemwiseLogBackward, ElemwiseLogForward, ElemwiseNegBackward,
114 ElemwiseNegForward, ElemwiseReciprocalBackward, ElemwiseReciprocalForward,
115 ElemwiseSignForward, ElemwiseSinBackward, ElemwiseSinForward, ElemwiseSinhBackward,
116 ElemwiseSinhForward, ElemwiseSqrtBackward, ElemwiseSqrtForward, ElemwiseTanBackward,
117 ElemwiseTanForward,
118 },
119 reduction::{
120 CumProdForward, CumSumForward, GlobalAvgPoolForward, GlobalMaxPoolForward,
121 ReduceL1Forward, ReduceL2Forward, ReduceLogSumExpForward, ReduceLogSumForward,
122 ReduceMaxForward, ReduceMeanForward, ReduceMinForward, ReduceProdForward,
123 ReduceSumForward, ReduceSumSquareForward,
124 },
125 upsample_nearest2d::{UpsampleNearest2dBackward, UpsampleNearest2dForward},
126 },
127};
128
129use crate::math::gemm::MatMulRuntimeOp;
130
131use crate::errors::Result;
132
133#[cfg(feature = "training")]
134use crate::nn::norm::batchnorm::{
135 BatchNorm2dNchwBackward, BatchNormNormalizeForward, BatchNormNormalizeRuntimeOp,
136 BatchNormStatsForward, BatchNormStatsRuntimeOp,
137};
138
139macro_rules! make_num_kernel {
152 ($K:ident ($($arg:expr),*), $node:expr) => {{
153 let (name, ks, rop) = match $node.dtype {
154 DtypeRepr::F32 => { let k = $K::<f32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
155 DtypeRepr::F64 => { let k = $K::<f64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
156 DtypeRepr::I8 => { let k = $K::<i8>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
157 DtypeRepr::I16 => { let k = $K::<i16>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
158 DtypeRepr::I32 => { let k = $K::<i32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
159 DtypeRepr::I64 => { let k = $K::<i64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
160 DtypeRepr::U8 => { let k = $K::<u8>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
161 DtypeRepr::U16 => { let k = $K::<u16>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
162 DtypeRepr::U32 => { let k = $K::<u32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
163 DtypeRepr::U64 => { let k = $K::<u64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
164 other => return Err(anyhow::anyhow!("{:?} is not a supported Num dtype for {}", other, stringify!($K))),
165 };
166 Box::new(KernelExecutable {
167 entry_point: format!("{}_entry_point", name),
168 name,
169 kernel_source: ks,
170 shape: $node.shape.clone(),
171 dtype: $node.dtype,
172 #[cfg(feature = "training")]
173 backward_kernel_source: String::new(),
174 #[cfg(feature = "training")]
175 backward_entry_point: String::new(),
176 runtime_op: rop,
177 })
178 }};
179 ($K:ident ($($arg:expr),*), $Bwd:ident ($($barg:expr),*), $node:expr) => {{
181 let (name, ks, rop) = match $node.dtype {
182 DtypeRepr::F32 => { let k = $K::<f32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
183 DtypeRepr::F64 => { let k = $K::<f64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
184 DtypeRepr::I8 => { let k = $K::<i8>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
185 DtypeRepr::I16 => { let k = $K::<i16>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
186 DtypeRepr::I32 => { let k = $K::<i32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
187 DtypeRepr::I64 => { let k = $K::<i64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
188 DtypeRepr::U8 => { let k = $K::<u8>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
189 DtypeRepr::U16 => { let k = $K::<u16>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
190 DtypeRepr::U32 => { let k = $K::<u32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
191 DtypeRepr::U64 => { let k = $K::<u64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
192 other => return Err(anyhow::anyhow!("{:?} is not a supported Num dtype for {}", other, stringify!($K))),
193 };
194 #[cfg(feature = "training")]
195 let (bwd_name, bwd_ks) = match $node.dtype {
196 DtypeRepr::F32 => { let k = $Bwd::<f32>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
197 DtypeRepr::F64 => { let k = $Bwd::<f64>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
198 DtypeRepr::I8 => { let k = $Bwd::<i8>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
199 DtypeRepr::I16 => { let k = $Bwd::<i16>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
200 DtypeRepr::I32 => { let k = $Bwd::<i32>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
201 DtypeRepr::I64 => { let k = $Bwd::<i64>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
202 DtypeRepr::U8 => { let k = $Bwd::<u8>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
203 DtypeRepr::U16 => { let k = $Bwd::<u16>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
204 DtypeRepr::U32 => { let k = $Bwd::<u32>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
205 DtypeRepr::U64 => { let k = $Bwd::<u64>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
206 other => return Err(anyhow::anyhow!("{:?} is not a supported Num dtype for {}", other, stringify!($Bwd))),
207 };
208 Box::new(KernelExecutable {
209 entry_point: format!("{}_entry_point", name),
210 name,
211 kernel_source: ks,
212 shape: $node.shape.clone(),
213 dtype: $node.dtype,
214 #[cfg(feature = "training")]
215 backward_kernel_source: bwd_ks,
216 #[cfg(feature = "training")]
217 backward_entry_point: format!("{}_entry_point", bwd_name),
218 runtime_op: rop,
219 })
220 }};
221}
222
223macro_rules! make_float_kernel {
226 ($K:ident ($($arg:expr),*), $node:expr) => {{
227 let (name, ks, rop) = match $node.dtype {
228 DtypeRepr::F32 => { let k = $K::<f32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
229 DtypeRepr::F64 => { let k = $K::<f64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
230 other => return Err(anyhow::anyhow!("{:?} is not a Float dtype for {}", other, stringify!($K))),
231 };
232 Box::new(KernelExecutable {
233 entry_point: format!("{}_entry_point", name),
234 name,
235 kernel_source: ks,
236 shape: $node.shape.clone(),
237 dtype: $node.dtype,
238 #[cfg(feature = "training")]
239 backward_kernel_source: String::new(),
240 #[cfg(feature = "training")]
241 backward_entry_point: String::new(),
242 runtime_op: rop,
243 })
244 }};
245 ($K:ident ($($arg:expr),*), $Bwd:ident ($($barg:expr),*), $node:expr) => {{
247 let (name, ks, rop) = match $node.dtype {
248 DtypeRepr::F32 => { let k = $K::<f32>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
249 DtypeRepr::F64 => { let k = $K::<f64>::new($($arg),*); let nm = k.name.to_string(); let src = k.source.clone(); let r: Arc<dyn RuntimeOp> = Arc::new(k); (nm, src, r) }
250 other => return Err(anyhow::anyhow!("{:?} is not a Float dtype for {}", other, stringify!($K))),
251 };
252 #[cfg(feature = "training")]
253 let (bwd_name, bwd_ks) = match $node.dtype {
254 DtypeRepr::F32 => { let k = $Bwd::<f32>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
255 DtypeRepr::F64 => { let k = $Bwd::<f64>::new($($barg),*); (k.name.to_string(), k.source.clone()) }
256 other => return Err(anyhow::anyhow!("{:?} is not a Float dtype for {}", other, stringify!($Bwd))),
257 };
258 Box::new(KernelExecutable {
259 entry_point: format!("{}_entry_point", name),
260 name,
261 kernel_source: ks,
262 shape: $node.shape.clone(),
263 dtype: $node.dtype,
264 #[cfg(feature = "training")]
265 backward_kernel_source: bwd_ks,
266 #[cfg(feature = "training")]
267 backward_entry_point: format!("{}_entry_point", bwd_name),
268 runtime_op: rop,
269 })
270 }};
271}
272
273fn exec_from(
276 shape: Shape,
277 dtype: DtypeRepr,
278 inst: teeny_core::model::KernelInstance,
279) -> Box<KernelExecutable> {
280 Box::new(KernelExecutable {
281 entry_point: format!("{}_entry_point", inst.name),
282 name: inst.name,
283 kernel_source: inst.source,
284 shape,
285 dtype,
286 #[cfg(feature = "training")]
287 backward_kernel_source: inst
288 .backward
289 .as_ref()
290 .map(|b| b.source.clone())
291 .unwrap_or_default(),
292 #[cfg(feature = "training")]
293 backward_entry_point: inst
294 .backward
295 .as_ref()
296 .map(|b| format!("{}_entry_point", b.name))
297 .unwrap_or_default(),
298 runtime_op: inst.runtime_op,
299 })
300}
301
302pub struct KernelExecutable {
311 pub name: String,
312 pub kernel_source: String,
313 pub entry_point: String,
314 pub shape: Shape,
315 pub dtype: DtypeRepr,
316 pub runtime_op: Arc<dyn RuntimeOp>,
319 #[cfg(feature = "training")]
321 pub backward_kernel_source: String,
322 #[cfg(feature = "training")]
324 pub backward_entry_point: String,
325}
326
327impl ExecutableOp for KernelExecutable {
328 fn name(&self) -> &str {
329 &self.name
330 }
331
332 fn is_input(&self) -> bool {
333 self.name == "input"
334 }
335
336 fn forward_kernel_source(&self) -> &str {
337 &self.kernel_source
338 }
339
340 fn forward_kernel_entry_point(&self) -> &str {
341 &self.entry_point
342 }
343
344 fn output_shape(&self) -> &Shape {
345 &self.shape
346 }
347
348 fn output_dtype(&self) -> DtypeRepr {
349 self.dtype
350 }
351
352 fn runtime_op(&self) -> Option<Arc<dyn RuntimeOp>> {
353 if self.is_input() {
354 None
355 } else {
356 Some(Arc::clone(&self.runtime_op))
357 }
358 }
359
360 #[cfg(feature = "training")]
361 fn backward_kernel_source(&self) -> &str {
362 &self.backward_kernel_source
363 }
364
365 #[cfg(feature = "training")]
366 fn backward_kernel_entry_point(&self) -> &str {
367 &self.backward_entry_point
368 }
369}
370
371macro_rules! impl_stub_runtime_op_num {
378 ($T:ident) => {
379 impl<D: teeny_core::dtype::Num + Send + Sync + 'static> RuntimeOp for $T<D> {
380 fn n_activation_inputs(&self) -> usize {
381 unimplemented!(concat!(stringify!($T), " has no runtime support"))
382 }
383 fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
384 unimplemented!()
385 }
386 fn pack_args(
387 &self,
388 _: &[(teeny_core::model::RawPtr, &[usize])],
389 _: &[teeny_core::model::RawPtr],
390 _: teeny_core::model::RawPtr,
391 _: &[usize],
392 _: i32,
393 _: &mut dyn teeny_core::device::program::ArgVisitor,
394 ) {
395 unimplemented!()
396 }
397 fn block(&self) -> [u32; 3] {
398 unimplemented!()
399 }
400 fn grid(&self, _: &[usize]) -> [u32; 3] {
401 unimplemented!()
402 }
403 }
404 };
405}
406
407macro_rules! impl_stub_runtime_op_float {
408 ($T:ident) => {
409 impl<D: teeny_core::dtype::Float + Send + Sync + 'static> RuntimeOp for $T<D> {
410 fn n_activation_inputs(&self) -> usize {
411 unimplemented!(concat!(stringify!($T), " has no runtime support"))
412 }
413 fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
414 unimplemented!()
415 }
416 fn pack_args(
417 &self,
418 _: &[(teeny_core::model::RawPtr, &[usize])],
419 _: &[teeny_core::model::RawPtr],
420 _: teeny_core::model::RawPtr,
421 _: &[usize],
422 _: i32,
423 _: &mut dyn teeny_core::device::program::ArgVisitor,
424 ) {
425 unimplemented!()
426 }
427 fn block(&self) -> [u32; 3] {
428 unimplemented!()
429 }
430 fn grid(&self, _: &[usize]) -> [u32; 3] {
431 unimplemented!()
432 }
433 }
434 };
435}
436
437impl_stub_runtime_op_float!(BatchNormForwardInference);
439impl_stub_runtime_op_float!(LayerNormForwardInference);
440impl_stub_runtime_op_float!(RmsNormForward);
441impl_stub_runtime_op_float!(GroupNormForwardInference);
442impl_stub_runtime_op_float!(InstanceNormForwardInference);
443
444impl_stub_runtime_op_num!(Conv3dForward);
446
447impl_stub_runtime_op_num!(Avgpool1dForward);
449impl_stub_runtime_op_num!(Avgpool3dForward);
450impl_stub_runtime_op_num!(Maxpool1dForward);
451impl_stub_runtime_op_num!(Maxpool3dForward);
452impl_stub_runtime_op_float!(Lppool1dForward);
453impl_stub_runtime_op_float!(Lppool2dForward);
454impl_stub_runtime_op_float!(Lppool3dForward);
455
456impl_stub_runtime_op_num!(ConstantPad1dForward);
458impl_stub_runtime_op_num!(ConstantPad2dForward);
459impl_stub_runtime_op_num!(ConstantPad3dForward);
460impl_stub_runtime_op_num!(ReflectionPad1dForward);
461impl_stub_runtime_op_num!(ReflectionPad2dForward);
462impl_stub_runtime_op_num!(ReflectionPad3dForward);
463impl_stub_runtime_op_num!(ReplicationPad1dForward);
464impl_stub_runtime_op_num!(ReplicationPad2dForward);
465impl_stub_runtime_op_num!(ReplicationPad3dForward);
466impl_stub_runtime_op_num!(CircularPad1dForward);
467impl_stub_runtime_op_num!(CircularPad2dForward);
468impl_stub_runtime_op_num!(CircularPad3dForward);
469
470impl_stub_runtime_op_float!(EluForward);
473impl_stub_runtime_op_float!(SeluForward);
474impl_stub_runtime_op_float!(CeluForward);
475impl_stub_runtime_op_float!(MishForward);
476impl_stub_runtime_op_float!(HardtanhForward);
477impl_stub_runtime_op_float!(Relu6Forward);
478impl_stub_runtime_op_float!(HardsigmoidForward);
479impl_stub_runtime_op_float!(HardswishForward);
480impl_stub_runtime_op_float!(HardshrinkForward);
481impl_stub_runtime_op_float!(LeakyReluForward);
482impl_stub_runtime_op_float!(ThresholdForward);
483impl_stub_runtime_op_float!(SoftsignForward);
484impl_stub_runtime_op_float!(SoftshrinkForward);
485impl_stub_runtime_op_float!(SoftplusForward);
486impl_stub_runtime_op_float!(LogsigmoidForward);
487impl_stub_runtime_op_float!(TanhForward);
488impl_stub_runtime_op_float!(TanhshrinkForward);
489
490struct InputRuntimeOp;
495
496impl RuntimeOp for InputRuntimeOp {
497 fn n_activation_inputs(&self) -> usize {
498 0
499 }
500 fn param_shapes(&self, _: &[&[usize]], _: &[usize]) -> Vec<Vec<usize>> {
501 Vec::new()
502 }
503 fn pack_args(
504 &self,
505 _: &[(teeny_core::model::RawPtr, &[usize])],
506 _: &[teeny_core::model::RawPtr],
507 _: teeny_core::model::RawPtr,
508 _: &[usize],
509 _: i32,
510 _: &mut dyn teeny_core::device::program::ArgVisitor,
511 ) {
512 }
513 fn block(&self) -> [u32; 3] {
514 [1, 1, 1]
515 }
516 fn grid(&self, _: &[usize]) -> [u32; 3] {
517 [0, 0, 0]
518 }
519}
520
521#[derive(Debug, Default)]
526pub struct TritonLowering {
527 sm_count: Option<u32>,
531}
532
533impl TritonLowering {
534 pub fn new() -> Self {
535 Self::default()
536 }
537
538 pub fn with_sm_count(mut self, sm_count: Option<u32>) -> Self {
542 self.sm_count = sm_count;
543 self
544 }
545}
546
547fn pick_adaptive_block_n(
558 tiled_dim: usize,
559 fixed_blocks: usize,
560 target_blocks: u32,
561 candidates: &[i32],
562) -> i32 {
563 for &c in candidates {
564 let n_tiles = tiled_dim.div_ceil(c.max(1) as usize);
565 if (fixed_blocks * n_tiles) as u64 >= target_blocks as u64 {
566 return c;
567 }
568 }
569 *candidates.last().expect("candidates must be non-empty")
570}
571
572#[cfg(test)]
573mod pick_adaptive_block_n_tests {
574 use super::pick_adaptive_block_n;
575
576 #[test]
577 fn keeps_largest_candidate_when_already_enough_blocks() {
578 let picked = pick_adaptive_block_n(256, 400, 512, &[16, 8, 4]);
581 assert_eq!(picked, 16);
582 }
583
584 #[test]
585 fn shrinks_for_occupancy_starved_shapes() {
586 let picked = pick_adaptive_block_n(256, 10, 512, &[16, 8, 4]);
589 assert_eq!(picked, 4);
590 }
591
592 #[test]
593 fn falls_back_to_smallest_candidate_when_target_unreachable() {
594 let picked = pick_adaptive_block_n(16, 1, 1_000_000, &[16, 8, 4]);
596 assert_eq!(picked, 4);
597 }
598}
599
600fn pick_gemm_tile_sizes(m: Option<usize>, n: usize, k: usize) -> (i32, i32, i32) {
624 let m = m.unwrap_or(64);
625 let block_k = if k >= 128 {
626 32
627 } else if k >= 32 {
628 16
629 } else {
630 8
631 };
632 let (block_m, block_n) = match (m, n) {
633 (m, n) if m >= 256 && n >= 128 => (256, 128),
634 (m, n) if m >= 128 && n >= 128 => (128, 128),
635 (m, _) if m >= 128 => (128, 64),
636 (_, n) if n >= 128 => (64, 128),
637 _ => (64, 64),
638 };
639 (block_m, block_n, block_k)
640}
641
642#[cfg(test)]
643mod pick_gemm_tile_sizes_tests {
644 use super::pick_gemm_tile_sizes;
645
646 #[test]
647 fn small_shape_gets_smallest_tiles() {
648 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 16), (64, 64, 8));
649 }
650
651 #[test]
652 fn large_m_and_n_get_the_largest_tile() {
653 assert_eq!(pick_gemm_tile_sizes(Some(512), 256, 256), (256, 128, 32));
654 }
655
656 #[test]
657 fn large_m_only_widens_block_m_not_block_n() {
658 assert_eq!(pick_gemm_tile_sizes(Some(512), 32, 64), (128, 64, 16));
659 }
660
661 #[test]
662 fn large_n_only_widens_block_n_not_block_m() {
663 assert_eq!(pick_gemm_tile_sizes(Some(32), 512, 64), (64, 128, 16));
664 }
665
666 #[test]
667 fn unknown_dynamic_batch_treated_as_small() {
668 assert_eq!(
670 pick_gemm_tile_sizes(None, 64, 16),
671 pick_gemm_tile_sizes(Some(64), 64, 16)
672 );
673 }
674
675 #[test]
676 fn block_k_never_drops_below_the_tensor_core_minimum() {
677 let (_, _, block_k) = pick_gemm_tile_sizes(Some(64), 64, 1);
678 assert!(block_k >= 8);
679 }
680
681 #[test]
682 fn block_k_grows_with_k() {
683 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 8).2, 8);
684 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 32).2, 16);
685 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 128).2, 32);
686 }
687}
688
689impl TritonLowering {
690 pub fn lower_with_mapping(
694 &self,
695 graph: &Graph,
696 mode: LoweringMode,
697 ) -> Result<(Dag<Box<dyn ExecutableOp>>, Vec<usize>)> {
698 let _ = mode; let node_indexes = graph.topological_sort();
700 let mut dag: Dag<Box<dyn ExecutableOp>> = Dag::new();
701 let mut graph_to_dag = vec![0usize; graph.nodes.len()];
703
704 for node_index in node_indexes {
705 let node = &graph.nodes[node_index];
706
707 #[cfg(feature = "training")]
711 if mode == LoweringMode::Training
712 && let Op::BatchNorm1d {
713 num_features,
714 eps,
715 momentum,
716 ..
717 }
718 | Op::BatchNorm3d {
719 num_features,
720 eps,
721 momentum,
722 ..
723 } = &node.op
724 {
725 let c = *num_features;
726 let eps_f32 = *eps as f32;
727 let momentum_f32 = *momentum as f32;
728 const BLOCK_N: i32 = 64;
729
730 let (stats_name, stats_src, stats_rop): (String, String, Arc<dyn RuntimeOp>) =
731 match node.dtype {
732 DtypeRepr::F32 => {
733 let k = BatchNormStatsForward::<f32>::new(BLOCK_N);
734 let src = k.source.clone();
735 let rop: Arc<dyn RuntimeOp> = Arc::new(
736 BatchNormStatsRuntimeOp::<f32>::new(BLOCK_N, eps_f32, momentum_f32),
737 );
738 (k.name.to_string(), src, rop)
739 }
740 DtypeRepr::F64 => {
741 let k = BatchNormStatsForward::<f64>::new(BLOCK_N);
742 let src = k.source.clone();
743 let rop: Arc<dyn RuntimeOp> = Arc::new(
744 BatchNormStatsRuntimeOp::<f64>::new(BLOCK_N, eps_f32, momentum_f32),
745 );
746 (k.name.to_string(), src, rop)
747 }
748 other => {
749 return Err(anyhow::anyhow!(
750 "{:?} is not a Float dtype for BatchNormStatsForward",
751 other
752 ));
753 }
754 };
755
756 let stats_node = Box::new(KernelExecutable {
757 entry_point: format!("{}_entry_point", stats_name),
758 name: stats_name,
759 kernel_source: stats_src,
760 shape: vec![Some(2 * c)],
761 dtype: node.dtype,
762 backward_kernel_source: String::new(),
763 backward_entry_point: String::new(),
764 runtime_op: stats_rop,
765 }) as Box<dyn ExecutableOp>;
766
767 let stats_dag_idx = dag.add_node(stats_node);
768 for &input_graph_idx in &node.inputs {
769 dag.add_edge(graph_to_dag[input_graph_idx], stats_dag_idx);
770 }
771
772 let (norm_name, norm_src, norm_bwd_src, norm_rop): (
773 String,
774 String,
775 String,
776 Arc<dyn RuntimeOp>,
777 ) = match node.dtype {
778 DtypeRepr::F32 => {
779 let k = BatchNormNormalizeForward::<f32>::new(BLOCK_N);
780 let src = k.source.clone();
781 let rop = BatchNormNormalizeRuntimeOp::<f32>::new(BLOCK_N);
782 let bwd_src = rop.backward_source().to_string();
783 (
784 k.name.to_string(),
785 src,
786 bwd_src,
787 Arc::new(rop) as Arc<dyn RuntimeOp>,
788 )
789 }
790 DtypeRepr::F64 => {
791 let k = BatchNormNormalizeForward::<f64>::new(BLOCK_N);
792 let src = k.source.clone();
793 let rop = BatchNormNormalizeRuntimeOp::<f64>::new(BLOCK_N);
794 let bwd_src = rop.backward_source().to_string();
795 (
796 k.name.to_string(),
797 src,
798 bwd_src,
799 Arc::new(rop) as Arc<dyn RuntimeOp>,
800 )
801 }
802 other => {
803 return Err(anyhow::anyhow!(
804 "{:?} is not a Float dtype for BatchNormNormalizeForward",
805 other
806 ));
807 }
808 };
809
810 let norm_node = Box::new(KernelExecutable {
811 entry_point: format!("{}_entry_point", norm_name),
812 name: norm_name,
813 kernel_source: norm_src,
814 shape: node.shape.clone(),
815 dtype: node.dtype,
816 backward_kernel_source: norm_bwd_src,
817 backward_entry_point: String::new(),
818 runtime_op: norm_rop,
819 }) as Box<dyn ExecutableOp>;
820
821 let norm_dag_idx = dag.add_node(norm_node);
822 for &input_graph_idx in &node.inputs {
824 dag.add_edge(graph_to_dag[input_graph_idx], norm_dag_idx);
825 }
826 dag.add_edge(stats_dag_idx, norm_dag_idx);
828
829 graph_to_dag[node_index] = norm_dag_idx;
830 continue;
831 }
832
833 if mode == LoweringMode::Inference
837 && let Op::Conv2d {
838 has_bias: true,
839 kernel_h,
840 kernel_w,
841 stride_h,
842 stride_w,
843 padding_h,
844 padding_w,
845 groups,
846 ..
847 } = &node.op
848 {
849 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
850 DtypeRepr::F32 => {
851 let k = Conv2dBiasForward::<f32>::new(
852 *kernel_h as i32,
853 *kernel_w as i32,
854 *stride_h as i32,
855 *stride_w as i32,
856 *padding_h as i32,
857 *padding_w as i32,
858 *groups as i32,
859 16,
860 );
861 let nm = k.name.to_string();
862 let src = k.source.clone();
863 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
864 (nm, src, rop)
865 }
866 DtypeRepr::F64 => {
867 let k = Conv2dBiasForward::<f64>::new(
868 *kernel_h as i32,
869 *kernel_w as i32,
870 *stride_h as i32,
871 *stride_w as i32,
872 *padding_h as i32,
873 *padding_w as i32,
874 *groups as i32,
875 16,
876 );
877 let nm = k.name.to_string();
878 let src = k.source.clone();
879 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
880 (nm, src, rop)
881 }
882 other => {
883 return Err(anyhow::anyhow!(
884 "{:?} is not supported for Conv2dBiasForward",
885 other
886 ));
887 }
888 };
889 let dag_idx = dag.add_node(Box::new(KernelExecutable {
890 entry_point: format!("{}_entry_point", name),
891 name,
892 kernel_source: ks,
893 shape: node.shape.clone(),
894 dtype: node.dtype,
895 #[cfg(feature = "training")]
896 backward_kernel_source: String::new(),
897 #[cfg(feature = "training")]
898 backward_entry_point: String::new(),
899 runtime_op: rop,
900 }) as Box<dyn ExecutableOp>);
901 for &input_graph_idx in &node.inputs {
902 dag.add_edge(graph_to_dag[input_graph_idx], dag_idx);
903 }
904 graph_to_dag[node_index] = dag_idx;
905 continue;
906 }
907
908 if let Op::Conv2d {
911 has_bias: true,
912 kernel_h,
913 kernel_w,
914 stride_h,
915 stride_w,
916 padding_h,
917 padding_w,
918 groups,
919 out_channels,
920 ..
921 } = &node.op
922 {
923 const BIAS_BLOCK_HW: i32 = 128;
924 let (conv_name, conv_ks, conv_rop): (String, String, Arc<dyn RuntimeOp>) =
925 match node.dtype {
926 DtypeRepr::F32 => {
927 let k = Conv2dForward::<f32>::new(
928 *kernel_h as i32,
929 *kernel_w as i32,
930 *stride_h as i32,
931 *stride_w as i32,
932 *padding_h as i32,
933 *padding_w as i32,
934 *groups as i32,
935 16,
936 );
937 let src = k.source.clone();
938 let rop: Arc<dyn RuntimeOp> = Arc::new(Conv2dForward::<f32>::new(
939 *kernel_h as i32,
940 *kernel_w as i32,
941 *stride_h as i32,
942 *stride_w as i32,
943 *padding_h as i32,
944 *padding_w as i32,
945 *groups as i32,
946 16,
947 ));
948 (k.name.to_string(), src, rop)
949 }
950 DtypeRepr::F64 => {
951 let k = Conv2dForward::<f64>::new(
952 *kernel_h as i32,
953 *kernel_w as i32,
954 *stride_h as i32,
955 *stride_w as i32,
956 *padding_h as i32,
957 *padding_w as i32,
958 *groups as i32,
959 16,
960 );
961 let src = k.source.clone();
962 let rop: Arc<dyn RuntimeOp> = Arc::new(Conv2dForward::<f64>::new(
963 *kernel_h as i32,
964 *kernel_w as i32,
965 *stride_h as i32,
966 *stride_w as i32,
967 *padding_h as i32,
968 *padding_w as i32,
969 *groups as i32,
970 16,
971 ));
972 (k.name.to_string(), src, rop)
973 }
974 other => {
975 return Err(anyhow::anyhow!(
976 "{:?} is not supported for Conv2dForward",
977 other
978 ));
979 }
980 };
981
982 #[cfg(feature = "training")]
983 let conv_bwd_ks = match node.dtype {
984 DtypeRepr::F32 => {
985 Conv2dBackward::<f32>::new(
986 *kernel_h as i32,
987 *kernel_w as i32,
988 *stride_h as i32,
989 *stride_w as i32,
990 *padding_h as i32,
991 *padding_w as i32,
992 *groups as i32,
993 16,
994 )
995 .source
996 }
997 DtypeRepr::F64 => {
998 Conv2dBackward::<f64>::new(
999 *kernel_h as i32,
1000 *kernel_w as i32,
1001 *stride_h as i32,
1002 *stride_w as i32,
1003 *padding_h as i32,
1004 *padding_w as i32,
1005 *groups as i32,
1006 16,
1007 )
1008 .source
1009 }
1010 _ => String::new(),
1011 };
1012
1013 let conv_dag_idx = dag.add_node(Box::new(KernelExecutable {
1014 entry_point: format!("{}_entry_point", conv_name),
1015 name: conv_name,
1016 kernel_source: conv_ks,
1017 shape: node.shape.clone(),
1018 dtype: node.dtype,
1019 #[cfg(feature = "training")]
1020 backward_kernel_source: conv_bwd_ks,
1021 #[cfg(feature = "training")]
1022 backward_entry_point: String::new(),
1023 runtime_op: conv_rop,
1024 }) as Box<dyn ExecutableOp>);
1025 for &input_graph_idx in &node.inputs {
1026 dag.add_edge(graph_to_dag[input_graph_idx], conv_dag_idx);
1027 }
1028
1029 let (bias_name, bias_ks, bias_rop): (String, String, Arc<dyn RuntimeOp>) =
1030 match node.dtype {
1031 DtypeRepr::F32 => {
1032 let r = NchwBiasAddRuntimeOp::<f32>::new(BIAS_BLOCK_HW);
1033 (
1034 r.kernel_name().to_string(),
1035 r.forward_source().to_string(),
1036 Arc::new(r),
1037 )
1038 }
1039 DtypeRepr::F64 => {
1040 let r = NchwBiasAddRuntimeOp::<f64>::new(BIAS_BLOCK_HW);
1041 (
1042 r.kernel_name().to_string(),
1043 r.forward_source().to_string(),
1044 Arc::new(r),
1045 )
1046 }
1047 other => {
1048 return Err(anyhow::anyhow!(
1049 "{:?} is not supported for NchwBiasAdd",
1050 other
1051 ));
1052 }
1053 };
1054
1055 #[cfg(feature = "training")]
1056 let bias_bwd_ks = match node.dtype {
1057 DtypeRepr::F32 => NchwBiasAddRuntimeOp::<f32>::new(BIAS_BLOCK_HW)
1058 .backward_source()
1059 .to_string(),
1060 DtypeRepr::F64 => NchwBiasAddRuntimeOp::<f64>::new(BIAS_BLOCK_HW)
1061 .backward_source()
1062 .to_string(),
1063 _ => String::new(),
1064 };
1065
1066 let biasadd_dag_idx = dag.add_node(Box::new(KernelExecutable {
1067 entry_point: format!("{}_entry_point", bias_name),
1068 name: bias_name,
1069 kernel_source: bias_ks,
1070 shape: node.shape.clone(),
1071 dtype: node.dtype,
1072 #[cfg(feature = "training")]
1073 backward_kernel_source: bias_bwd_ks,
1074 #[cfg(feature = "training")]
1075 backward_entry_point: String::new(),
1076 runtime_op: bias_rop,
1077 }) as Box<dyn ExecutableOp>);
1078 dag.add_edge(conv_dag_idx, biasadd_dag_idx);
1079 graph_to_dag[node_index] = biasadd_dag_idx;
1080 let _ = (conv_dag_idx, out_channels);
1083 continue;
1084 }
1085
1086 let executable: Box<dyn ExecutableOp> = match &node.op {
1087 Op::Input => Box::new(KernelExecutable {
1088 name: "input".to_string(),
1089 kernel_source: String::new(),
1090 entry_point: String::new(),
1091 shape: node.shape.clone(),
1092 dtype: node.dtype,
1093 #[cfg(feature = "training")]
1094 backward_kernel_source: String::new(),
1095 #[cfg(feature = "training")]
1096 backward_entry_point: String::new(),
1097 runtime_op: Arc::new(InputRuntimeOp),
1098 }),
1099
1100 Op::Linear { has_bias, .. } => {
1102 make_num_kernel!(
1103 LinearForward(*has_bias, 32, 64, 32, 8),
1104 LinearBackward(*has_bias, 32, 64, 32, 8),
1105 node
1106 )
1107 }
1108 Op::Flatten => make_num_kernel!(FlattenForward(32, 256), node),
1109
1110 Op::BatchNorm1d { .. } | Op::BatchNorm3d { .. } => {
1112 make_float_kernel!(BatchNormForwardInference(64), node)
1113 }
1114 Op::BatchNorm2d { eps, .. } => {
1115 let eps_f32 = *eps as f32;
1116 const BN2D_BLOCK_HW: i32 = 128;
1117 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
1118 DtypeRepr::F32 => {
1119 let r = BatchNorm2dNchwInferenceRuntimeOp::<f32>::new(
1120 BN2D_BLOCK_HW,
1121 eps_f32,
1122 );
1123 (
1124 r.kernel_name().to_string(),
1125 r.forward_source().to_string(),
1126 Arc::new(r),
1127 )
1128 }
1129 DtypeRepr::F64 => {
1130 let r = BatchNorm2dNchwInferenceRuntimeOp::<f64>::new(
1131 BN2D_BLOCK_HW,
1132 eps_f32,
1133 );
1134 (
1135 r.kernel_name().to_string(),
1136 r.forward_source().to_string(),
1137 Arc::new(r),
1138 )
1139 }
1140 other => {
1141 return Err(anyhow::anyhow!(
1142 "{:?} is not a Float dtype for BatchNorm2d",
1143 other
1144 ));
1145 }
1146 };
1147 #[cfg(feature = "training")]
1148 let bwd_ks = match node.dtype {
1149 DtypeRepr::F32 => BatchNorm2dNchwBackward::<f32>::new(BN2D_BLOCK_HW).source,
1150 DtypeRepr::F64 => BatchNorm2dNchwBackward::<f64>::new(BN2D_BLOCK_HW).source,
1151 _ => String::new(),
1152 };
1153 Box::new(KernelExecutable {
1154 entry_point: format!("{}_entry_point", name),
1155 name,
1156 kernel_source: ks,
1157 shape: node.shape.clone(),
1158 dtype: node.dtype,
1159 #[cfg(feature = "training")]
1160 backward_kernel_source: bwd_ks,
1161 #[cfg(feature = "training")]
1162 backward_entry_point: String::new(),
1163 runtime_op: rop,
1164 })
1165 }
1166 Op::LayerNorm { eps, .. } => {
1167 let eps_f32 = *eps as f32;
1168 const LN_BLOCK_N: i32 = 1024;
1169 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
1170 DtypeRepr::F32 => {
1171 let r =
1172 LayerNormForwardInferenceRuntimeOp::<f32>::new(LN_BLOCK_N, eps_f32);
1173 (
1174 r.kernel_name().to_string(),
1175 r.forward_source().to_string(),
1176 Arc::new(r),
1177 )
1178 }
1179 DtypeRepr::F64 => {
1180 let r =
1181 LayerNormForwardInferenceRuntimeOp::<f64>::new(LN_BLOCK_N, eps_f32);
1182 (
1183 r.kernel_name().to_string(),
1184 r.forward_source().to_string(),
1185 Arc::new(r),
1186 )
1187 }
1188 other => {
1189 return Err(anyhow::anyhow!(
1190 "{:?} is not a Float dtype for LayerNorm",
1191 other
1192 ));
1193 }
1194 };
1195 Box::new(KernelExecutable {
1196 entry_point: format!("{}_entry_point", name),
1197 name,
1198 kernel_source: ks,
1199 shape: node.shape.clone(),
1200 dtype: node.dtype,
1201 #[cfg(feature = "training")]
1202 backward_kernel_source: String::new(),
1203 #[cfg(feature = "training")]
1204 backward_entry_point: String::new(),
1205 runtime_op: rop,
1206 })
1207 }
1208 Op::RmsNorm { .. } => {
1209 make_float_kernel!(RmsNormForward(1024), node)
1210 }
1211 Op::GroupNorm { .. } => {
1212 make_float_kernel!(GroupNormForwardInference(256), node)
1213 }
1214 Op::InstanceNorm1d { .. }
1215 | Op::InstanceNorm2d { .. }
1216 | Op::InstanceNorm3d { .. } => {
1217 make_float_kernel!(InstanceNormForwardInference(256), node)
1218 }
1219
1220 Op::Conv1d {
1222 kernel_l,
1223 stride,
1224 padding,
1225 ..
1226 } => {
1227 make_num_kernel!(
1228 Conv1dForward(*kernel_l as i32, *stride as i32, *padding as i32, 32),
1229 node
1230 )
1231 }
1232 Op::Conv2d {
1233 kernel_h,
1234 kernel_w,
1235 stride_h,
1236 stride_w,
1237 padding_h,
1238 padding_w,
1239 groups,
1240 ..
1241 } => {
1242 make_num_kernel!(
1243 Conv2dForward(
1244 *kernel_h as i32,
1245 *kernel_w as i32,
1246 *stride_h as i32,
1247 *stride_w as i32,
1248 *padding_h as i32,
1249 *padding_w as i32,
1250 *groups as i32,
1251 16
1252 ),
1253 Conv2dBackward(
1254 *kernel_h as i32,
1255 *kernel_w as i32,
1256 *stride_h as i32,
1257 *stride_w as i32,
1258 *padding_h as i32,
1259 *padding_w as i32,
1260 *groups as i32,
1261 16
1262 ),
1263 node
1264 )
1265 }
1266 Op::Conv3d {
1267 kernel_d,
1268 kernel_h,
1269 kernel_w,
1270 stride_d,
1271 stride_h,
1272 stride_w,
1273 padding_d,
1274 padding_h,
1275 padding_w,
1276 ..
1277 } => {
1278 make_num_kernel!(
1279 Conv3dForward(
1280 *kernel_d as i32,
1281 *kernel_h as i32,
1282 *kernel_w as i32,
1283 *stride_d as i32,
1284 *stride_h as i32,
1285 *stride_w as i32,
1286 *padding_d as i32,
1287 *padding_h as i32,
1288 *padding_w as i32,
1289 8
1290 ),
1291 node
1292 )
1293 }
1294
1295 Op::Conv2dBnSilu {
1296 kernel_h,
1297 kernel_w,
1298 stride_h,
1299 stride_w,
1300 padding_h,
1301 padding_w,
1302 groups,
1303 in_channels,
1304 ..
1305 } => {
1306 if node.dtype != DtypeRepr::F32 {
1307 return Err(anyhow::anyhow!(
1308 "Conv2dBnSilu only supports f32 (got {:?})",
1309 node.dtype
1310 ));
1311 }
1312 let kh = *kernel_h as i32;
1313 let kw = *kernel_w as i32;
1314 let sh = *stride_h as i32;
1315 let sw = *stride_w as i32;
1316 let ph = *padding_h as i32;
1317 let pw = *padding_w as i32;
1318 let g = *groups as i32;
1319 let c_out = node.shape[1].unwrap_or(0);
1320 let oh = node.shape[2].unwrap_or(1);
1321 let ow = node.shape[3].unwrap_or(1);
1322 let is_depthwise = g as usize == *in_channels;
1323
1324 let use_gemm = kh == 1
1338 && kw == 1
1339 && sh == 1
1340 && sw == 1
1341 && ph == 0
1342 && pw == 0
1343 && !is_depthwise
1344 && g == 1
1345 && c_out >= 32;
1346
1347 let use_tiled = !is_depthwise && g == 1 && c_out >= 16 && !use_gemm;
1348
1349 if use_gemm {
1350 const GROUP_M: i32 = 8;
1351 let m = oh * ow;
1352 let (block_m, block_n_base, block_k) =
1357 pick_gemm_tile_sizes(Some(m), c_out, *in_channels);
1358 let block_n = match self.sm_count {
1365 Some(sm_count) => {
1366 let fixed_blocks = m.div_ceil(block_m as usize);
1367 pick_adaptive_block_n(
1368 c_out,
1369 fixed_blocks,
1370 4 * sm_count,
1371 &[block_n_base, 16, 8],
1372 )
1373 }
1374 None => block_n_base,
1375 };
1376 let k = Conv2dBnSiluGemmForward::new(block_m, block_n, block_k, GROUP_M);
1377 let nm = k.name.to_string();
1378 let ks = k.source.clone();
1379 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
1380 Box::new(KernelExecutable {
1381 entry_point: format!("{}_entry_point", nm),
1382 name: nm,
1383 kernel_source: ks,
1384 shape: node.shape.clone(),
1385 dtype: node.dtype,
1386 #[cfg(feature = "training")]
1387 backward_kernel_source: String::new(),
1388 #[cfg(feature = "training")]
1389 backward_entry_point: String::new(),
1390 runtime_op: rop,
1391 })
1392 } else if use_tiled {
1393 const BLOCK_OW: i32 = 16;
1394 const BLOCK_N_TILE: i32 = 16;
1395 let block_n_tile = match self.sm_count {
1400 Some(sm_count) => {
1401 let fixed_blocks = oh * ow.div_ceil(BLOCK_OW as usize);
1402 pick_adaptive_block_n(
1403 c_out,
1404 fixed_blocks,
1405 4 * sm_count,
1406 &[BLOCK_N_TILE, 8, 4],
1407 )
1408 }
1409 None => BLOCK_N_TILE,
1410 };
1411 let k = Conv2dBnSiluTiledForward::new(
1412 kh,
1413 kw,
1414 sh,
1415 sw,
1416 ph,
1417 pw,
1418 BLOCK_OW,
1419 block_n_tile,
1420 );
1421 let nm = k.name.to_string();
1422 let ks = k.source.clone();
1423 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
1424 Box::new(KernelExecutable {
1425 entry_point: format!("{}_entry_point", nm),
1426 name: nm,
1427 kernel_source: ks,
1428 shape: node.shape.clone(),
1429 dtype: node.dtype,
1430 #[cfg(feature = "training")]
1431 backward_kernel_source: String::new(),
1432 #[cfg(feature = "training")]
1433 backward_entry_point: String::new(),
1434 runtime_op: rop,
1435 })
1436 } else {
1437 const BLOCK_OW: i32 = 16;
1438 let k = Conv2dBnSiluForward::new(kh, kw, sh, sw, ph, pw, g, BLOCK_OW);
1439 let nm = k.name.to_string();
1440 let ks = k.source.clone();
1441 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
1442 Box::new(KernelExecutable {
1443 entry_point: format!("{}_entry_point", nm),
1444 name: nm,
1445 kernel_source: ks,
1446 shape: node.shape.clone(),
1447 dtype: node.dtype,
1448 #[cfg(feature = "training")]
1449 backward_kernel_source: String::new(),
1450 #[cfg(feature = "training")]
1451 backward_entry_point: String::new(),
1452 runtime_op: rop,
1453 })
1454 }
1455 }
1456
1457 Op::AvgPool1d { kernel_l, stride } => {
1459 make_num_kernel!(Avgpool1dForward(*kernel_l as i32, *stride as i32, 32), node)
1460 }
1461 Op::AvgPool2d {
1462 kernel_h,
1463 kernel_w,
1464 stride_h,
1465 stride_w,
1466 } => {
1467 make_num_kernel!(
1468 Avgpool2dForward(
1469 *kernel_h as i32,
1470 *kernel_w as i32,
1471 *stride_h as i32,
1472 *stride_w as i32,
1473 16
1474 ),
1475 node
1476 )
1477 }
1478 Op::AvgPool3d {
1479 kernel_d,
1480 kernel_h,
1481 kernel_w,
1482 stride_d,
1483 stride_h,
1484 stride_w,
1485 } => {
1486 make_num_kernel!(
1487 Avgpool3dForward(
1488 *kernel_d as i32,
1489 *kernel_h as i32,
1490 *kernel_w as i32,
1491 *stride_d as i32,
1492 *stride_h as i32,
1493 *stride_w as i32,
1494 8
1495 ),
1496 node
1497 )
1498 }
1499 Op::MaxPool1d { kernel_l, stride } => {
1500 make_num_kernel!(Maxpool1dForward(*kernel_l as i32, *stride as i32, 32), node)
1501 }
1502 Op::MaxPool2d {
1503 kernel_h,
1504 kernel_w,
1505 stride_h,
1506 stride_w,
1507 pad_h,
1508 pad_w,
1509 } => {
1510 make_num_kernel!(
1511 Maxpool2dForward(
1512 *kernel_h as i32,
1513 *kernel_w as i32,
1514 *stride_h as i32,
1515 *stride_w as i32,
1516 *pad_h as i32,
1517 *pad_w as i32,
1518 16
1519 ),
1520 Maxpool2dBackward(
1521 *kernel_h as i32,
1522 *kernel_w as i32,
1523 *stride_h as i32,
1524 *stride_w as i32,
1525 *pad_h as i32,
1526 *pad_w as i32,
1527 16
1528 ),
1529 node
1530 )
1531 }
1532 Op::MaxPool3d {
1533 kernel_d,
1534 kernel_h,
1535 kernel_w,
1536 stride_d,
1537 stride_h,
1538 stride_w,
1539 } => {
1540 make_num_kernel!(
1541 Maxpool3dForward(
1542 *kernel_d as i32,
1543 *kernel_h as i32,
1544 *kernel_w as i32,
1545 *stride_d as i32,
1546 *stride_h as i32,
1547 *stride_w as i32,
1548 8
1549 ),
1550 node
1551 )
1552 }
1553 Op::LpPool1d {
1554 kernel_l, stride, ..
1555 } => {
1556 make_float_kernel!(Lppool1dForward(*kernel_l as i32, *stride as i32, 32), node)
1557 }
1558 Op::LpPool2d {
1559 kernel_h,
1560 kernel_w,
1561 stride_h,
1562 stride_w,
1563 ..
1564 } => {
1565 make_float_kernel!(
1566 Lppool2dForward(
1567 *kernel_h as i32,
1568 *kernel_w as i32,
1569 *stride_h as i32,
1570 *stride_w as i32,
1571 16
1572 ),
1573 node
1574 )
1575 }
1576 Op::LpPool3d {
1577 kernel_d,
1578 kernel_h,
1579 kernel_w,
1580 stride_d,
1581 stride_h,
1582 stride_w,
1583 ..
1584 } => {
1585 make_float_kernel!(
1586 Lppool3dForward(
1587 *kernel_d as i32,
1588 *kernel_h as i32,
1589 *kernel_w as i32,
1590 *stride_d as i32,
1591 *stride_h as i32,
1592 *stride_w as i32,
1593 8
1594 ),
1595 node
1596 )
1597 }
1598
1599 Op::ConstantPad1d {
1601 pad_left,
1602 pad_right,
1603 ..
1604 } => {
1605 make_num_kernel!(
1606 ConstantPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1607 node
1608 )
1609 }
1610 Op::ConstantPad2d {
1611 pad_l,
1612 pad_r,
1613 pad_t,
1614 pad_b,
1615 ..
1616 } => {
1617 make_num_kernel!(
1618 ConstantPad2dForward(
1619 *pad_t as i32,
1620 *pad_b as i32,
1621 *pad_l as i32,
1622 *pad_r as i32,
1623 16
1624 ),
1625 node
1626 )
1627 }
1628 Op::ConstantPad3d {
1629 pad_d1,
1630 pad_d2,
1631 pad_h1,
1632 pad_h2,
1633 pad_w1,
1634 pad_w2,
1635 ..
1636 } => {
1637 make_num_kernel!(
1638 ConstantPad3dForward(
1639 *pad_d1 as i32,
1640 *pad_d2 as i32,
1641 *pad_h1 as i32,
1642 *pad_h2 as i32,
1643 *pad_w1 as i32,
1644 *pad_w2 as i32,
1645 8
1646 ),
1647 node
1648 )
1649 }
1650 Op::ReflectionPad1d {
1651 pad_left,
1652 pad_right,
1653 } => {
1654 make_num_kernel!(
1655 ReflectionPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1656 node
1657 )
1658 }
1659 Op::ReflectionPad2d {
1660 pad_l,
1661 pad_r,
1662 pad_t,
1663 pad_b,
1664 } => {
1665 make_num_kernel!(
1666 ReflectionPad2dForward(
1667 *pad_t as i32,
1668 *pad_b as i32,
1669 *pad_l as i32,
1670 *pad_r as i32,
1671 16
1672 ),
1673 node
1674 )
1675 }
1676 Op::ReflectionPad3d {
1677 pad_d1,
1678 pad_d2,
1679 pad_h1,
1680 pad_h2,
1681 pad_w1,
1682 pad_w2,
1683 } => {
1684 make_num_kernel!(
1685 ReflectionPad3dForward(
1686 *pad_d1 as i32,
1687 *pad_d2 as i32,
1688 *pad_h1 as i32,
1689 *pad_h2 as i32,
1690 *pad_w1 as i32,
1691 *pad_w2 as i32,
1692 8
1693 ),
1694 node
1695 )
1696 }
1697 Op::ReplicationPad1d {
1698 pad_left,
1699 pad_right,
1700 } => {
1701 make_num_kernel!(
1702 ReplicationPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1703 node
1704 )
1705 }
1706 Op::ReplicationPad2d {
1707 pad_l,
1708 pad_r,
1709 pad_t,
1710 pad_b,
1711 } => {
1712 make_num_kernel!(
1713 ReplicationPad2dForward(
1714 *pad_t as i32,
1715 *pad_b as i32,
1716 *pad_l as i32,
1717 *pad_r as i32,
1718 16
1719 ),
1720 node
1721 )
1722 }
1723 Op::ReplicationPad3d {
1724 pad_d1,
1725 pad_d2,
1726 pad_h1,
1727 pad_h2,
1728 pad_w1,
1729 pad_w2,
1730 } => {
1731 make_num_kernel!(
1732 ReplicationPad3dForward(
1733 *pad_d1 as i32,
1734 *pad_d2 as i32,
1735 *pad_h1 as i32,
1736 *pad_h2 as i32,
1737 *pad_w1 as i32,
1738 *pad_w2 as i32,
1739 8
1740 ),
1741 node
1742 )
1743 }
1744 Op::CircularPad1d {
1745 pad_left,
1746 pad_right,
1747 } => {
1748 make_num_kernel!(
1749 CircularPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1750 node
1751 )
1752 }
1753 Op::CircularPad2d {
1754 pad_l,
1755 pad_r,
1756 pad_t,
1757 pad_b,
1758 } => {
1759 make_num_kernel!(
1760 CircularPad2dForward(
1761 *pad_t as i32,
1762 *pad_b as i32,
1763 *pad_l as i32,
1764 *pad_r as i32,
1765 16
1766 ),
1767 node
1768 )
1769 }
1770 Op::CircularPad3d {
1771 pad_d1,
1772 pad_d2,
1773 pad_h1,
1774 pad_h2,
1775 pad_w1,
1776 pad_w2,
1777 } => {
1778 make_num_kernel!(
1779 CircularPad3dForward(
1780 *pad_d1 as i32,
1781 *pad_d2 as i32,
1782 *pad_h1 as i32,
1783 *pad_h2 as i32,
1784 *pad_w1 as i32,
1785 *pad_w2 as i32,
1786 8
1787 ),
1788 node
1789 )
1790 }
1791
1792 Op::Relu => make_num_kernel!(ReluForward(1024), ReluBackward(1024), node),
1794
1795 Op::Elu { .. } => exec_from(
1797 node.shape.clone(),
1798 node.dtype,
1799 EluForwardDispatch::dispatch(node.dtype, 1024)?,
1800 ),
1801 Op::Selu => exec_from(
1802 node.shape.clone(),
1803 node.dtype,
1804 SeluForwardDispatch::dispatch(node.dtype, 1024)?,
1805 ),
1806 Op::Celu { .. } => exec_from(
1807 node.shape.clone(),
1808 node.dtype,
1809 CeluForwardDispatch::dispatch(node.dtype, 1024)?,
1810 ),
1811 Op::Gelu => exec_from(
1812 node.shape.clone(),
1813 node.dtype,
1814 GeluForwardDispatch::dispatch(node.dtype, 1024)?,
1815 ),
1816 Op::Mish => exec_from(
1817 node.shape.clone(),
1818 node.dtype,
1819 MishForwardDispatch::dispatch(node.dtype, 1024)?,
1820 ),
1821 Op::Hardtanh { .. } => exec_from(
1822 node.shape.clone(),
1823 node.dtype,
1824 HardtanhForwardDispatch::dispatch(node.dtype, 1024)?,
1825 ),
1826 Op::Relu6 => exec_from(
1827 node.shape.clone(),
1828 node.dtype,
1829 Relu6ForwardDispatch::dispatch(node.dtype, 1024)?,
1830 ),
1831 Op::Hardsigmoid => exec_from(
1832 node.shape.clone(),
1833 node.dtype,
1834 HardsigmoidForwardDispatch::dispatch(node.dtype, 1024)?,
1835 ),
1836 Op::Hardswish => exec_from(
1837 node.shape.clone(),
1838 node.dtype,
1839 HardswishForwardDispatch::dispatch(node.dtype, 1024)?,
1840 ),
1841 Op::Hardshrink { .. } => exec_from(
1842 node.shape.clone(),
1843 node.dtype,
1844 HardshrinkForwardDispatch::dispatch(node.dtype, 1024)?,
1845 ),
1846 Op::LeakyRelu { .. } => exec_from(
1847 node.shape.clone(),
1848 node.dtype,
1849 LeakyReluForwardDispatch::dispatch(node.dtype, 1024)?,
1850 ),
1851 Op::Threshold { .. } => exec_from(
1852 node.shape.clone(),
1853 node.dtype,
1854 ThresholdForwardDispatch::dispatch(node.dtype, 1024)?,
1855 ),
1856 Op::Softsign => exec_from(
1857 node.shape.clone(),
1858 node.dtype,
1859 SoftsignForwardDispatch::dispatch(node.dtype, 1024)?,
1860 ),
1861 Op::Softshrink { .. } => exec_from(
1862 node.shape.clone(),
1863 node.dtype,
1864 SoftshrinkForwardDispatch::dispatch(node.dtype, 1024)?,
1865 ),
1866 Op::Softplus { .. } => exec_from(
1867 node.shape.clone(),
1868 node.dtype,
1869 SoftplusForwardDispatch::dispatch(node.dtype, 1024)?,
1870 ),
1871 Op::Sigmoid => exec_from(
1872 node.shape.clone(),
1873 node.dtype,
1874 SigmoidForwardDispatch::dispatch(node.dtype, 1024)?,
1875 ),
1876 Op::Silu => exec_from(
1877 node.shape.clone(),
1878 node.dtype,
1879 SiluForwardDispatch::dispatch(node.dtype, 1024)?,
1880 ),
1881 Op::Logsigmoid => exec_from(
1882 node.shape.clone(),
1883 node.dtype,
1884 LogsigmoidForwardDispatch::dispatch(node.dtype, 1024)?,
1885 ),
1886 Op::Tanh => exec_from(
1887 node.shape.clone(),
1888 node.dtype,
1889 TanhForwardDispatch::dispatch(node.dtype, 1024)?,
1890 ),
1891 Op::Tanhshrink => exec_from(
1892 node.shape.clone(),
1893 node.dtype,
1894 TanhshrinkForwardDispatch::dispatch(node.dtype, 1024)?,
1895 ),
1896
1897 Op::Softmax { .. } => {
1899 let n_cols = node.shape.last().and_then(|d| *d).unwrap_or(1024);
1901 let block_size = n_cols.next_power_of_two() as i32;
1902 make_float_kernel!(SoftmaxForward(block_size), node)
1903 }
1904
1905 Op::UpsampleNearest2d { scale_h, scale_w } => {
1907 make_num_kernel!(
1908 UpsampleNearest2dForward(*scale_h as i32, *scale_w as i32, 16),
1909 UpsampleNearest2dBackward(*scale_h as i32, *scale_w as i32, 16),
1910 node
1911 )
1912 }
1913
1914 Op::ChannelCat { .. } => {
1915 let n_inputs = node.inputs.len();
1916 let (name, fwd_src, bwd_src, rop): (
1917 String,
1918 String,
1919 String,
1920 Arc<dyn RuntimeOp>,
1921 ) = match node.dtype {
1922 DtypeRepr::F32 => {
1923 let r = ChannelCatRuntimeOp::<f32>::new(128, n_inputs);
1924 (
1925 r.kernel_name().to_string(),
1926 r.forward_source().to_string(),
1927 r.backward_source().to_string(),
1928 Arc::new(r),
1929 )
1930 }
1931 DtypeRepr::F64 => {
1932 let r = ChannelCatRuntimeOp::<f64>::new(128, n_inputs);
1933 (
1934 r.kernel_name().to_string(),
1935 r.forward_source().to_string(),
1936 r.backward_source().to_string(),
1937 Arc::new(r),
1938 )
1939 }
1940 DtypeRepr::I8 => {
1941 let r = ChannelCatRuntimeOp::<i8>::new(128, n_inputs);
1942 (
1943 r.kernel_name().to_string(),
1944 r.forward_source().to_string(),
1945 r.backward_source().to_string(),
1946 Arc::new(r),
1947 )
1948 }
1949 DtypeRepr::I16 => {
1950 let r = ChannelCatRuntimeOp::<i16>::new(128, n_inputs);
1951 (
1952 r.kernel_name().to_string(),
1953 r.forward_source().to_string(),
1954 r.backward_source().to_string(),
1955 Arc::new(r),
1956 )
1957 }
1958 DtypeRepr::I32 => {
1959 let r = ChannelCatRuntimeOp::<i32>::new(128, n_inputs);
1960 (
1961 r.kernel_name().to_string(),
1962 r.forward_source().to_string(),
1963 r.backward_source().to_string(),
1964 Arc::new(r),
1965 )
1966 }
1967 DtypeRepr::I64 => {
1968 let r = ChannelCatRuntimeOp::<i64>::new(128, n_inputs);
1969 (
1970 r.kernel_name().to_string(),
1971 r.forward_source().to_string(),
1972 r.backward_source().to_string(),
1973 Arc::new(r),
1974 )
1975 }
1976 DtypeRepr::U8 => {
1977 let r = ChannelCatRuntimeOp::<u8>::new(128, n_inputs);
1978 (
1979 r.kernel_name().to_string(),
1980 r.forward_source().to_string(),
1981 r.backward_source().to_string(),
1982 Arc::new(r),
1983 )
1984 }
1985 DtypeRepr::U16 => {
1986 let r = ChannelCatRuntimeOp::<u16>::new(128, n_inputs);
1987 (
1988 r.kernel_name().to_string(),
1989 r.forward_source().to_string(),
1990 r.backward_source().to_string(),
1991 Arc::new(r),
1992 )
1993 }
1994 DtypeRepr::U32 => {
1995 let r = ChannelCatRuntimeOp::<u32>::new(128, n_inputs);
1996 (
1997 r.kernel_name().to_string(),
1998 r.forward_source().to_string(),
1999 r.backward_source().to_string(),
2000 Arc::new(r),
2001 )
2002 }
2003 DtypeRepr::U64 => {
2004 let r = ChannelCatRuntimeOp::<u64>::new(128, n_inputs);
2005 (
2006 r.kernel_name().to_string(),
2007 r.forward_source().to_string(),
2008 r.backward_source().to_string(),
2009 Arc::new(r),
2010 )
2011 }
2012 other => {
2013 return Err(anyhow::anyhow!(
2014 "{:?} is not supported for ChannelCat",
2015 other
2016 ));
2017 }
2018 };
2019 Box::new(KernelExecutable {
2020 entry_point: format!("{}_entry_point", name),
2021 name,
2022 kernel_source: fwd_src,
2023 shape: node.shape.clone(),
2024 dtype: node.dtype,
2025 #[cfg(feature = "training")]
2026 backward_kernel_source: bwd_src,
2027 #[cfg(feature = "training")]
2028 backward_entry_point: String::new(),
2029 runtime_op: rop,
2030 })
2031 }
2032
2033 Op::ChannelChunk {
2034 chunk_c,
2035 chunk_offset,
2036 ..
2037 } => {
2038 let chunk_c = *chunk_c;
2039 let chunk_offset = *chunk_offset;
2040 let (name, fwd_src, bwd_src, rop): (
2041 String,
2042 String,
2043 String,
2044 Arc<dyn RuntimeOp>,
2045 ) = match node.dtype {
2046 DtypeRepr::F32 => {
2047 let r = ChannelChunkRuntimeOp::<f32>::new(128, chunk_c, chunk_offset);
2048 (
2049 r.kernel_name().to_string(),
2050 r.forward_source().to_string(),
2051 r.backward_source().to_string(),
2052 Arc::new(r),
2053 )
2054 }
2055 DtypeRepr::F64 => {
2056 let r = ChannelChunkRuntimeOp::<f64>::new(128, chunk_c, chunk_offset);
2057 (
2058 r.kernel_name().to_string(),
2059 r.forward_source().to_string(),
2060 r.backward_source().to_string(),
2061 Arc::new(r),
2062 )
2063 }
2064 DtypeRepr::I8 => {
2065 let r = ChannelChunkRuntimeOp::<i8>::new(128, chunk_c, chunk_offset);
2066 (
2067 r.kernel_name().to_string(),
2068 r.forward_source().to_string(),
2069 r.backward_source().to_string(),
2070 Arc::new(r),
2071 )
2072 }
2073 DtypeRepr::I16 => {
2074 let r = ChannelChunkRuntimeOp::<i16>::new(128, chunk_c, chunk_offset);
2075 (
2076 r.kernel_name().to_string(),
2077 r.forward_source().to_string(),
2078 r.backward_source().to_string(),
2079 Arc::new(r),
2080 )
2081 }
2082 DtypeRepr::I32 => {
2083 let r = ChannelChunkRuntimeOp::<i32>::new(128, chunk_c, chunk_offset);
2084 (
2085 r.kernel_name().to_string(),
2086 r.forward_source().to_string(),
2087 r.backward_source().to_string(),
2088 Arc::new(r),
2089 )
2090 }
2091 DtypeRepr::I64 => {
2092 let r = ChannelChunkRuntimeOp::<i64>::new(128, chunk_c, chunk_offset);
2093 (
2094 r.kernel_name().to_string(),
2095 r.forward_source().to_string(),
2096 r.backward_source().to_string(),
2097 Arc::new(r),
2098 )
2099 }
2100 DtypeRepr::U8 => {
2101 let r = ChannelChunkRuntimeOp::<u8>::new(128, chunk_c, chunk_offset);
2102 (
2103 r.kernel_name().to_string(),
2104 r.forward_source().to_string(),
2105 r.backward_source().to_string(),
2106 Arc::new(r),
2107 )
2108 }
2109 DtypeRepr::U16 => {
2110 let r = ChannelChunkRuntimeOp::<u16>::new(128, chunk_c, chunk_offset);
2111 (
2112 r.kernel_name().to_string(),
2113 r.forward_source().to_string(),
2114 r.backward_source().to_string(),
2115 Arc::new(r),
2116 )
2117 }
2118 DtypeRepr::U32 => {
2119 let r = ChannelChunkRuntimeOp::<u32>::new(128, chunk_c, chunk_offset);
2120 (
2121 r.kernel_name().to_string(),
2122 r.forward_source().to_string(),
2123 r.backward_source().to_string(),
2124 Arc::new(r),
2125 )
2126 }
2127 DtypeRepr::U64 => {
2128 let r = ChannelChunkRuntimeOp::<u64>::new(128, chunk_c, chunk_offset);
2129 (
2130 r.kernel_name().to_string(),
2131 r.forward_source().to_string(),
2132 r.backward_source().to_string(),
2133 Arc::new(r),
2134 )
2135 }
2136 other => {
2137 return Err(anyhow::anyhow!(
2138 "{:?} is not supported for ChannelChunk",
2139 other
2140 ));
2141 }
2142 };
2143 Box::new(KernelExecutable {
2144 entry_point: format!("{}_entry_point", name),
2145 name,
2146 kernel_source: fwd_src,
2147 shape: node.shape.clone(),
2148 dtype: node.dtype,
2149 #[cfg(feature = "training")]
2150 backward_kernel_source: bwd_src,
2151 #[cfg(feature = "training")]
2152 backward_entry_point: String::new(),
2153 runtime_op: rop,
2154 })
2155 }
2156
2157 Op::ChannelBiasAdd { c } => {
2158 let c = *c;
2159 let (name, fwd_src, bwd_src, rop): (
2160 String,
2161 String,
2162 String,
2163 Arc<dyn RuntimeOp>,
2164 ) = match node.dtype {
2165 DtypeRepr::F32 => {
2166 let r = ChannelBiasAddRuntimeOp::<f32>::new(128, c);
2167 (
2168 r.kernel_name().to_string(),
2169 r.forward_source().to_string(),
2170 r.backward_source().to_string(),
2171 Arc::new(r),
2172 )
2173 }
2174 DtypeRepr::F64 => {
2175 let r = ChannelBiasAddRuntimeOp::<f64>::new(128, c);
2176 (
2177 r.kernel_name().to_string(),
2178 r.forward_source().to_string(),
2179 r.backward_source().to_string(),
2180 Arc::new(r),
2181 )
2182 }
2183 other => {
2184 return Err(anyhow::anyhow!(
2185 "{:?} is not supported for ChannelBiasAdd",
2186 other
2187 ));
2188 }
2189 };
2190 Box::new(KernelExecutable {
2191 entry_point: format!("{}_entry_point", name),
2192 name,
2193 kernel_source: fwd_src,
2194 shape: node.shape.clone(),
2195 dtype: node.dtype,
2196 #[cfg(feature = "training")]
2197 backward_kernel_source: bwd_src,
2198 #[cfg(feature = "training")]
2199 backward_entry_point: String::new(),
2200 runtime_op: rop,
2201 })
2202 }
2203
2204 Op::Add => {
2205 make_num_kernel!(ElemwiseAddForward(128), ElemwiseAddBackward(128), node)
2206 }
2207
2208 Op::Abs => {
2210 make_num_kernel!(ElemwiseAbsForward(1024), ElemwiseAbsBackward(1024), node)
2211 }
2212 Op::Neg => {
2213 make_num_kernel!(ElemwiseNegForward(1024), ElemwiseNegBackward(1024), node)
2214 }
2215 Op::Sign => make_num_kernel!(ElemwiseSignForward(1024), node),
2216 Op::IsNaN => make_float_kernel!(ElemwiseIsnanForward(1024), node),
2217 Op::Ceil => make_float_kernel!(ElemwiseCeilForward(1024), node),
2218 Op::Floor => make_float_kernel!(ElemwiseFloorForward(1024), node),
2219 Op::Sqrt => {
2220 make_float_kernel!(ElemwiseSqrtForward(1024), ElemwiseSqrtBackward(1024), node)
2221 }
2222 Op::Reciprocal => make_float_kernel!(
2223 ElemwiseReciprocalForward(1024),
2224 ElemwiseReciprocalBackward(1024),
2225 node
2226 ),
2227 Op::Exp => {
2228 make_float_kernel!(ElemwiseExpForward(1024), ElemwiseExpBackward(1024), node)
2229 }
2230 Op::Log => {
2231 make_float_kernel!(ElemwiseLogForward(1024), ElemwiseLogBackward(1024), node)
2232 }
2233 Op::Erf => {
2234 make_float_kernel!(ElemwiseErfForward(1024), ElemwiseErfBackward(1024), node)
2235 }
2236 Op::Sin => {
2237 make_float_kernel!(ElemwiseSinForward(1024), ElemwiseSinBackward(1024), node)
2238 }
2239 Op::Cos => {
2240 make_float_kernel!(ElemwiseCosForward(1024), ElemwiseCosBackward(1024), node)
2241 }
2242 Op::Tan => {
2243 make_float_kernel!(ElemwiseTanForward(1024), ElemwiseTanBackward(1024), node)
2244 }
2245 Op::Asin => {
2246 make_float_kernel!(ElemwiseAsinForward(1024), ElemwiseAsinBackward(1024), node)
2247 }
2248 Op::Acos => {
2249 make_float_kernel!(ElemwiseAcosForward(1024), ElemwiseAcosBackward(1024), node)
2250 }
2251 Op::Atan => {
2252 make_float_kernel!(ElemwiseAtanForward(1024), ElemwiseAtanBackward(1024), node)
2253 }
2254 Op::Sinh => {
2255 make_float_kernel!(ElemwiseSinhForward(1024), ElemwiseSinhBackward(1024), node)
2256 }
2257 Op::Cosh => {
2258 make_float_kernel!(ElemwiseCoshForward(1024), ElemwiseCoshBackward(1024), node)
2259 }
2260 Op::Asinh => make_float_kernel!(
2261 ElemwiseAsinhForward(1024),
2262 ElemwiseAsinhBackward(1024),
2263 node
2264 ),
2265 Op::Acosh => make_float_kernel!(
2266 ElemwiseAcoshForward(1024),
2267 ElemwiseAcoshBackward(1024),
2268 node
2269 ),
2270 Op::Atanh => make_float_kernel!(
2271 ElemwiseAtanhForward(1024),
2272 ElemwiseAtanhBackward(1024),
2273 node
2274 ),
2275 Op::Round => {
2276 return Err(anyhow::anyhow!(
2277 "TODO: Op::Round — implement rounding kernel"
2278 ));
2279 }
2280
2281 Op::Mul => {
2283 make_num_kernel!(ElemwiseMulForward(1024), ElemwiseMulBackward(1024), node)
2284 }
2285 Op::Sub => {
2286 make_num_kernel!(ElemwiseSubForward(1024), ElemwiseSubBackward(1024), node)
2287 }
2288 Op::Div => {
2289 make_float_kernel!(ElemwiseDivForward(1024), ElemwiseDivBackward(1024), node)
2290 }
2291 Op::Pow => {
2292 make_float_kernel!(ElemwisePowForward(1024), ElemwisePowBackward(1024), node)
2293 }
2294 Op::Mod { .. } => make_float_kernel!(ElemwiseFmodForward(1024), node),
2295 Op::ElemMin => {
2296 make_num_kernel!(ElemwiseMinForward(1024), ElemwiseMinBackward(1024), node)
2297 }
2298 Op::ElemMax => {
2299 make_num_kernel!(ElemwiseMaxForward(1024), ElemwiseMaxBackward(1024), node)
2300 }
2301 Op::ElemMean => {
2302 make_float_kernel!(ElemwiseMeanForward(1024), ElemwiseMeanBackward(1024), node)
2303 }
2304 Op::ElemSum => {
2305 make_num_kernel!(ElemwiseSumForward(1024), ElemwiseSumBackward(1024), node)
2306 }
2307 Op::Equal => make_num_kernel!(ElemwiseEqualForward(1024), node),
2308 Op::Greater => make_num_kernel!(ElemwiseGreaterForward(1024), node),
2309 Op::GreaterOrEqual => make_num_kernel!(ElemwiseGreaterEqualForward(1024), node),
2310 Op::Less => make_num_kernel!(ElemwiseLessForward(1024), node),
2311 Op::LessOrEqual => make_num_kernel!(ElemwiseLessEqualForward(1024), node),
2312 Op::Where => make_float_kernel!(
2313 ElemwiseWhereForward(1024),
2314 ElemwiseWhereBackward(1024),
2315 node
2316 ),
2317 Op::Clip => {
2318 let (name, ks, bwd_ks, rop): (String, String, String, Arc<dyn RuntimeOp>) =
2319 match node.dtype {
2320 DtypeRepr::F32 => {
2321 let r = ClipRuntimeOp::<f32>::new(
2322 1024,
2323 f32::NEG_INFINITY,
2324 f32::INFINITY,
2325 );
2326 (
2327 r.kernel_name().to_string(),
2328 r.forward_source().to_string(),
2329 r.backward_source().to_string(),
2330 Arc::new(r),
2331 )
2332 }
2333 DtypeRepr::F64 => {
2334 let r = ClipRuntimeOp::<f64>::new(
2335 1024,
2336 f32::NEG_INFINITY,
2337 f32::INFINITY,
2338 );
2339 (
2340 r.kernel_name().to_string(),
2341 r.forward_source().to_string(),
2342 r.backward_source().to_string(),
2343 Arc::new(r),
2344 )
2345 }
2346 other => {
2347 return Err(anyhow::anyhow!(
2348 "{:?} is not supported for Clip",
2349 other
2350 ));
2351 }
2352 };
2353 Box::new(KernelExecutable {
2354 entry_point: format!("{}_entry_point", name),
2355 name,
2356 kernel_source: ks,
2357 shape: node.shape.clone(),
2358 dtype: node.dtype,
2359 #[cfg(feature = "training")]
2360 backward_kernel_source: bwd_ks,
2361 #[cfg(feature = "training")]
2362 backward_entry_point: String::new(),
2363 runtime_op: rop,
2364 })
2365 }
2366
2367 Op::ReduceSum { .. } => make_num_kernel!(ReduceSumForward(1024), node),
2369 Op::ReduceMean { .. } => make_float_kernel!(ReduceMeanForward(1024), node),
2370 Op::ReduceMax { .. } => make_num_kernel!(ReduceMaxForward(1024), node),
2371 Op::ReduceMin { .. } => make_num_kernel!(ReduceMinForward(1024), node),
2372 Op::ReduceProd { .. } => make_float_kernel!(ReduceProdForward(1024), node),
2373 Op::ReduceL1 { .. } => make_num_kernel!(ReduceL1Forward(1024), node),
2374 Op::ReduceL2 { .. } => make_float_kernel!(ReduceL2Forward(1024), node),
2375 Op::ReduceLogSum { .. } => make_float_kernel!(ReduceLogSumForward(1024), node),
2376 Op::ReduceLogSumExp { .. } => {
2377 make_float_kernel!(ReduceLogSumExpForward(1024), node)
2378 }
2379 Op::ReduceSumSquare { .. } => make_num_kernel!(ReduceSumSquareForward(1024), node),
2380 Op::CumSum { .. } => make_num_kernel!(CumSumForward(1024), node),
2381 Op::CumProd { .. } => make_num_kernel!(CumProdForward(1024), node),
2382 Op::GlobalAvgPool => make_float_kernel!(GlobalAvgPoolForward(1024), node),
2383 Op::GlobalMaxPool => make_float_kernel!(GlobalMaxPoolForward(1024), node),
2384 Op::ArgMax { .. } => {
2385 return Err(anyhow::anyhow!(
2386 "TODO: Op::ArgMax — I32Tensor output requires a custom kernel"
2387 ));
2388 }
2389 Op::ArgMin { .. } => {
2390 return Err(anyhow::anyhow!(
2391 "TODO: Op::ArgMin — I32Tensor output requires a custom kernel"
2392 ));
2393 }
2394
2395 Op::Swish => {
2397 let fwd = SwishForward::new(1024);
2398 let nm = fwd.name.to_string();
2399 let fwd_src = fwd.source.clone();
2400 #[cfg(feature = "training")]
2401 let bwd_src = SwishBackward::new(1024).source.clone();
2402 let rop: Arc<dyn RuntimeOp> = Arc::new(fwd);
2403 Box::new(KernelExecutable {
2404 entry_point: format!("{}_entry_point", nm),
2405 name: nm,
2406 kernel_source: fwd_src,
2407 shape: node.shape.clone(),
2408 dtype: node.dtype,
2409 #[cfg(feature = "training")]
2410 backward_kernel_source: bwd_src,
2411 #[cfg(feature = "training")]
2412 backward_entry_point: String::new(),
2413 runtime_op: rop,
2414 })
2415 }
2416 Op::PRelu => {
2417 let fwd = PreluForward::new(1024);
2418 let nm = fwd.name.to_string();
2419 let fwd_src = fwd.source.clone();
2420 #[cfg(feature = "training")]
2421 let bwd_src = crate::nn::activation::extra::PreluBackward::new(1024)
2422 .source
2423 .clone();
2424 let rop: Arc<dyn RuntimeOp> = Arc::new(fwd);
2425 Box::new(KernelExecutable {
2426 entry_point: format!("{}_entry_point", nm),
2427 name: nm,
2428 kernel_source: fwd_src,
2429 shape: node.shape.clone(),
2430 dtype: node.dtype,
2431 #[cfg(feature = "training")]
2432 backward_kernel_source: bwd_src,
2433 #[cfg(feature = "training")]
2434 backward_entry_point: String::new(),
2435 runtime_op: rop,
2436 })
2437 }
2438 Op::LogSoftmax { .. } => {
2439 let n_cols = node.shape.last().and_then(|d| *d).unwrap_or(1024);
2440 let block_size = n_cols.next_power_of_two() as i32;
2441 let fwd = LogSoftmaxForward::new(block_size);
2442 let nm = fwd.name.to_string();
2443 let fwd_src = fwd.source.clone();
2444 #[cfg(feature = "training")]
2445 let bwd_src = LogSoftmaxBackward::new(block_size).source.clone();
2446 let rop: Arc<dyn RuntimeOp> = Arc::new(fwd);
2447 Box::new(KernelExecutable {
2448 entry_point: format!("{}_entry_point", nm),
2449 name: nm,
2450 kernel_source: fwd_src,
2451 shape: node.shape.clone(),
2452 dtype: node.dtype,
2453 #[cfg(feature = "training")]
2454 backward_kernel_source: bwd_src,
2455 #[cfg(feature = "training")]
2456 backward_entry_point: String::new(),
2457 runtime_op: rop,
2458 })
2459 }
2460 Op::ThresholdedRelu { alpha } => {
2461 let alpha_f = *alpha as f32;
2462 let (name, ks, bwd_ks, rop): (String, String, String, Arc<dyn RuntimeOp>) =
2463 match node.dtype {
2464 DtypeRepr::F32 => {
2465 let r = ThresholdedReluRuntimeOp::new(1024, alpha_f);
2466 (
2467 r.kernel_name().to_string(),
2468 r.forward_source().to_string(),
2469 r.backward_source().to_string(),
2470 Arc::new(r),
2471 )
2472 }
2473 other => {
2474 return Err(anyhow::anyhow!(
2475 "{:?} is not supported for ThresholdedRelu",
2476 other
2477 ));
2478 }
2479 };
2480 Box::new(KernelExecutable {
2481 entry_point: format!("{}_entry_point", name),
2482 name,
2483 kernel_source: ks,
2484 shape: node.shape.clone(),
2485 dtype: node.dtype,
2486 #[cfg(feature = "training")]
2487 backward_kernel_source: bwd_ks,
2488 #[cfg(feature = "training")]
2489 backward_entry_point: String::new(),
2490 runtime_op: rop,
2491 })
2492 }
2493 Op::Shrink { lambd, bias } => {
2494 let lambd_f = *lambd as f32;
2495 let bias_f = *bias as f32;
2496 let (name, ks, bwd_ks, rop): (String, String, String, Arc<dyn RuntimeOp>) =
2497 match node.dtype {
2498 DtypeRepr::F32 => {
2499 let r = ShrinkRuntimeOp::new(1024, lambd_f, bias_f);
2500 (
2501 r.kernel_name().to_string(),
2502 r.forward_source().to_string(),
2503 r.backward_source().to_string(),
2504 Arc::new(r),
2505 )
2506 }
2507 other => {
2508 return Err(anyhow::anyhow!(
2509 "{:?} is not supported for Shrink",
2510 other
2511 ));
2512 }
2513 };
2514 Box::new(KernelExecutable {
2515 entry_point: format!("{}_entry_point", name),
2516 name,
2517 kernel_source: ks,
2518 shape: node.shape.clone(),
2519 dtype: node.dtype,
2520 #[cfg(feature = "training")]
2521 backward_kernel_source: bwd_ks,
2522 #[cfg(feature = "training")]
2523 backward_entry_point: String::new(),
2524 runtime_op: rop,
2525 })
2526 }
2527
2528 Op::MatMul | Op::Gemm { .. } => {
2530 let m = node.shape.first().copied().flatten();
2532 let n = node.shape.last().copied().flatten().unwrap_or(0);
2533 let k = node
2534 .inputs
2535 .first()
2536 .and_then(|&i| graph.nodes[i].shape.last().copied().flatten())
2537 .unwrap_or(0);
2538 let (block_m, block_n, block_k) = pick_gemm_tile_sizes(m, n, k);
2539
2540 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
2541 DtypeRepr::F32 => {
2542 let r = MatMulRuntimeOp::<f32>::new(block_m, block_n, block_k);
2543 (
2544 r.kernel_name().to_string(),
2545 r.forward_source().to_string(),
2546 Arc::new(r),
2547 )
2548 }
2549 DtypeRepr::F64 => {
2550 let r = MatMulRuntimeOp::<f64>::new(block_m, block_n, block_k);
2551 (
2552 r.kernel_name().to_string(),
2553 r.forward_source().to_string(),
2554 Arc::new(r),
2555 )
2556 }
2557 other => {
2558 return Err(anyhow::anyhow!(
2559 "{:?} is not a Float dtype for MatMul",
2560 other
2561 ));
2562 }
2563 };
2564 Box::new(KernelExecutable {
2565 entry_point: format!("{}_entry_point", name),
2566 name,
2567 kernel_source: ks,
2568 shape: node.shape.clone(),
2569 dtype: node.dtype,
2570 #[cfg(feature = "training")]
2571 backward_kernel_source: String::new(),
2572 #[cfg(feature = "training")]
2573 backward_entry_point: String::new(),
2574 runtime_op: rop,
2575 })
2576 }
2577
2578 Op::Lstm { .. } => {
2580 return Err(anyhow::anyhow!(
2581 "TODO: Op::Lstm — multi-step recurrent; implement as a custom loop kernel"
2582 ));
2583 }
2584 Op::Gru { .. } => {
2585 return Err(anyhow::anyhow!(
2586 "TODO: Op::Gru — multi-step recurrent; implement as a custom loop kernel"
2587 ));
2588 }
2589 Op::Rnn { .. } => {
2590 return Err(anyhow::anyhow!(
2591 "TODO: Op::Rnn — multi-step recurrent; implement as a custom loop kernel"
2592 ));
2593 }
2594 Op::RotaryEmbedding => {
2595 return Err(anyhow::anyhow!(
2596 "TODO: Op::RotaryEmbedding — implement RoPE kernel"
2597 ));
2598 }
2599 Op::MultiHeadAttention { .. } => {
2600 return Err(anyhow::anyhow!(
2601 "TODO: Op::MultiHeadAttention — use flash attention or a custom MHA kernel"
2602 ));
2603 }
2604 Op::FlexAttention { .. } => {
2605 return Err(anyhow::anyhow!(
2606 "TODO: Op::FlexAttention — implement a custom attention kernel supporting an arbitrary score-modification function"
2607 ));
2608 }
2609 Op::LinearAttention { .. } => {
2610 return Err(anyhow::anyhow!(
2611 "TODO: Op::LinearAttention — implement a linear-attention/gated-delta-rule kernel"
2612 ));
2613 }
2614 Op::CausalConvWithState { .. } => {
2615 return Err(anyhow::anyhow!(
2616 "TODO: Op::CausalConvWithState — implement a stateful causal-conv kernel"
2617 ));
2618 }
2619 Op::Reshape => {
2620 return Err(anyhow::anyhow!(
2621 "TODO: Op::Reshape — implement as a strided view or copy kernel"
2622 ));
2623 }
2624 Op::Transpose { .. } => {
2625 return Err(anyhow::anyhow!(
2626 "TODO: Op::Transpose — implement as a permuted-copy kernel"
2627 ));
2628 }
2629 Op::Squeeze { .. } | Op::Unsqueeze { .. } => {
2630 return Err(anyhow::anyhow!(
2631 "TODO: Op::Squeeze/Unsqueeze — implement as a zero-copy view"
2632 ));
2633 }
2634 Op::Concat { .. } => {
2635 return Err(anyhow::anyhow!(
2636 "TODO: Op::Concat — implement as a multi-input copy kernel"
2637 ));
2638 }
2639 Op::Split { .. } => {
2640 return Err(anyhow::anyhow!(
2641 "TODO: Op::Split — implement as a multi-output slice kernel"
2642 ));
2643 }
2644 Op::Slice => {
2645 return Err(anyhow::anyhow!(
2646 "TODO: Op::Slice — implement as a strided-copy kernel"
2647 ));
2648 }
2649 Op::Gather { .. } | Op::GatherElements { .. } | Op::GatherND { .. } => {
2650 return Err(anyhow::anyhow!(
2651 "TODO: Op::Gather — implement as an index-gather kernel"
2652 ));
2653 }
2654 Op::ScatterElements { .. }
2655 | Op::ScatterND
2656 | Op::Scatter { .. }
2657 | Op::TensorScatter => {
2658 return Err(anyhow::anyhow!(
2659 "TODO: Op::Scatter — implement as an index-scatter kernel"
2660 ));
2661 }
2662 Op::Tile => {
2663 return Err(anyhow::anyhow!(
2664 "TODO: Op::Tile — implement as a tiled-copy kernel"
2665 ));
2666 }
2667 Op::Expand => {
2668 return Err(anyhow::anyhow!(
2669 "TODO: Op::Expand — implement as a broadcast-copy kernel"
2670 ));
2671 }
2672 Op::ShapeOf { .. } | Op::SizeOf => {
2673 return Err(anyhow::anyhow!(
2674 "TODO: Op::ShapeOf/SizeOf — output is metadata, not tensor data"
2675 ));
2676 }
2677 Op::Compress { .. } | Op::NonZero => {
2678 return Err(anyhow::anyhow!(
2679 "TODO: Op::Compress/NonZero — variable-output ops require stream compaction"
2680 ));
2681 }
2682 Op::Range => {
2683 return Err(anyhow::anyhow!(
2684 "TODO: Op::Range — implement as a fill/arange kernel"
2685 ));
2686 }
2687 Op::Constant { .. } | Op::ConstantOfShape { .. } => {
2688 return Err(anyhow::anyhow!(
2689 "TODO: Op::Constant — inline constant; should be materialised before lowering"
2690 ));
2691 }
2692 Op::Trilu { .. } => {
2693 return Err(anyhow::anyhow!(
2694 "TODO: Op::Trilu — implement as a triangular mask kernel"
2695 ));
2696 }
2697 Op::Pad { .. } => {
2698 return Err(anyhow::anyhow!(
2699 "TODO: Op::Pad — implement as a generic N-D padding kernel"
2700 ));
2701 }
2702 Op::ReverseSequence { .. } => {
2703 return Err(anyhow::anyhow!(
2704 "TODO: Op::ReverseSequence — implement as a scatter-copy kernel"
2705 ));
2706 }
2707 Op::Einsum { .. } => {
2708 return Err(anyhow::anyhow!(
2709 "TODO: Op::Einsum — parse equation and emit a fused contraction kernel"
2710 ));
2711 }
2712 Op::Det => {
2713 return Err(anyhow::anyhow!(
2714 "TODO: Op::Det — implement via LU decomposition"
2715 ));
2716 }
2717 Op::QLinearMatMul
2718 | Op::MatMulInteger
2719 | Op::ConvInteger { .. }
2720 | Op::QLinearConv { .. } => {
2721 return Err(anyhow::anyhow!(
2722 "TODO: Op::Q* quantised matmul/conv — implement quantised compute kernels"
2723 ));
2724 }
2725 Op::ConvTranspose { .. } => {
2726 return Err(anyhow::anyhow!(
2727 "TODO: Op::ConvTranspose — implement transposed (gradient) convolution kernel"
2728 ));
2729 }
2730 Op::DeformConv { .. } => {
2731 return Err(anyhow::anyhow!(
2732 "TODO: Op::DeformConv — implement deformable convolution kernel"
2733 ));
2734 }
2735 Op::Col2Im { .. } => {
2736 return Err(anyhow::anyhow!(
2737 "TODO: Op::Col2Im — implement col2im (fold) kernel"
2738 ));
2739 }
2740 Op::Resize { .. } => {
2741 return Err(anyhow::anyhow!(
2742 "TODO: Op::Resize — implement nearest/bilinear/bicubic resize kernels"
2743 ));
2744 }
2745 Op::GridSample { .. } => {
2746 return Err(anyhow::anyhow!(
2747 "TODO: Op::GridSample — implement bilinear grid sample kernel"
2748 ));
2749 }
2750 Op::SpaceToDepth { .. } | Op::DepthToSpace { .. } => {
2751 return Err(anyhow::anyhow!(
2752 "TODO: Op::SpaceToDepth/DepthToSpace — implement pixel shuffle kernel"
2753 ));
2754 }
2755 Op::RoiAlign { .. } => {
2756 return Err(anyhow::anyhow!(
2757 "TODO: Op::RoiAlign — implement RoI-align pooling kernel"
2758 ));
2759 }
2760 Op::AffineGrid { .. } => {
2761 return Err(anyhow::anyhow!(
2762 "TODO: Op::AffineGrid — implement affine grid generator kernel"
2763 ));
2764 }
2765 Op::MaxUnpool { .. } => {
2766 return Err(anyhow::anyhow!(
2767 "TODO: Op::MaxUnpool — implement max-unpool (scatter with saved indices) kernel"
2768 ));
2769 }
2770 Op::CenterCropPad { .. } => {
2771 return Err(anyhow::anyhow!(
2772 "TODO: Op::CenterCropPad — implement center-crop-pad kernel"
2773 ));
2774 }
2775 Op::NonMaxSuppression { .. } => {
2776 return Err(anyhow::anyhow!(
2777 "TODO: Op::NonMaxSuppression — implement NMS kernel"
2778 ));
2779 }
2780 Op::TopK { .. } => {
2781 return Err(anyhow::anyhow!(
2782 "TODO: Op::TopK — implement radix sort / parallel selection kernel"
2783 ));
2784 }
2785 Op::Unique { .. } => {
2786 return Err(anyhow::anyhow!(
2787 "TODO: Op::Unique — implement stream-compaction unique kernel"
2788 ));
2789 }
2790 Op::EyeLike { .. }
2791 | Op::OneHot { .. }
2792 | Op::Bernoulli { .. }
2793 | Op::RandomUniformLike { .. } => {
2794 return Err(anyhow::anyhow!(
2795 "TODO: Op::EyeLike/OneHot/Bernoulli/RandomUniformLike — implement generation kernels"
2796 ));
2797 }
2798 Op::And | Op::Or | Op::Xor => {
2799 return Err(anyhow::anyhow!(
2800 "TODO: Op::And/Or/Xor — implement boolean logical kernels"
2801 ));
2802 }
2803 Op::BitShift { .. }
2804 | Op::BitwiseAnd
2805 | Op::BitwiseOr
2806 | Op::BitwiseXor
2807 | Op::BitwiseNot
2808 | Op::Not => {
2809 return Err(anyhow::anyhow!(
2810 "TODO: Op::Bitwise* — implement integer bitwise kernels"
2811 ));
2812 }
2813 Op::QuantizeLinear { .. }
2814 | Op::DequantizeLinear { .. }
2815 | Op::DynamicQuantizeLinear => {
2816 return Err(anyhow::anyhow!(
2817 "TODO: Op::Quantize/Dequantize — implement quantisation kernels"
2818 ));
2819 }
2820 Op::LRN { .. } => {
2821 return Err(anyhow::anyhow!(
2822 "TODO: Op::LRN — implement local response normalisation kernel"
2823 ));
2824 }
2825 Op::MeanVarianceNormalization { .. } | Op::LpNormalization { .. } => {
2826 return Err(anyhow::anyhow!(
2827 "TODO: Op::MvnNorm/LpNorm — implement normalisation kernels"
2828 ));
2829 }
2830 Op::Dft { .. }
2831 | Op::Stft
2832 | Op::MelWeightMatrix
2833 | Op::HannWindow { .. }
2834 | Op::BlackmanWindow { .. }
2835 | Op::HammingWindow { .. } => {
2836 return Err(anyhow::anyhow!(
2837 "TODO: Op::DFT/STFT/Window — implement signal processing kernels"
2838 ));
2839 }
2840 Op::NegativeLogLikelihoodLoss { .. } | Op::SoftmaxCrossEntropyLoss { .. } => {
2841 return Err(anyhow::anyhow!(
2842 "TODO: Op::NllLoss/SoftmaxCELoss — implement loss kernels"
2843 ));
2844 }
2845 Op::SequenceAt
2846 | Op::SequenceConstruct
2847 | Op::SequenceEmpty
2848 | Op::SequenceErase
2849 | Op::SequenceInsert
2850 | Op::SequenceLength
2851 | Op::SequenceMap
2852 | Op::SplitToSequence { .. }
2853 | Op::ConcatFromSequence { .. }
2854 | Op::OptionalGetElement
2855 | Op::OptionalHasElement
2856 | Op::Loop
2857 | Op::Scan { .. }
2858 | Op::If => {
2859 return Err(anyhow::anyhow!(
2860 "TODO: Op::Sequence/Control-flow — not lowerable to single Triton kernels"
2861 ));
2862 }
2863 Op::Adagrad | Op::Adam | Op::Momentum | Op::Gradient => {
2864 return Err(anyhow::anyhow!(
2865 "TODO: Op::OnnxOptimizer — use teenygrad's own optimizer kernels instead"
2866 ));
2867 }
2868 Op::StringNormalizer
2869 | Op::RegexFullMatch { .. }
2870 | Op::StringConcat
2871 | Op::StringSplit
2872 | Op::TfIdfVectorizer
2873 | Op::LabelEncoder
2874 | Op::ArrayFeatureExtractor
2875 | Op::Binarizer { .. }
2876 | Op::TreeEnsemble
2877 | Op::ImageDecoder => {
2878 return Err(anyhow::anyhow!(
2879 "TODO: Op::String/ClassicalML — not GPU-lowerable"
2880 ));
2881 }
2882 Op::CastLike | Op::BitCast { .. } => {
2883 return Err(anyhow::anyhow!(
2884 "TODO: Op::CastLike/BitCast — implement dtype-cast kernels"
2885 ));
2886 }
2887 Op::Cast { to } => {
2888 return Err(anyhow::anyhow!(
2889 "TODO: Op::Cast to {:?} — implement dtype-cast kernel",
2890 to
2891 ));
2892 }
2893 Op::Identity => {
2894 return Err(anyhow::anyhow!(
2895 "TODO: Op::Identity — implement zero-copy pass-through kernel"
2896 ));
2897 }
2898 Op::Dropout { .. } => {
2899 return Err(anyhow::anyhow!(
2900 "TODO: Op::Dropout — implement inference pass-through / training dropout kernel"
2901 ));
2902 }
2903 Op::IsInf { .. } => {
2904 return Err(anyhow::anyhow!("TODO: Op::IsInf — implement isinf kernel"));
2905 }
2906 Op::Hardmax { .. } => {
2907 return Err(anyhow::anyhow!(
2908 "TODO: Op::Hardmax — implement argmax + one-hot kernel"
2909 ));
2910 }
2911
2912 Op::Attention {
2913 c,
2914 num_heads,
2915 key_dim,
2916 } => {
2917 let _ = (c, num_heads, key_dim);
2921 return Err(anyhow::anyhow!(
2922 "Op::Attention reached the match arm — this should not happen"
2923 ));
2924 }
2925
2926 Op::Custom { data } => match data.0.lower() {
2927 Some((name, kernel_source, entry_point, runtime_op)) => {
2928 Box::new(KernelExecutable {
2929 name,
2930 kernel_source,
2931 entry_point,
2932 shape: node.shape.clone(),
2933 dtype: node.dtype,
2934 runtime_op,
2935 #[cfg(feature = "training")]
2936 backward_kernel_source: data.0.lower_backward_source(),
2937 #[cfg(feature = "training")]
2938 backward_entry_point: String::new(),
2939 })
2940 }
2941 None => {
2942 return Err(anyhow::anyhow!(
2943 "custom op '{}' is not handled — implement CustomOp::lower()",
2944 data.name()
2945 ));
2946 }
2947 },
2948
2949 Op::Fused { members } => {
2950 return Err(anyhow::anyhow!(
2956 "Op::Fused lowering is not implemented yet ({} member op(s)): \
2957 concatenate each member's kernel source and synthesize an \
2958 entry point that runs them in sequence (see spinorml-1fj.1)",
2959 members.len()
2960 ));
2961 }
2962 };
2963
2964 let dag_idx = dag.add_node(executable);
2965 graph_to_dag[node_index] = dag_idx;
2966
2967 for &input_graph_idx in &node.inputs {
2968 dag.add_edge(graph_to_dag[input_graph_idx], dag_idx);
2969 }
2970 }
2971
2972 Ok((dag, graph_to_dag))
2973 }
2974}
2975
2976impl<'a> Lowering<'a> for TritonLowering {
2977 fn lower(&self, graph: &Graph, mode: LoweringMode) -> Result<Dag<Box<dyn ExecutableOp>>> {
2978 TritonLowering::lower_with_mapping(self, graph, mode).map(|(dag, _)| dag)
2979 }
2980
2981 fn lower_with_mapping(
2982 &self,
2983 graph: &Graph,
2984 mode: LoweringMode,
2985 ) -> Result<(Dag<Box<dyn ExecutableOp>>, Vec<usize>)> {
2986 TritonLowering::lower_with_mapping(self, graph, mode)
2987 }
2988
2989 fn extra_dag_names(&self, graph: &Graph, graph_to_dag: &[usize]) -> Vec<(usize, String)> {
2994 let mut extra = Vec::new();
2995 for (graph_idx, node) in graph.nodes.iter().enumerate() {
2996 if let Op::Conv2d { has_bias: true, .. } = &node.op
2997 && let Some(name) = graph.names.get(&graph_idx)
2998 {
2999 let biasadd_dag_idx = graph_to_dag[graph_idx];
3000 if biasadd_dag_idx > 0 {
3001 extra.push((biasadd_dag_idx - 1, name.clone()));
3002 }
3003 }
3004 }
3005 extra
3006 }
3007}