[XPU][PHI Kernels] add index_put kernel for xpu (#56169)

lj970926 · web-flow · commit 332a73b18e06 · 2023-08-22T14:23:15.000+08:00
* add inverse kernel for xpu

* add more kernels

* add index_put kernel for xpu

* add index_put kernel for xpu

* remove unused headers

* refine test

* wait to avoid memory bugs for xpu

* refine inverse
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -325,6 +325,7 @@ XPUOpMap& get_kl2_ops() {
       {"fill_any_like",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
+                     phi::DataType::BOOL,
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32})},
       {"fill_diagonal_tensor",
@@ -449,6 +450,10 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
+      {"index_put",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"index_sample_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"index_sample",
        XPUKernelSet({phi::DataType::INT8,
@@ -465,6 +470,8 @@ XPUOpMap& get_kl2_ops() {
       {"instance_norm",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"instance_norm_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"inverse",
+       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT64})},
       {"iou_similarity", XPUKernelSet({phi::DataType::FLOAT32})},
       {"label_smooth", XPUKernelSet({phi::DataType::FLOAT32})},
       {"lamb", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
@@ -687,6 +694,12 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::INT64,
                      phi::DataType::FLOAT16,
                      phi::DataType::BOOL})},
+      {"set_value_with_tensor",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64,
+                     phi::DataType::BOOL})},
       {"set_value_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
@@ -966,7 +979,9 @@ XPUOpMap& get_kl2_ops() {
       {"cos", XPUKernelSet({phi::DataType::FLOAT32})},
       {"cos_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"linspace",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::INT32})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"randint", XPUKernelSet({phi::DataType::INT32, phi::DataType::INT64})},
       {"group_norm", XPUKernelSet({phi::DataType::FLOAT32})},
       {"group_norm_grad", XPUKernelSet({phi::DataType::FLOAT32})},
diff --git a/paddle/phi/kernels/funcs/index_put_utils.h b/paddle/phi/kernels/funcs/index_put_utils.h
@@ -105,6 +105,9 @@ std::vector<const phi::DenseTensor*> DealWithBoolIndices(
       }
       SplitWithNumKernel<int64_t, Context>(
           dev_ctx, nonzero_indices, rank, 1, integer_indices);
+#ifdef PADDLE_WITH_XPU
+      dev_ctx.Wait();
+#endif
 
     } else if ((indices_v[i]->dtype() == phi::DataType::INT64) ||
                (indices_v[i]->dtype() == phi::DataType::INT32)) {
diff --git a/paddle/phi/kernels/xpu/index_put_kernel.cc b/paddle/phi/kernels/xpu/index_put_kernel.cc
@@ -0,0 +1,148 @@
+// Copyright (c) 2023 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_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#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"
+
+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(vectorize<int64_t>(bd_dim)),
+                                     &expanded_index);
+    }
+
+    tmp_indices_v.emplace_back(expanded_index);
+  }
+
+  auto bd_dim_vec = 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(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);
+  dev_ctx.Wait();
+}
+
+template <typename T, typename Context>
+void IndexPutKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const std::vector<const DenseTensor*>& indices,
+                    const DenseTensor& value,
+                    bool accumulate,
+                    DenseTensor* out) {
+  PADDLE_ENFORCE_EQ(
+      x.dtype(),
+      value.dtype(),
+      phi::errors::InvalidArgument(
+          "The data type of tensor value must be same to the data type "
+          "of tensor x."));
+  PADDLE_ENFORCE_EQ(indices.empty(),
+                    false,
+                    phi::errors::InvalidArgument("Indices cannot be empty."));
+  const int64_t total_dims = x.dims().size();
+  PADDLE_ENFORCE_LE(
+      total_dims,
+      6,
+      errors::InvalidArgument("Dims of input tensor should be less than 7."));
+
+  // 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, &tmp_args);
+  if (int_indices_v.empty()) {
+    if (!out->initialized()) {
+      phi::Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+    }
+    return;
+  }
+
+  using XPUT = typename XPUTypeTrait<T>::Type;
+  auto out_data = dev_ctx.template Alloc<T>(out);
+  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 = vectorize<int64_t>(res_indices.dims());
+  auto x_shape = vectorize<int64_t>(x.dims());
+
+  const T* value_data = value.data<T>();
+
+  // Broadcast value
+  auto value_shape = vectorize<int64_t>(value.dims());
+  int64_t value_rank = bd_dims.size() + (x_shape.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(x_shape.begin() + int_indices_v.size(),
+            x_shape.end(),
+            value_shape_bd.begin() + index_shape.size() - 1);
+
+  DenseTensor value_bd(value.dtype());
+
+  if (value_shape != value_shape_bd) {
+    value_bd.Resize(make_ddim(value_shape_bd));
+    ExpandKernel<T, Context>(
+        dev_ctx, value, IntArray(value_shape_bd), &value_bd);
+    value_data = value_bd.data<T>();
+  }
+
+  int r =
+      xpu::index_put<XPUT, int64_t>(dev_ctx.x_context(),
+                                    reinterpret_cast<const XPUT*>(x.data<T>()),
+                                    reinterpret_cast<const XPUT*>(value_data),
+                                    res_indices.data<int64_t>(),
+                                    reinterpret_cast<XPUT*>(out_data),
+                                    x_shape,
+                                    index_shape,
+                                    accumulate);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "index_put");
+  dev_ctx.Wait();
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    index_put, XPU, ALL_LAYOUT, phi::IndexPutKernel, float, int, int64_t) {}
diff --git a/paddle/phi/kernels/xpu/inverse_kernel.cc b/paddle/phi/kernels/xpu/inverse_kernel.cc
@@ -0,0 +1,61 @@
+// Copyright (c) 2023 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/inverse_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void InverseKernel(const Context& dev_ctx,
+                   const DenseTensor& x,
+                   DenseTensor* out) {
+  using XPUT = typename XPUTypeTrait<T>::Type;
+  auto out_data = dev_ctx.template Alloc<T>(out);
+
+  int64_t x_dims_len = x.dims().size();
+  PADDLE_ENFORCE_GT(
+      x_dims_len,
+      1,
+      phi::errors::InvalidArgument(
+          "Dimensions of input should be greater than 1, but got %d.",
+          x_dims_len));
+
+  int64_t n = x.dims()[x_dims_len - 1];
+  int64_t batch = x_dims_len > 2 ? x.numel() / (n * n) : 1;
+  PADDLE_ENFORCE_LE(n * n * sizeof(T),
+                    8192,
+                    phi::errors::InvalidArgument(
+                        "The size of a single matrix (%d bytes) exceeds the "
+                        "maxinum numbers of bytes xpu supports (8192).",
+                        n * n * sizeof(T)));
+  auto RAII_GUARD = xpu::ctx_guard(dev_ctx.x_context());
+  auto* info_xpu = RAII_GUARD.alloc_l3_or_gm<int>(batch);
+  // Xpu inverse api has check for singularity itself.
+  int r = xpu::inverse<XPUT>(dev_ctx.x_context(),
+                             reinterpret_cast<const XPUT*>(x.data<T>()),
+                             reinterpret_cast<XPUT*>(out_data),
+                             info_xpu,
+                             batch,
+                             n);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "inverse");
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    inverse, XPU, ALL_LAYOUT, phi::InverseKernel, float, double) {}
diff --git a/paddle/phi/kernels/xpu/linspace_kernel.cc b/paddle/phi/kernels/xpu/linspace_kernel.cc
@@ -81,7 +81,7 @@ void LinspaceKernel(const Context& ctx,
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
-    linspace, XPU, ALL_LAYOUT, phi::LinspaceKernel, float, int32_t) {
+    linspace, XPU, ALL_LAYOUT, phi::LinspaceKernel, float, int32_t, int64_t) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
diff --git a/paddle/phi/kernels/xpu/set_value_kernel.cc b/paddle/phi/kernels/xpu/set_value_kernel.cc
@@ -429,4 +429,5 @@ PD_REGISTER_KERNEL(set_value_with_tensor,
                    float,
                    phi::dtype::float16,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   bool) {}
diff --git a/paddle/phi/kernels/xpu/split_kernel.cc b/paddle/phi/kernels/xpu/split_kernel.cc
@@ -34,7 +34,8 @@ void SplitKernel(const Context& dev_ctx,
   for (size_t j = 0; j < outs.size(); ++j) {
     dev_ctx.template Alloc<T>(outs[j]);
     out_ptrs.push_back(reinterpret_cast<XPUType*>(outs[j]->data<T>()));
-    split_lists.push_back(outs[j]->dims()[axis]);
+    split_lists.push_back(axis < outs[j]->dims().size() ? outs[j]->dims()[axis]
+                                                        : 1);
   }
   if (x.numel() == 0) {
     return;
diff --git a/test/xpu/test_fill_any_like_op_xpu.py b/test/xpu/test_fill_any_like_op_xpu.py
@@ -75,6 +75,8 @@ class TestFillAnyLikeOp5(TestFillAnyLikeOp):
         def set_value(self):
             if self.dtype == "float16":
                 self.value = 0.05
+            elif self.dtype == np.bool_:
+                self.value = 1.0
             else:
                 self.value = 5.0
 
diff --git a/test/xpu/test_index_put_op_xpu.py b/test/xpu/test_index_put_op_xpu.py
diff --git a/test/xpu/test_inverse_op_xpu.py b/test/xpu/test_inverse_op_xpu.py

Original file line number	Diff line number	Diff line change
`@@ -105,6 +105,9 @@ std::vector<const phi::DenseTensor*> DealWithBoolIndices(`
`105`	`105`	`}`
`106`	`106`	`SplitWithNumKernel<int64_t, Context>(`
`107`	`107`	`dev_ctx, nonzero_indices, rank, 1, integer_indices);`
	`108`	`+#ifdef PADDLE_WITH_XPU`
	`109`	`+ dev_ctx.Wait();`
	`110`	`+#endif`
`108`	`111`
`109`	`112`	`} else if ((indices_v[i]->dtype() == phi::DataType::INT64) \|\|`
`110`	`113`	`(indices_v[i]->dtype() == phi::DataType::INT32)) {`