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