tract-metal 0.23.6

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

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<PoolKindOp> 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<Vec<String>> {
        Ok(self.pool_spec.info())
    }

    op_as_typed_op!();
}

impl EvalOp for MetalPool {
    op_out_of_plan!();

    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
        let input = inputs[0].to_device_tensor()?;
        let output_shape = self.pool_spec.output_shape(input.shape())?;
        let output = tract_gpu::turn_handler::make_tensor_for_node(
            ctx,
            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<TVec<TypedFact>> {
        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<dyn TypedOp>))
});

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<dyn TypedOp>))
});

// 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<dyn TypedOp>))
});

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<dyn TypedOp>))
});

/// Referenced from the transform so the registrations below survive linking.
pub fn link_translators() {}