【Infer Symbolic Shape No.60,61,62】[BUAA] Add index_add,index_put,index_select op (#67267)

Luohongzhige · web-flow · commit 847b9c841ecc · 2024-08-21T15:16:10.000+08:00
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
@@ -1460,6 +1460,138 @@ bool YoloBoxOpInferSymbolicShape(
   return true;
 }
 
+bool IndexSelectOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &index_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+
+  std::vector<symbol::DimExpr> x_shape = x_shape_or_data.shape();
+  std::vector<symbol::DimExpr> index_shape = index_shape_or_data.shape();
+
+  int64_t dim = op->attribute<pir::Int32Attribute>("dim").data();
+
+  auto input_rank = x_shape.size();
+  auto index_rank = index_shape.size();
+  PADDLE_ENFORCE_EQ(
+      dim < static_cast<int64_t>(input_rank) &&
+          dim >= (0 - static_cast<int64_t>(input_rank)),
+      true,
+      common::errors::OutOfRange(
+          "Attr(dim) is out of range, It's expected "
+          "to be in range of [-%d, %d]. But received Attr(dim) = %d.",
+          input_rank,
+          input_rank - 1,
+          dim));
+
+  PADDLE_ENFORCE_EQ(index_rank == 1 || index_rank == 2,
+                    true,
+                    common::errors::InvalidArgument(
+                        "The 'shape' of Input(Index) must be 1-D tensor or 2-D "
+                        "tensor where second dimension is 1. "
+                        "But received: the 'shape' of Input(Index) is [%s], "
+                        "the dimension of Input(Index) is [%d].",
+                        index_shape,
+                        index_shape.size()));
+
+  if (index_rank == 2)
+    infer_context->AddEqualCstr(index_shape[1], symbol::DimExpr{1});
+
+  if (dim < 0) {
+    dim += input_rank;
+  }
+
+  std::vector<symbol::DimExpr> output_shape = x_shape;
+  output_shape[dim] = index_shape[0];
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_shape)});
+
+  return true;
+}
+
+bool IndexSelect_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return IndexSelectOpInferSymbolicShape(op, infer_context);
+}
+
+bool IndexAddOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &index_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &add_value_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+  std::vector<symbol::DimExpr> x_shape = x_shape_or_data.shape();
+  std::vector<symbol::DimExpr> index_shape = index_shape_or_data.shape();
+  std::vector<symbol::DimExpr> add_value_shape =
+      add_value_shape_or_data.shape();
+  int axis = op->attribute<pir::Int32Attribute>("axis").data();
+  int ndims_x = x_shape.size();
+
+  // Real axis calculation
+  int real_axis = axis >= 0 ? axis : axis + ndims_x;
+
+  // Check dimensions
+  PADDLE_ENFORCE_EQ(
+      index_shape.size(),
+      1,
+      common::errors::InvalidArgument("Index tensor must be 1-dimensional."));
+
+  PADDLE_ENFORCE_EQ(
+      x_shape.size(),
+      add_value_shape.size(),
+      common::errors::InvalidArgument(
+          "Input and addition value must have the same dimension."));
+
+  for (int i = 0; i < ndims_x; i++) {
+    if (i != real_axis) {
+      infer_context->AddEqualCstr(x_shape[i], add_value_shape[i]);
+    }
+  }
+
+  // Set the shape for the output
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+
+  return true;
+}
+
+bool IndexAdd_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return IndexAddOpInferSymbolicShape(op, infer_context);
+}
+
+bool IndexPutOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_shape = x_shape_or_data.shape();
+
+  PADDLE_ENFORCE_LT(
+      x_shape.size(),
+      7,
+      common::errors::InvalidArgument(
+          "The rank of input should be less than 7, but received %d.",
+          x_shape.size()));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+
+  return true;
+}
+
+bool IndexPut_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return IndexPutOpInferSymbolicShape(op, infer_context);
+}
+
 }  // namespace paddle::dialect
 
 namespace cinn::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
@@ -53,6 +53,11 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherNd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(HuberLoss)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Histogram)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexAdd)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexAdd_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexPut)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexPut_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelect)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AccuracyCheck)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSample)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelectStrided)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2531,6 +2531,7 @@
     data_type : x
   inplace : (x -> out)
   backward : index_add_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : index_put
   args : (Tensor x, Tensor[] indices, Tensor value, bool accumulate=false)
@@ -2542,6 +2543,7 @@
     data_type : x
   inplace : (x -> out)
   backward : index_put_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : index_sample
   args : (Tensor x, Tensor index)
@@ -2567,6 +2569,7 @@
   backward : index_select_grad
   data_transform :
     skip_transform : index
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : index_select_strided
   args : (Tensor x, int64_t index, int axis = 0)
