[CINN] Add symbolic backend for Flip op (#63654)

hxzd5568 · web-flow · commit 2cc225d52db5 · 2024-04-26T16:31:33.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/ir/ops.yaml b/paddle/cinn/hlir/dialect/operator/ir/ops.yaml
@@ -84,6 +84,15 @@
     func : reshape
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
+- op : reverse
+  args : (Tensor x, int[] axis)
+  output : Tensor
+  infer_meta :
+    func : ReverseInferMeta
+  kernel :
+    func : reverse
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+
 - op : scale
   args : (Tensor x, float scale=1.0, float bias=0.0, bool bias_after_scale=true)
   output : Tensor(out)
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
@@ -238,6 +238,35 @@ class ReshapeOpPattern
   }
 };
 
+class FlipOpPattern : public pir::OpRewritePattern<paddle::dialect::FlipOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::FlipOp>::OpRewritePattern;
+
+  bool Match(paddle::dialect::FlipOp op) const override {
+    const bool is_denied = CompatibleInfo::IsDeniedForCinn(*op.operation());
+    return !is_denied;
+  }
+
+  void Rewrite(paddle::dialect::FlipOp op,
+               pir::PatternRewriter &rewriter) const override {
+    std::vector<int> axis_value;
+    auto axis_attr =
+        op.attribute("axis").dyn_cast<pir::ArrayAttribute>().AsVector();
+    if (axis_attr.size() > 0) {
+      for (size_t i = 0; i < axis_attr.size(); ++i) {
+        PADDLE_ENFORCE(axis_attr[i].dyn_cast<::pir::Int32Attribute>(),
+                       ::common::errors::PreconditionNotMet(
+                           "Reqiured attr element must be Int32Attribute."));
+        axis_value.push_back(
+            axis_attr[i].dyn_cast<::pir::Int32Attribute>().data());
+      }
+    }
+    auto cinn_reverse = rewriter.Build<cinn::dialect::ReverseOp>(
+        op->operand_source(0), axis_value);
+    rewriter.ReplaceAllUsesWith(op.result(0), cinn_reverse.result(0));
+    rewriter.EraseOp(op);
+  }
+};
 class Pool2dOpPattern
     : public pir::OpRewritePattern<paddle::dialect::Pool2dOp> {
  public:
@@ -1004,6 +1033,7 @@ pir::RewritePatternSet PdOpToCinnOpPass::InitializePatterns(
   ps.Add<SplitWithNumOpPattern>(context);
   ps.Add<SplitOpPattern>(context);
   ps.Add<ExpandOpPattern>(context);
+  ps.Add<FlipOpPattern>(context);
   ps.Add<IsCloseOpPattern>(context);
   ps.Add<ElementwisePowOpPattern>(context);
   ps.Add<FullWithTensorOpPattern>(context);
diff --git a/paddle/cinn/hlir/op/transform.cc b/paddle/cinn/hlir/op/transform.cc
@@ -976,6 +976,56 @@ std::shared_ptr<OpStrategy> StrategyForReverse(
   return strategy;
 }
 
+std::shared_ptr<OpStrategy> StrategyForReverseSymbolic(
+    const framework::NodeAttr &attrs,
+    const std::vector<ir::Tensor> &inputs,
+    const std::vector<Type> &out_type,
+    const std::vector<std::vector<ir::Dim>> &output_shapes,
+    const Target &target) {
+  // check output shape
+  CHECK(!output_shapes.empty() && !output_shapes[0].empty())
+      << "Output shape is empty! Please check.\n";
+  // get axis[0, n_dim)
+  std::vector<int> axis;
+  if (attrs.attr_store.find("axis") != attrs.attr_store.end()) {
+    axis = absl::get<std::vector<int>>(attrs.attr_store.at("axis"));
+    for (auto &e : axis) {
+      if (e >= static_cast<int>(output_shapes[0].size()) ||
+          e < -1 * static_cast<int>(output_shapes[0].size())) {
+        PADDLE_THROW(::common::errors::InvalidArgument(
+            "axis is not in [0, n_dim), Please check."));
+      }
+      if (e < 0) {
+        e += output_shapes[0].size();
+      }
+    }
+  }
+
+  framework::CINNCompute reverse_compute([=](lang::Args args,
+                                             lang::RetValue *ret) {
+    CHECK(!args.empty())
+        << "The input argument of reverse compute is empty! Please check.\n";
+    CINNValuePack input_args = args[0];
+    CHECK(!input_args.empty())
+        << "at least one input tensor for reverse compute\n";
+    Expr A = input_args[0];
+    CHECK(A.as_tensor());
+
+    CHECK_EQ(input_args.size(), 2);
+    CHECK(input_args[1].is_string());
+    std::string tensor_name = input_args[1].operator std::string();
+    auto out = pe::Reverse(A.as_tensor_ref(), axis, tensor_name);
+    auto stages = CreateStages({A.as_tensor_ref(), out});
+    *ret = CINNValuePack{{CINNValue(out), CINNValue(stages)}};
+  });
+
+  auto strategy = std::make_shared<framework::OpStrategy>();
+  CHECK(out_type.size()) << "Out_type of reverse op is empty! Please check.";
+  strategy->AddImpl(
+      reverse_compute, lang::PackedFunc(), "strategy.reverse.x86", 1);
+  return strategy;
+}
+
 std::vector<framework::shape_t> InferShapeForReverse(
     const std::vector<framework::shape_t> &inputs_shape,
     const framework::AttrMapType &attrs) {
@@ -2164,6 +2214,8 @@ CINN_REGISTER_HELPER(transform_ops) {
       .set_num_outputs(1)
       .set_attr<cinn::hlir::framework::StrategyFunction>(
           "CINNStrategy", cinn::hlir::op::StrategyForReverse)
+      .set_attr<cinn::hlir::framework::StrategyFunctionSymbolic>(
+          "CINNStrategySymbolic", cinn::hlir::op::StrategyForReverseSymbolic)
       .set_attr("infershape",
                 MakeOpFunction(cinn::hlir::op::InferShapeForReverse))
       .set_attr("inferdtype",
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -105,6 +105,7 @@ OP_SAME_OPERANDS_AND_RESULT(Reciprocal_)
 OP_SAME_OPERANDS_AND_RESULT(Relu)
 OP_SAME_OPERANDS_AND_RESULT(Relu6)
 OP_SAME_OPERANDS_AND_RESULT(Relu_)
+OP_SAME_OPERANDS_AND_RESULT(Reverse)
 OP_SAME_OPERANDS_AND_RESULT(Roll)
 OP_SAME_OPERANDS_AND_RESULT(Round)
 OP_SAME_OPERANDS_AND_RESULT(Round_)
@@ -176,7 +177,8 @@ bool ScaleOpInferSymbolicShape(pir::Operation *op,
 }  // namespace paddle::dialect
 
 namespace cinn::dialect {
+using paddle::dialect::ReverseOpInferSymbolicShape;
 using paddle::dialect::ScaleOpInferSymbolicShape;
-}
+}  // namespace cinn::dialect
 
 #undef OP_SAME_OPERANDS_AND_RESULT
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -96,6 +96,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reciprocal_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Relu)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Relu6)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Relu_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reverse)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Roll)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Round)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Round_)
@@ -131,5 +132,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Sigmoid_)
 }  // namespace paddle::dialect
 
 namespace cinn::dialect {
+using paddle::dialect::ReverseOpInferSymbolicShape;
 using paddle::dialect::ScaleOpInferSymbolicShape;
-}
+}  // namespace cinn::dialect
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -2396,6 +2396,7 @@
     func : reverse
     data_type : x
   backward : reverse_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : rms_norm
   args : (Tensor x, Tensor bias, Tensor residual, Tensor norm_weight, Tensor norm_bias, float epsilon, int begin_norm_axis, float quant_scale, int quant_round_type, float quant_max_bound, float quant_min_bound)
diff --git a/test/ir/pir/cinn/symbolic/test_cinn_elementwise_symbolic.py b/test/ir/pir/cinn/symbolic/test_cinn_elementwise_symbolic.py
@@ -58,6 +58,10 @@ def isclose(x, y):
     )
 
 
+def reverse(x):
+    return paddle.flip(x, axis=[1])
+
+
 class CINNSubGraphNet(paddle.nn.Layer):
     def __init__(self, fn):
         super().__init__()
@@ -451,5 +455,41 @@ def test_eval_symbolic(self):
         np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
 
 
+class TestCinnSubGraphFlipOrReverse(unittest.TestCase):
+    """
+    Test Pir API + @to_static + CINN.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.x_shape = [32, 32]
+        self.x = paddle.randn(self.x_shape, dtype="float32")
+        self.x.stop_gradient = False
+
+    def check_jit_kernel_info(self, static_fn):
+        utils.check_jit_kernel_number(static_fn, 1)
+
+    def eval_symbolic(self, use_cinn):
+        paddle.seed(2022)
+        net = CINNSubGraphNet(reverse)
+        input_spec = [
+            InputSpec(shape=[None, 32], dtype='float32'),
+        ]
+        net = utils.apply_to_static(net, use_cinn, input_spec)
+        net.eval()
+        out = net(self.x)
+        if use_cinn:
+            self.check_jit_kernel_info(net.forward)
+        return out
+
+    def test_eval_symbolic(self):
+        cinn_out = self.eval_symbolic(use_cinn=True)
+        dy_out = self.eval_symbolic(use_cinn=False)
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
+
+
 if __name__ == '__main__':
     unittest.main()
