【Infer Symbolic Shape No.217】Add infer_symbol_shape for rank_attention (#68967)

AndPuQing · web-flow · commit 4aac772215aa · 2024-11-04T14:41:22.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -217,7 +217,6 @@
     'push_gpups_sparse',
     'push_gpups_sparse_',
     'random_routing',
-    'rank_attention',
     'rnn_',
     'row_conv',
     'seed',
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -3491,12 +3491,55 @@ bool PyramidHashOpInferSymbolicShape(
 //   return QuantizeLinearOpInferSymbolicShape(op, infer_context);
 // }
 
-// bool RankAttentionOpInferSymbolicShape(pir::Operation *op,
-//                                        pir::InferSymbolicShapeContext
-//                                        *infer_context) {
-//   // pass
-//   return true;
-// }
+bool RankAttentionOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &rank_offset_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &rank_param_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &offset_shape =
+      rank_offset_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &param_shape =
+      rank_param_shape_or_data.shape();
+
+  int max_rank = op->attribute<pir::Int32Attribute>("max_rank").data();
+  infer_context->AddEqualCstr((offset_shape[1] - 1) / 2,
+                              symbol::DimExpr(max_rank));
+
+  std::vector<symbol::DimExpr> out_shape = {x_shape[0], param_shape[1]};
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+
+  if (details::IsFakeValue(op->result(0))) {
+    infer_context->SetSymbolForValueByStaticShape(op->result(0));
+  } else {
+    std::vector<symbol::DimExpr> x_help_shape = {x_shape[0],
+                                                 x_shape[1] * max_rank};
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(x_help_shape)});
+  }
+
+  if (details::IsFakeValue(op->result(2))) {
+    infer_context->SetSymbolForValueByStaticShape(op->result(2));
+  } else {
+    std::vector<symbol::DimExpr> ins_rank_shape = {x_shape[0], 1};
+    infer_context->SetShapeOrDataForValue(
+        op->result(2),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(ins_rank_shape)});
+  }
+
+  return true;
+}
+
 bool RandomRoutingOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &top_value_shape =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -105,7 +105,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(PsroiPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(PyramidHash)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(QuantizeLinear)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(QuantizeLinear_)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(RankAttention)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(RankAttention)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RandomRouting)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RandomRouting_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RmsNorm)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3980,7 +3980,7 @@
     data_type : x
   backward : rank_attention_grad
   optional : ins_rank, input_help
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : read_file
   args : (str filename = "", DataType dtype=DataType::UINT8, Place place=CPUPlace())
diff --git a/test/legacy_test/test_rank_attention_op.py b/test/legacy_test/test_rank_attention_op.py
@@ -200,7 +200,7 @@ def setUp(self):
 
     def test_check_output_gpu(self):
         if core.is_compiled_with_cuda():
-            self.check_output_with_place(core.CUDAPlace(0))
+            self.check_output_with_place(core.CUDAPlace(0), check_pir=True)
 
     def test_check_grad_gpu(self):
         if core.is_compiled_with_cuda():
