diff --git a/core/src/ops/nn/resize.rs b/core/src/ops/nn/resize.rs index 9ba7a12149..cb5d2c8548 100644 --- a/core/src/ops/nn/resize.rs +++ b/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 { diff --git a/gpu/src/ops/mod.rs b/gpu/src/ops/mod.rs index ea5cdbc59b..1fc96e4fe2 100644 --- a/gpu/src/ops/mod.rs +++ b/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; 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 eb0973eb10..c3633442be 100644 --- a/metal/src/transform.rs +++ b/metal/src/transform.rs @@ -20,7 +20,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; @@ -281,6 +283,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())