chore: import upstream snapshot with attribution

paddle/pir/include/dialect/shape/utils/constraints_manager.h

@@ -0,0 +1,117 @@
+// 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
+//
+//     http://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 <unordered_set>
+
+#include "paddle/common/union_find_set.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
+
+namespace symbol {
+
+class IR_API ConstraintsManager {
+ public:
+  ConstraintsManager() = default;
+  ConstraintsManager(const ConstraintsManager&) = delete;
+  ConstraintsManager(ConstraintsManager&&) = delete;
+
+  void AddEqCstr(const DimExpr& lhs, const DimExpr& rhs);
+
+  bool IsEqual(const DimExpr& lhs, const DimExpr& rhs) const;
+
+  void AddGTOneCstr(const DimExpr& dim_expr);
+
+  bool IsGTOne(const DimExpr& dim_expr) const;
+
+  void AddBroadcastableCstr(const DimExpr& lhs, const DimExpr& rhs);
+
+  bool IsBroadcastable(const DimExpr& lhs, const DimExpr& rhs) const;
+
+  struct Range {
+    std::int64_t min;
+    std::int64_t max;
+    // TODO(Hongqing-work): Substitute 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)
+      const;
+
+  void EqualConstraintsVisitor(
+      const std::function<void(std::unordered_map<DimExpr, DimExpr>::iterator)>&
+          DoEach);
+
+  void GTOneConstraintsVisitor(
+      const std::function<void(GTOneConstraints::iterator)>& DoEach);
+
+  void GTOneConstraintsVisitor(
+      const std::function<void(GTOneConstraints::const_iterator)>& DoEach)
+      const;
+
+  void BroadcastableConstraintsVisitor(
+      const std::function<void(BroadcastableConstraints::iterator)>& DoEach);
+
+  void BroadcastableConstraintsVisitor(
+      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);
+
+  const EqualConstraints& equals() const { return equals_; }
+
+  const GTOneConstraints& gtones() const { return gtones_; }
+
+  const BroadcastableConstraints& broadcastables() const {
+    return broadcastables_;
+  }
+
+  const InputRangeConstraints& input_ranges() const { return input_ranges_; }
+
+ private:
+  void SubstituteInConstraint(const DimExpr& lhs, const DimExpr& rhs);
+
+ private:
+  EqualCallbackFunc equal_callback_func_ = nullptr;
+
+  EqualConstraints equals_;
+  GTOneConstraints gtones_;
+  BroadcastableConstraints broadcastables_;
+  InputRangeConstraints input_ranges_;
+};
+
+std::ostream& operator<<(std::ostream& os,
+                         const ConstraintsManager& constraints_manager);
+
+}  // namespace symbol

paddle/pir/include/dialect/shape/utils/dim_expr.h

@@ -0,0 +1,272 @@
+// Copyright (c) 2023 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
+//
+//     http://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 <functional>
+#include <memory>
+#include <optional>
+#include <ostream>
+#include <string>
+#include <variant>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/common/enforce.h"
+#include "paddle/common/overloaded.h"
+#include "paddle/pir/include/core/dll_decl.h"
+#include "paddle/pir/include/core/utils.h"
+
+namespace symbol {
+
+#define SYMBOL_NOT_IMPLEMENTED \
+  PADDLE_THROW(common::errors::Unimplemented("Not Implemented"))
+
+template <typename T>
+struct UnaryDimExpr {
+  explicit UnaryDimExpr(const T& d) : data(std::make_shared<Data>(d)) {}
+  struct Data {
+    explicit Data(const T& d) : data(d) {}
+    T data;
+  };
+
+  const Data& operator*() const { return *data; }
+  Data& operator*() { return *data; }
+  const Data* operator->() const { return data.get(); }
+  Data* operator->() { return data.get(); }
+
+  std::shared_ptr<Data> data;
+};
+
+template <typename T>
+struct BinaryDimExpr {
+  explicit BinaryDimExpr(const T& l, const T& r)
+      : data(std::make_shared<Data>(l, r)) {}
+
+  struct Data {
+    explicit Data(const T& l, const T& r) : lhs(l), rhs(r) {}
+    T lhs;
+    T rhs;
+  };
+
+  const Data& operator*() const { return *data; }
+  Data& operator*() { return *data; }
+  const Data* operator->() const { return data.get(); }
+  Data* operator->() { return data.get(); }
+
+  std::shared_ptr<Data> data;
+};
+
+template <typename T>
+class List final {
+ public:
+  List(const List&) = default;
+  List(List&&) = default;
+  List& operator=(const List&) = default;
+  List& operator=(List&&) = default;
+
+  using value_type = T;
+
+  explicit List() : vector_(std::make_shared<std::vector<T>>()) {}
+
+  template <
+      typename Arg,
+      std::enable_if_t<!std::is_same_v<std::decay_t<Arg>, List>, bool> = true>
+  explicit List(Arg&& arg)
+      : vector_(std::make_shared<std::vector<T>>(
+            std::vector<T>{std::forward<Arg>(arg)})) {}
+
+  template <typename Arg0, typename Arg1, typename... Args>
+  List(Arg0&& arg0, Arg1&& arg1, Args&&... args)
+      : vector_(std::make_shared<std::vector<T>>(
+            std::vector<T>{std::forward<Arg0>(arg0),
+                           std::forward<Arg1>(arg1),
+                           std::forward<Args>(args)...})) {}
+
+  bool operator==(const List& other) const {
+    if (&vector() == &other.vector()) {
+      return true;
+    }
+    return vector() == other.vector();
+  }
+
+  bool operator!=(const List& other) const { return !(*this == other); }
+
+  std::vector<T>& operator*() const { return *vector_; }
+  std::vector<T>* operator->() const { return vector_.get(); }
+
+  const std::vector<T>& vector() const { return *vector_; }
+
+  const auto& Get(std::size_t idx) const { return vector_->at(idx); }
+
+ private:
+  std::shared_ptr<std::vector<T>> vector_;
+};
+
+template <typename T>
+struct Negative final : public UnaryDimExpr<T> {
+  using UnaryDimExpr<T>::UnaryDimExpr;
+};
+
+template <typename T>
+struct Add final {
+  List<T> operands;
+};
+
+template <typename T>
+struct Mul final {
+  List<T> operands;
+};
+
+template <typename T>
+struct Div final : public BinaryDimExpr<T> {
+  using BinaryDimExpr<T>::BinaryDimExpr;
+};
+
+template <typename T>
+struct Max final {
+  List<T> operands;
+};
+
+template <typename T>
+struct Min final {
+  List<T> operands;
+};
+
+template <typename T>
+struct Broadcast final {
+  List<T> operands;
+};
+
+template <typename T>
+struct Equal final : public BinaryDimExpr<T> {
+  using BinaryDimExpr<T>::BinaryDimExpr;
+};
+
+template <typename T>
+struct Broadcastable final : public BinaryDimExpr<T> {
+  using BinaryDimExpr<T>::BinaryDimExpr;
+  bool operator==(const Broadcastable& other) const {
+    return (this->data->lhs == other.data->lhs &&
+            this->data->rhs == other.data->rhs) ||
+           (this->data->lhs == other.data->rhs &&
+            this->data->rhs == other.data->lhs);
+  }
+};
+
+class DimExpr;
+
+// DimExpr = std::int64_t
+//         | std::string
+//         | Negative DimExpr
+//         | Add DimExpr
+//         | Mul DimExpr
+//         | Div DimExpr
+//         | Max DimExpr
+//         | Min DimExpr
+//         | Broadcast DimExpr
+using DimExprBase = std::variant<std::int64_t,
+                                 std::string,
+                                 Negative<DimExpr>,
+                                 Add<DimExpr>,
+                                 Mul<DimExpr>,
+                                 Div<DimExpr>,
+                                 Max<DimExpr>,
+                                 Min<DimExpr>,
+                                 Broadcast<DimExpr>>;
+
+class IR_API DimExpr : public DimExprBase {
+ public:
+  using DimExprBase::DimExprBase;
+
+  template <typename T>
+  bool isa() const {
+    return std::holds_alternative<T>(*this);
+  }
+
+  template <typename T>
+  const T& dyn_cast() const {
+    return std::get<T>(*this);
+  }
+
+  template <typename T>
+  bool Has() const {
+    return std::holds_alternative<T>(*this);
+  }
+
+  template <typename T>
+  const T& Get() const {
+    return std::get<T>(*this);
+  }
+
+  const DimExprBase& variant() const {
+    return static_cast<const DimExprBase&>(*this);
+  }
+
+  DEFINE_MATCH_METHOD();
+
+  DimExpr operator+(const DimExpr& other) const;
+  DimExpr operator-(const DimExpr& other) const;
+  DimExpr operator*(const DimExpr& other) const;
+  DimExpr operator/(const DimExpr& other) const;
+  bool operator==(const DimExpr& other) const;
+  bool operator!=(const DimExpr& other) const;
+};
+
+// DimExprConstraint = Equal DimExpr
+//                   | Broadcastable DimExpr
+using DimExprConstraint = std::variant<Equal<DimExpr>, Broadcastable<DimExpr>>;
+
+IR_API std::string ToString(const DimExpr& dim_expr);
+
+IR_API std::ostream& operator<<(std::ostream&, const DimExpr& dim_expr);
+
+IR_API std::ostream& operator<<(std::ostream&,
+                                const std::vector<DimExpr>& dim_exprs);
+
+IR_API std::size_t GetHashValue(const DimExpr& dim_expr);
+
+}  // namespace symbol
+
+namespace std {
+
+template <>
+struct hash<symbol::DimExpr> {
+  std::size_t operator()(const symbol::DimExpr& dim_expr) const {
+    return symbol::GetHashValue(dim_expr);
+  }
+};
+
+template <>
+struct hash<std::vector<symbol::DimExpr>> {
+  std::size_t operator()(const std::vector<symbol::DimExpr>& dim_exprs) const {
+    std::size_t hash_value = 0;
+    const auto hash_func = std::hash<symbol::DimExpr>();
+    for (const auto& dim_expr : dim_exprs) {
+      hash_value = pir::detail::hash_combine(hash_value, hash_func(dim_expr));
+    }
+    return hash_value;
+  }
+};
+
+template <>
+struct hash<symbol::Broadcastable<symbol::DimExpr>> {
+  std::size_t operator()(
+      const symbol::Broadcastable<symbol::DimExpr>& broadcastable) const {
+    return pir::detail::hash_combine(GetHashValue(broadcastable.data->lhs),
+                                     GetHashValue(broadcastable.data->rhs));
+  }
+};
+
+}  // namespace std

paddle/pir/include/dialect/shape/utils/dim_expr_builder.h

@@ -0,0 +1,41 @@
+// Copyright (c) 2023 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
+//
+//     http://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/pir/include/core/dll_decl.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
+
+namespace symbol {
+
+class IR_API DimExprBuilder {
+ public:
+  DimExpr ConstSize(std::int64_t dim);
+  DimExpr Symbol(const std::string& symbol_name);
+  DimExpr Add(const DimExpr& lhs, const DimExpr& rhs);
+  DimExpr Any(const DimExpr& lhs, const DimExpr& rhs);
+  DimExpr Mul(const DimExpr& lhs, const DimExpr& rhs);
+  DimExpr Div(const DimExpr& lhs, const DimExpr& rhs);
+  DimExpr Max(const DimExpr& lhs, const DimExpr& rhs);
+  DimExpr Min(const DimExpr& lhs, const DimExpr& rhs);
+  DimExpr Broadcast(const DimExpr& lhs, const DimExpr& rhs);
+  std::vector<DimExpr> ConstShape(const std::vector<std::int64_t>& dims);
+
+  std::vector<DimExpr> Concat(const std::vector<DimExpr>& lhs,
+                              const std::vector<DimExpr>& rhs);
+  std::pair<std::vector<DimExpr>, std::vector<DimExpr>> SplitAt(
+      const std::vector<DimExpr>, int index);
+};
+
+}  // namespace symbol

paddle/pir/include/dialect/shape/utils/dim_expr_util.h

@@ -0,0 +1,53 @@
+// Copyright (c) 2023 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
+//
+//     http://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 <unordered_map>
+#include <unordered_set>
+
+#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
+
+namespace symbol {
+
+IR_API DimExpr SimplifyDimExpr(const DimExpr& dim_expr);
+
+IR_API DimExpr SubstituteDimExpr(
+    const DimExpr& dim_expr,
+    const std::unordered_map<DimExpr, DimExpr>& pattern_to_replacement);
+
+IR_API int GetDimExprPriority(const DimExpr& dim_expr);
+
+enum class PriorityComparisonStatus {
+  HIGHER,  // lhs has a higher priority than rhs
+  EQUAL,   // lhs and rhs have equal priority
+  LOWER    // lhs has a lower priority than rhs
+};
+IR_API PriorityComparisonStatus CompareDimExprPriority(const DimExpr& lhs,
+                                                       const DimExpr& rhs);
+
+enum class DimExprCompareResult {
+  GT,       // lhs is greater than rhs
+  EQ,       // lhs and rhs is equal
+  LT,       // lhs is less than rhs
+  UNKNOWN,  // lhs and rhs is not comparable
+};
+IR_API DimExprCompareResult Compare(const DimExpr& lhs, const DimExpr& rhs);
+
+IR_API std::unordered_set<std::string> CollectDimExprSymbols(
+    const DimExpr& dim_expr);
+
+IR_API int64_t CountExprSymbols(const DimExpr& dim_expr);
+
+}  // namespace symbol

paddle/pir/include/dialect/shape/utils/original_attributes_filter.h

@@ -0,0 +1,47 @@
+// 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
+//
+//     http://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 <map>
+#include <set>
+#include "paddle/pir/include/core/builtin_attribute.h"
+
+namespace pir {
+std::map<std::string, Attribute> GetOrderedOriginalAttributes(
+    std::string op_name,
+    const std::unordered_map<std::string, Attribute>& attributes);
+
+class OriginalAttributesFilter {
+ public:
+  static OriginalAttributesFilter& Instance();
+
+  OriginalAttributesFilter(const OriginalAttributesFilter&) = delete;
+  OriginalAttributesFilter(OriginalAttributesFilter&&) = delete;
+  OriginalAttributesFilter& operator=(const OriginalAttributesFilter&) = delete;
+
+  void SetOriginalAttributesMap(
+      const std::unordered_map<std::string, std::set<std::string>>&
+          original_attributes_map) {
+    original_attributes_map_ = original_attributes_map;
+  }
+
+ private:
+  OriginalAttributesFilter() {}
+  std::unordered_map<std::string, std::set<std::string>>
+      original_attributes_map_;
+  friend std::map<std::string, Attribute> GetOrderedOriginalAttributes(
+      std::string op_name,
+      const std::unordered_map<std::string, Attribute>& attributes);
+};
+}  // namespace pir

paddle/pir/include/dialect/shape/utils/shape_analysis.h

@@ -0,0 +1,302 @@
+// Copyright (c) 2023 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
+//
+//     http://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/include/core/builtin_attribute.h"
+#include "paddle/pir/include/core/builtin_op.h"
+#include "paddle/pir/include/core/builtin_type_interfaces.h"
+#include "paddle/pir/include/core/dll_decl.h"
+#include "paddle/pir/include/core/utils.h"
+#include "paddle/pir/include/dialect/shape/ir/shape_op.h"
+#include "paddle/pir/include/dialect/shape/utils/constraints_manager.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr_builder.h"
+#include "paddle/pir/include/dialect/shape/utils/shape_or_data_expr.h"
+
+namespace pir {
+using InferSymbolicShapeCacheValue = std::vector<symbol::ShapeOrDataDimExprs>;
+
+enum TransLayoutType { NCHW2NHWC, NHWC2NCHW, INVALID };
+
+/**
+ * This class represents information needed to determine the output
+ * shape of an operator, which includes the operator's name, input shapes, and
+ * attributes.
+ */
+class IR_API InferSymbolicShapeCacheKey {
+ public:
+  InferSymbolicShapeCacheKey(
+      const std::string& op_name,
+      const std::vector<symbol::ShapeOrDataDimExprs>& input_shape_or_datas,
+      const std::map<std::string, Attribute>& attributes)
+      : op_name_(op_name),
+        input_shape_or_datas_(input_shape_or_datas),
+        attributes_(attributes) {}
+  bool operator==(const InferSymbolicShapeCacheKey& other) const;
+  std::size_t GetHashValue() const;
+  friend std::ostream& operator<<(std::ostream& os,
+                                  const InferSymbolicShapeCacheKey& info);
+  friend class InferSymbolicShapeContext;
+
+ private:
+  std::string op_name_;
+  std::vector<symbol::ShapeOrDataDimExprs> input_shape_or_datas_;
+  std::map<std::string, Attribute> attributes_;
+  const std::vector<symbol::ShapeOrDataDimExprs>& GetInputShapeOrDatas() const;
+  void SetInputShapeOrDatas(
+      const std::vector<symbol::ShapeOrDataDimExprs>& input_shape_or_datas);
+};
+
+struct InputDynamicDimSpec {
+  std::string dim_name;
+  // input_bind = [(input_name, dim_index)]
+  std::vector<std::pair<std::string, int>> input_bind;
+  symbol::ConstraintsManager::Range range;
+};
+}  // namespace pir
+
+namespace std {
+template <>
+struct hash<pir::InferSymbolicShapeCacheKey> {
+  std::size_t operator()(const pir::InferSymbolicShapeCacheKey& obj) const {
+    return obj.GetHashValue();
+  }
+};
+}  // namespace std
+
+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::vector<InputDynamicDimSpec>& input_dynamic_dim_spec);
+
+  // Note: Only initialize the symbol info, the value info is not update.
+  void RegisterSymbolConstraintFromContext(
+      const InferSymbolicShapeContext& other);
+
+  const std::string GetNextSymName();
+
+  bool HasShapeOrDataForValue(Value val) const;
+
+  const symbol::ShapeOrDataDimExprs& GetShapeOrDataForValue(Value val) const;
+
+  void SetSymbolForValueByStaticShape(Value val);
+
+  void SetShapeOrDataForValue(Value val,
+                              const symbol::ShapeOrDataDimExprs& shape_or_data);
+
+  void AddEqualCstr(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs);
+
+  // Add equal constraints for each dim in lhs and rhs.
+  void AddEqualCstr(const std::vector<symbol::DimExpr>& lhs,
+                    const std::vector<symbol::DimExpr>& rhs);
+
+  bool IsEqual(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs) const;
+
+  // Returns true if:
+  //    lhs[i] == rhs[i] for all i
+  bool IsEqual(const std::vector<symbol::DimExpr>& lhs,
+               const std::vector<symbol::DimExpr>& rhs) const;
+
+  void AddGreatThanOneCstr(const symbol::DimExpr& dim_expr);
+
+  bool IsGreatThanOne(const symbol::DimExpr& dim_expr) const;
+
+  void AddBroadcastableCstr(const symbol::DimExpr& lhs,
+                            const symbol::DimExpr& rhs);
+
+  bool IsBroadcastable(const symbol::DimExpr& lhs,
+                       const symbol::DimExpr& rhs) const;
+
+  bool HasPredefinedRange(const symbol::DimExpr& dim_expr) const;
+
+  void PrintShapeOrDatas() const;
+
+  void SetOpInferSymbolicShapeCache(
+      const InferSymbolicShapeCacheKey& op_infer_cache_key,
+      InferSymbolicShapeCacheValue result_shape);
+
+  std::optional<InferSymbolicShapeCacheValue> GetOpInferSymbolicShapeCache(
+      const InferSymbolicShapeCacheKey& op_infer_cache_key) const;
+
+  void ClearOpInferSymbolicShapeCache();
+
+  const symbol::ConstraintsManager& constraints_manager() const {
+    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);
+
+  void SubstituteDimExpr(const symbol::DimExpr& origin,
+                         const symbol::DimExpr& substituted);
+
+  int64_t sym_idx_begin_ = 0;
+  int64_t next_sym_idx_ = 0;
+
+  std::unordered_map<uint64_t, symbol::ShapeOrDataDimExprs>
+      value_id_to_shape_or_data_;
+
+  symbol::ConstraintsManager constraints_manager_;
+
+  using DimExprSubstitutionPattern =
+      std::unordered_map<symbol::DimExpr, symbol::DimExpr>;
+  DimExprSubstitutionPattern substitution_pattern_;
+
+  std::unordered_map<InferSymbolicShapeCacheKey, InferSymbolicShapeCacheValue>
+      infer_symbolic_shape_cache_;
+
+  std::unordered_map<std::string, std::vector<DimIndexAndExpr>>
+      predefined_dimexpr_map_for_inputs_;
+
+  std::unordered_map<std::string, symbol::DimExpr>
+      input_dynamic_dim_name_spec_to_dimexpr_map_;
+};
+
+class IR_API ShapeConstraintIRAnalysis final
+    : public std::enable_shared_from_this<ShapeConstraintIRAnalysis> {
+ public:
+  ShapeConstraintIRAnalysis() = default;
+  ShapeConstraintIRAnalysis(const ShapeConstraintIRAnalysis&) = delete;
+  ShapeConstraintIRAnalysis(ShapeConstraintIRAnalysis&&) = delete;
+  void InitInferContext();
+
+  void RegisterSymbolConstraintFromShapeAnalysis(
+      const ShapeConstraintIRAnalysis& other);
+
+  const std::string GetNextSymName();
+
+  const symbol::ShapeOrDataDimExprs& GetShapeOrDataForValue(Value val);
+
+  void SetShapeOrDataForValue(Value val,
+                              const symbol::ShapeOrDataDimExprs& shape_or_data);
+
+  // Set ShapeOrData of `to` value by ShapeOrData of `from` value.
+  void ShareShapeOrData(Value from, Value to);
+
+  // Update Symbol Shape for value by layout transformation.
+  void UpdateShapeOrDataByTransLayout(Value val,
+                                      TransLayoutType trans_layout_type);
+
+  void AddEqualCstr(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs);
+
+  bool IsEqual(const symbol::DimExpr& lhs, const symbol::DimExpr& rhs) const;
+
+  bool IsGreatThanOne(const symbol::DimExpr& dim_expr) const;
+
+  bool IsBroadcastable(const symbol::DimExpr& lhs,
+                       const symbol::DimExpr& rhs) const;
+
+  // Used to debug
+  void PrintShapeOrDatas() const;
+
+  // Returns true if the two value have the same symbolic shape.
+  bool IsShapeEqual(Value lhs, Value rhs);
+
+  // Suppose:
+  //    lhs_dim_idxs = {ld0, ld1, ...}
+  //    rhs_dim_idxs = {rd0, rd1, ...}
+  // Returns true if:
+  //    lhs.shape[ld0] * lhs.shape[ld1] * ... ==
+  //    rhs.shape[rd0] * rhs.shape[rd1] * ...
+  bool IsProductEqual(Value lhs,
+                      const std::vector<int>& lhs_dim_idxs,
+                      Value rhs,
+                      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);
+
+  // Returns true if the two value have the same number elements.
+  bool IsSameNumel(Value lhs, Value rhs);
+
+  pir::PrintHooks PrintHook();
+
+  symbol::DimExpr GetProductDimExpr(Value lhs,
+                                    const std::vector<int>& lhs_dim_idxs);
+
+  const symbol::ConstraintsManager& constraints_manager() const {
+    return context_.constraints_manager();
+  }
+
+  void ClearOpInferSymbolicShapeCache() {
+    context_.ClearOpInferSymbolicShapeCache();
+  }
+
+  void SetInputDynamicDimSpec(
+      const std::vector<InputDynamicDimSpec>& input_dynamic_dim_spec);
+
+  void AppendInputDynamicDimSpec(
+      const std::vector<InputDynamicDimSpec>& input_dynamic_dim_spec);
+
+ private:
+  InferSymbolicShapeContext* MutInferSymbolicShapeContext() {
+    return &context_;
+  }
+
+  friend void InferSymExprForAllValues(ModuleOp module_op);
+
+  void SetSymbolForValueByStaticShape(Value val);
+
+  void InferShapeOrDataForValue(Value val);
+
+ private:
+  InferSymbolicShapeContext context_;
+  std::vector<InputDynamicDimSpec> input_dynamic_dim_spec_;
+};
+
+class IR_API ShapeAnalysisManager {
+ public:
+  static ShapeAnalysisManager& Instance();
+  ShapeConstraintIRAnalysis& Get(const pir::Program* program);
+
+  ShapeAnalysisManager(const ShapeAnalysisManager&) = delete;
+  ShapeAnalysisManager(ShapeAnalysisManager&&) = delete;
+  ShapeAnalysisManager& operator=(const ShapeAnalysisManager&) = delete;
+
+ private:
+  ShapeAnalysisManager() {}
+  std::unordered_map<uint64_t, std::shared_ptr<ShapeConstraintIRAnalysis>>
+      tables_;
+};
+
+#define OP_DECLARE_INFER_SYMBOLIC_SHAPE(name) \
+  bool name##OpInferSymbolicShape(            \
+      pir::Operation* op, pir::InferSymbolicShapeContext* infer_context);
+
+bool IsStaticShape(const Value& value);
+
+}  // namespace pir

paddle/pir/include/dialect/shape/utils/shape_or_data_expr.h

@@ -0,0 +1,281 @@
+// 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
+//
+//     http://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 <sstream>
+#include "paddle/common/overloaded.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr_util.h"
+
+namespace symbol {
+
+template <typename T>
+class ShapeOrData {
+ public:
+  explicit ShapeOrData(const std::vector<T>& shape)
+      : shape_(shape), data_(std::nullopt) {}
+  explicit ShapeOrData(const std::vector<T>& shape, const std::vector<T>& data)
+      : shape_(shape), data_(data) {
+    // Valid check
+    if (shape.size() == 0) {
+      PADDLE_ENFORCE_EQ(
+          data.size(),
+          1UL,
+          common::errors::InvalidArgument(
+              "When shape is 0-D, size of data should be 1, but got %d.",
+              data.size()));
+    } else if (shape.size() == 1) {
+      PADDLE_ENFORCE_EQ(shape[0].template Has<int64_t>(),
+                        true,
+                        common::errors::InvalidArgument(
+                            "When shape is 1-D, value of shape should be int"));
+      PADDLE_ENFORCE_EQ(
+          shape[0].template Get<int64_t>() == static_cast<int64_t>(data.size()),
+          true,
+          common::errors::InvalidArgument(
+              "When shape is 1-D, size of data should be the same as "
+              "value[%d] of shape, but got [%d].",
+              shape[0].template Get<std::int64_t>(),
+              data.size()));
+    } else {
+      int64_t numel = 1;
+      for (const auto& expr : shape) {
+        PADDLE_ENFORCE_EQ(expr.template isa<int64_t>(),
+                          true,
+                          ::common::errors::InvalidArgument(
+                              "When data has value, the expr of shape should "
+                              "be int, but got %s.",
+                              ToString(expr)));
+        numel *= expr.template Get<int64_t>();
+      }
+      PADDLE_ENFORCE_EQ(numel,
+                        data.size(),
+                        ::common::errors::InvalidArgument(
+                            "Size of data should be the same as "
+                            "product of value[%d] of shape, but got [%d].",
+                            numel,
+                            data.size()));
+    }
+  }
+
+  ShapeOrData() = default;
+  ShapeOrData(const ShapeOrData&) = default;
+  ShapeOrData(ShapeOrData&&) = default;
+  ShapeOrData& operator=(const ShapeOrData&) = default;
+  ShapeOrData& operator=(ShapeOrData&&) = default;
+
+  // Tensor's real shape
+  const std::vector<T>& shape() const { return shape_; }
+  // Specific for Tensor generated by shape-relevant ops
+  const std::optional<std::vector<T>>& data() const { return data_; }
+  void SetData(const std::vector<T>& data) { data_ = data; }
+
+  bool operator==(const ShapeOrData<T>& other) const {
+    if (data_.has_value() && !other.data_.has_value()) return false;
+    if (!data_.has_value() && other.data_.has_value()) return false;
+    if (shape_.size() != other.shape_.size()) return false;
+
+    if (data_.has_value() && other.data_.has_value()) {
+      if (data_.value().size() != other.data_.value().size()) return false;
+
+      for (size_t i = 0; i < data_.value().size(); ++i) {
+        DimExpr dim0 = symbol::SimplifyDimExpr(data_.value()[i]);
+        DimExpr dim1 = symbol::SimplifyDimExpr(other.data_.value()[i]);
+        if (dim0 != dim1) return false;
+      }
+    }
+
+    for (size_t i = 0; i < shape_.size(); ++i) {
+      DimExpr dim0 = symbol::SimplifyDimExpr(shape_[i]);
+      DimExpr dim1 = symbol::SimplifyDimExpr(other.shape_[i]);
+      if (dim0 != dim1) return false;
+    }
+
+    return true;
+  }
+
+  bool operator!=(const ShapeOrData<T>& other) const {
+    return !(*this == other);
+  }
+
+ private:
+  std::vector<T> shape_;
+  std::optional<std::vector<T>> data_;
+};
+
+using NullShapeOrDataDimExpr = std::monostate;
+using TensorShapeOrDataDimExprs = ShapeOrData<DimExpr>;
+using TensorListShapeOrDataDimExprs = std::vector<TensorShapeOrDataDimExprs>;
+
+/* TensorArray can append tensors dynamically. In a static graph, we only
+ * store the shape of one element as a hint, because we assume that all elements
+ * in the TensorArray have the same rank, and with equal constraints on specific
+ * dimensions. */
+class RankedTensorArrayShapeOrDataDimExprs {
+ public:
+  RankedTensorArrayShapeOrDataDimExprs() = default;
+  explicit RankedTensorArrayShapeOrDataDimExprs(
+      const std::vector<DimExpr>& shape)
+      : shape_hint_{shape} {}
+  const std::vector<DimExpr>& GetShapeHint() const { return shape_hint_; }
+  bool operator==(const RankedTensorArrayShapeOrDataDimExprs& other) const {
+    if (shape_hint_.size() != other.shape_hint_.size()) return false;
+    for (size_t i = 0; i < shape_hint_.size(); ++i) {
+      DimExpr dim0 = symbol::SimplifyDimExpr(shape_hint_[i]);
+      DimExpr dim1 = symbol::SimplifyDimExpr(other.shape_hint_[i]);
+      if (dim0 != dim1) return false;
+    }
+
+    return true;
+  }
+
+ private:
+  std::vector<DimExpr> shape_hint_;
+};
+
+using ShapeOrDataDimExprsBase =
+    std::variant<NullShapeOrDataDimExpr,
+                 TensorShapeOrDataDimExprs,
+                 TensorListShapeOrDataDimExprs,
+                 RankedTensorArrayShapeOrDataDimExprs>;
+
+class ShapeOrDataDimExprs : public ShapeOrDataDimExprsBase {
+ public:
+  ShapeOrDataDimExprs() = delete;
+  ShapeOrDataDimExprs(
+      const TensorShapeOrDataDimExprs& tensor_dim_exprs)  // NOLINT
+      : ShapeOrDataDimExprsBase(tensor_dim_exprs) {}
+  ShapeOrDataDimExprs(
+      const TensorListShapeOrDataDimExprs& tensor_list_dim_exprs)
+      : ShapeOrDataDimExprsBase(tensor_list_dim_exprs) {}
+
+  ShapeOrDataDimExprs(const RankedTensorArrayShapeOrDataDimExprs&
+                          tensor_array_dim_exprs)  // NOLINT
+      : ShapeOrDataDimExprsBase(tensor_array_dim_exprs) {}
+
+  ShapeOrDataDimExprs(const NullShapeOrDataDimExpr& null_dim_expr)  // NOLINT
+      : ShapeOrDataDimExprsBase(null_dim_expr) {}
+
+  template <typename T>
+  bool isa() const {
+    return std::holds_alternative<T>(*this);
+  }
+
+  template <typename T>
+  const T& dyn_cast() const {
+    return std::get<T>(*this);
+  }
+
+  const ShapeOrDataDimExprsBase& variant() const {
+    return static_cast<const ShapeOrDataDimExprsBase&>(*this);
+  }
+
+  DEFINE_MATCH_METHOD();
+
+  bool operator==(const ShapeOrDataDimExprs& other) const {
+    return this->variant() == other.variant();
+  }
+
+  bool operator!=(const ShapeOrDataDimExprs& other) const {
+    return !(*this == other);
+  }
+
+  const std::vector<DimExpr>& shape() const {
+    PADDLE_ENFORCE_EQ(std::holds_alternative<TensorShapeOrDataDimExprs>(*this),
+                      true,
+                      common::errors::PreconditionNotMet(
+                          "Shape of ShapeOrData is not a vector, "
+                          "check whether the value is a "
+                          "tensor-list or not."));
+    return std::get<TensorShapeOrDataDimExprs>(*this).shape();
+  }
+
+  const std::optional<std::vector<DimExpr>>& data() const {
+    PADDLE_ENFORCE_EQ(
+        std::holds_alternative<TensorShapeOrDataDimExprs>(*this),
+        true,
+        common::errors::PreconditionNotMet(
+            "Data of ShapeOrData is not a vector, check whether the value is a "
+            "tensor-list or not."));
+    return std::get<TensorShapeOrDataDimExprs>(*this).data();
+  }
+
+  void SetData(const std::vector<DimExpr>& data) {
+    PADDLE_ENFORCE_EQ(
+        std::holds_alternative<TensorShapeOrDataDimExprs>(*this),
+        true,
+        common::errors::PreconditionNotMet(
+            "Data of ShapeOrData is not a vector, check whether the value is a "
+            "tensor-list or not."));
+
+    std::get<TensorShapeOrDataDimExprs>(*this).SetData(data);
+  }
+};
+
+IR_API ShapeOrDataDimExprs SubstituteShapeOrData(
+    const ShapeOrDataDimExprs& shape_or_data,
+    const std::unordered_map<DimExpr, DimExpr>& substitution_pattern);
+
+IR_API std::ostream& operator<<(std::ostream&,
+                                const ShapeOrDataDimExprs& dim_expr);
+
+}  // namespace symbol
+
+namespace std {
+
+template <>
+struct hash<symbol::TensorShapeOrDataDimExprs> {
+  std::size_t operator()(const symbol::TensorShapeOrDataDimExprs& obj) const {
+    const auto hash_func = std::hash<std::vector<symbol::DimExpr>>();
+    std::size_t ret = hash_func(obj.shape());
+    ret = pir::detail::hash_combine(ret, obj.data().has_value());
+    if (obj.data().has_value()) {
+      ret = pir::detail::hash_combine(ret, hash_func(obj.data().value()));
+    }
+    return ret;
+  }
+};
+
+template <>
+struct hash<symbol::TensorListShapeOrDataDimExprs> {
+  std::size_t operator()(
+      const symbol::TensorListShapeOrDataDimExprs& obj) const {
+    const auto hash_func = std::hash<symbol::TensorShapeOrDataDimExprs>();
+    std::size_t ret = 0;
+    for (const auto& shape_or_data : obj) {
+      ret = pir::detail::hash_combine(ret, hash_func(shape_or_data));
+    }
+    return ret;
+  }
+};
+
+template <>
+struct hash<symbol::RankedTensorArrayShapeOrDataDimExprs> {
+  std::size_t operator()(
+      const symbol::RankedTensorArrayShapeOrDataDimExprs& obj) const {
+    return std::hash<std::vector<symbol::DimExpr>>()(obj.GetShapeHint());
+  }
+};
+
+template <>
+struct hash<symbol::ShapeOrDataDimExprs> {
+  std::size_t operator()(const symbol::ShapeOrDataDimExprs& obj) const {
+    return obj.Match([](const auto& impl) {
+      using T = std::decay_t<decltype(impl)>;
+      return std::hash<T>()(impl);
+    });
+  }
+};
+
+}  // namespace std