[CINN] use pass decomp flatten in dyshape (#64115)

phlrain · web-flow · commit 0177b51cc4ef · 2024-05-09T13:30:51.000+08:00
* disable flatten decomp in dyshape

* update

* revert unit test
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
@@ -934,6 +934,95 @@ class UnsqueezeOpPattern
   }
 };
 
+class FlattenOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::FlattenOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::FlattenOp>::OpRewritePattern;
+
+  bool Match(paddle::dialect::FlattenOp op) const override {
+    const bool is_denied = CompatibleInfo::IsDeniedForCinn(*op.operation());
+
+    bool is_dyshape = op->operand_source(0)
+                          .type()
+                          .dyn_cast<pir::ShapedTypeInterface>()
+                          .IsDynamicShape();
+    return !is_denied && is_dyshape;
+  }
+
+  void Rewrite(paddle::dialect::FlattenOp op,
+               pir::PatternRewriter &rewriter) const override {
+    int start_axis =
+        op.attribute("start_axis").dyn_cast<::pir::Int32Attribute>().data();
+    int end_axis =
+        op.attribute("stop_axis").dyn_cast<::pir::Int32Attribute>().data();
+
+    // build output shape
+    std::vector<pir::Value> out_shape;
+    auto x_rank = op->operand_source(0)
+                      .type()
+                      .dyn_cast<paddle::dialect::DenseTensorType>()
+                      .dims()
+                      .size();
+    auto x_shape =
+        rewriter.Build<paddle::dialect::ShapeOp>(op->operand_source(0))
+            .result(0);
+    for (size_t i = 0; i < x_rank;) {
+      if (i == static_cast<size_t>(start_axis)) {
+        auto new_single_dim =
+            rewriter
+                .Build<cinn::dialect::SliceOp>(x_shape,
+                                               std::vector<int64_t>({0}),
+                                               std::vector<int64_t>({i}),
+                                               std::vector<int64_t>({i + 1}),
+                                               std::vector<int64_t>({}),
+                                               std::vector<int64_t>({}))
+                .result(0);
+
+        for (auto t = start_axis + 1; t <= end_axis; ++t) {
+          auto dim_t =
+              rewriter
+                  .Build<cinn::dialect::SliceOp>(x_shape,
+                                                 std::vector<int64_t>({0}),
+                                                 std::vector<int64_t>({t}),
+                                                 std::vector<int64_t>({t + 1}),
+                                                 std::vector<int64_t>({}),
+                                                 std::vector<int64_t>({}))
+                  .result(0);
+          new_single_dim =
+              rewriter.Build<paddle::dialect::MultiplyOp>(new_single_dim, dim_t)
+                  .result(0);
+        }
+        out_shape.push_back(new_single_dim);
+        i = end_axis + 1;
+      } else {
+        auto t =
+            rewriter
+                .Build<cinn::dialect::SliceOp>(x_shape,
+                                               std::vector<int64_t>({0}),
+                                               std::vector<int64_t>({i}),
+                                               std::vector<int64_t>({i + 1}),
+                                               std::vector<int64_t>({}),
+                                               std::vector<int64_t>({}))
+                .result(0);
+        out_shape.push_back(t);
+        i++;
+      }
+    }
+
+    auto new_shape =
+        rewriter.Build<cinn::dialect::ConcatOp>(out_shape, -1).result(0);
+
+    auto reshape_op = rewriter.Build<paddle::dialect::ReshapeOp>(
+        op->operand_source(0), new_shape);
+
+    reshape_op.result(0).set_type(op.result(0).type());
+
+    rewriter.ReplaceAllUsesWith(op.result(0), reshape_op.result(0));
+
+    rewriter.EraseOp(op);
+  }
+};
+
 class SigmoidOpPattern
     : public pir::OpRewritePattern<paddle::dialect::SigmoidOp> {
  public:
@@ -1043,6 +1132,7 @@ pir::RewritePatternSet PdOpToCinnOpPass::InitializePatterns(
   ps.Add<UnsqueezeOpPattern>(context);
   ps.Add<SigmoidOpPattern>(context);
   ps.Add<GatherOpPattern>(context);
+  ps.Add<FlattenOpPattern>(context);
 
   return ps;
 }
diff --git a/paddle/fluid/primitive/base/decomp_trans.cc b/paddle/fluid/primitive/base/decomp_trans.cc
@@ -43,8 +43,8 @@ std::unordered_set<std::string> decomp_op_contain_none = {"pd_op.squeeze",
                                                           "pd_op.batch_norm",
                                                           "pd_op.batch_norm_"};
 //
-std::unordered_set<std::string> dynamic_shape_blacklist = {"pd_op.squeeze",
-                                                           "pd_op.unsqueeze"};
+std::unordered_set<std::string> dynamic_shape_blacklist = {
+    "pd_op.squeeze", "pd_op.unsqueeze", "pd_op.flatten"};
 
 namespace {
 std::set<std::string> StringSplit(const std::string& str) {
diff --git a/test/ir/pir/cinn/symbolic/test_sub_graph_flatten.py b/test/ir/pir/cinn/symbolic/test_sub_graph_flatten.py
@@ -0,0 +1,73 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle.static import InputSpec
+
+
+class FlattenCase(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+        self,
+        var_0,
+    ):
+        var_1 = paddle.flatten(var_0, 1, 2)
+        return var_1
+
+
+class TestSigmoid(unittest.TestCase):
+    def setUp(self):
+        self.inputs = (
+            paddle.rand(shape=[32, 16, 8, 256], dtype=paddle.float32),
+        )
+        self.net = FlattenCase()
+
+    def train(self, net, to_static, with_prim=False, with_cinn=False):
+        if to_static:
+            # paddle.set_flags({'FLAGS_prim_all': with_prim})
+            if with_cinn:
+                build_strategy = paddle.static.BuildStrategy()
+                build_strategy.build_cinn_pass = True
+                input_spec = [
+                    InputSpec(shape=[None, None, None, 256], dtype='float32')
+                ]
+                net = paddle.jit.to_static(
+                    net,
+                    build_strategy=build_strategy,
+                    input_spec=input_spec,
+                    full_graph=True,
+                )
+            else:
+                net = paddle.jit.to_static(net, full_graph=True)
+        paddle.seed(123)
+        outs = net(*self.inputs)
+        return outs
+
+    def test_ast_prim_cinn(self):
+        # st_out = self.train(self.net, to_static=True)
+        cinn_out = self.train(
+            self.net, to_static=True, with_prim=True, with_cinn=True
+        )
+        # for st, cinn in zip(
+        #     paddle.utils.flatten(st_out), paddle.utils.flatten(cinn_out)
+        # ):
+        #     np.testing.assert_allclose(st.numpy(), cinn.numpy(), atol=1e-8)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -382,32 +382,6 @@ def setUp(self):
         self.tol = 5e-6
 
 
-class TestPrimFlatten1(TestPrimBase):
-    def setUp(self):
-        np.random.seed(2023)
-        self.dtype = "float32"
-        self.x_shape = [3, 100, 100, 4]
-        self.init_x_shape = [3, None, None, 4]
-        self.x = np.random.random(self.x_shape).astype(self.dtype)
-        self.net = flatten_net
-        self.necessary_ops = "pd_op.flatten"
-        self.enable_cinn = False
-        self.tol = 1e-6
-
-
-class TestPrimFlatten2(TestPrimBase):
-    def setUp(self):
-        np.random.seed(2023)
-        self.dtype = "float32"
-        self.x_shape = [3, 100, 100, 640]
-        self.init_x_shape = [None, None, None, 640]
-        self.x = np.random.random(self.x_shape).astype(self.dtype)
-        self.net = flatten_net
-        self.necessary_ops = "pd_op.flatten"
-        self.enable_cinn = False
-        self.tol = 1e-6
-
-
 class TestPrimGroupNorm1(TestPrimBase):
     def setUp(self):
         np.random.seed(2023)
