【Infer Symbolic Shape】【BUAA】Add tdm_sampler op (#67461)

Guanhuachen2003 · web-flow · commit a3c2b562074b · 2024-08-21T14:31:43.000+08:00
* add tdmsampler

* modify codestyle

* Update multiary_infer_sym.cc

* Update multiary_infer_sym.h

* Update ops.yaml
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
@@ -2150,12 +2150,37 @@ bool SigmoidCrossEntropyWithLogits_OpInferSymbolicShape(
 //   return SyncBatchNormOpInferSymbolicShape(op, infer_context);
 // }
 
-// bool TdmSamplerOpInferSymbolicShape(pir::Operation *op,
-//                                     pir::InferSymbolicShapeContext
-//                                     *infer_context) {
-//   // pass
-//   return true;
-// }
+bool TdmSamplerOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_dims = x_shape_or_data.shape();
+
+  bool output_positive =
+      op->attribute<pir::BoolAttribute>("output_positive").data();
+  std::vector<int> neg_samples_num_list =
+      paddle::dialect::details::GetVectorAttr<int>(op, "neg_samples_num_list");
+
+  int64_t sample_res_length = 0;
+  for (int sample_nums : neg_samples_num_list) {
+    sample_res_length += sample_nums + static_cast<int64_t>(output_positive);
+  }
+
+  symbol::DimExpr batch_size = x_dims[0];
+  symbol::DimExpr sample_res_dim(sample_res_length);
+
+  std::vector<symbol::DimExpr> output_dims = {batch_size, sample_res_dim};
+  symbol::TensorShapeOrDataDimExprs output_shape(output_dims);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::ShapeOrDataDimExprs{output_shape});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1), symbol::ShapeOrDataDimExprs{output_shape});
+  infer_context->SetShapeOrDataForValue(
+      op->result(2), symbol::ShapeOrDataDimExprs{output_shape});
+
+  return true;
+}
 
 bool TrilinearInterpOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
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
@@ -110,7 +110,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SyncBatchNorm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SyncBatchNorm_)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(SaveCombine)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmSampler)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(TdmSampler)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TrilinearInterp)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(HsigmoidLoss)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ViterbiDecode)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4688,7 +4688,7 @@
     func : tdm_sampler
     data_type : x
   optional : labels
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : temporal_shift
   args : (Tensor x, int seg_num, float shift_ratio = 0.25f, str data_format = "NCHW")
