add backtrack for GetShapeOrDataForValue (#63367)

Hongqing-work · web-flow · commit bd9a547c7d51 · 2024-04-19T19:24:58.000+08:00
* add GetShapeOrDataForValue backtrack

* restore for insert_broadcast_pass

* fix merge bug

* add check for no value impl

* change recursion to iteration

* fix walker include dir after merge

* add lock for InferShapeOrDataForValue

* fix bug for static mutex

* delete lock and fix name

* fix merge bug
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
@@ -37,10 +37,15 @@ DimExprs4ValueT MakeDimExprs4Value(
   std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
   pass_manager->AddPass(pir::CreateShapeOptimizationPass());
   pass_manager->Run(program);
-  const auto* shape_analysis =
-      &pir::ShapeAnalysisManager::Instance().Get(program);
+  auto* shape_analysis = &pir::ShapeAnalysisManager::Instance().Get(program);
   return
       [shape_analysis](pir::Value value) -> const symbol::ShapeOrDataDimExprs& {
+        // TODO(Hongqing-work): define a default empty ShapeOrDataDimExprss
+        if (!value) {
+          static symbol::ShapeOrDataDimExprs empty{
+              symbol::TensorShapeOrDataDimExprs{}};
+          return empty;
+        }
         return shape_analysis->GetShapeOrDataForValue(value);
       };
 }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
@@ -236,10 +236,8 @@ int GetSharedSize(::pir::Operation* op) {
   return 0;
 }
 
-using ConditionFunction =
-    std::function<bool(::pir::Operation*,
-                       const GroupPtr&,
-                       const ::pir::ShapeConstraintIRAnalysis&)>;
+using ConditionFunction = std::function<bool(
+    ::pir::Operation*, const GroupPtr&, ::pir::ShapeConstraintIRAnalysis*)>;
 
 // Op Fusion Pass which performs Ops fusion, Ops are fused
 // "vertically", meaning producing Ops are fused into their consumers
@@ -354,7 +352,7 @@ class OpFusionPassHelper {
 
  private:
   void DoOpFusion() {
-    const auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(
+    auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(
         ops_.front()->GetParentProgram());
     for (auto consumer : ops_) {
       auto consumer_kind =
@@ -413,7 +411,7 @@ class OpFusionPassHelper {
           }
         }
 
-        if (!can_fuse || !CanFuse(producer, consumer, shape_analysis)) {
+        if (!can_fuse || !CanFuse(producer, consumer, &shape_analysis)) {
           continue;
         }
 
@@ -443,7 +441,7 @@ class OpFusionPassHelper {
           // VLOG(3) << "Insert Global Output Node : " << producer->id();
           consumer_fusion->output_ops.insert(producer);
         } else if (producer_data_used_num > 1 && producer->num_operands() > 0 &&
-                   is_same_size(producer, consumer_fusion, shape_analysis)) {
+                   is_same_size(producer, consumer_fusion, &shape_analysis)) {
           // producer is not a const value op.
           consumer_fusion->internal_ops.insert(producer);
         }
@@ -484,7 +482,7 @@ class OpFusionPassHelper {
           {OpPatternKind::kBroadcast,
            [](::pir::Operation* producer,
               const GroupPtr& consumer,
-              const ::pir::ShapeConstraintIRAnalysis& shape_analysis) -> bool {
+              ::pir::ShapeConstraintIRAnalysis* shape_analysis) -> bool {
              // NOTE, producer and consumer NEVER be same size
              if (is_same_size(producer, consumer, shape_analysis)) {
                return true;
@@ -598,7 +596,7 @@ class OpFusionPassHelper {
 
   bool CanFuse(::pir::Operation* producer,
                const ::pir::Operation* consumer,
-               const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+               ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
     auto& relation =
         fusion_relation_map_[hlir::framework::pir::CompatibleInfo::OpKind(
             *producer)];
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h
@@ -77,32 +77,30 @@ int GetSharedSize(::pir::Operation* op);
 inline bool always_fuse(
     ::pir::Operation* producer,
     const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {  // NOLINT
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {  // NOLINT
   return true;
 }
 
 inline bool no_fuse(::pir::Operation* producer,
                     const std::shared_ptr<Group>& consumer,
-                    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+                    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   return false;
 }
 
-inline bool is_same_shape(
-    ::pir::Operation* producer,
-    const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+inline bool is_same_shape(::pir::Operation* producer,
+                          const std::shared_ptr<Group>& consumer,
+                          ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   auto master_op = consumer->master_ops.begin();
-  return shape_analysis.IsShapeEqual(producer->result(0),
-                                     (*master_op)->result(0));
+  return shape_analysis->IsShapeEqual(producer->result(0),
+                                      (*master_op)->result(0));
 }
 
-inline bool is_same_size(
-    ::pir::Operation* producer,
-    const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+inline bool is_same_size(::pir::Operation* producer,
+                         const std::shared_ptr<Group>& consumer,
+                         ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   auto master_op = consumer->master_ops.begin();
-  return shape_analysis.IsSameNumel(producer->result(0),
-                                    (*master_op)->result(0));
+  return shape_analysis->IsSameNumel(producer->result(0),
+                                     (*master_op)->result(0));
 }
 
 inline bool without_last_dimension_in_reduce(
@@ -115,7 +113,7 @@ inline bool without_last_dimension_in_reduce(
 inline bool reduce_fuse_reduce(
     ::pir::Operation* producer,
     const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   ::pir::Operation* reducer = NULL;
   for (auto* master : consumer->master_ops) {
     if (hlir::framework::pir::CompatibleInfo::OpKind(*master) ==
@@ -227,7 +225,7 @@ inline bool is_horizontal_relation(::pir::Operation* producer,
 inline bool horizontal_or_vertical_reduce_relation(
     ::pir::Operation* producer,
     const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   // check is same shape with horizontal relation.
   if (is_same_size(producer, consumer, shape_analysis)) {
     return true;
@@ -298,7 +296,7 @@ inline bool horizontal_or_vertical_reduce_relation(
 inline bool horizontal_or_can_inline(
     ::pir::Operation* producer,
     const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   // horizontal relation.
   if (is_horizontal_relation(producer, consumer)) {
     if (is_same_size(producer, consumer, shape_analysis)) {
@@ -336,22 +334,22 @@ inline bool horizontal_or_can_inline(
 inline bool horizontal_with_same_size(
     ::pir::Operation* producer,
     const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   return is_horizontal_relation(producer, consumer) &&
          is_same_size(producer, consumer, shape_analysis);
 }
 
 inline std::vector<int64_t> GetBroadcastAxes(
     ::pir::Operation* bcast_op,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {  // NOLINT
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {  // NOLINT
   if (bcast_op->isa<cinn::dialect::BroadcastOp>()) {
     return GetVectorAttr(bcast_op, "broadcast_axes");
   } else if (bcast_op->isa<paddle::dialect::ExpandOp>()) {
     const auto& input_shape =
-        shape_analysis.GetShapeOrDataForValue(bcast_op->operand_source(0))
+        shape_analysis->GetShapeOrDataForValue(bcast_op->operand_source(0))
             .shape();
     const auto& output_shape =
-        shape_analysis.GetShapeOrDataForValue(bcast_op->result(0)).shape();
+        shape_analysis->GetShapeOrDataForValue(bcast_op->result(0)).shape();
     std::vector<int64_t> broadcast_axes(input_shape.size(), 0);
     size_t index_gap = output_shape.size() - input_shape.size();
     for (size_t i = 0; i < input_shape.size(); ++i) {
@@ -366,7 +364,7 @@ inline std::vector<int64_t> GetBroadcastAxes(
 inline bool reduce_fuse_broadcast(
     ::pir::Operation* producer,
     const std::shared_ptr<Group>& consumer,
-    const ::pir::ShapeConstraintIRAnalysis& shape_analysis) {
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
   if (is_horizontal_relation(producer, consumer)) {
     if (is_same_size(producer, consumer, shape_analysis)) {
       return true;
@@ -379,7 +377,7 @@ inline bool reduce_fuse_broadcast(
   // }
 
   const auto& rinput_shape =
-      shape_analysis.GetShapeOrDataForValue(producer->operand_source(0))
+      shape_analysis->GetShapeOrDataForValue(producer->operand_source(0))
           .shape();
   auto reduce_axes = GetVectorAttr(producer, "dim");
   auto keep_dim = producer->attributes()
@@ -429,7 +427,7 @@ inline bool reduce_fuse_broadcast(
       continue;
     }
     const auto& broadcast_shape =
-        shape_analysis.GetShapeOrDataForValue(op->result(0)).shape();
+        shape_analysis->GetShapeOrDataForValue(op->result(0)).shape();
     auto broadcast_axes = GetBroadcastAxes(op, shape_analysis);
 
     for (auto& axis : broadcast_axes) {
diff --git a/paddle/cinn/hlir/framework/pir/utils.cc b/paddle/cinn/hlir/framework/pir/utils.cc
@@ -288,7 +288,7 @@ bool IsSmallNumelOp(const ::pir::Operation& op) {
 }
 
 bool IsShapeComputeOp(const ::pir::Operation& op) {
-  const auto& shape_analysis = ::pir::ShapeAnalysisManager::Instance().Get(
+  auto& shape_analysis = ::pir::ShapeAnalysisManager::Instance().Get(
       op.GetParent()->parent_program());
   if (op.num_operands() == 0) {
     return false;
diff --git a/paddle/cinn/operator_fusion/group_cluster.h b/paddle/cinn/operator_fusion/group_cluster.h
@@ -58,8 +58,7 @@ inline std::vector<fusion::PatternNodePtr<T>> ClusterOps(
 
   pir::Program* program = ops.at(0)->GetParentProgram();
 
-  const auto* shape_analysis =
-      &pir::ShapeAnalysisManager::Instance().Get(program);
+  auto* shape_analysis = &pir::ShapeAnalysisManager::Instance().Get(program);
 
   VLOG(4) << "Start Create Policies and PolicyManager!";
   const auto& relative_judge_policy =
diff --git a/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc b/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc
@@ -259,7 +259,7 @@ symbol::DimExpr GetProductDimExprForValueDims(
   for (const auto& dim : dims) {
     dim_idx.emplace_back(dim.idx_);
   }
-  const auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(
+  auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(
       dims[0].v_.defining_op()->GetParentProgram());
   return shape_analysis.GetProductDimExpr(dims[0].v_, dim_idx);
 }
diff --git a/paddle/cinn/operator_fusion/policy/relative_judge_policy.h b/paddle/cinn/operator_fusion/policy/relative_judge_policy.h
@@ -101,7 +101,7 @@ static ValueDimRelation CreateOpRelativenessForElementWise(pir::Operation* op) {
 static std::vector<std::pair<size_t, size_t>> GetNonBroadCastDims(
     pir::Operation* op) {
   std::vector<std::pair<size_t, size_t>> res;
-  const auto* shape_analysis =
+  auto* shape_analysis =
       &pir::ShapeAnalysisManager::Instance().Get(op->GetParentProgram());
 
   const auto& broad_cast_value = GetBroadcastOpInputOuputValue(op);
@@ -262,7 +262,7 @@ template <typename T>
 class RelativeJudgePolicy final : public Policy<T> {
  public:
   RelativeJudgePolicy(const std::vector<pir::Operation*>& ops,
-                      const pir::ShapeConstraintIRAnalysis* shape_analysis)
+                      pir::ShapeConstraintIRAnalysis* shape_analysis)
       : axes_info_(ops, shape_analysis) {
     VLOG(4) << "[relative_judge_policy] Start AnalysisIndexExprRelation.";
     index_expr_map_ = AnalysisIndexExprRelation(ops);
diff --git a/paddle/cinn/operator_fusion/policy/shardable_axes_base.cc b/paddle/cinn/operator_fusion/policy/shardable_axes_base.cc
@@ -138,7 +138,7 @@ ShardableAxesSignature CreateSignatureForElementWise(pir::Operation* op) {
 }
 
 ShardableAxesSignature CreateSignatureForBroadcast(
-    pir::Operation* op, const pir::ShapeConstraintIRAnalysis* shape_analysis) {
+    pir::Operation* op, pir::ShapeConstraintIRAnalysis* shape_analysis) {
   ShardableAxesSignature result = ShardableAxesSignature();
 
   const auto& broad_cast_value = GetBroadcastOpInputOuputValue(op);
@@ -207,7 +207,7 @@ ShardableAxesSignature ShardableAxesInfoManager::CreateShardableSignature(
 
 ShardableAxesInfoManager::ShardableAxesInfoManager(
     const std::vector<pir::Operation*>& ops,
-    const pir::ShapeConstraintIRAnalysis* shape_analysis)
+    pir::ShapeConstraintIRAnalysis* shape_analysis)
     : ops_(ops), shape_analysis_(shape_analysis) {
   for (const auto& op : ops) {
     if (op->name() == "cf.yield") continue;
diff --git a/paddle/cinn/operator_fusion/policy/shardable_axes_base.h b/paddle/cinn/operator_fusion/policy/shardable_axes_base.h
@@ -33,9 +33,8 @@ struct ShardableAxesSignature {
 };
 
 struct ShardableAxesInfoManager {
-  ShardableAxesInfoManager(
-      const std::vector<pir::Operation*>& ops,
-      const pir::ShapeConstraintIRAnalysis* shape_analysis);
+  ShardableAxesInfoManager(const std::vector<pir::Operation*>& ops,
+                           pir::ShapeConstraintIRAnalysis* shape_analysis);
   ShardableAxesSignature GetSignature(pir::Operation* op);
   ShardableAxes GetAxes(pir::Value value);
   ShardableAxesSignature CreateShardableSignature(pir::Operation* op);
@@ -45,7 +44,7 @@ struct ShardableAxesInfoManager {
 
  private:
   const std::vector<pir::Operation*>& ops_;
-  const pir::ShapeConstraintIRAnalysis* shape_analysis_;
+  pir::ShapeConstraintIRAnalysis* shape_analysis_;
 
   std::unordered_map<pir::Operation*, ShardableAxesSignature> op_signature_map_;
   std::unordered_map<pir::Value, ShardableAxes> value_axes_map_;
diff --git a/paddle/cinn/operator_fusion/policy/shardable_axes_policy.h b/paddle/cinn/operator_fusion/policy/shardable_axes_policy.h
@@ -22,8 +22,8 @@ template <typename T>
 class ShardableAxesRRFusePolicy final : public Policy<T> {
  public:
   ShardableAxesRRFusePolicy(
-      const std::vector<pir::Operation*>& ops,               // NOLINT
-      const pir::ShapeConstraintIRAnalysis* shape_analysis)  // NOLINT
+      const std::vector<pir::Operation*>& ops,         // NOLINT
+      pir::ShapeConstraintIRAnalysis* shape_analysis)  // NOLINT
       : axes_info_(ops, shape_analysis) {}
   bool CanFuse(const PatternNodePtr<T>& upstream,
                const PatternNodePtr<T>& downstream) override;
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
@@ -486,6 +486,10 @@ bool ReshapeOpInferSymbolicShape(
           op->result(0),
           symbol::TensorShapeOrDataDimExprs(shape_data,
                                             x_dim_expr.data().value()));
+      shape_analysis->SetShapeOrDataForValue(
+          op->result(1),
+          CreateShapeOrDataForXShape(
+              shape_analysis->GetShapeOrDataForValue(op->operand_source(0))));
       return true;
     }
   }
diff --git a/paddle/pir/include/dialect/shape/utils/shape_analysis.h b/paddle/pir/include/dialect/shape/utils/shape_analysis.h
@@ -35,7 +35,9 @@ class IR_API ShapeConstraintIRAnalysis {
 
   bool HasShapeOrDataForValue(Value val) const;
 
-  const symbol::ShapeOrDataDimExprs& GetShapeOrDataForValue(Value val) const;
+  void InferShapeOrDataForValue(Value val);
+
+  const symbol::ShapeOrDataDimExprs& GetShapeOrDataForValue(Value val);
 
   void SetShapeOrDataForValue(Value val,
                               const symbol::ShapeOrDataDimExprs& shape_or_data);
@@ -58,7 +60,7 @@ class IR_API ShapeConstraintIRAnalysis {
   void PrintShapeOrDatas() const;
 
   // Returns true if the two value have the same symbolic shape.
-  bool IsShapeEqual(Value lhs, Value rhs) const;
+  bool IsShapeEqual(Value lhs, Value rhs);
 
   // Suppose:
   //    lhs_dim_idxs = {ld0, ld1, ...}
@@ -69,25 +71,21 @@ class IR_API ShapeConstraintIRAnalysis {
   bool IsProductEqual(Value lhs,
                       const std::vector<int>& lhs_dim_idxs,
                       Value rhs,
-                      const std::vector<int>& rhs_dim_idxs) const;
+                      const std::vector<int>& rhs_dim_idxs);
 
   // Returns true if:
   //    lhs.shape[lhs_from] * ... lhs.shape[lhs_to-1] ==
   //    rhs.shape[rhs_from] * ... rhs.shape[rhs_to-1]
-  bool IsProductEqual(Value lhs,
-                      int lhs_from,
-                      int lhs_to,
-                      Value rhs,
-                      int rhs_from,
-                      int rhs_to) const;
+  bool IsProductEqual(
+      Value lhs, int lhs_from, int lhs_to, Value rhs, int rhs_from, int rhs_to);
 
   // Returns true if the two value have the same number elements.
-  bool IsSameNumel(Value lhs, Value rhs) const;
+  bool IsSameNumel(Value lhs, Value rhs);
 
-  pir::PrintHooks PrintHook() const;
+  pir::PrintHooks PrintHook();
 
   symbol::DimExpr GetProductDimExpr(Value lhs,
-                                    const std::vector<int>& lhs_dim_idxs) const;
+                                    const std::vector<int>& lhs_dim_idxs);
 
  private:
   void SubstituteDimExpr(const symbol::DimExpr& origin,
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc

Original file line number	Diff line number	Diff line change
`@@ -288,7 +288,7 @@ bool IsSmallNumelOp(const ::pir::Operation& op) {`
`288`	`288`	`}`
`289`	`289`
`290`	`290`	`bool IsShapeComputeOp(const ::pir::Operation& op) {`
`291`		`- const auto& shape_analysis = ::pir::ShapeAnalysisManager::Instance().Get(`
	`291`	`+ auto& shape_analysis = ::pir::ShapeAnalysisManager::Instance().Get(`
`292`	`292`	`op.GetParent()->parent_program());`
`293`	`293`	`if (op.num_operands() == 0) {`
`294`	`294`	`return false;`
Original file line number	Diff line number	Diff line change
`@@ -259,7 +259,7 @@ symbol::DimExpr GetProductDimExprForValueDims(`
`259`	`259`	`for (const auto& dim : dims) {`
`260`	`260`	`dim_idx.emplace_back(dim.idx_);`
`261`	`261`	`}`
`262`		`- const auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(`
	`262`	`+ auto& shape_analysis = pir::ShapeAnalysisManager::Instance().Get(`
`263`	`263`	`dims[0].v_.defining_op()->GetParentProgram());`
`264`	`264`	`return shape_analysis.GetProductDimExpr(dims[0].v_, dim_idx);`
`265`	`265`	`}`