chore: import upstream snapshot with attribution

paddle/pir/include/pattern_rewrite/frozen_rewrite_pattern_set.h

@@ -0,0 +1,74 @@
+// 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.
+
+/// The design and code is mainly from MLIR, thanks to the great project.
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/pir/include/core/dll_decl.h"
+#include "paddle/pir/include/core/op_info.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_match.h"
+
+namespace pir {
+
+class IR_API FrozenRewritePatternSet {
+  using NativePatternListT = std::vector<std::unique_ptr<RewritePattern>>;
+
+ public:
+  using OpSpecificNativePatternListT =
+      std::unordered_map<OpInfo, std::vector<RewritePattern*>>;
+
+  FrozenRewritePatternSet();
+  FrozenRewritePatternSet(FrozenRewritePatternSet&& patterns) = default;
+  FrozenRewritePatternSet(const FrozenRewritePatternSet& patterns) = default;
+  FrozenRewritePatternSet& operator=(FrozenRewritePatternSet&& patterns) =
+      default;
+  FrozenRewritePatternSet& operator=(const FrozenRewritePatternSet& patterns) =
+      default;
+  ~FrozenRewritePatternSet() = default;
+
+  /// Freeze the patterns held in `patterns`, and take ownership.
+  FrozenRewritePatternSet(
+      RewritePatternSet&& patterns,
+      const std::vector<std::string>& disabled_pattern_labels = {},
+      const std::vector<std::string>& enabled_pattern_labels = {});
+
+  /// Return the op specific native patterns held by this list.
+  const OpSpecificNativePatternListT& op_specific_native_patterns() const {
+    return impl_->op_specific_native_pattern_map_;
+  }
+
+  /// Return the "match any" native patterns held by this list.
+  const NativePatternListT& match_any_op_native_patterns() const {
+    return impl_->match_any_op_native_patterns_;
+  }
+
+ private:
+  struct Impl {
+    OpSpecificNativePatternListT op_specific_native_pattern_map_;
+
+    NativePatternListT op_specific_native_patterns_;
+
+    NativePatternListT match_any_op_native_patterns_;
+  };
+
+  std::shared_ptr<Impl> impl_;
+};
+
+}  // namespace pir

paddle/pir/include/pattern_rewrite/pattern_applicator.h

@@ -0,0 +1,59 @@
+// 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.
+
+// The design and code is mainly from MLIR, thanks to the great project.
+
+#pragma once
+
+#include <functional>
+#include <unordered_map>
+
+#include "paddle/pir/include/core/op_info.h"
+#include "paddle/pir/include/core/operation.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_match.h"
+
+namespace pir {
+
+class FrozenRewritePatternSet;
+class RewritePattern;
+class Pattern;
+
+class PatternApplicator {
+ public:
+  using CostModel = std::function<PatternBenefit(const Pattern&)>;
+
+  explicit PatternApplicator(const FrozenRewritePatternSet& frozen_patter_list);
+  ~PatternApplicator() = default;
+
+  bool MatchAndRewrite(Operation* op,
+                       PatternRewriter& rewriter,  // NOLINT
+                       std::function<bool(const Pattern&)> can_apply = {},
+                       std::function<void(const Pattern&)> on_failure = {},
+                       std::function<bool(const Pattern&)> on_success = {});
+
+  void ApplyCostModel(const CostModel& model);
+
+  void ApplyDefaultCostModel() {
+    ApplyCostModel([](const Pattern& pattern) { return pattern.benefit(); });
+  }
+
+  void WalkAllPatterns(std::function<void(const Pattern&)> walk);
+
+ private:
+  const FrozenRewritePatternSet& frozen_pattern_list_;
+  std::unordered_map<OpInfo, std::vector<const RewritePattern*>> patterns_;
+  std::vector<const RewritePattern*> any_op_patterns_;
+};
+
+}  // namespace pir

paddle/pir/include/pattern_rewrite/pattern_match.h

@@ -0,0 +1,379 @@
+// 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.
+
+// The design is mainly from MLIR, very thanks to the great project.
+
+#pragma once
+
+#include <functional>
+#include <initializer_list>
+#include <memory>
+#include <optional>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "paddle/common/enforce.h"
+#include "paddle/pir/include/core/builder.h"
+#include "paddle/pir/include/core/dll_decl.h"
+#include "paddle/pir/include/core/ir_context.h"
+#include "paddle/pir/include/core/op_info.h"
+#include "paddle/pir/include/core/operation.h"
+#include "paddle/pir/include/core/type_id.h"
+#include "paddle/pir/include/core/type_name.h"
+#include "paddle/pir/include/core/value.h"
+
+namespace pir {
+
+// This class represents the benefit of a pattern. The most common
+// unit to use is the `number of operations` in the pattern.
+class IR_API PatternBenefit {
+ public:
+  PatternBenefit() = default;
+  PatternBenefit(uint32_t val) : val_(val) {}  // NOLINT
+
+  uint32_t benefit() { return val_; }
+
+  bool operator==(const PatternBenefit& rhs) const { return val_ == rhs.val_; }
+  bool operator!=(const PatternBenefit& rhs) const { return !(*this == rhs); }
+  bool operator<(const PatternBenefit& rhs) const { return val_ < rhs.val_; }
+  bool operator>(const PatternBenefit& rhs) const { return rhs < *this; }
+  bool operator<=(const PatternBenefit& rhs) const { return !(*this > rhs); }
+  bool operator>=(const PatternBenefit& rhs) const { return !(*this < rhs); }
+
+ private:
+  uint32_t val_{0};
+};
+
+// This class contains all of the data related to a Pattern, but not contains
+// any methods for the matching. This class is used to interface with the
+// metadata of a pattern, such as benefit or root operation.
+class IR_API Pattern {
+  enum class RootKind {
+    // The pattern root matches "any" operation.
+    Any,
+    // The pattern root is matched using a concrete operation.
+    OperationInfo,
+    // The pattern root is matched using an interface id.
+    InterfaceId,
+    // The pattern root is matched using a trait id.
+    TraitId
+  };
+
+ public:
+  const std::vector<OpInfo>& generated_ops() const { return generated_ops_; }
+
+  std::optional<OpInfo> root_kind() const {
+    if (root_kind_ == RootKind::OperationInfo)
+      return OpInfo::RecoverFromVoidPointer(root_val_);
+    return std::nullopt;
+  }
+
+  std::optional<TypeId> GetRootInterfaceID() const {
+    if (root_kind_ == RootKind::InterfaceId)
+      return TypeId::RecoverFromVoidPointer(root_val_);
+    return std::nullopt;
+  }
+
+  std::optional<TypeId> GetRootTraitID() const {
+    if (root_kind_ == RootKind::TraitId)
+      return TypeId::RecoverFromVoidPointer(root_val_);
+    return std::nullopt;
+  }
+
+  PatternBenefit benefit() const { return benefit_; }
+
+  IrContext* ir_context() const { return context_; }
+
+  std::string debug_name() const { return debug_name_; }
+
+  void SetDebugName(const std::string& name) { debug_name_ = name; }
+
+  const std::vector<std::string>& debug_labels() const { return debug_labels_; }
+
+  void AddDebugLabels(const std::vector<std::string>& labels) {
+    debug_labels_.insert(debug_labels_.end(), labels.begin(), labels.end());
+  }
+
+  void AddDebugLabels(const std::string& label) {
+    debug_labels_.push_back(label);
+  }
+
+ protected:
+  struct MatchAnyOpTypeTag {};
+  struct MatchInterfaceOpTypeTag {};
+  struct MatchTraitOpTypeTag {};
+
+  Pattern(const std::string& root_name,
+          PatternBenefit benefit,
+          IrContext* context,
+          const std::vector<std::string>& generated_names = {});
+
+  Pattern(MatchAnyOpTypeTag tag,
+          PatternBenefit benefit,
+          IrContext* context,
+          const std::vector<std::string>& generated_names = {});
+
+  Pattern(MatchInterfaceOpTypeTag tag,
+          TypeId interface_id,
+          PatternBenefit benefit,
+          IrContext* context,
+          const std::vector<std::string>& generated_names = {});
+
+  Pattern(MatchTraitOpTypeTag tag,
+          TypeId trait_id,
+          PatternBenefit benefit,
+          IrContext* context,
+          const std::vector<std::string>& generated_names = {});
+
+ private:
+  Pattern(void* root_val,
+          RootKind root_kind,
+          const std::vector<std::string>& generated_names,
+          PatternBenefit benefit,
+          IrContext* context);
+
+  void* root_val_;
+  RootKind root_kind_;
+
+  const PatternBenefit benefit_;
+  IrContext* context_;
+  // A list of the potential operations that may be generated when rewriting an
+  // op with this pattern.
+  std::vector<OpInfo> generated_ops_;
+
+  std::string debug_name_;
+  std::vector<std::string> debug_labels_;
+};
+
+class PatternRewriter;
+
+class IR_API RewritePattern : public Pattern {
+ public:
+  virtual ~RewritePattern();
+
+  virtual void Rewrite(Operation* op,
+                       PatternRewriter& rewriter) const {  // NOLINT
+    IR_THROW(
+        "need to implement either MatchAndRewrite or one of the rewrite "
+        "functions.");
+  }
+
+  virtual bool Match(Operation* op) const {
+    IR_THROW("need to implement either MatchAndRewrite or Match.");
+    return false;
+  }
+
+  virtual bool MatchAndRewrite(Operation* op,
+                               PatternRewriter& rewriter) const {  // NOLINT
+    if (Match(op)) {
+      Rewrite(op, rewriter);
+      return true;
+    }
+    return false;
+  }
+
+  virtual void Initialize() {}
+
+  template <typename T, typename... Args>
+  static std::unique_ptr<T> Create(Args&&... args) {
+    std::unique_ptr<T> pattern =
+        std::make_unique<T>(std::forward<Args>(args)...);
+    pattern->Initialize();
+
+    if (pattern->debug_name().empty())
+      pattern->SetDebugName(pir::get_type_name<T>());
+    return pattern;
+  }
+
+ protected:
+  using Pattern::Pattern;
+};
+
+namespace detail {
+// A wrapper around PatternWrite that allows for matching and rewriting
+// against an instance of a derived operation class or Interface.
+template <typename SourceOp>
+struct OpOrInterfaceRewritePatternBase : public RewritePattern {
+  using RewritePattern::RewritePattern;
+
+  void Rewrite(Operation* op,
+               PatternRewriter& rewriter) const final {  // NOLINT
+    Rewrite(op->dyn_cast<SourceOp>(), rewriter);
+  }
+
+  bool Match(Operation* op) const final {
+    return Match(op->dyn_cast<SourceOp>());
+  }
+  bool MatchAndRewrite(Operation* op,
+                       PatternRewriter& rewriter) const final {  // NOLINT
+    return MatchAndRewrite(op->dyn_cast<SourceOp>(), rewriter);
+  }
+
+  virtual void Rewrite(SourceOp op,
+                       PatternRewriter& rewriter) const {  // NOLINT
+    IR_THROW("must override Rewrite or MatchAndRewrite");
+  }
+  virtual bool Match(SourceOp op) const {
+    IR_THROW("must override Match or MatchAndRewrite");
+  }
+  virtual bool MatchAndRewrite(SourceOp op,
+                               PatternRewriter& rewriter) const {  // NOLINT
+    if (Match(op)) {
+      Rewrite(op, rewriter);
+      return true;
+    }
+    return false;
+  }
+};
+}  // namespace detail
+
+// OpRewritePattern is a wrapper around RewritePattern that allows for
+// matching and rewriting against an instance of a derived operation
+// class as opposed to a raw Operation.
+template <typename SourceOp>
+struct OpRewritePattern
+    : public detail::OpOrInterfaceRewritePatternBase<SourceOp> {
+  OpRewritePattern(IrContext* context,
+                   PatternBenefit benefit = 1,
+                   const std::vector<std::string>& generated_names = {})
+      : detail::OpOrInterfaceRewritePatternBase<SourceOp>(
+            SourceOp::name(), benefit, context, generated_names) {}
+};
+
+// TODO(wilber): Support OpInterfaceRewritePattern and OpTraitRewritePattern.
+// ...
+
+// This class provides a series of interfaces for modifying IR and tracking IR
+// changes. This class provides a unified API for IR modification.
+class RewriterBase : public Builder {
+ public:
+  // TODO(wilber): Supplementary methods of block and region.
+
+  virtual void ReplaceOp(Operation* op, const std::vector<Value>& new_values);
+
+  // This method erases an operation that is known to have no uses.
+  virtual void EraseOp(Operation* op);
+
+  IR_API void ReplaceAllUsesWith(Value from, Value to);
+
+ protected:
+  explicit RewriterBase(IrContext* ctx) : Builder(ctx) {}
+
+  virtual ~RewriterBase();
+
+  virtual void NotifyRootReplaced(Operation* op,
+                                  const std::vector<Value>& replacement) {}
+
+  virtual void NotifyOperationRemoved(Operation* op) {}
+
+  virtual void NotifyOperationInserted(Operation* op) {}
+
+  virtual void NotifyValueReplaced(Value from, Value to) {}
+
+  virtual void StartRootUpdate(Operation* op) {}
+
+  virtual void FinalizeRootUpdate(Operation* op) {}
+
+  virtual void CancelRootUpdate(Operation* op) {}
+
+  template <typename CallableT>
+  void UpdateRootInplace(Operation* root, CallableT&& callable) {
+    StartRootUpdate(root);
+    callable();
+    FinalizeRootUpdate(root);
+  }
+
+ private:
+  void operator=(const RewriterBase&) = delete;
+  RewriterBase(const RewriterBase&) = delete;
+};
+
+class PatternRewriter : public RewriterBase {
+ public:
+  using RewriterBase::RewriterBase;
+};
+
+// A pattern collection, easy to add patterns.
+class RewritePatternSet {
+  using NativePatternListT = std::vector<std::unique_ptr<RewritePattern>>;
+
+ public:
+  explicit RewritePatternSet(IrContext* context) : context_(context) {}
+
+  // Construct a RewritePatternSet with the given patterns.
+  RewritePatternSet(IrContext* context, std::unique_ptr<RewritePattern> pattern)
+      : context_(context) {
+    native_patterns_.emplace_back(std::move(pattern));
+  }
+
+  IrContext* ir_context() const { return context_; }
+
+  NativePatternListT& native_patterns() { return native_patterns_; }
+
+  void Clear() { native_patterns_.clear(); }
+
+  bool Empty() const { return native_patterns_.empty(); }
+
+  // 'add' methods for adding patterns to the set.
+  template <typename... Ts,
+            typename ConstructorArg,
+            typename... ConstructorArgs,
+            typename = std::enable_if_t<sizeof...(Ts) != 0>>
+  RewritePatternSet& Add(ConstructorArg&& arg, ConstructorArgs&&... args) {
+    (void)std::initializer_list<int>{
+        (AddImpl<Ts>({},
+                     std::forward<ConstructorArg>(arg),
+                     std::forward<ConstructorArgs>(args)...),
+         0)...};
+    return *this;
+  }
+
+  template <typename... Ts,
+            typename ConstructorArg,
+            typename... ConstructorArgs,
+            typename = std::enable_if_t<sizeof...(Ts) != 0>>
+  RewritePatternSet& AddWithLabel(const std::vector<std::string>& debug_labels,
+                                  ConstructorArg&& arg,
+                                  ConstructorArgs&&... args) {
+    (void)std::initializer_list<int>{
+        (AddImpl<Ts>(debug_labels,
+                     std::forward<ConstructorArg>(arg),
+                     std::forward<ConstructorArgs>(args)...),
+         0)...};
+    return *this;
+  }
+
+  RewritePatternSet& Add(std::unique_ptr<RewritePattern> pattern) {
+    native_patterns_.emplace_back(std::move(pattern));
+    return *this;
+  }
+
+ private:
+  template <typename T, typename... Args>
+  std::enable_if_t<std::is_base_of<RewritePattern, T>::value> AddImpl(
+      const std::vector<std::string>& debug_labels, Args&&... args) {
+    std::unique_ptr<T> pattern =
+        RewritePattern::Create<T>(std::forward<Args>(args)...);
+    pattern->AddDebugLabels(debug_labels);
+    native_patterns_.emplace_back(std::move(pattern));
+  }
+
+ private:
+  IrContext* const context_;
+
+  NativePatternListT native_patterns_;
+};
+
+}  // namespace pir

paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h

@@ -0,0 +1,86 @@
+// 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/core/region.h"
+
+namespace pir {
+
+using VALUE_REPLACED_HOOK_FUNC = std::function<void(pir::Value, pir::Value)>;
+
+class FrozenRewritePatternSet;
+
+/// This enum will control which ops will be added to the worklist during the
+/// match rewrite process
+enum class IR_API GreedyRewriteStrictness {
+  /// No restrictions wrt. any ops are processed.
+  AnyOp,
+  /// Only pre-existing and newly created ops are processed.
+  ExistingAndNewOps,
+  /// Only pre-existing ops are processed.
+  ExistingOps
+};
+
+/// Control over how the GreedyPatternRewriteDriver works.
+class IR_API GreedyRewriteConfig {
+ public:
+  /// Control the way op is added to the worklist: bottom-up or top-down.
+  bool use_top_down_traversal = false;
+
+  /// Control the maximum number of iterations in the process of applying the
+  /// pattern, use `kNolimit` to represent unlimited.
+  int64_t max_iterations = 10;
+
+  /// Control the upper limit of rewrite times during each iteration, use
+  /// kNoLimit to represent unlimited.
+  int64_t max_num_rewrites = kNoLimit;
+
+  /// Only the op inside this region will be added to the worklist.
+  Region* region{nullptr};
+
+  /// Limit which ops will be added to the worklist during the Match and Rewrite
+  /// process.
+  /// - AnyOp: all ops will be added to the worklist.
+  /// - ExistingAndNewOps: pre-existing ops and newly created ops are added to
+  /// the worklist.
+  /// - ExistingOps: only pre-existing ops are added to the worklist.
+  GreedyRewriteStrictness strict_mode = GreedyRewriteStrictness::AnyOp;
+
+  // Hook function for replacing the value.
+  VALUE_REPLACED_HOOK_FUNC value_replaced_hook = nullptr;
+
+  static constexpr int64_t kNoLimit = -1;
+};
+
+/// Perform the Match and Rewrite process in the specified region, greedily
+/// apply the Pattern with the highest benefit, and repeat this process until
+/// convergence or the upper limit of iterations.
+///
+/// Returns pair<bool,int64_t>
+// the first is true if the iteration converges and no patterns can be applied.
+// the second is the number of total match count.
+std::pair<bool, int64_t> IR_API
+ApplyPatternsGreedily(Region& region,  // NOLINT
+                      const FrozenRewritePatternSet& patterns,
+                      GreedyRewriteConfig config = GreedyRewriteConfig());
+
+/// Perform a match and rewrite process for all regions of a given op.
+IR_API std::pair<bool, int64_t> ApplyPatternsGreedily(
+    Operation* op,
+    const FrozenRewritePatternSet& patterns,
+    GreedyRewriteConfig config = GreedyRewriteConfig());
+
+}  // namespace pir