Add the new XDNN implementation. test=kunlun (#42683)

* Add the new XDNN implementation. test=kunlun

* Add the new XDNN implementation. test=kunlun

* Modify the code based on review, test=kunlun

cmake/external/xpu.cmake

@@ -9,7 +9,7 @@ SET(XPU_RT_LIB_NAME             "libxpurt.so")
 
 if(NOT DEFINED XPU_BASE_URL)
   SET(XPU_BASE_URL_WITHOUT_DATE "https://baidu-kunlun-product.cdn.bcebos.com/KL-SDK/klsdk-dev")
-  SET(XPU_BASE_URL "${XPU_BASE_URL_WITHOUT_DATE}/20220510")
+  SET(XPU_BASE_URL "${XPU_BASE_URL_WITHOUT_DATE}/20220511")
 else()
   SET(XPU_BASE_URL "${XPU_BASE_URL}")
 endif()
@@ -17,7 +17,7 @@ endif()
 # ubuntu and centos: use output by XDNN API team
 if(NOT DEFINED XPU_XDNN_BASE_URL)
   SET(XPU_XDNN_BASE_URL_WITHOUT_DATE "https://klx-sdk-release-public.su.bcebos.com/xdnn/dev")
-  SET(XPU_XDNN_BASE_URL "${XPU_XDNN_BASE_URL_WITHOUT_DATE}/20220510")
+  SET(XPU_XDNN_BASE_URL "${XPU_XDNN_BASE_URL_WITHOUT_DATE}/20220511")
 else()
   SET(XPU_XDNN_BASE_URL "${XPU_XDNN_BASE_URL}")
 endif()

paddle/fluid/operators/log_loss_op_xpu.cc

@@ -12,6 +12,7 @@ limitations under the License. */
 
 #include <memory>
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/platform/device/device_wrapper.h"
 namespace paddle {
 namespace operators {
 
@@ -21,7 +22,6 @@ template <typename DeviceContext, typename T, typename AttrType = T>
 class LogLossXPUKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
-    /*** TODO wait XDNN new interface
     auto* predict = ctx.Input<Tensor>("Predicted");
     auto* labels = ctx.Input<Tensor>("Labels");
     auto* loss = ctx.Output<Tensor>("Loss");
@@ -29,26 +29,15 @@ class LogLossXPUKernel : public framework::OpKernel<T> {
     loss->mutable_data<T>(ctx.GetPlace());
     int n = predict->numel();
     auto& dev_ctx = ctx.template device_context<DeviceContext>();
-        int r =
-            xpu::log_loss_fwd(dev_ctx.x_context(), n, epsilon,
-    predict->data<T>(),
-                              labels->data<T>(), loss->data<T>());
-        PADDLE_ENFORCE_EQ(
-            r, xpu::Error_t::SUCCESS,
-            platform::errors::External(
-                "XPU log_loss kernel return wrong value[%d], please check
-    whether "
-                "Baidu Kunlun Card is properly installed.",
-                r));
-    ***/
+    int r = xpu::log_loss(dev_ctx.x_context(), predict->data<T>(),
+                          labels->data<T>(), loss->data<T>(), n, epsilon);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "log_loss");
   }
 };
 template <typename DeviceContext, typename T, typename AttrType = T>
 class LogLossGradXPUKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
-    /*** TODO wait XDNN new interface
-
     auto* predict = ctx.Input<Tensor>("Predicted");
     auto* labels = ctx.Input<Tensor>("Labels");
     auto* dloss = ctx.Input<Tensor>(framework::GradVarName("Loss"));
@@ -60,28 +49,20 @@ class LogLossGradXPUKernel : public framework::OpKernel<T> {
     dpred->mutable_data<T>(ctx.GetPlace());
     int n = predict->numel();
     auto& dev_ctx = ctx.template device_context<DeviceContext>();
-        int r = xpu::log_loss_bwd(dev_ctx.x_context(), n, epsilon,
-                                  predict->data<T>(), labels->data<T>(),
-                                  dloss->data<T>(), dpred->data<T>());
-        PADDLE_ENFORCE_EQ(
-            r, xpu::Error_t::SUCCESS,
-            platform::errors::External(
-                "XPU log_loss kernel return wrong value[%d], please check
-    whether "
-                "Baidu Kunlun Card is properly installed.",
-                r));
-    ***/
+    int r = xpu::log_loss_grad(dev_ctx.x_context(), predict->data<T>(),
+                               labels->data<T>(), dloss->data<T>(),
+                               dpred->data<T>(), n, epsilon);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "log_loss_grad");
   }
 };
 
 }  // namespace operators
 }  // namespace paddle
-// namespace ops = paddle::operators;
-// REGISTER_OP_XPU_KERNEL(
-//     log_loss, ops::LogLossXPUKernel<paddle::platform::XPUDeviceContext,
-//     float>);
-// REGISTER_OP_XPU_KERNEL(
-//     log_loss_grad,
-//     ops::LogLossGradXPUKernel<paddle::platform::XPUDeviceContext, float>);
+namespace ops = paddle::operators;
+REGISTER_OP_XPU_KERNEL(
+    log_loss, ops::LogLossXPUKernel<paddle::platform::XPUDeviceContext, float>);
+REGISTER_OP_XPU_KERNEL(
+    log_loss_grad,
+    ops::LogLossGradXPUKernel<paddle::platform::XPUDeviceContext, float>);
 
 #endif

paddle/fluid/operators/optimizers/lamb_op_xpu.cc

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/optimizers/lamb_op.h"
 #include "gflags/gflags.h"
+#include "paddle/fluid/platform/device/device_wrapper.h"
 
 namespace paddle {
 namespace operators {
@@ -25,7 +26,6 @@ template <typename DeviceContext, typename T>
 class LambOpXPUKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
-    /*** TODO wait XDNN new interface
     using paddle::framework::LoDTensor;
     const auto* param_var = ctx.InputVar("Param");
     PADDLE_ENFORCE_EQ(param_var->IsType<framework::LoDTensor>(), true,
@@ -49,15 +49,12 @@ class LambOpXPUKernel : public framework::OpKernel<T> {
                                  "Moment1", "Lamb");
     auto& mom2 = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Moment2"), "Input",
                                  "Moment2", "Lamb");
-        auto& lr = GET_DATA_SAFELY(ctx.Input<LoDTensor>("LearningRate"),
-    "Input",
+    auto& lr = GET_DATA_SAFELY(ctx.Input<LoDTensor>("LearningRate"), "Input",
                                "LearningRate", "Lamb");
 
-        auto& beta1_pow = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Beta1Pow"),
-    "Input",
+    auto& beta1_pow = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Beta1Pow"), "Input",
                                       "Beta1Pow", "Lamb");
-        auto& beta2_pow = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Beta2Pow"),
-    "Input",
+    auto& beta2_pow = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Beta2Pow"), "Input",
                                       "Beta2Pow", "Lamb");
 
     auto& param_out = GET_DATA_SAFELY(ctx.Output<LoDTensor>("ParamOut"),
@@ -66,70 +63,38 @@ class LambOpXPUKernel : public framework::OpKernel<T> {
                                      "Output", "Moment1Out", "Lamb");
     auto& mom2_out = GET_DATA_SAFELY(ctx.Output<LoDTensor>("Moment2Out"),
                                      "Output", "Moment2Out", "Lamb");
-        auto& beta1_pow_out =
-    GET_DATA_SAFELY(ctx.Output<LoDTensor>("Beta1PowOut"),
+    auto& beta1_pow_out = GET_DATA_SAFELY(ctx.Output<LoDTensor>("Beta1PowOut"),
                                           "Output", "Beta1PowOut", "Lamb");
-        auto& beta2_pow_out =
-    GET_DATA_SAFELY(ctx.Output<LoDTensor>("Beta2PowOut"),
+    auto& beta2_pow_out = GET_DATA_SAFELY(ctx.Output<LoDTensor>("Beta2PowOut"),
                                           "Output", "Beta2PowOut", "Lamb");
     auto& dev_ctx = ctx.template device_context<DeviceContext>();
 
     if (grad_var->IsType<framework::LoDTensor>()) {
       auto& grad = *ctx.Input<LoDTensor>("Grad");
-          int r = xpu::lamb(dev_ctx.x_context(), grad.template data<T>(),
-                            mom1.template data<T>(), mom2.template data<T>(),
-                            param.template data<T>(), beta1_pow.template
-    data<T>(),
-                            beta2_pow.template data<T>(), beta1, beta2, epsilon,
-                            weight_decay, lr.template data<T>(),
+      int r = xpu::lamb(
+          dev_ctx.x_context(), grad.template data<T>(), mom1.template data<T>(),
+          mom2.template data<T>(), param.template data<T>(),
+          beta1_pow.template data<T>(), beta2_pow.template data<T>(),
           mom1_out.template mutable_data<T>(ctx.GetPlace()),
           mom2_out.template mutable_data<T>(ctx.GetPlace()),
           param_out.template mutable_data<T>(ctx.GetPlace()),
-                            beta1_pow_out.template
-    mutable_data<T>(ctx.GetPlace()),
-                            beta2_pow_out.template
-    mutable_data<T>(ctx.GetPlace()),
-                            param.numel());
+          beta1_pow_out.template mutable_data<T>(ctx.GetPlace()),
+          beta2_pow_out.template mutable_data<T>(ctx.GetPlace()), beta1, beta2,
+          epsilon, weight_decay, lr.template data<T>(), param.numel());
 
-          if (r == xpu::Error_t::INVALID_PARAM) {
-            PADDLE_ENFORCE_EQ(
-                r, xpu::Error_t::SUCCESS,
-                platform::errors::InvalidArgument(
-                    "XPU kernel error of LambOp, error message: INVALID_PARAM, "
-                    "please check your input & output."));
-          } else if (r == xpu::Error_t::RUNTIME_ERROR) {
-            PADDLE_ENFORCE_EQ(r, xpu::Error_t::SUCCESS,
-                              platform::errors::Unavailable(
-                                  "XPU kernel error of LambOp, error message: "
-                                  "RUNTIME_ERROR, please check whether Baidu "
-                                  "Kunlun Card is properly installed."));
-          } else if (r == xpu::Error_t::NO_ENOUGH_WORKSPACE) {
-            PADDLE_ENFORCE_EQ(r, xpu::Error_t::SUCCESS,
-                              platform::errors::ResourceExhausted(
-                                  "XPU kernel error of LambOp, error "
-                                  "message: NO_ENOUGH_WORKSPACE, XPU "
-                                  "has no enough memory."));
-          } else {
-            PADDLE_ENFORCE_EQ(r, xpu::Error_t::SUCCESS,
-                              platform::errors::ResourceExhausted(
-                                  "XPU kernel error of LambOp, error "
-                                  "message: OTHER "
-                                  "XPU API returns error code: %d.",
-                                  r));
-          }
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "lamb");
     } else {
       PADDLE_THROW(platform::errors::InvalidArgument(
           "Variable type not supported by lamb_op. Expect LoDTensor, "
           "but got %s",
           framework::ToTypeName(param_var->Type())));
     }
-    **/
   }
 };
 }  // namespace operators
 }  // namespace paddle
 
-// namespace ops = paddle::operators;
-// REGISTER_OP_XPU_KERNEL(
-//     lamb, ops::LambOpXPUKernel<paddle::platform::XPUDeviceContext, float>);
+namespace ops = paddle::operators;
+REGISTER_OP_XPU_KERNEL(
+    lamb, ops::LambOpXPUKernel<paddle::platform::XPUDeviceContext, float>);
 #endif

paddle/fluid/operators/optimizers/rmsprop_op_xpu.cc

@@ -17,6 +17,7 @@ limitations under the License. */
 #include <gflags/gflags.h>
 #include <iostream>
 #include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/platform/device/device_wrapper.h"
 
 namespace paddle {
 namespace operators {
@@ -40,7 +41,6 @@ template <typename DeviceContext, typename T>
 class RmspropOpXPUKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
-    /*** TODO wait XDNN new interface
     using paddle::framework::LoDTensor;
 
     // check Param & Grad tensor type
@@ -67,8 +67,7 @@ class RmspropOpXPUKernel : public framework::OpKernel<T> {
                                   "Param", "Rmsprop");
     auto& meanSquare = GET_DATA_SAFELY(ctx.Input<LoDTensor>("MeanSquare"),
                                        "Input", "MeanSquare", "Rmsprop");
-        auto& grad = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Grad"), "Input",
-    "Grad",
+    auto& grad = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Grad"), "Input", "Grad",
                                  "Rmsprop");
     auto& mom = GET_DATA_SAFELY(ctx.Input<LoDTensor>("Moment"), "Input",
                                 "Moment", "Rmsprop");
@@ -91,10 +90,8 @@ class RmspropOpXPUKernel : public framework::OpKernel<T> {
                                       "Output", "ParamOut", "Rmsprop");
     auto& mom_out = GET_DATA_SAFELY(ctx.Output<LoDTensor>("MomentOut"),
                                     "Output", "MomentOut", "Rmsprop");
-        auto& mom_sqrt_out =
-    GET_DATA_SAFELY(ctx.Output<LoDTensor>("MeanSquareOut"),
-                                             "Output", "MeanSquareOut",
-    "Rmsprop");
+    auto& mom_sqrt_out = GET_DATA_SAFELY(ctx.Output<LoDTensor>("MeanSquareOut"),
+                                         "Output", "MeanSquareOut", "Rmsprop");
     auto& dev_ctx = ctx.template device_context<DeviceContext>();
 
     ///// rmsprop优化算法
@@ -109,53 +106,23 @@ class RmspropOpXPUKernel : public framework::OpKernel<T> {
     /// const float* ms, const float* g, const float* mom,
     /// float epsilon, float rho, float momentum, float lr,
     /// float *ms_out, float *mom_out, float *p_out, int n)
-        int r = xpu::rmsprop(dev_ctx.x_context(), param.template data<T>(),
-                             meanSquare.template data<T>(), grad.template
-    data<T>(),
-                             mom.template data<T>(), epsilon, decay, momentum,
-    lr,
-                             mom_sqrt_out.template
-    mutable_data<T>(ctx.GetPlace()),
-                             mom_out.template mutable_data<T>(ctx.GetPlace()),
+    int r = xpu::rmsprop(dev_ctx.x_context(), grad.template data<T>(),
+                         param.template data<T>(),
+                         meanSquare.template data<T>(), mom.template data<T>(),
                          param_out.template mutable_data<T>(ctx.GetPlace()),
-                             param.numel());
+                         mom_sqrt_out.template mutable_data<T>(ctx.GetPlace()),
+                         mom_out.template mutable_data<T>(ctx.GetPlace()),
+                         epsilon, decay, momentum, lr, param.numel());
 
-        if (r == xpu::Error_t::INVALID_PARAM) {
-          PADDLE_ENFORCE_EQ(
-              r, xpu::Error_t::SUCCESS,
-              platform::errors::InvalidArgument(
-                  "XPU kernel error of RmspropOp, error message: INVALID_PARAM,
-    "
-                  "please check your input & output."));
-        } else if (r == xpu::Error_t::RUNTIME_ERROR) {
-          PADDLE_ENFORCE_EQ(r, xpu::Error_t::SUCCESS,
-                            platform::errors::Unavailable(
-                                "XPU kernel error of RmspropOp, error message: "
-                                "RUNTIME_ERROR, please check whether Baidu "
-                                "Kunlun Card is properly installed."));
-        } else if (r == xpu::Error_t::NO_ENOUGH_WORKSPACE) {
-          PADDLE_ENFORCE_EQ(r, xpu::Error_t::SUCCESS,
-                            platform::errors::ResourceExhausted(
-                                "XPU kernel error of RmspropOp, error "
-                                "message: NO_ENOUGH_WORKSPACE, XPU "
-                                "has no enough memory."));
-        } else {
-          PADDLE_ENFORCE_EQ(r, xpu::Error_t::SUCCESS,
-                            platform::errors::ResourceExhausted(
-                                "XPU kernel error of RmspropOp, error "
-                                "message: OTHER "
-                                "XPU API returns error code: %d.",
-                                r));
-        }
-    ***/
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "rmsprop");
   }
 };
 
 }  // namespace operators
 }  // namespace paddle
 
-// namespace ops = paddle::operators;
-// REGISTER_OP_XPU_KERNEL(
-//     rmsprop,
-//     ops::RmspropOpXPUKernel<paddle::platform::XPUDeviceContext, float>);
+namespace ops = paddle::operators;
+REGISTER_OP_XPU_KERNEL(
+    rmsprop,
+    ops::RmspropOpXPUKernel<paddle::platform::XPUDeviceContext, float>);
 #endif

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

@@ -145,6 +145,7 @@ XPUOpMap& get_kl1_ops() {
       {"hard_switch", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"iou_similarity",
        XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
+      {"lamb", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"layer_norm_grad",
        XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"layer_norm", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
@@ -174,6 +175,9 @@ XPUOpMap& get_kl1_ops() {
                                   pOpKernelType(vartype::INT32, XPUPlace()),
                                   pOpKernelType(vartype::INT64, XPUPlace()),
                                   pOpKernelType(vartype::FP32, XPUPlace())})},
+      {"log_loss_grad",
+       XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
+      {"log_loss", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"logsumexp", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"log", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"lookup_table_v2_grad",
@@ -232,6 +236,7 @@ XPUOpMap& get_kl1_ops() {
                                  pOpKernelType(vartype::INT32, XPUPlace()),
                                  pOpKernelType(vartype::BOOL, XPUPlace()),
                                  pOpKernelType(vartype::FP32, XPUPlace())})},
+      {"rmsprop", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"rnn_grad", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"rnn", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"roi_align_grad",