Write a pass to check infer symbolic. (#61222)

JiaWenxuan · web-flow · commit 18ed85d9ce7c · 2024-02-01T11:22:52.000+08:00
* Write a pass to check infer symbolic.

* fix some bugs

* fix some bugs

* use optional to handle null values

* fix some bugs

* rename the my unittest module
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.cc
@@ -0,0 +1,144 @@
+// 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.
+
+#pragma once
+
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.h"
+
+#include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
+#include "paddle/cinn/hlir/dialect/runtime/ir/runtime_dialect.h"
+#include "paddle/cinn/runtime/flags.h"
+#include "paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.h"
+#include "paddle/pir/core/builtin_type.h"
+#include "paddle/pir/dialect/shape/utils/dim_expr.h"
+#include "paddle/utils/flags.h"
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+namespace {
+
+std::string SprintShape(const std::vector<std::int64_t>& shape) {
+  std::string str = "[";
+  for (std::int64_t value : shape) {
+    str += std::to_string(value);
+    if (value != shape.back()) {
+      str += ", ";
+    }
+  }
+  return str + "]";
+}
+
+void PrintProgram(pir::ModuleOp m, const std::string& mgs) {
+  std::ostringstream print_stream;
+  print_stream << "\n\n";
+  m.program()->Print(print_stream);
+  print_stream << "\n\n";
+  VLOG(4) << "===================== " << mgs << " =====================\n"
+          << print_stream.str();
+}
+
+std::vector<std::int64_t> GetStaticValueShape(pir::Value value) {
+  const auto& dim = value.type().dyn_cast<::pir::DenseTensorType>().dims();
+  return ::common::vectorize(dim);
+}
+
+std::optional<std::vector<std::int64_t>> GetDynamicValueShape(
+    pir::Value value, const pir::ShapeConstraintIRAnalysis& shape_analysis) {
+  if (!shape_analysis.HasShapeOrDataForValue(value)) {
+    return std::nullopt;
+  }
+  const auto& dim_expr_dynamic_shapes =
+      shape_analysis.GetShapeOrDataForValue(value).shape();
+  std::vector<std::int64_t> dynamic_shapes{};
+  for (const auto& dim_expr_shape : dim_expr_dynamic_shapes) {
+    CHECK(dim_expr_shape.Has<std::int64_t>());
+    dynamic_shapes.push_back(dim_expr_shape.Get<std::int64_t>());
+  }
+  return dynamic_shapes;
+}
+
+void CompareStaticAndDynamicValueShape(
+    pir::Value value,
+    const pir::ShapeConstraintIRAnalysis& shape_analysis,
+    int op_index,
+    pir::ModuleOp module_op) {
+  std::vector<std::int64_t> static_value_shape = GetStaticValueShape(value);
+  std::optional<std::vector<std::int64_t>> opt_dynamic_value_shape =
+      GetDynamicValueShape(value, shape_analysis);
+  if (opt_dynamic_value_shape.has_value()) {
+    if (static_value_shape != opt_dynamic_value_shape.value()) {
+      VLOG(4) << "CheckInferSymbolic failed, in the fellowing program, the "
+              << op_index
+              << "th op : the shape is not equal\nthe static shape is: "
+              << SprintShape(static_value_shape)
+              << ", and the dynamic shape is: "
+              << SprintShape(opt_dynamic_value_shape.value());
+      PrintProgram(module_op, "CheckInferSymbolic");
+    }
+  } else {
+    VLOG(4) << "CheckInferSymbolic failed, in the fellowing program, the "
+            << op_index << "th op infer symbolic failed";
+    PrintProgram(module_op, "CheckInferSymbolic");
+  }
+}
+
+void CheckInferSymbolic(pir::ModuleOp module_op) {
+  VLOG(4) << "CheckInferSymbolic start";
+  int op_index = 0;
+  const auto& shape_analysis =
+      pir::ShapeAnalysisManager::Instance().Get(module_op.program());
+  for (uint32_t i = 0; i < module_op->num_regions(); i++) {
+    for (const auto& block : module_op->region(i)) {
+      for (const auto& op : block) {
+        for (std::size_t j = 0; j < op.num_operands(); ++j) {
+          CompareStaticAndDynamicValueShape(
+              op.operand_source(j), shape_analysis, op_index, module_op);
+        }
+        for (std::size_t j = 0; j < op.num_results(); ++j) {
+          CompareStaticAndDynamicValueShape(
+              op.result(j), shape_analysis, op_index, module_op);
+        }
+        op_index++;
+      }
+    }
+  }
+  VLOG(4) << "CheckInferSymbolic end";
+}
+
+class CheckInferSymbolicPass : public pir::Pass {
+ public:
+  CheckInferSymbolicPass() : pir::Pass("check_infer_symbolic_pass", 1) {}
+
+  void Run(pir::Operation* op) override {
+    pir::ModuleOp module_op = op->dyn_cast<pir::ModuleOp>();
+    CheckInferSymbolic(module_op);
+  }
+
+  bool CanApplyOn(pir::Operation* op) const override {
+    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<::pir::Pass> CreateCheckInferSymbolicPass() {
+  return std::make_unique<CheckInferSymbolicPass>();
+}
+
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.h
@@ -0,0 +1,29 @@
+// 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.
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include "paddle/pir/pass/pass.h"
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+// This is a helper pass for checking the symbolic inference accuracy.
+std::unique_ptr<::pir::Pass> CreateCheckInferSymbolicPass();
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -87,6 +87,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_static_dim_to_dynamic_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/divide_group_op_to_fusion_op_pass.h"
@@ -153,6 +154,7 @@ USE_PIR_PASS(batch_norm_act_fuse_pass);
 
 PHI_DECLARE_bool(print_ir);
 PHI_DECLARE_bool(pir_apply_shape_optimization_pass);
+PHI_DECLARE_bool(check_infer_symbolic);
 
 namespace paddle {
 namespace pybind {
@@ -1583,6 +1585,11 @@ void AddCinnPass(std::shared_ptr<PassManager> &pass_manager,  // NOLINT
   if (FLAGS_print_ir) {
     pass_manager->EnableIRPrinting();
   }
+
+  if (!has_dynamic_shape && FLAGS_check_infer_symbolic) {
+    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
+    pass_manager->AddPass(cinn::dialect::ir::CreateCheckInferSymbolicPass());
+  }
   pass_manager->AddPass(cinn::dialect::ir::CreatePdOpToCinnOpPass());
   pass_manager->AddPass(
       std::make_unique<cinn::dialect::ir::AddBroadcastToElementwisePass>());
diff --git a/paddle/phi/core/flags.cc b/paddle/phi/core/flags.cc
@@ -1403,6 +1403,21 @@ PHI_DEFINE_EXPORTED_bool(pir_apply_inplace_pass,
                          "Whether to apply inplace pass on lowering "
                          "::pir::Program to Kernel Dialect");
 
+/**
+ * Apply check infer symbolic pass FLAG
+ * Name: check_infer_symbolic_pass
+ * Since Version: 3.0.0
+ * Value Range: bool, default=false
+ * Example:
+ * Note: If Ture, will apply check_infer_symbolic pass.
+ */
+PHI_DEFINE_EXPORTED_bool(
+    check_infer_symbolic,
+    false,
+    "Whether to use check_infer_symbolic_pass. This pass can check "
+    "the symbolic inference accuracy by comparing the the value "
+    "shape between dynamic shape and static shape.");
+
 /**
  * Apply shape optimization pass to new IR FLAG
  * Name: pir_apply_shape_optimization_pass
diff --git a/test/ir/pir/cinn/symbolic/CMakeLists.txt b/test/ir/pir/cinn/symbolic/CMakeLists.txt
@@ -5,7 +5,8 @@ if(WITH_GPU)
     "test_*.py")
 
   list(REMOVE_ITEM CINN_PIR_SYMBOLIC_TEST test_cinn_reduce_symbolic_demo.py
-       test_sub_graph_for_backend.py test_sub_graph_for_frontend.py)
+       test_sub_graph_for_backend.py test_sub_graph_for_frontend.py
+       test_check_infer_symbolic.py)
 
   foreach(cinn_pir_test_name ${CINN_PIR_SYMBOLIC_TEST})
     string(REGEX REPLACE ".py" "" cinn_pir_test_name ${cinn_pir_test_name})
@@ -60,4 +61,16 @@ if(WITH_GPU)
   set_tests_properties(test_sub_graph_for_frontend PROPERTIES LABELS
                                                               "RUN_TYPE=CINN")
 
+  add_test(
+    NAME test_check_infer_symbolic
+    COMMAND
+      ${CMAKE_COMMAND} -E env
+      PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+      FLAGS_check_infer_symbolic=True FLAGS_cinn_bucket_compile=False
+      FLAGS_enable_pir_api=1 ${PYTHON_EXECUTABLE}
+      ${CMAKE_CURRENT_SOURCE_DIR}/test_check_infer_symbolic.py
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+  set_tests_properties(test_check_infer_symbolic PROPERTIES LABELS
+                                                            "RUN_TYPE=CINN")
+
 endif()
diff --git a/test/ir/pir/cinn/symbolic/test_check_infer_symbolic.py b/test/ir/pir/cinn/symbolic/test_check_infer_symbolic.py
@@ -0,0 +1,77 @@
+# 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 numpy as np
+
+import paddle
+
+
+def apply_to_static(net, use_cinn):
+    build_strategy = paddle.static.BuildStrategy()
+    build_strategy.build_cinn_pass = use_cinn
+    return paddle.jit.to_static(
+        net,
+        build_strategy=build_strategy,
+        full_graph=True,
+    )
+
+
+def exp_sub(x):
+    y = paddle.exp(x)
+    z = y - x
+    return z
+
+
+class CheckInferSymbolicNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.fn = exp_sub
+
+    def forward(self, x):
+        out = self.fn(x)
+        return out
+
+
+class TestCheckInferSymbolic(unittest.TestCase):
+    """
+    Test Pir API + @to_static + CINN.
+    """
+
+    def setUp(self):
+        paddle.seed(2022)
+        self.prepare_data()
+
+    def prepare_data(self):
+        self.shape = [64, 96]
+        self.x = paddle.randn(self.shape, dtype="float32")
+        self.x.stop_gradient = False
+
+    def eval(self, use_cinn):
+        paddle.seed(2022)
+        net = CheckInferSymbolicNet()
+        if use_cinn:
+            net = apply_to_static(net, use_cinn)
+        net.eval()
+        out = net(self.x)
+        return out
+
+    def test_eval(self):
+        cinn_out = self.eval(use_cinn=True)
+        dy_out = self.eval(use_cinn=False)
+        np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
+
+
+if __name__ == '__main__':
+    unittest.main()
