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) {
637 (m, n) if m >= 128 && n >= 128 => (128, 128),
638 (m, _) if m >= 128 => (128, 64),
639 (_, n) if n >= 128 => (64, 128),
640 _ => (64, 64),
641 };
642 (block_m, block_n, block_k)
643}
644
645#[cfg(test)]
646mod pick_gemm_tile_sizes_tests {
647 use super::pick_gemm_tile_sizes;
648
649 #[test]
650 fn small_shape_gets_smallest_tiles() {
651 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 16), (64, 64, 8));
652 }
653
654 #[test]
655 fn large_m_and_n_get_the_largest_tile() {
656 assert_eq!(pick_gemm_tile_sizes(Some(512), 256, 256), (128, 128, 32));
658 }
659
660 #[test]
661 fn large_m_only_widens_block_m_not_block_n() {
662 assert_eq!(pick_gemm_tile_sizes(Some(512), 32, 64), (128, 64, 16));
663 }
664
665 #[test]
666 fn large_n_only_widens_block_n_not_block_m() {
667 assert_eq!(pick_gemm_tile_sizes(Some(32), 512, 64), (64, 128, 16));
668 }
669
670 #[test]
671 fn unknown_dynamic_batch_treated_as_small() {
672 assert_eq!(
674 pick_gemm_tile_sizes(None, 64, 16),
675 pick_gemm_tile_sizes(Some(64), 64, 16)
676 );
677 }
678
679 #[test]
680 fn block_k_never_drops_below_the_tensor_core_minimum() {
681 let (_, _, block_k) = pick_gemm_tile_sizes(Some(64), 64, 1);
682 assert!(block_k >= 8);
683 }
684
685 #[test]
686 fn block_k_grows_with_k() {
687 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 8).2, 8);
688 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 32).2, 16);
689 assert_eq!(pick_gemm_tile_sizes(Some(64), 64, 128).2, 32);
690 }
691}
692
693impl TritonLowering {
694 pub fn lower_with_mapping(
698 &self,
699 graph: &Graph,
700 mode: LoweringMode,
701 ) -> Result<(Dag<Box<dyn ExecutableOp>>, Vec<usize>)> {
702 let _ = mode; let node_indexes = graph.topological_sort();
704 let mut dag: Dag<Box<dyn ExecutableOp>> = Dag::new();
705 let mut graph_to_dag = vec![0usize; graph.nodes.len()];
707
708 for node_index in node_indexes {
709 let node = &graph.nodes[node_index];
710
711 #[cfg(feature = "training")]
715 if mode == LoweringMode::Training
716 && let Op::BatchNorm1d {
717 num_features,
718 eps,
719 momentum,
720 ..
721 }
722 | Op::BatchNorm3d {
723 num_features,
724 eps,
725 momentum,
726 ..
727 } = &node.op
728 {
729 let c = *num_features;
730 let eps_f32 = *eps as f32;
731 let momentum_f32 = *momentum as f32;
732 const BLOCK_N: i32 = 64;
733
734 let (stats_name, stats_src, stats_rop): (String, String, Arc<dyn RuntimeOp>) =
735 match node.dtype {
736 DtypeRepr::F32 => {
737 let k = BatchNormStatsForward::<f32>::new(BLOCK_N);
738 let src = k.source.clone();
739 let rop: Arc<dyn RuntimeOp> = Arc::new(
740 BatchNormStatsRuntimeOp::<f32>::new(BLOCK_N, eps_f32, momentum_f32),
741 );
742 (k.name.to_string(), src, rop)
743 }
744 DtypeRepr::F64 => {
745 let k = BatchNormStatsForward::<f64>::new(BLOCK_N);
746 let src = k.source.clone();
747 let rop: Arc<dyn RuntimeOp> = Arc::new(
748 BatchNormStatsRuntimeOp::<f64>::new(BLOCK_N, eps_f32, momentum_f32),
749 );
750 (k.name.to_string(), src, rop)
751 }
752 other => {
753 return Err(anyhow::anyhow!(
754 "{:?} is not a Float dtype for BatchNormStatsForward",
755 other
756 ));
757 }
758 };
759
760 let stats_node = Box::new(KernelExecutable {
761 entry_point: format!("{}_entry_point", stats_name),
762 name: stats_name,
763 kernel_source: stats_src,
764 shape: vec![Some(2 * c)],
765 dtype: node.dtype,
766 backward_kernel_source: String::new(),
767 backward_entry_point: String::new(),
768 runtime_op: stats_rop,
769 }) as Box<dyn ExecutableOp>;
770
771 let stats_dag_idx = dag.add_node(stats_node);
772 for &input_graph_idx in &node.inputs {
773 dag.add_edge(graph_to_dag[input_graph_idx], stats_dag_idx);
774 }
775
776 let (norm_name, norm_src, norm_bwd_src, norm_rop): (
777 String,
778 String,
779 String,
780 Arc<dyn RuntimeOp>,
781 ) = match node.dtype {
782 DtypeRepr::F32 => {
783 let k = BatchNormNormalizeForward::<f32>::new(BLOCK_N);
784 let src = k.source.clone();
785 let rop = BatchNormNormalizeRuntimeOp::<f32>::new(BLOCK_N);
786 let bwd_src = rop.backward_source().to_string();
787 (
788 k.name.to_string(),
789 src,
790 bwd_src,
791 Arc::new(rop) as Arc<dyn RuntimeOp>,
792 )
793 }
794 DtypeRepr::F64 => {
795 let k = BatchNormNormalizeForward::<f64>::new(BLOCK_N);
796 let src = k.source.clone();
797 let rop = BatchNormNormalizeRuntimeOp::<f64>::new(BLOCK_N);
798 let bwd_src = rop.backward_source().to_string();
799 (
800 k.name.to_string(),
801 src,
802 bwd_src,
803 Arc::new(rop) as Arc<dyn RuntimeOp>,
804 )
805 }
806 other => {
807 return Err(anyhow::anyhow!(
808 "{:?} is not a Float dtype for BatchNormNormalizeForward",
809 other
810 ));
811 }
812 };
813
814 let norm_node = Box::new(KernelExecutable {
815 entry_point: format!("{}_entry_point", norm_name),
816 name: norm_name,
817 kernel_source: norm_src,
818 shape: node.shape.clone(),
819 dtype: node.dtype,
820 backward_kernel_source: norm_bwd_src,
821 backward_entry_point: String::new(),
822 runtime_op: norm_rop,
823 }) as Box<dyn ExecutableOp>;
824
825 let norm_dag_idx = dag.add_node(norm_node);
826 for &input_graph_idx in &node.inputs {
828 dag.add_edge(graph_to_dag[input_graph_idx], norm_dag_idx);
829 }
830 dag.add_edge(stats_dag_idx, norm_dag_idx);
832
833 graph_to_dag[node_index] = norm_dag_idx;
834 continue;
835 }
836
837 if mode == LoweringMode::Inference
841 && let Op::Conv2d {
842 has_bias: true,
843 kernel_h,
844 kernel_w,
845 stride_h,
846 stride_w,
847 padding_h,
848 padding_w,
849 groups,
850 ..
851 } = &node.op
852 {
853 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
854 DtypeRepr::F32 => {
855 let k = Conv2dBiasForward::<f32>::new(
856 *kernel_h as i32,
857 *kernel_w as i32,
858 *stride_h as i32,
859 *stride_w as i32,
860 *padding_h as i32,
861 *padding_w as i32,
862 *groups as i32,
863 16,
864 );
865 let nm = k.name.to_string();
866 let src = k.source.clone();
867 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
868 (nm, src, rop)
869 }
870 DtypeRepr::F64 => {
871 let k = Conv2dBiasForward::<f64>::new(
872 *kernel_h as i32,
873 *kernel_w as i32,
874 *stride_h as i32,
875 *stride_w as i32,
876 *padding_h as i32,
877 *padding_w as i32,
878 *groups as i32,
879 16,
880 );
881 let nm = k.name.to_string();
882 let src = k.source.clone();
883 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
884 (nm, src, rop)
885 }
886 other => {
887 return Err(anyhow::anyhow!(
888 "{:?} is not supported for Conv2dBiasForward",
889 other
890 ));
891 }
892 };
893 let dag_idx = dag.add_node(Box::new(KernelExecutable {
894 entry_point: format!("{}_entry_point", name),
895 name,
896 kernel_source: ks,
897 shape: node.shape.clone(),
898 dtype: node.dtype,
899 #[cfg(feature = "training")]
900 backward_kernel_source: String::new(),
901 #[cfg(feature = "training")]
902 backward_entry_point: String::new(),
903 runtime_op: rop,
904 }) as Box<dyn ExecutableOp>);
905 for &input_graph_idx in &node.inputs {
906 dag.add_edge(graph_to_dag[input_graph_idx], dag_idx);
907 }
908 graph_to_dag[node_index] = dag_idx;
909 continue;
910 }
911
912 if let Op::Conv2d {
915 has_bias: true,
916 kernel_h,
917 kernel_w,
918 stride_h,
919 stride_w,
920 padding_h,
921 padding_w,
922 groups,
923 out_channels,
924 ..
925 } = &node.op
926 {
927 const BIAS_BLOCK_HW: i32 = 128;
928 let (conv_name, conv_ks, conv_rop): (String, String, Arc<dyn RuntimeOp>) =
929 match node.dtype {
930 DtypeRepr::F32 => {
931 let k = Conv2dForward::<f32>::new(
932 *kernel_h as i32,
933 *kernel_w as i32,
934 *stride_h as i32,
935 *stride_w as i32,
936 *padding_h as i32,
937 *padding_w as i32,
938 *groups as i32,
939 16,
940 );
941 let src = k.source.clone();
942 let rop: Arc<dyn RuntimeOp> = Arc::new(Conv2dForward::<f32>::new(
943 *kernel_h as i32,
944 *kernel_w as i32,
945 *stride_h as i32,
946 *stride_w as i32,
947 *padding_h as i32,
948 *padding_w as i32,
949 *groups as i32,
950 16,
951 ));
952 (k.name.to_string(), src, rop)
953 }
954 DtypeRepr::F64 => {
955 let k = Conv2dForward::<f64>::new(
956 *kernel_h as i32,
957 *kernel_w as i32,
958 *stride_h as i32,
959 *stride_w as i32,
960 *padding_h as i32,
961 *padding_w as i32,
962 *groups as i32,
963 16,
964 );
965 let src = k.source.clone();
966 let rop: Arc<dyn RuntimeOp> = Arc::new(Conv2dForward::<f64>::new(
967 *kernel_h as i32,
968 *kernel_w as i32,
969 *stride_h as i32,
970 *stride_w as i32,
971 *padding_h as i32,
972 *padding_w as i32,
973 *groups as i32,
974 16,
975 ));
976 (k.name.to_string(), src, rop)
977 }
978 other => {
979 return Err(anyhow::anyhow!(
980 "{:?} is not supported for Conv2dForward",
981 other
982 ));
983 }
984 };
985
986 #[cfg(feature = "training")]
987 let conv_bwd_ks = match node.dtype {
988 DtypeRepr::F32 => {
989 Conv2dBackward::<f32>::new(
990 *kernel_h as i32,
991 *kernel_w as i32,
992 *stride_h as i32,
993 *stride_w as i32,
994 *padding_h as i32,
995 *padding_w as i32,
996 *groups as i32,
997 16,
998 )
999 .source
1000 }
1001 DtypeRepr::F64 => {
1002 Conv2dBackward::<f64>::new(
1003 *kernel_h as i32,
1004 *kernel_w as i32,
1005 *stride_h as i32,
1006 *stride_w as i32,
1007 *padding_h as i32,
1008 *padding_w as i32,
1009 *groups as i32,
1010 16,
1011 )
1012 .source
1013 }
1014 _ => String::new(),
1015 };
1016
1017 let conv_dag_idx = dag.add_node(Box::new(KernelExecutable {
1018 entry_point: format!("{}_entry_point", conv_name),
1019 name: conv_name,
1020 kernel_source: conv_ks,
1021 shape: node.shape.clone(),
1022 dtype: node.dtype,
1023 #[cfg(feature = "training")]
1024 backward_kernel_source: conv_bwd_ks,
1025 #[cfg(feature = "training")]
1026 backward_entry_point: String::new(),
1027 runtime_op: conv_rop,
1028 }) as Box<dyn ExecutableOp>);
1029 for &input_graph_idx in &node.inputs {
1030 dag.add_edge(graph_to_dag[input_graph_idx], conv_dag_idx);
1031 }
1032
1033 let (bias_name, bias_ks, bias_rop): (String, String, Arc<dyn RuntimeOp>) =
1034 match node.dtype {
1035 DtypeRepr::F32 => {
1036 let r = NchwBiasAddRuntimeOp::<f32>::new(BIAS_BLOCK_HW);
1037 (
1038 r.kernel_name().to_string(),
1039 r.forward_source().to_string(),
1040 Arc::new(r),
1041 )
1042 }
1043 DtypeRepr::F64 => {
1044 let r = NchwBiasAddRuntimeOp::<f64>::new(BIAS_BLOCK_HW);
1045 (
1046 r.kernel_name().to_string(),
1047 r.forward_source().to_string(),
1048 Arc::new(r),
1049 )
1050 }
1051 other => {
1052 return Err(anyhow::anyhow!(
1053 "{:?} is not supported for NchwBiasAdd",
1054 other
1055 ));
1056 }
1057 };
1058
1059 #[cfg(feature = "training")]
1060 let bias_bwd_ks = match node.dtype {
1061 DtypeRepr::F32 => NchwBiasAddRuntimeOp::<f32>::new(BIAS_BLOCK_HW)
1062 .backward_source()
1063 .to_string(),
1064 DtypeRepr::F64 => NchwBiasAddRuntimeOp::<f64>::new(BIAS_BLOCK_HW)
1065 .backward_source()
1066 .to_string(),
1067 _ => String::new(),
1068 };
1069
1070 let biasadd_dag_idx = dag.add_node(Box::new(KernelExecutable {
1071 entry_point: format!("{}_entry_point", bias_name),
1072 name: bias_name,
1073 kernel_source: bias_ks,
1074 shape: node.shape.clone(),
1075 dtype: node.dtype,
1076 #[cfg(feature = "training")]
1077 backward_kernel_source: bias_bwd_ks,
1078 #[cfg(feature = "training")]
1079 backward_entry_point: String::new(),
1080 runtime_op: bias_rop,
1081 }) as Box<dyn ExecutableOp>);
1082 dag.add_edge(conv_dag_idx, biasadd_dag_idx);
1083 graph_to_dag[node_index] = biasadd_dag_idx;
1084 let _ = (conv_dag_idx, out_channels);
1087 continue;
1088 }
1089
1090 let executable: Box<dyn ExecutableOp> = match &node.op {
1091 Op::Input => Box::new(KernelExecutable {
1092 name: "input".to_string(),
1093 kernel_source: String::new(),
1094 entry_point: String::new(),
1095 shape: node.shape.clone(),
1096 dtype: node.dtype,
1097 #[cfg(feature = "training")]
1098 backward_kernel_source: String::new(),
1099 #[cfg(feature = "training")]
1100 backward_entry_point: String::new(),
1101 runtime_op: Arc::new(InputRuntimeOp),
1102 }),
1103
1104 Op::Linear { has_bias, .. } => {
1106 make_num_kernel!(
1107 LinearForward(*has_bias, 32, 64, 32, 8),
1108 LinearBackward(*has_bias, 32, 64, 32, 8),
1109 node
1110 )
1111 }
1112 Op::Flatten => make_num_kernel!(FlattenForward(32, 256), node),
1113
1114 Op::BatchNorm1d { .. } | Op::BatchNorm3d { .. } => {
1116 make_float_kernel!(BatchNormForwardInference(64), node)
1117 }
1118 Op::BatchNorm2d { eps, .. } => {
1119 let eps_f32 = *eps as f32;
1120 const BN2D_BLOCK_HW: i32 = 128;
1121 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
1122 DtypeRepr::F32 => {
1123 let r = BatchNorm2dNchwInferenceRuntimeOp::<f32>::new(
1124 BN2D_BLOCK_HW,
1125 eps_f32,
1126 );
1127 (
1128 r.kernel_name().to_string(),
1129 r.forward_source().to_string(),
1130 Arc::new(r),
1131 )
1132 }
1133 DtypeRepr::F64 => {
1134 let r = BatchNorm2dNchwInferenceRuntimeOp::<f64>::new(
1135 BN2D_BLOCK_HW,
1136 eps_f32,
1137 );
1138 (
1139 r.kernel_name().to_string(),
1140 r.forward_source().to_string(),
1141 Arc::new(r),
1142 )
1143 }
1144 other => {
1145 return Err(anyhow::anyhow!(
1146 "{:?} is not a Float dtype for BatchNorm2d",
1147 other
1148 ));
1149 }
1150 };
1151 #[cfg(feature = "training")]
1152 let bwd_ks = match node.dtype {
1153 DtypeRepr::F32 => BatchNorm2dNchwBackward::<f32>::new(BN2D_BLOCK_HW).source,
1154 DtypeRepr::F64 => BatchNorm2dNchwBackward::<f64>::new(BN2D_BLOCK_HW).source,
1155 _ => String::new(),
1156 };
1157 Box::new(KernelExecutable {
1158 entry_point: format!("{}_entry_point", name),
1159 name,
1160 kernel_source: ks,
1161 shape: node.shape.clone(),
1162 dtype: node.dtype,
1163 #[cfg(feature = "training")]
1164 backward_kernel_source: bwd_ks,
1165 #[cfg(feature = "training")]
1166 backward_entry_point: String::new(),
1167 runtime_op: rop,
1168 })
1169 }
1170 Op::LayerNorm { eps, .. } => {
1171 let eps_f32 = *eps as f32;
1172 const LN_BLOCK_N: i32 = 1024;
1173 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
1174 DtypeRepr::F32 => {
1175 let r =
1176 LayerNormForwardInferenceRuntimeOp::<f32>::new(LN_BLOCK_N, eps_f32);
1177 (
1178 r.kernel_name().to_string(),
1179 r.forward_source().to_string(),
1180 Arc::new(r),
1181 )
1182 }
1183 DtypeRepr::F64 => {
1184 let r =
1185 LayerNormForwardInferenceRuntimeOp::<f64>::new(LN_BLOCK_N, eps_f32);
1186 (
1187 r.kernel_name().to_string(),
1188 r.forward_source().to_string(),
1189 Arc::new(r),
1190 )
1191 }
1192 other => {
1193 return Err(anyhow::anyhow!(
1194 "{:?} is not a Float dtype for LayerNorm",
1195 other
1196 ));
1197 }
1198 };
1199 Box::new(KernelExecutable {
1200 entry_point: format!("{}_entry_point", name),
1201 name,
1202 kernel_source: ks,
1203 shape: node.shape.clone(),
1204 dtype: node.dtype,
1205 #[cfg(feature = "training")]
1206 backward_kernel_source: String::new(),
1207 #[cfg(feature = "training")]
1208 backward_entry_point: String::new(),
1209 runtime_op: rop,
1210 })
1211 }
1212 Op::RmsNorm { .. } => {
1213 make_float_kernel!(RmsNormForward(1024), node)
1214 }
1215 Op::GroupNorm { .. } => {
1216 make_float_kernel!(GroupNormForwardInference(256), node)
1217 }
1218 Op::InstanceNorm1d { .. }
1219 | Op::InstanceNorm2d { .. }
1220 | Op::InstanceNorm3d { .. } => {
1221 make_float_kernel!(InstanceNormForwardInference(256), node)
1222 }
1223
1224 Op::Conv1d {
1226 kernel_l,
1227 stride,
1228 padding,
1229 ..
1230 } => {
1231 make_num_kernel!(
1232 Conv1dForward(*kernel_l as i32, *stride as i32, *padding as i32, 32),
1233 node
1234 )
1235 }
1236 Op::Conv2d {
1237 kernel_h,
1238 kernel_w,
1239 stride_h,
1240 stride_w,
1241 padding_h,
1242 padding_w,
1243 groups,
1244 ..
1245 } => {
1246 make_num_kernel!(
1247 Conv2dForward(
1248 *kernel_h as i32,
1249 *kernel_w as i32,
1250 *stride_h as i32,
1251 *stride_w as i32,
1252 *padding_h as i32,
1253 *padding_w as i32,
1254 *groups as i32,
1255 16
1256 ),
1257 Conv2dBackward(
1258 *kernel_h as i32,
1259 *kernel_w as i32,
1260 *stride_h as i32,
1261 *stride_w as i32,
1262 *padding_h as i32,
1263 *padding_w as i32,
1264 *groups as i32,
1265 16
1266 ),
1267 node
1268 )
1269 }
1270 Op::Conv3d {
1271 kernel_d,
1272 kernel_h,
1273 kernel_w,
1274 stride_d,
1275 stride_h,
1276 stride_w,
1277 padding_d,
1278 padding_h,
1279 padding_w,
1280 ..
1281 } => {
1282 make_num_kernel!(
1283 Conv3dForward(
1284 *kernel_d as i32,
1285 *kernel_h as i32,
1286 *kernel_w as i32,
1287 *stride_d as i32,
1288 *stride_h as i32,
1289 *stride_w as i32,
1290 *padding_d as i32,
1291 *padding_h as i32,
1292 *padding_w as i32,
1293 8
1294 ),
1295 node
1296 )
1297 }
1298
1299 Op::Conv2dBnSilu {
1300 kernel_h,
1301 kernel_w,
1302 stride_h,
1303 stride_w,
1304 padding_h,
1305 padding_w,
1306 groups,
1307 in_channels,
1308 ..
1309 } => {
1310 if node.dtype != DtypeRepr::F32 {
1311 return Err(anyhow::anyhow!(
1312 "Conv2dBnSilu only supports f32 (got {:?})",
1313 node.dtype
1314 ));
1315 }
1316 let kh = *kernel_h as i32;
1317 let kw = *kernel_w as i32;
1318 let sh = *stride_h as i32;
1319 let sw = *stride_w as i32;
1320 let ph = *padding_h as i32;
1321 let pw = *padding_w as i32;
1322 let g = *groups as i32;
1323 let c_out = node.shape[1].unwrap_or(0);
1324 let oh = node.shape[2].unwrap_or(1);
1325 let ow = node.shape[3].unwrap_or(1);
1326 let is_depthwise = g as usize == *in_channels;
1327
1328 let use_gemm = kh == 1
1342 && kw == 1
1343 && sh == 1
1344 && sw == 1
1345 && ph == 0
1346 && pw == 0
1347 && !is_depthwise
1348 && g == 1
1349 && c_out >= 32;
1350
1351 let use_tiled = !is_depthwise && g == 1 && c_out >= 16 && !use_gemm;
1352
1353 if use_gemm {
1354 const GROUP_M: i32 = 8;
1355 let m = oh * ow;
1356 let (block_m, block_n_base, block_k) =
1361 pick_gemm_tile_sizes(Some(m), c_out, *in_channels);
1362 let block_n = match self.sm_count {
1369 Some(sm_count) => {
1370 let fixed_blocks = m.div_ceil(block_m as usize);
1371 pick_adaptive_block_n(
1372 c_out,
1373 fixed_blocks,
1374 4 * sm_count,
1375 &[block_n_base, 16, 8],
1376 )
1377 }
1378 None => block_n_base,
1379 };
1380 let k = Conv2dBnSiluGemmForward::new(block_m, block_n, block_k, GROUP_M);
1381 let nm = k.name.to_string();
1382 let ks = k.source.clone();
1383 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
1384 Box::new(KernelExecutable {
1385 entry_point: format!("{}_entry_point", nm),
1386 name: nm,
1387 kernel_source: ks,
1388 shape: node.shape.clone(),
1389 dtype: node.dtype,
1390 #[cfg(feature = "training")]
1391 backward_kernel_source: String::new(),
1392 #[cfg(feature = "training")]
1393 backward_entry_point: String::new(),
1394 runtime_op: rop,
1395 })
1396 } else if use_tiled {
1397 const BLOCK_OW: i32 = 16;
1398 const BLOCK_N_TILE: i32 = 16;
1399 let block_n_tile = match self.sm_count {
1404 Some(sm_count) => {
1405 let fixed_blocks = oh * ow.div_ceil(BLOCK_OW as usize);
1406 pick_adaptive_block_n(
1407 c_out,
1408 fixed_blocks,
1409 4 * sm_count,
1410 &[BLOCK_N_TILE, 8, 4],
1411 )
1412 }
1413 None => BLOCK_N_TILE,
1414 };
1415 let k = Conv2dBnSiluTiledForward::new(
1416 kh,
1417 kw,
1418 sh,
1419 sw,
1420 ph,
1421 pw,
1422 BLOCK_OW,
1423 block_n_tile,
1424 );
1425 let nm = k.name.to_string();
1426 let ks = k.source.clone();
1427 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
1428 Box::new(KernelExecutable {
1429 entry_point: format!("{}_entry_point", nm),
1430 name: nm,
1431 kernel_source: ks,
1432 shape: node.shape.clone(),
1433 dtype: node.dtype,
1434 #[cfg(feature = "training")]
1435 backward_kernel_source: String::new(),
1436 #[cfg(feature = "training")]
1437 backward_entry_point: String::new(),
1438 runtime_op: rop,
1439 })
1440 } else {
1441 const BLOCK_OW: i32 = 16;
1442 let k = Conv2dBnSiluForward::new(kh, kw, sh, sw, ph, pw, g, BLOCK_OW);
1443 let nm = k.name.to_string();
1444 let ks = k.source.clone();
1445 let rop: Arc<dyn RuntimeOp> = Arc::new(k);
1446 Box::new(KernelExecutable {
1447 entry_point: format!("{}_entry_point", nm),
1448 name: nm,
1449 kernel_source: ks,
1450 shape: node.shape.clone(),
1451 dtype: node.dtype,
1452 #[cfg(feature = "training")]
1453 backward_kernel_source: String::new(),
1454 #[cfg(feature = "training")]
1455 backward_entry_point: String::new(),
1456 runtime_op: rop,
1457 })
1458 }
1459 }
1460
1461 Op::AvgPool1d { kernel_l, stride } => {
1463 make_num_kernel!(Avgpool1dForward(*kernel_l as i32, *stride as i32, 32), node)
1464 }
1465 Op::AvgPool2d {
1466 kernel_h,
1467 kernel_w,
1468 stride_h,
1469 stride_w,
1470 } => {
1471 make_num_kernel!(
1472 Avgpool2dForward(
1473 *kernel_h as i32,
1474 *kernel_w as i32,
1475 *stride_h as i32,
1476 *stride_w as i32,
1477 16
1478 ),
1479 node
1480 )
1481 }
1482 Op::AvgPool3d {
1483 kernel_d,
1484 kernel_h,
1485 kernel_w,
1486 stride_d,
1487 stride_h,
1488 stride_w,
1489 } => {
1490 make_num_kernel!(
1491 Avgpool3dForward(
1492 *kernel_d as i32,
1493 *kernel_h as i32,
1494 *kernel_w as i32,
1495 *stride_d as i32,
1496 *stride_h as i32,
1497 *stride_w as i32,
1498 8
1499 ),
1500 node
1501 )
1502 }
1503 Op::MaxPool1d { kernel_l, stride } => {
1504 make_num_kernel!(Maxpool1dForward(*kernel_l as i32, *stride as i32, 32), node)
1505 }
1506 Op::MaxPool2d {
1507 kernel_h,
1508 kernel_w,
1509 stride_h,
1510 stride_w,
1511 pad_h,
1512 pad_w,
1513 } => {
1514 make_num_kernel!(
1515 Maxpool2dForward(
1516 *kernel_h as i32,
1517 *kernel_w as i32,
1518 *stride_h as i32,
1519 *stride_w as i32,
1520 *pad_h as i32,
1521 *pad_w as i32,
1522 16
1523 ),
1524 Maxpool2dBackward(
1525 *kernel_h as i32,
1526 *kernel_w as i32,
1527 *stride_h as i32,
1528 *stride_w as i32,
1529 *pad_h as i32,
1530 *pad_w as i32,
1531 16
1532 ),
1533 node
1534 )
1535 }
1536 Op::MaxPool3d {
1537 kernel_d,
1538 kernel_h,
1539 kernel_w,
1540 stride_d,
1541 stride_h,
1542 stride_w,
1543 } => {
1544 make_num_kernel!(
1545 Maxpool3dForward(
1546 *kernel_d as i32,
1547 *kernel_h as i32,
1548 *kernel_w as i32,
1549 *stride_d as i32,
1550 *stride_h as i32,
1551 *stride_w as i32,
1552 8
1553 ),
1554 node
1555 )
1556 }
1557 Op::LpPool1d {
1558 kernel_l, stride, ..
1559 } => {
1560 make_float_kernel!(Lppool1dForward(*kernel_l as i32, *stride as i32, 32), node)
1561 }
1562 Op::LpPool2d {
1563 kernel_h,
1564 kernel_w,
1565 stride_h,
1566 stride_w,
1567 ..
1568 } => {
1569 make_float_kernel!(
1570 Lppool2dForward(
1571 *kernel_h as i32,
1572 *kernel_w as i32,
1573 *stride_h as i32,
1574 *stride_w as i32,
1575 16
1576 ),
1577 node
1578 )
1579 }
1580 Op::LpPool3d {
1581 kernel_d,
1582 kernel_h,
1583 kernel_w,
1584 stride_d,
1585 stride_h,
1586 stride_w,
1587 ..
1588 } => {
1589 make_float_kernel!(
1590 Lppool3dForward(
1591 *kernel_d as i32,
1592 *kernel_h as i32,
1593 *kernel_w as i32,
1594 *stride_d as i32,
1595 *stride_h as i32,
1596 *stride_w as i32,
1597 8
1598 ),
1599 node
1600 )
1601 }
1602
1603 Op::ConstantPad1d {
1605 pad_left,
1606 pad_right,
1607 ..
1608 } => {
1609 make_num_kernel!(
1610 ConstantPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1611 node
1612 )
1613 }
1614 Op::ConstantPad2d {
1615 pad_l,
1616 pad_r,
1617 pad_t,
1618 pad_b,
1619 ..
1620 } => {
1621 make_num_kernel!(
1622 ConstantPad2dForward(
1623 *pad_t as i32,
1624 *pad_b as i32,
1625 *pad_l as i32,
1626 *pad_r as i32,
1627 16
1628 ),
1629 node
1630 )
1631 }
1632 Op::ConstantPad3d {
1633 pad_d1,
1634 pad_d2,
1635 pad_h1,
1636 pad_h2,
1637 pad_w1,
1638 pad_w2,
1639 ..
1640 } => {
1641 make_num_kernel!(
1642 ConstantPad3dForward(
1643 *pad_d1 as i32,
1644 *pad_d2 as i32,
1645 *pad_h1 as i32,
1646 *pad_h2 as i32,
1647 *pad_w1 as i32,
1648 *pad_w2 as i32,
1649 8
1650 ),
1651 node
1652 )
1653 }
1654 Op::ReflectionPad1d {
1655 pad_left,
1656 pad_right,
1657 } => {
1658 make_num_kernel!(
1659 ReflectionPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1660 node
1661 )
1662 }
1663 Op::ReflectionPad2d {
1664 pad_l,
1665 pad_r,
1666 pad_t,
1667 pad_b,
1668 } => {
1669 make_num_kernel!(
1670 ReflectionPad2dForward(
1671 *pad_t as i32,
1672 *pad_b as i32,
1673 *pad_l as i32,
1674 *pad_r as i32,
1675 16
1676 ),
1677 node
1678 )
1679 }
1680 Op::ReflectionPad3d {
1681 pad_d1,
1682 pad_d2,
1683 pad_h1,
1684 pad_h2,
1685 pad_w1,
1686 pad_w2,
1687 } => {
1688 make_num_kernel!(
1689 ReflectionPad3dForward(
1690 *pad_d1 as i32,
1691 *pad_d2 as i32,
1692 *pad_h1 as i32,
1693 *pad_h2 as i32,
1694 *pad_w1 as i32,
1695 *pad_w2 as i32,
1696 8
1697 ),
1698 node
1699 )
1700 }
1701 Op::ReplicationPad1d {
1702 pad_left,
1703 pad_right,
1704 } => {
1705 make_num_kernel!(
1706 ReplicationPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1707 node
1708 )
1709 }
1710 Op::ReplicationPad2d {
1711 pad_l,
1712 pad_r,
1713 pad_t,
1714 pad_b,
1715 } => {
1716 make_num_kernel!(
1717 ReplicationPad2dForward(
1718 *pad_t as i32,
1719 *pad_b as i32,
1720 *pad_l as i32,
1721 *pad_r as i32,
1722 16
1723 ),
1724 node
1725 )
1726 }
1727 Op::ReplicationPad3d {
1728 pad_d1,
1729 pad_d2,
1730 pad_h1,
1731 pad_h2,
1732 pad_w1,
1733 pad_w2,
1734 } => {
1735 make_num_kernel!(
1736 ReplicationPad3dForward(
1737 *pad_d1 as i32,
1738 *pad_d2 as i32,
1739 *pad_h1 as i32,
1740 *pad_h2 as i32,
1741 *pad_w1 as i32,
1742 *pad_w2 as i32,
1743 8
1744 ),
1745 node
1746 )
1747 }
1748 Op::CircularPad1d {
1749 pad_left,
1750 pad_right,
1751 } => {
1752 make_num_kernel!(
1753 CircularPad1dForward(*pad_left as i32, *pad_right as i32, 32),
1754 node
1755 )
1756 }
1757 Op::CircularPad2d {
1758 pad_l,
1759 pad_r,
1760 pad_t,
1761 pad_b,
1762 } => {
1763 make_num_kernel!(
1764 CircularPad2dForward(
1765 *pad_t as i32,
1766 *pad_b as i32,
1767 *pad_l as i32,
1768 *pad_r as i32,
1769 16
1770 ),
1771 node
1772 )
1773 }
1774 Op::CircularPad3d {
1775 pad_d1,
1776 pad_d2,
1777 pad_h1,
1778 pad_h2,
1779 pad_w1,
1780 pad_w2,
1781 } => {
1782 make_num_kernel!(
1783 CircularPad3dForward(
1784 *pad_d1 as i32,
1785 *pad_d2 as i32,
1786 *pad_h1 as i32,
1787 *pad_h2 as i32,
1788 *pad_w1 as i32,
1789 *pad_w2 as i32,
1790 8
1791 ),
1792 node
1793 )
1794 }
1795
1796 Op::Relu => make_num_kernel!(ReluForward(1024), ReluBackward(1024), node),
1798
1799 Op::Elu { .. } => exec_from(
1801 node.shape.clone(),
1802 node.dtype,
1803 EluForwardDispatch::dispatch(node.dtype, 1024)?,
1804 ),
1805 Op::Selu => exec_from(
1806 node.shape.clone(),
1807 node.dtype,
1808 SeluForwardDispatch::dispatch(node.dtype, 1024)?,
1809 ),
1810 Op::Celu { .. } => exec_from(
1811 node.shape.clone(),
1812 node.dtype,
1813 CeluForwardDispatch::dispatch(node.dtype, 1024)?,
1814 ),
1815 Op::Gelu => exec_from(
1816 node.shape.clone(),
1817 node.dtype,
1818 GeluForwardDispatch::dispatch(node.dtype, 1024)?,
1819 ),
1820 Op::Mish => exec_from(
1821 node.shape.clone(),
1822 node.dtype,
1823 MishForwardDispatch::dispatch(node.dtype, 1024)?,
1824 ),
1825 Op::Hardtanh { .. } => exec_from(
1826 node.shape.clone(),
1827 node.dtype,
1828 HardtanhForwardDispatch::dispatch(node.dtype, 1024)?,
1829 ),
1830 Op::Relu6 => exec_from(
1831 node.shape.clone(),
1832 node.dtype,
1833 Relu6ForwardDispatch::dispatch(node.dtype, 1024)?,
1834 ),
1835 Op::Hardsigmoid => exec_from(
1836 node.shape.clone(),
1837 node.dtype,
1838 HardsigmoidForwardDispatch::dispatch(node.dtype, 1024)?,
1839 ),
1840 Op::Hardswish => exec_from(
1841 node.shape.clone(),
1842 node.dtype,
1843 HardswishForwardDispatch::dispatch(node.dtype, 1024)?,
1844 ),
1845 Op::Hardshrink { .. } => exec_from(
1846 node.shape.clone(),
1847 node.dtype,
1848 HardshrinkForwardDispatch::dispatch(node.dtype, 1024)?,
1849 ),
1850 Op::LeakyRelu { .. } => exec_from(
1851 node.shape.clone(),
1852 node.dtype,
1853 LeakyReluForwardDispatch::dispatch(node.dtype, 1024)?,
1854 ),
1855 Op::Threshold { .. } => exec_from(
1856 node.shape.clone(),
1857 node.dtype,
1858 ThresholdForwardDispatch::dispatch(node.dtype, 1024)?,
1859 ),
1860 Op::Softsign => exec_from(
1861 node.shape.clone(),
1862 node.dtype,
1863 SoftsignForwardDispatch::dispatch(node.dtype, 1024)?,
1864 ),
1865 Op::Softshrink { .. } => exec_from(
1866 node.shape.clone(),
1867 node.dtype,
1868 SoftshrinkForwardDispatch::dispatch(node.dtype, 1024)?,
1869 ),
1870 Op::Softplus { .. } => exec_from(
1871 node.shape.clone(),
1872 node.dtype,
1873 SoftplusForwardDispatch::dispatch(node.dtype, 1024)?,
1874 ),
1875 Op::Sigmoid => exec_from(
1876 node.shape.clone(),
1877 node.dtype,
1878 SigmoidForwardDispatch::dispatch(node.dtype, 1024)?,
1879 ),
1880 Op::Silu => exec_from(
1881 node.shape.clone(),
1882 node.dtype,
1883 SiluForwardDispatch::dispatch(node.dtype, 1024)?,
1884 ),
1885 Op::Logsigmoid => exec_from(
1886 node.shape.clone(),
1887 node.dtype,
1888 LogsigmoidForwardDispatch::dispatch(node.dtype, 1024)?,
1889 ),
1890 Op::Tanh => exec_from(
1891 node.shape.clone(),
1892 node.dtype,
1893 TanhForwardDispatch::dispatch(node.dtype, 1024)?,
1894 ),
1895 Op::Tanhshrink => exec_from(
1896 node.shape.clone(),
1897 node.dtype,
1898 TanhshrinkForwardDispatch::dispatch(node.dtype, 1024)?,
1899 ),
1900
1901 Op::Softmax { .. } => {
1903 let n_cols = node.shape.last().and_then(|d| *d).unwrap_or(1024);
1905 let block_size = n_cols.next_power_of_two() as i32;
1906 make_float_kernel!(SoftmaxForward(block_size), node)
1907 }
1908
1909 Op::UpsampleNearest2d { scale_h, scale_w } => {
1911 make_num_kernel!(
1912 UpsampleNearest2dForward(*scale_h as i32, *scale_w as i32, 16),
1913 UpsampleNearest2dBackward(*scale_h as i32, *scale_w as i32, 16),
1914 node
1915 )
1916 }
1917
1918 Op::ChannelCat { .. } => {
1919 let n_inputs = node.inputs.len();
1920 let (name, fwd_src, bwd_src, rop): (
1921 String,
1922 String,
1923 String,
1924 Arc<dyn RuntimeOp>,
1925 ) = match node.dtype {
1926 DtypeRepr::F32 => {
1927 let r = ChannelCatRuntimeOp::<f32>::new(128, n_inputs);
1928 (
1929 r.kernel_name().to_string(),
1930 r.forward_source().to_string(),
1931 r.backward_source().to_string(),
1932 Arc::new(r),
1933 )
1934 }
1935 DtypeRepr::F64 => {
1936 let r = ChannelCatRuntimeOp::<f64>::new(128, n_inputs);
1937 (
1938 r.kernel_name().to_string(),
1939 r.forward_source().to_string(),
1940 r.backward_source().to_string(),
1941 Arc::new(r),
1942 )
1943 }
1944 DtypeRepr::I8 => {
1945 let r = ChannelCatRuntimeOp::<i8>::new(128, n_inputs);
1946 (
1947 r.kernel_name().to_string(),
1948 r.forward_source().to_string(),
1949 r.backward_source().to_string(),
1950 Arc::new(r),
1951 )
1952 }
1953 DtypeRepr::I16 => {
1954 let r = ChannelCatRuntimeOp::<i16>::new(128, n_inputs);
1955 (
1956 r.kernel_name().to_string(),
1957 r.forward_source().to_string(),
1958 r.backward_source().to_string(),
1959 Arc::new(r),
1960 )
1961 }
1962 DtypeRepr::I32 => {
1963 let r = ChannelCatRuntimeOp::<i32>::new(128, n_inputs);
1964 (
1965 r.kernel_name().to_string(),
1966 r.forward_source().to_string(),
1967 r.backward_source().to_string(),
1968 Arc::new(r),
1969 )
1970 }
1971 DtypeRepr::I64 => {
1972 let r = ChannelCatRuntimeOp::<i64>::new(128, n_inputs);
1973 (
1974 r.kernel_name().to_string(),
1975 r.forward_source().to_string(),
1976 r.backward_source().to_string(),
1977 Arc::new(r),
1978 )
1979 }
1980 DtypeRepr::U8 => {
1981 let r = ChannelCatRuntimeOp::<u8>::new(128, n_inputs);
1982 (
1983 r.kernel_name().to_string(),
1984 r.forward_source().to_string(),
1985 r.backward_source().to_string(),
1986 Arc::new(r),
1987 )
1988 }
1989 DtypeRepr::U16 => {
1990 let r = ChannelCatRuntimeOp::<u16>::new(128, n_inputs);
1991 (
1992 r.kernel_name().to_string(),
1993 r.forward_source().to_string(),
1994 r.backward_source().to_string(),
1995 Arc::new(r),
1996 )
1997 }
1998 DtypeRepr::U32 => {
1999 let r = ChannelCatRuntimeOp::<u32>::new(128, n_inputs);
2000 (
2001 r.kernel_name().to_string(),
2002 r.forward_source().to_string(),
2003 r.backward_source().to_string(),
2004 Arc::new(r),
2005 )
2006 }
2007 DtypeRepr::U64 => {
2008 let r = ChannelCatRuntimeOp::<u64>::new(128, n_inputs);
2009 (
2010 r.kernel_name().to_string(),
2011 r.forward_source().to_string(),
2012 r.backward_source().to_string(),
2013 Arc::new(r),
2014 )
2015 }
2016 other => {
2017 return Err(anyhow::anyhow!(
2018 "{:?} is not supported for ChannelCat",
2019 other
2020 ));
2021 }
2022 };
2023 Box::new(KernelExecutable {
2024 entry_point: format!("{}_entry_point", name),
2025 name,
2026 kernel_source: fwd_src,
2027 shape: node.shape.clone(),
2028 dtype: node.dtype,
2029 #[cfg(feature = "training")]
2030 backward_kernel_source: bwd_src,
2031 #[cfg(feature = "training")]
2032 backward_entry_point: String::new(),
2033 runtime_op: rop,
2034 })
2035 }
2036
2037 Op::ChannelChunk {
2038 chunk_c,
2039 chunk_offset,
2040 ..
2041 } => {
2042 let chunk_c = *chunk_c;
2043 let chunk_offset = *chunk_offset;
2044 let (name, fwd_src, bwd_src, rop): (
2045 String,
2046 String,
2047 String,
2048 Arc<dyn RuntimeOp>,
2049 ) = match node.dtype {
2050 DtypeRepr::F32 => {
2051 let r = ChannelChunkRuntimeOp::<f32>::new(128, chunk_c, chunk_offset);
2052 (
2053 r.kernel_name().to_string(),
2054 r.forward_source().to_string(),
2055 r.backward_source().to_string(),
2056 Arc::new(r),
2057 )
2058 }
2059 DtypeRepr::F64 => {
2060 let r = ChannelChunkRuntimeOp::<f64>::new(128, chunk_c, chunk_offset);
2061 (
2062 r.kernel_name().to_string(),
2063 r.forward_source().to_string(),
2064 r.backward_source().to_string(),
2065 Arc::new(r),
2066 )
2067 }
2068 DtypeRepr::I8 => {
2069 let r = ChannelChunkRuntimeOp::<i8>::new(128, chunk_c, chunk_offset);
2070 (
2071 r.kernel_name().to_string(),
2072 r.forward_source().to_string(),
2073 r.backward_source().to_string(),
2074 Arc::new(r),
2075 )
2076 }
2077 DtypeRepr::I16 => {
2078 let r = ChannelChunkRuntimeOp::<i16>::new(128, chunk_c, chunk_offset);
2079 (
2080 r.kernel_name().to_string(),
2081 r.forward_source().to_string(),
2082 r.backward_source().to_string(),
2083 Arc::new(r),
2084 )
2085 }
2086 DtypeRepr::I32 => {
2087 let r = ChannelChunkRuntimeOp::<i32>::new(128, chunk_c, chunk_offset);
2088 (
2089 r.kernel_name().to_string(),
2090 r.forward_source().to_string(),
2091 r.backward_source().to_string(),
2092 Arc::new(r),
2093 )
2094 }
2095 DtypeRepr::I64 => {
2096 let r = ChannelChunkRuntimeOp::<i64>::new(128, chunk_c, chunk_offset);
2097 (
2098 r.kernel_name().to_string(),
2099 r.forward_source().to_string(),
2100 r.backward_source().to_string(),
2101 Arc::new(r),
2102 )
2103 }
2104 DtypeRepr::U8 => {
2105 let r = ChannelChunkRuntimeOp::<u8>::new(128, chunk_c, chunk_offset);
2106 (
2107 r.kernel_name().to_string(),
2108 r.forward_source().to_string(),
2109 r.backward_source().to_string(),
2110 Arc::new(r),
2111 )
2112 }
2113 DtypeRepr::U16 => {
2114 let r = ChannelChunkRuntimeOp::<u16>::new(128, chunk_c, chunk_offset);
2115 (
2116 r.kernel_name().to_string(),
2117 r.forward_source().to_string(),
2118 r.backward_source().to_string(),
2119 Arc::new(r),
2120 )
2121 }
2122 DtypeRepr::U32 => {
2123 let r = ChannelChunkRuntimeOp::<u32>::new(128, chunk_c, chunk_offset);
2124 (
2125 r.kernel_name().to_string(),
2126 r.forward_source().to_string(),
2127 r.backward_source().to_string(),
2128 Arc::new(r),
2129 )
2130 }
2131 DtypeRepr::U64 => {
2132 let r = ChannelChunkRuntimeOp::<u64>::new(128, chunk_c, chunk_offset);
2133 (
2134 r.kernel_name().to_string(),
2135 r.forward_source().to_string(),
2136 r.backward_source().to_string(),
2137 Arc::new(r),
2138 )
2139 }
2140 other => {
2141 return Err(anyhow::anyhow!(
2142 "{:?} is not supported for ChannelChunk",
2143 other
2144 ));
2145 }
2146 };
2147 Box::new(KernelExecutable {
2148 entry_point: format!("{}_entry_point", name),
2149 name,
2150 kernel_source: fwd_src,
2151 shape: node.shape.clone(),
2152 dtype: node.dtype,
2153 #[cfg(feature = "training")]
2154 backward_kernel_source: bwd_src,
2155 #[cfg(feature = "training")]
2156 backward_entry_point: String::new(),
2157 runtime_op: rop,
2158 })
2159 }
2160
2161 Op::ChannelBiasAdd { c } => {
2162 let c = *c;
2163 let (name, fwd_src, bwd_src, rop): (
2164 String,
2165 String,
2166 String,
2167 Arc<dyn RuntimeOp>,
2168 ) = match node.dtype {
2169 DtypeRepr::F32 => {
2170 let r = ChannelBiasAddRuntimeOp::<f32>::new(128, c);
2171 (
2172 r.kernel_name().to_string(),
2173 r.forward_source().to_string(),
2174 r.backward_source().to_string(),
2175 Arc::new(r),
2176 )
2177 }
2178 DtypeRepr::F64 => {
2179 let r = ChannelBiasAddRuntimeOp::<f64>::new(128, c);
2180 (
2181 r.kernel_name().to_string(),
2182 r.forward_source().to_string(),
2183 r.backward_source().to_string(),
2184 Arc::new(r),
2185 )
2186 }
2187 other => {
2188 return Err(anyhow::anyhow!(
2189 "{:?} is not supported for ChannelBiasAdd",
2190 other
2191 ));
2192 }
2193 };
2194 Box::new(KernelExecutable {
2195 entry_point: format!("{}_entry_point", name),
2196 name,
2197 kernel_source: fwd_src,
2198 shape: node.shape.clone(),
2199 dtype: node.dtype,
2200 #[cfg(feature = "training")]
2201 backward_kernel_source: bwd_src,
2202 #[cfg(feature = "training")]
2203 backward_entry_point: String::new(),
2204 runtime_op: rop,
2205 })
2206 }
2207
2208 Op::Add => {
2209 make_num_kernel!(ElemwiseAddForward(128), ElemwiseAddBackward(128), node)
2210 }
2211
2212 Op::Abs => {
2214 make_num_kernel!(ElemwiseAbsForward(1024), ElemwiseAbsBackward(1024), node)
2215 }
2216 Op::Neg => {
2217 make_num_kernel!(ElemwiseNegForward(1024), ElemwiseNegBackward(1024), node)
2218 }
2219 Op::Sign => make_num_kernel!(ElemwiseSignForward(1024), node),
2220 Op::IsNaN => make_float_kernel!(ElemwiseIsnanForward(1024), node),
2221 Op::Ceil => make_float_kernel!(ElemwiseCeilForward(1024), node),
2222 Op::Floor => make_float_kernel!(ElemwiseFloorForward(1024), node),
2223 Op::Sqrt => {
2224 make_float_kernel!(ElemwiseSqrtForward(1024), ElemwiseSqrtBackward(1024), node)
2225 }
2226 Op::Reciprocal => make_float_kernel!(
2227 ElemwiseReciprocalForward(1024),
2228 ElemwiseReciprocalBackward(1024),
2229 node
2230 ),
2231 Op::Exp => {
2232 make_float_kernel!(ElemwiseExpForward(1024), ElemwiseExpBackward(1024), node)
2233 }
2234 Op::Log => {
2235 make_float_kernel!(ElemwiseLogForward(1024), ElemwiseLogBackward(1024), node)
2236 }
2237 Op::Erf => {
2238 make_float_kernel!(ElemwiseErfForward(1024), ElemwiseErfBackward(1024), node)
2239 }
2240 Op::Sin => {
2241 make_float_kernel!(ElemwiseSinForward(1024), ElemwiseSinBackward(1024), node)
2242 }
2243 Op::Cos => {
2244 make_float_kernel!(ElemwiseCosForward(1024), ElemwiseCosBackward(1024), node)
2245 }
2246 Op::Tan => {
2247 make_float_kernel!(ElemwiseTanForward(1024), ElemwiseTanBackward(1024), node)
2248 }
2249 Op::Asin => {
2250 make_float_kernel!(ElemwiseAsinForward(1024), ElemwiseAsinBackward(1024), node)
2251 }
2252 Op::Acos => {
2253 make_float_kernel!(ElemwiseAcosForward(1024), ElemwiseAcosBackward(1024), node)
2254 }
2255 Op::Atan => {
2256 make_float_kernel!(ElemwiseAtanForward(1024), ElemwiseAtanBackward(1024), node)
2257 }
2258 Op::Sinh => {
2259 make_float_kernel!(ElemwiseSinhForward(1024), ElemwiseSinhBackward(1024), node)
2260 }
2261 Op::Cosh => {
2262 make_float_kernel!(ElemwiseCoshForward(1024), ElemwiseCoshBackward(1024), node)
2263 }
2264 Op::Asinh => make_float_kernel!(
2265 ElemwiseAsinhForward(1024),
2266 ElemwiseAsinhBackward(1024),
2267 node
2268 ),
2269 Op::Acosh => make_float_kernel!(
2270 ElemwiseAcoshForward(1024),
2271 ElemwiseAcoshBackward(1024),
2272 node
2273 ),
2274 Op::Atanh => make_float_kernel!(
2275 ElemwiseAtanhForward(1024),
2276 ElemwiseAtanhBackward(1024),
2277 node
2278 ),
2279 Op::Round => {
2280 return Err(anyhow::anyhow!(
2281 "TODO: Op::Round — implement rounding kernel"
2282 ));
2283 }
2284
2285 Op::Mul => {
2287 make_num_kernel!(ElemwiseMulForward(1024), ElemwiseMulBackward(1024), node)
2288 }
2289 Op::Sub => {
2290 make_num_kernel!(ElemwiseSubForward(1024), ElemwiseSubBackward(1024), node)
2291 }
2292 Op::Div => {
2293 make_float_kernel!(ElemwiseDivForward(1024), ElemwiseDivBackward(1024), node)
2294 }
2295 Op::Pow => {
2296 make_float_kernel!(ElemwisePowForward(1024), ElemwisePowBackward(1024), node)
2297 }
2298 Op::Mod { .. } => make_float_kernel!(ElemwiseFmodForward(1024), node),
2299 Op::ElemMin => {
2300 make_num_kernel!(ElemwiseMinForward(1024), ElemwiseMinBackward(1024), node)
2301 }
2302 Op::ElemMax => {
2303 make_num_kernel!(ElemwiseMaxForward(1024), ElemwiseMaxBackward(1024), node)
2304 }
2305 Op::ElemMean => {
2306 make_float_kernel!(ElemwiseMeanForward(1024), ElemwiseMeanBackward(1024), node)
2307 }
2308 Op::ElemSum => {
2309 make_num_kernel!(ElemwiseSumForward(1024), ElemwiseSumBackward(1024), node)
2310 }
2311 Op::Equal => make_num_kernel!(ElemwiseEqualForward(1024), node),
2312 Op::Greater => make_num_kernel!(ElemwiseGreaterForward(1024), node),
2313 Op::GreaterOrEqual => make_num_kernel!(ElemwiseGreaterEqualForward(1024), node),
2314 Op::Less => make_num_kernel!(ElemwiseLessForward(1024), node),
2315 Op::LessOrEqual => make_num_kernel!(ElemwiseLessEqualForward(1024), node),
2316 Op::Where => make_float_kernel!(
2317 ElemwiseWhereForward(1024),
2318 ElemwiseWhereBackward(1024),
2319 node
2320 ),
2321 Op::Clip => {
2322 let (name, ks, bwd_ks, rop): (String, String, String, Arc<dyn RuntimeOp>) =
2323 match node.dtype {
2324 DtypeRepr::F32 => {
2325 let r = ClipRuntimeOp::<f32>::new(
2326 1024,
2327 f32::NEG_INFINITY,
2328 f32::INFINITY,
2329 );
2330 (
2331 r.kernel_name().to_string(),
2332 r.forward_source().to_string(),
2333 r.backward_source().to_string(),
2334 Arc::new(r),
2335 )
2336 }
2337 DtypeRepr::F64 => {
2338 let r = ClipRuntimeOp::<f64>::new(
2339 1024,
2340 f32::NEG_INFINITY,
2341 f32::INFINITY,
2342 );
2343 (
2344 r.kernel_name().to_string(),
2345 r.forward_source().to_string(),
2346 r.backward_source().to_string(),
2347 Arc::new(r),
2348 )
2349 }
2350 other => {
2351 return Err(anyhow::anyhow!(
2352 "{:?} is not supported for Clip",
2353 other
2354 ));
2355 }
2356 };
2357 Box::new(KernelExecutable {
2358 entry_point: format!("{}_entry_point", name),
2359 name,
2360 kernel_source: ks,
2361 shape: node.shape.clone(),
2362 dtype: node.dtype,
2363 #[cfg(feature = "training")]
2364 backward_kernel_source: bwd_ks,
2365 #[cfg(feature = "training")]
2366 backward_entry_point: String::new(),
2367 runtime_op: rop,
2368 })
2369 }
2370
2371 Op::ReduceSum { .. } => make_num_kernel!(ReduceSumForward(1024), node),
2373 Op::ReduceMean { .. } => make_float_kernel!(ReduceMeanForward(1024), node),
2374 Op::ReduceMax { .. } => make_num_kernel!(ReduceMaxForward(1024), node),
2375 Op::ReduceMin { .. } => make_num_kernel!(ReduceMinForward(1024), node),
2376 Op::ReduceProd { .. } => make_float_kernel!(ReduceProdForward(1024), node),
2377 Op::ReduceL1 { .. } => make_num_kernel!(ReduceL1Forward(1024), node),
2378 Op::ReduceL2 { .. } => make_float_kernel!(ReduceL2Forward(1024), node),
2379 Op::ReduceLogSum { .. } => make_float_kernel!(ReduceLogSumForward(1024), node),
2380 Op::ReduceLogSumExp { .. } => {
2381 make_float_kernel!(ReduceLogSumExpForward(1024), node)
2382 }
2383 Op::ReduceSumSquare { .. } => make_num_kernel!(ReduceSumSquareForward(1024), node),
2384 Op::CumSum { .. } => make_num_kernel!(CumSumForward(1024), node),
2385 Op::CumProd { .. } => make_num_kernel!(CumProdForward(1024), node),
2386 Op::GlobalAvgPool => make_float_kernel!(GlobalAvgPoolForward(1024), node),
2387 Op::GlobalMaxPool => make_float_kernel!(GlobalMaxPoolForward(1024), node),
2388 Op::ArgMax { .. } => {
2389 return Err(anyhow::anyhow!(
2390 "TODO: Op::ArgMax — I32Tensor output requires a custom kernel"
2391 ));
2392 }
2393 Op::ArgMin { .. } => {
2394 return Err(anyhow::anyhow!(
2395 "TODO: Op::ArgMin — I32Tensor output requires a custom kernel"
2396 ));
2397 }
2398
2399 Op::Swish => {
2401 let fwd = SwishForward::new(1024);
2402 let nm = fwd.name.to_string();
2403 let fwd_src = fwd.source.clone();
2404 #[cfg(feature = "training")]
2405 let bwd_src = SwishBackward::new(1024).source.clone();
2406 let rop: Arc<dyn RuntimeOp> = Arc::new(fwd);
2407 Box::new(KernelExecutable {
2408 entry_point: format!("{}_entry_point", nm),
2409 name: nm,
2410 kernel_source: fwd_src,
2411 shape: node.shape.clone(),
2412 dtype: node.dtype,
2413 #[cfg(feature = "training")]
2414 backward_kernel_source: bwd_src,
2415 #[cfg(feature = "training")]
2416 backward_entry_point: String::new(),
2417 runtime_op: rop,
2418 })
2419 }
2420 Op::PRelu => {
2421 let fwd = PreluForward::new(1024);
2422 let nm = fwd.name.to_string();
2423 let fwd_src = fwd.source.clone();
2424 #[cfg(feature = "training")]
2425 let bwd_src = crate::nn::activation::extra::PreluBackward::new(1024)
2426 .source
2427 .clone();
2428 let rop: Arc<dyn RuntimeOp> = Arc::new(fwd);
2429 Box::new(KernelExecutable {
2430 entry_point: format!("{}_entry_point", nm),
2431 name: nm,
2432 kernel_source: fwd_src,
2433 shape: node.shape.clone(),
2434 dtype: node.dtype,
2435 #[cfg(feature = "training")]
2436 backward_kernel_source: bwd_src,
2437 #[cfg(feature = "training")]
2438 backward_entry_point: String::new(),
2439 runtime_op: rop,
2440 })
2441 }
2442 Op::LogSoftmax { .. } => {
2443 let n_cols = node.shape.last().and_then(|d| *d).unwrap_or(1024);
2444 let block_size = n_cols.next_power_of_two() as i32;
2445 let fwd = LogSoftmaxForward::new(block_size);
2446 let nm = fwd.name.to_string();
2447 let fwd_src = fwd.source.clone();
2448 #[cfg(feature = "training")]
2449 let bwd_src = LogSoftmaxBackward::new(block_size).source.clone();
2450 let rop: Arc<dyn RuntimeOp> = Arc::new(fwd);
2451 Box::new(KernelExecutable {
2452 entry_point: format!("{}_entry_point", nm),
2453 name: nm,
2454 kernel_source: fwd_src,
2455 shape: node.shape.clone(),
2456 dtype: node.dtype,
2457 #[cfg(feature = "training")]
2458 backward_kernel_source: bwd_src,
2459 #[cfg(feature = "training")]
2460 backward_entry_point: String::new(),
2461 runtime_op: rop,
2462 })
2463 }
2464 Op::ThresholdedRelu { alpha } => {
2465 let alpha_f = *alpha as f32;
2466 let (name, ks, bwd_ks, rop): (String, String, String, Arc<dyn RuntimeOp>) =
2467 match node.dtype {
2468 DtypeRepr::F32 => {
2469 let r = ThresholdedReluRuntimeOp::new(1024, alpha_f);
2470 (
2471 r.kernel_name().to_string(),
2472 r.forward_source().to_string(),
2473 r.backward_source().to_string(),
2474 Arc::new(r),
2475 )
2476 }
2477 other => {
2478 return Err(anyhow::anyhow!(
2479 "{:?} is not supported for ThresholdedRelu",
2480 other
2481 ));
2482 }
2483 };
2484 Box::new(KernelExecutable {
2485 entry_point: format!("{}_entry_point", name),
2486 name,
2487 kernel_source: ks,
2488 shape: node.shape.clone(),
2489 dtype: node.dtype,
2490 #[cfg(feature = "training")]
2491 backward_kernel_source: bwd_ks,
2492 #[cfg(feature = "training")]
2493 backward_entry_point: String::new(),
2494 runtime_op: rop,
2495 })
2496 }
2497 Op::Shrink { lambd, bias } => {
2498 let lambd_f = *lambd as f32;
2499 let bias_f = *bias as f32;
2500 let (name, ks, bwd_ks, rop): (String, String, String, Arc<dyn RuntimeOp>) =
2501 match node.dtype {
2502 DtypeRepr::F32 => {
2503 let r = ShrinkRuntimeOp::new(1024, lambd_f, bias_f);
2504 (
2505 r.kernel_name().to_string(),
2506 r.forward_source().to_string(),
2507 r.backward_source().to_string(),
2508 Arc::new(r),
2509 )
2510 }
2511 other => {
2512 return Err(anyhow::anyhow!(
2513 "{:?} is not supported for Shrink",
2514 other
2515 ));
2516 }
2517 };
2518 Box::new(KernelExecutable {
2519 entry_point: format!("{}_entry_point", name),
2520 name,
2521 kernel_source: ks,
2522 shape: node.shape.clone(),
2523 dtype: node.dtype,
2524 #[cfg(feature = "training")]
2525 backward_kernel_source: bwd_ks,
2526 #[cfg(feature = "training")]
2527 backward_entry_point: String::new(),
2528 runtime_op: rop,
2529 })
2530 }
2531
2532 Op::MatMul | Op::Gemm { .. } => {
2534 let m = node.shape.first().copied().flatten();
2536 let n = node.shape.last().copied().flatten().unwrap_or(0);
2537 let k = node
2538 .inputs
2539 .first()
2540 .and_then(|&i| graph.nodes[i].shape.last().copied().flatten())
2541 .unwrap_or(0);
2542 let (block_m, block_n, block_k) = pick_gemm_tile_sizes(m, n, k);
2543
2544 let (name, ks, rop): (String, String, Arc<dyn RuntimeOp>) = match node.dtype {
2545 DtypeRepr::F32 => {
2546 let r = MatMulRuntimeOp::<f32>::new(block_m, block_n, block_k);
2547 (
2548 r.kernel_name().to_string(),
2549 r.forward_source().to_string(),
2550 Arc::new(r),
2551 )
2552 }
2553 DtypeRepr::F64 => {
2554 let r = MatMulRuntimeOp::<f64>::new(block_m, block_n, block_k);
2555 (
2556 r.kernel_name().to_string(),
2557 r.forward_source().to_string(),
2558 Arc::new(r),
2559 )
2560 }
2561 other => {
2562 return Err(anyhow::anyhow!(
2563 "{:?} is not a Float dtype for MatMul",
2564 other
2565 ));
2566 }
2567 };
2568 Box::new(KernelExecutable {
2569 entry_point: format!("{}_entry_point", name),
2570 name,
2571 kernel_source: ks,
2572 shape: node.shape.clone(),
2573 dtype: node.dtype,
2574 #[cfg(feature = "training")]
2575 backward_kernel_source: String::new(),
2576 #[cfg(feature = "training")]
2577 backward_entry_point: String::new(),
2578 runtime_op: rop,
2579 })
2580 }
2581
2582 Op::Lstm { .. } => {
2584 return Err(anyhow::anyhow!(
2585 "TODO: Op::Lstm — multi-step recurrent; implement as a custom loop kernel"
2586 ));
2587 }
2588 Op::Gru { .. } => {
2589 return Err(anyhow::anyhow!(
2590 "TODO: Op::Gru — multi-step recurrent; implement as a custom loop kernel"
2591 ));
2592 }
2593 Op::Rnn { .. } => {
2594 return Err(anyhow::anyhow!(
2595 "TODO: Op::Rnn — multi-step recurrent; implement as a custom loop kernel"
2596 ));
2597 }
2598 Op::RotaryEmbedding => {
2599 return Err(anyhow::anyhow!(
2600 "TODO: Op::RotaryEmbedding — implement RoPE kernel"
2601 ));
2602 }
2603 Op::MultiHeadAttention { .. } => {
2604 return Err(anyhow::anyhow!(
2605 "TODO: Op::MultiHeadAttention — use flash attention or a custom MHA kernel"
2606 ));
2607 }
2608 Op::FlexAttention { .. } => {
2609 return Err(anyhow::anyhow!(
2610 "TODO: Op::FlexAttention — implement a custom attention kernel supporting an arbitrary score-modification function"
2611 ));
2612 }
2613 Op::LinearAttention { .. } => {
2614 return Err(anyhow::anyhow!(
2615 "TODO: Op::LinearAttention — implement a linear-attention/gated-delta-rule kernel"
2616 ));
2617 }
2618 Op::CausalConvWithState { .. } => {
2619 return Err(anyhow::anyhow!(
2620 "TODO: Op::CausalConvWithState — implement a stateful causal-conv kernel"
2621 ));
2622 }
2623 Op::Reshape => {
2624 return Err(anyhow::anyhow!(
2625 "TODO: Op::Reshape — implement as a strided view or copy kernel"
2626 ));
2627 }
2628 Op::Transpose { .. } => {
2629 return Err(anyhow::anyhow!(
2630 "TODO: Op::Transpose — implement as a permuted-copy kernel"
2631 ));
2632 }
2633 Op::Squeeze { .. } | Op::Unsqueeze { .. } => {
2634 return Err(anyhow::anyhow!(
2635 "TODO: Op::Squeeze/Unsqueeze — implement as a zero-copy view"
2636 ));
2637 }
2638 Op::Concat { .. } => {
2639 return Err(anyhow::anyhow!(
2640 "TODO: Op::Concat — implement as a multi-input copy kernel"
2641 ));
2642 }
2643 Op::Split { .. } => {
2644 return Err(anyhow::anyhow!(
2645 "TODO: Op::Split — implement as a multi-output slice kernel"
2646 ));
2647 }
2648 Op::Slice => {
2649 return Err(anyhow::anyhow!(
2650 "TODO: Op::Slice — implement as a strided-copy kernel"
2651 ));
2652 }
2653 Op::Gather { .. } | Op::GatherElements { .. } | Op::GatherND { .. } => {
2654 return Err(anyhow::anyhow!(
2655 "TODO: Op::Gather — implement as an index-gather kernel"
2656 ));
2657 }
2658 Op::ScatterElements { .. }
2659 | Op::ScatterND
2660 | Op::Scatter { .. }
2661 | Op::TensorScatter => {
2662 return Err(anyhow::anyhow!(
2663 "TODO: Op::Scatter — implement as an index-scatter kernel"
2664 ));
2665 }
2666 Op::Tile => {
2667 return Err(anyhow::anyhow!(
2668 "TODO: Op::Tile — implement as a tiled-copy kernel"
2669 ));
2670 }
2671 Op::Expand => {
2672 return Err(anyhow::anyhow!(
2673 "TODO: Op::Expand — implement as a broadcast-copy kernel"
2674 ));
2675 }
2676 Op::ShapeOf { .. } | Op::SizeOf => {
2677 return Err(anyhow::anyhow!(
2678 "TODO: Op::ShapeOf/SizeOf — output is metadata, not tensor data"
2679 ));
2680 }
2681 Op::Compress { .. } | Op::NonZero => {
2682 return Err(anyhow::anyhow!(
2683 "TODO: Op::Compress/NonZero — variable-output ops require stream compaction"
2684 ));
2685 }
2686 Op::Range => {
2687 return Err(anyhow::anyhow!(
2688 "TODO: Op::Range — implement as a fill/arange kernel"
2689 ));
2690 }
2691 Op::Constant { .. } | Op::ConstantOfShape { .. } => {
2692 return Err(anyhow::anyhow!(
2693 "TODO: Op::Constant — inline constant; should be materialised before lowering"
2694 ));
2695 }
2696 Op::Trilu { .. } => {
2697 return Err(anyhow::anyhow!(
2698 "TODO: Op::Trilu — implement as a triangular mask kernel"
2699 ));
2700 }
2701 Op::Pad { .. } => {
2702 return Err(anyhow::anyhow!(
2703 "TODO: Op::Pad — implement as a generic N-D padding kernel"
2704 ));
2705 }
2706 Op::ReverseSequence { .. } => {
2707 return Err(anyhow::anyhow!(
2708 "TODO: Op::ReverseSequence — implement as a scatter-copy kernel"
2709 ));
2710 }
2711 Op::Einsum { .. } => {
2712 return Err(anyhow::anyhow!(
2713 "TODO: Op::Einsum — parse equation and emit a fused contraction kernel"
2714 ));
2715 }
2716 Op::Det => {
2717 return Err(anyhow::anyhow!(
2718 "TODO: Op::Det — implement via LU decomposition"
2719 ));
2720 }
2721 Op::QLinearMatMul
2722 | Op::MatMulInteger
2723 | Op::ConvInteger { .. }
2724 | Op::QLinearConv { .. } => {
2725 return Err(anyhow::anyhow!(
2726 "TODO: Op::Q* quantised matmul/conv — implement quantised compute kernels"
2727 ));
2728 }
2729 Op::ConvTranspose { .. } => {
2730 return Err(anyhow::anyhow!(
2731 "TODO: Op::ConvTranspose — implement transposed (gradient) convolution kernel"
2732 ));
2733 }
2734 Op::DeformConv { .. } => {
2735 return Err(anyhow::anyhow!(
2736 "TODO: Op::DeformConv — implement deformable convolution kernel"
2737 ));
2738 }
2739 Op::Col2Im { .. } => {
2740 return Err(anyhow::anyhow!(
2741 "TODO: Op::Col2Im — implement col2im (fold) kernel"
2742 ));
2743 }
2744 Op::Resize { .. } => {
2745 return Err(anyhow::anyhow!(
2746 "TODO: Op::Resize — implement nearest/bilinear/bicubic resize kernels"
2747 ));
2748 }
2749 Op::GridSample { .. } => {
2750 return Err(anyhow::anyhow!(
2751 "TODO: Op::GridSample — implement bilinear grid sample kernel"
2752 ));
2753 }
2754 Op::SpaceToDepth { .. } | Op::DepthToSpace { .. } => {
2755 return Err(anyhow::anyhow!(
2756 "TODO: Op::SpaceToDepth/DepthToSpace — implement pixel shuffle kernel"
2757 ));
2758 }
2759 Op::RoiAlign { .. } => {
2760 return Err(anyhow::anyhow!(
2761 "TODO: Op::RoiAlign — implement RoI-align pooling kernel"
2762 ));
2763 }
2764 Op::AffineGrid { .. } => {
2765 return Err(anyhow::anyhow!(
2766 "TODO: Op::AffineGrid — implement affine grid generator kernel"
2767 ));
2768 }
2769 Op::MaxUnpool { .. } => {
2770 return Err(anyhow::anyhow!(
2771 "TODO: Op::MaxUnpool — implement max-unpool (scatter with saved indices) kernel"
2772 ));
2773 }
2774 Op::CenterCropPad { .. } => {
2775 return Err(anyhow::anyhow!(
2776 "TODO: Op::CenterCropPad — implement center-crop-pad kernel"
2777 ));
2778 }
2779 Op::NonMaxSuppression { .. } => {
2780 return Err(anyhow::anyhow!(
2781 "TODO: Op::NonMaxSuppression — implement NMS kernel"
2782 ));
2783 }
2784 Op::TopK { .. } => {
2785 return Err(anyhow::anyhow!(
2786 "TODO: Op::TopK — implement radix sort / parallel selection kernel"
2787 ));
2788 }
2789 Op::Unique { .. } => {
2790 return Err(anyhow::anyhow!(
2791 "TODO: Op::Unique — implement stream-compaction unique kernel"
2792 ));
2793 }
2794 Op::EyeLike { .. }
2795 | Op::OneHot { .. }
2796 | Op::Bernoulli { .. }
2797 | Op::RandomUniformLike { .. } => {
2798 return Err(anyhow::anyhow!(
2799 "TODO: Op::EyeLike/OneHot/Bernoulli/RandomUniformLike — implement generation kernels"
2800 ));
2801 }
2802 Op::And | Op::Or | Op::Xor => {
2803 return Err(anyhow::anyhow!(
2804 "TODO: Op::And/Or/Xor — implement boolean logical kernels"
2805 ));
2806 }
2807 Op::BitShift { .. }
2808 | Op::BitwiseAnd
2809 | Op::BitwiseOr
2810 | Op::BitwiseXor
2811 | Op::BitwiseNot
2812 | Op::Not => {
2813 return Err(anyhow::anyhow!(
2814 "TODO: Op::Bitwise* — implement integer bitwise kernels"
2815 ));
2816 }
2817 Op::QuantizeLinear { .. }
2818 | Op::DequantizeLinear { .. }
2819 | Op::DynamicQuantizeLinear => {
2820 return Err(anyhow::anyhow!(
2821 "TODO: Op::Quantize/Dequantize — implement quantisation kernels"
2822 ));
2823 }
2824 Op::LRN { .. } => {
2825 return Err(anyhow::anyhow!(
2826 "TODO: Op::LRN — implement local response normalisation kernel"
2827 ));
2828 }
2829 Op::MeanVarianceNormalization { .. } | Op::LpNormalization { .. } => {
2830 return Err(anyhow::anyhow!(
2831 "TODO: Op::MvnNorm/LpNorm — implement normalisation kernels"
2832 ));
2833 }
2834 Op::Dft { .. }
2835 | Op::Stft
2836 | Op::MelWeightMatrix
2837 | Op::HannWindow { .. }
2838 | Op::BlackmanWindow { .. }
2839 | Op::HammingWindow { .. } => {
2840 return Err(anyhow::anyhow!(
2841 "TODO: Op::DFT/STFT/Window — implement signal processing kernels"
2842 ));
2843 }
2844 Op::NegativeLogLikelihoodLoss { .. } | Op::SoftmaxCrossEntropyLoss { .. } => {
2845 return Err(anyhow::anyhow!(
2846 "TODO: Op::NllLoss/SoftmaxCELoss — implement loss kernels"
2847 ));
2848 }
2849 Op::SequenceAt
2850 | Op::SequenceConstruct
2851 | Op::SequenceEmpty
2852 | Op::SequenceErase
2853 | Op::SequenceInsert
2854 | Op::SequenceLength
2855 | Op::SequenceMap
2856 | Op::SplitToSequence { .. }
2857 | Op::ConcatFromSequence { .. }
2858 | Op::OptionalGetElement
2859 | Op::OptionalHasElement
2860 | Op::Loop
2861 | Op::Scan { .. }
2862 | Op::If => {
2863 return Err(anyhow::anyhow!(
2864 "TODO: Op::Sequence/Control-flow — not lowerable to single Triton kernels"
2865 ));
2866 }
2867 Op::Adagrad | Op::Adam | Op::Momentum | Op::Gradient => {
2868 return Err(anyhow::anyhow!(
2869 "TODO: Op::OnnxOptimizer — use teenygrad's own optimizer kernels instead"
2870 ));
2871 }
2872 Op::StringNormalizer
2873 | Op::RegexFullMatch { .. }
2874 | Op::StringConcat
2875 | Op::StringSplit
2876 | Op::TfIdfVectorizer
2877 | Op::LabelEncoder
2878 | Op::ArrayFeatureExtractor
2879 | Op::Binarizer { .. }
2880 | Op::TreeEnsemble
2881 | Op::ImageDecoder => {
2882 return Err(anyhow::anyhow!(
2883 "TODO: Op::String/ClassicalML — not GPU-lowerable"
2884 ));
2885 }
2886 Op::CastLike | Op::BitCast { .. } => {
2887 return Err(anyhow::anyhow!(
2888 "TODO: Op::CastLike/BitCast — implement dtype-cast kernels"
2889 ));
2890 }
2891 Op::Cast { to } => {
2892 return Err(anyhow::anyhow!(
2893 "TODO: Op::Cast to {:?} — implement dtype-cast kernel",
2894 to
2895 ));
2896 }
2897 Op::Identity => {
2898 return Err(anyhow::anyhow!(
2899 "TODO: Op::Identity — implement zero-copy pass-through kernel"
2900 ));
2901 }
2902 Op::Dropout { .. } => {
2903 return Err(anyhow::anyhow!(
2904 "TODO: Op::Dropout — implement inference pass-through / training dropout kernel"
2905 ));
2906 }
2907 Op::IsInf { .. } => {
2908 return Err(anyhow::anyhow!("TODO: Op::IsInf — implement isinf kernel"));
2909 }
2910 Op::Hardmax { .. } => {
2911 return Err(anyhow::anyhow!(
2912 "TODO: Op::Hardmax — implement argmax + one-hot kernel"
2913 ));
2914 }
2915
2916 Op::Attention {
2917 c,
2918 num_heads,
2919 key_dim,
2920 } => {
2921 let _ = (c, num_heads, key_dim);
2925 return Err(anyhow::anyhow!(
2926 "Op::Attention reached the match arm — this should not happen"
2927 ));
2928 }
2929
2930 Op::Custom { data } => match data.0.lower() {
2931 Some((name, kernel_source, entry_point, runtime_op)) => {
2932 Box::new(KernelExecutable {
2933 name,
2934 kernel_source,
2935 entry_point,
2936 shape: node.shape.clone(),
2937 dtype: node.dtype,
2938 runtime_op,
2939 #[cfg(feature = "training")]
2940 backward_kernel_source: data.0.lower_backward_source(),
2941 #[cfg(feature = "training")]
2942 backward_entry_point: String::new(),
2943 })
2944 }
2945 None => {
2946 return Err(anyhow::anyhow!(
2947 "custom op '{}' is not handled — implement CustomOp::lower()",
2948 data.name()
2949 ));
2950 }
2951 },
2952
2953 Op::Fused { members } => {
2954 return Err(anyhow::anyhow!(
2960 "Op::Fused lowering is not implemented yet ({} member op(s)): \
2961 concatenate each member's kernel source and synthesize an \
2962 entry point that runs them in sequence (see spinorml-1fj.1)",
2963 members.len()
2964 ));
2965 }
2966 };
2967
2968 let dag_idx = dag.add_node(executable);
2969 graph_to_dag[node_index] = dag_idx;
2970
2971 for &input_graph_idx in &node.inputs {
2972 dag.add_edge(graph_to_dag[input_graph_idx], dag_idx);
2973 }
2974 }
2975
2976 Ok((dag, graph_to_dag))
2977 }
2978}
2979
2980impl<'a> Lowering<'a> for TritonLowering {
2981 fn lower(&self, graph: &Graph, mode: LoweringMode) -> Result<Dag<Box<dyn ExecutableOp>>> {
2982 TritonLowering::lower_with_mapping(self, graph, mode).map(|(dag, _)| dag)
2983 }
2984
2985 fn lower_with_mapping(
2986 &self,
2987 graph: &Graph,
2988 mode: LoweringMode,
2989 ) -> Result<(Dag<Box<dyn ExecutableOp>>, Vec<usize>)> {
2990 TritonLowering::lower_with_mapping(self, graph, mode)
2991 }
2992
2993 fn extra_dag_names(&self, graph: &Graph, graph_to_dag: &[usize]) -> Vec<(usize, String)> {
2998 let mut extra = Vec::new();
2999 for (graph_idx, node) in graph.nodes.iter().enumerate() {
3000 if let Op::Conv2d { has_bias: true, .. } = &node.op
3001 && let Some(name) = graph.names.get(&graph_idx)
3002 {
3003 let biasadd_dag_idx = graph_to_dag[graph_idx];
3004 if biasadd_dag_idx > 0 {
3005 extra.push((biasadd_dag_idx - 1, name.clone()));
3006 }
3007 }
3008 }
3009 extra
3010 }
3011}