lift_tensor/
ops.rs

1use serde::{Deserialize, Serialize};
2
3/// FP8 quantisation format variants.
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
5pub enum Fp8Format {
6    E4M3,
7    E5M2,
8}
9
10/// Aggregation type for GNN message passing and pooling.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
12pub enum AggregationType {
13    Sum,
14    Mean,
15    Max,
16    Min,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
20pub enum TensorOp {
21    // ── Arithmetic ──
22    Add,
23    Sub,
24    Mul,
25    Div,
26    Neg,
27    MatMul,
28    Linear,
29    Conv2D,
30    Embedding,
31
32    // ── Activations ──
33    ReLU,
34    GeLU,
35    SiLU,
36    Sigmoid,
37    Softmax,
38    Tanh,
39    LeakyReLU,
40    ELU,
41    Mish,
42    HardSwish,
43    HardSigmoid,
44
45    // ── Normalisation ──
46    LayerNorm,
47    RMSNorm,
48    BatchNorm,
49    GroupNorm,
50    InstanceNorm,
51
52    // ── Shape operations ──
53    Reshape,
54    Transpose,
55    Concat,
56    Split,
57    Gather,
58    Scatter,
59    Squeeze,
60    Unsqueeze,
61    Permute,
62    Expand,
63    Slice,
64    Pad,
65    Tile,
66
67    // ── Constants ──
68    Constant,
69    Zeros,
70    Ones,
71    Arange,
72    Full,
73
74    // ── Attention variants ──
75    Attention,
76    MultiHeadAttention,
77    MultiQueryAttention,
78    GroupedQueryAttention,
79    FlashAttention,
80    SlidingWindowAttention,
81    CrossAttention,
82    PagedAttention,
83
84    // ── MoE (Mixture of Experts) ──
85    MoEDispatch,
86    MoECombine,
87
88    // ── Convolution variants ──
89    Conv1D,
90    Conv3D,
91    ConvTranspose2D,
92    DepthwiseConv2D,
93    DilatedConv2D,
94
95    // ── Pooling ──
96    MaxPool2D,
97    AvgPool2D,
98    AdaptiveAvgPool2D,
99    GlobalAvgPool,
100
101    // ── Recurrent ──
102    LSTMCell,
103    GRUCell,
104    RNNCell,
105
106    // ── Advanced math ──
107    Einsum,
108    FFT,
109    IFFT,
110    SVD,
111    Eig,
112    Solve,
113    TopK,
114    Sort,
115    Cumsum,
116    Where,
117    Clamp,
118
119    // ── Sparse ──
120    SparseMatMul,
121    SparseEmbedding,
122
123    // ── Quantisation ──
124    Quantize,
125    Dequantize,
126    QuantizeInt4,
127    DequantizeInt4,
128    QuantizeFp8,
129    DequantizeFp8,
130
131    // ── Diffusion / Generative ──
132    UNetDownBlock,
133    UNetUpBlock,
134    TimestepEmbedding,
135
136    // ── GNN (Graph Neural Networks) ──
137    GNNMessagePassing,
138    GNNGlobalPooling,
139
140    // ── Memory management ──
141    Checkpoint,
142    Offload,
143    GradAccumulate,
144
145    // ── Gradient operations ──
146    GradMatMul,
147    GradReLU,
148    GradSoftmax,
149    GradLayerNorm,
150    GradAttention,
151    GradConv2D,
152    GradLinear,
153    GradGeLU,
154
155    // ── Parallelism ──
156    ParallelSplit,
157    ParallelAllReduce,
158    PipelineSend,
159    PipelineReceive,
160
161    // ── Fused operations ──
162    FusedMatMulBiasReLU,
163    FusedMatMulBias,
164    FusedLinearGeLU,
165    FusedAttentionLayerNorm,
166    FusedLinearSiLU,
167    FusedConvBatchNormReLU,
168}
169
170impl TensorOp {
171    pub fn name(&self) -> &'static str {
172        match self {
173            // Arithmetic
174            Self::Add => "tensor.add",
175            Self::Sub => "tensor.sub",
176            Self::Mul => "tensor.mul",
177            Self::Div => "tensor.div",
178            Self::Neg => "tensor.neg",
179            Self::MatMul => "tensor.matmul",
180            Self::Linear => "tensor.linear",
181            Self::Conv2D => "tensor.conv2d",
182            Self::Embedding => "tensor.embedding",
183            // Activations
184            Self::ReLU => "tensor.relu",
185            Self::GeLU => "tensor.gelu",
186            Self::SiLU => "tensor.silu",
187            Self::Sigmoid => "tensor.sigmoid",
188            Self::Softmax => "tensor.softmax",
189            Self::Tanh => "tensor.tanh",
190            Self::LeakyReLU => "tensor.leaky_relu",
191            Self::ELU => "tensor.elu",
192            Self::Mish => "tensor.mish",
193            Self::HardSwish => "tensor.hard_swish",
194            Self::HardSigmoid => "tensor.hard_sigmoid",
195            // Normalisation
196            Self::LayerNorm => "tensor.layernorm",
197            Self::RMSNorm => "tensor.rmsnorm",
198            Self::BatchNorm => "tensor.batchnorm",
199            Self::GroupNorm => "tensor.groupnorm",
200            Self::InstanceNorm => "tensor.instancenorm",
201            // Shape
202            Self::Reshape => "tensor.reshape",
203            Self::Transpose => "tensor.transpose",
204            Self::Concat => "tensor.concat",
205            Self::Split => "tensor.split",
206            Self::Gather => "tensor.gather",
207            Self::Scatter => "tensor.scatter",
208            Self::Squeeze => "tensor.squeeze",
209            Self::Unsqueeze => "tensor.unsqueeze",
210            Self::Permute => "tensor.permute",
211            Self::Expand => "tensor.expand",
212            Self::Slice => "tensor.slice",
213            Self::Pad => "tensor.pad",
214            Self::Tile => "tensor.tile",
215            // Constants
216            Self::Constant => "tensor.constant",
217            Self::Zeros => "tensor.zeros",
218            Self::Ones => "tensor.ones",
219            Self::Arange => "tensor.arange",
220            Self::Full => "tensor.full",
221            // Attention variants
222            Self::Attention => "tensor.attention",
223            Self::MultiHeadAttention => "tensor.multi_head_attention",
224            Self::MultiQueryAttention => "tensor.multi_query_attention",
225            Self::GroupedQueryAttention => "tensor.grouped_query_attention",
226            Self::FlashAttention => "tensor.flash_attention",
227            Self::SlidingWindowAttention => "tensor.sliding_window_attention",
228            Self::CrossAttention => "tensor.cross_attention",
229            Self::PagedAttention => "tensor.paged_attention",
230            // MoE
231            Self::MoEDispatch => "tensor.moe_dispatch",
232            Self::MoECombine => "tensor.moe_combine",
233            // Conv variants
234            Self::Conv1D => "tensor.conv1d",
235            Self::Conv3D => "tensor.conv3d",
236            Self::ConvTranspose2D => "tensor.conv_transpose2d",
237            Self::DepthwiseConv2D => "tensor.depthwise_conv2d",
238            Self::DilatedConv2D => "tensor.dilated_conv2d",
239            // Pooling
240            Self::MaxPool2D => "tensor.maxpool2d",
241            Self::AvgPool2D => "tensor.avgpool2d",
242            Self::AdaptiveAvgPool2D => "tensor.adaptive_avgpool2d",
243            Self::GlobalAvgPool => "tensor.global_avgpool",
244            // Recurrent
245            Self::LSTMCell => "tensor.lstm_cell",
246            Self::GRUCell => "tensor.gru_cell",
247            Self::RNNCell => "tensor.rnn_cell",
248            // Advanced math
249            Self::Einsum => "tensor.einsum",
250            Self::FFT => "tensor.fft",
251            Self::IFFT => "tensor.ifft",
252            Self::SVD => "tensor.svd",
253            Self::Eig => "tensor.eig",
254            Self::Solve => "tensor.solve",
255            Self::TopK => "tensor.topk",
256            Self::Sort => "tensor.sort",
257            Self::Cumsum => "tensor.cumsum",
258            Self::Where => "tensor.where",
259            Self::Clamp => "tensor.clamp",
260            // Sparse
261            Self::SparseMatMul => "tensor.sparse_matmul",
262            Self::SparseEmbedding => "tensor.sparse_embedding",
263            // Quantisation
264            Self::Quantize => "tensor.quantize",
265            Self::Dequantize => "tensor.dequantize",
266            Self::QuantizeInt4 => "tensor.quantize_int4",
267            Self::DequantizeInt4 => "tensor.dequantize_int4",
268            Self::QuantizeFp8 => "tensor.quantize_fp8",
269            Self::DequantizeFp8 => "tensor.dequantize_fp8",
270            // Diffusion / Generative
271            Self::UNetDownBlock => "tensor.unet_down_block",
272            Self::UNetUpBlock => "tensor.unet_up_block",
273            Self::TimestepEmbedding => "tensor.timestep_embedding",
274            // GNN
275            Self::GNNMessagePassing => "tensor.gnn_message_passing",
276            Self::GNNGlobalPooling => "tensor.gnn_global_pooling",
277            // Memory management
278            Self::Checkpoint => "tensor.checkpoint",
279            Self::Offload => "tensor.offload",
280            Self::GradAccumulate => "tensor.grad_accumulate",
281            // Gradient operations
282            Self::GradMatMul => "tensor.grad_matmul",
283            Self::GradReLU => "tensor.grad_relu",
284            Self::GradSoftmax => "tensor.grad_softmax",
285            Self::GradLayerNorm => "tensor.grad_layernorm",
286            Self::GradAttention => "tensor.grad_attention",
287            Self::GradConv2D => "tensor.grad_conv2d",
288            Self::GradLinear => "tensor.grad_linear",
289            Self::GradGeLU => "tensor.grad_gelu",
290            // Parallelism
291            Self::ParallelSplit => "tensor.parallel_split",
292            Self::ParallelAllReduce => "tensor.parallel_allreduce",
293            Self::PipelineSend => "tensor.pipeline_send",
294            Self::PipelineReceive => "tensor.pipeline_receive",
295            // Fused operations
296            Self::FusedMatMulBiasReLU => "tensor.fused_matmul_bias_relu",
297            Self::FusedMatMulBias => "tensor.fused_matmul_bias",
298            Self::FusedLinearGeLU => "tensor.fused_linear_gelu",
299            Self::FusedAttentionLayerNorm => "tensor.fused_attention_layernorm",
300            Self::FusedLinearSiLU => "tensor.fused_linear_silu",
301            Self::FusedConvBatchNormReLU => "tensor.fused_conv_batchnorm_relu",
302        }
303    }
304
305    pub fn from_name(name: &str) -> Option<Self> {
306        match name {
307            "tensor.add" => Some(Self::Add),
308            "tensor.sub" => Some(Self::Sub),
309            "tensor.mul" => Some(Self::Mul),
310            "tensor.div" => Some(Self::Div),
311            "tensor.neg" => Some(Self::Neg),
312            "tensor.matmul" => Some(Self::MatMul),
313            "tensor.linear" => Some(Self::Linear),
314            "tensor.conv2d" => Some(Self::Conv2D),
315            "tensor.embedding" => Some(Self::Embedding),
316            "tensor.relu" => Some(Self::ReLU),
317            "tensor.gelu" => Some(Self::GeLU),
318            "tensor.silu" => Some(Self::SiLU),
319            "tensor.sigmoid" => Some(Self::Sigmoid),
320            "tensor.softmax" => Some(Self::Softmax),
321            "tensor.tanh" => Some(Self::Tanh),
322            "tensor.leaky_relu" => Some(Self::LeakyReLU),
323            "tensor.elu" => Some(Self::ELU),
324            "tensor.mish" => Some(Self::Mish),
325            "tensor.hard_swish" => Some(Self::HardSwish),
326            "tensor.hard_sigmoid" => Some(Self::HardSigmoid),
327            "tensor.layernorm" => Some(Self::LayerNorm),
328            "tensor.rmsnorm" => Some(Self::RMSNorm),
329            "tensor.batchnorm" => Some(Self::BatchNorm),
330            "tensor.groupnorm" => Some(Self::GroupNorm),
331            "tensor.instancenorm" => Some(Self::InstanceNorm),
332            "tensor.reshape" => Some(Self::Reshape),
333            "tensor.transpose" => Some(Self::Transpose),
334            "tensor.concat" => Some(Self::Concat),
335            "tensor.split" => Some(Self::Split),
336            "tensor.gather" => Some(Self::Gather),
337            "tensor.scatter" => Some(Self::Scatter),
338            "tensor.squeeze" => Some(Self::Squeeze),
339            "tensor.unsqueeze" => Some(Self::Unsqueeze),
340            "tensor.permute" => Some(Self::Permute),
341            "tensor.expand" => Some(Self::Expand),
342            "tensor.slice" => Some(Self::Slice),
343            "tensor.pad" => Some(Self::Pad),
344            "tensor.tile" => Some(Self::Tile),
345            "tensor.constant" => Some(Self::Constant),
346            "tensor.zeros" => Some(Self::Zeros),
347            "tensor.ones" => Some(Self::Ones),
348            "tensor.arange" => Some(Self::Arange),
349            "tensor.full" => Some(Self::Full),
350            "tensor.attention" => Some(Self::Attention),
351            "tensor.multi_head_attention" => Some(Self::MultiHeadAttention),
352            "tensor.multi_query_attention" => Some(Self::MultiQueryAttention),
353            "tensor.grouped_query_attention" => Some(Self::GroupedQueryAttention),
354            "tensor.flash_attention" => Some(Self::FlashAttention),
355            "tensor.sliding_window_attention" => Some(Self::SlidingWindowAttention),
356            "tensor.cross_attention" => Some(Self::CrossAttention),
357            "tensor.paged_attention" => Some(Self::PagedAttention),
358            "tensor.moe_dispatch" => Some(Self::MoEDispatch),
359            "tensor.moe_combine" => Some(Self::MoECombine),
360            "tensor.conv1d" => Some(Self::Conv1D),
361            "tensor.conv3d" => Some(Self::Conv3D),
362            "tensor.conv_transpose2d" => Some(Self::ConvTranspose2D),
363            "tensor.depthwise_conv2d" => Some(Self::DepthwiseConv2D),
364            "tensor.dilated_conv2d" => Some(Self::DilatedConv2D),
365            "tensor.maxpool2d" => Some(Self::MaxPool2D),
366            "tensor.avgpool2d" => Some(Self::AvgPool2D),
367            "tensor.adaptive_avgpool2d" => Some(Self::AdaptiveAvgPool2D),
368            "tensor.global_avgpool" => Some(Self::GlobalAvgPool),
369            "tensor.lstm_cell" => Some(Self::LSTMCell),
370            "tensor.gru_cell" => Some(Self::GRUCell),
371            "tensor.rnn_cell" => Some(Self::RNNCell),
372            "tensor.einsum" => Some(Self::Einsum),
373            "tensor.fft" => Some(Self::FFT),
374            "tensor.ifft" => Some(Self::IFFT),
375            "tensor.svd" => Some(Self::SVD),
376            "tensor.eig" => Some(Self::Eig),
377            "tensor.solve" => Some(Self::Solve),
378            "tensor.topk" => Some(Self::TopK),
379            "tensor.sort" => Some(Self::Sort),
380            "tensor.cumsum" => Some(Self::Cumsum),
381            "tensor.where" => Some(Self::Where),
382            "tensor.clamp" => Some(Self::Clamp),
383            "tensor.sparse_matmul" => Some(Self::SparseMatMul),
384            "tensor.sparse_embedding" => Some(Self::SparseEmbedding),
385            "tensor.quantize" => Some(Self::Quantize),
386            "tensor.dequantize" => Some(Self::Dequantize),
387            "tensor.quantize_int4" => Some(Self::QuantizeInt4),
388            "tensor.dequantize_int4" => Some(Self::DequantizeInt4),
389            "tensor.quantize_fp8" => Some(Self::QuantizeFp8),
390            "tensor.dequantize_fp8" => Some(Self::DequantizeFp8),
391            "tensor.unet_down_block" => Some(Self::UNetDownBlock),
392            "tensor.unet_up_block" => Some(Self::UNetUpBlock),
393            "tensor.timestep_embedding" => Some(Self::TimestepEmbedding),
394            "tensor.gnn_message_passing" => Some(Self::GNNMessagePassing),
395            "tensor.gnn_global_pooling" => Some(Self::GNNGlobalPooling),
396            "tensor.checkpoint" => Some(Self::Checkpoint),
397            "tensor.offload" => Some(Self::Offload),
398            "tensor.grad_accumulate" => Some(Self::GradAccumulate),
399            "tensor.grad_matmul" => Some(Self::GradMatMul),
400            "tensor.grad_relu" => Some(Self::GradReLU),
401            "tensor.grad_softmax" => Some(Self::GradSoftmax),
402            "tensor.grad_layernorm" => Some(Self::GradLayerNorm),
403            "tensor.grad_attention" => Some(Self::GradAttention),
404            "tensor.grad_conv2d" => Some(Self::GradConv2D),
405            "tensor.grad_linear" => Some(Self::GradLinear),
406            "tensor.grad_gelu" => Some(Self::GradGeLU),
407            "tensor.parallel_split" => Some(Self::ParallelSplit),
408            "tensor.parallel_allreduce" => Some(Self::ParallelAllReduce),
409            "tensor.pipeline_send" => Some(Self::PipelineSend),
410            "tensor.pipeline_receive" => Some(Self::PipelineReceive),
411            "tensor.fused_matmul_bias_relu" => Some(Self::FusedMatMulBiasReLU),
412            "tensor.fused_matmul_bias" => Some(Self::FusedMatMulBias),
413            "tensor.fused_linear_gelu" => Some(Self::FusedLinearGeLU),
414            "tensor.fused_attention_layernorm" => Some(Self::FusedAttentionLayerNorm),
415            "tensor.fused_linear_silu" => Some(Self::FusedLinearSiLU),
416            "tensor.fused_conv_batchnorm_relu" => Some(Self::FusedConvBatchNormReLU),
417            _ => None,
418        }
419    }
420
421    pub fn num_inputs(&self) -> (usize, usize) {
422        match self {
423            // Unary (1 input)
424            Self::Neg
425            | Self::ReLU
426            | Self::GeLU
427            | Self::SiLU
428            | Self::Sigmoid
429            | Self::Tanh
430            | Self::LeakyReLU
431            | Self::ELU
432            | Self::Mish
433            | Self::HardSwish
434            | Self::HardSigmoid
435            | Self::Reshape
436            | Self::Transpose
437            | Self::Squeeze
438            | Self::Unsqueeze
439            | Self::Permute
440            | Self::Expand
441            | Self::Slice
442            | Self::Pad
443            | Self::Tile
444            | Self::Quantize
445            | Self::Dequantize
446            | Self::QuantizeInt4
447            | Self::DequantizeInt4
448            | Self::QuantizeFp8
449            | Self::DequantizeFp8
450            | Self::Offload
451            | Self::Checkpoint
452            | Self::GradReLU
453            | Self::GradGeLU
454            | Self::Softmax
455            | Self::Cumsum
456            | Self::Sort
457            | Self::TopK
458            | Self::FFT
459            | Self::IFFT
460            | Self::SVD
461            | Self::Eig
462            | Self::GlobalAvgPool
463            | Self::AdaptiveAvgPool2D
464            | Self::GNNGlobalPooling => (1, 1),
465
466            // Binary (2 inputs)
467            Self::Add
468            | Self::Sub
469            | Self::Mul
470            | Self::Div
471            | Self::MatMul
472            | Self::SparseMatMul
473            | Self::GradMatMul
474            | Self::Embedding
475            | Self::SparseEmbedding
476            | Self::Conv2D
477            | Self::Conv1D
478            | Self::Conv3D
479            | Self::ConvTranspose2D
480            | Self::DepthwiseConv2D
481            | Self::DilatedConv2D
482            | Self::MaxPool2D
483            | Self::AvgPool2D
484            | Self::Solve
485            | Self::GradConv2D
486            | Self::Concat => (2, 2),
487
488            // Ternary (3 inputs)
489            Self::Linear
490            | Self::FusedMatMulBias
491            | Self::FusedLinearGeLU
492            | Self::FusedMatMulBiasReLU
493            | Self::FusedLinearSiLU
494            | Self::Where
495            | Self::Clamp
496            | Self::GradLinear => (3, 3),
497
498            // Attention (3-4 inputs: Q, K, V, optional mask)
499            Self::Attention
500            | Self::MultiHeadAttention
501            | Self::MultiQueryAttention
502            | Self::GroupedQueryAttention
503            | Self::FlashAttention
504            | Self::SlidingWindowAttention
505            | Self::CrossAttention
506            | Self::GradAttention => (3, 4),
507            Self::PagedAttention => (3, 5),
508            Self::FusedAttentionLayerNorm => (3, 5),
509
510            // Normalisation (variable: input + scale + bias)
511            Self::LayerNorm
512            | Self::RMSNorm
513            | Self::GroupNorm
514            | Self::InstanceNorm
515            | Self::GradLayerNorm => (2, 3),
516            Self::BatchNorm | Self::FusedConvBatchNormReLU => (3, 5),
517
518            // Recurrent (2 inputs: input, hidden state)
519            Self::LSTMCell | Self::GRUCell | Self::RNNCell => (2, 2),
520
521            // GNN (2 inputs: node features, edge index)
522            Self::GNNMessagePassing => (2, 3),
523
524            // Diffusion blocks (2-3 inputs)
525            Self::UNetDownBlock | Self::UNetUpBlock => (2, 3),
526            Self::TimestepEmbedding => (1, 1),
527
528            // MoE
529            Self::MoEDispatch => (2, 3),
530            Self::MoECombine => (2, 3),
531
532            // Constants (0 inputs)
533            Self::Constant | Self::Zeros | Self::Ones | Self::Arange | Self::Full => (0, 0),
534
535            // Einsum (variable)
536            Self::Einsum => (1, usize::MAX),
537
538            // Parallelism / memory
539            Self::GradAccumulate
540            | Self::GradSoftmax
541            | Self::ParallelSplit
542            | Self::ParallelAllReduce
543            | Self::PipelineSend
544            | Self::PipelineReceive
545            | Self::Gather
546            | Self::Scatter
547            | Self::Split => (1, usize::MAX),
548        }
549    }
550
551    /// Returns the asymptotic FLOPs formula as a human-readable string.
552    pub fn flops_formula(&self) -> &'static str {
553        match self {
554            Self::MatMul | Self::SparseMatMul => "2*M*N*K",
555            Self::Linear => "2*M*N*K + N (bias)",
556            Self::Add | Self::Sub | Self::Mul | Self::Div => "N (element count)",
557            Self::ReLU
558            | Self::Sigmoid
559            | Self::Tanh
560            | Self::LeakyReLU
561            | Self::ELU
562            | Self::HardSigmoid => "N",
563            Self::GeLU | Self::SiLU | Self::Mish | Self::HardSwish => "~8*N",
564            Self::Softmax => "5*N (exp + sum + div)",
565            Self::LayerNorm | Self::RMSNorm | Self::GroupNorm | Self::InstanceNorm => "7*N",
566            Self::BatchNorm => "5*N",
567            Self::Conv2D | Self::DepthwiseConv2D | Self::DilatedConv2D => "2*Cout*Cin*Kh*Kw*Oh*Ow",
568            Self::Conv1D => "2*Cout*Cin*K*Oout",
569            Self::Conv3D => "2*Cout*Cin*Kd*Kh*Kw*Od*Oh*Ow",
570            Self::Attention
571            | Self::MultiHeadAttention
572            | Self::GroupedQueryAttention
573            | Self::MultiQueryAttention
574            | Self::FlashAttention
575            | Self::SlidingWindowAttention
576            | Self::CrossAttention => "2*B*H*(S^2*D + S*D^2)",
577            Self::LSTMCell => "4*(input_size+hidden)*hidden*2",
578            Self::GRUCell => "3*(input_size+hidden)*hidden*2",
579            Self::RNNCell => "(input_size+hidden)*hidden*2",
580            Self::FFT | Self::IFFT => "5*N*log2(N)",
581            Self::Einsum => "depends on equation",
582            Self::MaxPool2D | Self::AvgPool2D | Self::AdaptiveAvgPool2D | Self::GlobalAvgPool => {
583                "N (comparisons or additions)"
584            }
585            Self::Reshape
586            | Self::Transpose
587            | Self::Squeeze
588            | Self::Unsqueeze
589            | Self::Permute
590            | Self::Expand
591            | Self::Slice
592            | Self::Pad
593            | Self::Tile
594            | Self::Concat
595            | Self::Split
596            | Self::Gather
597            | Self::Scatter => "0 (no compute)",
598            _ => "varies",
599        }
600    }
601
602    /// Returns `true` if this op performs no arithmetic (zero FLOPs).
603    #[inline]
604    pub fn is_zero_flop(&self) -> bool {
605        matches!(
606            self,
607            Self::Reshape
608                | Self::Transpose
609                | Self::Squeeze
610                | Self::Unsqueeze
611                | Self::Permute
612                | Self::Expand
613                | Self::Slice
614                | Self::Pad
615                | Self::Tile
616                | Self::Concat
617                | Self::Split
618                | Self::Gather
619                | Self::Scatter
620                | Self::Constant
621                | Self::Zeros
622                | Self::Ones
623                | Self::Arange
624                | Self::Full
625                | Self::Checkpoint
626                | Self::Offload
627                | Self::PipelineSend
628                | Self::PipelineReceive
629                | Self::ParallelSplit
630                | Self::ParallelAllReduce
631        )
632    }
633
634    /// Returns `true` if this is an element-wise (unary or binary) activation.
635    #[inline]
636    pub fn is_activation(&self) -> bool {
637        matches!(
638            self,
639            Self::ReLU
640                | Self::GeLU
641                | Self::SiLU
642                | Self::Sigmoid
643                | Self::Tanh
644                | Self::LeakyReLU
645                | Self::ELU
646                | Self::Mish
647                | Self::HardSwish
648                | Self::HardSigmoid
649        )
650    }
651
652    /// Returns `true` if this is an attention variant.
653    #[inline]
654    pub fn is_attention(&self) -> bool {
655        matches!(
656            self,
657            Self::Attention
658                | Self::MultiHeadAttention
659                | Self::MultiQueryAttention
660                | Self::GroupedQueryAttention
661                | Self::FlashAttention
662                | Self::SlidingWindowAttention
663                | Self::CrossAttention
664                | Self::PagedAttention
665        )
666    }
667
668    /// Returns `true` if this is a convolution variant.
669    #[inline]
670    pub fn is_convolution(&self) -> bool {
671        matches!(
672            self,
673            Self::Conv1D
674                | Self::Conv2D
675                | Self::Conv3D
676                | Self::ConvTranspose2D
677                | Self::DepthwiseConv2D
678                | Self::DilatedConv2D
679        )
680    }
681
682    /// Returns `true` if this is a normalisation op.
683    #[inline]
684    pub fn is_normalisation(&self) -> bool {
685        matches!(
686            self,
687            Self::LayerNorm
688                | Self::RMSNorm
689                | Self::BatchNorm
690                | Self::GroupNorm
691                | Self::InstanceNorm
692        )
693    }
694
695    /// Returns `true` if this is a fused operation.
696    #[inline]
697    pub fn is_fused(&self) -> bool {
698        matches!(
699            self,
700            Self::FusedMatMulBiasReLU
701                | Self::FusedMatMulBias
702                | Self::FusedLinearGeLU
703                | Self::FusedAttentionLayerNorm
704                | Self::FusedLinearSiLU
705                | Self::FusedConvBatchNormReLU
706        )
707    }
708
709    /// Returns `true` if this is a gradient (backward) operation.
710    #[inline]
711    pub fn is_gradient(&self) -> bool {
712        matches!(
713            self,
714            Self::GradMatMul
715                | Self::GradReLU
716                | Self::GradSoftmax
717                | Self::GradLayerNorm
718                | Self::GradAttention
719                | Self::GradConv2D
720                | Self::GradLinear
721                | Self::GradGeLU
722        )
723    }
724}
725
726#[cfg(test)]
727mod tests {
728    use super::*;
729
730    #[test]
731    fn test_op_name_roundtrip() {
732        for op in &[
733            TensorOp::MatMul,
734            TensorOp::ReLU,
735            TensorOp::Attention,
736            TensorOp::Softmax,
737        ] {
738            let name = op.name();
739            let recovered = TensorOp::from_name(name).unwrap();
740            assert_eq!(op, &recovered);
741        }
742    }
743
744    #[test]
745    fn test_all_ops_have_names() {
746        let ops = vec![
747            TensorOp::Add,
748            TensorOp::Sub,
749            TensorOp::Mul,
750            TensorOp::Div,
751            TensorOp::MatMul,
752            TensorOp::Linear,
753            TensorOp::ReLU,
754            TensorOp::GeLU,
755            TensorOp::Softmax,
756            TensorOp::LayerNorm,
757            TensorOp::Attention,
758        ];
759        for op in ops {
760            assert!(!op.name().is_empty());
761            assert!(TensorOp::from_name(op.name()).is_some());
762        }
763    }
764}
lift_tensor/ops.rs

lift_tensor/
ops.rs
