[PIR]Move Operation::operand_index() into Operand::index() (#59791)

Aurelius84 · web-flow · commit 65df4729284c · 2023-12-08T10:22:54.000+08:00
* [PIR]Move Operation::operand_index() into Operand::index()

* fix order

* add IR_ENFORCE

* fix header
diff --git a/paddle/fluid/pir/transforms/transform_general_functions.cc b/paddle/fluid/pir/transforms/transform_general_functions.cc
@@ -84,8 +84,7 @@ std::vector<std::pair<Operation*, int32_t>> GetUseOpsForOutput(Operation* op,
   auto result = op->result(index);
   std::vector<std::pair<Operation*, int32_t>> use_ops;
   for (auto it = result.use_begin(); it != result.use_end(); ++it) {
-    use_ops.push_back(
-        std::make_pair(it->owner(), it->owner()->operand_index(*it)));
+    use_ops.push_back(std::make_pair(it->owner(), it->index()));
   }
   return use_ops;
 }
diff --git a/paddle/pir/core/op_operand.cc b/paddle/pir/core/op_operand.cc
@@ -53,6 +53,11 @@ Operation *OpOperand::owner() const {
   return impl_->owner();
 }
 
+uint32_t OpOperand::index() const {
+  CHECK_OPOPEREND_NULL_IMPL(index);
+  return impl_->index();
+}
+
 void OpOperand::RemoveFromUdChain() {
   CHECK_OPOPEREND_NULL_IMPL(RemoveFromUdChain);
   return impl_->RemoveFromUdChain();
diff --git a/paddle/pir/core/op_operand.h b/paddle/pir/core/op_operand.h
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 #pragma once
-
+#include <cstdint>
 #include "paddle/pir/core/dll_decl.h"
 
 namespace pir {
@@ -57,6 +57,8 @@ class IR_API OpOperand {
 
   Operation *owner() const;
 
+  uint32_t index() const;
+
   void RemoveFromUdChain();
 
   friend Operation;
diff --git a/paddle/pir/core/op_operand_impl.cc b/paddle/pir/core/op_operand_impl.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include "paddle/pir/core/op_operand_impl.h"
+#include "paddle/pir/core/operation.h"
 #include "paddle/pir/core/value_impl.h"
 
 namespace pir {
@@ -40,6 +41,13 @@ OpOperandImpl::OpOperandImpl(pir::Value source, pir::Operation *owner)
   InsertToUdChain();
 }
 
+uint32_t OpOperandImpl::index() const {
+  const char *start =
+      reinterpret_cast<const char *>(owner_) + sizeof(Operation);
+  const char *end = reinterpret_cast<const char *>(this);
+  return (end - start) / sizeof(OpOperandImpl);
+}
+
 void OpOperandImpl::InsertToUdChain() {
   prev_use_addr_ = source_.impl()->first_use_addr();
   next_use_ = source_.impl()->first_use();
diff --git a/paddle/pir/core/op_operand_impl.h b/paddle/pir/core/op_operand_impl.h
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+#include <cstdint>
 #include "paddle/pir/core/value.h"
 
 namespace pir {
@@ -33,6 +34,8 @@ class OpOperandImpl {
 
   void set_source(Value value);
 
+  uint32_t index() const;
+
   /// Remove this op_operand from the current use list.
   void RemoveFromUdChain();
 
diff --git a/paddle/pir/core/operation.cc b/paddle/pir/core/operation.cc
@@ -245,17 +245,6 @@ std::vector<Value> Operation::operands_source() const {
   return res;
 }
 
-int32_t Operation::operand_index(const OpOperand &op_operand) const {
-  int32_t res = -1;
-  for (uint32_t i = 0; i < num_operands(); ++i) {
-    if (op_operand == operand(i)) {
-      res = i;
-      break;
-    }
-  }
-  return res;
-}
-
 ///
 /// \brief op successor related public interfaces
 ///
diff --git a/paddle/pir/core/operation.h b/paddle/pir/core/operation.h
@@ -100,7 +100,6 @@ class IR_API alignas(8) Operation final
   std::vector<OpOperand> operands();
   Value operand_source(uint32_t index) const;
   std::vector<Value> operands_source() const;
-  int32_t operand_index(const OpOperand &op_operand) const;
 
   ///
   /// \brief op successor related public interfaces

Original file line number	Diff line number	Diff line change
`@@ -84,8 +84,7 @@ std::vector<std::pair<Operation, int32_t>> GetUseOpsForOutput(Operation op,`
`84`	`84`	`auto result = op->result(index);`
`85`	`85`	`std::vector<std::pair<Operation*, int32_t>> use_ops;`
`86`	`86`	`for (auto it = result.use_begin(); it != result.use_end(); ++it) {`
`87`		`- use_ops.push_back(`
`88`		`- std::make_pair(it->owner(), it->owner()->operand_index(*it)));`
	`87`	`+ use_ops.push_back(std::make_pair(it->owner(), it->index()));`
`89`	`88`	`}`
`90`	`89`	`return use_ops;`
`91`	`90`	`}`