[CINN]Add input shape constraint (#67384)

Hongqing-work · web-flow · commit 4dca33e8fa56 · 2024-08-21T10:48:11.000+08:00
* [CINN]Add input shape constraint

* change data structure

* fix_warning
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/convert_dynamic_to_static_dim_pass.cc
@@ -108,7 +108,7 @@ class DynamicToStaticConverter {
     VisitEachValue(fusion_op_, [&](pir::Value value) {
       updated |= UpdateValueShape(value);
     });
-    shape_analysis_->Init();
+    shape_analysis_->InitInferContext();
     return updated;
   }
 
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/local_infer_symbolic_util.cc b/paddle/cinn/hlir/dialect/operator/transforms/local_infer_symbolic_util.cc
@@ -104,7 +104,7 @@ void InitLocalShapeAnalysis(const pir::Operation& op,
 std::shared_ptr<pir::ShapeConstraintIRAnalysis> MakeOpShapeAnalysis(
     const pir::Operation* op, const DimExprs4ValueT& GraphDimExprs4Value) {
   auto shape_analysis = std::make_shared<pir::ShapeConstraintIRAnalysis>();
-  shape_analysis->Init();
+  shape_analysis->InitInferContext();
   InitLocalShapeAnalysis(*op, shape_analysis.get(), GraphDimExprs4Value);
   return shape_analysis;
 }
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.cc
@@ -197,4 +197,63 @@ void BuildCstrEqForTensorListAlongAxis(
   }
 }
 
+std::vector<symbol::DimExpr> GetSymShapeForInputValue(
+    const std::string &input_name,
+    const pir::Value &value,
+    pir::InferSymbolicShapeContext *infer_context) {
+  const common::DDim &result_dims =
+      value.type().dyn_cast<pir::DenseTensorType>().dims();
+  const auto &predefined_dim_index_to_expr = [&]() {
+    std::unordered_map<int, symbol::DimExpr> index_to_expr;
+    if (infer_context->HasPredefinedDimExprForInputName(input_name)) {
+      const auto &dim_index_and_exprs =
+          infer_context->GetPredefinedDimExprForInputName(input_name);
+      for (const auto &item : dim_index_and_exprs) {
+        index_to_expr[item.index] = item.dim_expr;
+      }
+    }
+    return index_to_expr;
+  }();
+
+  const auto &CheckStaticDimMatchConstraints =
+      [&](const symbol::DimExpr &predefined_dim_expr,
+          const int64_t &static_dim,
+          const int &dim_index) {
+        if (static_dim == -1) {
+          // no need to check
+          return;
+        }
+        PADDLE_ENFORCE_EQ(
+            infer_context->HasPredefinedRange(predefined_dim_expr),
+            false,
+            common::errors::InvalidArgument(
+                "Dim with static shape can not set range. The input value name "
+                "is %s, dim index is %d, static value is %d.",
+                input_name,
+                dim_index,
+                static_dim));
+        infer_context->AddEqualCstr(predefined_dim_expr,
+                                    symbol::DimExpr{static_dim});
+      };
+
+  const auto &GetDimExpr = [&](const int64_t &static_dim,
+                               const int &dim_index) -> symbol::DimExpr {
+    if (predefined_dim_index_to_expr.find(dim_index) !=
+        predefined_dim_index_to_expr.end()) {
+      symbol::DimExpr dim_expr = predefined_dim_index_to_expr.at(dim_index);
+      CheckStaticDimMatchConstraints(dim_expr, static_dim, dim_index);
+      return dim_expr;
+    }
+    if (static_dim != -1) {
+      return symbol::DimExpr{static_dim};
+    }
+    return symbol::DimExpr{infer_context->GetNextSymName()};
+  };
+
+  std::vector<symbol::DimExpr> result_dim_exprs;
+  for (int i = 0; i < result_dims.size(); ++i) {
+    result_dim_exprs.emplace_back(GetDimExpr(result_dims[i], i));
+  }
+  return result_dim_exprs;
+}
 }  // namespace paddle::dialect::details
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_utils.h
@@ -171,4 +171,8 @@ void BuildCstrEqForTensorListAlongAxis(
     const std::vector<pir::Value> &values,
     int axis);
 
+std::vector<symbol::DimExpr> GetSymShapeForInputValue(
+    const std::string &input_name,
+    const pir::Value &value,
+    pir::InferSymbolicShapeContext *infer_context);
 }  // namespace paddle::dialect::details
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc
@@ -105,26 +105,10 @@ bool AssignValue_OpInferSymbolicShape(
 
 bool DataOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
-  const auto &attributes = op->attributes();
-  pir::Attribute attr = attributes.at("shape");
-
-  const std::vector<symbol::DimExpr> sym_dims = [&] {
-    std::vector<symbol::DimExpr> sym_dims;
-    const std::vector<int64_t> &dims =
-        attr.dyn_cast<paddle::dialect::IntArrayAttribute>().data().GetData();
-    for (auto dim : dims) {
-      symbol::DimExpr dim_expr;
-      if (dim == pir::ShapedTypeInterface::kDynamic) {
-        symbol::DimExpr symbolic_dim_expr(infer_context->GetNextSymName());
-        dim_expr = symbolic_dim_expr;
-      } else {
-        symbol::DimExpr numeric_dim_expr(dim);
-        dim_expr = numeric_dim_expr;
-      }
-      sym_dims.push_back(dim_expr);
-    }
-    return sym_dims;
-  }();
+  std::string name =
+      op->attributes().at("name").dyn_cast<pir::StrAttribute>().AsString();
+  const std::vector<symbol::DimExpr> sym_dims =
+      details::GetSymShapeForInputValue(name, op->result(0), infer_context);
 
   auto IsNumelLEKMaxRank = [](pir::Value value) {
     const auto &dims = value.type().dyn_cast<pir::DenseTensorType>().dims();
@@ -204,21 +188,12 @@ bool EmptyOpInferSymbolicShape(pir::Operation *op,
 
 bool FeedOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
-  const common::DDim &result_dims =
-      op->result(0).type().dyn_cast<pir::DenseTensorType>().dims();
-  std::vector<symbol::DimExpr> out_dims;
-  for (int i = 0; i < result_dims.size(); i++) {
-    if (result_dims[i] == -1) {
-      out_dims.emplace_back(infer_context->GetNextSymName());
-    } else {
-      out_dims.emplace_back(result_dims[i]);
-    }
-  }
-
+  std::string name =
+      op->attributes().at("name").dyn_cast<pir::StrAttribute>().AsString();
+  const auto &symbolic_shape =
+      details::GetSymShapeForInputValue(name, op->result(0), infer_context);
   infer_context->SetShapeOrDataForValue(
-      op->result(0),
-      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(out_dims)});
-
+      op->result(0), symbol::TensorShapeOrDataDimExprs(symbolic_shape));
   return true;
 }
 
diff --git a/paddle/pir/include/dialect/shape/utils/constraints_manager.h b/paddle/pir/include/dialect/shape/utils/constraints_manager.h
@@ -39,9 +39,23 @@ class IR_API ConstraintsManager {
 
   bool IsBroadcastable(const DimExpr& lhs, const DimExpr& rhs) const;
 
+  struct Range {
+    std::int64_t min;
+    std::int64_t max;
+    // TODO(Hongqing-work): Subsitute INT32_MAX with a more meaningful value.
+    Range() : min(1), max(INT32_MAX) {}
+    Range(int min_val, int max_val) : min(min_val), max(max_val) {}
+  };
+  void AddInputRangeCstr(const DimExpr& dim_expr, const Range& range);
+
+  bool IsBoundedInput(const DimExpr& dim_expr) const;
+
+  const Range& GetRangeOfBoundedInput(const DimExpr& dim_expr) const;
+
   using EqualConstraints = common::UnionFindSet<DimExpr>;
   using GTOneConstraints = std::unordered_set<DimExpr>;
   using BroadcastableConstraints = std::unordered_set<Broadcastable<DimExpr>>;
+  using InputRangeConstraints = std::unordered_map<DimExpr, Range>;
 
   void VisitEqualClusters(
       const std::function<void(const std::vector<DimExpr>&)>& DoEachCluster)
@@ -65,6 +79,13 @@ class IR_API ConstraintsManager {
       const std::function<void(BroadcastableConstraints::const_iterator)>&
           DoEach) const;
 
+  void InputRangeConstraintsVisitor(
+      const std::function<void(InputRangeConstraints::iterator)>& DoEach);
+
+  void InputRangeConstraintsVisitor(
+      const std::function<void(InputRangeConstraints::const_iterator)>& DoEach)
+      const;
+
   using EqualCallbackFunc = std::function<void(const DimExpr&, const DimExpr&)>;
   void SetEqualCallbackFunc(EqualCallbackFunc equal_callback_func);
 
@@ -76,6 +97,8 @@ class IR_API ConstraintsManager {
     return broadcastables_;
   }
 
+  const InputRangeConstraints& input_ranges() const { return input_ranges_; }
+
  private:
   void SubstituteInConstraint(const DimExpr& lhs, const DimExpr& rhs);
 
@@ -85,6 +108,7 @@ class IR_API ConstraintsManager {
   EqualConstraints equals_;
   GTOneConstraints gtones_;
   BroadcastableConstraints broadcastables_;
+  InputRangeConstraints input_ranges_;
 };
 
 std::ostream& operator<<(std::ostream& os,
diff --git a/paddle/pir/include/dialect/shape/utils/shape_analysis.h b/paddle/pir/include/dialect/shape/utils/shape_analysis.h
@@ -56,6 +56,12 @@ class IR_API InferSymbolicShapeCacheKey {
   void SetInputShapeOrDatas(
       const std::vector<symbol::ShapeOrDataDimExprs>& input_shape_or_datas);
 };
+struct ConstraintsForInputDimExpr {
+  symbol::DimExpr dim_expr;
+  // bind_info = [(input_name, dim_index)]
+  std::vector<std::pair<std::string, int>> bind_info;
+  symbol::ConstraintsManager::Range range;
+};
 }  // namespace pir
 
 namespace std {
@@ -75,7 +81,8 @@ class IR_API InferSymbolicShapeContext {
   InferSymbolicShapeContext() = default;
   InferSymbolicShapeContext(const InferSymbolicShapeContext&) = delete;
   InferSymbolicShapeContext(InferSymbolicShapeContext&&) = delete;
-  void Init();
+  void Init(
+      const std::vector<ConstraintsForInputDimExpr>& input_shape_constraints);
 
   // Note: Only initialize the symbol info, the value info is not update.
   void RegisterSymbolConstraintFromContext(
@@ -106,6 +113,8 @@ class IR_API InferSymbolicShapeContext {
   bool IsBroadcastable(const symbol::DimExpr& lhs,
                        const symbol::DimExpr& rhs) const;
 
+  bool HasPredefinedRange(const symbol::DimExpr& dim_expr) const;
+
   void PrintShapeOrDatas() const;
 
   void SetOpInferSymbolicShapeCache(
@@ -119,6 +128,18 @@ class IR_API InferSymbolicShapeContext {
     return constraints_manager_;
   }
 
+  struct DimIndexAndExpr {
+    int index;
+    symbol::DimExpr dim_expr;
+    DimIndexAndExpr(int index_val, const symbol::DimExpr& dim_expr_val)
+        : index(index_val), dim_expr(dim_expr_val) {}
+  };
+
+  bool HasPredefinedDimExprForInputName(const std::string& input_name) const;
+
+  const std::vector<DimIndexAndExpr> GetPredefinedDimExprForInputName(
+      const std::string& input_name) const;
+
  private:
   symbol::ShapeOrDataDimExprs SimplifyBroadcastForShapeOrData(
       const symbol::ShapeOrDataDimExprs& shape_or_data);
@@ -140,6 +161,9 @@ class IR_API InferSymbolicShapeContext {
 
   std::unordered_map<InferSymbolicShapeCacheKey, InferSymbolicShapeCacheValue>
       infer_symbolic_shape_cache_;
+
+  std::unordered_map<std::string, std::vector<DimIndexAndExpr>>
+      predefined_dimexpr_map_for_inputs_;
 };
 
 class IR_API ShapeConstraintIRAnalysis final
@@ -148,7 +172,7 @@ class IR_API ShapeConstraintIRAnalysis final
   ShapeConstraintIRAnalysis() = default;
   ShapeConstraintIRAnalysis(const ShapeConstraintIRAnalysis&) = delete;
   ShapeConstraintIRAnalysis(ShapeConstraintIRAnalysis&&) = delete;
-  void Init();
+  void InitInferContext();
 
   void RegisterSymbolConstraintFromShapeAnalysis(
       const ShapeConstraintIRAnalysis& other);
@@ -202,6 +226,11 @@ class IR_API ShapeConstraintIRAnalysis final
     return context_.constraints_manager();
   }
 
+  void SetInputShapeConstraints(
+      const std::vector<ConstraintsForInputDimExpr>& input_shape_constraints) {
+    input_shape_constraints_ = input_shape_constraints;
+  }
+
  private:
   InferSymbolicShapeContext* MutInferSymbolicShapeContext() {
     return &context_;
@@ -215,6 +244,7 @@ class IR_API ShapeConstraintIRAnalysis final
 
  private:
   InferSymbolicShapeContext context_;
+  std::vector<ConstraintsForInputDimExpr> input_shape_constraints_;
 };
 
 class IR_API ShapeAnalysisManager {
diff --git a/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc b/paddle/pir/src/dialect/shape/transforms/shape_optimization_pass.cc
@@ -399,7 +399,7 @@ void InferSymExprForAllValues(ModuleOp module_op) {
         return symbol_shape_map;
       }();
 
-  shape_analysis.Init();
+  shape_analysis.InitInferContext();
   // init the kwarg symbol shape info
   for (const auto& kv : symbol_shape_map) {
     infer_context->SetShapeOrDataForValue(kv.first, kv.second);
diff --git a/paddle/pir/src/dialect/shape/utils/constraints_manager.cc b/paddle/pir/src/dialect/shape/utils/constraints_manager.cc
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
`@@ -108,7 +108,7 @@ class DynamicToStaticConverter {`
`108`	`108`	`VisitEachValue(fusion_op_, [&](pir::Value value) {`
`109`	`109`	`updated \|= UpdateValueShape(value);`
`110`	`110`	`});`
`111`		`- shape_analysis_->Init();`
	`111`	`+ shape_analysis_->InitInferContext();`
`112`	`112`	`return updated;`
`113`	`113`	`}`
`114`	`114`
Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,7 @@ void InitLocalShapeAnalysis(const pir::Operation& op,`
`104`	`104`	`std::shared_ptr<pir::ShapeConstraintIRAnalysis> MakeOpShapeAnalysis(`
`105`	`105`	`const pir::Operation* op, const DimExprs4ValueT& GraphDimExprs4Value) {`
`106`	`106`	`auto shape_analysis = std::make_shared<pir::ShapeConstraintIRAnalysis>();`
`107`		`- shape_analysis->Init();`
	`107`	`+ shape_analysis->InitInferContext();`
`108`	`108`	`InitLocalShapeAnalysis(*op, shape_analysis.get(), GraphDimExprs4Value);`
`109`	`109`	`return shape_analysis;`
`110`	`110`	`}`