ruda_tensor_device/dispatch/
neural.rs

1use crate::{DeviceBackend, DeviceRuntime, FloatElement, IntElement, element::BoolElement};
2use rudnn::convolution::tensor::ConvTranspose2dStrategy;
3use ruda_tensor::tensor::{BoolTensor, FloatTensor, IntTensor};
4use ruda_tensor::{
5    TensorMetadata,
6    ops::{
7        AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward,
8        DeformConvOptions, InterpolateOptions, MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps,
9    },
10};
11
12impl<R, F, I, BT> ModuleOps<Self> for DeviceBackend<R, F, I, BT>
13where
14    R: DeviceRuntime,
15    F: FloatElement,
16    I: IntElement,
17    BT: BoolElement,
18{
19    fn conv1d(
20        x: FloatTensor<Self>,
21        weight: FloatTensor<Self>,
22        bias: Option<FloatTensor<Self>>,
23        options: ConvOptions<1>,
24    ) -> FloatTensor<Self> {
25        rudnn::convolution::tensor::conv_forward::<R, 1>(x, weight, bias, options, Default::default()).unwrap()
26    }
27
28    fn conv1d_x_backward(
29        x: FloatTensor<Self>,
30        weight: FloatTensor<Self>,
31        output_grad: FloatTensor<Self>,
32        options: ConvOptions<1>,
33    ) -> FloatTensor<Self> {
34        rudnn::convolution::tensor::conv_data_backward(
35            output_grad,
36            weight,
37            x.shape(),
38            options,
39            Default::default(),
40        )
41        .unwrap()
42    }
43
44    fn conv1d_weight_backward(
45        x: FloatTensor<Self>,
46        weight: FloatTensor<Self>,
47        output_grad: FloatTensor<Self>,
48        options: ConvOptions<1>,
49    ) -> FloatTensor<Self> {
50        rudnn::convolution::tensor::conv_weight_backward::<R, 1>(
51            x,
52            output_grad,
53            weight.shape(),
54            options,
55            Default::default(),
56        )
57        .unwrap()
58    }
59
60    fn conv2d(
61        x: FloatTensor<Self>,
62        weight: FloatTensor<Self>,
63        bias: Option<FloatTensor<Self>>,
64        options: ConvOptions<2>,
65    ) -> FloatTensor<Self> {
66        rudnn::convolution::tensor::conv_forward::<R, 2>(x, weight, bias, options, Default::default()).unwrap()
67    }
68
69    fn conv2d_x_backward(
70        x: FloatTensor<Self>,
71        weight: FloatTensor<Self>,
72        output_grad: FloatTensor<Self>,
73        options: ConvOptions<2>,
74    ) -> FloatTensor<Self> {
75        rudnn::convolution::tensor::conv_data_backward(
76            output_grad,
77            weight,
78            x.shape(),
79            options,
80            Default::default(),
81        )
82        .unwrap()
83    }
84
85    fn conv2d_weight_backward(
86        x: FloatTensor<Self>,
87        weight: FloatTensor<Self>,
88        output_grad: FloatTensor<Self>,
89        options: ConvOptions<2>,
90    ) -> FloatTensor<Self> {
91        rudnn::convolution::tensor::conv_weight_backward::<R, 2>(
92            x,
93            output_grad,
94            weight.shape(),
95            options,
96            Default::default(),
97        )
98        .unwrap()
99    }
100
101    fn deform_conv2d(
102        x: FloatTensor<Self>,
103        offset: FloatTensor<Self>,
104        weight: FloatTensor<Self>,
105        mask: Option<FloatTensor<Self>>,
106        bias: Option<FloatTensor<Self>>,
107        options: DeformConvOptions<2>,
108    ) -> FloatTensor<Self> {
109        rudnn::convolution::tensor::deform_conv2d(x, offset, weight, mask, bias, options).unwrap()
110    }
111
112    fn deform_conv2d_backward(
113        x: FloatTensor<Self>,
114        offset: FloatTensor<Self>,
115        weight: FloatTensor<Self>,
116        mask: Option<FloatTensor<Self>>,
117        bias: Option<FloatTensor<Self>>,
118        output_grad: FloatTensor<Self>,
119        options: DeformConvOptions<2>,
120    ) -> DeformConv2dBackward<Self> {
121        let (x, o, w, m, b) = rudnn::convolution::tensor::deform_conv2d_backward(
122            x,
123            offset,
124            weight,
125            mask,
126            bias,
127            output_grad,
128            options,
129        )
130        .unwrap();
131        DeformConv2dBackward::new(x, o, w, m, b)
132    }
133
134    fn conv3d(
135        x: FloatTensor<Self>,
136        weight: FloatTensor<Self>,
137        bias: Option<FloatTensor<Self>>,
138        options: ConvOptions<3>,
139    ) -> FloatTensor<Self> {
140        rudnn::convolution::tensor::conv_forward::<R, 3>(x, weight, bias, options, Default::default()).unwrap()
141    }
142
143    fn conv3d_x_backward(
144        x: FloatTensor<Self>,
145        weight: FloatTensor<Self>,
146        output_grad: FloatTensor<Self>,
147        options: ConvOptions<3>,
148    ) -> FloatTensor<Self> {
149        rudnn::convolution::tensor::conv_data_backward(
150            output_grad,
151            weight,
152            x.shape(),
153            options,
154            Default::default(),
155        )
156        .unwrap()
157    }
158
159    fn conv3d_weight_backward(
160        x: FloatTensor<Self>,
161        weight: FloatTensor<Self>,
162        output_grad: FloatTensor<Self>,
163        options: ConvOptions<3>,
164    ) -> FloatTensor<Self> {
165        rudnn::convolution::tensor::conv_weight_backward::<R, 3>(
166            x,
167            output_grad,
168            weight.shape(),
169            options,
170            Default::default(),
171        )
172        .unwrap()
173    }
174
175    fn conv_transpose2d(
176        x: FloatTensor<Self>,
177        weight: FloatTensor<Self>,
178        bias: Option<FloatTensor<Self>>,
179        options: ConvTransposeOptions<2>,
180    ) -> FloatTensor<Self> {
181        rudnn::convolution::tensor::conv_transpose2d(x, weight, bias, options, ConvTranspose2dStrategy::default())
182            .unwrap()
183    }
184
185    fn conv_transpose3d(
186        x: FloatTensor<Self>,
187        weight: FloatTensor<Self>,
188        bias: Option<FloatTensor<Self>>,
189        options: ConvTransposeOptions<3>,
190    ) -> FloatTensor<Self> {
191        rudnn::convolution::tensor::conv_transpose3d(x, weight, bias, options).expect("Kernel to never fail")
192    }
193
194    fn avg_pool2d(
195        x: FloatTensor<Self>,
196        kernel_size: [usize; 2],
197        stride: [usize; 2],
198        padding: [usize; 2],
199        count_include_pad: bool,
200        ceil_mode: bool,
201    ) -> FloatTensor<Self> {
202        rudnn::pooling::avg_pool2d(
203            x,
204            kernel_size,
205            stride,
206            padding,
207            count_include_pad,
208            ceil_mode,
209        )
210    }
211
212    fn avg_pool2d_backward(
213        x: FloatTensor<Self>,
214        grad: FloatTensor<Self>,
215        kernel_size: [usize; 2],
216        stride: [usize; 2],
217        padding: [usize; 2],
218        count_include_pad: bool,
219        ceil_mode: bool,
220    ) -> FloatTensor<Self> {
221        rudnn::pooling::avg_pool2d_backward(
222            x,
223            grad,
224            kernel_size,
225            stride,
226            padding,
227            count_include_pad,
228            ceil_mode,
229        )
230    }
231
232    fn max_pool2d(
233        x: FloatTensor<Self>,
234        kernel_size: [usize; 2],
235        stride: [usize; 2],
236        padding: [usize; 2],
237        dilation: [usize; 2],
238        ceil_mode: bool,
239    ) -> FloatTensor<Self> {
240        rudnn::pooling::max_pool2d(x, kernel_size, stride, padding, dilation, ceil_mode)
241    }
242
243    fn max_pool2d_with_indices(
244        x: FloatTensor<Self>,
245        kernel_size: [usize; 2],
246        stride: [usize; 2],
247        padding: [usize; 2],
248        dilation: [usize; 2],
249        ceil_mode: bool,
250    ) -> MaxPool2dWithIndices<Self> {
251        let (output, indices) = rudnn::pooling::max_pool2d_with_indices(
252            x,
253            kernel_size,
254            stride,
255            padding,
256            dilation,
257            ceil_mode,
258            I::dtype(),
259        );
260
261        MaxPool2dWithIndices::new(output, indices)
262    }
263
264    fn max_pool2d_with_indices_backward(
265        x: FloatTensor<Self>,
266        kernel_size: [usize; 2],
267        stride: [usize; 2],
268        padding: [usize; 2],
269        dilation: [usize; 2],
270        ceil_mode: bool,
271        output_grad: FloatTensor<Self>,
272        indices: IntTensor<Self>,
273    ) -> MaxPool2dBackward<Self> {
274        MaxPool2dBackward::new(rudnn::pooling::max_pool2d_with_indices_backward(
275            x,
276            output_grad,
277            indices,
278            kernel_size,
279            stride,
280            padding,
281            dilation,
282            ceil_mode,
283        ))
284    }
285
286    fn adaptive_avg_pool2d(x: FloatTensor<Self>, output_size: [usize; 2]) -> FloatTensor<Self> {
287        rudnn::pooling::adaptive_avg_pool2d(x, output_size)
288    }
289
290    fn adaptive_avg_pool2d_backward(
291        x: FloatTensor<Self>,
292        grad: FloatTensor<Self>,
293    ) -> FloatTensor<Self> {
294        rudnn::pooling::adaptive_avg_pool2d_backward(x, grad)
295    }
296
297    fn interpolate(
298        x: FloatTensor<Self>,
299        output_size: [usize; 2],
300        options: InterpolateOptions,
301    ) -> FloatTensor<Self> {
302        rudnn::interpolation::interpolate(x, output_size, options)
303    }
304
305    fn interpolate_backward(
306        x: FloatTensor<Self>,
307        grad: FloatTensor<Self>,
308        output_size: [usize; 2],
309        options: InterpolateOptions,
310    ) -> FloatTensor<Self> {
311        rudnn::interpolation::interpolate_backward(x, grad, output_size, options)
312    }
313
314    fn attention(
315        query: FloatTensor<Self>,
316        key: FloatTensor<Self>,
317        value: FloatTensor<Self>,
318        mask: Option<BoolTensor<Self>>,
319        attn_bias: Option<FloatTensor<Self>>,
320        options: AttentionModuleOptions,
321    ) -> FloatTensor<Self> {
322        // Fall back to naive attention for features the flash kernel doesn't support.
323        if attn_bias.is_some() || options.softcap.is_some() || options.scale.is_some() {
324            return ruda_tensor::ops::attention::attention_fallback::<Self>(
325                query, key, value, mask, attn_bias, options,
326            );
327        }
328
329        rudnn::attention::tensor::attention(
330            query,
331            key,
332            value,
333            mask,
334            attn_bias,
335            options,
336            Default::default(),
337        )
338        .expect("Kernel to never fail")
339    }
340
341    fn has_ctc_loss_backward() -> bool {
342        true
343    }
344
345    fn ctc_loss(
346        log_probs: FloatTensor<Self>,
347        targets: IntTensor<Self>,
348        input_lengths: IntTensor<Self>,
349        target_lengths: IntTensor<Self>,
350        blank: usize,
351    ) -> FloatTensor<Self> {
352        rudnn::ctc::ctc_loss(log_probs, targets, input_lengths, target_lengths, blank)
353    }
354
355    fn ctc_loss_backward(
356        log_probs: FloatTensor<Self>,
357        targets: IntTensor<Self>,
358        input_lengths: IntTensor<Self>,
359        target_lengths: IntTensor<Self>,
360        grad_loss: FloatTensor<Self>,
361        blank: usize,
362    ) -> FloatTensor<Self> {
363        let (log_alpha_full, log_beta_full, nll) = rudnn::ctc::ctc_alpha_beta(
364            log_probs.clone(),
365            targets.clone(),
366            input_lengths.clone(),
367            target_lengths,
368            blank,
369        );
370        ruda_tensor::ops::ctc::ctc_grad_from_alpha_beta_default::<Self>(
371            log_probs,
372            targets,
373            input_lengths,
374            grad_loss,
375            log_alpha_full,
376            log_beta_full,
377            nll,
378            blank,
379        )
380    }
381
382    fn rfft(
383        signal: FloatTensor<Self>,
384        dim: usize,
385        n: Option<usize>,
386    ) -> (FloatTensor<Self>, FloatTensor<Self>) {
387        rufft::tensor::rfft(signal, dim, n)
388    }
389
390    fn irfft(
391        spectrum_re: FloatTensor<Self>,
392        spectrum_im: FloatTensor<Self>,
393        dim: usize,
394        n: Option<usize>,
395    ) -> FloatTensor<Self> {
396        rufft::tensor::irfft(spectrum_re, spectrum_im, dim, n)
397    }
398}
ruda_tensor_device/dispatch/neural.rs

ruda_tensor_device/dispatch/
neural.rs
