Compilation optimization for reduce (#57082)

tianhaodongbd · web-flow · commit 185cd5dbc301 · 2023-09-27T14:51:05.000+08:00
diff --git a/paddle/phi/kernels/gpu/broadcast_tensors_grad_kernel.cu b/paddle/phi/kernels/gpu/broadcast_tensors_grad_kernel.cu
@@ -22,6 +22,7 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
 #include "paddle/phi/kernels/primitive/functor_primitives.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -88,12 +89,12 @@ void BroadcastTensorsGradKernel(const Context& ctx,
       phi::Copy(ctx, *input_tensor, ctx.GetPlace(), false, output_tensor);
     } else {
       // reduce_sum implementation on CUDA
-      funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-          ctx,
-          *input_tensor,
-          output_tensor,
-          kps::IdentityFunctor<T>(),
-          reduce_dims_vec);
+      phi::SumKernel<T, Context>(ctx,
+                                 *input_tensor,
+                                 reduce_dims_vec,
+                                 output_tensor->dtype(),
+                                 false,
+                                 output_tensor);
     }
   }
 }
diff --git a/paddle/phi/kernels/gpu/elementwise_grad.h b/paddle/phi/kernels/gpu/elementwise_grad.h
@@ -21,6 +21,7 @@ limitations under the License. */
 #include "paddle/phi/kernels/funcs/broadcast_function.h"
 #include "paddle/phi/kernels/funcs/elementwise_grad_base.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -31,8 +32,8 @@ void ReduceWrapper(const GPUContext &dev_ctx,
                    DenseTensor *dst) {
   std::vector<int> reduce_dims =
       funcs::GetReduceDim(dst->dims(), src->dims(), axis);
-  funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-      dev_ctx, *src, dst, kps::IdentityFunctor<T>(), reduce_dims);
+  phi::SumKernel<T, GPUContext>(
+      dev_ctx, *src, reduce_dims, src->dtype(), false, dst);
 }
 
 template <typename T, typename Functor>
@@ -169,8 +170,8 @@ void DefaultElementwiseAddGrad(const GPUContext &ctx,
       }
       std::vector<int> reduce_dims =
           funcs::GetReduceDim(x.dims(), out.dims(), axis);
-      funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-          ctx, dout, dx, kps::IdentityFunctor<T>(), reduce_dims);
+      phi::SumKernel<T, GPUContext>(
+          ctx, dout, reduce_dims, dout.dtype(), false, dx);
     }
   }
   // dy
@@ -183,8 +184,8 @@ void DefaultElementwiseAddGrad(const GPUContext &ctx,
     } else {
       std::vector<int> reduce_dims =
           funcs::GetReduceDim(y.dims(), out.dims(), axis);
-      funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-          ctx, dout, dy, kps::IdentityFunctor<T>(), reduce_dims);
+      phi::SumKernel<T, GPUContext>(
+          ctx, dout, reduce_dims, dout.dtype(), false, dy);
     }
   }
 }
@@ -280,8 +281,8 @@ void default_elementwise_sub_grad(const GPUContext &ctx,
       }
       std::vector<int> reduce_dims =
           funcs::GetReduceDim(x.dims(), out.dims(), axis);
-      funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-          ctx, dout, dx, kps::IdentityFunctor<T>(), reduce_dims);
+      phi::SumKernel<T, GPUContext>(
+          ctx, dout, reduce_dims, dout.dtype(), false, dx);
     }
   }
   // dy
diff --git a/paddle/phi/kernels/gpu/expand_as_grad_kernel.cu b/paddle/phi/kernels/gpu/expand_as_grad_kernel.cu
@@ -18,6 +18,7 @@
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -46,8 +47,8 @@ void ExpandAsGradKernel(const Context& context,
   } else {
     std::vector<int> reduce_dims = funcs::GetReduceDim(in_dims, out_dims, -1);
 
-    funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-        context, out_grad, in_grad, kps::IdentityFunctor<T>(), reduce_dims);
+    phi::SumKernel<T, Context>(
+        context, out_grad, reduce_dims, out_grad.dtype(), false, in_grad);
   }
 }
 
diff --git a/paddle/phi/kernels/gpu/expand_grad_kernel.cu b/paddle/phi/kernels/gpu/expand_grad_kernel.cu
@@ -18,6 +18,7 @@
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/reduce_function.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -33,8 +34,8 @@ void ExpandGradKernel(const Context& ctx,
   } else {
     std::vector<int> reduce_dims =
         funcs::GetReduceDim(x_grad->dims(), out_grad.dims(), -1);
-    funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
-        ctx, out_grad, x_grad, kps::IdentityFunctor<T>(), reduce_dims);
+    phi::SumKernel<T, Context>(
+        ctx, out_grad, reduce_dims, out_grad.dtype(), false, x_grad);
   }
 }
 
diff --git a/paddle/phi/kernels/kps/reduce_kernel.cu b/paddle/phi/kernels/kps/reduce_kernel.cu
@@ -16,6 +16,7 @@
 #include "paddle/phi/core/enforce.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/gpu/reduce.h"
+#include "paddle/phi/kernels/legacy/reduce_max_kernel.h"
 #include "paddle/phi/kernels/prod_kernel.h"
 #include "paddle/phi/kernels/reduce_all_kernel.h"
 #include "paddle/phi/kernels/reduce_amin_kernel.h"
@@ -102,9 +103,7 @@ void MaxKernel(const Context& dev_ctx,
                bool keep_dim,
                DenseTensor* out) {
   bool reduce_all = recompute_reduce_all(x, dims);
-  auto out_dtype = x.dtype();
-  phi::Reduce<T, kps::MaxFunctor, kps::IdentityFunctor>(
-      dev_ctx, x, reduce_all, dims.GetData(), keep_dim, out_dtype, out);
+  phi::MaxRawKernel<T, Context>(dev_ctx, x, dims, keep_dim, reduce_all, out);
 }
 
 template <typename T, typename Context>

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ limitations under the License. */`
`21`	`21`	`#include "paddle/phi/kernels/funcs/broadcast_function.h"`
`22`	`22`	`#include "paddle/phi/kernels/funcs/elementwise_grad_base.h"`
`23`	`23`	`#include "paddle/phi/kernels/funcs/reduce_function.h"`
	`24`	`+#include "paddle/phi/kernels/reduce_sum_kernel.h"`
`24`	`25`
`25`	`26`	`namespace phi {`
`26`	`27`
`@@ -31,8 +32,8 @@ void ReduceWrapper(const GPUContext &dev_ctx,`
`31`	`32`	`DenseTensor *dst) {`
`32`	`33`	`std::vector<int> reduce_dims =`
`33`	`34`	`funcs::GetReduceDim(dst->dims(), src->dims(), axis);`
`34`		`- funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(`
`35`		`- dev_ctx, *src, dst, kps::IdentityFunctor<T>(), reduce_dims);`
	`35`	`+ phi::SumKernel<T, GPUContext>(`
	`36`	`+ dev_ctx, *src, reduce_dims, src->dtype(), false, dst);`
`36`	`37`	`}`
`37`	`38`
`38`	`39`	`template <typename T, typename Functor>`
`@@ -169,8 +170,8 @@ void DefaultElementwiseAddGrad(const GPUContext &ctx,`
`169`	`170`	`}`
`170`	`171`	`std::vector<int> reduce_dims =`
`171`	`172`	`funcs::GetReduceDim(x.dims(), out.dims(), axis);`
`172`		`- funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(`
`173`		`- ctx, dout, dx, kps::IdentityFunctor<T>(), reduce_dims);`
	`173`	`+ phi::SumKernel<T, GPUContext>(`
	`174`	`+ ctx, dout, reduce_dims, dout.dtype(), false, dx);`
`174`	`175`	`}`
`175`	`176`	`}`
`176`	`177`	`// dy`
`@@ -183,8 +184,8 @@ void DefaultElementwiseAddGrad(const GPUContext &ctx,`
`183`	`184`	`} else {`
`184`	`185`	`std::vector<int> reduce_dims =`
`185`	`186`	`funcs::GetReduceDim(y.dims(), out.dims(), axis);`
`186`		`- funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(`
`187`		`- ctx, dout, dy, kps::IdentityFunctor<T>(), reduce_dims);`
	`187`	`+ phi::SumKernel<T, GPUContext>(`
	`188`	`+ ctx, dout, reduce_dims, dout.dtype(), false, dy);`
`188`	`189`	`}`
`189`	`190`	`}`
`190`	`191`	`}`
`@@ -280,8 +281,8 @@ void default_elementwise_sub_grad(const GPUContext &ctx,`
`280`	`281`	`}`
`281`	`282`	`std::vector<int> reduce_dims =`
`282`	`283`	`funcs::GetReduceDim(x.dims(), out.dims(), axis);`
`283`		`- funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(`
`284`		`- ctx, dout, dx, kps::IdentityFunctor<T>(), reduce_dims);`
	`284`	`+ phi::SumKernel<T, GPUContext>(`
	`285`	`+ ctx, dout, reduce_dims, dout.dtype(), false, dx);`
`285`	`286`	`}`
`286`	`287`	`}`
`287`	`288`	`// dy`
Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@`
`18`	`18`	`#include "paddle/phi/core/dense_tensor.h"`
`19`	`19`	`#include "paddle/phi/core/kernel_registry.h"`
`20`	`20`	`#include "paddle/phi/kernels/funcs/reduce_function.h"`
	`21`	`+#include "paddle/phi/kernels/reduce_sum_kernel.h"`
`21`	`22`
`22`	`23`	`namespace phi {`
`23`	`24`
`@@ -46,8 +47,8 @@ void ExpandAsGradKernel(const Context& context,`
`46`	`47`	`} else {`
`47`	`48`	`std::vector<int> reduce_dims = funcs::GetReduceDim(in_dims, out_dims, -1);`
`48`	`49`
`49`		`- funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(`
`50`		`- context, out_grad, in_grad, kps::IdentityFunctor<T>(), reduce_dims);`
	`50`	`+ phi::SumKernel<T, Context>(`
	`51`	`+ context, out_grad, reduce_dims, out_grad.dtype(), false, in_grad);`
`51`	`52`	`}`
`52`	`53`	`}`
`53`	`54`