[Prim][PIR] Add pd.rint op (#74012)

Author: huangjiyi

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc

@@ -139,6 +139,8 @@ OP_SAME_OPERANDS_AND_RESULT(Relu)
 OP_SAME_OPERANDS_AND_RESULT(Relu6)
 OP_SAME_OPERANDS_AND_RESULT(Relu_)
 OP_SAME_OPERANDS_AND_RESULT(Reverse)
+OP_SAME_OPERANDS_AND_RESULT(Rint)
+OP_SAME_OPERANDS_AND_RESULT(Rint_)
 OP_SAME_OPERANDS_AND_RESULT(Roll)
 OP_SAME_OPERANDS_AND_RESULT(Round)
 OP_SAME_OPERANDS_AND_RESULT(Round_)

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h

@@ -131,6 +131,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Relu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Relu6)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Relu_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reverse)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Rint)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Rint_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Roll)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Round)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Round_)

paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h

@@ -1525,8 +1525,14 @@ std::tuple<Tensor, Tensor, Tensor> rms_norm_decomp(
  auto quant_max_bound_scalar =
  full_scalar<T>(quant_max_bound, out.dtype(), out.place());
  auto scale_out = out * quant_scale_scalar * quant_max_bound_scalar;
+ Tensor round_out;
+ if (quant_round_type == 0) {
+ round_out = backend::rint<T>(scale_out);
+ } else {
+ round_out = round<T>(scale_out);
+ }
  auto clip_out = clip_decomp<T>(
- round<T>(scale_out), quant_min_bound_scalar, quant_max_bound_scalar);
+ round_out, quant_min_bound_scalar, quant_max_bound_scalar);
  if (fabs(quant_max_bound - 127.0f) < 0.000001) {
  out = cast<T>(clip_out, phi::DataType::INT8);
  } else if (fabs(quant_max_bound - 448.0f) < 0.000001) {

paddle/phi/kernels/activation_grad_kernel.h

@@ -302,6 +302,7 @@ DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid);
 DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Sqrt);
 DECLARE_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu6);
 
+DECLARE_ACTIVATION_GRAD_KERNEL_NODEP(Rint);
 DECLARE_ACTIVATION_GRAD_KERNEL_NODEP(Round);
 DECLARE_ACTIVATION_GRAD_KERNEL_NODEP(Floor);
 DECLARE_ACTIVATION_GRAD_KERNEL_NODEP(Ceil);

paddle/phi/kernels/cpu/activation_grad_kernel.cc

@@ -150,6 +150,7 @@ DEFINE_CPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu, ReluGradFunctor);
 DEFINE_CPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Tanh, TanhGradFunctor);
 DEFINE_CPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid, SigmoidGradFunctor);
 
+DEFINE_CPU_ACTIVATION_GRAD_KERNEL_NODEP(Rint, ZeroGradFunctor);
 DEFINE_CPU_ACTIVATION_GRAD_KERNEL_NODEP(Round, ZeroGradFunctor);
 DEFINE_CPU_ACTIVATION_GRAD_KERNEL_NODEP(Floor, ZeroGradFunctor);
 DEFINE_CPU_ACTIVATION_GRAD_KERNEL_NODEP(Ceil, ZeroGradFunctor);
@@ -488,6 +489,15 @@ PD_REGISTER_ACTIVATION_GRAD_KERNEL(celu_grad, CeluGradKernel)
 PD_REGISTER_ACTIVATION_DOUBLE_GRAD_KERNEL(celu_double_grad,
  CeluDoubleGradKernel)
 
+PD_REGISTER_KERNEL(rint_grad,
+ CPU,
+ ALL_LAYOUT,
+ phi::RintGradKernel,
+ float,
+ double,
+ int,
+ int64_t) {}
+
 PD_REGISTER_KERNEL(round_grad,
  CPU,
  ALL_LAYOUT,

paddle/phi/kernels/cpu/activation_kernel.cc

@@ -97,6 +97,7 @@ DEFINE_CPU_ACTIVATION_KERNEL(LogSigmoid, LogSigmoidFunctor)
 DEFINE_CPU_ACTIVATION_KERNEL(Floor, FloorFunctor)
 DEFINE_CPU_ACTIVATION_KERNEL(Ceil, CeilFunctor)
 DEFINE_CPU_ACTIVATION_KERNEL(Negative, NegativeFunctor)
+DEFINE_CPU_ACTIVATION_KERNEL(Rint, RintFunctor)
 
 DEFINE_CPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log, LogFunctor)
 DEFINE_CPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log2, Log2Functor)
@@ -257,6 +258,9 @@ PD_REGISTER_ACTIVATION_KERNEL(floor, FloorKernel)
 PD_REGISTER_ACTIVATION_KERNEL(ceil, CeilKernel)
 PD_REGISTER_ACTIVATION_KERNEL(celu, CeluKernel)
 
+PD_REGISTER_KERNEL(
+ rint, CPU, ALL_LAYOUT, phi::RintKernel, int, int64_t, float, double) {}
+
 PD_REGISTER_KERNEL(round,
  CPU,
  ALL_LAYOUT,

paddle/phi/kernels/funcs/activation_functor.h

@@ -3005,6 +3005,26 @@ struct FloorFunctor : public BaseActivationFunctor<T> {
  }
 };
 
+// rint(x) = [x]
+template <typename T, typename Enable = void>
+struct RintFunctor : public BaseActivationFunctor<T> {
+ template <typename Device, typename X, typename Out>
+ void operator()(Device d, X x, Out out) const {
+ out.device(d) = x.unaryExpr([](const T& val) {
+ return (std::isnan(val) || std::isinf(val)) ? val : std::rint(val);
+ });
+ }
+};
+
+template <typename T>
+struct RintFunctor<T, std::enable_if_t<std::is_integral_v<T>>>
+ : public BaseActivationFunctor<T> {
+ template <typename Device, typename X, typename Out>
+ void operator()(Device d, X x, Out out) const {
+ out.device(d) = x;
+ }
+};
+
 // round(x) = [x]
 template <typename T, typename Enable = void>
 struct RoundFunctor : public BaseActivationFunctor<T> {
@@ -5410,6 +5430,25 @@ struct CudaFloorFunctor : public BaseActivationFunctor<T> {
  }
 };
 
+template <typename T, typename Enable = void>
+struct CudaRintFunctor : public BaseActivationFunctor<T> {
+ using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
+
+ // rint(x) = rint(x)
+ __device__ __forceinline__ T operator()(const T arg_x) const {
+ MPType x = static_cast<MPType>(arg_x);
+ if (isnan(x) || isinf(x)) return arg_x;
+ return static_cast<T>(std::rint(x));
+ }
+};
+
+template <typename T>
+struct CudaRintFunctor<T, std::enable_if_t<std::is_integral_v<T>>>
+ : public BaseActivationFunctor<T> {
+ // rint(x) = x
+ __device__ __forceinline__ T operator()(const T arg_x) const { return arg_x; }
+};
+
 template <typename T, typename Enable = void>
 struct CudaRoundFunctor : public BaseActivationFunctor<T> {
  using MPType = typename phi::dtype::MPTypeTrait<T>::Type;

paddle/phi/kernels/gpu/activation_grad_kernel.cu

@@ -178,6 +178,7 @@ DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu, CudaReluGradFunctor);
 DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Tanh, CudaTanhGradFunctor);
 DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid, CudaSigmoidGradFunctor);
 
+DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Rint, CudaZeroGradFunctor);
 DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Round, CudaZeroGradFunctor);
 DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Floor, CudaZeroGradFunctor);
 DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Ceil, CudaZeroGradFunctor);
@@ -558,6 +559,16 @@ PD_REGISTER_ACTIVATION_GRAD_KERNEL(ceil_grad, CeilGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(celu_grad, CeluGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(celu_double_grad, CeluDoubleGradKernel)
 
+PD_REGISTER_KERNEL(rint_grad,
+ GPU,
+ ALL_LAYOUT,
+ phi::RintGradKernel,
+ int,
+ int64_t,
+ float,
+ double,
+ phi::dtype::float16,
+ phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(round_grad,
  GPU,
  ALL_LAYOUT,

paddle/phi/kernels/gpu/activation_kernel.cu

@@ -116,6 +116,7 @@ DEFINE_GPU_ACTIVATION_KERNEL(Sigmoid, CudaSigmoidFunctor)
 DEFINE_GPU_ACTIVATION_KERNEL(LogSigmoid, CudaLogSigmoidFunctor)
 DEFINE_GPU_ACTIVATION_KERNEL(Floor, CudaFloorFunctor)
 DEFINE_GPU_ACTIVATION_KERNEL(Ceil, CudaCeilFunctor)
+DEFINE_GPU_ACTIVATION_KERNEL(Rint, CudaRintFunctor)
 
 DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log, CudaLogFunctor)
 DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log2, CudaLog2Functor)
@@ -352,6 +353,16 @@ PD_REGISTER_ACTIVATION_KERNEL(celu, CeluKernel)
 PD_REGISTER_ACTIVATION_KERNEL(selu, SeluKernel)
 PD_REGISTER_ACTIVATION_KERNEL(logit, LogitCUDAKernel)
 
+PD_REGISTER_KERNEL(rint,
+ GPU,
+ ALL_LAYOUT,
+ phi::RintKernel,
+ int,
+ int64_t,
+ float,
+ double,
+ phi::dtype::float16,
+ phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(round,
  GPU,
  ALL_LAYOUT,

paddle/phi/ops/yaml/backward.yaml

@@ -2896,6 +2896,17 @@
  output : Tensor(x_grad)
  invoke : reverse(out_grad, axis)
 
+- backward_op : rint_grad
+ forward : rint(Tensor x) -> Tensor(out)
+ args : (Tensor out_grad)
+ output : Tensor(x_grad)
+ infer_meta :
+ func : UnchangedInferMeta
+ param: [out_grad]
+ kernel :
+ func : rint_grad
+ inplace : (out_grad -> x_grad)
+
 - backward_op : rms_norm_grad
  forward : rms_norm (Tensor x, Tensor bias, Tensor residual, Tensor norm_weight, Tensor norm_bias, float epsilon, int begin_norm_axis, float quant_scale, int quant_round_type, float quant_max_bound, float quant_min_bound) -> Tensor(out), Tensor(residual_out), Tensor(inv_var)
  args : (Tensor x, Tensor bias, Tensor residual, Tensor norm_weight, Tensor norm_bias, Tensor inv_var, Tensor out_grad, float epsilon, int begin_norm_axis, float quant_scale)