support KL label smooth (#35177)

* support KL label smooth

* update UT for KL label_smooth

cmake/external/xpu.cmake

@@ -35,7 +35,7 @@ ELSE ()
 ENDIF()
 
 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}/20210818")
+SET(XPU_BASE_URL "${XPU_BASE_URL_WITHOUT_DATE}/20210826")
 SET(XPU_XRE_URL  "${XPU_BASE_URL}/${XPU_XRE_DIR_NAME}.tar.gz" CACHE STRING "" FORCE)
 SET(XPU_XDNN_URL "${XPU_BASE_URL}/${XPU_XDNN_DIR_NAME}.tar.gz" CACHE STRING "" FORCE)
 SET(XPU_XCCL_URL "${XPU_BASE_URL_WITHOUT_DATE}/20210623/${XPU_XCCL_DIR_NAME}.tar.gz" CACHE STRING "" FORCE)

paddle/fluid/framework/operator.cc

@@ -1254,10 +1254,10 @@ void OperatorWithKernel::ChooseKernel(const RuntimeContext& ctx,
   }
 #endif
 #ifdef PADDLE_WITH_XPU
-  if ((kernel_iter == kernels.end() &&
-       is_xpu_place(expected_kernel_key.place_) &&
-       !paddle::platform::is_xpu_support_op(type_, expected_kernel_key)) ||
-      paddle::platform::is_in_xpu_black_list(type_)) {
+  if (is_xpu_place(expected_kernel_key.place_) &&
+      (kernel_iter == kernels.end() ||
+       !paddle::platform::is_xpu_support_op(type_, expected_kernel_key) ||
+       paddle::platform::is_in_xpu_black_list(type_))) {
     VLOG(3) << "missing XPU kernel: " << type_
             << ", expected_kernel_key:" << expected_kernel_key
             << ", fallbacking to CPU one!";

paddle/fluid/imperative/prepared_operator.cc

@@ -131,10 +131,10 @@ PreparedOp PrepareImpl(const NameVarMap<VarType>& ins,
   auto& kernels = kernels_iter->second;
   auto kernel_iter = kernels.find(expected_kernel_key);
 #ifdef PADDLE_WITH_XPU
-  if ((kernel_iter == kernels.end() &&
-       is_xpu_place(expected_kernel_key.place_) &&
-       !paddle::platform::is_xpu_support_op(op.Type(), expected_kernel_key)) ||
-      paddle::platform::is_in_xpu_black_list(op.Type())) {
+  if (is_xpu_place(expected_kernel_key.place_) &&
+      (kernel_iter == kernels.end() ||
+       !paddle::platform::is_xpu_support_op(op.Type(), expected_kernel_key) ||
+       paddle::platform::is_in_xpu_black_list(op.Type()))) {
     VLOG(3) << "missing XPU kernel: " << op.Type()
             << ", expected_kernel_key:" << expected_kernel_key
             << ", fallbacking to CPU one!";

paddle/fluid/operators/label_smooth_op_xpu.cc

+/* Copyright (c) 2016 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
+#include "paddle/fluid/operators/label_smooth_op.h"
+#include "paddle/fluid/framework/op_registry.h"
+
+namespace paddle {
+namespace operators {
+
+template <typename DeviceContext, typename T>
+class LabelSmoothXPUKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const {
+    auto* out_t = ctx.Output<framework::LoDTensor>("Out");
+    auto* in_t = ctx.Input<framework::LoDTensor>("X");
+    auto* dist_t = ctx.Input<framework::Tensor>("PriorDist");
+    auto label_dim = in_t->dims()[in_t->dims().size() - 1];
+    auto ptr = out_t->mutable_data<T>(ctx.GetPlace());
+
+    auto epsilon = ctx.Attr<float>("epsilon");
+    auto& dev_ctx = ctx.template device_context<DeviceContext>();
+    if (dist_t) {
+      PADDLE_THROW(
+          platform::errors::External("XPU doesn't support dist label smooth"));
+    } else {
+      int r = xpu::label_smooth<T>(dev_ctx.x_context(), in_t->data<T>(), ptr,
+                                   in_t->numel(), epsilon, label_dim);
+      PADDLE_ENFORCE_EQ(
+          r, XPU_SUCCESS,
+          platform::errors::External("XPU API(label_smooth) return wrong "
+                                     "value[%d %s]",
+                                     r, XPUAPIErrorMsg[r]));
+    }
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+REGISTER_OP_XPU_KERNEL(
+    label_smooth,
+    ops::LabelSmoothXPUKernel<paddle::platform::XPUDeviceContext, float>);
+#endif

paddle/fluid/platform/xpu/xpu2_op_list.h

@@ -29,6 +29,8 @@ using XPUOpMap = std::unordered_map<std::string, XPUKernelSet>;
 XPUOpMap& get_kl2_ops() {
   // KL1支持的op，通过op_name, data_type, place来索引
   static XPUOpMap s_xpu2_kernels{
+      {"label_smooth",
+       XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"mul", XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace()),
                             pOpKernelType(vartype::FP16, XPUPlace())})},
       {"elementwise_sub",

python/paddle/fluid/tests/unittests/xpu/test_label_smooth_op_xpu.py

+#   Copyright (c) 2018 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.
+
+from __future__ import print_function
+
+import unittest
+import paddle
+import numpy as np
+import sys
+sys.path.append("..")
+from op_test_xpu import XPUOpTest
+
+paddle.enable_static()
+
+
+class TestLabelSmoothOp(XPUOpTest):
+    def config(self):
+        self.op_type = "label_smooth"
+        self.epsilon = 0.1
+        self.use_xpu = True
+        batch_size, self.label_dim = 10, 12
+        self.label = np.zeros((batch_size, self.label_dim)).astype("float32")
+        nonzero_index = np.random.randint(self.label_dim, size=(batch_size))
+        self.label[np.arange(batch_size), nonzero_index] = 1
+
+    def setUp(self):
+        self.config()
+        smoothed_label = (1 - self.epsilon
+                          ) * self.label + self.epsilon / self.label_dim
+        self.inputs = {'X': self.label}
+        self.attrs = {'epsilon': self.epsilon}
+        self.outputs = {'Out': smoothed_label}
+
+    def test_check_output(self):
+        if not paddle.is_compiled_with_xpu():
+            return
+        self.check_output_with_place(paddle.XPUPlace(0), atol=1e-6)
+
+    def test_check_grad(self):
+        return
+
+
+class TestLabelSmoothOp3D(TestLabelSmoothOp):
+    def setUp(self):
+        super(TestLabelSmoothOp3D, self).setUp()
+        self.inputs['X'] = self.inputs['X'].reshape(
+            [2, -1, self.inputs['X'].shape[-1]])
+        self.outputs['Out'] = self.outputs['Out'].reshape(self.inputs['X']
+                                                          .shape)
+
+
+if __name__ == '__main__':
+    unittest.main()