Fix

Author: co63oc

paddle/phi/kernels/cpu/send_u_recv_grad_kernel.cc

@@ -19,6 +19,7 @@
 
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cpu/graph_send_recv_funcs.h"
+#include "paddle/phi/kernels/full_kernel.h"
 
 namespace phi {
 
@@ -131,6 +132,11 @@ void SendURecvGradKernel(const Context& dev_ctx,
  dev_ctx.template Alloc<T>(x_grad);
  return;
  }
+ if (src_index.numel() == 0) {
+ phi::Full<T, Context>(
+ dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+ return;
+ }
  auto index_type = src_index.dtype();
  if (index_type == phi::DataType::INT32) {
  GraphSendRecvGradOpKernelLaunchHelper<Context, T, int32_t>(

paddle/phi/kernels/cpu/send_u_recv_kernel.cc

@@ -112,7 +112,11 @@ void GraphSendRecvOpKernelLaunchHelper(const Context& dev_ctx,
  T* p_output = out->data<T>();
  const size_t& memset_bytes = memset_size * sizeof(T);
  memset(p_output, 0, memset_bytes);
- if (x.numel() == 0) {
+ if (x.numel() == 0 || src_index.numel() == 0) {
+ if (out->numel() != 0) {
+ phi::Full<T, Context>(
+ dev_ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+ }
  if (dst_count) {
  int64_t input_size = out_size <= 0 ? src_dims[0] : out_size;
  // dst_count shape [-1], need to Resize

paddle/phi/kernels/gpu/send_u_recv_grad_kernel.cu

@@ -20,6 +20,7 @@
 #include "paddle/common/hostdevice.h"
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/gpu/graph_send_recv_funcs.h"
 
 namespace phi {
@@ -108,6 +109,11 @@ void SendURecvGradKernel(const Context& dev_ctx,
  dev_ctx.template Alloc<T>(x_grad);
  return;
  }
+ if (src_index.numel() == 0) {
+ phi::Full<T, Context>(
+ dev_ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+ return;
+ }
  auto index_type = src_index.dtype();
  if (index_type == phi::DataType::INT32) {
  GraphSendRecvGradOpCUDAKernelLaunchHelper<Context, T, int32_t>(

paddle/phi/kernels/gpu/send_u_recv_kernel.cu

@@ -70,7 +70,11 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& dev_ctx,
  } else if (reduce_op == "MIN") {
  constant_functor(dev_ctx, out, std::numeric_limits<T>::max());
  }
- if (x.numel() == 0) {
+ if (x.numel() == 0 || src_index.numel() == 0) {
+ if (out->numel() != 0) {
+ phi::Full<T, Context>(
+ dev_ctx, phi::IntArray(common::vectorize(out->dims())), 0, out);
+ }
  if (dst_count) {
  int64_t input_size = out_size <= 0 ? src_dims[0] : out_size;
  // dst_count shape [-1], need to Resize