[MLIR] Move `replaceAllUsesExcept` from LoopUtil.h to Value.h.

[bondhugula]

Thanks!

mlir/include/mlir/IR/Value.h

141

Nit: with the new value -> with newValue

[rriddle]

mlir/include/mlir/IR/Value.h

19

Please don't add includes here, if it is truly necessary move it to the .cpp file.

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

111–112

nit: Change the type of newIndex to AddIOp to avoid needing to recover the defining op.

[bondhugula]

mlir/include/mlir/IR/Value.h

19

This include is actually not necessary. SmallPtrSet and its impl is already fwd decl'ed in mlir/Support/LLVM.h.

[pifon2a]

Addressed the comments.

mlir/include/mlir/IR/Value.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_VALUE_H
 #define MLIR_IR_VALUE_H
 
 #include "mlir/IR/Types.h"
 #include "mlir/IR/UseDefLists.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 class AsmState;
 class BlockArgument;
 class Operation;
 class OpResult;
 class Region;
 class Value;
public:
   /// Drop all uses of this object from their respective owners.
   void dropAllUses() const;
 
   /// Replace all uses of 'this' value with the new value, updating anything in
   /// the IR that uses 'this' to use the other value instead.  When this returns
   /// there are zero uses of 'this'.
   void replaceAllUsesWith(Value newValue) 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
+  /// listed in 'exceptions' .
+  void
+  replaceAllUsesExcept(Value newValue,
+                       const SmallPtrSetImpl<Operation *> &exceptions) const;
+
   //===--------------------------------------------------------------------===//
   // Uses
 
   /// This class implements an iterator over the uses of a value.
   using use_iterator = FilteredValueUseIterator<OpOperand>;
   using use_range = iterator_range<use_iterator>;
 
   use_iterator use_begin() const;

mlir/include/mlir/Transforms/LoopUtils.h

 ///      // full_tile
 ///    else
 ///      // partial tile
 ///
 LogicalResult
 separateFullTiles(MutableArrayRef<AffineForOp> nest,
                   SmallVectorImpl<AffineForOp> *fullTileNest = nullptr);
 
-/// Replaces all uses of `orig` with `replacement` except if the user is listed
-/// in `exceptions`.
-void replaceAllUsesExcept(Value orig, Value replacement,
-                          const SmallPtrSetImpl<Operation *> &exceptions);
-
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_LOOP_UTILS_H

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

 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/FoldUtils.h"
-#include "mlir/Transforms/LoopUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "linalg-fusion"
 
 using namespace mlir;
 using namespace mlir::edsc;
static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op,
   // IV arguments according to the tiling of the consumer, i.e. offset them by
   // the values computed in `loopRanges`.
   if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(clonedOp)) {
     auto &block = indexedGenericOp.region().front();
 
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPointToStart(&block);
     for (unsigned i = 0, e = indexedGenericOp.getNumLoops(); i < e; ++i) {
       Value oldIndex = block.getArgument(i);
-      Value newIndex = b.create<AddIOp>(indexedGenericOp.getLoc(), oldIndex,
+      AddIOp newIndex = b.create<AddIOp>(indexedGenericOp.getLoc(), oldIndex,
                                          loopRanges[i].offset);
-      replaceAllUsesExcept(
-          oldIndex, newIndex,
-          SmallPtrSet<Operation *, 1>{newIndex.getDefiningOp()});
+      oldIndex.replaceAllUsesExcept(newIndex,
+                                    SmallPtrSet<Operation *, 1>{newIndex});
     }
   }
   return clonedOp;
 }
 
 struct ViewDimension {
   Value view;
   unsigned dimension;

mlir/lib/IR/Value.cpp

 //===- Value.cpp - MLIR Value Classes -------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/StandardTypes.h"
+#include "llvm/ADT/SmallPtrSet.h"
 using namespace mlir;
 
 /// Construct a value.
 Value::Value(detail::BlockArgumentImpl *impl)
     : ownerAndKind(impl, Kind::BlockArgument) {}
 Value::Value(Operation *op, unsigned resultNo) {
   assert(op->getNumResults() > resultNo && "invalid result number");
   if (LLVM_LIKELY(canPackResultInline(resultNo))) {
if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->replaceAllUsesWith(newValue);
   Operation *owner = cast<OpResult>().getOwner();
   IRMultiObjectWithUseList<OpOperand> *useList = owner;
   if (owner->hasSingleResult)
     return useList->replaceAllUsesWith(newValue);
   useList->replaceAllUsesWith(*this, 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())) {
+    if (exceptions.count(use.getOwner()) == 0)
+      use.set(newValue);
+  }
+}
+
 //===--------------------------------------------------------------------===//
 // Uses
 
 auto Value::use_begin() const -> use_iterator {
   if (BlockArgument arg = dyn_cast<BlockArgument>())
     return arg.getImpl()->use_begin();
   Operation *owner = cast<OpResult>().getOwner();
   return owner->hasSingleResult ? use_iterator(owner->use_begin())

mlir/lib/Transforms/Utils/LoopUtils.cpp

 #include "mlir/IR/Function.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SetVector.h"
-#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "LoopUtils"
 
 using namespace mlir;
 using llvm::SetVector;
 using llvm::SmallMapVector;
Value scaled =
       isStepOne ? inductionVar
                 : insideLoopBuilder.create<MulIOp>(loc, inductionVar, step);
   Value shifted =
       isZeroBased ? scaled
                   : insideLoopBuilder.create<AddIOp>(loc, scaled, lowerBound);
 
   SmallPtrSet<Operation *, 2> preserve{scaled.getDefiningOp(),
                                        shifted.getDefiningOp()};
-  replaceAllUsesExcept(inductionVar, shifted, preserve);
+  inductionVar.replaceAllUsesExcept(shifted, preserve);
   return {/*lowerBound=*/newLowerBound, /*upperBound=*/newUpperBound,
           /*step=*/newStep};
 }
 
 /// Transform a loop with a strictly positive step
 ///   for %i = %lb to %ub step %s
 /// into a 0-based loop with step 1
 ///   for %ii = 0 to ceildiv(%ub - %lb, %s) step 1 {
elseBlock->getOperations().splice(
       firstLoop.getOperation()->getBlock()->getOperations(),
       Block::iterator(firstLoop));
 
   if (fullTileNest)
     *fullTileNest = std::move(fullTileLoops);
 
   return success();
 }
-
-void mlir::replaceAllUsesExcept(
-    Value orig, Value replacement,
-    const SmallPtrSetImpl<Operation *> &exceptions) {
-  for (auto &use : llvm::make_early_inc_range(orig.getUses())) {
-    if (exceptions.count(use.getOwner()) == 0)
-      use.set(replacement);
-  }
-}