tract_cuda/ops/
conv.rs

1use crate::kernels::conv::{ConvGeneric, ConvKernel, ConvKernelScratch};
2use crate::kernels::conv_cudnn::ConvCudnn;
3use std::cell::RefCell;
4use std::collections::HashMap;
5use tract_core::internal::*;
6use tract_core::ops::cnn::Conv;
7use tract_gpu::ops::change_axes::GpuAxisOp;
8use tract_gpu::tensor::DeviceTensorExt;
9
10pub fn wire_cuda_conv(
11    source: &TypedModel,
12    node: &TypedNode,
13    target: &mut TypedModel,
14    inputs: &[OutletId],
15    op: &Conv,
16) -> TractResult<TVec<OutletId>> {
17    let facts = source.node_input_facts(node.id)?;
18    let data_shape = op.pool_spec.data_format.shape(&facts[0].shape)?;
19    let hw_rank = data_shape.hw_rank();
20    let is_f16 = facts[0].datum_type.is::<f16>();
21    if facts.iter().all(|f| f.datum_type.is::<f32>() || f.datum_type.is::<f16>())
22        && hw_rank <= if is_f16 { 2 } else { 6 }
23        && op
24            .pool_spec
25            .computed_padding(data_shape.hw_dims())
26            .iter()
27            .all(|paddings| paddings.pad_before == paddings.pad_after)
28    {
29        let prefix = &node.name;
30        let bias = &facts[2];
31        let need_bias = !(bias.konst.is_some() && bias.konst.as_ref().unwrap().is_all_zero()?);
32        let conv_name = format!("{prefix}.conv");
33        let mut conv_wire = target.wire_node(
34            if need_bias { &conv_name } else { &node.name },
35            CudaConv { op: op.clone(), kernel: Box::new(ConvCudnn) },
36            &inputs[0..2],
37        )?[0];
38        if need_bias {
39            let mut needed_shape = tvec![1.to_dim(); node.outputs[0].fact.rank()];
40            needed_shape[data_shape.c_axis()] = op.pool_spec.output_channels.to_dim();
41            let reshaped = target.wire_node(
42                format!("{prefix}.bias_reshaped"),
43                GpuAxisOp::new(AxisOp::Reshape(0, bias.shape.to_tvec(), needed_shape)),
44                &[inputs[2]],
45            )?[0];
46            conv_wire = target.wire_node(
47                prefix,
48                crate::kernels::binary::cuda_bin_op(Box::new(tract_core::ops::math::Add)),
49                &[conv_wire, reshaped],
50            )?[0];
51        }
52        Ok(tvec!(conv_wire))
53    } else {
54        target.wire_node(
55            &node.name,
56            CudaConv { op: op.clone(), kernel: Box::new(ConvGeneric) },
57            inputs,
58        )
59    }
60}
61
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct CudaConv {
64    op: Conv,
65    kernel: Box<dyn ConvKernel>,
66}
67
68impl Op for CudaConv {
69    fn name(&self) -> StaticName {
70        "CudaConv".into()
71    }
72
73    fn info(&self) -> TractResult<Vec<String>> {
74        let mut info = self.op.info()?;
75        info.push(format!("kernel: {}", self.kernel.name()));
76        Ok(info)
77    }
78
79    op_as_typed_op!();
80}
81
82impl TypedOp for CudaConv {
83    as_op!();
84
85    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
86        tract_gpu::utils::facts_to_device_facts(inputs, |facts| {
87            let zero = facts[0].datum_type.scalar_fact();
88            let mut facts: TVec<&TypedFact> = facts.into();
89            if facts.len() == 2 {
90                facts.push(&zero);
91            }
92            self.op.output_facts(&facts)
93        })
94        .with_context(|| format!("Error while computing facts for Conv/{:?}", self.kernel.name()))
95    }
96}
97
98// cudnn descriptors are not Send and carry no meaning the model depends on, so
99// they are scratch the op keeps itself: a thread-local keyed by the session and
100// node the EvalContext names, emptied for a session in `drop_session`.
101thread_local! {
102    static CUDA_CONV_SCRATCH: RefCell<HashMap<(SessionId, usize), Box<dyn ConvKernelScratch>>> =
103        RefCell::new(HashMap::new());
104}
105
106impl EvalOp for CudaConv {
107    not_out_of_plan!();
108
109    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
110        let inputs =
111            inputs.iter().map(|it| it.to_device_tensor()).collect::<TractResult<TVec<_>>>()?;
112        let output_shape = self.op.pool_spec.output_shape(inputs[0].shape())?;
113        let output = tract_gpu::turn_handler::make_tensor_for_node(
114            ctx,
115            inputs[0].datum_type(),
116            &output_shape.shape,
117        )?;
118
119        if output.len() > 0 {
120            crate::with_cuda_stream(|stream| {
121                CUDA_CONV_SCRATCH.with_borrow_mut(|cache| {
122                    let scratch = cache
123                        .entry((ctx.session, ctx.node_id))
124                        .or_insert_with(|| self.kernel.state());
125                    self.kernel.dispatch(
126                        &mut **scratch,
127                        ctx.node_id,
128                        &self.op,
129                        stream,
130                        inputs[0],
131                        inputs[1],
132                        inputs.get(2).cloned(),
133                        &output,
134                    )
135                })
136            })?;
137        }
138        Ok(tvec!(output.into_tensor().into_tvalue()))
139    }
140
141    fn drop_session(&self, session: SessionId, node_id: usize) {
142        CUDA_CONV_SCRATCH.with_borrow_mut(|cache| cache.remove(&(session, node_id)));
143    }
144}
tract_cuda/ops/conv.rs

tract_cuda/ops/
conv.rs
