reused RawKernel in elementwise_kernel (PaddlePaddle#57721)

danleifeng · Sep 28, 2023 · f34bf3c · f34bf3c
1 parent d9ad40c
commit f34bf3c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 45 deletions.
diff --git a/paddle/phi/kernels/kps/elementwise_kernel.cu b/paddle/phi/kernels/kps/elementwise_kernel.cu
@@ -20,7 +20,10 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/elementwise_add_kernel.h"
 #include "paddle/phi/kernels/impl/elementwise_kernel_impl.h"
+#include "paddle/phi/kernels/legacy/elementwise_add_kernel.h"
 #include "paddle/phi/kernels/legacy/elementwise_kernel.h"
+#include "paddle/phi/kernels/legacy/elementwise_multipy_kernel.h"
+#include "paddle/phi/kernels/legacy/elementwise_subtract_kernel.h"
 
 namespace phi {
 
@@ -29,63 +32,23 @@ void SubtractKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     DenseTensor* out) {
-  std::vector<const DenseTensor*> inputs;
-  inputs.reserve(2);
-  std::vector<DenseTensor*> outputs;
-  outputs.reserve(1);
-  inputs.emplace_back(&x);
-  inputs.emplace_back(&y);
-  outputs.emplace_back(out);
-  dev_ctx.template Alloc<T>(out);
-  funcs::BroadcastKernel<T>(
-      dev_ctx, inputs, &outputs, funcs::SubtractFunctor<T>(), -1);
+  phi::SubtractRawKernel<T, Context>(dev_ctx, x, y, -1, out);
 }
 
 template <typename T, typename Context>
 void MultiplyKernel(const Context& dev_ctx,
                     const DenseTensor& x,
                     const DenseTensor& y,
                     DenseTensor* out) {
-  std::vector<const DenseTensor*> inputs;
-  inputs.reserve(2);
-  std::vector<DenseTensor*> outputs;
-  outputs.reserve(1);
-  inputs.emplace_back(&x);
-  inputs.emplace_back(&y);
-  outputs.emplace_back(out);
-  dev_ctx.template Alloc<T>(out);
-  funcs::BroadcastKernel<T>(
-      dev_ctx, inputs, &outputs, funcs::MultiplyFunctor<T>(), -1);
+  phi::MultiplyRawKernel<T, Context>(dev_ctx, x, y, -1, out);
 }
 
 template <typename T, typename Context>
 void DivideKernel(const Context& dev_ctx,
                   const DenseTensor& x,
                   const DenseTensor& y,
                   DenseTensor* out) {
-  std::vector<const DenseTensor*> inputs;
-  inputs.reserve(2);
-  std::vector<DenseTensor*> outputs;
-  outputs.reserve(1);
-  inputs.emplace_back(&x);
-  inputs.emplace_back(&y);
-  outputs.emplace_back(out);
-  dev_ctx.template Alloc<T>(out);
-  funcs::BroadcastKernel<T>(
-      dev_ctx, inputs, &outputs, funcs::DivideFunctor<T>(), -1);
-}
-
-template <typename T, typename Context>
-void AddKernelImpl(const Context& dev_ctx,
-                   const DenseTensor& x,
-                   const DenseTensor& y,
-                   int axis,
-                   DenseTensor* out) {
-  std::vector<const DenseTensor*> inputs = {&x, &y};
-  std::vector<DenseTensor*> outputs = {out};
-  dev_ctx.template Alloc<T>(out);
-  funcs::BroadcastKernel<T>(
-      dev_ctx, inputs, &outputs, funcs::AddFunctor<T>(), axis);
+  phi::DivideRawKernel<T, Context>(dev_ctx, x, y, -1, out);
 }
 
 template <typename T, typename Context>
@@ -127,7 +90,7 @@ void AddKernel(const Context& dev_ctx,
     MultiPrecisionAddKernelImpl<float, Context>(dev_ctx, x, y, out);
   } else {
 #endif
-    AddKernelImpl<T, Context>(dev_ctx, x, y, -1, out);
+    phi::AddRawKernel<T, Context>(dev_ctx, x, y, -1, out);
 #ifdef PADDLE_WITH_CUDA
   }
 #endif
@@ -138,7 +101,7 @@ void GradAddKernel(const Context& dev_ctx,
                    const DenseTensor& x,
                    const DenseTensor& y,
                    DenseTensor* out) {
-  AddKernelImpl<T>(dev_ctx, x, y, -1, out);
+  phi::AddRawKernel<T>(dev_ctx, x, y, -1, out);
 }
 
 template <typename T, typename Context>

diff --git a/paddle/phi/kernels/legacy/elementwise_kernel.h b/paddle/phi/kernels/legacy/elementwise_kernel.h
@@ -19,6 +19,13 @@
 
 namespace phi {
 
+template <typename T, typename Context>
+void DivideRawKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const DenseTensor& y,
+                     int axis,
+                     DenseTensor* out);
+
 template <typename T, typename Context>
 void MaximumRawKernel(const Context& dev_ctx,
                       const DenseTensor& x,