tract_cuda/kernels/
conv.rs

1use crate::context::{TractCudaStream, cuda_context};
2use crate::kernels::launch_args::TractLaunchArgs;
3use crate::kernels::{WARP_SIZE, get_cuda_view};
4use cudarc::driver::{LaunchArgs, LaunchConfig, PushKernelArg};
5use downcast_rs::{Downcast, impl_downcast};
6use std::any::Any;
7use std::fmt::Debug;
8use tract_core::dyn_clone::{self, DynClone};
9use tract_core::internal::dyn_eq::DynEq;
10use tract_core::internal::*;
11use tract_core::ops::cnn::Conv;
12use tract_gpu::tensor::DeviceTensor;
13
14pub trait ConvKernelScratch: Debug + Downcast {}
15impl_downcast!(ConvKernelScratch);
16
17pub trait ConvKernel: 'static + Send + Sync + Debug + DynClone + DynEq {
18    fn name(&self) -> StaticName;
19    #[allow(clippy::too_many_arguments)]
20    fn state(&self) -> Box<dyn ConvKernelScratch>;
21    #[allow(clippy::too_many_arguments)]
22    fn dispatch(
23        &self,
24        state: &mut dyn ConvKernelScratch,
25        node_id: usize,
26        op: &Conv,
27        stream: &TractCudaStream,
28        input: &DeviceTensor,
29        weights: &DeviceTensor,
30        bias: Option<&DeviceTensor>,
31        output: &DeviceTensor,
32    ) -> TractResult<()>;
33}
34dyn_clone::clone_trait_object!(ConvKernel);
35dyn_eq::eq_trait_object!(ConvKernel);
36
37impl ConvKernelScratch for () {}
38
39#[derive(Hash, Clone, Debug, PartialEq, Eq)]
40pub struct ConvGeneric;
41
42impl ConvKernel for ConvGeneric {
43    fn name(&self) -> StaticName {
44        "Generic".into()
45    }
46
47    fn state(&self) -> Box<dyn ConvKernelScratch> {
48        Box::new(())
49    }
50
51    fn dispatch(
52        &self,
53        _state: &mut dyn ConvKernelScratch,
54        _node_id: usize,
55        op: &Conv,
56        stream: &TractCudaStream,
57        input: &DeviceTensor,
58        weights: &DeviceTensor,
59        bias: Option<&DeviceTensor>,
60        output: &DeviceTensor,
61    ) -> TractResult<()> {
62        let input_shape = op.pool_spec.data_format.shape(input.shape())?;
63
64        let ctx = cuda_context();
65        let dt_name = if input.datum_type() == DatumType::F16 { "f16" } else { "f32" };
66        let func_name = format!("conv{}d_{}_generic", input_shape.hw_rank(), dt_name);
67        let func = ctx.load_pipeline(crate::kernels::LibraryName::Cnn, func_name)?;
68        let null = stream.null::<u8>()?;
69        let null_view = null.as_view();
70
71        let mut launcher = TractLaunchArgs::new(stream, &func);
72
73        let input = get_cuda_view(input);
74
75        launcher.push_view(&input);
76        launcher.push_i32(*input_shape.n().unwrap_or(&1));
77        launcher.push_i32(*input_shape.c());
78        launcher.push_slice_i32(input_shape.hw_dims());
79
80        launcher.push_i32(*input_shape.n_stride().unwrap_or(&0));
81        launcher.push_i32(*input_shape.c_stride());
82        launcher.push_slice_i32(input_shape.hw_strides());
83
84        let kfmt = op.kernel_fmt;
85        let co_per_group = op.pool_spec.output_channels / op.group;
86        let ci_per_group = op.pool_spec.input_channels / op.group;
87
88        let weights_view = get_cuda_view(weights);
89        launcher.push_view(&weights_view);
90        // split go_i_h_w in g_o_i_h_w
91        launcher.push_i32(op.group);
92        launcher.push_i32(co_per_group);
93        launcher.push_slice_i32(&weights.shape()[1..]);
94
95        let group_stride = weights.strides()[0] as usize * co_per_group;
96        launcher.push_i32(group_stride);
97        launcher.push_slice_i32(weights.strides());
98
99        let mut bias_view = None;
100        if let Some(bias) = &bias {
101            bias_view = Some(get_cuda_view(bias));
102            launcher.push_view(bias_view.as_ref().unwrap());
103            launcher.push_i32(if bias.rank() == 0 {
104                0 // scalar bias: stride = 0 is broadcasting
105            } else {
106                1
107            });
108        } else {
109            launcher.push_view(&null_view);
110            launcher.push_i32(0);
111        }
112
113        let padding = op.pool_spec.computed_padding(input_shape.hw_dims());
114        for d in 0..input_shape.hw_rank() {
115            launcher.push_i32(padding[d].pad_before);
116        }
117
118        let strides = op.pool_spec.strides();
119        launcher.push_slice_i32(&strides);
120
121        let dilations = op.pool_spec.dilations();
122        launcher.push_slice_i32(&dilations);
123
124        let output_shape = op.pool_spec.data_format.shape(output.shape())?;
125        let output = get_cuda_view(output);
126        launcher.push_view(&output);
127        launcher.push_i32(*output_shape.n().unwrap_or(&1));
128        launcher.push_i32(*output_shape.c());
129        launcher.push_slice_i32(output_shape.hw_dims());
130
131        launcher.push_i32(*output_shape.n_stride().unwrap_or(&0));
132        launcher.push_i32(*output_shape.c_stride());
133        launcher.push_slice_i32(output_shape.hw_strides());
134
135        let cfg = LaunchConfig {
136            grid_dim: (
137                output_shape.hw_dims().iter().product::<usize>().div_ceil(WARP_SIZE) as u32,
138                *output_shape.c() as u32,
139                input_shape.n().copied().unwrap_or(1) as u32,
140            ),
141            block_dim: (WARP_SIZE as u32, 1, 1),
142            shared_mem_bytes: 0,
143        };
144
145        launcher.launch(cfg)
146    }
147}
tract_cuda/kernels/conv.rs

tract_cuda/kernels/
conv.rs
