【Infer Symbolic Shape No.218】Add infer_symbol_shape for reindexgraph (#68973)

AndPuQing · web-flow · commit cead7f59d4f0 · 2024-10-30T17:55:12.000+08:00
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
@@ -2067,12 +2067,71 @@ bool GraphKhopSamplerOpInferSymbolicShape(
   return true;
 }
 
-// bool GraphReindexOpInferSymbolicShape(pir::Operation *op,
-//                                           pir::InferSymbolicShapeContext
-//                                           *infer_context) {
-//   // pass
-//   return true;
-// }
+bool ReindexGraphOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const symbol::ShapeOrDataDimExprs &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const symbol::ShapeOrDataDimExprs &neighbors_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const symbol::ShapeOrDataDimExprs &count_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+  const symbol::ShapeOrDataDimExprs &hashtable_value_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(3));
+  const symbol::ShapeOrDataDimExprs &hashtable_index_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(4));
+
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &neighbors_shape =
+      neighbors_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &count_shape = count_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &hashtable_value_shape =
+      hashtable_value_shape_or_data.shape();
+  const std::vector<symbol::DimExpr> &hashtable_index_shape =
+      hashtable_index_shape_or_data.shape();
+
+  auto GraphReindexShapeCheck = [&](const std::vector<symbol::DimExpr> &shape,
+                                    const std::string &tensor_name) {
+    if (shape.size() == 2) {
+      infer_context->AddEqualCstr(shape[1], symbol::DimExpr{1});
+    } else {
+      PADDLE_ENFORCE_EQ(
+          shape.size(),
+          1,
+          common::errors::InvalidArgument(
+              "The %s should be 1D, when it is not 2D, but we get %d",
+              tensor_name,
+              shape.size()));
+    }
+  };
+
+  GraphReindexShapeCheck(x_shape, "X");
+  GraphReindexShapeCheck(neighbors_shape, "Neighbors");
+  GraphReindexShapeCheck(count_shape, "Count");
+
+  if (!paddle::dialect::details::IsFakeValue(op->operand_source(3)) &&
+      !paddle::dialect::details::IsFakeValue(op->operand_source(4))) {
+    GraphReindexShapeCheck(hashtable_value_shape, "HashTable_Value");
+    GraphReindexShapeCheck(hashtable_index_shape, "HashTable_Index");
+  }
+
+  VLOG(3) << "ReindexGraphOpInferSymbolicShape";
+
+  symbol::DimExpr unique_dim_sym = infer_context->GetNextSymName();
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({neighbors_shape[0]})});
+  infer_context->SetShapeOrDataForValue(
+      op->result(1),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({neighbors_shape[0]})});
+  infer_context->SetShapeOrDataForValue(
+      op->result(2),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs({unique_dim_sym})});
+
+  return true;
+}
 
 bool GraphSampleNeighborsOpInferSymbolicShape(
     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
@@ -66,7 +66,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedGemmEpilogue)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FusedMultiTransformer)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GenerateProposals)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GraphKhopSampler)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(GraphReindex)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GraphSampleNeighbors)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Gru)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GruUnit)
@@ -114,6 +113,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiPool)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RoiAlign)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SpectralNorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SequenceConv)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReindexGraph)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseAttention)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Stack)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SigmoidCrossEntropyWithLogits)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4071,7 +4071,7 @@
     data_type : x
   optional : hashtable_value, hashtable_index
   traits : paddle::dialect::ForwardOnlyTrait
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : relu
   args : (Tensor x)
