From eeacf8e2841315d7e8cec8f072ca8d39d6165abb Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Sun, 2 Aug 2026 17:54:01 +0100 Subject: [PATCH 1/4] core,onnx: only lower a nearest Resize to Tile when it replicates pixels The nearest-neighbour Resize declutter lowered any integer scale to Reshape -> Tile -> Reshape, but that replication pattern only matches some coordinate transform and tie-break pairs, so half_pixel or ceil rounding silently produced shifted output. Both declutters now probe the rounding of each upsampled axis and lower only when it really is pixel replication. --- core/src/ops/nn/resize.rs | 52 +++++++++++++++++++++++++++++++++++++ onnx-opl/src/resize.rs | 12 ++++++++- test-rt/suite-onnx/node.txt | 1 + 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/core/src/ops/nn/resize.rs b/core/src/ops/nn/resize.rs index cb21b68121..8f4d98426c 100644 --- a/core/src/ops/nn/resize.rs +++ b/core/src/ops/nn/resize.rs @@ -314,11 +314,50 @@ impl TypedOp for Resize { scales.iter().zip(&int_scales).all(|(&s, &i)| (s - i as f32).abs() <= 1e-5 && i != 0) ); rule_if!(int_scales.iter().any(|&s| s != 1)); + let input_shape = &model.outlet_fact(node.inputs[0])?.shape; + for (axis, &scale) in int_scales.iter().enumerate().filter(|&(_, &s)| s > 1) { + let Some(len_in) = probe_length(&self.coord_transformer, &input_shape[axis]) else { + return Ok(None); + }; + rule_if!(is_pixel_replication(&self.coord_transformer, len_in, scale, |frac| self + .nearest + .prefers_right(frac))); + } lower_nearest_integer_upsample(model, node, &int_scales) } } +/// Whether nearest-neighbour resampling by `scale` reads input `x / scale` for +/// every output `x`, which is the pattern [`lower_nearest_integer_upsample`] +/// produces. Only some coordinate transform and tie-break pairs round that way, +/// so both Resize declutters must check before lowering. +pub fn is_pixel_replication( + coord_transformer: &CoordTransformer, + len_in: usize, + scale: usize, + prefers_right: impl Fn(f32) -> bool, +) -> bool { + let len_out = len_in * scale; + (0..len_out).all(|x| { + let x_in = coord_transformer.transform(x, scale as f32, len_in, len_out); + let x_floor = x_in.floor() as isize; + let picked = + (x_floor + prefers_right(x_in - x_floor as f32) as isize).clamp(0, len_in as isize - 1); + picked == (x / scale) as isize + }) +} + +/// An axis length to probe [`is_pixel_replication`] on. `HalfPixel` and +/// `Asymmetric` map coordinates without consulting the axis lengths, so a +/// symbolic axis can still be probed on a stand-in; the others cannot. +pub fn probe_length(coord_transformer: &CoordTransformer, len: &TDim) -> Option { + len.to_usize().ok().or(match coord_transformer { + CoordTransformer::HalfPixel | CoordTransformer::Asymmetric => Some(4), + _ => None, + }) +} + /// Lowers a nearest-neighbour integer upsample to Reshape → Tile → Reshape: each /// upsampled axis is split into a size-1 axis, tiled by its scale, then merged /// back. Shared by the core and ONNX Resize declutters. @@ -446,4 +485,17 @@ mod tests { ); assert_eq!(out.shape(), &[4, 4]); } + + fn replicates(coord_transformer: CoordTransformer, nearest: Nearest, scale: usize) -> bool { + is_pixel_replication(&coord_transformer, 4, scale, |frac| nearest.prefers_right(frac)) + } + + #[test] + fn only_some_nearest_modes_replicate_pixels() { + assert!(replicates(CoordTransformer::Asymmetric, Nearest::Floor, 2)); + assert!(replicates(CoordTransformer::HalfPixel, Nearest::RoundPreferCeil, 2)); + assert!(replicates(CoordTransformer::HalfPixel, Nearest::RoundPreferCeil, 3)); + assert!(!replicates(CoordTransformer::HalfPixel, Nearest::Floor, 2)); + assert!(!replicates(CoordTransformer::Asymmetric, Nearest::RoundPreferCeil, 2)); + } } diff --git a/onnx-opl/src/resize.rs b/onnx-opl/src/resize.rs index df1fc29552..84c178ad8e 100644 --- a/onnx-opl/src/resize.rs +++ b/onnx-opl/src/resize.rs @@ -1,6 +1,7 @@ use tract_nnef::internal::*; use tract_nnef::tract_core::ops::nn::resize::{ - self, CoordTransformer, Interpolator, cubic_kernel, lower_nearest_integer_upsample, + self, CoordTransformer, Interpolator, cubic_kernel, is_pixel_replication, + lower_nearest_integer_upsample, probe_length, }; /// Nearest-neighbour tie-breaking, the full ONNX set. `Floor` and @@ -309,6 +310,15 @@ impl TypedOp for Resize { scales.iter().zip(&int_scales).all(|(&s, &i)| (s - i as f32).abs() <= 1e-5 && i != 0) ); rule_if!(int_scales.iter().any(|&s| s != 1)); + let input_shape = &model.outlet_fact(node.inputs[0])?.shape; + for (axis, &scale) in int_scales.iter().enumerate().filter(|&(_, &s)| s > 1) { + let Some(len_in) = probe_length(&self.coord_transformer, &input_shape[axis]) else { + return Ok(None); + }; + rule_if!(is_pixel_replication(&self.coord_transformer, len_in, scale, |frac| self + .nearest + .prefers_right(frac))); + } lower_nearest_integer_upsample(model, node, &int_scales) } diff --git a/test-rt/suite-onnx/node.txt b/test-rt/suite-onnx/node.txt index b12694c7c2..08cf225252 100644 --- a/test-rt/suite-onnx/node.txt +++ b/test-rt/suite-onnx/node.txt @@ -500,6 +500,7 @@ test_resize_upsample_scales_cubic_asymmetric test_resize_upsample_scales_cubic_A_n0p5_exclude_outside input:X test_resize_upsample_scales_linear_align_corners input:X not-nnef test_resize_upsample_sizes_cubic input:X +test_resize_upsample_sizes_nearest_ceil_half_pixel input:X test_rnn_seq_length test_round test_scan9_sum From 8ee1b2faffe9d070d7a0a135b94d6be7bcf3fa3f Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Sun, 2 Aug 2026 18:03:39 +0100 Subject: [PATCH 2/4] core,onnx: complete Resize for opset 18 and 19 Resize silently ignored antialias and keep_aspect_ratio_policy, carried axes as a dead field, and rejected half_pixel_symmetric and tf_crop_and_resize, so opset-18 and -19 models were either wrong or refused outright. The op now implements all of them, and resamples through a per-axis plan of precomputed taps and weights applied over contiguous runs rather than recomputing a dynamic-rank index for every output element. --- core/src/ops/nn/resize.rs | 301 +++++++++++++++------ onnx-opl/src/resize.rs | 523 +++++++++++++++++++++++++----------- onnx/src/ops/resize.rs | 65 +++-- test-rt/suite-onnx/node.txt | 14 +- 4 files changed, 614 insertions(+), 289 deletions(-) diff --git a/core/src/ops/nn/resize.rs b/core/src/ops/nn/resize.rs index 8f4d98426c..5c3f9f0f93 100644 --- a/core/src/ops/nn/resize.rs +++ b/core/src/ops/nn/resize.rs @@ -1,7 +1,7 @@ use crate::internal::*; use crate::ops::array::Tile; -/// Maps an output coordinate back to the input axis. The four ONNX coordinate +/// Maps an output coordinate back to the input axis. The ONNX coordinate /// transformation modes that have a well-defined inverse without an input ROI. #[derive(Clone, Debug, Hash, PartialEq, Eq)] pub enum CoordTransformer { @@ -9,6 +9,7 @@ pub enum CoordTransformer { AlignCorners, Asymmetric, PytorchHalfPixel, + HalfPixelSymmetric, } impl CoordTransformer { @@ -28,9 +29,14 @@ impl CoordTransformer { if len_out > 1 { (x_out as f32 + 0.5) / scale - 0.5 } else { - 0.0 + -0.5 } } + CoordTransformer::HalfPixelSymmetric => { + let adjustment = len_out as f32 / (scale * len_in as f32); + let offset = len_in as f32 / 2.0 * (1.0 - adjustment); + offset + (x_out as f32 + 0.5) / scale - 0.5 + } } } @@ -40,6 +46,7 @@ impl CoordTransformer { CoordTransformer::AlignCorners => "align_corners", CoordTransformer::Asymmetric => "asymmetric", CoordTransformer::PytorchHalfPixel => "pytorch_half_pixel", + CoordTransformer::HalfPixelSymmetric => "half_pixel_symmetric", } } @@ -49,6 +56,7 @@ impl CoordTransformer { "align_corners" => CoordTransformer::AlignCorners, "asymmetric" => CoordTransformer::Asymmetric, "pytorch_half_pixel" => CoordTransformer::PytorchHalfPixel, + "half_pixel_symmetric" => CoordTransformer::HalfPixelSymmetric, s => bail!("coordinate_transformation_mode: {s}"), }) } @@ -82,6 +90,47 @@ impl Interpolator { } } +/// Number of input taps feeding one output element. Antialiasing widens the +/// footprint of `Linear` and `Cubic` by `1/scale` when downscaling, which is +/// what turns the interpolation into a low-pass filter. +pub fn window_size(interpolator: &Interpolator, antialias: bool, scale: f32) -> usize { + let support = match interpolator { + Interpolator::Nearest | Interpolator::Linear => 1.0f32, + Interpolator::Cubic => 2.0, + }; + if !antialias || scale >= 1.0 || matches!(interpolator, Interpolator::Nearest) { + return 2 * support as usize; + } + let first = (-support / scale).floor() as isize + 1; + (2 - 2 * first) as usize +} + +/// Triangular kernel, low-passed by `scale` when antialiasing. +pub fn linear_weights(r: f32, scale: f32, antialias: bool, weights: &mut [f32]) { + let scale = if antialias { scale.min(1.0) } else { 1.0 }; + fill_weights(r, scale, weights, |x| (1.0 - x.abs()).clamp(0.0, 1.0)); +} + +/// Cubic convolution kernel of coefficient `a`, low-passed by `scale` when +/// antialiasing. +pub fn cubic_weights(r: f32, scale: f32, a: f32, antialias: bool, weights: &mut [f32]) { + let scale = if antialias { scale.min(1.0) } else { 1.0 }; + fill_weights(r, scale, weights, |x| cubic_kernel(x, a)); +} + +/// Evaluates `kernel` at the offset of each tap in the window from the source +/// coordinate, then renormalizes when the kernel was stretched by `scale`. +fn fill_weights(r: f32, scale: f32, weights: &mut [f32], kernel: impl Fn(f32) -> f32) { + let first = 1.0 - (weights.len() / 2) as f32; + for (k, w) in weights.iter_mut().enumerate() { + *w = kernel((first + k as f32 - r) * scale); + } + if scale != 1.0 { + let sum: f32 = weights.iter().sum(); + weights.iter_mut().for_each(|w| *w /= sum); + } +} + /// Standard Catmull-Rom-family cubic convolution kernel with coefficient `a`. pub fn cubic_kernel(s: f32, a: f32) -> f32 { let abs_s = s.abs(); @@ -128,6 +177,122 @@ impl Nearest { } } +/// Resampling plan for one axis: for each output index, the `window` input taps +/// it reads (already clamped into the axis) and their weights. `extrapolated` +/// flags output indices that map outside the input, which are filled with a +/// constant instead of being interpolated. +#[derive(Clone, Debug)] +pub struct AxisPlan { + pub window: usize, + pub indices: Vec, + pub weights: Vec, + pub extrapolated: Vec, +} + +/// Builds the resampling plan for one axis. `coord` maps an output index to a +/// source coordinate, or `None` when it falls outside the input. `weights` +/// fills a window for a fractional offset in `(0, 1]`, the tap at index `k` +/// sitting `k + 1 - window / 2` cells right of the source cell. With +/// `exclude_outside` the taps falling off the axis are dropped and the +/// remaining weights renormalized; otherwise they read the edge value. +pub fn plan_axis( + len_in: usize, + len_out: usize, + window: usize, + exclude_outside: bool, + coord: impl Fn(usize) -> Option, + weights: impl Fn(f32, &mut [f32]), +) -> AxisPlan { + let mut plan = AxisPlan { + window, + indices: vec![0; window * len_out], + weights: vec![0.0; window * len_out], + extrapolated: vec![false; len_out], + }; + for x in 0..len_out { + let Some(x_in) = coord(x) else { + plan.extrapolated[x] = true; + continue; + }; + let cell = x_in.ceil() - 1.0; + let taps = &mut plan.weights[x * window..][..window]; + weights(x_in - cell, taps); + let first = cell as isize + 1 - (window / 2) as isize; + for (k, tap) in taps.iter_mut().enumerate() { + let raw = first + k as isize; + if exclude_outside && (raw < 0 || raw >= len_in as isize) { + *tap = 0.0; + } + plan.indices[x * window + k] = raw.clamp(0, len_in as isize - 1) as usize; + } + if exclude_outside { + let sum: f32 = taps.iter().sum(); + if sum != 0.0 { + taps.iter_mut().for_each(|w| *w /= sum); + } + } + } + plan +} + +/// Whether `plan` is exactly pixel replication — every output reading the +/// single input `x / scale` — which is the pattern +/// [`lower_nearest_integer_upsample`] produces. Only some coordinate transform +/// and tie-break pairs round that way, so the plan itself is the arbiter. +pub fn is_pixel_replication(plan: &AxisPlan, scale: usize) -> bool { + !plan.extrapolated.contains(&true) + && plan + .indices + .chunks_exact(plan.window) + .zip(plan.weights.chunks_exact(plan.window)) + .enumerate() + .all(|(x, (indices, weights))| { + let mut taps = indices.iter().zip(weights).filter(|(_, w)| **w != 0.0); + taps.next().is_some_and(|(i, w)| *w == 1.0 && *i == x / scale) + && taps.next().is_none() + }) +} + +/// Resamples `axis` of a row-major `input` of shape `shape` into `output`, +/// which must be sized for the same shape with `axis` set to the plan length. +pub fn resample_axis( + input: &[f32], + shape: &[usize], + axis: usize, + plan: &AxisPlan, + extrapolation_value: f32, + output: &mut [f32], +) { + let len_in = shape[axis]; + let len_out = plan.extrapolated.len(); + let inner: usize = shape[axis + 1..].iter().product(); + let window = plan.window; + if len_in * inner == 0 || len_out * inner == 0 { + return; + } + for (src, dst) in + input.chunks_exact(len_in * inner).zip(output.chunks_exact_mut(len_out * inner)) + { + for (x, dst) in dst.chunks_exact_mut(inner).enumerate() { + if plan.extrapolated[x] { + dst.fill(extrapolation_value); + continue; + } + dst.fill(0.0); + let indices = &plan.indices[x * window..][..window]; + let weights = &plan.weights[x * window..][..window]; + for (&i, &w) in indices.iter().zip(weights) { + if w == 0.0 { + continue; + } + for (d, s) in dst.iter_mut().zip(&src[i * inner..][..inner]) { + *d += w * s; + } + } + } + } +} + /// Resamples `input` along the axes given by `scales`/`sizes`, the clean subset /// of ONNX Resize: `interpolator` × `coord_transformer` × `nearest`, fixed /// `cubic_coeff_a = -0.75`, no ROI and no `exclude_outside`. The ONNX op carries @@ -186,6 +351,24 @@ impl Resize { input_sizes, ); } + + fn plan_axis(&self, scale: f32, len_in: usize, len_out: usize) -> AxisPlan { + let window = window_size(&self.interpolator, false, scale); + let coord = |x| Some(self.coord_transformer.transform(x, scale, len_in, len_out)); + match self.interpolator { + Interpolator::Linear => plan_axis(len_in, len_out, window, false, coord, |r, w| { + linear_weights(r, scale, false, w) + }), + Interpolator::Cubic => plan_axis(len_in, len_out, window, false, coord, |r, w| { + cubic_weights(r, scale, -0.75, false, w) + }), + Interpolator::Nearest => plan_axis(len_in, len_out, window, false, coord, |r, w| { + let right = r == 1.0 || self.nearest.prefers_right(r); + w[0] = !right as u8 as f32; + w[1] = right as u8 as f32; + }), + } + } } impl Op for Resize { @@ -217,67 +400,21 @@ impl EvalOp for Resize { output_shape.iter().zip(inputs[0].shape()).map(|(o, i)| *o as f32 / *i as f32).collect() }; let input = inputs.remove(0).into_tensor(); - let mut data = if input.datum_type() == f32::datum_type() { - input.into_plain_array::()? - } else { - input.cast_to::()?.into_owned().into_plain_array::()? - }; - for (axis, scale) in scales.into_iter().enumerate().filter(|(_, s)| *s != 1.0) { - let mut new_shape: TVec = data.shape().into(); - new_shape[axis] = output_shape[axis]; - let input_len = data.shape()[axis]; - data = match self.interpolator { - Interpolator::Cubic => { - let a = -0.75f32; - tract_ndarray::ArrayD::from_shape_fn(&*new_shape, |co_o| -> f32 { - let x_out = co_o[axis]; - let x_in = self.coord_transformer.transform( - x_out, - scale, - input_len, - new_shape[axis], - ); - let x_floor = x_in.floor() as isize; - let t = x_in - x_floor as f32; - let mut co_i = co_o; - let mut acc = 0.0f32; - for j in -1..=2isize { - let w = cubic_kernel(t - j as f32, a); - let idx = (x_floor + j).clamp(0, input_len as isize - 1) as usize; - co_i[axis] = idx; - acc += w * data[&co_i]; - } - acc - }) - } - _ => tract_ndarray::ArrayD::from_shape_fn(&*new_shape, |co_o| -> f32 { - let x_out = co_o[axis]; - let x_in = - self.coord_transformer.transform(x_out, scale, input_len, new_shape[axis]); - let mut co_i = co_o; - let x_floor = x_in.floor() as isize; - let x_left = x_floor.clamp(0, input_len as isize - 1) as usize; - co_i[axis] = x_left; - let y_left = data[&co_i]; - let x_right = (x_floor + 1).clamp(0, input_len as isize - 1) as usize; - co_i[axis] = x_right; - let y_right = data[&co_i]; - let x_frac = x_in - x_floor as f32; - match self.interpolator { - Interpolator::Linear => y_left * (1.0 - x_frac) + y_right * x_frac, - Interpolator::Nearest => { - if self.nearest.prefers_right(x_frac) { - y_right - } else { - y_left - } - } - Interpolator::Cubic => unreachable!(), - } - }), + let input = input.cast_to::()?; + let mut shape: TVec = input.shape().into(); + let mut data: Vec = input.try_as_plain()?.as_slice::()?.to_vec(); + for (axis, scale) in scales.into_iter().enumerate() { + let (len_in, len_out) = (shape[axis], output_shape[axis]); + if len_in == len_out && scale == 1.0 { + continue; } + let plan = self.plan_axis(scale, len_in, len_out); + let mut resampled = vec![0f32; data.len() / len_in * len_out]; + resample_axis(&data, &shape, axis, &plan, 0.0, &mut resampled); + data = resampled; + shape[axis] = len_out; } - let out = data.into_tensor(); + let out = tract_ndarray::ArrayD::from_shape_vec(&*shape, data)?.into_tensor(); let out = if out.datum_type() == input_dt { out } else { out.cast_to_dt(input_dt)?.into_owned() }; Ok(tvec!(out.into_tvalue())) @@ -319,38 +456,19 @@ impl TypedOp for Resize { let Some(len_in) = probe_length(&self.coord_transformer, &input_shape[axis]) else { return Ok(None); }; - rule_if!(is_pixel_replication(&self.coord_transformer, len_in, scale, |frac| self - .nearest - .prefers_right(frac))); + rule_if!(is_pixel_replication( + &self.plan_axis(scale as f32, len_in, len_in * scale), + scale + )); } lower_nearest_integer_upsample(model, node, &int_scales) } } -/// Whether nearest-neighbour resampling by `scale` reads input `x / scale` for -/// every output `x`, which is the pattern [`lower_nearest_integer_upsample`] -/// produces. Only some coordinate transform and tie-break pairs round that way, -/// so both Resize declutters must check before lowering. -pub fn is_pixel_replication( - coord_transformer: &CoordTransformer, - len_in: usize, - scale: usize, - prefers_right: impl Fn(f32) -> bool, -) -> bool { - let len_out = len_in * scale; - (0..len_out).all(|x| { - let x_in = coord_transformer.transform(x, scale as f32, len_in, len_out); - let x_floor = x_in.floor() as isize; - let picked = - (x_floor + prefers_right(x_in - x_floor as f32) as isize).clamp(0, len_in as isize - 1); - picked == (x / scale) as isize - }) -} - -/// An axis length to probe [`is_pixel_replication`] on. `HalfPixel` and -/// `Asymmetric` map coordinates without consulting the axis lengths, so a -/// symbolic axis can still be probed on a stand-in; the others cannot. +/// An axis length to build a probe plan on. `HalfPixel` and `Asymmetric` map +/// coordinates without consulting the axis lengths, so a symbolic axis can +/// still be probed on a stand-in; the others cannot. pub fn probe_length(coord_transformer: &CoordTransformer, len: &TDim) -> Option { len.to_usize().ok().or(match coord_transformer { CoordTransformer::HalfPixel | CoordTransformer::Asymmetric => Some(4), @@ -487,7 +605,14 @@ mod tests { } fn replicates(coord_transformer: CoordTransformer, nearest: Nearest, scale: usize) -> bool { - is_pixel_replication(&coord_transformer, 4, scale, |frac| nearest.prefers_right(frac)) + let op = Resize { + coord_transformer, + interpolator: Interpolator::Nearest, + nearest, + optional_scales_input: Some(1), + optional_sizes_input: None, + }; + is_pixel_replication(&op.plan_axis(scale as f32, 4, 4 * scale), scale) } #[test] diff --git a/onnx-opl/src/resize.rs b/onnx-opl/src/resize.rs index 84c178ad8e..743fe190cc 100644 --- a/onnx-opl/src/resize.rs +++ b/onnx-opl/src/resize.rs @@ -1,7 +1,8 @@ use tract_nnef::internal::*; use tract_nnef::tract_core::ops::nn::resize::{ - self, CoordTransformer, Interpolator, cubic_kernel, is_pixel_replication, - lower_nearest_integer_upsample, probe_length, + self, AxisPlan, CoordTransformer, Interpolator, cubic_weights, is_pixel_replication, + linear_weights, lower_nearest_integer_upsample, plan_axis, probe_length, resample_axis, + window_size, }; /// Nearest-neighbour tie-breaking, the full ONNX set. `Floor` and @@ -45,14 +46,70 @@ impl Nearest { } } +/// ONNX `coordinate_transformation_mode`. Every mode but `tf_crop_and_resize` +/// inverts without an input ROI and is shared with tract-core. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub enum CoordTransform { + Plain(CoordTransformer), + TfCropAndResize, +} + +impl CoordTransform { + pub fn as_str(&self) -> &'static str { + match self { + CoordTransform::Plain(t) => t.as_str(), + CoordTransform::TfCropAndResize => "tf_crop_and_resize", + } + } + + pub fn parse(s: &str) -> TractResult { + Ok(match s { + "tf_crop_and_resize" => CoordTransform::TfCropAndResize, + s => CoordTransform::Plain(CoordTransformer::parse(s)?), + }) + } +} + +/// ONNX `keep_aspect_ratio_policy`: reconciles the requested `sizes` into a +/// single scale shared by every resized axis. Ignored when `scales` drives the +/// resize. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] +pub enum AspectRatio { + Stretch, + NotLarger, + NotSmaller, +} + +impl AspectRatio { + pub fn as_str(&self) -> &'static str { + match self { + AspectRatio::Stretch => "stretch", + AspectRatio::NotLarger => "not_larger", + AspectRatio::NotSmaller => "not_smaller", + } + } + + pub fn parse(s: &str) -> TractResult { + Ok(match s { + "stretch" => AspectRatio::Stretch, + "not_larger" => AspectRatio::NotLarger, + "not_smaller" => AspectRatio::NotSmaller, + s => bail!("keep_aspect_ratio_policy: {s}"), + }) + } +} + #[derive(Clone, Debug, Hash, PartialEq, Eq)] pub struct Resize { pub axes: Option>, - pub coord_transformer: CoordTransformer, + pub coord_transformer: CoordTransform, pub interpolator: Interpolator, pub nearest: Nearest, + pub antialias: bool, pub cubic_coeff_a_bits: u32, pub exclude_outside: bool, + pub extrapolation_value_bits: u32, + pub keep_aspect_ratio_policy: AspectRatio, pub optional_roi_input: Option, pub optional_scales_input: Option, pub optional_sizes_input: Option, @@ -63,42 +120,62 @@ impl Resize { f32::from_bits(self.cubic_coeff_a_bits) } + pub fn extrapolation_value(&self) -> f32 { + f32::from_bits(self.extrapolation_value_bits) + } + + /// The axes `scales`, `sizes` and `roi` describe, negatives resolved. Axes + /// left out keep their input length. + pub fn resized_axes(&self, rank: usize) -> TractResult> { + let Some(axes) = &self.axes else { return Ok((0..rank).collect()) }; + axes.iter() + .map(|a| { + let a = if *a < 0 { a + rank as i64 } else { *a }; + ensure!((0..rank as i64).contains(&a), "Resize axes {axes:?} out of rank {rank}"); + Ok(a as usize) + }) + .collect() + } + pub fn compute_output_shape( &self, input_shape: &[D], input_scale: Option<&Tensor>, input_sizes: Option<&Tensor>, ) -> TractResult> { - if let Some(scale) = input_scale - && scale.len() == input_shape.len() - { - let mut shape = tvec!(); - for (i, s) in input_shape - .iter() - .zip(scale.cast_to::()?.try_as_plain()?.as_slice::()?.iter()) - { - if s.round() == *s { - shape.push(i.clone() * (*s as usize)); + let axes = self.resized_axes(input_shape.len())?; + let mut shape: TVec = input_shape.into(); + if let Some(scale) = input_scale.filter(|s| s.len() == axes.len()) { + let scale = scale.cast_to::()?; + for (&axis, s) in axes.iter().zip(scale.try_as_plain()?.as_slice::()?) { + let i = &input_shape[axis]; + shape[axis] = if s.round() == *s { + i.clone() * (*s as usize) } else if let Ok(i) = i.to_usize() { - shape.push(((i as f32 * s) as usize).into()); + ((i as f32 * s) as usize).into() } else { bail!( "Can not compute output shape. inputs are {input_shape:?} and scale {scale:?}" ) - } + }; } return Ok(shape); } - if let Some(sizes) = input_sizes - && sizes.len() == input_shape.len() - { - return sizes - .cast_to::()? - .try_as_plain()? - .as_slice::()? - .iter() - .map(|i| i.try_into()) - .collect(); + if let Some(sizes) = input_sizes.filter(|s| s.len() == axes.len()) { + let sizes = sizes.cast_to::()?; + let sizes = sizes.try_as_plain()?.as_slice::()?; + if self.keep_aspect_ratio_policy == AspectRatio::Stretch { + for (&axis, s) in axes.iter().zip(sizes) { + shape[axis] = s.try_into()?; + } + } else { + let scale = self.aspect_ratio_scale(input_shape, &axes, sizes)?; + for &axis in &axes { + let len = input_shape[axis].to_usize()?; + shape[axis] = ((scale * len as f32 + 0.5) as usize).into(); + } + } + return Ok(shape); } bail!( "Neither sizes nor scales makes sense: input_shape: {:?}, scale: {:?}, sizes: {:?}", @@ -108,13 +185,122 @@ impl Resize { ); } + /// The scale every resized axis takes under a non-`Stretch` policy: the + /// smallest requested ratio to stay within `sizes`, the largest to cover it. + fn aspect_ratio_scale( + &self, + input_shape: &[D], + axes: &[usize], + sizes: &[TDim], + ) -> TractResult { + let mut ratios: TVec = tvec!(); + for (&axis, size) in axes.iter().zip(sizes) { + ratios.push(size.to_usize()? as f32 / input_shape[axis].to_usize()? as f32); + } + let pick = if self.keep_aspect_ratio_policy == AspectRatio::NotLarger { + f32::min + } else { + f32::max + }; + ratios.into_iter().reduce(pick).context("Resize sizes must not be empty") + } + + /// Per-axis scale and output length over the full rank. + fn resolve( + &self, + input_shape: &[usize], + scales: Option<&Tensor>, + sizes: Option<&Tensor>, + ) -> TractResult<(TVec, TVec)> { + let axes = self.resized_axes(input_shape.len())?; + let output_shape = self.compute_output_shape(input_shape, scales, sizes)?; + let mut per_axis: TVec = tvec!(1.0; input_shape.len()); + if let Some(scales) = scales.filter(|s| s.len() == axes.len()) { + let scales = scales.cast_to::()?; + for (&axis, s) in axes.iter().zip(scales.try_as_plain()?.as_slice::()?) { + per_axis[axis] = *s; + } + } else if self.keep_aspect_ratio_policy != AspectRatio::Stretch { + let sizes = + sizes.context("Resize aspect ratio policy needs sizes")?.cast_to::()?; + let scale = + self.aspect_ratio_scale(input_shape, &axes, sizes.try_as_plain()?.as_slice()?)?; + for &axis in &axes { + per_axis[axis] = scale; + } + } else { + for &axis in &axes { + per_axis[axis] = output_shape[axis] as f32 / input_shape[axis] as f32; + } + } + Ok((per_axis, output_shape)) + } + + /// Normalized ROI `(start, end)` per axis, `(0, 1)` for the axes left out. + fn roi(&self, rank: usize, roi: Option<&Tensor>) -> TractResult> { + let axes = self.resized_axes(rank)?; + let Some(roi) = roi.filter(|r| r.len() == 2 * axes.len()) else { + bail!("Resize in tf_crop_and_resize mode needs a roi of 2 x {} elements", axes.len()) + }; + let roi = roi.cast_to::()?; + let roi = roi.try_as_plain()?.as_slice::()?; + let mut per_axis: TVec<(f32, f32)> = tvec!((0.0, 1.0); rank); + for (i, &axis) in axes.iter().enumerate() { + per_axis[axis] = (roi[i], roi[i + axes.len()]); + } + Ok(per_axis) + } + + fn plan_axis(&self, scale: f32, len_in: usize, len_out: usize, roi: (f32, f32)) -> AxisPlan { + let window = window_size(&self.interpolator, self.antialias, scale); + let coord: Box Option> = match &self.coord_transformer { + CoordTransform::Plain(t) => { + let t = t.clone(); + Box::new(move |x| Some(t.transform(x, scale, len_in, len_out))) + } + CoordTransform::TfCropAndResize => { + let last = len_in as f32 - 1.0; + let span = last * (roi.1 - roi.0); + let width = scale * len_in as f32; + Box::new(move |x| { + let offset = + if width == 1.0 { span / 2.0 } else { x as f32 * span / (width - 1.0) }; + let x = offset + roi.0 * last; + (x >= 0.0 && x <= last).then_some(x) + }) + } + }; + let exclude = self.exclude_outside; + let (antialias, a) = (self.antialias, self.cubic_coeff_a()); + match self.interpolator { + Interpolator::Linear => plan_axis(len_in, len_out, window, exclude, coord, |r, w| { + linear_weights(r, scale, antialias, w) + }), + Interpolator::Cubic => plan_axis(len_in, len_out, window, exclude, coord, |r, w| { + cubic_weights(r, scale, a, antialias, w) + }), + Interpolator::Nearest => plan_axis(len_in, len_out, window, exclude, coord, |r, w| { + let right = r == 1.0 || self.nearest.prefers_right(r); + w[0] = !right as u8 as f32; + w[1] = right as u8 as f32; + }), + } + } + /// The clean subset reachable by `tract_core::ops::nn::resize::Resize`: - /// default `cubic_coeff_a`, no `exclude_outside`, no ROI and a nearest mode - /// core understands. `None` keeps the op as an ONNX edge-case op. + /// default `cubic_coeff_a`, no `exclude_outside`, no antialiasing, no ROI, + /// a stretching aspect ratio and a nearest mode core understands. `None` + /// keeps the op as an ONNX edge-case op. fn as_core(&self) -> Option { - if self.exclude_outside || self.optional_roi_input.is_some() { + if self.exclude_outside + || self.antialias + || self.keep_aspect_ratio_policy != AspectRatio::Stretch + { return None; } + let CoordTransform::Plain(coord_transformer) = &self.coord_transformer else { + return None; + }; if self.interpolator == Interpolator::Cubic && self.cubic_coeff_a() != -0.75 { return None; } @@ -129,13 +315,42 @@ impl Resize { _ => return None, }; Some(resize::Resize { - coord_transformer: self.coord_transformer.clone(), + coord_transformer: coord_transformer.clone(), interpolator: self.interpolator.clone(), nearest, optional_scales_input: Some(1), optional_sizes_input: None, }) } + + /// Spreads a per-`axes` `scales`/`sizes` constant over the full rank, so the + /// core op — which carries no `axes` of its own — can take the node over. + fn full_rank_aux( + &self, + input_shape: &ShapeFact, + axes: &[usize], + konst: &Tensor, + sizes: bool, + ) -> TractResult> { + if sizes { + let mut full: Vec = vec![0; input_shape.rank()]; + for (slot, dim) in full.iter_mut().zip(input_shape.iter()) { + *slot = dim.to_usize()? as i64; + } + let konst = konst.cast_to::()?; + for (&axis, v) in axes.iter().zip(konst.try_as_plain()?.as_slice::()?) { + full[axis] = *v; + } + Ok(tract_ndarray::arr1(&full).into_arc_tensor()) + } else { + let mut full = vec![1.0f32; input_shape.rank()]; + let konst = konst.cast_to::()?; + for (&axis, v) in axes.iter().zip(konst.try_as_plain()?.as_slice::()?) { + full[axis] = *v; + } + Ok(tract_ndarray::arr1(&full).into_arc_tensor()) + } + } } impl Op for Resize { @@ -153,97 +368,34 @@ impl EvalOp for Resize { fn eval(&self, mut inputs: TVec) -> TractResult> { let input_dt = inputs[0].datum_type(); - let scales = self.optional_scales_input.and_then(|ix| inputs.get(ix)); - let sizes = self.optional_sizes_input.and_then(|ix| inputs.get(ix)); - let output_shape = self.compute_output_shape( - inputs[0].shape(), - scales.map(|t| &**t), - sizes.map(|t| &**t), - )?; - let scales: TVec = if let Some(scales) = scales.filter(|s| s.len() == inputs[0].rank()) - { - scales.try_as_plain()?.as_slice::()?.into() + let rank = inputs[0].rank(); + let tf_crop = self.coord_transformer == CoordTransform::TfCropAndResize; + let roi = if tf_crop { + self.roi(rank, self.optional_roi_input.and_then(|ix| inputs.get(ix)).map(|t| &**t))? } else { - output_shape.iter().zip(inputs[0].shape()).map(|(o, i)| *o as f32 / *i as f32).collect() + tvec!((0.0, 1.0); rank) }; + let (scales, output_shape) = self.resolve( + inputs[0].shape(), + self.optional_scales_input.and_then(|ix| inputs.get(ix)).map(|t| &**t), + self.optional_sizes_input.and_then(|ix| inputs.get(ix)).map(|t| &**t), + )?; let input = inputs.remove(0).into_tensor(); - let mut data = if input.datum_type() == f32::datum_type() { - input.into_plain_array::()? - } else { - input.cast_to::()?.into_owned().into_plain_array::()? - }; - for (axis, scale) in scales.into_iter().enumerate().filter(|(_, s)| *s != 1.0) { - let mut new_shape: TVec = data.shape().into(); - new_shape[axis] = output_shape[axis]; - let input_len = data.shape()[axis]; - data = match self.interpolator { - Interpolator::Cubic => { - let a = self.cubic_coeff_a(); - let exclude = self.exclude_outside; - tract_ndarray::ArrayD::from_shape_fn(&*new_shape, |co_o| -> f32 { - let x_out = co_o[axis]; - let x_in = self.coord_transformer.transform( - x_out, - scale, - input_len, - new_shape[axis], - ); - let x_floor = x_in.floor() as isize; - let t = x_in - x_floor as f32; - let mut co_i = co_o; - let mut weights = [0.0f32; 4]; - let mut values = [0.0f32; 4]; - for (i, j) in (-1..=2isize).enumerate() { - let raw_idx = x_floor + j; - let w = cubic_kernel(t - j as f32, a); - if exclude && (raw_idx < 0 || raw_idx >= input_len as isize) { - weights[i] = 0.0; - } else { - weights[i] = w; - let idx = raw_idx.clamp(0, input_len as isize - 1) as usize; - co_i[axis] = idx; - values[i] = data[&co_i]; - } - } - if exclude { - let sum: f32 = weights.iter().sum(); - if sum != 0.0 { - for w in &mut weights { - *w /= sum; - } - } - } - weights.iter().zip(values.iter()).map(|(w, v)| w * v).sum() - }) - } - _ => tract_ndarray::ArrayD::from_shape_fn(&*new_shape, |co_o| -> f32 { - let x_out = co_o[axis]; - let x_in = - self.coord_transformer.transform(x_out, scale, input_len, new_shape[axis]); - let mut co_i = co_o; - let x_floor = x_in.floor() as isize; - let x_left = x_floor.clamp(0, input_len as isize - 1) as usize; - co_i[axis] = x_left; - let y_left = data[&co_i]; - let x_right = (x_floor + 1).clamp(0, input_len as isize - 1) as usize; - co_i[axis] = x_right; - let y_right = data[&co_i]; - let x_frac = x_in - x_floor as f32; - match self.interpolator { - Interpolator::Linear => y_left * (1.0 - x_frac) + y_right * x_frac, - Interpolator::Nearest => { - if self.nearest.prefers_right(x_frac) { - y_right - } else { - y_left - } - } - Interpolator::Cubic => unreachable!(), - } - }), + let input = input.cast_to::()?; + let mut shape: TVec = input.shape().into(); + let mut data: Vec = input.try_as_plain()?.as_slice::()?.to_vec(); + for (axis, scale) in scales.into_iter().enumerate() { + let (len_in, len_out) = (shape[axis], output_shape[axis]); + if len_in == len_out && scale == 1.0 && !tf_crop { + continue; } + let plan = self.plan_axis(scale, len_in, len_out, roi[axis]); + let mut resampled = vec![0f32; data.len() / len_in * len_out]; + resample_axis(&data, &shape, axis, &plan, self.extrapolation_value(), &mut resampled); + data = resampled; + shape[axis] = len_out; } - let out = data.into_tensor(); + let out = tract_ndarray::ArrayD::from_shape_vec(&*shape, data)?.into_tensor(); let out = if out.datum_type() == input_dt { out } else { out.cast_to_dt(input_dt)?.into_owned() }; Ok(tvec!(out.into_tvalue())) @@ -254,7 +406,6 @@ impl TypedOp for Resize { as_op!(); fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult> { - let _roi = self.optional_roi_input.and_then(|ix| inputs.get(ix)); let scales = self.optional_scales_input.and_then(|ix| inputs.get(ix)); let sizes = self.optional_sizes_input.and_then(|ix| inputs.get(ix)); let output_shape = self.compute_output_shape( @@ -270,29 +421,41 @@ impl TypedOp for Resize { model: &TypedModel, node: &TypedNode, ) -> TractResult> { + let input_fact = model.outlet_fact(node.inputs[0])?; + let rank = input_fact.rank(); + let axes = self.resized_axes(rank)?; if let Some(mut core_op) = self.as_core() { - let rank = model.outlet_fact(node.inputs[0])?.rank(); - let is_active = |ix: usize| -> bool { + let konst = |ix: usize| -> Option> { model .outlet_fact(node.inputs[ix]) - .ok() - .and_then(|f| f.konst.as_ref()) - .map(|k| k.len() == rank) - .unwrap_or(false) + .ok()? + .konst + .clone() + .filter(|k| k.len() == axes.len()) }; let active = self .optional_scales_input - .filter(|&ix| is_active(ix)) + .filter(|&ix| konst(ix).is_some()) .map(|ix| (ix, false)) .or_else(|| { - self.optional_sizes_input.filter(|&ix| is_active(ix)).map(|ix| (ix, true)) + self.optional_sizes_input.filter(|&ix| konst(ix).is_some()).map(|ix| (ix, true)) }); if let Some((ix, use_sizes)) = active { core_op.optional_scales_input = (!use_sizes).then_some(1); core_op.optional_sizes_input = use_sizes.then_some(1); let mut patch = TypedModelPatch::default(); let data = patch.tap_model(model, node.inputs[0])?; - let aux = patch.tap_model(model, node.inputs[ix])?; + let aux = if axes.len() == rank { + patch.tap_model(model, node.inputs[ix])? + } else { + let full = self.full_rank_aux( + &input_fact.shape, + &axes, + &konst(ix).unwrap(), + use_sizes, + )?; + patch.add_const(format!("{}.resize_aux", node.name), full)? + }; let wire = patch.wire_node(&node.name, core_op, &[data, aux])?; patch.shunt_outside(model, node.id.into(), wire[0])?; return Ok(Some(patch)); @@ -303,6 +466,7 @@ impl TypedOp for Resize { rule_if_some!(scales_input = self.optional_scales_input); let scales_fact = model.outlet_fact(node.inputs[scales_input])?; rule_if_some!(scales_tensor = &scales_fact.konst); + rule_if!(scales_tensor.len() == rank); let scales: Vec = scales_tensor.cast_to::()?.try_as_plain()?.as_slice::()?.to_vec(); let int_scales: Vec = scales.iter().map(|&s| s.round() as usize).collect(); @@ -310,14 +474,17 @@ impl TypedOp for Resize { scales.iter().zip(&int_scales).all(|(&s, &i)| (s - i as f32).abs() <= 1e-5 && i != 0) ); rule_if!(int_scales.iter().any(|&s| s != 1)); - let input_shape = &model.outlet_fact(node.inputs[0])?.shape; + let CoordTransform::Plain(coord_transformer) = &self.coord_transformer else { + return Ok(None); + }; for (axis, &scale) in int_scales.iter().enumerate().filter(|&(_, &s)| s > 1) { - let Some(len_in) = probe_length(&self.coord_transformer, &input_shape[axis]) else { + let Some(len_in) = probe_length(coord_transformer, &input_fact.shape[axis]) else { return Ok(None); }; - rule_if!(is_pixel_replication(&self.coord_transformer, len_in, scale, |frac| self - .nearest - .prefers_right(frac))); + rule_if!(is_pixel_replication( + &self.plan_axis(scale as f32, len_in, len_in * scale, (0.0, 1.0)), + scale + )); } lower_nearest_integer_upsample(model, node, &int_scales) @@ -340,68 +507,102 @@ fn parameters() -> Vec { vec![ TypeName::Scalar.tensor().named("input"), TypeName::Scalar.tensor().named("scales"), + TypeName::Scalar.tensor().named("roi").default(false), TypeName::String.named("coord_transformer").default("half_pixel"), TypeName::String.named("interpolator").default("nearest"), TypeName::String.named("nearest_mode").default("floor"), TypeName::Scalar.named("cubic_coeff_a").default(-0.75f32), TypeName::Logical.named("exclude_outside").default(false), + TypeName::Logical.named("antialias").default(false), + TypeName::Scalar.named("extrapolation_value").default(0.0f32), ] } fn dump(ast: &mut IntoAst, node: &TypedNode, op: &Resize) -> TractResult>> { let input = ast.mapping[&node.inputs[0]].clone(); - let scales = if let Some(scales_ix) = op.optional_scales_input { - ast.mapping[&node.inputs[scales_ix]].clone() - } else if let Some(sizes_ix) = op.optional_sizes_input { - let input_shape = ast.model.outlet_fact(node.inputs[0])?.shape.to_tvec(); - let sizes_fact = ast.model.outlet_fact(node.inputs[sizes_ix])?; - let sizes = - sizes_fact.konst.as_ref().context("sizes must be a constant for NNEF export")?; - let sizes = sizes.cast_to::()?; - let sizes = sizes.try_as_plain()?.as_slice::()?; - let scales: Vec = input_shape - .iter() - .zip(sizes.iter()) - .map(|(i, s)| i.to_usize().map(|i| *s / i as f32).unwrap_or(1.0)) - .collect(); - let scales_tensor = tract_ndarray::arr1(&scales).into_arc_tensor(); - ast.konst_variable(format!("{}.scales", node.name), &scales_tensor)? + let input_shape = ast.model.outlet_fact(node.inputs[0])?.shape.to_tvec(); + let axes = op.resized_axes(input_shape.len())?; + let passthrough = op.optional_scales_input.filter(|&ix| { + axes.len() == input_shape.len() + && ast + .model + .outlet_fact(node.inputs[ix]) + .map(|f| f.shape.volume() == axes.len().to_dim()) + .unwrap_or(false) + }); + let scales = if let Some(ix) = passthrough { + ast.mapping[&node.inputs[ix]].clone() } else { - bail!("Resize op has neither scales nor sizes input") + let output_shape = &node.outputs[0].fact.shape; + let mut scales = vec![1.0f32; input_shape.len()]; + for &axis in &axes { + let (i, o) = (input_shape[axis].to_usize()?, output_shape[axis].to_usize()?); + scales[axis] = o as f32 / i as f32; + } + let scales = tract_ndarray::arr1(&scales).into_arc_tensor(); + ast.konst_variable(format!("{}.scales", node.name), &scales)? }; - Ok(Some(invocation( - "tract_onnx_resize", - &[input, scales], - &[ - ("coord_transformer", string(op.coord_transformer.as_str())), - ("interpolator", string(op.interpolator.as_str())), - ("nearest_mode", string(op.nearest.as_str())), - ("cubic_coeff_a", numeric(op.cubic_coeff_a())), - ("exclude_outside", logical(op.exclude_outside)), - ], - ))) + let mut args = vec![ + ("coord_transformer", string(op.coord_transformer.as_str())), + ("interpolator", string(op.interpolator.as_str())), + ("nearest_mode", string(op.nearest.as_str())), + ("cubic_coeff_a", numeric(op.cubic_coeff_a())), + ("exclude_outside", logical(op.exclude_outside)), + ("antialias", logical(op.antialias)), + ("extrapolation_value", numeric(op.extrapolation_value())), + ]; + if op.coord_transformer == CoordTransform::TfCropAndResize { + let ix = op.optional_roi_input.context("tf_crop_and_resize needs a roi input")?; + let roi = ast.model.outlet_fact(node.inputs[ix])?; + let roi = roi.konst.as_ref().context("roi must be a constant for NNEF export")?; + let roi = full_rank_roi(&axes, input_shape.len(), roi)?; + let roi = ast.konst_variable(format!("{}.roi", node.name), &roi)?; + args.push(("roi", (*roi).clone())); + } + Ok(Some(invocation("tract_onnx_resize", &[input, scales], &args))) +} + +/// Spreads a per-`axes` ROI over the full rank, the layout the NNEF op expects. +fn full_rank_roi(axes: &[usize], rank: usize, roi: &Tensor) -> TractResult> { + let roi = roi.cast_to::()?; + let roi = roi.try_as_plain()?.as_slice::()?; + let mut full = vec![0.0f32; rank]; + full.extend(std::iter::repeat_n(1.0f32, rank)); + for (i, &axis) in axes.iter().enumerate() { + full[axis] = roi[i]; + full[rank + axis] = roi[i + axes.len()]; + } + Ok(tract_ndarray::arr1(&full).into_arc_tensor()) } fn load(builder: &mut ModelBuilder, invocation: &ResolvedInvocation) -> TractResult { let input = invocation.named_arg_as(builder, "input")?; let scales = invocation.named_arg_as(builder, "scales")?; + let roi = invocation.optional_named_arg_as::(builder, "roi")?; let coord_transformer: String = invocation.named_arg_as(builder, "coord_transformer")?; let interpolator: String = invocation.named_arg_as(builder, "interpolator")?; let nearest_mode: String = invocation.named_arg_as(builder, "nearest_mode")?; let cubic_coeff_a: f32 = invocation.named_arg_as(builder, "cubic_coeff_a")?; let exclude_outside: bool = invocation.named_arg_as(builder, "exclude_outside")?; + let antialias: bool = invocation.named_arg_as(builder, "antialias")?; + let extrapolation_value: f32 = invocation.named_arg_as(builder, "extrapolation_value")?; let op = Resize { axes: None, - coord_transformer: CoordTransformer::parse(&coord_transformer)?, + coord_transformer: CoordTransform::parse(&coord_transformer)?, interpolator: Interpolator::parse(&interpolator)?, nearest: Nearest::parse(&nearest_mode)?, + antialias, cubic_coeff_a_bits: cubic_coeff_a.to_bits(), exclude_outside, - optional_roi_input: None, + extrapolation_value_bits: extrapolation_value.to_bits(), + keep_aspect_ratio_policy: AspectRatio::Stretch, + optional_roi_input: roi.map(|_| 2), optional_scales_input: Some(1), optional_sizes_input: None, }; - builder.wire(op, &[input, scales]) + let mut wires = tvec!(input, scales); + wires.extend(roi); + builder.wire(op, &wires) } diff --git a/onnx/src/ops/resize.rs b/onnx/src/ops/resize.rs index 7d348a0d34..700d35d438 100644 --- a/onnx/src/ops/resize.rs +++ b/onnx/src/ops/resize.rs @@ -1,9 +1,9 @@ use crate::model::ParsingContext; use crate::pb::*; use tract_hir::internal::*; -use tract_nnef::tract_core::ops::nn::resize::{CoordTransformer, Interpolator}; +use tract_nnef::tract_core::ops::nn::resize::Interpolator; use tract_nnef::tract_num_traits::Zero as _; -use tract_onnx_opl::resize::{Nearest, Resize}; +use tract_onnx_opl::resize::{AspectRatio, CoordTransform, Nearest, Resize}; pub fn resize( ctx: &ParsingContext, @@ -21,45 +21,30 @@ pub fn resize( fn resize_10(node: &NodeProto) -> TractResult { Ok(Resize { - axes: None, optional_roi_input: None, optional_scales_input: Some(1), optional_sizes_input: None, - coord_transformer: coord_transformer_from_node(node)?, - interpolator: interpolator_from_node(node)?, - nearest: nearest_from_node(node)?, - cubic_coeff_a_bits: cubic_coeff_a_from_node(node)?, - exclude_outside: exclude_outside_from_node(node)?, + ..common(node)? }) } fn resize_11(node: &NodeProto) -> TractResult { let mut options = crate::model::optional_inputs(node).skip(3); Ok(Resize { - axes: None, optional_roi_input: Some(1), optional_scales_input: Some(2), optional_sizes_input: options.next().unwrap(), - coord_transformer: coord_transformer_from_node(node)?, - interpolator: interpolator_from_node(node)?, - nearest: nearest_from_node(node)?, - cubic_coeff_a_bits: cubic_coeff_a_from_node(node)?, - exclude_outside: exclude_outside_from_node(node)?, + ..common(node)? }) } fn resize_13(node: &NodeProto) -> TractResult { let mut options = crate::model::optional_inputs(node).skip(1); Ok(Resize { - axes: None, optional_roi_input: options.next().unwrap(), optional_scales_input: options.next().unwrap(), optional_sizes_input: options.next().unwrap(), - coord_transformer: coord_transformer_from_node(node)?, - interpolator: interpolator_from_node(node)?, - nearest: nearest_from_node(node)?, - cubic_coeff_a_bits: cubic_coeff_a_from_node(node)?, - exclude_outside: exclude_outside_from_node(node)?, + ..common(node)? }) } @@ -67,19 +52,40 @@ fn resize_18(node: &NodeProto) -> TractResult { let mut options = crate::model::optional_inputs(node).skip(1); Ok(Resize { axes: node.get_attr_opt_vec("axes")?, + antialias: node.get_attr_opt::("antialias")?.unwrap_or(0) != 0, + keep_aspect_ratio_policy: AspectRatio::parse( + node.get_attr_opt("keep_aspect_ratio_policy")?.unwrap_or("stretch"), + )?, optional_roi_input: options.next().unwrap(), optional_scales_input: options.next().unwrap(), optional_sizes_input: options.next().unwrap(), + ..common(node)? + }) +} + +/// The attributes every Resize opset shares. `axes`, `antialias` and +/// `keep_aspect_ratio_policy` arrive with opset 18 and stay at their neutral +/// value below it. +fn common(node: &NodeProto) -> TractResult { + let extrapolation_value: f32 = node.get_attr_opt("extrapolation_value")?.unwrap_or(0.0); + Ok(Resize { + axes: None, + antialias: false, + keep_aspect_ratio_policy: AspectRatio::Stretch, coord_transformer: coord_transformer_from_node(node)?, interpolator: interpolator_from_node(node)?, nearest: nearest_from_node(node)?, cubic_coeff_a_bits: cubic_coeff_a_from_node(node)?, exclude_outside: exclude_outside_from_node(node)?, + extrapolation_value_bits: extrapolation_value.to_bits(), + optional_roi_input: None, + optional_scales_input: None, + optional_sizes_input: None, }) } -fn coord_transformer_from_node(node: &NodeProto) -> TractResult { - CoordTransformer::parse( +fn coord_transformer_from_node(node: &NodeProto) -> TractResult { + CoordTransform::parse( node.get_attr_opt("coordinate_transformation_mode")?.unwrap_or("half_pixel"), ) } @@ -130,7 +136,7 @@ impl Expansion for ResizeInference { } else if op.optional_sizes_input.is_some() { rules_with_sizes(op, s, inputs, outputs) } else { - todo!() + bail!("Resize with neither scales nor sizes") } } @@ -154,7 +160,9 @@ fn rules_with_scales<'r, 'p: 'r, 's: 'r>( let scales = &inputs[scales_input]; s.equals(&scales.datum_type, f32::datum_type())?; s.equals(&scales.rank, 1)?; - s.equals(&scales.shape[0], inputs[0].rank.bex().to_dim())?; + s.given(&inputs[0].rank, move |s, rank| { + s.equals(&scales.shape[0], op.resized_axes(rank as usize)?.len().to_dim()) + })?; s.given_2(&inputs[0].shape, &inputs[scales_input].value, move |s, input_shape, scales| { let output_size = op.compute_output_shape(&input_shape, Some(scales.as_ref()), None)?; let rank = input_shape.len(); @@ -173,10 +181,13 @@ fn rules_with_sizes<'r, 'p: 'r, 's: 'r>( ) -> InferenceResult { let sizes = &inputs[op.optional_sizes_input.unwrap()]; s.equals(&sizes.rank, 1)?; - s.equals(&sizes.shape[0], inputs[0].rank.bex().to_dim())?; s.given(&inputs[0].rank, move |s, rank| { - for i in 0..(rank as usize) { - s.equals(&outputs[0].shape[i], sizes.value[i].bex().to_dim())?; + s.equals(&sizes.shape[0], op.resized_axes(rank as usize)?.len().to_dim()) + })?; + s.given_2(&inputs[0].shape, &sizes.value, move |s, input_shape, sizes| { + let output_size = op.compute_output_shape(&input_shape, None, Some(sizes.as_ref()))?; + for (i, dim) in output_size.iter().enumerate() { + s.equals(&outputs[0].shape[i], dim.to_dim())?; } Ok(()) }) diff --git a/test-rt/suite-onnx/node.txt b/test-rt/suite-onnx/node.txt index 08cf225252..a14d1f7bba 100644 --- a/test-rt/suite-onnx/node.txt +++ b/test-rt/suite-onnx/node.txt @@ -488,19 +488,7 @@ test_reshape_reordered_dims test_reshape_reordered_last_dims input:data test_reshape_zero_and_negative_dim input:data test_reshape_zero_dim input:data -test_resize_downsample_scales_cubic input:X -test_resize_downsample_scales_cubic_align_corners input:X -test_resize_downsample_scales_cubic_A_n0p5_exclude_outside input:X -test_resize_downsample_scales_linear input:X -test_resize_downsample_sizes_cubic input:X -test_resize_downsample_sizes_linear_pytorch_half_pixel input:X -test_resize_upsample_scales_cubic input:X -test_resize_upsample_scales_cubic_align_corners input:X -test_resize_upsample_scales_cubic_asymmetric input:X -test_resize_upsample_scales_cubic_A_n0p5_exclude_outside input:X -test_resize_upsample_scales_linear_align_corners input:X not-nnef -test_resize_upsample_sizes_cubic input:X -test_resize_upsample_sizes_nearest_ceil_half_pixel input:X +test_resize.* input:X test_rnn_seq_length test_round test_scan9_sum From 690e5afa0baf2d823cd7293120f7dd14d1fee041 Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Sun, 2 Aug 2026 18:44:02 +0100 Subject: [PATCH 3/4] metal: run Resize on the GPU Resize had no Metal kernel, so every node round-tripped the tensor through the host in the middle of a GPU graph. A resample-one-axis kernel now consumes the same per-axis tap-and-weight plan the CPU op builds, which makes it independent of the interpolator; the plan is baked at translation time, so the node keeps only its data input and the scales/sizes TDim constant is dropped. --- core/src/ops/nn/resize.rs | 3 +- gpu/src/ops/mod.rs | 1 + gpu/src/ops/resize.rs | 109 +++++++++++++++++ metal/src/kernels/array/array_ops.metal | 49 ++++++++ metal/src/kernels/array/mod.rs | 10 ++ metal/src/kernels/array/resize.rs | 156 ++++++++++++++++++++++++ metal/src/transform.rs | 18 ++- 7 files changed, 344 insertions(+), 2 deletions(-) create mode 100644 gpu/src/ops/resize.rs create mode 100644 metal/src/kernels/array/resize.rs diff --git a/core/src/ops/nn/resize.rs b/core/src/ops/nn/resize.rs index 5c3f9f0f93..aa5367b0fd 100644 --- a/core/src/ops/nn/resize.rs +++ b/core/src/ops/nn/resize.rs @@ -352,7 +352,8 @@ impl Resize { ); } - fn plan_axis(&self, scale: f32, len_in: usize, len_out: usize) -> AxisPlan { + /// The resampling plan for one axis under this op's interpolator. + pub fn plan_axis(&self, scale: f32, len_in: usize, len_out: usize) -> AxisPlan { let window = window_size(&self.interpolator, false, scale); let coord = |x| Some(self.coord_transformer.transform(x, scale, len_in, len_out)); match self.interpolator { diff --git a/gpu/src/ops/mod.rs b/gpu/src/ops/mod.rs index ec0f7d32dd..8aa7fe3925 100644 --- a/gpu/src/ops/mod.rs +++ b/gpu/src/ops/mod.rs @@ -15,6 +15,7 @@ pub mod leaky_relu; pub mod pad; pub mod pulse; pub mod reduce; +pub mod resize; pub mod rms_norm; pub mod rotate_half; pub mod scaled_masked_softmax; diff --git a/gpu/src/ops/resize.rs b/gpu/src/ops/resize.rs new file mode 100644 index 0000000000..da24455729 --- /dev/null +++ b/gpu/src/ops/resize.rs @@ -0,0 +1,109 @@ +use crate::tensor::{DeviceTensor, DeviceTensorExt, IntoDevice}; +use derive_new::new; +use tract_core::internal::*; + +/// Resamples one axis: `output[.., x, ..] = sum_k weights[x, k] * input[.., indices[x, k], ..]`, +/// with `indices` already clamped into the axis by the host-built plan. +pub type DispatchResizeAxisFn = fn( + input: &DeviceTensor, + axis: usize, + indices: &DeviceTensor, + weights: &DeviceTensor, + window: usize, + output: &DeviceTensor, +) -> TractResult<()>; + +/// Resize against a plan baked at translation time, one dispatch per resampled +/// axis. The plan makes the op independent of the interpolator: nearest, linear +/// and cubic differ only in window size and weights, so translation is limited +/// to nodes whose shapes and scales are known then. The scales/sizes input is +/// kept for arity but no longer read. +#[derive(Clone, new)] +pub struct GpuResize { + pub axes: TVec, + pub windows: TVec, + pub plans: TVec<(Arc, Arc)>, + pub output_shape: TVec, + pub backend_name: &'static str, + pub dispatch: DispatchResizeAxisFn, +} + +impl std::fmt::Debug for GpuResize { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{}Resize", self.backend_name) + } +} + +impl PartialEq for GpuResize { + fn eq(&self, other: &Self) -> bool { + self.backend_name == other.backend_name + && self.axes == other.axes + && self.windows == other.windows + && self.output_shape == other.output_shape + } +} +impl Eq for GpuResize {} + +impl std::hash::Hash for GpuResize { + fn hash(&self, state: &mut H) { + self.backend_name.hash(state); + self.axes.hash(state); + self.windows.hash(state); + self.output_shape.hash(state); + } +} + +impl Op for GpuResize { + fn name(&self) -> StaticName { + format!("{}Resize", self.backend_name).into() + } + fn info(&self) -> TractResult> { + Ok(vec![format!("axes={:?} windows={:?}", self.axes, self.windows)]) + } + op_as_typed_op!(); +} + +impl EvalOp for GpuResize { + fn is_stateless(&self) -> bool { + true + } + + fn eval_with_session( + &self, + node_id: usize, + session: &TurnState, + inputs: TVec, + ) -> TractResult> { + let data = inputs[0].to_device_tensor()?; + let dt = data.datum_type(); + let mut shape: TVec = data.shape().into(); + let mut current = data.clone(); + for (step, (&axis, &window)) in self.axes.iter().zip(&self.windows).enumerate() { + let (indices, weights) = &self.plans[step]; + let indices = indices.as_ref().clone().into_device()?; + let weights = weights.as_ref().clone().into_device()?; + shape[axis] = self.output_shape[axis]; + let last = step + 1 == self.axes.len(); + let output = if last { + crate::session_handler::make_tensor_for_node(session, node_id, dt, &shape)? + } else { + DeviceTensor::uninitialized_dt(dt, &shape)? + }; + (self.dispatch)(¤t, axis, &indices, &weights, window, &output)?; + current = output; + } + Ok(tvec!(current.into_tensor().into_tvalue())) + } +} + +impl TypedOp for GpuResize { + fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult> { + crate::utils::facts_to_device_facts(inputs, |facts| { + ensure!(facts.len() == 1); + let shape: TVec = self.output_shape.iter().map(|d| d.to_dim()).collect(); + Ok(tvec!(facts[0].datum_type.fact(&shape))) + }) + .with_context(|| format!("Error while computing facts for {:?}", self.name())) + } + as_op!(); +} diff --git a/metal/src/kernels/array/array_ops.metal b/metal/src/kernels/array/array_ops.metal index 2a5353c3d0..4fef1e6d77 100644 --- a/metal/src/kernels/array/array_ops.metal +++ b/metal/src/kernels/array/array_ops.metal @@ -336,6 +336,51 @@ typedef decltype(gather) gather_t; template [[host_name( \ "array_ops::gather_" #tname)]] [[kernel]] gather_t gather; +// Resample one axis against a host-built plan: every output index owns a window +// of `window` input indices (already clamped) and their weights, so nearest, +// linear, cubic and their antialiased variants all run through this kernel. +// The axes around the resampled one are flattened into `outer` and `inner`. +// +// params layout: [outer, len_in, len_out, inner, window] +template +[[kernel]] void resize_axis(device const void *input_b [[buffer(0)]], + device void *output_b [[buffer(1)]], + device const int32_t *indices [[buffer(2)]], + device const float *weights [[buffer(3)]], + constant const int32_t *params [[buffer(4)]], + uint3 tpig [[thread_position_in_grid]]) { + const int32_t len_in = params[1]; + const int32_t len_out = params[2]; + const int32_t inner = params[3]; + const int32_t window = params[4]; + + const int32_t i = (int32_t)tpig.x; + const int32_t x = (int32_t)tpig.y; + const int32_t o = (int32_t)tpig.z; + + if (i >= inner || x >= len_out) + return; + + device const T *input = (device const T *)input_b; + device T *output = (device T *)output_b; + + const int32_t in_base = o * len_in * inner + i; + float acc = 0.0f; + for (int32_t k = 0; k < window; ++k) { + const float w = weights[x * window + k]; + if (w != 0.0f) + acc += w * (float)input[in_base + indices[x * window + k] * inner]; + } + output[o * len_out * inner + x * inner + i] = (T)acc; +} + +typedef decltype(resize_axis) resize_axis_t; + +#define INSTANTIATE_RESIZE_AXIS(tname, type) \ + template [[host_name( \ + "array_ops::resize_axis_" #tname)]] [[kernel]] resize_axis_t \ + resize_axis; + // Copy kernels: only u8/u16/u32/u64 (copy is type-size based) INSTANTIATE_COPY(u8, uint8_t) INSTANTIATE_COPY(u16, uint16_t) @@ -379,3 +424,7 @@ INSTANTIATE_DIAG_GATHER(f16, half) // Axis Gather: f32 and f16 only (indices are int64). INSTANTIATE_GATHER(f32, float) INSTANTIATE_GATHER(f16, half) + +// Axis resample: f32 and f16 only (plan indices are int32, weights f32). +INSTANTIATE_RESIZE_AXIS(f32, float) +INSTANTIATE_RESIZE_AXIS(f16, half) diff --git a/metal/src/kernels/array/mod.rs b/metal/src/kernels/array/mod.rs index d9a51b1dc2..9e2cd212c2 100644 --- a/metal/src/kernels/array/mod.rs +++ b/metal/src/kernels/array/mod.rs @@ -3,6 +3,7 @@ mod copy; mod diag_gather; mod dispatch; mod gather; +mod resize; mod rotate_half; pub use cast::Cast; @@ -13,6 +14,9 @@ pub use diag_gather::metal_diag_gather_dispatch; pub use dispatch::metal_copy_nd_dispatch; pub use gather::Gather; pub use gather::metal_gather_dispatch; +pub use resize::ResizeAxis; +pub use resize::metal_resize; +pub use resize::metal_resize_axis_dispatch; pub use rotate_half::RotateHalf; pub use rotate_half::metal_rotate_half_dispatch; @@ -50,5 +54,11 @@ pub fn all_functions() -> Vec { .flat_map(|dt| Gather.kernel_name(dt).into_iter()), ); + functions.extend( + tract_gpu::tensor::DeviceTensor::SUPPORTED_DT + .into_iter() + .flat_map(|dt| ResizeAxis.kernel_name(dt).into_iter()), + ); + functions.into_iter().collect() } diff --git a/metal/src/kernels/array/resize.rs b/metal/src/kernels/array/resize.rs new file mode 100644 index 0000000000..1561aa36dd --- /dev/null +++ b/metal/src/kernels/array/resize.rs @@ -0,0 +1,156 @@ +use crate::encoder::EncoderExt; +use crate::{LibraryName, MetalStream}; +use anyhow::ensure; +use metal::MTLSize; +use std::fmt; +use tract_core::internal::*; +use tract_core::ops::nn::resize::Resize as CoreResize; +use tract_gpu::ops::resize::GpuResize; +use tract_gpu::tensor::DeviceTensor; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ResizeAxis; + +impl fmt::Display for ResizeAxis { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{self:?}") + } +} + +impl ResizeAxis { + pub fn is_supported_dt(dt: DatumType) -> bool { + matches!(dt, DatumType::F32 | DatumType::F16) + } + + pub fn kernel_name(&self, dt: DatumType) -> TractResult { + ensure!(Self::is_supported_dt(dt), "Unsupported dt {:?} for metal resize op", dt); + let tname = DeviceTensor::tname(dt)?; + Ok(format!("array_ops::resize_axis_{tname}")) + } + + #[allow(clippy::too_many_arguments)] + pub fn dispatch_eval( + &self, + stream: &MetalStream, + input: &DeviceTensor, + axis: usize, + indices: &DeviceTensor, + weights: &DeviceTensor, + window: usize, + output: &DeviceTensor, + ) -> TractResult<()> { + stream.retain_tensor(input); + stream.retain_tensor(indices); + stream.retain_tensor(weights); + stream.retain_tensor(output); + + ensure!(input.rank() > axis); + ensure!(output.datum_type() == input.datum_type()); + ensure!(indices.datum_type() == i32::datum_type()); + ensure!(weights.datum_type() == f32::datum_type()); + + let len_in = input.shape()[axis]; + let len_out = output.shape()[axis]; + let inner: usize = input.shape()[axis + 1..].iter().product(); + let outer: usize = input.shape()[..axis].iter().product(); + ensure!(indices.len() == len_out * window && weights.len() == len_out * window); + + let params: [i32; 5] = + [outer as i32, len_in as i32, len_out as i32, inner as i32, window as i32]; + + let pipeline = + stream.load_pipeline(LibraryName::ArrayOps, &self.kernel_name(input.datum_type())?)?; + let command_buffer = stream.command_buffer(); + command_buffer.encode(|encoder| { + encoder.set_compute_pipeline_state(&pipeline); + encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); + encoder.set_metal_tensor(1, output, metal::MTLResourceUsage::Write); + encoder.set_metal_tensor(2, indices, metal::MTLResourceUsage::Read); + encoder.set_metal_tensor(3, weights, metal::MTLResourceUsage::Read); + encoder.set_slice(4, ¶ms); + let grid_size = MTLSize { width: inner as _, height: len_out as _, depth: outer as _ }; + let group_size = MTLSize { width: 1, height: 1, depth: 1 }; + encoder.dispatch_thread_groups(grid_size, group_size); + }); + Ok(()) + } +} + +pub fn metal_resize_axis_dispatch( + input: &DeviceTensor, + axis: usize, + indices: &DeviceTensor, + weights: &DeviceTensor, + window: usize, + output: &DeviceTensor, +) -> TractResult<()> { + crate::with_metal_stream(|stream| { + ResizeAxis.dispatch_eval(stream, input, axis, indices, weights, window, output) + }) +} + +/// Bakes the resampling plan for every non-identity axis. Needs concrete input +/// dims and a constant scales/sizes input, so a symbolically-shaped or +/// dynamically-scaled Resize stays on CPU. +fn baked_plans( + source: &TypedModel, + node: &TypedNode, + op: &CoreResize, +) -> TractResult> { + let facts = source.node_input_facts(node.id)?; + let Some(input_shape) = facts[0].shape.as_concrete() else { return Ok(None) }; + let aux = |ix: Option| ix.and_then(|ix| facts.get(ix)?.konst.as_deref()); + let scales_konst = aux(op.optional_scales_input); + let sizes_konst = aux(op.optional_sizes_input); + if scales_konst.is_none() && sizes_konst.is_none() { + return Ok(None); + } + let output_shape: TVec = + op.compute_output_shape(input_shape, scales_konst, sizes_konst)?; + + let scales: TVec = match scales_konst.filter(|s| s.len() == input_shape.len()) { + Some(scales) => scales.cast_to::()?.try_as_plain()?.as_slice::()?.into(), + None => output_shape.iter().zip(input_shape).map(|(o, i)| *o as f32 / *i as f32).collect(), + }; + + let (mut axes, mut windows, mut plans) = (tvec!(), tvec!(), tvec!()); + for (axis, &scale) in scales.iter().enumerate() { + let (len_in, len_out) = (input_shape[axis], output_shape[axis]); + if len_in == len_out && scale == 1.0 { + continue; + } + let plan = op.plan_axis(scale, len_in, len_out); + let indices: Vec = plan.indices.iter().map(|&i| i as i32).collect(); + plans.push(( + tract_ndarray::arr1(&indices).into_arc_tensor(), + tract_ndarray::arr1(&plan.weights).into_arc_tensor(), + )); + axes.push(axis); + windows.push(plan.window); + } + if axes.is_empty() { + return Ok(None); + } + Ok(Some(GpuResize::new( + axes, + windows, + plans, + output_shape, + "Metal", + metal_resize_axis_dispatch, + ))) +} + +/// Translates a core Resize whose plan can be baked. Wired by `transform.rs` +/// rather than `register_metal_op!` because the resulting op drops the +/// scales/sizes input, whose `TDim` datum type has no device equivalent. +pub fn metal_resize( + source: &TypedModel, + node: &TypedNode, +) -> TractResult>> { + let Some(op) = node.op_as::() else { return Ok(None) }; + let facts = source.node_input_facts(node.id)?; + rule_if!(facts[0].is_plain()); + rule_if!(ResizeAxis::is_supported_dt(facts[0].datum_type)); + Ok(baked_plans(source, node, op)?.map(|op| Box::new(op) as Box)) +} diff --git a/metal/src/transform.rs b/metal/src/transform.rs index 65510ee538..1d19a5d075 100644 --- a/metal/src/transform.rs +++ b/metal/src/transform.rs @@ -19,7 +19,9 @@ use tract_core::transform::ModelTransform; use tract_gpu::fact::{DeviceFact, DeviceTypedFactExt}; use tract_gpu::rewrite_rules::rewire_syncs::rewire_syncs; use tract_gpu::rewrite_rules::rms_norm::remove_rms_norm_cast; -use tract_gpu::sync::{DeviceSyncKind, sync_inputs_if_required, sync_model_outputs_if_required}; +use tract_gpu::sync::{ + DeviceSync, DeviceSyncKind, sync_inputs_if_required, sync_model_outputs_if_required, +}; use tract_gpu::tensor::{DeviceTensor, IntoDevice}; use tract_gpu::utils::as_quant_fact; @@ -244,6 +246,20 @@ impl Translate, TypedFact, Box> for Met ops::conv::wire_metal_conv(source, node, target, &device_inputs, conv)?; return sync_model_outputs_if_required(source, node, target, outlet_ids); } + // Resize bakes its plan, so it keeps only the data input: the scales / + // sizes input it drops is a TDim const with no device equivalent. + if let Some(gpu_op) = crate::kernels::array::metal_resize(source, node)? { + let mut input = mapping[&node.inputs[0]]; + if target.outlet_fact(input)?.as_device_fact().is_none() { + input = target.wire_node( + format!("{}.to-device-0", node.name), + DeviceSync::new(DeviceSyncKind::ToDevice), + &[input], + )?[0]; + } + let outlet_ids = target.wire_node(node.name.clone(), gpu_op, &[input])?; + return sync_model_outputs_if_required(source, node, target, outlet_ids); + } // Const: inline conversion, not a GPU op if let Some(op) = node.op_as::() && DeviceTensor::is_supported_dt(op.val().datum_type()) From 08d80314045d0ce8e84f272c6cd8720bc7f3c0ba Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Sun, 2 Aug 2026 19:15:10 +0100 Subject: [PATCH 4/4] metal: run 2D MaxPool on the GPU MaxPool had no Metal kernel, so it round-tripped its tensor through the host in the middle of a GPU graph. A kernel now pools two spatial axes with every axis stride passed explicitly, which covers NCHW and NHWC alike, and skips positions in the padding so an all-padding window keeps the CPU op's lowest-value result. The geometry is resolved at translation time; index outputs, other ranks and symbolic shapes stay on CPU. --- core/src/ops/cnn/mod.rs | 2 +- gpu/src/ops/max_pool.rs | 107 +++++++++++++++++++ gpu/src/ops/mod.rs | 1 + metal/src/kernels/nn/max_pool.rs | 165 ++++++++++++++++++++++++++++++ metal/src/kernels/nn/mod.rs | 9 ++ metal/src/kernels/nn/nn_ops.metal | 69 +++++++++++++ 6 files changed, 352 insertions(+), 1 deletion(-) create mode 100644 gpu/src/ops/max_pool.rs create mode 100644 metal/src/kernels/nn/max_pool.rs diff --git a/core/src/ops/cnn/mod.rs b/core/src/ops/cnn/mod.rs index 540e975c03..8bbbe7b3a2 100644 --- a/core/src/ops/cnn/mod.rs +++ b/core/src/ops/cnn/mod.rs @@ -11,7 +11,7 @@ mod sumpool; pub use self::conv::{Conv, KernelFormat}; pub use self::deconv::Deconv; -pub use self::maxpool::MaxPool; +pub use self::maxpool::{MaxPool, OptMaxPool}; pub use self::padding::PaddingSpec; pub use self::patch_axis::PatchAxis; pub use self::patches::{Patch, PatchSpec}; diff --git a/gpu/src/ops/max_pool.rs b/gpu/src/ops/max_pool.rs new file mode 100644 index 0000000000..319303850e --- /dev/null +++ b/gpu/src/ops/max_pool.rs @@ -0,0 +1,107 @@ +use crate::tensor::{DeviceTensor, DeviceTensorExt}; +use derive_new::new; +use tract_core::internal::*; + +/// Geometry of a two-spatial-axis max pooling, with every axis stride given +/// explicitly so one kernel serves both NCHW and NHWC. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct MaxPool2dGeometry { + pub batch: usize, + pub channels: usize, + pub input_hw: (usize, usize), + pub output_hw: (usize, usize), + pub kernel: (usize, usize), + pub strides: (usize, usize), + pub dilations: (usize, usize), + pub padding: (usize, usize), + pub input_strides: [usize; 4], + pub output_strides: [usize; 4], +} + +pub type DispatchMaxPool2dFn = + fn(input: &DeviceTensor, geo: &MaxPool2dGeometry, output: &DeviceTensor) -> TractResult<()>; + +/// Max pooling over two spatial axes, geometry resolved at translation time. +/// Windows lying entirely in the padding take the datum type's lowest value, +/// matching the CPU op. +#[derive(Clone, new)] +pub struct GpuMaxPool { + pub geometry: MaxPool2dGeometry, + pub output_shape: TVec, + pub backend_name: &'static str, + pub dispatch: DispatchMaxPool2dFn, +} + +impl std::fmt::Debug for GpuMaxPool { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{}MaxPool", self.backend_name) + } +} + +impl PartialEq for GpuMaxPool { + fn eq(&self, other: &Self) -> bool { + self.backend_name == other.backend_name + && self.geometry == other.geometry + && self.output_shape == other.output_shape + } +} +impl Eq for GpuMaxPool {} + +impl std::hash::Hash for GpuMaxPool { + fn hash(&self, state: &mut H) { + self.backend_name.hash(state); + self.geometry.hash(state); + self.output_shape.hash(state); + } +} + +impl Op for GpuMaxPool { + fn name(&self) -> StaticName { + format!("{}MaxPool", self.backend_name).into() + } + fn info(&self) -> TractResult> { + Ok(vec![format!( + "kernel={:?} strides={:?} dilations={:?} padding={:?}", + self.geometry.kernel, + self.geometry.strides, + self.geometry.dilations, + self.geometry.padding + )]) + } + op_as_typed_op!(); +} + +impl EvalOp for GpuMaxPool { + fn is_stateless(&self) -> bool { + true + } + + fn eval_with_session( + &self, + node_id: usize, + session: &TurnState, + inputs: TVec, + ) -> TractResult> { + let input = inputs[0].to_device_tensor()?; + let output = crate::session_handler::make_tensor_for_node( + session, + node_id, + input.datum_type(), + &self.output_shape, + )?; + (self.dispatch)(input, &self.geometry, &output)?; + Ok(tvec!(output.into_tensor().into_tvalue())) + } +} + +impl TypedOp for GpuMaxPool { + fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult> { + crate::utils::facts_to_device_facts(inputs, |facts| { + ensure!(facts.len() == 1); + let shape: TVec = self.output_shape.iter().map(|d| d.to_dim()).collect(); + Ok(tvec!(facts[0].datum_type.fact(&shape))) + }) + .with_context(|| format!("Error while computing facts for {:?}", self.name())) + } + as_op!(); +} diff --git a/gpu/src/ops/mod.rs b/gpu/src/ops/mod.rs index 8aa7fe3925..b03baafa43 100644 --- a/gpu/src/ops/mod.rs +++ b/gpu/src/ops/mod.rs @@ -12,6 +12,7 @@ pub mod gather; pub mod gelu_approximate; pub mod iff; pub mod leaky_relu; +pub mod max_pool; pub mod pad; pub mod pulse; pub mod reduce; diff --git a/metal/src/kernels/nn/max_pool.rs b/metal/src/kernels/nn/max_pool.rs new file mode 100644 index 0000000000..7ff980e6e5 --- /dev/null +++ b/metal/src/kernels/nn/max_pool.rs @@ -0,0 +1,165 @@ +use crate::encoder::EncoderExt; +use crate::{LibraryName, MetalStream}; +use anyhow::ensure; +use metal::MTLSize; +use std::fmt; +use tract_core::internal::*; +use tract_core::ops::cnn::pools::PoolSpec; +use tract_core::ops::cnn::{MaxPool, OptMaxPool}; +use tract_gpu::ops::max_pool::{GpuMaxPool, MaxPool2dGeometry}; +use tract_gpu::tensor::DeviceTensor; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct MaxPool2d; + +impl fmt::Display for MaxPool2d { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{self:?}") + } +} + +impl MaxPool2d { + pub fn is_supported_dt(dt: DatumType) -> bool { + matches!(dt, DatumType::F32 | DatumType::F16) + } + + pub fn kernel_name(&self, dt: DatumType) -> TractResult { + ensure!(Self::is_supported_dt(dt), "Unsupported dt {:?} for metal max pool op", dt); + let tname = DeviceTensor::tname(dt)?; + Ok(format!("nn_ops::max_pool_2d_{tname}")) + } + + pub fn dispatch_eval( + &self, + stream: &MetalStream, + input: &DeviceTensor, + geo: &MaxPool2dGeometry, + output: &DeviceTensor, + ) -> TractResult<()> { + stream.retain_tensor(input); + stream.retain_tensor(output); + ensure!(output.datum_type() == input.datum_type()); + + let params: [i32; 22] = [ + geo.batch as i32, + geo.channels as i32, + geo.input_hw.0 as i32, + geo.input_hw.1 as i32, + geo.output_hw.0 as i32, + geo.output_hw.1 as i32, + geo.kernel.0 as i32, + geo.kernel.1 as i32, + geo.strides.0 as i32, + geo.strides.1 as i32, + geo.dilations.0 as i32, + geo.dilations.1 as i32, + geo.padding.0 as i32, + geo.padding.1 as i32, + geo.input_strides[0] as i32, + geo.input_strides[1] as i32, + geo.input_strides[2] as i32, + geo.input_strides[3] as i32, + geo.output_strides[0] as i32, + geo.output_strides[1] as i32, + geo.output_strides[2] as i32, + geo.output_strides[3] as i32, + ]; + + let pipeline = + stream.load_pipeline(LibraryName::NNOps, &self.kernel_name(input.datum_type())?)?; + let command_buffer = stream.command_buffer(); + command_buffer.encode(|encoder| { + encoder.set_compute_pipeline_state(&pipeline); + encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); + encoder.set_metal_tensor(1, output, metal::MTLResourceUsage::Write); + encoder.set_slice(2, ¶ms); + let grid_size = MTLSize { + width: geo.output_hw.1 as _, + height: geo.output_hw.0 as _, + depth: (geo.batch * geo.channels) as _, + }; + let group_size = MTLSize { width: 1, height: 1, depth: 1 }; + encoder.dispatch_thread_groups(grid_size, group_size); + }); + Ok(()) + } +} + +pub fn metal_max_pool_dispatch( + input: &DeviceTensor, + geo: &MaxPool2dGeometry, + output: &DeviceTensor, +) -> TractResult<()> { + crate::with_metal_stream(|stream| MaxPool2d.dispatch_eval(stream, input, geo, output)) +} + +/// Resolves an `OptMaxPool` to a bakeable two-spatial-axis geometry. Index +/// outputs, ranks other than 4 and symbolic shapes stay on CPU. +fn geometry( + source: &TypedModel, + node: &TypedNode, + pool_spec: &PoolSpec, + with_index_outputs: Option, +) -> TractResult> { + if with_index_outputs.is_some() || pool_spec.kernel_shape.len() != 2 { + return Ok(None); + } + let facts = source.node_input_facts(node.id)?; + let Some(input_shape) = facts[0].shape.as_concrete() else { return Ok(None) }; + if input_shape.len() != 4 { + return Ok(None); + } + let dims: TVec = input_shape.iter().map(|d| d.to_dim()).collect(); + let geo = pool_spec.compute_geo(&dims)?.to_concrete(input_shape)?.into_owned(); + let (input_shape, output_shape) = (&geo.input_shape, &geo.output_shape); + let patch = &geo.patch; + if patch.pad_before.len() != 2 { + return Ok(None); + } + // The kernel subtracts one pad per axis, so an asymmetric `pad_after` is + // only reachable through the bounds check it already does. + let strides = |shape: &tract_core::ops::nn::DataShape| -> [usize; 4] { + [ + *shape.n_stride().unwrap_or(&0), + *shape.c_stride(), + shape.hw_strides()[0], + shape.hw_strides()[1], + ] + }; + let hw = |shape: &tract_core::ops::nn::DataShape| (shape.hw_dims()[0], shape.hw_dims()[1]); + + let geometry = MaxPool2dGeometry { + batch: *input_shape.n().unwrap_or(&1), + channels: *input_shape.c(), + input_hw: hw(input_shape), + output_hw: hw(output_shape), + kernel: (patch.spec.kernel_shape[0], patch.spec.kernel_shape[1]), + strides: (patch.spec.strides[0], patch.spec.strides[1]), + dilations: (patch.spec.dilations[0], patch.spec.dilations[1]), + padding: (patch.pad_before[0], patch.pad_before[1]), + input_strides: strides(input_shape), + output_strides: strides(output_shape), + }; + Ok(Some(GpuMaxPool::new( + geometry, + output_shape.shape.clone(), + "Metal", + metal_max_pool_dispatch, + ))) +} + +crate::register_metal_op!(MaxPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + rule_if!(facts[0].is_plain()); + rule_if!(MaxPool2d::is_supported_dt(facts[0].datum_type)); + Ok(geometry(source, node, &op.pool_spec, op.with_index_outputs)? + .map(|op| Box::new(op) as Box)) +}); + +crate::register_metal_op!(OptMaxPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + rule_if!(facts[0].is_plain()); + rule_if!(MaxPool2d::is_supported_dt(facts[0].datum_type)); + Ok(geometry(source, node, &op.pool_spec, op.with_index_outputs)? + .map(|op| Box::new(op) as Box)) +}); diff --git a/metal/src/kernels/nn/mod.rs b/metal/src/kernels/nn/mod.rs index 9ff4b13ea8..124d8e6ac1 100644 --- a/metal/src/kernels/nn/mod.rs +++ b/metal/src/kernels/nn/mod.rs @@ -1,6 +1,7 @@ pub mod apply_rope; pub mod gelu_approximate; pub mod leaky_relu; +pub mod max_pool; pub mod reduce; pub mod rms_norm; pub mod scaled_masked_softmax; @@ -12,6 +13,8 @@ pub use gelu_approximate::GeluApproximate; pub use gelu_approximate::metal_gelu_approximate_dispatch; pub use leaky_relu::LeakyRelu; pub use leaky_relu::metal_leaky_relu_dispatch; +pub use max_pool::MaxPool2d; +pub use max_pool::metal_max_pool_dispatch; pub use reduce::{Reducer, metal_reduce_launch}; pub use rms_norm::RmsNorm; pub use rms_norm::metal_rms_norm_dispatch; @@ -40,6 +43,12 @@ pub fn all_functions() -> Vec { .flat_map(|dt| Softmax.kernel_name(dt).into_iter()), ); + functions.extend( + tract_gpu::tensor::DeviceTensor::SUPPORTED_DT + .into_iter() + .flat_map(|dt| MaxPool2d.kernel_name(dt).into_iter()), + ); + functions.extend(tract_gpu::tensor::DeviceTensor::SUPPORTED_DT.into_iter().flat_map(|dt| { [false, true] .into_iter() diff --git a/metal/src/kernels/nn/nn_ops.metal b/metal/src/kernels/nn/nn_ops.metal index 5d27078345..6796e2276f 100644 --- a/metal/src/kernels/nn/nn_ops.metal +++ b/metal/src/kernels/nn/nn_ops.metal @@ -1,3 +1,4 @@ +#include #include #include @@ -785,3 +786,71 @@ template [[host_name("nn_ops::apply_rope_nd3_f16")]] [[kernel]] apply_rope_nd3_t apply_rope_nd3; template [[host_name("nn_ops::apply_rope_nd4_f16")]] [[kernel]] apply_rope_nd4_t apply_rope_nd4; + +// Max pooling over two spatial axes. Strides are passed for every axis, so +// NCHW and NHWC share this kernel; positions falling in the padding are skipped +// rather than compared, matching the CPU op's `min_value()` seed for a window +// that is entirely padding. Each thread owns one output element. +// +// params layout: [n, c, in_h, in_w, out_h, out_w, k_h, k_w, s_h, s_w, +// d_h, d_w, pad_h, pad_w, +// in_n, in_c, in_hs, in_ws, out_n, out_c, out_hs, out_ws] +template +[[kernel]] void max_pool_2d(device const void *input_b [[buffer(0)]], + device void *output_b [[buffer(1)]], + constant const int32_t *params [[buffer(2)]], + uint3 tpig [[thread_position_in_grid]]) { + const int32_t c = params[1]; + const int32_t in_h = params[2]; + const int32_t in_w = params[3]; + const int32_t out_h = params[4]; + const int32_t out_w = params[5]; + const int32_t k_h = params[6]; + const int32_t k_w = params[7]; + const int32_t s_h = params[8]; + const int32_t s_w = params[9]; + const int32_t d_h = params[10]; + const int32_t d_w = params[11]; + const int32_t pad_h = params[12]; + const int32_t pad_w = params[13]; + + const int32_t ow = (int32_t)tpig.x; + const int32_t oh = (int32_t)tpig.y; + const int32_t nc = (int32_t)tpig.z; + + if (ow >= out_w || oh >= out_h) + return; + + const int32_t n_idx = nc / c; + const int32_t c_idx = nc % c; + + device const T *input = (device const T *)input_b; + device T *output = (device T *)output_b; + + const int32_t in_base = n_idx * params[14] + c_idx * params[15]; + T best = numeric_limits::lowest(); + for (int32_t kh = 0; kh < k_h; ++kh) { + const int32_t ih = oh * s_h - pad_h + kh * d_h; + if (ih < 0 || ih >= in_h) + continue; + for (int32_t kw = 0; kw < k_w; ++kw) { + const int32_t iw = ow * s_w - pad_w + kw * d_w; + if (iw < 0 || iw >= in_w) + continue; + const T v = input[in_base + ih * params[16] + iw * params[17]]; + if (v > best) + best = v; + } + } + output[n_idx * params[18] + c_idx * params[19] + oh * params[20] + ow * params[21]] = best; +} + +typedef decltype(max_pool_2d) max_pool_2d_t; + +#define INSTANTIATE_MAX_POOL_2D(tname, type) \ + template [[host_name( \ + "nn_ops::max_pool_2d_" #tname)]] [[kernel]] max_pool_2d_t \ + max_pool_2d; + +INSTANTIATE_MAX_POOL_2D(f32, float) +INSTANTIATE_MAX_POOL_2D(f16, half)