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 index ff1127deae03d..0fae5cb439bc7 100644 --- 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 @@ -1460,19 +1460,67 @@ bool RoiAlignOpInferSymbolicShape( // return MergedMomentumOpInferSymbolicShape(op, infer_context); // } -// bool MoeOpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext *infer_context) -// { -// // pass -// return true; -// } +bool MoeOpInferSymbolicShape(pir::Operation *op, + pir::InferSymbolicShapeContext *infer_context) { + const auto &x_shape = + infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape(); + infer_context->SetShapeOrDataForValue( + op->result(0), + symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)}); + return true; +} -// bool MulticlassNMS3OpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext -// *infer_context) { -// // pass -// return true; -// } +bool MulticlassNms3OpInferSymbolicShape( + pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { + const auto &bboxes_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(0)); + const auto &scores_shape_or_data = + infer_context->GetShapeOrDataForValue(op->operand_source(1)); + + const std::vector &box_dims = bboxes_shape_or_data.shape(); + const std::vector &score_dims = scores_shape_or_data.shape(); + const size_t score_size = score_dims.size(); + // TODO(Jeff114514): + // add constraint for score_size == 2 || score_size == 3 + // add constraint for box_dims.size() == 3 + if (score_size == 3) { + // TODO(Jeff114514): + // add constraint for box_dims[2] == 4 || box_dims[2] == 8 || box_dims[2] == + // 16 || box_dims[2] == 24 || box_dims[2] == 32 + infer_context->AddEqualCstr(box_dims[1], score_dims[2]); + } else { + infer_context->AddEqualCstr(box_dims[2], symbol::DimExpr(4)); + infer_context->AddEqualCstr(box_dims[1], score_dims[1]); + } + + const auto &next_symbol_out_and_index = infer_context->GetNextSymName(); + + std::vector out_shape; + out_shape.emplace_back(next_symbol_out_and_index); + out_shape.emplace_back(box_dims[2] + 2); + + std::vector index_shape; + index_shape.emplace_back(next_symbol_out_and_index); + index_shape.emplace_back(1); + + std::vector nms_rois_num_shape; + nms_rois_num_shape.emplace_back(infer_context->GetNextSymName()); + + infer_context->SetShapeOrDataForValue( + op->result(0), + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(out_shape)}); + infer_context->SetShapeOrDataForValue( + op->result(1), + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(index_shape)}); + infer_context->SetShapeOrDataForValue( + op->result(2), + symbol::ShapeOrDataDimExprs{ + symbol::TensorShapeOrDataDimExprs(nms_rois_num_shape)}); + + return true; +} bool MeshgridOpInferSymbolicShape( pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { @@ -1532,8 +1580,7 @@ bool MovingAverageAbsMaxScale_OpInferSymbolicShape( } // bool NceOpInferSymbolicShape(pir::Operation *op, -// pir::InferSymbolicShapeContext *infer_context) -// { +// pir::InferSymbolicShapeContext *infer_context){ // // pass // return true; // } 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 index dc88574ee155b..fc65d076bfba1 100644 --- 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 @@ -67,10 +67,10 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logspace) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedAdam_) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum) // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MergedMomentum_) -// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) -// OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNMS3) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(MulticlassNms3) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MemoryEfficientAttention) OP_DECLARE_INFER_SYMBOLIC_SHAPE(Meshgrid) +OP_DECLARE_INFER_SYMBOLIC_SHAPE(Moe) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MovingAverageAbsMaxScale) OP_DECLARE_INFER_SYMBOLIC_SHAPE(MovingAverageAbsMaxScale_) OP_DECLARE_INFER_SYMBOLIC_SHAPE(NearestInterp) diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml index d3c61cc36b8f9..87694d0bac812 100755 --- a/paddle/phi/ops/yaml/ops.yaml +++ b/paddle/phi/ops/yaml/ops.yaml @@ -3283,6 +3283,7 @@ func : multiclass_nms3 data_type : scores optional : rois_num, nms_rois_num + interfaces : paddle::dialect::InferSymbolicShapeInterface - op : multinomial args : (Tensor x, Scalar(int) num_samples = 1, bool replacement = false) @@ -5111,6 +5112,7 @@ func: MoeInferMeta kernel: func: moe + interfaces: paddle::dialect::InferSymbolicShapeInterface - op: number_count args: (Tensor numbers, int upper_range)