lift_tensor/
shape.rs

1use crate::ops::TensorOp;
2use lift_core::attributes::Attributes;
3use lift_core::types::{Dimension, TensorTypeInfo};
4
5#[derive(Debug)]
6pub struct ShapeInference;
7
8/// Reads a spatial conv/pool parameter (`stride`, `padding`, `dilation`) from
9/// `attrs`, falling back to `default` when absent. A single integer applies
10/// to every spatial axis (symmetric kernels only — matching every example
11/// and test in this codebase, none of which use per-axis values).
12fn spatial_param(attrs: Option<&Attributes>, key: &str, default: i64) -> i64 {
13    attrs.and_then(|a| a.get_integer(key)).unwrap_or(default)
14}
15
16/// Standard convolution/pooling output-length formula:
17/// `floor((in + 2*padding - dilation*(kernel-1) - 1) / stride) + 1`.
18fn conv_output_dim(input: i64, kernel: i64, stride: i64, padding: i64, dilation: i64) -> i64 {
19    let numerator = input + 2 * padding - dilation * (kernel - 1) - 1;
20    (numerator.max(0) / stride.max(1)) + 1
21}
22
23impl ShapeInference {
24    pub fn infer_output_shape(
25        op: &TensorOp,
26        inputs: &[&TensorTypeInfo],
27        attrs: Option<&Attributes>,
28    ) -> Result<Vec<TensorTypeInfo>, String> {
29        match op {
30            // ── Binary element-wise (broadcast) ──
31            TensorOp::Add | TensorOp::Sub | TensorOp::Mul | TensorOp::Div => {
32                if inputs.len() != 2 {
33                    return Err(format!("{} requires 2 inputs", op.name()));
34                }
35                let result = broadcast_shapes(&inputs[0].shape, &inputs[1].shape)?;
36                Ok(vec![TensorTypeInfo {
37                    shape: result,
38                    dtype: inputs[0].dtype,
39                    layout: inputs[0].layout,
40                }])
41            }
42
43            // ── Unary shape-preserving ──
44            TensorOp::Neg
45            | TensorOp::ReLU
46            | TensorOp::GeLU
47            | TensorOp::SiLU
48            | TensorOp::Sigmoid
49            | TensorOp::Tanh
50            | TensorOp::LeakyReLU
51            | TensorOp::ELU
52            | TensorOp::Mish
53            | TensorOp::HardSwish
54            | TensorOp::HardSigmoid
55            | TensorOp::Softmax
56            | TensorOp::Cumsum
57            | TensorOp::Quantize
58            | TensorOp::Dequantize
59            | TensorOp::QuantizeInt4
60            | TensorOp::DequantizeInt4
61            | TensorOp::QuantizeFp8
62            | TensorOp::DequantizeFp8
63            | TensorOp::Checkpoint
64            | TensorOp::Offload
65            | TensorOp::GradReLU
66            | TensorOp::GradGeLU
67            | TensorOp::GradSoftmax => {
68                if inputs.is_empty() {
69                    return Err(format!("{} requires at least 1 input", op.name()));
70                }
71                Ok(vec![inputs[0].clone()])
72            }
73
74            // ── Normalisation (shape-preserving) ──
75            TensorOp::LayerNorm
76            | TensorOp::RMSNorm
77            | TensorOp::BatchNorm
78            | TensorOp::GroupNorm
79            | TensorOp::InstanceNorm
80            | TensorOp::GradLayerNorm => {
81                if inputs.is_empty() {
82                    return Err(format!("{} requires at least 1 input", op.name()));
83                }
84                Ok(vec![inputs[0].clone()])
85            }
86
87            // ── MatMul ──
88            TensorOp::MatMul | TensorOp::SparseMatMul => {
89                if inputs.len() != 2 {
90                    return Err("matmul requires 2 inputs".into());
91                }
92                let a = &inputs[0].shape;
93                let b = &inputs[1].shape;
94                if a.len() < 2 || b.len() < 2 {
95                    return Err("matmul inputs must be at least 2D".into());
96                }
97                let m = a[a.len() - 2].clone();
98                let n = b[b.len() - 1].clone();
99
100                let k_a = &a[a.len() - 1];
101                let k_b = &b[b.len() - 2];
102                if let (Some(ka), Some(kb)) = (k_a.static_value(), k_b.static_value()) {
103                    if ka != kb {
104                        return Err(format!("matmul inner dimension mismatch: {} vs {}", ka, kb));
105                    }
106                }
107
108                let mut result_shape = Vec::new();
109                let batch_a = &a[..a.len() - 2];
110                let batch_b = &b[..b.len() - 2];
111                let batch = broadcast_shapes(batch_a, batch_b)?;
112                result_shape.extend(batch);
113                result_shape.push(m);
114                result_shape.push(n);
115
116                Ok(vec![TensorTypeInfo {
117                    shape: result_shape,
118                    dtype: inputs[0].dtype,
119                    layout: inputs[0].layout,
120                }])
121            }
122
123            // ── Linear ──
124            TensorOp::Linear => {
125                if inputs.len() < 2 {
126                    return Err("linear requires at least 2 inputs (x, W)".into());
127                }
128                let x = &inputs[0].shape;
129                let w = &inputs[1].shape;
130                if x.is_empty() || w.len() != 2 {
131                    return Err("linear: x must be at least 1D, W must be 2D".into());
132                }
133                let mut result_shape = x[..x.len() - 1].to_vec();
134                result_shape.push(w[1].clone());
135
136                Ok(vec![TensorTypeInfo {
137                    shape: result_shape,
138                    dtype: inputs[0].dtype,
139                    layout: inputs[0].layout,
140                }])
141            }
142
143            // ── Conv2D ──
144            TensorOp::Conv2D | TensorOp::DepthwiseConv2D | TensorOp::DilatedConv2D => {
145                if inputs.len() < 2 {
146                    return Err("conv2d requires at least 2 inputs (input, kernel)".into());
147                }
148                let input = &inputs[0].shape;
149                let kernel = &inputs[1].shape;
150                if input.len() != 4 || kernel.len() != 4 {
151                    return Err("conv2d: input and kernel must be 4D (NCHW)".into());
152                }
153
154                // `quantum.dilated_conv2d`'s dilation and any non-default
155                // stride/padding come from attrs — without them (the common
156                // case in tests), this reduces to the stride=1/padding=0/
157                // dilation=1 formula `in - kernel + 1`.
158                let stride = spatial_param(attrs, "stride", 1);
159                let padding = spatial_param(attrs, "padding", 0);
160                let dilation = spatial_param(
161                    attrs,
162                    "dilation",
163                    if matches!(op, TensorOp::DilatedConv2D) {
164                        2
165                    } else {
166                        1
167                    },
168                );
169
170                let n = input[0].clone();
171                let cout = kernel[0].clone();
172                let h_out = match (&input[2], &kernel[2]) {
173                    (Dimension::Constant(ih), Dimension::Constant(kh)) => Dimension::Constant(
174                        conv_output_dim(*ih as i64, *kh as i64, stride, padding, dilation) as usize,
175                    ),
176                    _ => Dimension::Symbolic("H_out".into()),
177                };
178                let w_out = match (&input[3], &kernel[3]) {
179                    (Dimension::Constant(iw), Dimension::Constant(kw)) => Dimension::Constant(
180                        conv_output_dim(*iw as i64, *kw as i64, stride, padding, dilation) as usize,
181                    ),
182                    _ => Dimension::Symbolic("W_out".into()),
183                };
184
185                Ok(vec![TensorTypeInfo {
186                    shape: vec![n, cout, h_out, w_out],
187                    dtype: inputs[0].dtype,
188                    layout: inputs[0].layout,
189                }])
190            }
191
192            // ── Conv1D ──
193            TensorOp::Conv1D => {
194                if inputs.len() < 2 {
195                    return Err("conv1d requires at least 2 inputs".into());
196                }
197                let input = &inputs[0].shape;
198                let kernel = &inputs[1].shape;
199                if input.len() != 3 || kernel.len() != 3 {
200                    return Err("conv1d: input [N,C,L] and kernel [Cout,Cin,K]".into());
201                }
202                let stride = spatial_param(attrs, "stride", 1);
203                let padding = spatial_param(attrs, "padding", 0);
204                let dilation = spatial_param(attrs, "dilation", 1);
205                let n = input[0].clone();
206                let cout = kernel[0].clone();
207                let l_out = match (&input[2], &kernel[2]) {
208                    (Dimension::Constant(il), Dimension::Constant(kl)) => Dimension::Constant(
209                        conv_output_dim(*il as i64, *kl as i64, stride, padding, dilation) as usize,
210                    ),
211                    _ => Dimension::Symbolic("L_out".into()),
212                };
213                Ok(vec![TensorTypeInfo {
214                    shape: vec![n, cout, l_out],
215                    dtype: inputs[0].dtype,
216                    layout: inputs[0].layout,
217                }])
218            }
219
220            // ── Conv3D ──
221            TensorOp::Conv3D => {
222                if inputs.len() < 2 {
223                    return Err("conv3d requires at least 2 inputs".into());
224                }
225                let input = &inputs[0].shape;
226                let kernel = &inputs[1].shape;
227                if input.len() != 5 || kernel.len() != 5 {
228                    return Err("conv3d: input [N,C,D,H,W] and kernel [Cout,Cin,Kd,Kh,Kw]".into());
229                }
230                let stride = spatial_param(attrs, "stride", 1);
231                let padding = spatial_param(attrs, "padding", 0);
232                let dilation = spatial_param(attrs, "dilation", 1);
233                let n = input[0].clone();
234                let cout = kernel[0].clone();
235                let dims: Vec<Dimension> = (2..5)
236                    .map(|i| match (&input[i], &kernel[i]) {
237                        (Dimension::Constant(iv), Dimension::Constant(kv)) => Dimension::Constant(
238                            conv_output_dim(*iv as i64, *kv as i64, stride, padding, dilation)
239                                as usize,
240                        ),
241                        _ => Dimension::Symbolic(format!("dim{}_out", i)),
242                    })
243                    .collect();
244                Ok(vec![TensorTypeInfo {
245                    shape: vec![n, cout, dims[0].clone(), dims[1].clone(), dims[2].clone()],
246                    dtype: inputs[0].dtype,
247                    layout: inputs[0].layout,
248                }])
249            }
250
251            // ── Pooling ──
252            // The second input, if present, is a kernel-shaped tensor whose
253            // spatial dims give the pooling window (mirroring how conv reads
254            // its window from the kernel tensor) — this matches
255            // test_shape_max_pool2d. Default stride is the window size
256            // (non-overlapping pooling, the standard framework default when
257            // stride is unset), overridable via a `stride` attr; `padding`
258            // defaults to 0. With only 1 input (no window given), the
259            // pooling window is unknown, so the shape is left unchanged.
260            TensorOp::MaxPool2D | TensorOp::AvgPool2D => {
261                if inputs.is_empty() {
262                    return Err(format!("{} requires at least 1 input", op.name()));
263                }
264                let Some(kernel) = inputs.get(1) else {
265                    return Ok(vec![inputs[0].clone()]);
266                };
267                let input = &inputs[0].shape;
268                let kh = kernel.shape.first().and_then(|d| d.static_value());
269                let kw = kernel.shape.get(1).and_then(|d| d.static_value());
270                if input.len() < 4 {
271                    return Err(format!("{}: input must be 4D [N,C,H,W]", op.name()));
272                }
273                let padding = spatial_param(attrs, "padding", 0);
274                let mut out = input.clone();
275                if let (Some(kh), Dimension::Constant(ih)) = (kh, &input[2]) {
276                    let stride = spatial_param(attrs, "stride", kh as i64);
277                    out[2] = Dimension::Constant(conv_output_dim(
278                        *ih as i64, kh as i64, stride, padding, 1,
279                    ) as usize);
280                }
281                if let (Some(kw), Dimension::Constant(iw)) = (kw, &input[3]) {
282                    let stride = spatial_param(attrs, "stride", kw as i64);
283                    out[3] = Dimension::Constant(conv_output_dim(
284                        *iw as i64, kw as i64, stride, padding, 1,
285                    ) as usize);
286                }
287                Ok(vec![TensorTypeInfo {
288                    shape: out,
289                    dtype: inputs[0].dtype,
290                    layout: inputs[0].layout,
291                }])
292            }
293
294            // Adaptive pooling targets a caller-specified output size rather
295            // than deriving one from a kernel/stride, and no attrs schema
296            // for that output size exists yet in this codebase — simplified
297            // to the input shape unchanged, same as MaxPool2D/AvgPool2D
298            // without a kernel input, until one is added.
299            TensorOp::AdaptiveAvgPool2D => {
300                if inputs.is_empty() {
301                    return Err("adaptive_avgpool2d requires 1 input".into());
302                }
303                Ok(vec![inputs[0].clone()])
304            }
305
306            TensorOp::GlobalAvgPool => {
307                if inputs.is_empty() {
308                    return Err("global_avgpool requires 1 input".into());
309                }
310                let shape = &inputs[0].shape;
311                if shape.len() < 3 {
312                    return Err("global_avgpool: input must be at least 3D [N,C,...]".into());
313                }
314                // [N, C, ...] -> [N, C, 1, 1, ...]
315                let mut out = vec![shape[0].clone(), shape[1].clone()];
316                for _ in 2..shape.len() {
317                    out.push(Dimension::Constant(1));
318                }
319                Ok(vec![TensorTypeInfo {
320                    shape: out,
321                    dtype: inputs[0].dtype,
322                    layout: inputs[0].layout,
323                }])
324            }
325
326            // ── Attention variants ──
327            TensorOp::Attention
328            | TensorOp::MultiHeadAttention
329            | TensorOp::MultiQueryAttention
330            | TensorOp::GroupedQueryAttention
331            | TensorOp::FlashAttention
332            | TensorOp::SlidingWindowAttention
333            | TensorOp::CrossAttention
334            | TensorOp::PagedAttention
335            | TensorOp::GradAttention => {
336                if inputs.len() < 3 {
337                    return Err("attention requires at least 3 inputs (Q, K, V)".into());
338                }
339                Ok(vec![inputs[0].clone()])
340            }
341
342            // ── Recurrent ──
343            TensorOp::LSTMCell => {
344                if inputs.len() < 2 {
345                    return Err("lstm_cell requires input and hidden state".into());
346                }
347                // Returns (h_new, c_new) with same shape as hidden
348                Ok(vec![inputs[1].clone(), inputs[1].clone()])
349            }
350
351            TensorOp::GRUCell | TensorOp::RNNCell => {
352                if inputs.len() < 2 {
353                    return Err(format!("{} requires input and hidden state", op.name()));
354                }
355                Ok(vec![inputs[1].clone()])
356            }
357
358            // ── Shape / zero-flop ops ──
359            TensorOp::Reshape
360            | TensorOp::Transpose
361            | TensorOp::Squeeze
362            | TensorOp::Unsqueeze
363            | TensorOp::Permute
364            | TensorOp::Expand
365            | TensorOp::Slice
366            | TensorOp::Pad
367            | TensorOp::Tile => {
368                // These need target shape from attributes; passthrough for now
369                if inputs.is_empty() {
370                    return Err(format!("{} requires at least 1 input", op.name()));
371                }
372                Ok(vec![inputs[0].clone()])
373            }
374
375            // ── Concat ──
376            TensorOp::Concat => {
377                if inputs.is_empty() {
378                    return Err("concat requires at least 1 input".into());
379                }
380                Ok(vec![inputs[0].clone()])
381            }
382
383            // ── TopK / Sort ──
384            TensorOp::TopK | TensorOp::Sort => {
385                if inputs.is_empty() {
386                    return Err(format!("{} requires 1 input", op.name()));
387                }
388                Ok(vec![inputs[0].clone()])
389            }
390
391            // ── FFT / IFFT ──
392            TensorOp::FFT | TensorOp::IFFT => {
393                if inputs.is_empty() {
394                    return Err(format!("{} requires 1 input", op.name()));
395                }
396                Ok(vec![inputs[0].clone()])
397            }
398
399            // ── SVD: returns U, S, V ──
400            TensorOp::SVD => {
401                if inputs.is_empty() {
402                    return Err("svd requires 1 input".into());
403                }
404                Ok(vec![inputs[0].clone()])
405            }
406
407            // ── Where: condition, x, y -> x ──
408            TensorOp::Where | TensorOp::Clamp => {
409                if inputs.len() < 2 {
410                    return Err(format!("{} requires at least 2 inputs", op.name()));
411                }
412                Ok(vec![inputs[0].clone()])
413            }
414
415            _ => {
416                // For ops not yet handled, passthrough first input or empty
417                if !inputs.is_empty() {
418                    Ok(vec![inputs[0].clone()])
419                } else {
420                    Ok(Vec::new())
421                }
422            }
423        }
424    }
425
426    pub fn compute_flops(
427        op: &TensorOp,
428        inputs: &[&TensorTypeInfo],
429        attrs: Option<&Attributes>,
430    ) -> Option<u64> {
431        match op {
432            TensorOp::MatMul | TensorOp::SparseMatMul => {
433                if inputs.len() != 2 {
434                    return None;
435                }
436                let a = &inputs[0].shape;
437                let b = &inputs[1].shape;
438                let m = a.get(a.len().checked_sub(2)?)?.static_value()? as u64;
439                let k = a.last()?.static_value()? as u64;
440                let n = b.last()?.static_value()? as u64;
441                let batch: u64 = a[..a.len() - 2]
442                    .iter()
443                    .filter_map(|d| d.static_value())
444                    .map(|v| v as u64)
445                    .product::<u64>()
446                    .max(1);
447                Some(2 * batch * m * n * k)
448            }
449
450            TensorOp::Add | TensorOp::Sub | TensorOp::Mul | TensorOp::Div => {
451                if inputs.is_empty() {
452                    return None;
453                }
454                Some(element_count(&inputs[0].shape)? as u64)
455            }
456
457            TensorOp::ReLU
458            | TensorOp::Sigmoid
459            | TensorOp::Tanh
460            | TensorOp::LeakyReLU
461            | TensorOp::ELU
462            | TensorOp::HardSigmoid => {
463                if inputs.is_empty() {
464                    return None;
465                }
466                Some(element_count(&inputs[0].shape)? as u64)
467            }
468
469            TensorOp::GeLU | TensorOp::SiLU | TensorOp::Mish | TensorOp::HardSwish => {
470                if inputs.is_empty() {
471                    return None;
472                }
473                let n = element_count(&inputs[0].shape)? as u64;
474                Some(8 * n)
475            }
476
477            TensorOp::Softmax => {
478                if inputs.is_empty() {
479                    return None;
480                }
481                let n = element_count(&inputs[0].shape)? as u64;
482                Some(5 * n)
483            }
484
485            TensorOp::LayerNorm
486            | TensorOp::RMSNorm
487            | TensorOp::GroupNorm
488            | TensorOp::InstanceNorm => {
489                if inputs.is_empty() {
490                    return None;
491                }
492                let n = element_count(&inputs[0].shape)? as u64;
493                Some(7 * n)
494            }
495
496            TensorOp::BatchNorm => {
497                if inputs.is_empty() {
498                    return None;
499                }
500                let n = element_count(&inputs[0].shape)? as u64;
501                Some(5 * n)
502            }
503
504            TensorOp::Linear => {
505                if inputs.len() < 2 {
506                    return None;
507                }
508                let x = &inputs[0].shape;
509                let w = &inputs[1].shape;
510                let m: u64 = x[..x.len() - 1]
511                    .iter()
512                    .filter_map(|d| d.static_value())
513                    .map(|v| v as u64)
514                    .product::<u64>()
515                    .max(1);
516                let k = x.last()?.static_value()? as u64;
517                let n = w.last()?.static_value()? as u64;
518                Some(2 * m * n * k + n)
519            }
520
521            TensorOp::Conv2D | TensorOp::DepthwiseConv2D | TensorOp::DilatedConv2D => {
522                if inputs.len() < 2 {
523                    return None;
524                }
525                let kernel = &inputs[1].shape;
526                let cout = kernel[0].static_value()? as u64;
527                let cin = kernel[1].static_value()? as u64;
528                let kh = kernel[2].static_value()? as u64;
529                let kw = kernel[3].static_value()? as u64;
530                let input = &inputs[0].shape;
531                let n = input[0].static_value()? as u64;
532                let ih = input[2].static_value()? as u64;
533                let iw = input[3].static_value()? as u64;
534                let stride = spatial_param(attrs, "stride", 1);
535                let padding = spatial_param(attrs, "padding", 0);
536                let dilation = spatial_param(
537                    attrs,
538                    "dilation",
539                    if matches!(op, TensorOp::DilatedConv2D) {
540                        2
541                    } else {
542                        1
543                    },
544                );
545                let oh = conv_output_dim(ih as i64, kh as i64, stride, padding, dilation) as u64;
546                let ow = conv_output_dim(iw as i64, kw as i64, stride, padding, dilation) as u64;
547                Some(2 * n * cout * cin * kh * kw * oh * ow)
548            }
549
550            TensorOp::Conv1D => {
551                if inputs.len() < 2 {
552                    return None;
553                }
554                let kernel = &inputs[1].shape;
555                let cout = kernel[0].static_value()? as u64;
556                let cin = kernel[1].static_value()? as u64;
557                let k = kernel[2].static_value()? as u64;
558                let input = &inputs[0].shape;
559                let n = input[0].static_value()? as u64;
560                let il = input[2].static_value()? as u64;
561                let stride = spatial_param(attrs, "stride", 1);
562                let padding = spatial_param(attrs, "padding", 0);
563                let dilation = spatial_param(attrs, "dilation", 1);
564                let ol = conv_output_dim(il as i64, k as i64, stride, padding, dilation) as u64;
565                Some(2 * n * cout * cin * k * ol)
566            }
567
568            TensorOp::Conv3D => {
569                if inputs.len() < 2 {
570                    return None;
571                }
572                let kernel = &inputs[1].shape;
573                let cout = kernel.first()?.static_value()? as u64;
574                let cin = kernel.get(1)?.static_value()? as u64;
575                let kd = kernel.get(2)?.static_value()? as u64;
576                let kh = kernel.get(3)?.static_value()? as u64;
577                let kw = kernel.get(4)?.static_value()? as u64;
578                let input = &inputs[0].shape;
579                let n = input.first()?.static_value()? as u64;
580                let id = input.get(2)?.static_value()? as u64;
581                let ih = input.get(3)?.static_value()? as u64;
582                let iw = input.get(4)?.static_value()? as u64;
583                let stride = spatial_param(attrs, "stride", 1);
584                let padding = spatial_param(attrs, "padding", 0);
585                let dilation = spatial_param(attrs, "dilation", 1);
586                let od = conv_output_dim(id as i64, kd as i64, stride, padding, dilation) as u64;
587                let oh = conv_output_dim(ih as i64, kh as i64, stride, padding, dilation) as u64;
588                let ow = conv_output_dim(iw as i64, kw as i64, stride, padding, dilation) as u64;
589                Some(2 * n * cout * cin * kd * kh * kw * od * oh * ow)
590            }
591
592            // Attention variants: 2*B*H*(S^2*D + S*D^2)
593            TensorOp::Attention
594            | TensorOp::MultiHeadAttention
595            | TensorOp::MultiQueryAttention
596            | TensorOp::GroupedQueryAttention
597            | TensorOp::FlashAttention
598            | TensorOp::SlidingWindowAttention
599            | TensorOp::CrossAttention => {
600                if inputs.is_empty() {
601                    return None;
602                }
603                let shape = &inputs[0].shape;
604                if shape.len() < 3 {
605                    return None;
606                }
607                let b = shape[0].static_value().unwrap_or(1) as u64;
608                let s = shape[shape.len() - 2].static_value()? as u64;
609                let d = shape.last()?.static_value()? as u64;
610                let h = if shape.len() >= 4 {
611                    shape[1].static_value().unwrap_or(1) as u64
612                } else {
613                    1
614                };
615                Some(4 * b * h * s * s * d)
616            }
617
618            // Recurrent
619            TensorOp::LSTMCell => {
620                // 4 * (input_size + hidden_size) * hidden_size * 2
621                if inputs.len() < 2 {
622                    return None;
623                }
624                let input_size = inputs[0].shape.last()?.static_value()? as u64;
625                let hidden_size = inputs[1].shape.last()?.static_value()? as u64;
626                Some(8 * (input_size + hidden_size) * hidden_size)
627            }
628
629            TensorOp::GRUCell => {
630                if inputs.len() < 2 {
631                    return None;
632                }
633                let input_size = inputs[0].shape.last()?.static_value()? as u64;
634                let hidden_size = inputs[1].shape.last()?.static_value()? as u64;
635                Some(6 * (input_size + hidden_size) * hidden_size)
636            }
637
638            TensorOp::RNNCell => {
639                if inputs.len() < 2 {
640                    return None;
641                }
642                let input_size = inputs[0].shape.last()?.static_value()? as u64;
643                let hidden_size = inputs[1].shape.last()?.static_value()? as u64;
644                Some(2 * (input_size + hidden_size) * hidden_size)
645            }
646
647            // FFT: 5*N*log2(N)
648            TensorOp::FFT | TensorOp::IFFT => {
649                if inputs.is_empty() {
650                    return None;
651                }
652                let n = element_count(&inputs[0].shape)? as u64;
653                if n == 0 {
654                    return Some(0);
655                }
656                let log2n = (n as f64).log2().ceil() as u64;
657                Some(5 * n * log2n)
658            }
659
660            // Pooling
661            TensorOp::MaxPool2D
662            | TensorOp::AvgPool2D
663            | TensorOp::AdaptiveAvgPool2D
664            | TensorOp::GlobalAvgPool => {
665                if inputs.is_empty() {
666                    return None;
667                }
668                Some(element_count(&inputs[0].shape)? as u64)
669            }
670
671            // Zero-flop ops
672            _ if op.is_zero_flop() => Some(0),
673
674            _ => None,
675        }
676    }
677
678    /// Sums every input's bytes plus the output's bytes — memory traffic
679    /// includes writing the result, not just reading the operands. Every op
680    /// family used to only get this for `MatMul`/`SparseMatMul`; every other
681    /// op (Conv*, Attention, Norm, pooling, activations, ...) silently
682    /// omitted the output entirely, understating real traffic (e.g. by
683    /// ~46% for a typical Conv2D).
684    pub fn compute_memory_bytes(
685        op: &TensorOp,
686        inputs: &[&TensorTypeInfo],
687        attrs: Option<&Attributes>,
688    ) -> Option<u64> {
689        let input_bytes: u64 = inputs
690            .iter()
691            .filter_map(|i| tensor_bytes(i).map(|b| b as u64))
692            .sum();
693        let output_bytes: u64 = Self::infer_output_shape(op, inputs, attrs)
694            .ok()
695            .map(|outs| {
696                outs.iter()
697                    .filter_map(|o| tensor_info_bytes(o).map(|b| b as u64))
698                    .sum()
699            })
700            .unwrap_or(0);
701        Some(input_bytes + output_bytes)
702    }
703}
704
705fn broadcast_shapes(a: &[Dimension], b: &[Dimension]) -> Result<Vec<Dimension>, String> {
706    let max_rank = a.len().max(b.len());
707    let mut result = Vec::with_capacity(max_rank);
708
709    for i in 0..max_rank {
710        let da = if i < a.len() {
711            Some(&a[a.len() - 1 - i])
712        } else {
713            None
714        };
715        let db = if i < b.len() {
716            Some(&b[b.len() - 1 - i])
717        } else {
718            None
719        };
720
721        let dim = match (da, db) {
722            (Some(a_dim), Some(b_dim)) => match (a_dim.static_value(), b_dim.static_value()) {
723                (Some(a_val), Some(b_val)) => {
724                    if a_val == b_val {
725                        Dimension::Constant(a_val)
726                    } else if a_val == 1 {
727                        Dimension::Constant(b_val)
728                    } else if b_val == 1 {
729                        Dimension::Constant(a_val)
730                    } else {
731                        return Err(format!("Shape broadcast error: {} vs {}", a_val, b_val));
732                    }
733                }
734                _ => Dimension::Symbolic("broadcast".into()),
735            },
736            (Some(d), None) | (None, Some(d)) => d.clone(),
737            (None, None) => unreachable!(),
738        };
739        result.push(dim);
740    }
741
742    result.reverse();
743    Ok(result)
744}
745
746fn element_count(shape: &[Dimension]) -> Option<usize> {
747    let mut count = 1usize;
748    for dim in shape {
749        count = count.checked_mul(dim.static_value()?)?;
750    }
751    Some(count)
752}
753
754fn tensor_bytes(info: &TensorTypeInfo) -> Option<usize> {
755    Some(element_count(&info.shape)? * info.dtype.byte_size())
756}
757
758fn tensor_info_bytes(info: &TensorTypeInfo) -> Option<usize> {
759    tensor_bytes(info)
760}
761
762#[cfg(test)]
763mod tests {
764    use super::*;
765    use lift_core::attributes::Attribute;
766    use lift_core::types::{DataType, MemoryLayout};
767
768    fn make_tensor(shape: Vec<usize>, dtype: DataType) -> TensorTypeInfo {
769        TensorTypeInfo {
770            shape: shape.into_iter().map(Dimension::Constant).collect(),
771            dtype,
772            layout: MemoryLayout::Contiguous,
773        }
774    }
775
776    #[test]
777    fn test_matmul_shape() {
778        let a = make_tensor(vec![2, 3, 4], DataType::FP32);
779        let b = make_tensor(vec![2, 4, 5], DataType::FP32);
780        let result =
781            ShapeInference::infer_output_shape(&TensorOp::MatMul, &[&a, &b], None).unwrap();
782        assert_eq!(result.len(), 1);
783        let shape = &result[0].shape;
784        assert_eq!(shape.len(), 3);
785        assert_eq!(shape[0].static_value(), Some(2));
786        assert_eq!(shape[1].static_value(), Some(3));
787        assert_eq!(shape[2].static_value(), Some(5));
788    }
789
790    #[test]
791    fn test_matmul_dimension_mismatch() {
792        let a = make_tensor(vec![3, 4], DataType::FP32);
793        let b = make_tensor(vec![5, 6], DataType::FP32);
794        let result = ShapeInference::infer_output_shape(&TensorOp::MatMul, &[&a, &b], None);
795        assert!(result.is_err());
796    }
797
798    #[test]
799    fn test_matmul_flops() {
800        let a = make_tensor(vec![2, 3], DataType::FP32);
801        let b = make_tensor(vec![3, 4], DataType::FP32);
802        let flops = ShapeInference::compute_flops(&TensorOp::MatMul, &[&a, &b], None);
803        assert_eq!(flops, Some(2 * 2 * 4 * 3)); // 2*M*N*K
804    }
805
806    #[test]
807    fn test_relu_shape() {
808        let a = make_tensor(vec![2, 3, 4], DataType::FP32);
809        let result = ShapeInference::infer_output_shape(&TensorOp::ReLU, &[&a], None).unwrap();
810        assert_eq!(result[0].shape, a.shape);
811    }
812
813    #[test]
814    fn test_linear_shape() {
815        let x = make_tensor(vec![1, 784], DataType::FP32);
816        let w = make_tensor(vec![784, 64], DataType::FP32);
817        let b = make_tensor(vec![64], DataType::FP32);
818        let result =
819            ShapeInference::infer_output_shape(&TensorOp::Linear, &[&x, &w, &b], None).unwrap();
820        assert_eq!(result[0].shape[0].static_value(), Some(1));
821        assert_eq!(result[0].shape[1].static_value(), Some(64));
822    }
823
824    #[test]
825    fn test_conv2d_shape() {
826        let input = make_tensor(vec![1, 3, 28, 28], DataType::FP32);
827        let kernel = make_tensor(vec![16, 3, 5, 5], DataType::FP32);
828        let result =
829            ShapeInference::infer_output_shape(&TensorOp::Conv2D, &[&input, &kernel], None)
830                .unwrap();
831        assert_eq!(result[0].shape[0].static_value(), Some(1));
832        assert_eq!(result[0].shape[1].static_value(), Some(16));
833        assert_eq!(result[0].shape[2].static_value(), Some(24)); // 28-5+1
834        assert_eq!(result[0].shape[3].static_value(), Some(24));
835    }
836
837    /// Regression test: Conv2D used to ignore stride/padding/dilation
838    /// entirely (no attrs were even passed in), always computing
839    /// `in - kernel + 1` regardless of what the op actually specified.
840    #[test]
841    fn test_conv2d_shape_honours_stride_padding_dilation() {
842        let input = make_tensor(vec![1, 3, 28, 28], DataType::FP32);
843        let kernel = make_tensor(vec![16, 3, 3, 3], DataType::FP32);
844
845        // stride=2, padding=1, dilation=1: out = floor((28+2-2-1)/2)+1 = 14
846        let mut attrs = Attributes::new();
847        attrs.set("stride", Attribute::Integer(2));
848        attrs.set("padding", Attribute::Integer(1));
849        let result =
850            ShapeInference::infer_output_shape(&TensorOp::Conv2D, &[&input, &kernel], Some(&attrs))
851                .unwrap();
852        assert_eq!(result[0].shape[2].static_value(), Some(14));
853        assert_eq!(result[0].shape[3].static_value(), Some(14));
854
855        let flops =
856            ShapeInference::compute_flops(&TensorOp::Conv2D, &[&input, &kernel], Some(&attrs))
857                .unwrap();
858        assert_eq!(flops, 2 * 16 * 3 * 3 * 3 * 14 * 14);
859    }
860
861    /// Regression test: DilatedConv2D used to compute the exact same shape
862    /// as a plain Conv2D, silently ignoring dilation. With no attrs given it
863    /// now defaults to dilation=2 (the point of the op), producing a
864    /// different, smaller output than Conv2D would for the same input.
865    #[test]
866    fn test_dilated_conv2d_differs_from_plain_conv2d_by_default() {
867        let input = make_tensor(vec![1, 3, 28, 28], DataType::FP32);
868        let kernel = make_tensor(vec![16, 3, 3, 3], DataType::FP32);
869
870        let plain = ShapeInference::infer_output_shape(&TensorOp::Conv2D, &[&input, &kernel], None)
871            .unwrap();
872        let dilated =
873            ShapeInference::infer_output_shape(&TensorOp::DilatedConv2D, &[&input, &kernel], None)
874                .unwrap();
875
876        assert_eq!(plain[0].shape[2].static_value(), Some(26)); // 28-3+1
877                                                                // dilation=2: out = floor((28 - 2*(3-1) - 1)/1)+1 = 24
878        assert_eq!(dilated[0].shape[2].static_value(), Some(24));
879        assert_ne!(
880            plain[0].shape[2], dilated[0].shape[2],
881            "DilatedConv2D must not silently behave like Conv2D"
882        );
883    }
884
885    /// Regression test: MaxPool2D/AvgPool2D used to return the input shape
886    /// unchanged ("simplified"), never reducing spatial dims at all.
887    #[test]
888    fn test_maxpool2d_reduces_spatial_dims() {
889        let input = make_tensor(vec![1, 64, 32, 32], DataType::FP32);
890        let kernel = make_tensor(vec![2, 2], DataType::FP32);
891        let result =
892            ShapeInference::infer_output_shape(&TensorOp::MaxPool2D, &[&input, &kernel], None)
893                .unwrap();
894        assert_eq!(result[0].shape[0].static_value(), Some(1));
895        assert_eq!(result[0].shape[1].static_value(), Some(64));
896        // Default stride = kernel size (non-overlapping): 32/2 = 16.
897        assert_eq!(result[0].shape[2].static_value(), Some(16));
898        assert_eq!(result[0].shape[3].static_value(), Some(16));
899    }
900
901    /// Regression test: compute_memory_bytes used to add the output's bytes
902    /// only for MatMul/SparseMatMul — every other op (Conv2D here) silently
903    /// omitted the output from the memory-traffic total.
904    #[test]
905    fn test_compute_memory_bytes_includes_output_for_every_op() {
906        let input = make_tensor(vec![1, 3, 28, 28], DataType::FP32); // 2352 elems
907        let kernel = make_tensor(vec![16, 3, 5, 5], DataType::FP32); // 1200 elems
908        let mem = ShapeInference::compute_memory_bytes(&TensorOp::Conv2D, &[&input, &kernel], None)
909            .unwrap();
910        // input + kernel + output(1,16,24,24) elements, all FP32 (4 bytes/elem).
911        let expected = (2352 + 1200 + 16 * 24 * 24) * 4;
912        assert_eq!(mem, expected as u64);
913    }
914}
lift_tensor/shape.rs

lift_tensor/
shape.rs
