[Operator] add origin_id for kernel op (#67725)

chen2016013 · web-flow · commit cfd6fd08bb84 · 2024-08-28T15:42:55.000+08:00
* add origin_id for kernel op

* update

* update

* Update ir_printer.cc

* update
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -1673,7 +1673,7 @@ void PirInterpreter::TraceRunInstructionList(
     InstructionBase* instr_node = vec_instruction_base_.at(instr_id).get();
 
     VLOG(6) << "Run InstructionBase " << instr_node->Name() << "[" << instr_id
-            << "]";
+            << "], op id: " << instr_node->Operation()->id();
     RunInstructionBase(instr_node);
 
     if (UNLIKELY(exception_holder_.IsCaught())) {
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -2146,8 +2146,11 @@ void HandleForSpecialOp(
 
   pir::OpInfo op_info = ctx->GetRegisteredOpInfo(op_item->name());
   // Generate new op
+
   pir::Operation* op = pir::Operation::Create(
       vec_inputs, op_item->attributes(), op_output_types, op_info);
+  op->set_attribute("origin_id", pir::Int64Attribute::get(ctx, op_item->id()));
+
   block->push_back(op);
   (*map_op_pair)[op_item] = op;
   // only deal with single output
@@ -2317,6 +2320,9 @@ void HandleForCustomOp(
     op_attribute.emplace("is_inplace", pir::BoolAttribute::get(ctx, true));
   }
 
+  op_attribute.emplace("origin_id",
+                       pir::Int64Attribute::get(ctx, op_item->id()));
+
   VLOG(6) << "Lower custom op: " << op_item->name()
           << " to : " << CustomKernelOp::name();
 
@@ -3003,6 +3009,9 @@ pir::Operation* BuildKernelOp(
     op_attribute.emplace("is_inplace", pir::BoolAttribute::get(ctx, true));
   }
 
+  op_attribute.emplace("origin_id",
+                       pir::Int64Attribute::get(ctx, op_item->id()));
+
   pir::Operation* op = nullptr;
 #ifdef PADDLE_WITH_DNNL
   if (op_item->HasTrait<OneDNNTrait>()) {
diff --git a/paddle/pir/include/core/operation.h b/paddle/pir/include/core/operation.h
@@ -288,7 +288,7 @@ class IR_API alignas(8) Operation final
   const uint32_t num_operands_ = 0;
   const uint32_t num_regions_ = 0;
   const uint32_t num_successors_ = 0;
-  const uint64_t id_;
+  const uint64_t id_ = -1;
 
   detail::BlockOperandImpl *block_operands_{nullptr};
   Region *regions_{nullptr};
diff --git a/paddle/pir/src/core/ir_printer.cc b/paddle/pir/src/core/ir_printer.cc
@@ -17,6 +17,7 @@
 #include <string>
 #include <unordered_map>
 
+#include "glog/logging.h"
 #include "paddle/common/flags.h"
 #include "paddle/pir/include/core/block.h"
 #include "paddle/pir/include/core/builtin_attribute.h"
@@ -189,7 +190,7 @@ void IrPrinter::PrintOperationWithNoRegion(Operation* op) {
 
   os << " \"" << op->name() << "\"";
 
-  if (FLAGS_pir_debug) {
+  if (VLOG_IS_ON(1) || FLAGS_pir_debug) {
     os << " [id:" << op->id() << "]";
   }
 
