【Infer Symbolic Shape No.200】[BUAA] Addgather_tree (#67625)

Luohongzhige · web-flow · commit 2c7d753b56f6 · 2024-08-26T15:41:08.000+08:00
* gathertree

* fixed

* fixed

* fixed

* fixed
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -764,12 +764,34 @@ bool GatherNdOpInferSymbolicShape(
   return true;
 }
 
-// bool GatherTreeOpInferSymbolicShape(pir::Operation *op,
-//                                     pir::InferSymbolicShapeContext
-//                                     *infer_context) {
-//   // pass
-//   return true;
-// }
+bool GatherTreeOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &ids_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const symbol::ShapeOrDataDimExprs &parents_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  const std::vector<symbol::DimExpr> &ids_shape = ids_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &parents_shape =
+      parents_shape_or_data.shape();
+
+  PADDLE_ENFORCE_EQ(ids_shape.size() == parents_shape.size(),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The shape of Input(Parents) must be same with the "
+                        "shape of Input(Ids)."));
+  size_t rank = ids_shape.size();
+  for (size_t i = 0; i < rank; ++i) {
+    infer_context->AddEqualCstr(ids_shape[i], parents_shape[i]);
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(ids_shape)});
+
+  return true;
+}
 
 bool HuberLossOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -49,7 +49,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedSoftmaxMask)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(GridSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gather)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherNd)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherTree)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherTree)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(HuberLoss)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Histogram)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2225,7 +2225,7 @@
   kernel :
     func : gather_tree
     data_type : ids
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : gaussian
   args : (IntArray shape, float mean, float std, int seed, DataType dtype, Place place={})
