diff --git a/mlir/include/mlir/IR/UseDefLists.h b/mlir/include/mlir/IR/UseDefLists.h --- a/mlir/include/mlir/IR/UseDefLists.h +++ b/mlir/include/mlir/IR/UseDefLists.h @@ -131,6 +131,7 @@ 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 --- a/mlir/lib/IR/Value.cpp +++ b/mlir/lib/IR/Value.cpp @@ -102,7 +102,10 @@ void Value::dropAllUses() const { if (BlockArgument arg = dyn_cast()) return arg.getImpl()->dropAllUses(); - return cast().getOwner()->dropAllUses(*this); + Operation *owner = cast().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::replaceAllUsesWith(Value newValue) const { if (BlockArgument arg = dyn_cast()) return arg.getImpl()->replaceAllUsesWith(newValue); - IRMultiObjectWithUseList *useList = cast().getOwner(); + Operation *owner = cast().getOwner(); + IRMultiObjectWithUseList *useList = owner; + if (owner->hasSingleResult) + return useList->replaceAllUsesWith(newValue); useList->replaceAllUsesWith(*this, newValue); } @@ -121,21 +127,25 @@ auto Value::use_begin() const -> use_iterator { if (BlockArgument arg = dyn_cast()) return arg.getImpl()->use_begin(); - return cast().getOwner()->use_begin(*this); + Operation *owner = cast().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()) return arg.getImpl()->hasOneUse(); - return cast().getOwner()->hasOneUse(*this); + Operation *owner = cast().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()) return arg.getImpl()->use_empty(); - return cast().getOwner()->use_empty(*this); + Operation *owner = cast().getOwner(); + return owner->hasSingleResult ? owner->use_empty() : owner->use_empty(*this); } //===----------------------------------------------------------------------===//