Resize PR 3: Run Resize on the Metal GPU

core/src/ops/nn/resize.rs

@@ -356,7 +356,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 {

gpu/src/ops/mod.rs

@@ -17,6 +17,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;

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<usize>,
+    pub windows: TVec<usize>,
+    pub plans: TVec<(Arc<Tensor>, Arc<Tensor>)>,
+    pub output_shape: TVec<usize>,
+    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<H: std::hash::Hasher>(&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<Vec<String>> {
+        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<TValue>,
+    ) -> TractResult<TVec<TValue>> {
+        let data = inputs[0].to_device_tensor()?;
+        let dt = data.datum_type();
+        let mut shape: TVec<usize> = 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)(&current, 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<TVec<TypedFact>> {
+        crate::utils::facts_to_device_facts(inputs, |facts| {
+            ensure!(facts.len() == 1);
+            let shape: TVec<TDim> = 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!();
+}

metal/src/kernels/array/array_ops.metal

@@ -336,6 +336,51 @@ typedef decltype(gather<float>) gather_t;
     template [[host_name(                                                      \
         "array_ops::gather_" #tname)]] [[kernel]] gather_t gather<type>;
 
+// 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 <typename T>
+[[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<float>) resize_axis_t;
+
+#define INSTANTIATE_RESIZE_AXIS(tname, type)                                   \
+    template [[host_name(                                                      \
+        "array_ops::resize_axis_" #tname)]] [[kernel]] resize_axis_t           \
+        resize_axis<type>;
+
 // 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)

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<String> {
             .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()
 }

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<String> {
+        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, &params);
+            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<Option<GpuResize>> {
+    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<usize>| 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<usize> =
+        op.compute_output_shape(input_shape, scales_konst, sizes_konst)?;
+
+    let scales: TVec<f32> = match scales_konst.filter(|s| s.len() == input_shape.len()) {
+        Some(scales) => scales.cast_to::<f32>()?.try_as_plain()?.as_slice::<f32>()?.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<i32> = 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<Option<Box<dyn TypedOp>>> {
+    let Some(op) = node.op_as::<CoreResize>() 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<dyn TypedOp>))
+}