【Infer Symbolic Shape No.171,176-179】【BUAA】yolo_box yolo_box_head yolo_box_post updata_loss_scaling_ tdm_child (#67121)

tlxd · web-flow · commit ade576434a31 · 2024-08-14T11:22:38.000+08:00
* change op yolo_box yolo_box_head yolo_box_post updata_loss_scaling_ tdm_child

* update multiary_infer_sym.cc multiary_infer_sym.h binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc

* Update binary_infer_sym.cc
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
@@ -1140,12 +1140,27 @@ bool TopPSamplingOpInferSymbolicShape(
   return true;
 }
 
-// bool TdmChildOpInferSymbolicShape(pir::Operation *op,
-//                                   pir::InferSymbolicShapeContext
-//                                   *infer_context) {
-//   // pass
-//   return true;
-// }
+bool TdmChildOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &input_dims = x_shape_or_data.shape();
+  int child_nums = op->attribute<pir::Int32Attribute>("child_nums").data();
+
+  std::vector<symbol::DimExpr> output_dims = input_dims;
+  output_dims.push_back(symbol::DimExpr(child_nums));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_dims)});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_dims)});
+
+  return true;
+}
 
 // bool UnpoolOpInferSymbolicShape(pir::Operation *op,
 //                                 pir::InferSymbolicShapeContext
@@ -1154,19 +1169,36 @@ bool TopPSamplingOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool YoloBoxOpInferSymbolicShape(pir::Operation *op,
-//                                  pir::InferSymbolicShapeContext
-//                                  *infer_context) {
-//   // pass
-//   return true;
-// }
+bool YoloBoxOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto anchors =
+      paddle::dialect::details::GetVectorAttr<int>(op, "anchors");
+  int class_num = op->attribute<pir::Int32Attribute>("class_num").data();
 
-// bool YoloBoxHeadOpInferSymbolicShape(pir::Operation *op,
-//                                      pir::InferSymbolicShapeContext
-//                                      *infer_context) {
-//   // pass
-//   return true;
-// }
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  int anchor_num = static_cast<int>(anchors.size() / 2);
+
+  symbol::DimExpr box_num = symbol::DimExpr(0);
+  box_num = x_shape[2] * x_shape[3] * symbol::DimExpr(anchor_num);
+
+  std::vector<symbol::DimExpr> boxes_shape = {
+      x_shape[0], box_num, symbol::DimExpr(4)};
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(boxes_shape)});
+
+  std::vector<symbol::DimExpr> scores_shape = {
+      x_shape[0], box_num, symbol::DimExpr(class_num)};
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(scores_shape)});
+
+  return true;
+}
 
 }  // namespace paddle::dialect
 
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
@@ -67,10 +67,10 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(SegmentPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Swiglu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TakeAlongAxis)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TopPSampling)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmChild)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmChild)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unpool)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBox)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxHead)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBox)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxHead)
 
 }  // namespace paddle::dialect
 
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
@@ -1812,18 +1812,64 @@ bool FakeChannelWiseDequantizeMaxAbsOpInferSymbolicShape(
   return true;
 }
 
-// bool UpdateLossScaling_OpInferSymbolicShape(pir::Operation *op,
-//                                             pir::InferSymbolicShapeContext
-//                                             *infer_context) {
-//   // pass
-//   return true;
-// }
+bool UpdateLossScaling_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::TensorListShapeOrDataDimExprs &shape_data_list =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0))
+          .dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
 
-// bool YoloBoxPostOpInferSymbolicShape(pir::Operation *op,
-//                                      pir::InferSymbolicShapeContext
-//                                      *infer_context) {
-//   // pass
-//   return true;
-// }
+  const symbol::ShapeOrDataDimExprs sym_shape_dim_exprs = [&] {
+    symbol::TensorListShapeOrDataDimExprs shape_dim_exprs_list;
+
+    for (auto &shape_data : shape_data_list) {
+      auto shape_dim_exprs =
+          symbol::TensorShapeOrDataDimExprs(shape_data.shape());
+      shape_dim_exprs_list.emplace_back(shape_dim_exprs);
+    }
+
+    return symbol::ShapeOrDataDimExprs(shape_dim_exprs_list);
+  }();
+
+  infer_context->SetShapeOrDataForValue(op->result(0), sym_shape_dim_exprs);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)})});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(2),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)})});
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(3),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({symbol::DimExpr(1)})});
+
+  return true;
+}
+
+bool YoloBoxPostOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &image_shape_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(3));
+
+  const symbol::DimExpr &batch = image_shape_shape_or_data.shape()[0];
+
+  std::vector<symbol::DimExpr> dim_out = {symbol::DimExpr(1),
+                                          symbol::DimExpr(6)};
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(dim_out)});
+
+  std::vector<symbol::DimExpr> dim_nms_rois_num = {batch};
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(dim_nms_rois_num)});
+
+  return true;
+}
 
 }  // namespace paddle::dialect
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
@@ -90,7 +90,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(HsigmoidLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Where)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Where_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeChannelWiseDequantizeMaxAbs)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(UpdateLossScaling_)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxPost)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(UpdateLossScaling_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxPost)
 
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -196,6 +196,7 @@ OP_SAME_OPERANDS_AND_RESULT(MergeSelectedRows)
 OP_SAME_OPERANDS_AND_RESULT(Renorm)
 OP_SAME_OPERANDS_AND_RESULT(Renorm_)
 OP_SAME_OPERANDS_AND_RESULT(TanhShrink)
+OP_SAME_OPERANDS_AND_RESULT(YoloBoxHead)
 
 bool ScaleOpInferSymbolicShape(pir::Operation *op,
                                pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -189,6 +189,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergeSelectedRows)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Renorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Renorm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TanhShrink)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(YoloBoxHead)
 
 }  // namespace paddle::dialect
 
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4612,6 +4612,7 @@
   kernel:
     func: tdm_child
     data_type: x
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : tdm_sampler
   args: (Tensor x, Tensor travel, Tensor layer, bool output_positive=true, int[] neg_samples_num_list={}, int[] layer_offset_lod={}, int seed = 0, int dtype=2)
@@ -4933,6 +4934,7 @@
   data_transform :
     skip_transform : found_infinite
   inplace : (x -> out), (prev_loss_scaling -> loss_scaling), (in_good_steps -> out_good_steps), (in_bad_steps -> out_bad_steps)
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : view_dtype
   args : (Tensor input, DataType dtype)
@@ -5050,6 +5052,7 @@
   kernel :
     func : yolo_box
     data_type : x
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : yolo_box_head
   args : (Tensor x, int[] anchors, int class_num)
@@ -5059,6 +5062,7 @@
   kernel :
     func : yolo_box_head
     data_type : x
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : yolo_box_post
   args : (Tensor boxes0, Tensor boxes1, Tensor boxes2, Tensor image_shape, Tensor image_scale, int[] anchors0, int[] anchors1, int[] anchors2, int class_num, float conf_thresh, int downsample_ratio0, int downsample_ratio1, int downsample_ratio2, bool clip_bbox, float scale_x_y, float nms_threshold)
@@ -5068,6 +5072,7 @@
   kernel :
     func : yolo_box_post
     data_type : boxes0
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : yolo_loss
   args : (Tensor x, Tensor gt_box, Tensor gt_label, Tensor gt_score, int[] anchors={}, int[] anchor_mask={}, int class_num =1 , float ignore_thresh=0.7, int downsample_ratio=32, bool use_label_smooth=true, float scale_x_y=1.0)
