[PIR] Support Operation::Clone Interface

paddle/cinn/hlir/framework/pir/group.h

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 #pragma once
+#include <memory>
 #include <string>
+#include <unordered_map>
 #include <vector>
+#include "glog/logging.h"
 
 #include "paddle/cinn/adt/graph_symbolic_dim_infer_ctx.h"
 #include "paddle/cinn/hlir/framework/op.h"
@@ -34,8 +37,17 @@ namespace framework {
 namespace pir {
 using framework::OpPatternKind;
 
-// TODO(Aurelius84): Need to be replaced with CinnGroupOp
 struct Group {
+  // Control the clone strategy for Group.
+  class Options {
+   public:
+    Options() : only_clone_ops(true) {}
+    bool OnlyCloneOps() const { return only_clone_ops; }
+
+   private:
+    bool only_clone_ops = false;
+  };
+
  public:
   Group() = default;
   Group(const Group&) = delete;
@@ -47,6 +59,36 @@ struct Group {
   explicit Group(std::initializer_list<::pir::Operation*> group_ops)
       : ops(group_ops) {}
 
+  std::shared_ptr<Group> Clone(::pir::Block* target_block,
+                               ::pir::IrMapping& ir_mapping,
+                               const Options& option = Options()) const {
+    CHECK_EQ(option.OnlyCloneOps(), true)
+        << "Only Support Clone Group ops information.";
+    std::vector<::pir::Operation*> new_ops;
+    // Mapper from original to new ops.
+    std::unordered_map<::pir::Operation*, ::pir::Operation*> ops_mapper;
+    ::pir::CloneOptions clone_options(false, true);
+    for (auto* op : this->ops_set) {
+      auto* new_op = op->Clone(ir_mapping, clone_options);
+      // NOTE(dev): Must call MoveTo to deal with ownership, otherwise it
+      // will lead memory-leak.
+      new_op->MoveTo(target_block, target_block->end());
+      new_ops.push_back(new_op);
+      ops_mapper[op] = new_op;
+    }
+    // Construct Base information for new Group
+    auto new_group = std::make_shared<Group>(new_ops);
+    this->CollectOps();
+    for (auto& iter : this->input_ops) {
+      new_group->input_ops[ops_mapper[iter.first]] = iter.second;
+    }
+    for (auto* op : this->output_ops) {
+      new_group->output_ops.insert(ops_mapper[op]);
+    }
+
+    return new_group;
+  }
+
   // distance to last group.
   int depth{0};
   int max_depth{0};

paddle/pir/core/ir_mapping.h

@@ -0,0 +1,37 @@
+// 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_map>
+#include "paddle/common/enforce.h"
+#include "paddle/pir/core/block.h"
+
+namespace pir {
+
+class IrMapping {
+ public:
+  void Add(Value from, Value to) { value_map_[from] = to; }
+
+  Value Lookup(Value from) const {
+    IR_ENFORCE(value_map_.count(from) > 0, "Not Found Value in IRMapping.");
+    return value_map_.at(from);
+  }
+  void Earse(Value from) { value_map_.erase(from); }
+
+  void Clear() { value_map_.clear(); }
+
+ private:
+  std::unordered_map<Value, Value> value_map_;
+};
+
+}  // namespace pir

paddle/pir/core/operation.cc

@@ -137,6 +137,32 @@ Operation *Operation::Create(const std::vector<Value> &inputs,
   return op;
 }
 
+Operation *Operation::Clone(IrMapping &ir_mapping, CloneOptions options) {
+  IR_ENFORCE(options.IsCloneRegions() || num_regions_ > 0,
+             "Operation CloneOperands is unimplemented currently.");
+  IR_ENFORCE(num_successors_ == 0,
+             "Operation::Clone is not unimplemented for multiple successors.");
+
+  auto inputs = operands_source();
+  if (options.IsCloneOperands()) {
+    // replace value by IRMapping inplacely.
+    for (auto &value : inputs) {
+      value = ir_mapping.Lookup(value);
+    }
+  }
+
+  std::vector<Type> output_types;
+  for (auto &result : results()) {
+    output_types.push_back(result.type());
+  }
+  auto *new_op = Create(inputs, attributes_, output_types, info_, num_regions_);
+  // record outputs mapping info
+  for (uint32_t i = 0; i < num_results_; ++i) {
+    ir_mapping.Add(result(i), new_op->result(i));
+  }
+  return new_op;
+}
+
 // Call destructors for Region , OpResults, Operation, and OpOperands in
 // sequence, and finally free memory.
 void Operation::Destroy() {

paddle/pir/core/operation.h

@@ -20,12 +20,12 @@
 #include "paddle/common/enforce.h"
 #include "paddle/common/macros.h"
 #include "paddle/pir/core/block.h"
+#include "paddle/pir/core/ir_mapping.h"
 #include "paddle/pir/core/iterator.h"
 #include "paddle/pir/core/op_info.h"
 #include "paddle/pir/core/operation_utils.h"
 #include "paddle/pir/core/type.h"
 #include "paddle/pir/core/visitors.h"
-
 namespace pir {
 class OpBase;
 class Program;
@@ -37,6 +37,20 @@ class OpResultImpl;
 class OpOperendImpl;
 }  // namespace detail
 
+class CloneOptions {
+ public:
+  CloneOptions() : clone_regions_{false}, clone_operands_{false} {}
+  CloneOptions(bool clone_regions, bool clone_operands)
+      : clone_regions_(clone_regions), clone_operands_(clone_operands) {}
+
+  bool IsCloneRegions() const { return clone_regions_; }
+  bool IsCloneOperands() const { return clone_operands_; }
+
+ private:
+  bool clone_regions_{true};
+  bool clone_operands_{true};
+};
+
 class IR_API alignas(8) Operation final
     : public DoubleLevelContainer<Operation> {
  public:
@@ -53,6 +67,12 @@ class IR_API alignas(8) Operation final
                            size_t num_regions = 0,
                            const std::vector<Block *> &successors = {});
   static Operation *Create(OperationArgument &&op_argument);
+
+  ///
+  /// \brief Deep copy all information and create a new operation.
+  ///
+  Operation *Clone(IrMapping &ir_mapping,
+                   CloneOptions options = CloneOptions());
   ///
   /// \brief Destroy the operation objects and free memory by create().
   ///