[CINN]delete redundant InferSymbolicShape in passes (#64176)

Hongqing-work · web-flow · commit b047fa1c8ee3 · 2024-05-15T20:30:09.000+08:00
* [CINN]delete redundant InferSymbolicShape in passes

* delete shape_optimization_pass.h in shape_analysis.h

* change name

* fix code style
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_pass.cc
@@ -136,23 +136,6 @@ class BlockDimExprsAsserter {
   DimExprs4ValueT MakeOpDimExprs4Value(const pir::Operation* op) {
     auto shape_analysis = std::make_shared<pir::ShapeConstraintIRAnalysis>();
     InitLocalShapeAnalysis(*op, shape_analysis.get());
-
-    pir::Operation* mut_op = const_cast<pir::Operation*>(op);
-    auto interface =
-        mut_op->dyn_cast<paddle::dialect::InferSymbolicShapeInterface>();
-    if (!interface) {
-      PADDLE_THROW(phi::errors::Unimplemented(
-          op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
-    } else {
-      // TODO(Hongqing-work): delete this after the shape analysis reconstruct
-      // is done.
-      bool infer_result = interface.InferSymbolicShape(
-          shape_analysis->GetInferSymbolicShapeContext());
-      PADDLE_ENFORCE_EQ(infer_result,
-                        true,
-                        ::common::errors::PreconditionNotMet(
-                            "InferSymbolicShape for %s failed.", op->name()));
-    }
     return [shape_analysis](
                pir::Value value) -> const symbol::ShapeOrDataDimExprs& {
       return shape_analysis->GetShapeOrDataForValue(value);
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_util.cc b/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_util.cc
@@ -22,6 +22,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.h"
 #include "paddle/common/flags.h"
 #include "paddle/pir/include/dialect/shape/transforms/shape_optimization_pass.h"
+#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 
 COMMON_DECLARE_bool(check_infer_symbolic);
 PD_DECLARE_bool(prim_all);
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc
@@ -207,24 +207,6 @@ std::vector<pir::Operation*> GetSubGraphFromOutputToInputsValue(
   return ops;
 }
 
-void InferSymbolicShapeForSubgraph(
-    const std::vector<pir::Operation*>& ops,
-    pir::ShapeConstraintIRAnalysis* shape_analysis) {
-  for (auto* op : ops) {
-    auto infer_symbolic_shape_interface =
-        op->dyn_cast<paddle::dialect::InferSymbolicShapeInterface>();
-    if (infer_symbolic_shape_interface) {
-      // TODO(Hongqing-work): delete this after the shape analysis reconstruct
-      // is done.
-      infer_symbolic_shape_interface.InferSymbolicShape(
-          shape_analysis->GetInferSymbolicShapeContext());
-    } else {
-      PADDLE_THROW(phi::errors::Unimplemented(
-          op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
-    }
-  }
-}
-
 void UpdateLocalShapeAnalysis(
     const std::vector<pir::Value>& input_tensors,
     pir::Value shape,
@@ -261,10 +243,6 @@ void UpdateLocalShapeAnalysis(
           input_tensor, symbol::TensorShapeOrDataDimExprs(new_shape));
     }
   }
-  // infer new symbol shape for shape value
-  std::vector<pir::Operation*> sub_graph_ops =
-      GetSubGraphFromOutputToInputsValue(input_tensors, shape);
-  InferSymbolicShapeForSubgraph(sub_graph_ops, shape_analysis);
 }
 
 std::optional<pir::Value> GetOutOfRewrittenGenerateShapeOp(
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.cc
@@ -36,21 +36,12 @@ namespace {
 
 pir::Value GetOutputDimTensor(pir::PatternRewriter* rewriter,
                               pir::Value x,
-                              pir::Value y,
-                              pir::ShapeConstraintIRAnalysis* shape_analysis) {
+                              pir::Value y) {
   pir::Operation* x_shape_op = rewriter->Build<paddle::dialect::ShapeOp>(x);
   pir::Operation* y_shape_op = rewriter->Build<paddle::dialect::ShapeOp>(y);
   pir::Operation* shape_broadcast_op =
       rewriter->Build<paddle::dialect::ShapeBroadcastOp>(x_shape_op->result(0),
                                                          y_shape_op->result(0));
-  for (auto* op : std::vector{x_shape_op, y_shape_op, shape_broadcast_op}) {
-    auto infer_symbolic_shape_interface =
-        op->dyn_cast<paddle::dialect::InferSymbolicShapeInterface>();
-    // TODO(Hongqing-work): delete this after the shape analysis reconstruct is
-    // done.
-    infer_symbolic_shape_interface.InferSymbolicShape(
-        shape_analysis->GetInferSymbolicShapeContext());
-  }
   return shape_broadcast_op->result(0);
 }
 
@@ -67,19 +58,16 @@ bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
     return false;
   }
 
-  pir::Value output_dim_tensor =
-      GetOutputDimTensor(rewriter, x, y, &shape_analysis);
+  pir::Value output_dim_tensor = GetOutputDimTensor(rewriter, x, y);
   if (x_shape.shape() != out_shape.shape()) {
     pir::Value broadcasted_x =
         rewriter->Build<paddle::dialect::ExpandOp>(x, output_dim_tensor).out();
     op->operand(0).set_source(broadcasted_x);
-    shape_analysis.SetShapeOrDataForValue(broadcasted_x, out_shape);
   }
   if (y_shape.shape() != out_shape.shape()) {
     pir::Value broadcasted_y =
         rewriter->Build<paddle::dialect::ExpandOp>(y, output_dim_tensor).out();
     op->operand(1).set_source(broadcasted_y);
-    shape_analysis.SetShapeOrDataForValue(broadcasted_y, out_shape);
   }
   return true;
 }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/collect_sym_expr.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/collect_sym_expr.cc
@@ -207,10 +207,6 @@ CreateGroupShapeOrDataExprs(
 
   // process the result values of each op.
   for (auto* op : group->ops()) {
-    // TODO(Hongqing-work): delete this after the shape analysis reconstruct is
-    // done.
-    InferSymbolicShapeForOperation(
-        op, local_shape_analysis.GetInferSymbolicShapeContext());
     for (size_t i = 0; i < op->num_results(); ++i) {
       auto result = op->result(i);
       if (result && !value2shape.count(result)) {
diff --git a/paddle/common/union_find_set.h b/paddle/common/union_find_set.h
@@ -65,7 +65,7 @@ class UnionFindSet {
 
   bool HasSameRoot(const T& p, const T& q) const { return Find(p) == Find(q); }
 
-  std::unordered_map<T, T>* GetMap() { return &parent_; }
+  std::unordered_map<T, T>* MutMap() { return &parent_; }
 
  private:
   std::unordered_map<T, T> parent_;
diff --git a/paddle/pir/include/dialect/shape/transforms/shape_optimization_pass.h b/paddle/pir/include/dialect/shape/transforms/shape_optimization_pass.h
@@ -15,19 +15,22 @@
 #pragma once
 
 #include <memory>
+#include "paddle/pir/include/core/builtin_op.h"
 #include "paddle/pir/include/core/dll_decl.h"
-#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 #include "paddle/pir/include/pass/pass_manager.h"
 
 namespace pir {
 
 class Pass;
 
+class InferSymbolicShapeContext;
+
 IR_API std::unique_ptr<Pass> CreateShapeOptimizationPass();
 
 void InferSymExprForBlock(const Block &block,
                           InferSymbolicShapeContext *infer_context);
 
+void InferSymExprForAllValues(ModuleOp module_op);
 }  // namespace pir
 
 namespace pir::shape {
diff --git a/paddle/pir/include/dialect/shape/utils/shape_analysis.h b/paddle/pir/include/dialect/shape/utils/shape_analysis.h
@@ -27,8 +27,13 @@
 
 namespace pir {
 
+void InferSymExprForAllValues(ModuleOp module_op);
+
 class IR_API InferSymbolicShapeContext {
  public:
+  InferSymbolicShapeContext() = default;
+  InferSymbolicShapeContext(const InferSymbolicShapeContext&) = delete;
+  InferSymbolicShapeContext(InferSymbolicShapeContext&&) = delete;
   void Init();
 
   const std::string GetNextSymName();
@@ -134,17 +139,17 @@ class IR_API ShapeConstraintIRAnalysis final
   symbol::DimExpr GetProductDimExpr(Value lhs,
                                     const std::vector<int>& lhs_dim_idxs);
 
-  // TODO(hongqing-work): make it a private component only for infer friend
-  // class
-  InferSymbolicShapeContext* GetInferSymbolicShapeContext() {
-    return &context_;
-  }
-
   const symbol::ConstraintsManager& constraints_manager() const {
     return context_.constraints_manager();
   }
 
  private:
+  InferSymbolicShapeContext* MutInferSymbolicShapeContext() {
+    return &context_;
+  }
+
+  friend void InferSymExprForAllValues(ModuleOp module_op);
+
   void SetStaticShapeForValue(Value val);
 
   void InferShapeOrDataForValue(Value val);
diff --git a/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc b/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc
@@ -21,6 +21,7 @@
 #include "paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/infer_symbolic_shape.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_attribute.h"
 #include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
+#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
 #include "paddle/pir/include/pass/pass_manager.h"
 #include "paddle/pir/include/pass/pass_registry.h"
 
@@ -226,18 +227,6 @@ void CheckInferSymWithInferMeta(
   }
 }
 
-void InferSymExprForAllValues(ModuleOp module_op) {
-  ShapeConstraintIRAnalysis& shape_analysis =
-      ShapeAnalysisManager::Instance().Get(module_op.program());
-  shape_analysis.Init();
-  auto infer_context = shape_analysis.GetInferSymbolicShapeContext();
-  for (uint32_t i = 0; i < module_op->num_regions(); i++) {
-    for (auto& block : module_op->region(i)) {
-      InferSymExprForBlock(block, infer_context);
-    }
-  }
-}
-
 class ShapeOptimizationPass : public pir::Pass {
  public:
   ShapeOptimizationPass() : pir::Pass("shape_optimization_pass", 0) {}
@@ -253,7 +242,7 @@ class ShapeOptimizationPass : public pir::Pass {
             "ShapeOptimizationPass should run on module op."));
     PrintProgram(module_op, "Origin Program");
 
-    InferSymExprForAllValues(module_op);
+    ::pir::InferSymExprForAllValues(module_op);
     // Runner is for Canonicalizer.
     PassPipelineRunner runner = [](pir::PassManager& pm, pir::ModuleOp m) {
       pm.EnableIRPrinting();
@@ -349,6 +338,18 @@ void InferSymExprForBlock(const Block& block,
   }
 }
 
+void InferSymExprForAllValues(ModuleOp module_op) {
+  ShapeConstraintIRAnalysis& shape_analysis =
+      ShapeAnalysisManager::Instance().Get(module_op.program());
+  shape_analysis.Init();
+  auto infer_context = shape_analysis.MutInferSymbolicShapeContext();
+  for (uint32_t i = 0; i < module_op->num_regions(); i++) {
+    for (auto& block : module_op->region(i)) {
+      InferSymExprForBlock(block, infer_context);
+    }
+  }
+}
+
 std::unique_ptr<Pass> CreateShapeOptimizationPass() {
   return std::make_unique<ShapeOptimizationPass>();
 }
diff --git a/paddle/pir/src/dialect/shape/utils/constraints_manager.cc b/paddle/pir/src/dialect/shape/utils/constraints_manager.cc
@@ -263,7 +263,7 @@ void ConstraintsManager::SubstituteInConstraint(const DimExpr& origin,
 
 template <typename DoEachT>
 void ConstraintsManager::EqualConstraintsVisitor(const DoEachT& DoEach) {
-  auto equals_parents = equals_.GetMap();
+  auto equals_parents = equals_.MutMap();
   for (auto it = equals_parents->begin(); it != equals_parents->end(); it++) {
     DoEach(it);
   }

Original file line number	Diff line number	Diff line change
`@@ -263,7 +263,7 @@ void ConstraintsManager::SubstituteInConstraint(const DimExpr& origin,`
`263`	`263`
`264`	`264`	`template <typename DoEachT>`
`265`	`265`	`void ConstraintsManager::EqualConstraintsVisitor(const DoEachT& DoEach) {`
`266`		`- auto equals_parents = equals_.GetMap();`
	`266`	`+ auto equals_parents = equals_.MutMap();`
`267`	`267`	`for (auto it = equals_parents->begin(); it != equals_parents->end(); it++) {`
`268`	`268`	`DoEach(it);`
`269`	`269`	`}`