【PIR OpTest Fix No.34】 fix test_rank_attention_op

paddle/fluid/pir/dialect/op_generator/ops_api_gen.py

@@ -160,6 +160,7 @@
     'max_pool2d_v2',
     'partial_sum',
     'random_routing',
+    'rank_attention',
     'recv_v2',
     'rnn_',
     'row_conv',

paddle/fluid/pir/dialect/operator/ir/ops.yaml

@@ -1362,6 +1362,17 @@
     data_type : dtype
     backend : place
 
+- op : rank_attention
+  args : (Tensor x, Tensor rank_offset, Tensor rank_param, int max_rank = 3, int max_size = 0)
+  output : Tensor(input_help), Tensor(out), Tensor(ins_rank)
+  infer_meta :
+    func : RankAttentionInferMeta
+  kernel :
+    func : rank_attention
+    data_type : x
+  backward : rank_attention_grad
+  optional : ins_rank, input_help
+  W
 - op : read_file
   args : (str filename = "", DataType dtype=DataType::UINT8, Place place=CPUPlace())
   output : Tensor(out)

paddle/fluid/pir/dialect/operator/ir/ops_backward.yaml

@@ -657,6 +657,16 @@
     func : prod_grad
   composite: prod_grad(x, out, out_grad, dims, keep_dim, reduce_all, x_grad)
 
+- backward_op : rank_attention_grad
+  forward : rank_attention (Tensor x, Tensor rank_offset, Tensor rank_param, int max_rank = 3, int max_size = 0) -> Tensor(input_help), Tensor(out), Tensor(ins_rank)
+  args : (Tensor x, Tensor rank_offset, Tensor rank_param, Tensor input_help, Tensor ins_rank, Tensor out_grad, int max_rank = 3, int max_size = 0)
+  output : Tensor(rank_param_grad)
+  infer_meta :
+    func : RankAttentionGradInferMeta
+  kernel :
+    func : rank_attention_grad
+    data_type : out_grad
+
 - backward_op : repeat_interleave_grad
   forward : repeat_interleave(Tensor x, int repeats, int axis) -> Tensor(out)
   args : (Tensor x, Tensor out_grad, int repeats, int axis)

paddle/fluid/pir/dialect/operator/utils/utils.cc

@@ -69,6 +69,8 @@ const std::unordered_set<std::string> LegacyOpList = {
     SparseMomentumOp::name(),
     GetTensorFromSelectedRowsOp::name(),
     TdmSamplerOp::name(),
+    RankAttentionOp::name(),
+    RankAttentionGradOp::name(),
     RowConvOp::name(),
     RowConvGradOp::name(),
     SoftReluOp::name(),

paddle/phi/api/yaml/op_compat.yaml

@@ -3882,6 +3882,15 @@
   outputs:
     out : Out
 
+- op: rank_attention
+  backward: rank_attention_grad
+  inputs:
+    {x : X, rank_offset : RankOffset, rank_param : RankParam}
+  outputs:
+    {input_help : InputHelp, out : Out, ins_rank: InsRank}
+  attrs:
+    {max_rank : MaxRank, max_size : MaxSize}
+
 - op: read_from_array
   inputs:
     array : X

paddle/phi/infermeta/backward.cc

@@ -1044,6 +1044,19 @@ void PsroiPoolGradInferMeta(const MetaTensor& x,
   dx->share_meta(x);
 }
 
+void RankAttentionGradInferMeta(const MetaTensor& x,
+                                const MetaTensor& rank_offset,
+                                const MetaTensor& rank_param,
+                                const MetaTensor& input_help,
+                                const MetaTensor& ins_rank,
+                                const MetaTensor& out_grad,
+                                int max_rank,
+                                int max_size,
+                                MetaTensor* rank_param_grad) {
+  rank_param_grad->set_dims(rank_param.dims());
+  rank_param_grad->set_dtype(rank_param.dtype());
+}
+
 void RealAndImagGradInferMeta(const MetaTensor& out_grad, MetaTensor* dx) {
   dx->set_dims(out_grad.dims());
   dx->set_dtype(dtype::ToComplex(out_grad.dtype()));

paddle/phi/infermeta/backward.h

@@ -430,6 +430,16 @@ void PsroiPoolGradInferMeta(const MetaTensor& x,
                             float spatial_scale,
                             MetaTensor* dx);
 
+void RankAttentionGradInferMeta(const MetaTensor& x,
+                                const MetaTensor& rank_offset,
+                                const MetaTensor& rank_param,
+                                const MetaTensor& input_help,
+                                const MetaTensor& ins_rank,
+                                const MetaTensor& out_grad,
+                                int max_rank,
+                                int max_size,
+                                MetaTensor* rank_param_grad);
+
 void RealAndImagGradInferMeta(const MetaTensor& out_grad, MetaTensor* dx);
 
 void ReshapeDoubleGradInferMeta(const MetaTensor& out_grad,

paddle/phi/infermeta/ternary.cc

@@ -1134,6 +1134,45 @@ void RandomRoutingInferMeta(const MetaTensor& prob,
   out->share_lod(topk_idx);
 }
 
+void RankAttentionInferMeta(const MetaTensor& x,
+                            const MetaTensor& rank_offset,
+                            const MetaTensor& rank_param,
+                            int max_rank,
+                            int max_size,
+                            MetaTensor* input_help,
+                            MetaTensor* out,
+                            MetaTensor* ins_rank) {
+  auto x_dims = x.dims();
+  auto ins_num = x_dims[0];
+  auto param_dims = rank_param.dims();
+  auto para_col = param_dims[1];
+  auto rank_offset_dims = rank_offset.dims();
+  auto x_fea_dim = x_dims[1];
+  auto block_matrix_row = max_rank * x_fea_dim;
+
+  PADDLE_ENFORCE_EQ(
+      (rank_offset_dims[1] - 1) / 2,
+      max_rank,
+      phi::errors::InvalidArgument("Input(RankOffset) has wrong columns, "
+                                   "except columns to be %d, but got %d",
+                                   max_rank,
+                                   (rank_offset_dims[1] - 1) / 2));
+
+  std::vector<int64_t> out_dims({ins_num, para_col});
+  out->set_dims(common::make_ddim(out_dims));
+  out->set_dtype(x.dtype());
+
+  std::vector<int64_t> input_help_dims({ins_num, block_matrix_row});
+  input_help->set_dims(common::make_ddim(input_help_dims));
+  input_help->set_dtype(x.dtype());
+
+  std::vector<int64_t> ins_rank_dims({ins_num, 1});
+  ins_rank->set_dims(common::make_ddim(ins_rank_dims));
+  ins_rank->set_dtype(x.dtype());
+
+  out->share_lod(x);
+}
+
 void RoiAlignInferMeta(const MetaTensor& x,
                        const MetaTensor& boxes,
                        const MetaTensor& boxes_num,

paddle/phi/infermeta/ternary.h

@@ -210,6 +210,15 @@ void RandomRoutingInferMeta(const MetaTensor& prob,
                             const MetaTensor& topk_idx,
                             MetaTensor* out);
 
+void RankAttentionInferMeta(const MetaTensor& x,
+                            const MetaTensor& rank_offset,
+                            const MetaTensor& rank_param,
+                            int max_rank,
+                            int max_size,
+                            MetaTensor* input_help,
+                            MetaTensor* out,
+                            MetaTensor* ins_rank);
+
 void RoiAlignInferMeta(const MetaTensor& x,
                        const MetaTensor& boxes,
                        const MetaTensor& boxes_num,

test/white_list/pir_op_test_white_list

@@ -221,6 +221,7 @@ test_qr_op
 test_randint_op
 test_randperm_op
 test_range
+test_rank_attention_op
 test_reduce_op
 test_reduce_op_static_build
 test_repeat_interleave_op