[MLIR] Add replaceUsesWithIf on Operation

[bondhugula]

Summary

Add replaceUsesWithIf on Operation along the lines of
Value::replaceUsesWithIf. This had been missing on Operation and is
convenient to replace multi-result operations' results conditionally.

[bondhugula]

Fix spacing.

[lattner]

Seems reasonable to me

mlir/include/mlir/IR/Operation.h

public:
   void replaceUsesOfWith(Value from, Value to);
 
   /// Replace all uses of results of this operation with the provided 'values'.
   template <typename ValuesT>
   void replaceAllUsesWith(ValuesT &&values) {
     getResults().replaceAllUsesWith(std::forward<ValuesT>(values));
   }
 
+  /// Replace uses of results of this operation with the provided `values` if
+  /// the given callback returns true.
+  template <typename ValuesT>
+  void replaceUsesWithIf(ValuesT &&values,
+                         function_ref<bool(OpOperand &)> shouldReplace) {
+    getResults().replaceUsesWithIf(std::forward<ValuesT>(values),
+                                   shouldReplace);
+  }
+
   /// Destroys this operation and its subclass data.
   void destroy();
 
   /// This drops all operand uses from this operation, which is an essential
   /// step in breaking cyclic dependences between references when they are to
   /// be deleted.
   void dropAllReferences();
 

mlir/include/mlir/IR/ValueRange.h

replaceAllUsesWith(ValuesT &&values) {
 
     for (auto it : llvm::zip(*this, values))
       std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
   }
 
   /// Replace all uses of results of this range with results of 'op'.
   void replaceAllUsesWith(Operation *op);
 
+  /// Replace uses of results of this range with the provided 'values' if the
+  /// given callback returns true. The size of `values` must match the size of
+  /// this range.
+  template <typename ValuesT>
+  std::enable_if_t<!std::is_convertible<ValuesT, Operation *>::value>
+  replaceUsesWithIf(ValuesT &&values,
+                    function_ref<bool(OpOperand &)> shouldReplace) {
+    assert(static_cast<size_t>(std::distance(values.begin(), values.end())) ==
+               size() &&
+           "expected 'values' to correspond 1-1 with the number of results");
+
+    for (auto it : llvm::zip(*this, values))
+      std::get<0>(it).replaceUsesWithIf(std::get<1>(it), shouldReplace);
+  }
+
+  /// Replace uses of results of this range with results of `op` if the given
+  /// callback returns true.
+  void replaceUsesWithIf(Operation *op,
+                         function_ref<bool(OpOperand &)> shouldReplace);
+
   //===--------------------------------------------------------------------===//
   // Users
   //===--------------------------------------------------------------------===//
 
   using user_iterator = ValueUserIterator<use_iterator, OpOperand>;
   using user_range = iterator_range<user_iterator>;
 
   /// Returns a range of all users.

mlir/lib/IR/OperationSupport.cpp

void ResultRange::UseIterator::skipOverResultsWithNoUsers() {
   else
     use = (*it).use_begin();
 }
 
 void ResultRange::replaceAllUsesWith(Operation *op) {
   replaceAllUsesWith(op->getResults());
 }
 
+void ResultRange::replaceUsesWithIf(
+    Operation *op, function_ref<bool(OpOperand &)> shouldReplace) {
+  replaceUsesWithIf(op->getResults(), shouldReplace);
+}
+
 //===----------------------------------------------------------------------===//
 // ValueRange
 
 ValueRange::ValueRange(ArrayRef<Value> values)
     : ValueRange(values.data(), values.size()) {}
 ValueRange::ValueRange(OperandRange values)
     : ValueRange(values.begin().getBase(), values.size()) {}
 ValueRange::ValueRange(ResultRange values)

mlir/lib/Transforms/CSE.cpp

void CSE::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
   if (hasSSADominance) {
     // If the region has SSA dominance, then we are guaranteed to have not
     // visited any use of the current operation.
     op->replaceAllUsesWith(existing);
     opsToErase.push_back(op);
   } else {
     // When the region does not have SSA dominance, we need to check if we
     // have visited a use before replacing any use.
-    for (auto it : llvm::zip(op->getResults(), existing->getResults())) {
-      std::get<0>(it).replaceUsesWithIf(
-          std::get<1>(it), [&](OpOperand &operand) {
-            return !knownValues.count(operand.getOwner());
-          });
-    }
+    op->replaceUsesWithIf(existing->getResults(), [&](OpOperand &operand) {
 
     // There may be some remaining uses of the operation.
     if (op->use_empty())
       opsToErase.push_back(op);
   }
 
   // If the existing operation has an unknown location and the current
   // operation doesn't, then set the existing op's location to that of the