[PHI] Support int64 for paddle.kthvalue param k

paddle/fluid/framework/infershape_utils.cc

@@ -771,7 +771,7 @@ CompatInferMetaContext BuildInferMetaContext(InferShapeContext* ctx,
  }
  infer_meta_context.EmplaceBackAttr(std::move(scalar_list));
  } break;
- case proto::AttrType::SCALARS: {
+ case framework::proto::AttrType::SCALARS: {
  const auto& vec = PADDLE_GET_CONST(
  std::vector<paddle::experimental::Scalar>, attr);
  std::vector<phi::Scalar> scalar_list{vec.begin(), vec.end()};
@@ -805,8 +805,21 @@ CompatInferMetaContext BuildInferMetaContext(InferShapeContext* ctx,
  infer_meta_context.EmplaceBackAttr(PADDLE_GET_CONST(bool, attr));
  break;
  case phi::AttributeType::INT64:
- infer_meta_context.EmplaceBackAttr(
- PADDLE_GET_CONST(int64_t, attr));
+ switch (AttrTypeID(attr)) {
+ case framework::proto::AttrType::LONG:
+ infer_meta_context.EmplaceBackAttr(
+ PADDLE_GET_CONST(int64_t, attr));
+ break;
+ case framework::proto::AttrType::INT: {
+ const auto val = PADDLE_GET_CONST(int, attr);
+ infer_meta_context.EmplaceBackAttr(static_cast<int64_t>(val));
+ } break;
+ default:
+ PADDLE_THROW(common::errors::Unimplemented(
+ "Unsupported cast op attribute `%s` to int64_t when "
+ "construct InferMetaContext.",
+ attr_names[i]));
+ }
  break;
  case phi::AttributeType::INT32S:
  infer_meta_context.EmplaceBackAttr(

paddle/fluid/framework/operator.cc

@@ -3504,8 +3504,21 @@ void OperatorWithKernel::BuildPhiKernelContext(
  PADDLE_GET_CONST(bool, attr_iter->second));
  break;
  case phi::AttributeType::INT64:
- phi_kernel_context->EmplaceBackAttr(
- PADDLE_GET_CONST(int64_t, attr_iter->second));
+ switch (AttrTypeID(attr_iter->second)) {
+ case proto::AttrType::LONG:
+ phi_kernel_context->EmplaceBackAttr(
+ PADDLE_GET_CONST(int64_t, attr_iter->second));
+ break;
+ case proto::AttrType::INT: {
+ const auto val = PADDLE_GET_CONST(int, attr_iter->second);
+ phi_kernel_context->EmplaceBackAttr(static_cast<int64_t>(val));
+ } break;
+ default:
+ PADDLE_THROW(common::errors::Unimplemented(
+ "Unsupported cast op attribute `%s` to int64_t when "
+ "construct KernelContext.",
+ attr_names[i]));
+ }
  break;
  case phi::AttributeType::INT32S: // NOLINT
  phi_kernel_context->EmplaceBackAttr(

paddle/fluid/pir/serialize_deserialize/patch/0.yaml

@@ -0,0 +1,6 @@
+op_patches:
+ - op_name : pd_op.kthvalue
+ actions:
+ - action : modify_attr
+ object : k
+ type : pir::Int64Attribute

paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h

@@ -3154,7 +3154,7 @@ template <typename T>
 void kthvalue_grad(const Tensor& x,
  const Tensor& indices,
  const Tensor& out_grad,
- int k,
+ int64_t k UNUSED,
  int axis,
  bool keepdim,
  Tensor* x_grad) {

paddle/phi/infermeta/unary.cc

@@ -2486,7 +2486,7 @@ void IsfiniteInferMeta(const MetaTensor& x, MetaTensor* out) {
 }
 
 void KthvalueInferMeta(const MetaTensor& x,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  MetaTensor* out,
@@ -2523,7 +2523,7 @@ void KthvalueInferMeta(const MetaTensor& x,
  k,
  1,
  common::errors::InvalidArgument(
- "the k in the kthvalue must >= 1, but received %d .", k));
+ "the k in the kthvalue must >= 1, but received %lld .", k));
  PADDLE_ENFORCE_GE(input_dims.size(),
  0,
  common::errors::InvalidArgument(
@@ -2533,7 +2533,7 @@ void KthvalueInferMeta(const MetaTensor& x,
  input_dims[axis],
  k,
  common::errors::InvalidArgument(
- "input of kthvalue must have >= %d columns in axis of %d",
+ "input of kthvalue must have >= %lld columns in axis of %d",
  k,
  axis));
  }

paddle/phi/infermeta/unary.h

@@ -403,7 +403,7 @@ void IsEmptyInferMeta(const MetaTensor& x, MetaTensor* out);
 void IsfiniteInferMeta(const MetaTensor& input, MetaTensor* out);
 
 void KthvalueInferMeta(const MetaTensor& x,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  MetaTensor* out,

paddle/phi/kernels/cpu/kthvalue_grad_kernel.cc

@@ -49,7 +49,7 @@ void KthvalueGradKernel(const Context& dev_ctx,
  const DenseTensor& x,
  const DenseTensor& indices,
  const DenseTensor& d_out,
- int k UNUSED,
+ int64_t k UNUSED,
  int axis,
  bool keepdim,
  DenseTensor* d_x) {

paddle/phi/kernels/cpu/kthvalue_kernel.cc

@@ -27,7 +27,7 @@ static void getKthvalue(Type input_height,
  const DenseTensor* input,
  T* t_out,
  Type* t_indices,
- const int& k) {
+ const int64_t& k) {
  bool partial_sort_flag = (k * 64) < input_width;
 #ifdef PADDLE_WITH_MKLML
 #pragma omp parallel for
@@ -75,7 +75,7 @@ static void getKthvalue(Type input_height,
 template <typename T, typename Context>
 void KthvalueKernel(const Context& dev_ctx,
  const DenseTensor& x,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  DenseTensor* output,
@@ -98,7 +98,7 @@ void KthvalueKernel(const Context& dev_ctx,
  1,
  common::errors::InvalidArgument(
  "the k in the kthvalue must less equal than the "
- "elements number of the input X, but received %d .",
+ "elements number of the input X, but received %lld .",
  k));
 
  phi::Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, output);

paddle/phi/kernels/gpu/kthvalue_grad_kernel.cu

@@ -37,7 +37,7 @@ void KthvalueGradKernel(const Context& dev_ctx,
  const DenseTensor& x,
  const DenseTensor& indices,
  const DenseTensor& d_out,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  DenseTensor* d_x) {

paddle/phi/kernels/gpu/kthvalue_kernel.cu

@@ -158,14 +158,11 @@ bool SortKthvalue(const phi::GPUContext& dev_ctx,
 template <typename T, typename Context>
 void KthvalueKernel(const Context& dev_ctx,
  const DenseTensor& x,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  DenseTensor* output,
  DenseTensor* indices) {
- // TODO(cangtianhuang): support int64_t k
- k = static_cast<int64_t>(k);
-
  if (x.numel() == 0) {
  phi::Full<T, Context>(
  dev_ctx, phi::IntArray(common::vectorize(output->dims())), NAN, output);
@@ -186,7 +183,7 @@ void KthvalueKernel(const Context& dev_ctx,
  1,
  common::errors::InvalidArgument(
  "the k in the kthvalue must less equal than the "
- "elements number of the input X, but received %d .",
+ "elements number of the input X, but received %lld .",
  k));
 
  phi::Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, output);

paddle/phi/kernels/kthvalue_grad_kernel.h

@@ -23,7 +23,7 @@ void KthvalueGradKernel(const Context& dev_ctx,
  const DenseTensor& x,
  const DenseTensor& indices,
  const DenseTensor& d_out,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  DenseTensor* d_x);

paddle/phi/kernels/kthvalue_kernel.h

@@ -22,7 +22,7 @@ namespace phi {
 template <typename T, typename Context>
 void KthvalueKernel(const Context& dev_ctx,
  const DenseTensor& x,
- int k,
+ int64_t k,
  int axis,
  bool keepdim,
  DenseTensor* out,

paddle/phi/ops/yaml/backward.yaml

@@ -1849,8 +1849,8 @@
  data_type : out_grad
 
 - backward_op : kthvalue_grad
- forward : kthvalue(Tensor x, int k, int axis, bool keepdim) -> Tensor(out), Tensor(indices)
- args : (Tensor x, Tensor indices, Tensor out_grad, int k, int axis, bool keepdim)
+ forward : kthvalue(Tensor x, int64_t k, int axis, bool keepdim) -> Tensor(out), Tensor(indices)
+ args : (Tensor x, Tensor indices, Tensor out_grad, int64_t k, int axis, bool keepdim)
  output : Tensor(x_grad)
  infer_meta :
  func : UnchangedInferMeta

paddle/phi/ops/yaml/legacy/backward_exclude.yaml

@@ -23,6 +23,7 @@
 - fused_softmax_mask_grad
 - fused_softmax_mask_upper_triangle_grad
 - hsigmoid_loss_grad
+- kthvalue_grad
 - lp_pool2d_grad
 - max_grad
 - mean_double_grad

paddle/phi/ops/yaml/legacy/ops_exclude.yaml

@@ -47,6 +47,7 @@
 - gaussian
 - hsigmoid_loss
 - increment
+- kthvalue
 - linspace
 - logspace
 - lp_pool2d

paddle/phi/ops/yaml/legacy/static_backward.yaml

@@ -179,6 +179,17 @@
  param : [x, out_grad]
  inplace : (out_grad -> x_grad)
 
+- backward_op : kthvalue_grad
+ forward : kthvalue(Tensor x, int k, int axis, bool keepdim) -> Tensor(out), Tensor(indices)
+ args : (Tensor x, Tensor indices, Tensor out_grad, int k, int axis, bool keepdim)
+ output : Tensor(x_grad)
+ infer_meta :
+ func : UnchangedInferMeta
+ param: [x]
+ kernel :
+ func : kthvalue_grad
+ data_type : out_grad
+
 - backward_op : legacy_bilinear_interp_grad
  forward : legacy_bilinear_interp (Tensor x, Tensor out_size, Tensor[] size_tensor, Tensor scale_tensor, str data_format="NCHW", int out_d=0, int out_h=0, int out_w=0, float scale=0.0, str interp_method="bilinear", bool align_corners=true, int align_mode=1) -> Tensor(output)
  args : (Tensor x, Tensor out_size, Tensor[] size_tensor, Tensor scale_tensor, Tensor output_grad, str data_format, int out_d, int out_h, int out_w, float scale, str interp_method, bool align_corners, int align_mode)

paddle/phi/ops/yaml/legacy/static_ops.yaml

@@ -403,6 +403,16 @@
  data_type: x
  traits : paddle::dialect::ForwardOnlyTrait
 
+- op : kthvalue
+ args : (Tensor x, int k = 1, int axis = -1, bool keepdim = false)
+ output : Tensor(out), Tensor(indices)
+ infer_meta :
+ func : KthvalueInferMeta
+ kernel :
+ func : kthvalue
+ backward : kthvalue_grad
+ interfaces : paddle::dialect::InferSymbolicShapeInterface, paddle::dialect::LayoutTransformationInterface
+
 - op : legacy_bilinear_interp
  args : (Tensor x, Tensor out_size, Tensor[] size_tensor, Tensor scale_tensor, str data_format="NCHW", int out_d=0, int out_h=0, int out_w=0, float scale=0.0, str interp_method="bilinear", bool align_corners=true, int align_mode=1)
  output : Tensor(output)

paddle/phi/ops/yaml/ops.yaml

@@ -2976,7 +2976,7 @@
  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : kthvalue
- args : (Tensor x, int k = 1, int axis = -1, bool keepdim = false)
+ args : (Tensor x, int64_t k = 1, int axis = -1, bool keepdim = false)
  output : Tensor(out), Tensor(indices)
  infer_meta :
  func : KthvalueInferMeta