Fast softmax

src/targets/gpu/kernels/include/migraphx/kernels/softmax.hpp

@@ -33,11 +33,9 @@ template <index_int Axis, class Input, class Output>
 __device__ void softmax(Input input, Output output)
 {
     reduce::block::run<reduce::with_axis<Input, Axis>>([&](auto, auto r) {
-        auto batch_max = r.reduce(op::max{}, lowest{}, op::id{})(input);
-        auto batch_sum =
-            r.reduce(op::sum{}, 0, [&](auto x) { return migraphx::exp(x - batch_max); })(input);
-        r.inner([&](auto& y, auto x) { y = migraphx::exp(x - batch_max) / batch_sum; })(output,
-                                                                                        input);
+        auto batch_sum = r.reduce(
+            op::sum{}, 0, [](auto x) { return migraphx::convert<float>(migraphx::exp(x)); })(input);
+        r.inner([&](auto& y, auto x) { y = migraphx::exp(x) / batch_sum; })(output, input);
     });
 }