fix fp16 reduce_mean

PaddlePaddle · Dec 16, 2021 · 3c18889 · 3c18889
1 parent 25c35ba
commit 3c18889
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/paddle/fluid/operators/elementwise/elementwise_min_op.h b/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:

diff --git a/paddle/fluid/operators/reduce_ops/reduce_mean_op.cu b/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>,