Fix paddle.dist & paddle.nn.functional.normalize

paddle/phi/kernels/funcs/reduce_grad_functions.h

@@ -38,10 +38,10 @@ void ReduceGradFunctor(const Context& dev_ctx,
  auto x_dims = input0.dims();
  auto reduced_dims_v = common::vectorize(x_dims);
  std::vector<int> dims_ref = dims;
- Eigen::array<int, D> broadcast_dim;
+ Eigen::array<int64_t, D> broadcast_dim;
  for (size_t i = 0; i < D; ++i) broadcast_dim[i] = 1;
 
- int broad_cast_times = 1;
+ int64_t broad_cast_times = 1;
  for (size_t i = 0; i < dims_ref.size(); ++i) {
  if (dims_ref[i] < 0) {
  dims_ref[i] = x_rank + dims_ref[i];
@@ -142,7 +142,7 @@ void LaunchReduceGradKernel(const Context& dev_ctx,
  auto& place = *dev_ctx.eigen_device();
  // *dev_ctx.eigen_device();
  auto broadcast_dim =
- Eigen::array<int, 1>({{static_cast<int>(input0->numel())}});
+ Eigen::array<int64_t, 1>({{static_cast<int64_t>(input0->numel())}});
  functor(place,
  &x,
  &x_reduce,

paddle/phi/kernels/gpu/dist_kernel.cu

@@ -63,6 +63,18 @@ struct PowFunctor {
  Ty p_order_;
 };
 
+template <typename Tx,
+ typename Ty,
+ typename Tout> // Tx is high precision, Tout is low/out precision
+struct PowFunctorHighPrecision {
+ HOSTDEVICE explicit inline PowFunctorHighPrecision(const Ty& p_order)
+ : p_order_(p_order) {}
+ HOSTDEVICE inline Tx operator()(const Tx x) const {
+ return static_cast<Tout>(pow(static_cast<Ty>(x), p_order_));
+ }
+ Ty p_order_;
+};
+
 template <typename T, typename Functor>
 __global__ void ReduceSumWithSubtract(
  const T* x, const T* y, T* out, int64_t N, Functor func) {
@@ -126,16 +138,17 @@ void DistKernel(const Context& dev_ctx,
  DenseTensor intermediate;
  const T* x_ptr = x.data<T>();
  const T* y_ptr = y.data<T>();
+
  T* o_ptr = dev_ctx.template Alloc<T>(out);
  auto stream = dev_ctx.stream();
 
  auto xdim = x.dims();
  if (xdim == y.dims()) { // same shape
- auto n = x.numel();
+ int64_t n = x.numel();
+
  auto config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, n);
  intermediate.Resize(common::make_ddim({config.block_per_grid.x}));
  T* i_ptr = dev_ctx.template Alloc<T>(&intermediate);
-
  std::vector<int64_t> axis_dims = {static_cast<int64_t>(-1)};
  std::vector<int> reduce_axis =
  funcs::details::GetReduceDim(axis_dims, xdim.size(), true);
@@ -166,15 +179,23 @@ void DistKernel(const Context& dev_ctx,
  ReduceSumWithSubtract<T>
  <<<config.block_per_grid.x, config.thread_per_block.x, 0, stream>>>(
  x_ptr, y_ptr, i_ptr, n, OtherOrderFunctor<T, MT>(p_order));
- phi::funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<MT>>(
- dev_ctx, intermediate, out, kps::IdentityFunctor<MT>(), reduce_axis);
-
- const DenseTensor* tmp_norm = out;
- std::vector<const DenseTensor*> ins = {tmp_norm};
+ DenseTensor out_other;
+ out_other.Resize(out->dims());
+ dev_ctx.template Alloc<MT>(&out_other);
+
+ phi::funcs::
+ ReduceKernel<T, MT, kps::AddFunctor, kps::IdentityFunctor<MT>>(
+ dev_ctx,
+ intermediate,
+ &out_other,
+ kps::IdentityFunctor<MT>(),
+ reduce_axis);
+ std::vector<const DenseTensor*> ins = {&out_other};
  std::vector<DenseTensor*> outs = {out};
- MT p_order_ = static_cast<MT>(static_cast<MT>(1.) / p_order);
+
+ MT p_order_ = static_cast<MT>(1.f / p_order);
  phi::funcs::ElementwiseKernel<T>(
- dev_ctx, ins, &outs, PowFunctor<T, MT>(p_order_));
+ dev_ctx, ins, &outs, PowFunctorHighPrecision<MT, MT, T>(p_order_));
  }
 
  } else {

paddle/phi/kernels/gpu/p_norm_grad_kernel.cu

@@ -42,10 +42,12 @@ struct AbsMaxAndMinGradFunctor {
 
 template <typename T>
 struct PNormGradFunctor {
+ using MT = typename phi::dtype::MPTypeTrait<T>::Type;
  HOSTDEVICE explicit inline PNormGradFunctor(float porder, float eps) {
- this->porder = static_cast<T>(porder - 1.);
- this->eps = static_cast<T>(eps);
+ this->porder = static_cast<MT>(porder - 1.);
+ this->eps = static_cast<MT>(eps);
  }
+
  template <typename Context,
  typename X,
  typename Y,
@@ -59,12 +61,33 @@ struct PNormGradFunctor {
  DY* dy,
  const Dim& dim,
  int size) {
+ auto x_mt = x->template cast<MT>();
+ auto y_mt = y->template cast<MT>();
+ auto dy_mt = dy->template cast<MT>();
+
+ auto norm_pow = y_mt.pow(-this->porder);
+ auto mask_norm_nonzero = (y_mt != static_cast<MT>(0)).template cast<MT>();
+
+ // Set to 0 where porder < 0 and x == 0
+ MT zero = static_cast<MT>(0);
+ auto mask_x_zero = (x_mt == zero).template cast<MT>();
+
+ MT is_porder_negative =
+ this->porder < zero ? static_cast<MT>(1) : static_cast<MT>(0);
+ auto invalid_mask = (mask_x_zero * is_porder_negative);
+ auto safe_pow =
+ x_mt.abs().pow(this->porder) * (static_cast<MT>(1) - invalid_mask);
+
  dx->device(place) =
- (*x).abs().pow(this->porder) * (*x).sign() * dy->broadcast(dim) *
- (*y + y->constant(eps)).pow(-this->porder).broadcast(dim);
+ (safe_pow * x_mt.sign() * dy_mt.broadcast(dim) *
+ norm_pow.broadcast(dim) *
+ mask_norm_nonzero.broadcast(dim) // Mask out positions where norm == 0
+ )
+ .template cast<T>();
  }
- T porder;
- T eps;
+
+ MT porder;
+ MT eps;
 };
 
 template <typename T, typename Context>

paddle/phi/kernels/gpu/p_norm_kernel.cu

@@ -124,31 +124,38 @@ void PNormKernel(const Context& dev_ctx,
  phi::funcs::ElementwiseKernel<T>(
  dev_ctx, ins, &outs, UnsignedPowFunctor<T>(1. / porder));
 #else
+ DenseTensor out_temp;
+ out_temp.Resize(out_norm->dims());
+ dev_ctx.template Alloc<MT>(&out_temp);
+
  if (porder == 1.0) {
  // fast 1-norm
  phi::funcs::ReduceKernel<T, T, kps::AddFunctor, FabsFunctor<T>>(
  dev_ctx, *in_x, out_norm, FabsFunctor<T>(), reduce_axis);
  } else if (porder == 2.0) {
  // fast 2-norm
- phi::funcs::ReduceKernel<T, T, kps::AddFunctor, SquareFunctor<T>>(
- dev_ctx, *in_x, out_norm, SquareFunctor<T>(), reduce_axis);
+ phi::funcs::ReduceKernel<T, MT, kps::AddFunctor, SquareFunctor<MT>>(
+ dev_ctx, *in_x, &out_temp, SquareFunctor<MT>(), reduce_axis);
  } else if (porder == 3.0) {
  // fast 3-norm
- phi::funcs::ReduceKernel<T, T, kps::AddFunctor, FabsCubicFunctor<T>>(
- dev_ctx, *in_x, out_norm, FabsCubicFunctor<T>(), reduce_axis);
+ phi::funcs::ReduceKernel<T, MT, kps::AddFunctor, FabsCubicFunctor<MT>>(
+ dev_ctx, *in_x, &out_temp, FabsCubicFunctor<MT>(), reduce_axis);
  } else {
  // vanilla norm
- phi::funcs::ReduceKernel<T, T, kps::AddFunctor, UnsignedPowFunctor<T>>(
- dev_ctx, *in_x, out_norm, UnsignedPowFunctor<T>(porder), reduce_axis);
+ phi::funcs::ReduceKernel<T, MT, kps::AddFunctor, UnsignedPowFunctor<MT>>(
+ dev_ctx,
+ *in_x,
+ &out_temp,
+ UnsignedPowFunctor<MT>(porder),
+ reduce_axis);
  }
 
  if (porder != 1.0) {
- // save computation when porder is 1.0
- const DenseTensor* tmp_norm = out_norm;
- std::vector<const DenseTensor*> ins = {tmp_norm};
+ std::vector<const DenseTensor*> ins = {&out_temp};
  std::vector<DenseTensor*> outs = {out_norm};
+ MT p_order_ = static_cast<MT>(1.f / porder);
  phi::funcs::ElementwiseKernel<T>(
- dev_ctx, ins, &outs, UnsignedPowFunctor<T>(1. / porder));
+ dev_ctx, ins, &outs, UnsignedPowFunctor<MT>(p_order_));
  }
 #endif
  }