1use crate::context::{TractCudaStream, cuda_context};
2use crate::kernels::launch_args::TractLaunchArgs;
3use crate::kernels::{LibraryName, MAX_THREADS, get_cuda_view, launch_args, utils};
4use cudarc::driver::{CudaStream, LaunchConfig, PushKernelArg};
5use tract_core::internal::*;
6use tract_gpu::tensor::DeviceTensor;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub struct Softmax;
10
11impl Softmax {
12 pub fn is_supported_dt(dt: DatumType) -> bool {
13 matches!(dt, DatumType::F32 | DatumType::F16)
14 }
15
16 pub fn kernel_name(&self, dt: DatumType, n_cols: usize) -> TractResult<String> {
17 ensure!(Self::is_supported_dt(dt), "Unsupported dt {:?} for cuda softmaxop", dt);
18 let tname = DeviceTensor::tname(dt)?;
19 if n_cols < MAX_THREADS {
20 Ok(format!("softmax_small_{tname}"))
21 } else {
22 Ok(format!("softmax_{tname}"))
23 }
24 }
25
26 pub fn eval(
27 &self,
28 stream: &TractCudaStream,
29 input: &DeviceTensor,
30 axis: usize,
31 ) -> TractResult<DeviceTensor> {
32 let output = unsafe { DeviceTensor::uninitialized_dt(input.datum_type(), input.shape())? };
33 self.dispatch_eval(stream, input, axis, &output)?;
34 stream.synchronize()?;
35 Ok(output)
36 }
37
38 pub fn dispatch_eval(
39 &self,
40 stream: &TractCudaStream,
41 input: &DeviceTensor,
42 axis: usize,
43 output: &DeviceTensor,
44 ) -> TractResult<()> {
45 ensure!(output.shape() == input.shape());
46 ensure!(output.datum_type() == input.datum_type());
47
48 let shape_nd3 = utils::reshape_to_rank_3(input.shape(), axis);
49 let strides_nd3 = Tensor::natural_strides(&shape_nd3);
50
51 let i_view = get_cuda_view(input);
52 let o_view = get_cuda_view(output);
53
54 let func = cuda_context()
55 .load_pipeline(LibraryName::NN, self.kernel_name(input.datum_type(), shape_nd3[1])?)?;
56 let mut launch_args = TractLaunchArgs::new(stream, &func);
57 launch_args.push_view(&i_view);
58 launch_args.push_view(&o_view);
59 launch_args.push_slice_i32(&shape_nd3);
60 launch_args.push_slice_i32(&strides_nd3);
61
62 let cfg = LaunchConfig {
63 grid_dim: ((shape_nd3[0] * shape_nd3[2]) as _, 1, 1),
64 block_dim: if shape_nd3[1] < MAX_THREADS {
65 (32, 1, 1)
66 } else {
67 (MAX_THREADS as _, 1, 1)
68 },
69 shared_mem_bytes: 0,
70 };
71
72 launch_args.launch(cfg)
73 }
74}
75
76pub fn cuda_softmax_dispatch(
77 input: &DeviceTensor,
78 axis: usize,
79 output: &DeviceTensor,
80) -> TractResult<()> {
81 crate::with_cuda_stream(|stream| Softmax.dispatch_eval(stream, input, axis, output))
82}
83
84crate::register_cuda_op!(tract_core::ops::nn::Softmax, |source, node, op| {
85 rule_if!(Softmax::is_supported_dt(source.node_input_facts(node.id)?[0].datum_type));
86 Ok(Some(Box::new(tract_gpu::ops::softmax::GpuSoftmax::from_tract_core(
87 op,
88 "Cuda",
89 cuda_softmax_dispatch,
90 )?)))
91});
92
93#[cfg(test)]
94mod tests {
95
96 use super::*;
97 use derive_new::new;
98 use num_traits::AsPrimitive;
99 use num_traits::Float;
100 use proptest::collection::vec;
101 use proptest::prelude::*;
102 use tract_core::internal::Tensor;
103 use tract_core::ops::nn::Softmax as TractSoftmax;
104 use tract_core::ops::nn::SoftmaxKind;
105 use tract_gpu::tensor::IntoDevice;
106
107 #[test]
108 fn test_softmax_f32() -> TractResult<()> {
109 crate::with_cuda_stream(|stream| {
110 let m = 2;
111 let k = 3;
112 let axis = 1;
113
114 let a = Tensor::from_shape(&[m, k], &(0..m * k).map(|f| f as f32).collect::<Vec<_>>())?
115 .into_device()?;
116
117 let cpu_softmax = TractSoftmax {
118 axes: tvec![axis],
119 quant_output_dt: None,
120 kind: SoftmaxKind::Softmax,
121 };
122
123 let cpu_output = cpu_softmax
124 .eval(&EvalContext::out_of_plan(), tvec![a.to_host()?.into_tvalue()])?[0]
125 .clone()
126 .into_tensor();
127 let cuda_output = Softmax.eval(stream, &a, axis)?;
128
129 cpu_output
130 .close_enough(&cuda_output.to_host()?.into_tensor(), Approximation::Approximate)?;
131 Ok(())
132 })
133 }
134
135 #[test]
136 fn test_softmax_f32_2() -> TractResult<()> {
137 crate::with_cuda_stream(|stream| {
138 let shape = [8, 4, 3];
139 let num_elements = shape.iter().product();
140 let axis = 0;
141
142 let a = Tensor::from_shape(
143 &shape,
144 &(0..num_elements).map(|f| f as f32 / 1000.0).collect::<Vec<_>>(),
145 )?
146 .into_device()?;
147
148 let cpu_softmax = TractSoftmax {
149 axes: tvec![axis],
150 quant_output_dt: None,
151 kind: SoftmaxKind::Softmax,
152 };
153
154 let cpu_output = cpu_softmax
155 .eval(&EvalContext::out_of_plan(), tvec![a.to_host()?.into_tvalue()])?[0]
156 .clone()
157 .into_tensor();
158 let cuda_output = Softmax.eval(stream, &a, axis)?;
159 cpu_output
160 .close_enough(&cuda_output.to_host()?.into_tensor(), Approximation::Approximate)?;
161 Ok(())
162 })
163 }
164
165 #[test]
166 fn test_softmax_f16() -> TractResult<()> {
167 crate::with_cuda_stream(|stream| {
168 let m = 4;
169 let k = 4;
170 let axis = 1;
171
172 let a = Tensor::from_shape(
173 &[m, k],
174 &(0..m * k).map(|f| -> f16 { f.as_() }).collect::<Vec<_>>(),
175 )?
176 .into_device()?;
177
178 let cpu_softmax = TractSoftmax {
179 axes: tvec![axis],
180 quant_output_dt: None,
181 kind: SoftmaxKind::Softmax,
182 };
183
184 let cpu_output = cpu_softmax
185 .eval(&EvalContext::out_of_plan(), tvec![a.to_host()?.into_tvalue()])?[0]
186 .clone()
187 .into_tensor();
188 let cuda_output = Softmax.eval(stream, &a, axis)?;
189 cpu_output
190 .close_enough(&cuda_output.to_host()?.into_tensor(), Approximation::Approximate)?;
191 Ok(())
192 })
193 }
194
195 proptest::proptest! {
196 #[test]
197 fn softmax_prop_f32(pb in any::<SoftmaxProblem<f32>>()) {
198 fn run(pb: SoftmaxProblem<f32>) -> TractResult<()> {
199 let out = pb.run()?;
200 let reference = pb.reference()?;
201
202 out.close_enough(&reference, Approximation::Approximate)
203 .with_context(|| format!("Cpu: {:?}, Cuda: {:?}", reference.dump(true), out.dump(true)))
204 }
205 run(pb).map_err(|e| TestCaseError::Fail(format!("{:?}", e).into()))?;
206 }
207
208 #[test]
209 fn softmax_prop_f16(pb in any::<SoftmaxProblem<f16>>()) {
210 fn run(pb: SoftmaxProblem<f16>) -> TractResult<()> {
211 let out = pb.run()?;
212 let reference = pb.reference()?;
213
214 out.close_enough(&reference, Approximation::Approximate)
215 .with_context(|| format!("Cpu: {:?}, Cuda: {:?}", reference.dump(true), out.dump(true)))
216 }
217
218 run(pb).map_err(|e| TestCaseError::Fail(format!("{:?}", e).into()))?;
219 }
220 }
221
222 #[derive(Debug, new)]
223 pub struct SoftmaxProblem<F: Datum + Float>
224 where
225 F: Datum + Float,
226 usize: AsPrimitive<F>,
227 {
228 pub shape: Vec<usize>,
229 pub axis: usize,
230 pub input: Vec<F>,
231 }
232
233 impl<F> Arbitrary for SoftmaxProblem<F>
234 where
235 F: Datum + Float,
236 usize: AsPrimitive<F>,
237 {
238 type Parameters = ();
239 type Strategy = BoxedStrategy<Self>;
240
241 fn arbitrary_with(_: ()) -> Self::Strategy {
242 (0usize..3, 0usize..3)
243 .prop_flat_map(|(left, right)| {
244 let axis = left;
245 let shape_len = usize::min(left + right + 1, 4);
246 let shape = 1usize..10;
247 (vec(shape, shape_len..=shape_len), Just(axis))
248 })
249 .prop_map(|(shape, axis)| {
250 let input = (0..shape.iter().product::<usize>())
251 .map(|f| f.as_() / 1000.as_())
252 .collect::<Vec<_>>();
253 Self { shape, axis, input }
254 })
255 .boxed()
256 }
257 }
258
259 impl<F> SoftmaxProblem<F>
260 where
261 F: Datum + Float + std::ops::AddAssign,
262 usize: AsPrimitive<F>,
263 {
264 pub fn reference(&self) -> TractResult<Tensor> {
265 let a = Tensor::from_shape(self.shape.as_slice(), &self.input)?;
266
267 let cpu_softmax = TractSoftmax {
268 axes: tvec![self.axis],
269 quant_output_dt: None,
270 kind: SoftmaxKind::Softmax,
271 };
272 let cpu_output = cpu_softmax
273 .eval(&EvalContext::out_of_plan(), tvec![a.into_tvalue()])?[0]
274 .clone()
275 .into_tensor();
276 Ok(cpu_output)
277 }
278
279 pub fn run(&self) -> TractResult<Tensor> {
280 crate::with_cuda_stream(|stream| {
281 let a = Tensor::from_shape(self.shape.as_slice(), &self.input)?.into_device()?;
282 let cuda_output = Softmax.eval(stream, &a, self.axis)?;
283 Ok(cuda_output.to_host()?.into_tensor())
284 })
285 }
286 }
287}