【Infer Symbolic Shape BUAA No.63】Add index_select_strided op (#66735)

crazyxiaoxi · web-flow · commit 3d0e24d6b7c8 · 2024-08-07T10:23:43.000+08:00
* test_index_select_strided

* testop
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
@@ -802,6 +802,31 @@ bool IscloseOpInferSymbolicShape(
   return true;
 }
 
+bool IndexSelectStridedOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> input_dims = x_shape_or_data.shape();
+
+  int dim = op->attribute<pir::Int32Attribute>("dim").data();
+
+  if (dim < 0) {
+    dim += input_dims.size();
+  }
+
+  std::vector<symbol::DimExpr> output_dims(input_dims.begin(),
+                                           input_dims.end());
+  output_dims.erase(output_dims.begin() + dim);
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(output_dims)});
+
+  // No need to add any constraints here as we are simply removing a dimension.
+  return true;
+}
+
 bool AccuracyCheckOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   // The shape of output is the same as input `values` (op->operand_source(1))
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
@@ -46,6 +46,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(GatherNd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AccuracyCheck)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSample)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelectStrided)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Kron)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lstsq)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(MatrixRankTol)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2527,6 +2527,7 @@
     func : index_select_strided
     data_type : x
   backward : index_select_strided_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : instance_norm
   args : (Tensor x, Tensor scale, Tensor bias, float epsilon=1e-5)
