【Infer Symbolic Shape BUAA No.10 No.73 No.78】Add mode , margin_cross_entropy , box_clip. op

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc

@@ -91,6 +91,33 @@ bool AllcloseOpInferSymbolicShape(
   return true;
 }
 
+bool BoxClipOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const auto &im_info_shape =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1)).shape();
+
+  // Check rank and dimensions of input tensors
+  const auto &three = symbol::DimExpr{3};
+  const auto &four = symbol::DimExpr{4};
+  infer_context->AddEqualCstr(input_shape[input_shape.size() - 1], four);
+  PADDLE_ENFORCE_EQ(im_info_shape.size(),
+                    2,
+                    common::errors::InvalidArgument(
+                        "The rank of Input(im_info) in BoxClipOp must be 2. "
+                        "But received rank = %d",
+                        im_info_shape.size()));
+  infer_context->AddEqualCstr(im_info_shape[1], three);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(input_shape)});
+
+  return true;
+}
+
 bool Atan2OpInferSymbolicShape(pir::Operation *op,
                                pir::InferSymbolicShapeContext *infer_context) {
   const auto x_shape =
@@ -748,6 +775,46 @@ bool MatmulOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool MarginCrossEntropyOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &logits_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &labels_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  std::vector<symbol::DimExpr> logits_dims = logits_shape_or_data.shape();
+  std::vector<symbol::DimExpr> labels_dims = labels_shape_or_data.shape();
+
+  size_t logits_rank = logits_dims.size();
+  auto axis = logits_rank - 1;
+
+  for (size_t i = 0; i < logits_rank; i++) {
+    if (i != axis) {
+      infer_context->AddEqualCstr(logits_dims[i], labels_dims[i]);
+    }
+  }
+
+  const auto &one = symbol::DimExpr{1};
+
+  if (labels_dims.size() > 1) {
+    infer_context->AddEqualCstr(labels_dims[axis - 1], one);
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(logits_dims)});
+
+  logits_dims[axis] = symbol::DimExpr(1);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(logits_dims)});
+
+  return true;
+}
+
 bool MvOpInferSymbolicShape(pir::Operation *op,
                             pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h

@@ -22,6 +22,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Allclose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Atan2)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BceLoss_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(BoxClip)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Binomial_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bincount)
@@ -54,6 +55,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(MaskedSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Matmul)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(MarginCrossEntropy)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mv)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Nextafter)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(PullBoxSparse)

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc

@@ -1125,6 +1125,42 @@ bool MaxOpInferSymbolicShape(pir::Operation *op,
 
   return details::ReduceInferDim(op, infer_context, axis, keepdim, reduce_all);
 }
+bool ModeOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &x_shape = x_shape_or_data.shape();
+
+  int axis = op->attribute<pir::Int32Attribute>("axis").data();
+  bool keepdim = op->attribute<pir::BoolAttribute>("keepdim").data();
+
+  int dim_size = x_shape.size();
+
+  if (axis < 0) {
+    axis += dim_size;
+  }
+
+  std::vector<symbol::DimExpr> out_dims;
+  for (int i = 0; i < axis; i++) {
+    out_dims.emplace_back(x_shape[i]);
+  }
+  if (keepdim && dim_size > 0) {
+    out_dims.emplace_back(symbol::DimExpr(1));
+  }
+  for (int i = axis + 1; i < dim_size; i++) {
+    out_dims.emplace_back(x_shape[i]);
+  }
+
+  symbol::TensorShapeOrDataDimExprs out_shape(out_dims);
+
+  infer_context->SetShapeOrDataForValue(op->result(0),
+                                        symbol::ShapeOrDataDimExprs{out_shape});
+
+  infer_context->SetShapeOrDataForValue(op->result(1),
+                                        symbol::ShapeOrDataDimExprs{out_shape});
+
+  return true;
+}
 
 bool MaxoutOpInferSymbolicShape(pir::Operation *op,
                                 pir::InferSymbolicShapeContext *infer_context) {

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h

@@ -73,6 +73,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(L1Norm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LpPool2d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logcumsumexp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsumexp)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mode)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Max)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Maxout)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Min)

paddle/phi/ops/yaml/ops.yaml

@@ -652,6 +652,7 @@
     func: BoxClipInferMeta
   kernel:
     func: box_clip
+  interfaces: paddle::dialect::InferSymbolicShapeInterface
 
 - op : box_coder
   args : (Tensor prior_box, Tensor prior_box_var, Tensor target_box, str code_type = "encode_center_size", bool box_normalized = true, int axis = 0, float[] variance = {})
@@ -3012,6 +3013,7 @@
     func : margin_cross_entropy
     data_type : logits
   backward : margin_cross_entropy_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : masked_multihead_attention_
   args : (Tensor x, Tensor cache_kv, Tensor bias, Tensor src_mask, Tensor cum_offsets, Tensor sequence_lengths, Tensor rotary_tensor, Tensor beam_cache_offset, Tensor qkv_out_scale, Tensor out_shift, Tensor out_smooth, int seq_len, int rotary_emb_dims, bool use_neox_rotary_style=false, str compute_dtype = "default", float out_scale=-1, int quant_round_type=1, float quant_max_bound=127.0, float quant_min_bound=-127.0)
@@ -3231,6 +3233,7 @@
   kernel :
     func : mode
   backward : mode_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : momentum_
   args : (Tensor param, Tensor grad, Tensor velocity, Tensor learning_rate, Tensor master_param, float mu, bool use_nesterov = false, str regularization_method = "", float regularization_coeff = 0.0f, bool multi_precision = false, float rescale_grad = 1.0f)