ruda_tensor_host/ops/
module.rs

1//! Module operations for the Host backend.
2//!
3//! These operations power neural network modules like convolutions and pooling.
4
5use crate::Host;
6use ruda_tensor::{
7    ops::{
8        AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward,
9        DeformConvOptions, InterpolateOptions,
10        MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps,
11    },
12    tensor::{BoolTensor, FloatTensor, IntTensor},
13};
14
15pub(crate) use ruda_core::tensor::host::cast::{cast_from_f32, cast_to_f32};
16
17impl ModuleOps<Host> for Host {
18    fn conv1d(
19        x: FloatTensor<Host>,
20        weight: FloatTensor<Host>,
21        bias: Option<FloatTensor<Host>>,
22        options: ConvOptions<1>,
23    ) -> FloatTensor<Host> {
24        rudnn_host::convolution::dispatch::conv1d(x, weight, bias, options)
25    }
26
27    fn conv2d(
28        x: FloatTensor<Host>,
29        weight: FloatTensor<Host>,
30        bias: Option<FloatTensor<Host>>,
31        options: ConvOptions<2>,
32    ) -> FloatTensor<Host> {
33        rudnn_host::convolution::dispatch::conv2d(x, weight, bias, options)
34    }
35
36    fn deform_conv2d(
37        x: FloatTensor<Host>,
38        offset: FloatTensor<Host>,
39        weight: FloatTensor<Host>,
40        mask: Option<FloatTensor<Host>>,
41        bias: Option<FloatTensor<Host>>,
42        options: DeformConvOptions<2>,
43    ) -> FloatTensor<Host> {
44        rudnn_host::convolution::dispatch::deform_conv2d(x, offset, weight, mask, bias, options)
45    }
46
47    fn deform_conv2d_backward(
48        x: FloatTensor<Host>,
49        offset: FloatTensor<Host>,
50        weight: FloatTensor<Host>,
51        mask: Option<FloatTensor<Host>>,
52        bias: Option<FloatTensor<Host>>,
53        output_grad: FloatTensor<Host>,
54        options: DeformConvOptions<2>,
55    ) -> DeformConv2dBackward<Host> {
56        let (x_grad, offset_grad, weight_grad, mask_grad, bias_grad) = rudnn_host::convolution::dispatch::deform_conv2d_backward(x, offset, weight, mask, bias, output_grad, options);
57        DeformConv2dBackward::new(x_grad, offset_grad, weight_grad, mask_grad, bias_grad)
58    }
59
60    fn conv3d(
61        x: FloatTensor<Host>,
62        weight: FloatTensor<Host>,
63        bias: Option<FloatTensor<Host>>,
64        options: ConvOptions<3>,
65    ) -> FloatTensor<Host> {
66        rudnn_host::convolution::dispatch::conv3d(x, weight, bias, options)
67    }
68
69    fn conv_transpose1d(
70        x: FloatTensor<Host>,
71        weight: FloatTensor<Host>,
72        bias: Option<FloatTensor<Host>>,
73        options: ConvTransposeOptions<1>,
74    ) -> FloatTensor<Host> {
75        rudnn_host::convolution::dispatch::conv_transpose1d(x, weight, bias, options)
76    }
77
78    fn conv_transpose2d(
79        x: FloatTensor<Host>,
80        weight: FloatTensor<Host>,
81        bias: Option<FloatTensor<Host>>,
82        options: ConvTransposeOptions<2>,
83    ) -> FloatTensor<Host> {
84        rudnn_host::convolution::dispatch::conv_transpose2d(x, weight, bias, options)
85    }
86
87    fn conv_transpose3d(
88        x: FloatTensor<Host>,
89        weight: FloatTensor<Host>,
90        bias: Option<FloatTensor<Host>>,
91        options: ConvTransposeOptions<3>,
92    ) -> FloatTensor<Host> {
93        rudnn_host::convolution::dispatch::conv_transpose3d(x, weight, bias, options)
94    }
95
96    fn avg_pool2d(
97        x: FloatTensor<Host>,
98        kernel_size: [usize; 2],
99        stride: [usize; 2],
100        padding: [usize; 2],
101        count_include_pad: bool,
102        ceil_mode: bool,
103    ) -> FloatTensor<Host> {
104        rudnn_host::pool::dispatch::avg_pool2d(x, kernel_size, stride, padding, count_include_pad, ceil_mode)
105    }
106
107    fn avg_pool2d_backward(
108        x: FloatTensor<Host>,
109        grad: FloatTensor<Host>,
110        kernel_size: [usize; 2],
111        stride: [usize; 2],
112        padding: [usize; 2],
113        count_include_pad: bool,
114        _divisor_override: bool,
115    ) -> FloatTensor<Host> {
116        rudnn_host::pool::dispatch::avg_pool2d_backward(x, grad, kernel_size, stride, padding, count_include_pad, _divisor_override)
117    }
118
119    fn adaptive_avg_pool2d(x: FloatTensor<Host>, output_size: [usize; 2]) -> FloatTensor<Host> {
120        rudnn_host::pool::dispatch::adaptive_avg_pool2d(x, output_size)
121    }
122
123    fn adaptive_avg_pool2d_backward(
124        x: FloatTensor<Host>,
125        grad: FloatTensor<Host>,
126    ) -> FloatTensor<Host> {
127        rudnn_host::pool::dispatch::adaptive_avg_pool2d_backward(x, grad)
128    }
129
130    fn max_pool2d(
131        x: FloatTensor<Host>,
132        kernel_size: [usize; 2],
133        stride: [usize; 2],
134        padding: [usize; 2],
135        dilation: [usize; 2],
136        ceil_mode: bool,
137    ) -> FloatTensor<Host> {
138        rudnn_host::pool::dispatch::max_pool2d(x, kernel_size, stride, padding, dilation, ceil_mode)
139    }
140
141    fn max_pool2d_with_indices(
142        x: FloatTensor<Host>,
143        kernel_size: [usize; 2],
144        stride: [usize; 2],
145        padding: [usize; 2],
146        dilation: [usize; 2],
147        ceil_mode: bool,
148    ) -> MaxPool2dWithIndices<Host> {
149        let (output, indices) = rudnn_host::pool::dispatch::max_pool2d_with_indices(x, kernel_size, stride, padding, dilation, ceil_mode);
150        MaxPool2dWithIndices::new(output, indices)
151    }
152
153    fn max_pool2d_with_indices_backward(
154        x: FloatTensor<Host>,
155        _kernel_size: [usize; 2],
156        _stride: [usize; 2],
157        _padding: [usize; 2],
158        _dilation: [usize; 2],
159        _ceil_mode: bool,
160        output_grad: FloatTensor<Host>,
161        indices: IntTensor<Host>,
162    ) -> MaxPool2dBackward<Host> {
163        let x_grad = rudnn_host::pool::dispatch::max_pool2d_with_indices_backward(x, _kernel_size, _stride, _padding, _dilation, _ceil_mode, output_grad, indices);
164        MaxPool2dBackward::new(x_grad)
165    }
166
167    fn interpolate(
168        x: FloatTensor<Host>,
169        output_size: [usize; 2],
170        options: InterpolateOptions,
171    ) -> FloatTensor<Host> {
172        rudnn_host::interpolate::dispatch::interpolate(x, output_size, options)
173    }
174
175    fn interpolate_backward(
176        x: FloatTensor<Host>,
177        grad: FloatTensor<Host>,
178        output_size: [usize; 2],
179        options: InterpolateOptions,
180    ) -> FloatTensor<Host> {
181        rudnn_host::interpolate::dispatch::interpolate_backward(x, grad, output_size, options)
182    }
183
184    fn attention(
185        query: FloatTensor<Host>,
186        key: FloatTensor<Host>,
187        value: FloatTensor<Host>,
188        mask: Option<BoolTensor<Host>>,
189        attn_bias: Option<FloatTensor<Host>>,
190        options: AttentionModuleOptions,
191    ) -> FloatTensor<Host> {
192        crate::ops::attention::attention(query, key, value, mask, attn_bias, options)
193    }
194
195    fn rfft(
196        signal: FloatTensor<Host>,
197        dim: usize,
198        n: Option<usize>,
199    ) -> (FloatTensor<Host>, FloatTensor<Host>) {
200        rufft_host::dispatch::rfft(signal, dim, n)
201    }
202
203    fn irfft(
204        spectrum_re: FloatTensor<Host>,
205        spectrum_im: FloatTensor<Host>,
206        dim: usize,
207        n: Option<usize>,
208    ) -> FloatTensor<Host> {
209        rufft_host::dispatch::irfft(spectrum_re, spectrum_im, dim, n)
210    }
211
212    fn embedding(weights: FloatTensor<Host>, indices: IntTensor<Host>) -> FloatTensor<Host> {
213        rudnn_host::embedding::embedding(weights, indices)
214    }
215
216    fn layer_norm(
217        tensor: FloatTensor<Host>,
218        gamma: FloatTensor<Host>,
219        beta: Option<FloatTensor<Host>>,
220        epsilon: f64,
221    ) -> FloatTensor<Host> {
222        crate::ops::activation::layer_norm(tensor, gamma, beta, epsilon)
223    }
224
225    fn embedding_backward(
226        weights: FloatTensor<Host>,
227        output_grad: FloatTensor<Host>,
228        indices: IntTensor<Host>,
229    ) -> FloatTensor<Host> {
230        rudnn_host::embedding::embedding_backward(weights, output_grad, indices)
231    }
232}
ruda_tensor_host/ops/module.rs

ruda_tensor_host/ops/
module.rs
