[XPU] support index_put grad (#67173)

lj970926 · web-flow · commit 7464d9e0519e · 2024-08-13T10:41:06.000+08:00
* [XPU] support index_put grad

* bug fix
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -576,6 +576,10 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
+      {"index_put_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"index_sample_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"index_sample",
        XPUKernelSet({phi::DataType::INT8,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -570,6 +570,10 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
+      {"index_put_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"index_sample_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"index_sample",
        XPUKernelSet({phi::DataType::INT8,
diff --git a/paddle/phi/kernels/xpu/index_put_grad_kernel.cc b/paddle/phi/kernels/xpu/index_put_grad_kernel.cc
@@ -0,0 +1,144 @@
+// Copyright (c) 2024 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
+//
+//     https://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.
+
+#include "paddle/phi/kernels/index_put_grad_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/funcs/index_put_utils.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/xpu/index_put_xpu_utils.h"
+
+namespace phi {
+template <typename T, typename Context>
+void IndexPutGradKernel(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        const std::vector<const DenseTensor*>& indices_v,
+                        const DenseTensor& value,
+                        const DenseTensor& out_grad,
+                        bool accumulate,
+                        DenseTensor* x_grad,
+                        DenseTensor* value_grad) {
+  PADDLE_ENFORCE_EQ(
+      x.dtype(),
+      value.dtype(),
+      common::errors::InvalidArgument(
+          "The data type of tensor value must be same to the data type "
+          "of tensor x."));
+  // All bool indices are converted to integers currently
+  std::vector<DenseTensor> tmp_args;
+  std::vector<const DenseTensor*> int_indices_v =
+      funcs::DealWithBoolIndices<T, Context>(dev_ctx, indices_v, &tmp_args);
+
+  if (int_indices_v.empty()) {
+    if (x_grad) {
+      phi::Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    }
+    if (value_grad) {
+      FullKernel<T, Context>(dev_ctx,
+                             common::vectorize(value_grad->dims()),
+                             0.0f,
+                             value_grad->dtype(),
+                             value_grad);
+    }
+    return;
+  }
+
+  auto bd_dims = funcs::BroadCastTensorsDims(int_indices_v);
+  DenseTensor res_indices(DataType::INT64);
+  // Broadcast and merge indices
+  XPUDealWithIndices<Context>(dev_ctx, int_indices_v, bd_dims, &res_indices);
+  auto index_shape = common::vectorize<int64_t>(res_indices.dims());
+  xpu::VectorParam<int64_t> index_param = {
+      nullptr, res_indices.numel(), res_indices.data<int64_t>()};
+  auto xshape = common::vectorize<int64_t>(x.dims());
+  xpu::VectorParam<int64_t> xshape_param = {
+      xshape.data(), static_cast<int64_t>(xshape.size()), nullptr};
+
+  int64_t value_rank = bd_dims.size() + (xshape.size() - int_indices_v.size());
+  std::vector<int64_t> value_shape_bd(value_rank);
+  std::copy(index_shape.begin(), index_shape.end() - 1, value_shape_bd.begin());
+  std::copy(xshape.begin() + int_indices_v.size(),
+            xshape.end(),
+            value_shape_bd.begin() + index_shape.size() - 1);
+  auto value_shape = common::vectorize<int64_t>(value_grad->dims());
+  int ret = xpu::SUCCESS;
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  if (x_grad) {
+    phi::Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    if (!accumulate) {
+      DenseTensor zero_tensor(x_grad->dtype());
+      FullKernel<T, Context>(
+          dev_ctx, value_shape_bd, 0.0f, zero_tensor.dtype(), &zero_tensor);
+      ret = xpu::scatter_nd<XPUType, int64_t>(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(x_grad->data<T>()),
+          reinterpret_cast<const XPUType*>(zero_tensor.data<T>()),
+          reinterpret_cast<XPUType*>(x_grad->data<T>()),
+          index_param,
+          xshape_param,
+          index_shape,
+          false);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, "scatter_nd");
+    }
+  }
+  if (value_grad) {
+    dev_ctx.template Alloc<T>(value_grad);
+    if (value_shape != value_shape_bd) {
+      std::vector<int64_t> compress_dims;
+      std::vector<int64_t> dims_without_1;
+      funcs::CalCompressedDimsWith1AndWithout1(
+          &value_shape_bd, &value_shape, &compress_dims, &dims_without_1);
+      DenseTensor value_grad_bd(value_grad->dtype());
+      value_grad_bd.Resize(common::make_ddim(value_shape_bd));
+      dev_ctx.template Alloc<T>(&value_grad_bd);
+      ret = xpu::gather_nd<XPUType, int64_t>(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(out_grad.data<T>()),
+          res_indices.data<int64_t>(),
+          reinterpret_cast<XPUType*>(value_grad_bd.data<T>()),
+          xshape_param,
+          index_shape);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, "gather_nd");
+      IntArray v_axis(compress_dims);
+      auto pre_dims = value_grad->dims();
+      SumKernel<T>(dev_ctx,
+                   value_grad_bd,
+                   v_axis,
+                   value_grad->dtype(),
+                   false,
+                   value_grad);
+      value_grad->Resize(pre_dims);
+    } else {
+      ret = xpu::gather_nd<XPUType, int64_t>(
+          dev_ctx.x_context(),
+          reinterpret_cast<const XPUType*>(out_grad.data<T>()),
+          res_indices.data<int64_t>(),
+          reinterpret_cast<XPUType*>(value_grad->data<T>()),
+          xshape_param,
+          index_shape);
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, "gather_nd");
+    }
+  }
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(index_put_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::IndexPutGradKernel,
+                   float,
+                   int,
+                   int64_t) {}
diff --git a/paddle/phi/kernels/xpu/index_put_kernel.cc b/paddle/phi/kernels/xpu/index_put_kernel.cc
@@ -18,59 +18,9 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/expand_kernel.h"
 #include "paddle/phi/kernels/funcs/index_put_utils.h"
-#include "paddle/phi/kernels/stack_kernel.h"
+#include "paddle/phi/kernels/xpu/index_put_xpu_utils.h"
 
 namespace phi {
-template <typename Context>
-void XPUDealWithIndices(const Context& dev_ctx,
-                        const std::vector<const DenseTensor*>& int_indices_v,
-                        DDim bd_dim,
-                        DenseTensor* out) {
-  std::vector<DenseTensor> tmp_indices_v;
-  for (size_t i = 0; i < int_indices_v.size(); ++i) {
-    // Use int64 for all indices Because XPU needs to merge all indices into a
-    // single tensor. Same with CPU and GPU.
-    DenseTensor casted_index;
-    if (int_indices_v[i]->dtype() == DataType::INT32) {
-      casted_index =
-          phi::Cast<int, Context>(dev_ctx, *int_indices_v[i], DataType::INT64);
-    } else {
-      casted_index = *int_indices_v[i];
-    }
-
-    DenseTensor expanded_index(DataType::INT64);
-    if (casted_index.dims() == bd_dim) {
-      expanded_index = casted_index;
-    } else {
-      expanded_index.Resize(bd_dim);
-      ExpandKernel<int64_t, Context>(
-          dev_ctx,
-          casted_index,
-          IntArray(common::vectorize<int64_t>(bd_dim)),
-          &expanded_index);
-    }
-
-    tmp_indices_v.emplace_back(expanded_index);
-  }
-
-  auto bd_dim_vec = common::vectorize<int64_t>(bd_dim);
-  std::vector<int64_t> stacked_dim_vec(bd_dim.size() + 1);
-  std::copy(bd_dim_vec.begin(), bd_dim_vec.end(), stacked_dim_vec.begin());
-  stacked_dim_vec.back() = int_indices_v.size();
-  out->Resize(common::make_ddim(stacked_dim_vec));
-
-  std::vector<const DenseTensor*> tmp_indices_ptr(tmp_indices_v.size(),
-                                                  nullptr);
-  for (size_t i = 0; i < tmp_indices_ptr.size(); ++i) {
-    tmp_indices_ptr[i] = &tmp_indices_v[i];
-  }
-
-  StackKernel<int64_t, Context>(dev_ctx, tmp_indices_ptr, -1, out);
-  if (dev_ctx.x_context()->xpu_stream) {
-    dev_ctx.Wait();
-  }
-}
-
 template <typename T, typename Context>
 void IndexPutKernel(const Context& dev_ctx,
                     const DenseTensor& x,
diff --git a/paddle/phi/kernels/xpu/index_put_xpu_utils.h b/paddle/phi/kernels/xpu/index_put_xpu_utils.h
@@ -0,0 +1,69 @@
+// Copyright (c) 2024 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
+//
+//     https://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.
+
+#pragma once
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+
+#include "paddle/phi/kernels/expand_kernel.h"
+#include "paddle/phi/kernels/stack_kernel.h"
+
+namespace phi {
+template <typename Context>
+void XPUDealWithIndices(const Context& dev_ctx,
+                        const std::vector<const DenseTensor*>& int_indices_v,
+                        DDim bd_dim,
+                        DenseTensor* out) {
+  std::vector<DenseTensor> tmp_indices_v;
+  for (size_t i = 0; i < int_indices_v.size(); ++i) {
+    // Use int64 for all indices Because XPU needs to merge all indices into a
+    // single tensor. Same with CPU and GPU.
+    DenseTensor casted_index;
+    if (int_indices_v[i]->dtype() == DataType::INT32) {
+      casted_index =
+          phi::Cast<int, Context>(dev_ctx, *int_indices_v[i], DataType::INT64);
+    } else {
+      casted_index = *int_indices_v[i];
+    }
+
+    DenseTensor expanded_index(DataType::INT64);
+    if (casted_index.dims() == bd_dim) {
+      expanded_index = casted_index;
+    } else {
+      expanded_index.Resize(bd_dim);
+      ExpandKernel<int64_t, Context>(
+          dev_ctx,
+          casted_index,
+          IntArray(common::vectorize<int64_t>(bd_dim)),
+          &expanded_index);
+    }
+
+    tmp_indices_v.emplace_back(expanded_index);
+  }
+
+  auto bd_dim_vec = common::vectorize<int64_t>(bd_dim);
+  std::vector<int64_t> stacked_dim_vec(bd_dim.size() + 1);
+  std::copy(bd_dim_vec.begin(), bd_dim_vec.end(), stacked_dim_vec.begin());
+  stacked_dim_vec.back() = int_indices_v.size();
+  out->Resize(common::make_ddim(stacked_dim_vec));
+
+  std::vector<const DenseTensor*> tmp_indices_ptr(tmp_indices_v.size(),
+                                                  nullptr);
+  for (size_t i = 0; i < tmp_indices_ptr.size(); ++i) {
+    tmp_indices_ptr[i] = &tmp_indices_v[i];
+  }
+
+  StackKernel<int64_t, Context>(dev_ctx, tmp_indices_ptr, -1, out);
+}
+}  // namespace phi
diff --git a/test/xpu/test_index_put_op_xpu.py b/test/xpu/test_index_put_op_xpu.py
@@ -87,7 +87,7 @@ def __init__(self):
     class TestXPUIndexPutOp(XPUOpTest):
         def setUp(self):
             self.op_type = "index_put"
-            self.x_dtype = self.in_type
+            self.dtype = self.in_type
             self.mixed_indices = False
             self.is_all_false = False
             self.place = paddle.XPUPlace(0)
@@ -104,11 +104,11 @@ def set_case(self):
 
         def init_data(self):
             x_np = ((np.random.random(self.x_shape) - 0.5) * 10.0).astype(
-                self.x_dtype
+                self.dtype
             )
             value_np = (
                 (np.random.random(self.value_shape) - 0.5) * 10.0
-            ).astype(self.x_dtype)
+            ).astype(self.dtype)
 
             if self.mixed_indices:
                 tmp_indices_np1 = gen_indices_np(
