diff --git a/core/src/ops/cnn/mod.rs b/core/src/ops/cnn/mod.rs index 540e975c03..0abc831a0a 100644 --- a/core/src/ops/cnn/mod.rs +++ b/core/src/ops/cnn/mod.rs @@ -11,12 +11,12 @@ 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}; pub use self::pools::PoolSpec; -pub use self::sumpool::SumPool; +pub use self::sumpool::{OptSumPool, SumPool}; use super::array::MultiBroadcastTo; diff --git a/metal/src/kernels/mod.rs b/metal/src/kernels/mod.rs index 036c030a64..8fce4884e9 100644 --- a/metal/src/kernels/mod.rs +++ b/metal/src/kernels/mod.rs @@ -31,7 +31,7 @@ const GGML: &str = include_str!("matmul/ggml_gemm/ggml_mm_mv.metal"); const BASIC_MAT_MUL: &str = include_str!("matmul/basic/basic_mat_mul.metal"); const ARRAY_OPS: &str = include_str!("array/array_ops.metal"); const BIN_OPS: &str = include_str!("bin_ops.metal"); -const NN_OPS: &str = include_str!("nn/nn_ops.metal"); +const NN_OPS: &str = concat!(include_str!("nn/nn_ops.metal"), include_str!("nn/pool.metal")); const CONV_OPS: &str = include_str!("conv.metal"); const ELEMENT_WISE_OPS: &str = include_str!("element_wise.metal"); const FFT_OPS: &str = include_str!("fft.metal"); diff --git a/metal/src/kernels/nn/mod.rs b/metal/src/kernels/nn/mod.rs index 9ff4b13ea8..8f4ffd99c7 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 pool; pub mod reduce; pub mod rms_norm; pub mod scaled_masked_softmax; diff --git a/metal/src/kernels/nn/pool.metal b/metal/src/kernels/nn/pool.metal new file mode 100644 index 0000000000..daa8a128e7 --- /dev/null +++ b/metal/src/kernels/nn/pool.metal @@ -0,0 +1,167 @@ +#include +using namespace metal; + +// 2D pooling over a channels-last tensor. One thread owns one (n, oh, ow, c), +// so consecutive threads walk the contiguous channel axis and every window read +// is coalesced. +// +// Buffer layout: +// 0: input [N, iH, iW, C] +// 1: output [N, oH, oW, C] +// 2: params (see PoolParams) +struct PoolParams { + int n; + int ih; + int iw; + int c; + int oh; + int ow; + int kh; + int kw; + int stride_h; + int stride_w; + int pad_h; + int pad_w; + int dil_h; + int dil_w; + // Divide the sum by the window area including padding, rather than by the + // number of positions that actually landed inside the input. + int count_include_pad; + int normalize; +}; + +template +[[kernel]] void max_pool_2d( + const device T* input [[buffer(0)]], + device T* output [[buffer(1)]], + const constant PoolParams& p [[buffer(2)]], + uint3 gid [[thread_position_in_grid]]) { + const int c = int(gid.x); + const int ow = int(gid.y); + const int rest = int(gid.z); + const int oh = rest % p.oh; + const int n = rest / p.oh; + if (c >= p.c || ow >= p.ow || n >= p.n) { + return; + } + + const int h_start = oh * p.stride_h - p.pad_h; + const int w_start = ow * p.stride_w - p.pad_w; + + // An all-padding window has no value to take, and tract's CPU op leaves + // -inf there too. + T best = T(-INFINITY); + for (int kh = 0; kh < p.kh; ++kh) { + const int ih = h_start + kh * p.dil_h; + if (ih < 0 || ih >= p.ih) { + continue; + } + for (int kw = 0; kw < p.kw; ++kw) { + const int iw = w_start + kw * p.dil_w; + if (iw < 0 || iw >= p.iw) { + continue; + } + const int64_t idx = + ((int64_t(n) * p.ih + ih) * p.iw + iw) * p.c + c; + best = max(best, input[idx]); + } + } + const int64_t out_idx = ((int64_t(n) * p.oh + oh) * p.ow + ow) * p.c + c; + output[out_idx] = best; +} + +template +[[kernel]] void sum_pool_2d( + const device T* input [[buffer(0)]], + device T* output [[buffer(1)]], + const constant PoolParams& p [[buffer(2)]], + uint3 gid [[thread_position_in_grid]]) { + const int c = int(gid.x); + const int ow = int(gid.y); + const int rest = int(gid.z); + const int oh = rest % p.oh; + const int n = rest / p.oh; + if (c >= p.c || ow >= p.ow || n >= p.n) { + return; + } + + const int h_start = oh * p.stride_h - p.pad_h; + const int w_start = ow * p.stride_w - p.pad_w; + + float acc = 0.0f; + int counted = 0; + for (int kh = 0; kh < p.kh; ++kh) { + const int ih = h_start + kh * p.dil_h; + if (ih < 0 || ih >= p.ih) { + continue; + } + for (int kw = 0; kw < p.kw; ++kw) { + const int iw = w_start + kw * p.dil_w; + if (iw < 0 || iw >= p.iw) { + continue; + } + const int64_t idx = + ((int64_t(n) * p.ih + ih) * p.iw + iw) * p.c + c; + acc += float(input[idx]); + counted += 1; + } + } + if (p.normalize) { + const int divisor = p.count_include_pad ? (p.kh * p.kw) : max(counted, 1); + acc /= float(divisor); + } + const int64_t out_idx = ((int64_t(n) * p.oh + oh) * p.ow + ow) * p.c + c; + output[out_idx] = T(acc); +} + +// Channels-first variant: the contiguous axis is now width, so one thread owns +// one (n, c, oh, ow) and consecutive threads walk the row. +template +[[kernel]] void max_pool_2d_nchw( + const device T* input [[buffer(0)]], + device T* output [[buffer(1)]], + const constant PoolParams& p [[buffer(2)]], + uint3 gid [[thread_position_in_grid]]) { + const int ow = int(gid.x); + const int oh = int(gid.y); + const int rest = int(gid.z); + const int c = rest % p.c; + const int n = rest / p.c; + if (ow >= p.ow || oh >= p.oh || n >= p.n) { + return; + } + + const int h_start = oh * p.stride_h - p.pad_h; + const int w_start = ow * p.stride_w - p.pad_w; + const int64_t plane = (int64_t(n) * p.c + c); + + T best = T(-INFINITY); + for (int kh = 0; kh < p.kh; ++kh) { + const int ih = h_start + kh * p.dil_h; + if (ih < 0 || ih >= p.ih) { + continue; + } + for (int kw = 0; kw < p.kw; ++kw) { + const int iw = w_start + kw * p.dil_w; + if (iw < 0 || iw >= p.iw) { + continue; + } + best = max(best, input[(plane * p.ih + ih) * p.iw + iw]); + } + } + output[(plane * p.oh + oh) * p.ow + ow] = best; +} + +#define instantiate_pool(name, tname, itype) \ + template [[host_name(#name "_" #tname)]] [[kernel]] void name( \ + const device itype* input [[buffer(0)]], \ + device itype* output [[buffer(1)]], \ + const constant PoolParams& p [[buffer(2)]], \ + uint3 gid [[thread_position_in_grid]]); + +instantiate_pool(max_pool_2d, f32, float) +instantiate_pool(max_pool_2d, f16, half) +instantiate_pool(sum_pool_2d, f32, float) +instantiate_pool(sum_pool_2d, f16, half) +instantiate_pool(max_pool_2d_nchw, f32, float) +instantiate_pool(max_pool_2d_nchw, f16, half) diff --git a/metal/src/kernels/nn/pool.rs b/metal/src/kernels/nn/pool.rs new file mode 100644 index 0000000000..c76336b937 --- /dev/null +++ b/metal/src/kernels/nn/pool.rs @@ -0,0 +1,329 @@ +//! 2D max and sum pooling on channels-last tensors (see pool.metal). +//! +//! Pooling had no Metal kernel, so a pooled model bounced back to the host at +//! every pool: on Inception v3 that was 14 device syncs and a fifth of the +//! runtime. + +use crate::encoder::EncoderExt; +use crate::{LibraryName, MetalStream}; +use anyhow::ensure; +use metal::MTLSize; +use tract_core::internal::*; +use tract_core::ops::cnn::PoolSpec; +use tract_gpu::tensor::DeviceTensor; + +/// Mirror of `PoolParams` in pool.metal — keep field order in sync. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +struct PoolParams { + n: i32, + ih: i32, + iw: i32, + c: i32, + oh: i32, + ow: i32, + kh: i32, + kw: i32, + stride_h: i32, + stride_w: i32, + pad_h: i32, + pad_w: i32, + dil_h: i32, + dil_w: i32, + count_include_pad: i32, + normalize: i32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PoolKind { + Max, + /// `normalize` turns the sum into an average. + Sum { + count_include_pad: bool, + normalize: bool, + }, +} + +/// Whether the kernel covers this pooling: rank-4 f16/f32 with static geometry +/// and no index output. Max pooling has a channels-first variant too; sum +/// pooling stays channels-last. +pub fn metal_pool_supported(pool_spec: &PoolSpec, kind: PoolKind, fact: &TypedFact) -> bool { + let layout_ok = pool_spec.data_format.c_is_last() + || (kind == PoolKind::Max && pool_spec.data_format.has_n()); + matches!(fact.datum_type, DatumType::F16 | DatumType::F32) + && layout_ok + && fact.rank() == 4 + && pool_spec.kernel_shape.len() == 2 + && fact.shape.as_concrete().is_some() +} + +pub fn dispatch_metal_pool( + stream: &MetalStream, + pool_spec: &PoolSpec, + kind: PoolKind, + input: &DeviceTensor, + output: &DeviceTensor, +) -> TractResult<()> { + let dt = input.datum_type(); + let tname = match dt { + DatumType::F32 => "f32", + DatumType::F16 => "f16", + _ => bail!("Metal pool: F32/F16 only, got {dt:?}"), + }; + let in_shape = pool_spec.data_format.shape(input.shape())?; + let out_shape = pool_spec.data_format.shape(output.shape())?; + ensure!(in_shape.hw_rank() == 2, "Metal pool is 2D only"); + + let strides = pool_spec.strides(); + let dilations = pool_spec.dilations(); + let padding = pool_spec.computed_padding(in_shape.hw_dims()); + let (count_include_pad, normalize) = match kind { + PoolKind::Max => (false, false), + PoolKind::Sum { count_include_pad, normalize } => (count_include_pad, normalize), + }; + let params = PoolParams { + n: *in_shape.n().unwrap_or(&1) as i32, + ih: in_shape.hw_dims()[0] as i32, + iw: in_shape.hw_dims()[1] as i32, + c: *in_shape.c() as i32, + oh: out_shape.hw_dims()[0] as i32, + ow: out_shape.hw_dims()[1] as i32, + kh: pool_spec.kernel_shape[0] as i32, + kw: pool_spec.kernel_shape[1] as i32, + stride_h: strides[0] as i32, + stride_w: strides[1] as i32, + pad_h: padding[0].pad_before as i32, + pad_w: padding[1].pad_before as i32, + dil_h: dilations[0] as i32, + dil_w: dilations[1] as i32, + count_include_pad: count_include_pad as i32, + normalize: normalize as i32, + }; + + let channels_last = pool_spec.data_format.c_is_last(); + let base = match kind { + PoolKind::Max if channels_last => "max_pool_2d", + PoolKind::Max => "max_pool_2d_nchw", + PoolKind::Sum { .. } => "sum_pool_2d", + }; + ensure!(channels_last || kind == PoolKind::Max, "Metal sum pool is channels-last only"); + let pipeline = stream.load_pipeline(LibraryName::NNOps, &format!("{base}_{tname}"))?; + + stream.retain_tensor(input); + stream.retain_tensor(output); + + 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, std::slice::from_ref(¶ms)); + let (fastest, height, depth) = if channels_last { + (params.c, params.ow, params.oh * params.n) + } else { + (params.ow, params.oh, params.c * params.n) + }; + let group_w = 32u64.min(fastest as u64).max(1); + encoder.dispatch_threads( + MTLSize { width: fastest as _, height: height as _, depth: depth as _ }, + MTLSize { width: group_w, height: 1, depth: 1 }, + ); + }); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::MetalTransform; + use crate::utils::with_borrowed_metal_stream; + use tract_core::ops::cnn::{MaxPool, PaddingSpec, SumPool}; + use tract_core::ops::nn::DataFormat; + use tract_core::transform::ModelTransform; + use tract_gpu::tensor::IntoDevice; + + fn ramp(dt: DatumType, shape: &[usize]) -> TractResult { + let len: usize = shape.iter().product(); + let v: Vec = (0..len).map(|i| ((i * 17 % 61) as f32 - 30.0) / 8.0).collect(); + Ok(Tensor::from_shape(shape, &v)?.cast_to_dt(dt)?.into_owned()) + } + + fn spec(k: usize, stride: usize, padding: PaddingSpec, c: usize) -> PoolSpec { + spec_fmt(DataFormat::NHWC, k, stride, padding, c) + } + + fn spec_fmt( + data_format: DataFormat, + k: usize, + stride: usize, + padding: PaddingSpec, + c: usize, + ) -> PoolSpec { + PoolSpec::new(data_format, tvec![k, k], padding, None, Some(tvec![stride, stride]), c, c) + } + + /// Run the CPU op and the Metal kernel over the same input. + fn check( + dt: DatumType, + shape: &[usize], + pool_spec: PoolSpec, + kind: PoolKind, + cpu: Box, + ) -> TractResult<()> { + let input = ramp(dt, shape)?; + let mut model = TypedModel::default(); + let i = model.add_source("i", dt.fact(shape))?; + let out = model.wire_node("pool", cpu, &[i])?; + model.select_output_outlets(&out)?; + let expected = model + .clone() + .into_optimized()? + .into_runnable()? + .run(tvec![input.clone().into_tvalue()])?[0] + .clone() + .into_tensor(); + let got = with_borrowed_metal_stream(|stream| { + let i = input.clone().into_device()?; + let o_shape = pool_spec.output_shape(&input.shape().to_vec())?; + let o = unsafe { DeviceTensor::uninitialized_dt(dt, &o_shape.shape)? }; + dispatch_metal_pool(stream, &pool_spec, kind, &i, &o)?; + stream.wait_until_completed()?; + Ok(o.to_host()?.into_tensor()) + })?; + expected + .close_enough(&got, Approximation::Approximate) + .with_context(|| format!("{kind:?} dt={dt:?} shape={shape:?}")) + } + + #[test] + fn max_pool_3x3_stride2() -> TractResult<()> { + let s = spec(3, 2, PaddingSpec::Valid, 32); + check( + DatumType::F32, + &[1, 16, 16, 32], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn max_pool_nchw() -> TractResult<()> { + let s = spec_fmt(DataFormat::NCHW, 2, 2, PaddingSpec::Valid, 24); + check( + DatumType::F32, + &[1, 24, 16, 16], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn max_pool_nchw_same_padding() -> TractResult<()> { + let s = spec_fmt(DataFormat::NCHW, 3, 2, PaddingSpec::SameUpper, 12); + check( + DatumType::F16, + &[2, 12, 11, 13], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn max_pool_same_padding() -> TractResult<()> { + let s = spec(3, 1, PaddingSpec::SameUpper, 48); + check( + DatumType::F32, + &[1, 13, 11, 48], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn avg_pool_3x3_same() -> TractResult<()> { + let s = spec(3, 1, PaddingSpec::SameUpper, 64); + check( + DatumType::F32, + &[1, 12, 12, 64], + s.clone(), + PoolKind::Sum { count_include_pad: false, normalize: true }, + Box::new(SumPool::new(s, false, true)), + ) + } + + #[test] + fn avg_pool_count_include_pad() -> TractResult<()> { + let s = spec(3, 2, PaddingSpec::SameUpper, 16); + check( + DatumType::F32, + &[2, 9, 9, 16], + s.clone(), + PoolKind::Sum { count_include_pad: true, normalize: true }, + Box::new(SumPool::new(s, true, true)), + ) + } + + #[test] + fn sum_pool_no_normalize() -> TractResult<()> { + let s = spec(2, 2, PaddingSpec::Valid, 32); + check( + DatumType::F32, + &[1, 8, 8, 32], + s.clone(), + PoolKind::Sum { count_include_pad: false, normalize: false }, + Box::new(SumPool::new(s, false, false)), + ) + } + + #[test] + fn pool_f16() -> TractResult<()> { + let s = spec(3, 2, PaddingSpec::SameUpper, 96); + check( + DatumType::F16, + &[1, 14, 14, 96], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + // The pools must actually land on the GPU once the transform has run. + #[test] + fn pools_route_through_metal_transform() -> TractResult<()> { + let dt = DatumType::F32; + let shape = [1usize, 14, 14, 32]; + let input = ramp(dt, &shape)?; + let mut model = TypedModel::default(); + let i = model.add_source("i", dt.fact(&shape))?; + let m = model.wire_node( + "max", + MaxPool::new(spec(3, 1, PaddingSpec::SameUpper, 32), None), + &[i], + )?[0]; + let a = model.wire_node( + "avg", + SumPool::new(spec(3, 1, PaddingSpec::SameUpper, 32), false, true), + &[m], + )?; + model.select_output_outlets(&a)?; + let cpu = model + .clone() + .into_optimized()? + .into_runnable()? + .run(tvec![input.clone().into_tvalue()])?; + let metal = MetalTransform::default().transform_into(model.into_optimized()?)?; + let n_pools = + metal.nodes().iter().filter(|n| n.op_is::()).count(); + assert_eq!(n_pools, 2, "both pools should be on the GPU"); + let got = metal.into_runnable()?.run(tvec![input.into_tvalue()])?; + cpu[0] + .clone() + .into_tensor() + .close_enough(&got[0].clone().into_tensor(), Approximation::Approximate)?; + Ok(()) + } +} diff --git a/metal/src/ops/mod.rs b/metal/src/ops/mod.rs index 4c0defe450..26306fe79e 100644 --- a/metal/src/ops/mod.rs +++ b/metal/src/ops/mod.rs @@ -1,6 +1,7 @@ pub mod conv; pub mod fused_axis_op; pub mod gemm; +pub mod pool; pub use fused_axis_op::MetalFusedAxisOp; pub use gemm::MetalGemm; diff --git a/metal/src/ops/pool.rs b/metal/src/ops/pool.rs new file mode 100644 index 0000000000..08f3c628ba --- /dev/null +++ b/metal/src/ops/pool.rs @@ -0,0 +1,137 @@ +use crate::kernels::nn::pool::{PoolKind, dispatch_metal_pool, metal_pool_supported}; +use tract_core::internal::*; +use tract_core::ops::cnn::{MaxPool, OptMaxPool, OptSumPool, PoolSpec, SumPool}; +use tract_gpu::tensor::DeviceTensorExt; + +/// Metal device op for `OptMaxPool` / `OptSumPool` over a channels-last tensor. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct MetalPool { + pub pool_spec: PoolSpec, + pub kind: PoolKindOp, +} + +/// `PoolKind` without the borrowed geometry, so the op stays hashable. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum PoolKindOp { + Max, + Sum { count_include_pad: bool, normalize: bool }, +} + +impl From for PoolKind { + fn from(k: PoolKindOp) -> Self { + match k { + PoolKindOp::Max => PoolKind::Max, + PoolKindOp::Sum { count_include_pad, normalize } => { + PoolKind::Sum { count_include_pad, normalize } + } + } + } +} + +impl Op for MetalPool { + fn name(&self) -> StaticName { + match self.kind { + PoolKindOp::Max => "MetalMaxPool".into(), + PoolKindOp::Sum { .. } => "MetalSumPool".into(), + } + } + + fn info(&self) -> TractResult> { + Ok(self.pool_spec.info()) + } + + op_as_typed_op!(); +} + +impl EvalOp for MetalPool { + 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_shape = self.pool_spec.output_shape(input.shape())?; + let output = tract_gpu::session_handler::make_tensor_for_node( + session, + node_id, + input.datum_type(), + &output_shape.shape, + )?; + if output.len() > 0 { + crate::with_metal_stream(|stream| { + dispatch_metal_pool(stream, &self.pool_spec, self.kind.into(), input, &output) + })?; + } + Ok(tvec!(output.into_tensor().into_tvalue())) + } +} + +impl TypedOp for MetalPool { + as_op!(); + + fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult> { + tract_gpu::utils::facts_to_device_facts(inputs, |facts| { + let shape = self.pool_spec.output_shape(&facts[0].shape)?; + Ok(tvec!(facts[0].datum_type.fact(shape.shape))) + }) + .with_context(|| "Error while computing facts for MetalPool") + } +} + +crate::register_metal_op!(OptMaxPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + // an index output would need a second buffer the kernel does not write + if op.with_index_outputs.is_some() + || !metal_pool_supported(&op.pool_spec, PoolKind::Max, facts[0]) + { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { pool_spec: op.pool_spec.clone(), kind: PoolKindOp::Max }) + as Box)) +}); + +crate::register_metal_op!(OptSumPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + let kind = PoolKind::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }; + if !metal_pool_supported(&op.pool_spec, kind, facts[0]) { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { + pool_spec: op.pool_spec.clone(), + kind: PoolKindOp::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }, + }) as Box)) +}); + +// The metal transform runs before optimization, so a model still carries the +// unoptimized pools; the Opt* forms are registered too for callers that +// optimize first. +crate::register_metal_op!(MaxPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + if op.with_index_outputs.is_some() + || !metal_pool_supported(&op.pool_spec, PoolKind::Max, facts[0]) + { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { pool_spec: op.pool_spec.clone(), kind: PoolKindOp::Max }) + as Box)) +}); + +crate::register_metal_op!(SumPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + let kind = PoolKind::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }; + if !metal_pool_supported(&op.pool_spec, kind, facts[0]) { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { + pool_spec: op.pool_spec.clone(), + kind: PoolKindOp::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }, + }) as Box)) +}); + +/// Referenced from the transform so the registrations below survive linking. +pub fn link_translators() {} diff --git a/metal/src/transform.rs b/metal/src/transform.rs index c3633442be..4dfd82fb0e 100644 --- a/metal/src/transform.rs +++ b/metal/src/transform.rs @@ -146,6 +146,11 @@ impl ModelTransform for MetalTransform { } fn transform(&self, model: &mut TypedModel) -> TractResult<()> { + // The pool translators live in `ops::pool`, which nothing else calls + // into; without a reference the linker drops the module and with it the + // inventory registrations. + crate::ops::pool::link_translators(); + self.transform_up_to_phase(model, usize::MAX) } }