[CINN] Fix bug of SplitWithNumOpInferSymbolicShape And Add unit test (#65266)

gongshaotian · web-flow · commit d5b838be0058 · 2024-06-21T12:50:17.000+08:00
* 1.add SymbolChecker in OpTest 2.bind ShapeAnalysis and ShapeOrData to python 3.add log

* bind shape_analysis and shape_or_data to python level

* modify get_shape_constraint_ir_analysis as lambda func

* 1.delete print_shape_or_data 2. modify get_shape_constraint_ir_analysis as a lambda func

* delete using in pir.cc

* delete note

* Complete prototype implementation of SymbolInferTest

* fix bug and refine log

* delete note

* delete compare fun in shape_or_data_expr.h

* dele print and refine compare func

* 1.fix bug of SplitWithNumOp in symbol infer along with assign op 2. add unit test for SplitWithNumOp

* delete AssignOpInferSymbolicShape

* refine logic of SplitWithNumOpInferSymbolicShape

* replace Unimplemented as InvalidArgument

* refine error message

* delete print
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
@@ -825,13 +825,26 @@ bool SplitOpInferSymbolicShape(pir::Operation *op,
 
 bool SplitWithNumOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
-  int64_t axis = op->operand_source(1)
-                     .defining_op<paddle::dialect::FullOp>()
-                     .attributes()
-                     .at("value")
-                     .dyn_cast<paddle::dialect::ScalarAttribute>()
-                     .data()
-                     .to<int64_t>();
+  const symbol::ShapeOrDataDimExprs &axis_shape_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  PADDLE_ENFORCE_EQ(
+      axis_shape_data.data().has_value(),
+      true,
+      phi::errors::InvalidArgument(
+          "In InferSymbolicShape, axis of SplitWithNumOp is null"));
+  const std::vector<symbol::DimExpr> &axis_data =
+      axis_shape_data.data().value();
+  PADDLE_ENFORCE_EQ(
+      axis_data.size() == 1,
+      true,
+      phi::errors::InvalidArgument(
+          "In SplitWithNumOp, data of axis should be one dimension"));
+  if (!axis_data[0].isa<int64_t>()) {
+    PADDLE_THROW(
+        phi::errors::InvalidArgument("The type of axis must be int64_t"));
+  }
+  int64_t axis = axis_data[0].dyn_cast<int64_t>();
+
   const auto &attributes = op->attributes();
   int num = attributes.at("num").dyn_cast<pir::Int32Attribute>().data();
   const auto &x_s_or_d =
diff --git a/test/ir/pir/cinn/symbolic/CMakeLists.txt b/test/ir/pir/cinn/symbolic/CMakeLists.txt
@@ -23,7 +23,8 @@ if(WITH_GPU)
     test_llama_mlp_dy.py
     test_while_st.py
     test_infer_sym_shape_utils.py
-    test_dyshape_cast.py)
+    test_dyshape_cast.py
+    test_dyshape_split_with_num.py)
 
   foreach(cinn_pir_test_name ${CINN_PIR_SYMBOLIC_TEST})
     string(REGEX REPLACE ".py" "" cinn_pir_test_name ${cinn_pir_test_name})
@@ -222,4 +223,17 @@ if(WITH_GPU)
     WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
   set_tests_properties(test_dyshape_cast PROPERTIES LABELS "RUN_TYPE=CINN")
 
+  add_test(
+    NAME test_dyshape_split_with_num
+    COMMAND
+      ${CMAKE_COMMAND} -E env
+      PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+      FLAGS_prim_all=true FLAGS_cinn_bucket_compile=True
+      FLAGS_group_schedule_tiling_first=1 FLAGS_enable_pir_api=1
+      ${PYTHON_EXECUTABLE}
+      ${CMAKE_CURRENT_SOURCE_DIR}/test_dyshape_split_with_num.py
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+  set_tests_properties(test_dyshape_split_with_num PROPERTIES LABELS
+                                                              "RUN_TYPE=CINN")
+
 endif()
diff --git a/test/ir/pir/cinn/symbolic/test_dyshape_split_with_num.py b/test/ir/pir/cinn/symbolic/test_dyshape_split_with_num.py
@@ -0,0 +1,79 @@
+# 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 sys
+import unittest
+from os.path import dirname
+
+import numpy as np
+
+import paddle
+from paddle import nn
+from paddle.static import InputSpec
+
+sys.path.append(dirname(dirname(__file__)))
+
+import utils
+
+
+class SplitWithNumLayer(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        axis = paddle.full(shape=[1], fill_value=3, dtype='int32')
+        axis = paddle.assign(axis)
+        num = 4
+        return paddle.split(x, num, axis=axis)
+
+
+class TestSplitWithNum(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(2024)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.shape = [4, 5, 6, 8]
+        self.x = paddle.randn(self.shape, dtype='float16')
+        self.x.stop_gradient = True
+
+    def check_jit_kernel_info(self, static_fn):
+        utils.check_jit_kernel_number(static_fn, 0)
+        # utils.check_jit_kernel_structure(static_fn, {utils.JIT_KERNEL_NAME: })
+
+    def eval(self, use_cinn):
+        net = SplitWithNumLayer()
+        input_spec = [InputSpec([None, None, None, 8], dtype='float16')]
+        net = utils.apply_to_static(net, use_cinn, input_spec)
+        net.eval()
+        outs = net(self.x)
+        if use_cinn:
+            self.check_jit_kernel_info(net.forward)
+        return outs
+
+    def test_eval(self):
+        cinn_outs = self.eval(use_cinn=True)
+        dygraph_outs = self.eval(use_cinn=False)
+        assert len(cinn_outs) == len(dygraph_outs)
+        for i in range(len(cinn_outs)):
+            np.testing.assert_allclose(
+                cinn_outs[i].numpy(),
+                dygraph_outs[i].numpy(),
+                atol=1e-6,
+                rtol=1e-6,
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()
