[NPU] GatherND and Gather question fix. (PaddlePaddle#1288)

Author: xuanyuanminzheng

backends/npu/kernels/gather_kernel.cc

@@ -118,11 +118,8 @@ void GatherGradKernel(const Context& dev_ctx,
   zeroslike_xout.Resize(x.dims());
 
   // step3: scatter(x_grad)
-  const auto& runner_scatter = NpuOpRunner("TensorScatterUpdate",
-                                           {zeroslike_xout, *p_index, out_grad},
-                                           {*x_grad},
-                                           {});
-  runner_scatter.Run(stream);
+  EXEC_NPU_CMD(
+      aclnnScatterNd, dev_ctx, zeroslike_xout, *p_index, out_grad, *x_grad);
 }
 
 }  // namespace custom_kernel

backends/npu/kernels/gather_nd_kernel.cc

@@ -17,11 +17,64 @@
 
 namespace custom_kernel {
 
+template <typename T, typename Context>
+void AclopGatherNdKernel(const Context &dev_ctx,
+                         const phi::DenseTensor &x,
+                         const phi::DenseTensor &index,
+                         phi::DenseTensor *out) {
+  dev_ctx.template Alloc<T>(out);
+  auto stream = dev_ctx.stream();
+
+  if (x.numel() == 0) return;
+
+  if (index.numel() == 0) {
+    int diff = out->dims().size() - x.dims().size();
+    if (diff == 0) {
+      TensorCopy(dev_ctx, x, false, out);
+    } else {
+      std::vector<int64_t> new_dims(diff, 1);
+      for (size_t i = 0; i < x.dims().size(); ++i) {
+        new_dims.emplace_back(x.dims()[i]);
+      }
+
+      phi::DenseTensor x_tmp(x);
+      x_tmp.Resize(phi::make_ddim(new_dims));
+
+      NpuOpRunner runner;
+      runner.SetType("BroadcastTo")
+          .AddInput(x_tmp)
+          .AddInput(dev_ctx, phi::vectorize(out->dims()))
+          .AddOutput(*out);
+      runner.Run(stream);
+    }
+    return;
+  }
+
+  const auto &index_type = index.dtype();
+  bool index_type_match =
+      index_type == phi::DataType::INT32 || index_type == phi::DataType::INT64;
+  PADDLE_ENFORCE_EQ(
+      index_type_match,
+      true,
+      phi::errors::InvalidArgument("Index holds the wrong type, it holds [%s],"
+                                   "but desires to be [%s] or [%s]",
+                                   index_type,
+                                   phi::DataType::INT32,
+                                   phi::DataType::INT64));
+
+  const auto &runner = NpuOpRunner("GatherNd", {x, index}, {*out}, {});
+  runner.Run(stream);
+}
+
 template <typename T, typename Context>
 void GatherNdKernel(const Context &dev_ctx,
                     const phi::DenseTensor &x,
                     const phi::DenseTensor &index,
                     phi::DenseTensor *out) {
+  DO_COMPATIBILITY(
+      aclnnGatherNd,
+      (custom_kernel::AclopGatherNdKernel<T, Context>(dev_ctx, x, index, out)));
+
   dev_ctx.template Alloc<T>(out);
   auto stream = dev_ctx.stream();
 
@@ -62,8 +115,8 @@ void GatherNdKernel(const Context &dev_ctx,
                                    phi::DataType::INT32,
                                    phi::DataType::INT64));
 
-  const auto &runner = NpuOpRunner("GatherNd", {x, index}, {*out}, {});
-  runner.Run(stream);
+  bool negativeIndexSupport = false;
+  EXEC_NPU_CMD(aclnnGatherNd, dev_ctx, x, index, negativeIndexSupport, *out);
 }
 
 template <typename T, typename Context>
@@ -134,6 +187,7 @@ PD_REGISTER_PLUGIN_KERNEL(gather_nd,
                           ALL_LAYOUT,
                           custom_kernel::GatherNdKernel,
                           int64_t,
+                          int32_t,
                           float,
                           phi::dtype::float16,
                           phi::dtype::bfloat16) {}
@@ -143,6 +197,7 @@ PD_REGISTER_PLUGIN_KERNEL(gather_nd_grad,
                           ALL_LAYOUT,
                           custom_kernel::GatherNdGradKernel,
                           int64_t,
+                          int32_t,
                           float,
                           phi::dtype::float16,
                           phi::dtype::bfloat16) {}

backends/npu/kernels/index_sample_kernel.cc

@@ -23,6 +23,12 @@ void ConcatKernel(const Context& dev_ctx,
                   const phi::Scalar& axis_scalar,
                   phi::DenseTensor* out);
 
+template <typename T, typename Context>
+void GatherNdKernel(const Context& dev_ctx,
+                    const phi::DenseTensor& x,
+                    const phi::DenseTensor& index,
+                    phi::DenseTensor* out);
+
 template <typename T, typename Context>
 void IndexSampleGather(const Context& dev_ctx,
                        const phi::DenseTensor* index,
@@ -150,12 +156,20 @@ void IndexSampleGather(const Context& dev_ctx,
     TensorFromVector(dev_ctx, gather_index_vec, dev_ctx, &gather_index);
     gather_index.Resize({batch_size, index_length, 2});
 
-    NpuOpRunner runner;
-    runner.SetType("GatherNd")
-        .AddInput(*input)
-        .AddInput(gather_index)
-        .AddOutput(*out);
-    runner.Run(dev_ctx.stream());
+    auto dtype = input->dtype();
+    if (dtype == phi::DataType::FLOAT32) {
+      custom_kernel::GatherNdKernel<float, Context>(
+          dev_ctx, *input, gather_index, out);
+    } else if (dtype == phi::DataType::INT32) {
+      custom_kernel::GatherNdKernel<int32_t, Context>(
+          dev_ctx, *input, gather_index, out);
+    } else if (dtype == phi::DataType::INT64) {
+      custom_kernel::GatherNdKernel<int64_t, Context>(
+          dev_ctx, *input, gather_index, out);
+    } else if (dtype == phi::DataType::FLOAT16) {
+      custom_kernel::GatherNdKernel<phi::dtype::float16, Context>(
+          dev_ctx, *input, gather_index, out);
+    }
   }
 }
 

backends/npu/kernels/scatter_nd_add_kernel.cc

@@ -17,6 +17,12 @@
 
 namespace custom_kernel {
 
+template <typename T, typename Context>
+void GatherNdKernel(const Context& dev_ctx,
+                    const phi::DenseTensor& x,
+                    const phi::DenseTensor& index,
+                    phi::DenseTensor* out);
+
 template <typename T, typename Context>
 void ScatterNdAddKernel(const Context& dev_ctx,
                         const phi::DenseTensor& x,
@@ -90,9 +96,9 @@ void ScatterNdAddGradKernel(const Context& dev_ctx,
   }
   if (updates_grad) {
     dev_ctx.template Alloc<T>(updates_grad);
-    const auto& runner =
-        NpuOpRunner("GatherNd", {out_grad, index}, {*updates_grad}, {});
-    runner.Run(dev_ctx.stream());
+
+    custom_kernel::GatherNdKernel<T, Context>(
+        dev_ctx, out_grad, index, updates_grad);
   }
 }