[mlir] Optimize OpResult use case for single result operations.

Summary:
Operation represents all of the uses of each result with one use list, so manipulating the use list of a specific result requires filtering the main use list. This revision adds an optimization for the case of single result operations to avoid this filtering.

Differential Revision: https://reviews.llvm.org/D73430

diff --git a/mlir/include/mlir/IR/UseDefLists.h b/mlir/include/mlir/IR/UseDefLists.h
index 49a7e5a6e3df83a8d9952a111739e4925e16281b..f4433384106c65678cf3d958338a815cc114927b 100644 (file)
--- a/mlir/include/mlir/IR/UseDefLists.h
+++ b/mlir/include/mlir/IR/UseDefLists.h
@@ -131,6 +131,7 @@ public:
     for (OperandType &use : llvm::make_early_inc_range(getUses(oldValue)))
       use.set(newValue);
   }
+  using BaseType::replaceAllUsesWith;
 
   //===--------------------------------------------------------------------===//
   // Uses

diff --git a/mlir/lib/IR/Value.cpp b/mlir/lib/IR/Value.cpp
index 921bcefcdd2e66459bddf75fae29b1ace06e1dbd..2f821c94366d6ccc92cb531beb0d92e414206f22 100644 (file)
--- a/mlir/lib/IR/Value.cpp
+++ b/mlir/lib/IR/Value.cpp
@@ -102,7 +102,10 @@ IRObjectWithUseList<OpOperand> *Value::getUseList() const {
 void Value::dropAllUses() const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->dropAllUses();
-  return cast<OpResult>().getOwner()->dropAllUses(*this);
+  Operation *owner = cast<OpResult>().getOwner();
+  if (owner->hasSingleResult)
+    return owner->dropAllUses();
+  return owner->dropAllUses(*this);
 }
 
 /// Replace all uses of 'this' value with the new value, updating anything in
@@ -111,7 +114,10 @@ void Value::dropAllUses() const {
 void Value::replaceAllUsesWith(Value newValue) const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->replaceAllUsesWith(newValue);
-  IRMultiObjectWithUseList<OpOperand> *useList = cast<OpResult>().getOwner();
+  Operation *owner = cast<OpResult>().getOwner();
+  IRMultiObjectWithUseList<OpOperand> *useList = owner;
+  if (owner->hasSingleResult)
+    return useList->replaceAllUsesWith(newValue);
   useList->replaceAllUsesWith(*this, newValue);
 }
 
@@ -121,21 +127,25 @@ void Value::replaceAllUsesWith(Value newValue) const {
 auto Value::use_begin() const -> use_iterator {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->use_begin();
-  return cast<OpResult>().getOwner()->use_begin(*this);
+  Operation *owner = cast<OpResult>().getOwner();
+  return owner->hasSingleResult ? use_iterator(owner->use_begin())
+                                : owner->use_begin(*this);
 }
 
 /// Returns true if this value has exactly one use.
 bool Value::hasOneUse() const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->hasOneUse();
-  return cast<OpResult>().getOwner()->hasOneUse(*this);
+  Operation *owner = cast<OpResult>().getOwner();
+  return owner->hasSingleResult ? owner->hasOneUse() : owner->hasOneUse(*this);
 }
 
 /// Returns true if this value has no uses.
 bool Value::use_empty() const {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->use_empty();
-  return cast<OpResult>().getOwner()->use_empty(*this);
+  Operation *owner = cast<OpResult>().getOwner();
+  return owner->hasSingleResult ? owner->use_empty() : owner->use_empty(*this);
 }
 
 //===----------------------------------------------------------------------===//