[PIR][DynamicShape] Fix bug in slice op's InferSymbolicShape (#62247)

lanxianghit · web-flow · commit cbe8810bbea2 · 2024-03-02T08:19:07.000+08:00
* Fix bug in slice op's InferSymbolicShape

* add more tests

* fix ci
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc
@@ -16,6 +16,17 @@
 
 namespace paddle::dialect::details {
 
+std::optional<std::vector<int64_t>> VecExpr2Int64(const ExprVec &expr_vec) {
+  std::vector<int64_t> int64vec;
+  for (auto item : expr_vec) {
+    if (!item.isa<int64_t>()) {
+      return std::nullopt;
+    }
+    int64vec.push_back(item.Get<int64_t>());
+  }
+  return int64vec;
+}
+
 bool ReduceInferDim(pir::Operation *op,
                     pir::ShapeConstraintIRAnalysis *shape_analysis,
                     const std::vector<int64_t> &axis,
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h
@@ -17,6 +17,12 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
+// To make codes shorter
+using ExprVec = std::vector<symbol::DimExpr>;
+using ShapeOrData = symbol::ShapeOrDataDimExprs;
+using TensorExprs = symbol::TensorShapeOrDataDimExprs;
+using TensorListExprs = symbol::TensorListShapeOrDataDimExprs;
+
 namespace paddle::dialect::details {
 template <typename T>
 struct AttributeTrait;
@@ -60,6 +66,8 @@ std::vector<T> GetVectorAttr(const ::pir::Operation *op,
   return vec_res;
 }
 
+std::optional<std::vector<int64_t>> VecExpr2Int64(const ExprVec &expr_vec);
+
 bool ReduceInferDim(pir::Operation *op,
                     pir::ShapeConstraintIRAnalysis *shape_analysis,
                     const std::vector<int64_t> &axis,
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc
@@ -19,11 +19,6 @@
 
 namespace paddle::dialect {
 
-// To make codes shorter
-using ShapeOrData = symbol::ShapeOrDataDimExprs;
-using TensorExprs = symbol::TensorShapeOrDataDimExprs;
-using TensorListExprs = symbol::TensorListShapeOrDataDimExprs;
-
 bool DataOpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const auto &attributes = op->attributes();
@@ -270,9 +265,104 @@ bool FullIntArrayOpInferSymbolicShape(
   return true;
 }
 
+inline void CheckAndUpdateSliceAttrs(
+    const ExprVec &in_dims,
+    const std::vector<int64_t> &axes,
+    ExprVec *starts_p,
+    ExprVec *ends_p,
+    std::vector<int64_t> *infer_flags = nullptr) {
+  auto vec_int64 = details::VecExpr2Int64(*starts_p);
+  IR_ENFORCE(vec_int64.has_value(),
+             "for slice op, all the elements in `starts` must be int64_t");
+  std::vector<int64_t> starts_int = vec_int64.value();
+
+  vec_int64 = details::VecExpr2Int64(*ends_p);
+  IR_ENFORCE(vec_int64.has_value(),
+             "for slice op, all the elements in `ends` must be int64_t");
+  std::vector<int64_t> ends_int = vec_int64.value();
+
+  ExprVec &starts = *starts_p;
+  ExprVec &ends = *ends_p;
+  auto IsMaxInt = [](const symbol::DimExpr &expr) {
+    return expr.isa<int64_t>() &&
+           expr.Get<int64_t>() ==
+               static_cast<int64_t>(std::numeric_limits<int>::max());
+  };
+
+  for (size_t i = 0; i < axes.size(); ++i) {
+    int64_t axis = axes[i];
+
+    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
+      PADDLE_THROW(
+          phi::errors::Unimplemented("SliceOpInferSymbolicShape CAN NOT "
+                                     "deal with -1 in infer_flags now"));
+    }
+
+    // For both start and end can be negtive or positive, we need to handle the
+    // following different arrangements.
+    ends[i] = IsMaxInt(ends[i]) ? in_dims[axis] : ends[i];
+
+    bool both_negative_or_positive = (starts_int[i] >= 0 && ends_int[i] >= 0) ||
+                                     (starts_int[i] <= 0 && ends_int[i] <= 0);
+    bool start_negative_end_positive = starts_int[i] <= 0 && ends_int[i] >= 0;
+    bool start_positive_end_negative = starts_int[i] >= 0 && ends_int[i] <= 0;
+
+    if (both_negative_or_positive) {
+      continue;
+    } else if (start_negative_end_positive) {
+      starts[i] = starts[i] + in_dims[axis];
+    } else if (start_positive_end_negative) {
+      starts[i] = starts[i] - in_dims[axis];
+    } else {
+      LOG(FATAL) << "Dead code";
+    }
+  }
+}
+
+inline ExprVec GetSliceDims(const ExprVec &in_dims,
+                            const std::vector<int64_t> &axes,
+                            const ExprVec &starts,
+                            const ExprVec &ends,
+                            std::vector<int64_t> *infer_flags = nullptr) {
+  ExprVec slice_dims(in_dims);
+
+  for (size_t i = 0; i < axes.size(); ++i) {
+    int64_t axis = axes[i];
+
+    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
+      PADDLE_THROW(
+          phi::errors::Unimplemented("SliceOpInferSymbolicShape CAN NOT "
+                                     "deal with -1 in infer_flags now"));
+    }
+
+    slice_dims[axis] = ends[i] - starts[i];
+  }
+
+  return slice_dims;
+}
+
+inline ExprVec GetDecreasedDims(const ExprVec &slice_dims,
+                                const std::vector<int64_t> &decrease_axes) {
+  ExprVec decreased_dims(slice_dims);
+  std::vector<uint8_t> decrease_flag(slice_dims.size(), 0);
+  if (decrease_axes.size() > 0) {
+    for (size_t i = 0; i < decrease_axes.size(); ++i) {
+      int64_t axis = decrease_axes[i];
+      decrease_flag[axis] = 1;
+    }
+    ExprVec new_shape;
+    for (size_t i = 0; i < slice_dims.size(); ++i) {
+      if (decrease_flag[i] == 0) {
+        new_shape.emplace_back(slice_dims[i]);
+      }
+    }
+    decreased_dims = new_shape;
+  }
+  return decreased_dims;
+}
+
 bool SliceOpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  // TODO(zhangbopd): Not implemented yet.
   pir::Value operand_source = op->operand_source(0);
   pir::Value operand_starts = op->operand_source(1);
   pir::Value operand_ends = op->operand_source(2);
@@ -285,107 +375,76 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
   const symbol::ShapeOrDataDimExprs &ends_shape_data =
       shape_analysis->GetShapeOrDataForValue(operand_ends);
 
-  // Currently, we DO NOT support the case that any element in `axes` `starts`
-  // or `ends` is a Symbol.
   const std::vector<int64_t> axes = [&] {
-    const auto &attributes = op->attributes();
-    pir::Attribute attr_axes = attributes.at("axes");
-
-    const auto &axes_vec = attr_axes.dyn_cast<pir::ArrayAttribute>().AsVector();
-    std::vector<int64_t> axes;
+    std::vector<int64_t> axes_vec = details::GetVectorAttr(op, "axes");
     int64_t rank = int64_t(operand_shape_or_data.shape().size());
-    for (auto item : axes_vec) {
-      int64_t axis = item.dyn_cast<pir::Int64Attribute>().data();
-      axes.emplace_back(axis >= 0 ? axis : std::max(int64_t(0), axis + rank));
+    for (size_t i = 0; i < axes_vec.size(); i++) {
+      int64_t axis = axes_vec[i];
+      axes_vec[i] = axis >= 0 ? axis : std::max(int64_t(0), axis + rank);
     }
-    return axes;
+    return axes_vec;
   }();
 
-  const std::vector<int64_t> starts = [&] {
-    std::vector<int64_t> starts;
-    for (auto item : starts_shape_data.data().value()) {
-      IR_ENFORCE(item.isa<int64_t>(),
-                 "Currently, we DO NOT support the case that any element in "
-                 "`starts` is a Symbol.");
-      starts.push_back(item.Get<int64_t>());
-    }
-    return starts;
-  }();
+  // Currently, we DO NOT support any element in `starts` is a Symbol.
+  ExprVec starts = starts_shape_data.data().value();
+  ExprVec ends = ends_shape_data.data().value();
 
-  const std::vector<int64_t> ends = [&] {
-    std::vector<int64_t> ends;
-    for (auto item : ends_shape_data.data().value()) {
-      IR_ENFORCE(item.isa<int64_t>(),
-                 "Currently, we DO NOT support the case that any element in "
-                 "`ends` is a Symbol.");
-      ends.push_back(item.Get<int64_t>());
+  std::vector<int64_t> infer_flags = [op, &axes] {
+    std::vector<int64_t> infer_flags_t =
+        details::GetVectorAttr(op, "infer_flags");
+    if (infer_flags_t.empty()) {
+      infer_flags_t = std::vector<int64_t>(axes.size(), 1);
     }
-    return ends;
+    return infer_flags_t;
   }();
 
-  // When `pd.slice` is operating on a tensor which is produced by a `pd.shape`
-  // op, the reseult should be written into data.
-  const auto &GetDataDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
-    const std::vector<symbol::DimExpr> out_data = [&] {
-      std::vector<symbol::DimExpr> out_data;
-      const int64_t start =
-          starts[0] < 0
-              ? starts[0] + operand_shape_or_data.data().value().size()
-              : starts[0];
-      const int64_t end =
-          static_cast<int64_t>(std::numeric_limits<int>::max()) == ends[0]
-              ? operand_shape_or_data.data().value().size()
-              : ends[0];
-
-      for (int64_t i = start; i < end; i++) {
-        out_data.push_back(operand_shape_or_data.data().value()[i]);
-      }
-      return out_data;
-    }();
-    const std::vector<symbol::DimExpr> shape{std::int64_t(out_data.size())};
-    return symbol::ShapeOrDataDimExprs{
-        symbol::TensorShapeOrDataDimExprs(shape, out_data)};
-  };
+  const std::vector<int64_t> decrease_axis =
+      details::GetVectorAttr(op, "decrease_axis");
 
-  // Othewise, the reseult should be written into the shape.
   const auto &GetShapeDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
-    std::vector<symbol::DimExpr> out_shape = operand_shape_or_data.shape();
+    const ExprVec &in_dims = operand_shape_or_data.shape();
+    CheckAndUpdateSliceAttrs(in_dims, axes, &starts, &ends, &infer_flags);
+    ExprVec slice_dims =
+        GetSliceDims(in_dims, axes, starts, ends, &infer_flags);
+    ExprVec out_dims = GetDecreasedDims(slice_dims, decrease_axis);
 
-    const std::vector<symbol::DimExpr> &dim_expr_starts =
-        starts_shape_data.data().value();
-    const std::vector<symbol::DimExpr> &dim_expr_ends =
-        ends_shape_data.data().value();
+    return symbol::ShapeOrDataDimExprs{
+        symbol::TensorShapeOrDataDimExprs(out_dims)};
+  };
 
-    // For both start and end can be negtive or positive, we need to handle the
-    // following different arrangements.
-    auto IsMaxInt = [](const symbol::DimExpr &expr) {
-      return expr.isa<int64_t>() &&
-             expr.Get<int64_t>() ==
-                 static_cast<int64_t>(std::numeric_limits<int>::max());
-    };
-    for (size_t i = 0; i < axes.size(); ++i) {
-      const int64_t axis = axes[i];
-      auto end =
-          IsMaxInt(dim_expr_ends[i]) ? out_shape[axis] : dim_expr_ends[i];
-
-      bool both_negative_or_positive =
-          (starts[i] >= 0 && ends[i] >= 0) || (starts[i] <= 0 && ends[i] <= 0);
-      bool start_negative_end_positive = starts[i] <= 0 && ends[i] >= 0;
-      bool start_positive_end_negative = starts[i] >= 0 && ends[i] <= 0;
-
-      if (both_negative_or_positive) {
-        out_shape[axis] = end - dim_expr_starts[i];
-      } else if (start_negative_end_positive) {
-        out_shape[axis] = end - dim_expr_starts[i] - out_shape[axis];
-      } else if (start_positive_end_negative) {
-        out_shape[axis] = out_shape[axis] - dim_expr_starts[i] + end;
-      } else {
-        LOG(FATAL) << "Dead code";
-      }
+  // When `pd.slice` is operating on a tensor which is produced by a `pd.shape`
+  // op, the reseult should be written into data.
+  const auto &GetDataDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
+    std::vector<symbol::DimExpr> out_data;
+
+    // Currently, we DO NOT support the case that any element in `axes` `starts`
+    // or `ends` is a Symbol.
+    auto vec_int64 = details::VecExpr2Int64(starts);
+    IR_ENFORCE(vec_int64.has_value(),
+               "for slice op, all the elements in `starts` must be int64_t");
+    std::vector<int64_t> starts_int = vec_int64.value();
+
+    vec_int64 = details::VecExpr2Int64(ends);
+    IR_ENFORCE(vec_int64.has_value(),
+               "for slice op, all the elements in `ends` must be int64_t");
+    std::vector<int64_t> ends_int = vec_int64.value();
+
+    const int64_t start =
+        starts_int[0] < 0
+            ? starts_int[0] + operand_shape_or_data.data().value().size()
+            : starts_int[0];
+    const int64_t end =
+        static_cast<int64_t>(std::numeric_limits<int>::max()) == ends_int[0]
+            ? operand_shape_or_data.data().value().size()
+            : ends_int[0];
+
+    for (int64_t i = start; i < end; i++) {
+      out_data.push_back(operand_shape_or_data.data().value()[i]);
     }
 
+    const std::vector<symbol::DimExpr> shape{std::int64_t(out_data.size())};
     return symbol::ShapeOrDataDimExprs{
-        symbol::TensorShapeOrDataDimExprs(out_shape)};
+        symbol::TensorShapeOrDataDimExprs(shape, out_data)};
   };
 
   symbol::ShapeOrDataDimExprs shape_data =
diff --git a/test/cpp/pir/shape_dialect/shape_optimization_test.cc b/test/cpp/pir/shape_dialect/shape_optimization_test.cc
@@ -122,10 +122,10 @@ TEST(shape_optimization, shape_optimization_pass) {
             "Mul(Mul(Mul(Mul(1, S1), 128), 32), 1 / (128))");
   EXPECT_EQ(cast_res.shape()[3], 2);
 
-  EXPECT_EQ(symbol::ToString(relu_res.shape()[0]), "Add(Add(S2, -2), -2)");
-  EXPECT_EQ(symbol::ToString(relu_res.shape()[1]), "Add(Add(S3, -2), -2)");
-  EXPECT_EQ(symbol::ToString(relu_res.shape()[2]), "Add(Add(S4, -2), -2)");
-  EXPECT_EQ(symbol::ToString(relu_res.shape()[3]), "Add(Add(S5, -2), -2)");
+  EXPECT_EQ(symbol::ToString(relu_res.shape()[0]), "Add(-2, -Add(2, -S2))");
+  EXPECT_EQ(symbol::ToString(relu_res.shape()[1]), "Add(-2, -Add(2, -S3))");
+  EXPECT_EQ(symbol::ToString(relu_res.shape()[2]), "Add(-2, -Add(2, -S4))");
+  EXPECT_EQ(symbol::ToString(relu_res.shape()[3]), "Add(-2, -Add(2, -S5))");
 
   EXPECT_EQ(subtract_res.shape()[0], 1);
   EXPECT_EQ(subtract_res.shape()[1], 64);
diff --git a/test/ir/pir/cinn/symbolic/test_op_infer_sym_shape.py b/test/ir/pir/cinn/symbolic/test_op_infer_sym_shape.py
@@ -459,5 +459,63 @@ def test_eval_symbolic(self):
         return True
 
 
+class SliceNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        out = x[:, -1, :]
+        out = x[1:3, 0:2, 2:4]
+
+        axes = [0, 1, 2]
+        starts = [-3, 0, 2]
+        ends = [3, 2, 4]
+        out = paddle.slice(x, axes=axes, starts=starts, ends=ends)
+
+        return out
+
+
+class TestSliceOpInferSymbolicShape(TestBase):
+    def prepare_data(self):
+        self.cases = [np.random.rand(4, 5, 6)]
+
+        self.expected = [
+            [
+                'shape[S0, S2], data[NULL]',
+                'shape[2, 2, 2], data[NULL]',
+                'shape[Add(3, -Add(-3, S0)), 2, 2]',
+            ]
+        ]
+
+    def test_eval_symbolic(self):
+        net = SliceNet()
+
+        for i in range(len(self.cases)):
+            x = self.cases[i]
+            x_spec = InputSpec(
+                shape=[None for index in range(len(x.shape))], dtype='float32'
+            )
+
+            input_spec = [x_spec]
+            net = apply_to_static(net, False, input_spec)
+            net.eval()
+
+            # check the infer result
+            sym_shape_str_list = get_sym_shape_str_for_op(
+                net, input_spec, 'pd_op.slice'
+            )
+            np.testing.assert_equal(
+                len(sym_shape_str_list), len(self.expected[i])
+            )
+            for j in range(len(sym_shape_str_list)):
+                np.testing.assert_equal(
+                    sym_shape_str_list[j].find(self.expected[i][j]),
+                    0,
+                    f'in case i,j = {i},{j}: output shape ({sym_shape_str_list[0]}) is not expected {(self.expected[i][j])}',
+                )
+
+        return True
+
+
 if __name__ == '__main__':
     unittest.main()
