[CINN] 【Infer Symbolic Shape BUAA 】Add crop op (#67957)

crazyxiaoxi · Jeff114514 · web-flow · commit f09a9a2bf1b3 · 2024-09-19T15:49:21.000+08:00
* first try

* error

* fix

* rerun

* codestyle

* fix comment

* fix commit

---------

Co-authored-by: Jeff114514 &lt;928430299@qq.com&gt;
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -760,12 +760,72 @@ bool ChannelShuffleOpInferSymbolicShape(
   return true;
 }
 
-// bool CropOpInferSymbolicShape(pir::Operation *op,
-//                               pir::InferSymbolicShapeContext *infer_context)
-//                               {
-//   // pass
-//   return true;
-// }
+bool CropOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const std::vector<symbol::DimExpr> &x_shape = x_shape_or_data.shape();
+  std::vector<symbol::DimExpr> offsets;
+  std::vector<symbol::DimExpr> in_shape;
+  std::vector<symbol::DimExpr> out_dims;
+
+  if (op->HasAttribute("offsets")) {
+    std::vector<int64_t> offsets_ =
+        paddle::dialect::details::GetVectorAttr<int64_t>(op, "offsets");
+    for (const auto &i : offsets_) offsets.emplace_back(symbol::DimExpr{i});
+  } else {
+    const auto &offsets_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(2));
+    offsets = details::GetOrCreateExprVecFromData(offsets_shape_or_data,
+                                                  infer_context);
+  }
+
+  if (op->HasAttribute("shape")) {
+    std::vector<int64_t> shape_ =
+        paddle::dialect::details::GetVectorAttr<int64_t>(op, "shape");
+    for (const auto &i : shape_) in_shape.emplace_back(symbol::DimExpr{i});
+  } else {
+    const auto &in_shape_or_data =
+        infer_context->GetShapeOrDataForValue(op->operand_source(1));
+    in_shape =
+        details::GetOrCreateExprVecFromData(in_shape_or_data, infer_context);
+  }
+
+  PADDLE_ENFORCE_EQ(in_shape.size(),
+                    x_shape.size(),
+                    phi::errors::InvalidArgument(
+                        "The number of elements (%d) of attribute 'shape' for "
+                        "CropTensor must be equal to the number of "
+                        "dimensions (%d) of the input.",
+                        in_shape.size(),
+                        x_shape.size()));
+  PADDLE_ENFORCE_EQ(
+      offsets.size(),
+      x_shape.size(),
+      phi::errors::InvalidArgument(
+          "The number of elements (%d) of attribute 'offsets' for "
+          "CropTensor must be equal to the number of "
+          "dimensions (%d) of the input.",
+          offsets.size(),
+          x_shape.size()));
+
+  for (size_t i = 0; i < in_shape.size(); ++i) {
+    if (in_shape[i].isa<int64_t>()) {
+      if (in_shape[i].Get<int64_t>() == -1) {
+        out_dims.push_back(symbol::DimExpr(x_shape[i] - offsets[i]));
+      } else {
+        out_dims.push_back(symbol::DimExpr(in_shape[i]));
+      }
+    } else {
+      out_dims.push_back(in_shape[i]);
+    }
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
+  return true;
+}
 
 bool DecodeJpegOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -47,7 +47,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cumsum_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ChannelShuffle)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(CheckNumerics)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(Crop)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Crop)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DecodeJpeg)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Diag)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(DiagEmbed)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1112,6 +1112,7 @@
     func : crop
     data_type : x
   backward : crop_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : cross
   args : (Tensor x, Tensor y, int axis = 9)
diff --git a/test/legacy_test/test_crop_tensor_op.py b/test/legacy_test/test_crop_tensor_op.py
@@ -109,6 +109,9 @@ def initTestCase(self):
         self.offsets = [1, 5, 3]
         self.shape_by_input = True
 
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
 
 class TestCase4(TestCropTensorOp):
     def initTestCase(self):
@@ -117,6 +120,9 @@ def initTestCase(self):
         self.offsets = [0, 0, 1, 0]
         self.shape_by_input = True
 
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
 
 class TestCase5(TestCropTensorOp):
     def initTestCase(self):
@@ -134,6 +140,9 @@ def initTestCase(self):
         self.shape_by_input = True
         self.offset_by_input = True
 
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
 
 class TestCropTensorOpTensorAttr(OpTest):
     def setUp(self):
@@ -183,7 +192,7 @@ def initTestCase(self):
         self.shape_attr = [0, 0]
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_symbol_infer=False)
 
     def test_check_grad_normal(self):
         self.check_grad(["X"], "Out", check_pir=True)
@@ -214,6 +223,9 @@ def initTestCase(self):
         self.ShapeTensor = False
         self.OffsetsTensor = True
 
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=True)
+
 
 class TestCropTensorOpTensorAttrCase4(TestCropTensorOpTensorAttr):
     def initTestCase(self):
@@ -224,6 +236,9 @@ def initTestCase(self):
         self.offsets_attr = [-1, -1, 3]
         self.OffsetsTensor = True
 
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=True)
+
 
 class TestCropTensorException(unittest.TestCase):
     @test_with_pir_api
