[KP] Unify .cu and .xpu files with .kps files

cmake/operators.cmake

@@ -73,6 +73,12 @@ function(op_library TARGET)
             if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET}.cu)
                 list(APPEND cu_srcs ${TARGET}.cu)
             endif()
+            # rename in KP: .kps -> .cu
+            if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET}.kps)
+                file(COPY ${TARGET}.kps DESTINATION  ${CMAKE_CURRENT_BINARY_DIR})
+                file(RENAME ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.kps ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.cu)
+                list(APPEND cu_srcs ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.cu)
+            endif()
             if (WITH_NV_JETSON)
                 list(REMOVE_ITEM cu_srcs "decode_jpeg_op.cu")
             endif()
@@ -96,6 +102,12 @@ function(op_library TARGET)
             if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET}.cu)
                 list(APPEND hip_srcs ${TARGET}.cu)
             endif()
+            # rename in KP: .kps -> .cu
+            if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET}.kps)
+                file(COPY ${TARGET}.kps DESTINATION  ${CMAKE_CURRENT_BINARY_DIR})
+                file(RENAME ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.kps ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.cu)
+                list(APPEND hip_srcs ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.cu)
+            endif()
             if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET}.part.cu)
                 set(PART_CUDA_KERNEL_FILES ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET}.part.cu
                         ${PART_CUDA_KERNEL_FILES} PARENT_SCOPE)

paddle/fluid/operators/elementwise/elementwise_add_op.h

@@ -13,22 +13,40 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
-
+#ifdef __xpu__
+#include <memory>
+#include <string>
+#include "paddle/fluid/operators/elementwise/elementwise_op.h"
+#include "paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h"
+#include "paddle/fluid/operators/elementwise/elementwise_xpu.h"
+#include "paddle/fluid/platform/device/device_wrapper.h"
+#else
 #include <algorithm>
 #include <utility>
 #include "paddle/fluid/operators/elementwise/elementwise_op.h"
 
 // only can include the headers in paddle/phi/include dirs
 #include "paddle/phi/kernels/elementwise_grad_kernel.h"
 #include "paddle/phi/kernels/math_kernel.h"
+#endif
 
 namespace paddle {
 namespace operators {
 
 template <typename DeviceContext, typename T>
 class ElementwiseAddKernel : public framework::OpKernel<T> {
  public:
-  void Compute(const framework::ExecutionContext &ctx) const override {
+  void Compute(const framework::ExecutionContext& ctx) const override {
+#ifdef __xpu__
+    std::vector<const framework::Tensor*> ins;
+    std::vector<framework::Tensor*> outs;
+    int axis = PackTensorsIntoVector<T>(ctx, &ins, &outs);
+    const auto& xpu_ctx =
+        ctx.template device_context<paddle::platform::XPUDeviceContext>();
+    paddle::operators::LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T,
+                                                   T, kps::AddFunctor<T>, 1>(
+        xpu_ctx, ins, &outs, axis, kps::AddFunctor<T>());
+#else
     auto *x = ctx.Input<framework::LoDTensor>("X");
     auto *y = ctx.Input<framework::LoDTensor>("Y");
     auto *z = ctx.Output<framework::LoDTensor>("Out");
@@ -40,6 +58,7 @@ class ElementwiseAddKernel : public framework::OpKernel<T> {
         static_cast<const typename framework::ConvertToPtenContext<
             DeviceContext>::TYPE &>(dev_ctx),
         *x, *y, axis, z);
+#endif
   }
 };
 

paddle/fluid/operators/elementwise/elementwise_add_op.kps

@@ -1,14 +1,19 @@
 /* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
+
     http://www.apache.org/licenses/LICENSE-2.0
+
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
 
+#ifdef PADDLE_WITH_XPU_KP
+
 // Please do not modify the following code
 #if defined(__CUDA_ARCH__)
 #undef __CUDA_ARCH__
@@ -26,163 +31,31 @@ limitations under the License. */
 #undef __NVCC__
 #endif
 
-#ifdef PADDLE_WITH_XPU_KP
 #include <xpu/runtime.h>                // NOLINT
 #include "xpu/kernel/cluster_header.h"  // NOLINT
 #include "xpu/kernel/debug.h"           // NOLINT
 #include "xpu/kernel/math.h"            // NOLINT
 
-#include <memory>
-#include <string>
 #include "paddle/fluid/operators/elementwise/elementwise_add_op.h"
-#include "paddle/fluid/operators/elementwise/elementwise_op.h"
-#include "paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h"
-#include "paddle/fluid/operators/elementwise/elementwise_xpu.h"
-#include "paddle/fluid/platform/device/device_wrapper.h"
-
-namespace paddle {
-namespace operators {
-
-template <typename T>
-class ElementwiseAddXPUKPKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& ctx) const override {
-    std::vector<const framework::Tensor*> ins;
-    std::vector<framework::Tensor*> outs;
-    int axis = PackTensorsIntoVector<T>(ctx, &ins, &outs);
-    const auto& xpu_ctx =
-        ctx.template device_context<paddle::platform::XPUDeviceContext>();
-    paddle::operators::LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T,
-                                                   T, kps::AddFunctor<T>, 1>(
-        xpu_ctx, ins, &outs, axis, kps::AddFunctor<T>());
-  }
-};
-
-static std::vector<int> get_rdims(const std::vector<int>& xdims,
-                                  const std::vector<int>& ydims) {
-  std::vector<int> rdims;
-  for (size_t i = 0; i < xdims.size(); i++) {
-    if (xdims[i] != ydims[i]) {
-      rdims.push_back(i);
-    }
-  }
-  return rdims;
-}
-
-template <typename T>
-class ElementwiseAddGradXPUKPKernel : public ElemwiseGradKernel<T> {
-  using XPUType = typename XPUTypeTrait<T>::Type;
-
- public:
-  void Compute(const framework::ExecutionContext& ctx) const override {
-    ElemwiseGradKernel<T>::Compute(ctx);
-    auto* x = ctx.Input<framework::Tensor>("X");
-    auto* y = ctx.Input<framework::Tensor>("Y");
-    auto* dz = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
-    auto* dx = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
-    auto* dy = ctx.Output<framework::Tensor>(framework::GradVarName("Y"));
-    const framework::DDim& x_dims = x->dims();
-    const framework::DDim& y_dims = y->dims();
-    const framework::DDim& dz_dims = dz->dims();
-    int axis = ctx.Attr<int>("axis");
-    axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
-    int max_dim = std::max(x_dims.size(), y_dims.size());
-    PADDLE_ENFORCE_GE(
-        axis, 0,
-        platform::errors::InvalidArgument(
-            "Axis should be great than or equal to 0, but received axis is %d.",
-            axis));
-    PADDLE_ENFORCE_LT(
-        axis, max_dim,
-        platform::errors::InvalidArgument(
-            "Axis should be less than %d, but received axis is %d.", max_dim,
-            axis));
-
-    std::vector<int> x_dims_vec(max_dim, 1);
-    std::vector<int> y_dims_vec(max_dim, 1);
-    std::vector<int> z_dims_vec(max_dim, 1);
-    if (x_dims.size() == max_dim) {
-      for (int i = 0; i < max_dim; i++) {
-        x_dims_vec[i] = x_dims[i];
-      }
-    } else {
-      for (int i = 0; i < x_dims.size(); i++) {
-        x_dims_vec[i + axis] = x_dims[i];
-      }
-    }
-
-    if (y_dims.size() == max_dim) {
-      for (int i = 0; i < max_dim; i++) {
-        y_dims_vec[i] = y_dims[i];
-      }
-    } else {
-      for (int i = 0; i < y_dims.size(); i++) {
-        y_dims_vec[i + axis] = y_dims[i];
-      }
-    }
-
-    for (int i = 0; i < max_dim; i++) {
-      z_dims_vec[i] = dz_dims[i];
-    }
-    std::vector<int> rdims_for_x;
-    std::vector<int> rdims_for_y;
-    rdims_for_x = get_rdims(x_dims_vec, z_dims_vec);
-    rdims_for_y = get_rdims(y_dims_vec, z_dims_vec);
-    const T* dz_data = dz->data<T>();
-    auto& dev_ctx =
-        ctx.template device_context<paddle::platform::XPUDeviceContext>();
-
-    if (dx != nullptr) {
-      T* dx_data = dx->mutable_data<T>(ctx.GetPlace());
-      if (rdims_for_x.size() == 0) {
-        if (dx_data != dz_data) {
-          framework::TensorCopy(
-              *dz, ctx.GetPlace(),
-              ctx.template device_context<platform::DeviceContext>(), dx);
-        }
-      } else {
-        // For inplace strategy, dx will be stored in addr of dz, which makes
-        // the result of dy wrong.
-        if (dx->IsSharedBufferWith(*dz)) {
-          dx->clear();
-          dx->mutable_data<T>(x->dims(), ctx.GetPlace());
-        }
-
-        int ret = xpu::reduce_sum<XPUType>(
-            dev_ctx.x_context(), reinterpret_cast<const XPUType*>(dz_data),
-            reinterpret_cast<XPUType*>(dx_data), z_dims_vec, rdims_for_x);
-        PADDLE_ENFORCE_XDNN_SUCCESS(ret, "reduce_sum ");
-      }
-    }
-
-    if (dy != nullptr) {
-      T* dy_data = dy->mutable_data<T>(ctx.GetPlace());
-      if (rdims_for_y.size() == 0) {
-        if (dy_data != dz_data) {
-          framework::TensorCopy(
-              *dz, ctx.GetPlace(),
-              ctx.template device_context<platform::DeviceContext>(), dy);
-        }
-      } else {
-        int ret = xpu::reduce_sum<XPUType>(
-            dev_ctx.x_context(), reinterpret_cast<const XPUType*>(dz_data),
-            reinterpret_cast<XPUType*>(dy_data), z_dims_vec, rdims_for_y);
-        PADDLE_ENFORCE_XDNN_SUCCESS(ret, "reduce_sum ");
-      }
-    }
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
+#else
+#include "paddle/fluid/operators/elementwise/elementwise_add_op.h"
+#include "paddle/phi/kernels/gpu/elementwise.h"
+#endif
 
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;
 
+#ifdef PADDLE_WITH_XPU_KP
 REGISTER_OP_KERNEL(elementwise_add, KP, plat::XPUPlace,
-                   ops::ElementwiseAddXPUKPKernel<float>);
-
-REGISTER_OP_KERNEL(elementwise_add_grad, KP, plat::XPUPlace,
-                   ops::ElementwiseAddGradXPUKPKernel<float>);
-
-#endif  // PADDLE_WITH_XPU_KP
+                   ops::ElementwiseAddKernel<plat::XPUDeviceContext, float>);
+#else
+REGISTER_OP_CUDA_KERNEL(
+    grad_add, ops::ElementwiseAddKernel<plat::CUDADeviceContext, float>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, double>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, int>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, int64_t>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, plat::float16>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, plat::bfloat16>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, plat::complex<float>>,
+    ops::ElementwiseAddKernel<plat::CUDADeviceContext, plat::complex<double>>);
+#endif

paddle/fluid/platform/device/xpu/xpu_op_kpfirst_list.h

@@ -27,7 +27,10 @@ using XPUKernelSet =
 using XPUOpMap = std::unordered_map<std::string, XPUKernelSet>;
 
 XPUOpMap& get_kp_ops() {
-  static XPUOpMap s_xpu_kp_kernels{};
+  static XPUOpMap s_xpu_kp_kernels{
+      {"elementwise_add",
+       XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
+  };
 
   return s_xpu_kp_kernels;
 }