[mlir][NFC] Add helper for common pattern of replaceAllUsesExcept

[silvas]

Summary

This covers the extremely common case of replacing all uses of a Value
with a new op that is itself a user of the original Value.

This should also be a little bit more efficient than the
SmallPtrSet<Operation *, 1>{op} idiom that was being used before.

[bondhugula]

Nice.

mlir/lib/IR/Value.cpp

77

Nit: auto -> OpOperand
(should still fit in one line)

[silvas]

mlir/lib/IR/Value.cpp

77

That would be good but auto is consistent with the other replaceAllUsesExcept functions

[silvas]

mlir/lib/IR/Value.cpp

77

Sorry, I see it is ambiguous with replaceUsesWithIf as well. Will change both replaceAllUsesExcept to OpOperand

mlir/include/mlir/IR/Value.h

public:
 
   /// Replace all uses of 'this' value with 'newValue', updating anything in the
   /// IR that uses 'this' to use the other value instead except if the user is
   /// listed in 'exceptions' .
   void
   replaceAllUsesExcept(Value newValue,
                        const SmallPtrSetImpl<Operation *> &exceptions) const;
 
+  /// Replace all uses of 'this' value with 'newValue', updating anything in the
+  /// IR that uses 'this' to use the other value instead except if the user is
+  /// 'exceptedUser'.
+  void replaceAllUsesExcept(Value newValue, Operation *exceptedUser) const;
+
   /// Replace all uses of 'this' value with 'newValue' if the given callback
   /// returns true.
   void replaceUsesWithIf(Value newValue,
                          function_ref<bool(OpOperand &)> shouldReplace);
 
   /// Returns true if the value is used outside of the given block.
   bool isUsedOutsideOfBlock(Block *block);
 

mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp

for (unsigned i = 0, e = steps.size(); i < e; ++i) {
     // canonicalizations.
     OperandRange lbOperands = op.getLowerBoundsOperands();
     OperandRange dimOperands = lbOperands.take_front(nDims);
     OperandRange symbolOperands = lbOperands.drop_front(nDims);
     SmallVector<Value, 8> applyOperands{dimOperands};
     applyOperands.push_back(iv);
     applyOperands.append(symbolOperands.begin(), symbolOperands.end());
     auto apply = builder.create<AffineApplyOp>(op.getLoc(), map, applyOperands);
-    iv.replaceAllUsesExcept(apply, SmallPtrSet<Operation *, 1>{apply});
+    iv.replaceAllUsesExcept(apply, apply);
   }
 
   SmallVector<int64_t, 8> newSteps(op.getNumDims(), 1);
   op.setSteps(newSteps);
   auto newLowerMap = AffineMap::get(
       /*dimCount=*/0, /*symbolCount=*/0, lbExprs, op.getContext());
   op.setLowerBounds({}, newLowerMap);
   auto newUpperMap = AffineMap::get(ranges.getNumDims(), ranges.getNumSymbols(),
static void normalizeAffineFor(AffineForOp op) {
   for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
     lbOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
 
   AffineExpr origIVExpr = opBuilder.getAffineDimExpr(lb.getMap().getNumDims());
   AffineExpr newIVExpr = origIVExpr * origLoopStep + origLbMap.getResult(0);
   AffineMap ivMap = AffineMap::get(origLbMap.getNumDims() + 1,
                                    origLbMap.getNumSymbols(), newIVExpr);
   Operation *newIV = opBuilder.create<AffineApplyOp>(loc, ivMap, lbOperands);
-  op.getInductionVar().replaceAllUsesExcept(newIV->getResult(0),
-                                            SmallPtrSet<Operation *, 1>{newIV});
+  op.getInductionVar().replaceAllUsesExcept(newIV->getResult(0), newIV);
 }
 
 namespace {
 
 /// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1.
 /// As currently implemented, this pass cannot fail, but it might skip over ops
 /// that are already in a normalized form.
 struct AffineLoopNormalizePass

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

if (producer.hasIndexSemantics()) {
     for (IndexOp indexOp : block.getOps<IndexOp>()) {
       OpBuilder::InsertionGuard g(builder);
       builder.setInsertionPointAfter(indexOp);
       AffineExpr index, offset;
       bindDims(builder.getContext(), index, offset);
       AffineApplyOp applyOp = builder.create<AffineApplyOp>(
           indexOp.getLoc(), index + offset,
           ValueRange{indexOp.getResult(), loopRanges[indexOp.dim()].offset});
-      indexOp.getResult().replaceAllUsesExcept(
-          applyOp, SmallPtrSet<Operation *, 1>{applyOp});
+      indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
     }
   }
 
   return clonedOp;
 }
 
 /// Get the loop range for a dimension `dim` based on the `shapedOperand`. It is
 /// expected to be defined by a subview op or a subtensor op.

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

for (IndexOp indexOp : block.getOps<linalg::IndexOp>()) {
     // replace all uses of the previous value.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPointAfter(indexOp);
     AffineExpr index, iv;
     bindDims(b.getContext(), index, iv);
     AffineApplyOp applyOp = b.create<AffineApplyOp>(
         indexOp.getLoc(), index + iv,
         ValueRange{indexOp.getResult(), ivs[rangeIndex->second]});
-    indexOp.getResult().replaceAllUsesExcept(
-        applyOp.getResult(), SmallPtrSet<Operation *, 1>{applyOp});
+    indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
   }
 }
 
 template <typename LoopTy>
 static Optional<TiledLinalgOp>
 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ValueRange tileSizes,
                  const LinalgTilingOptions &options) {
   auto nLoops = op.getNumLoops();

mlir/lib/Dialect/SCF/Transforms/ParallelLoopTiling.cpp

mlir::scf::tileParallelLoop(ParallelOp op, ArrayRef<int64_t> tileSizes) {
 
   // Insert computation for new index vectors and replace uses.
   b.setInsertionPointToStart(innerLoop.getBody());
   for (auto ivs :
        llvm::zip(innerLoop.getInductionVars(), outerLoop.getInductionVars())) {
     Value inner_index = std::get<0>(ivs);
     AddIOp newIndex =
         b.create<AddIOp>(op.getLoc(), std::get<0>(ivs), std::get<1>(ivs));
-    inner_index.replaceAllUsesExcept(
-        newIndex, SmallPtrSet<Operation *, 1>{newIndex.getOperation()});
+    inner_index.replaceAllUsesExcept(newIndex, newIndex);
   }
 
   op.erase();
   return std::make_pair(outerLoop, innerLoop);
 }
 
 namespace {
 struct ParallelLoopTiling

mlir/lib/IR/Value.cpp

void Value::replaceAllUsesWith(Value newValue) const {
   return getUseList()->replaceAllUsesWith(newValue);
 }
 
 /// Replace all uses of 'this' value with the new value, updating anything in
 /// the IR that uses 'this' to use the other value instead except if the user is
 /// listed in 'exceptions' .
 void Value::replaceAllUsesExcept(
     Value newValue, const SmallPtrSetImpl<Operation *> &exceptions) const {
-  for (auto &use : llvm::make_early_inc_range(getUses())) {
+  for (OpOperand &use : llvm::make_early_inc_range(getUses())) {
     if (exceptions.count(use.getOwner()) == 0)
       use.set(newValue);
   }
 }
 
+/// Replace all uses of 'this' value with 'newValue', updating anything in the
+/// IR that uses 'this' to use the other value instead except if the user is
+/// 'exceptedUser'.
+void Value::replaceAllUsesExcept(Value newValue,
+                                 Operation *exceptedUser) const {
+  for (OpOperand &use : llvm::make_early_inc_range(getUses())) {
+    if (use.getOwner() != exceptedUser)
+      use.set(newValue);
+  }
+}
+
 /// Replace all uses of 'this' value with 'newValue' if the given callback
 /// returns true.
 void Value::replaceUsesWithIf(Value newValue,
                               function_ref<bool(OpOperand &)> shouldReplace) {
   for (OpOperand &use : llvm::make_early_inc_range(getUses()))
     if (shouldReplace(use))
       use.set(newValue);
 }