fix fp16 reduce_mean

Author: sneaxiy

paddle/fluid/operators/elementwise/elementwise_min_op.h

@@ -19,6 +19,7 @@ limitations under the License. */
 #include "paddle/fluid/operators/elementwise/elementwise_op.h"
 #include "paddle/fluid/operators/elementwise/elementwise_op_function.h"
 #include "paddle/fluid/platform/eigen_ext.h"
+#include "paddle/fluid/platform/float16.h"
 
 namespace paddle {
 namespace operators {
@@ -56,17 +57,39 @@ class ElementwiseFMinKernel : public framework::OpKernel<T> {
 template <typename T>
 struct MinGradDx {
  HOSTDEVICE T operator()(T x, T y, T out, T dout) const {
- return dout * static_cast<T>(x < y);
+ return dout * (x < y);
  }
 };
 
 template <typename T>
 struct MinGradDy {
  HOSTDEVICE T operator()(T x, T y, T out, T dout) const {
- return dout * static_cast<T>(x >= y);
+ return dout * (x >= y);
  }
 };
 
+#ifdef PADDLE_CUDA_FP16
+template <>
+struct MinGradDx<platform::float16> {
+ HOSTDEVICE platform::float16 operator()(platform::float16 x,
+ platform::float16 y,
+ platform::float16 out,
+ platform::float16 dout) const {
+ return x < y ? dout : static_cast<platform::float16>(0);
+ }
+};
+
+template <>
+struct MinGradDy<platform::float16> {
+ HOSTDEVICE platform::float16 operator()(platform::float16 x,
+ platform::float16 y,
+ platform::float16 out,
+ platform::float16 dout) const {
+ return x >= y ? dout : static_cast<platform::float16>(0);
+ }
+};
+#endif
+
 template <typename DeviceContext, typename T>
 class ElementwiseMinGradKernel : public ElemwiseGradKernel<T> {
  public:

paddle/fluid/operators/reduce_ops/reduce_mean_op.cu

@@ -17,6 +17,7 @@
 #include "paddle/fluid/operators/reduce_ops/reduce_op.h"
 
 REGISTER_OP_CUDA_KERNEL(
+ reduce_mean,
  ops::ReduceCudaKernel<bool, kps::AddFunctor, kps::DivideFunctor>,
  ops::ReduceCudaKernel<paddle::platform::float16, kps::AddFunctor,
  kps::DivideFunctor>,