[MLIR] NFC Clean up doc comments on memref replacement utility

[bondhugula]

Summary

NFC. Clean up stale doc comments on memref replacement utility and some
variable renaming in it to avoid confusion.

mlir/include/mlir/Transforms/Utils.h

 ///
 /// `indexRemap` remaps indices of the old memref access to a new set of indices
 /// that are used to index the memref. Additional input operands to indexRemap
 /// can be optionally provided in `extraOperands`, and they occupy the start
 /// of its input list. `indexRemap`'s dimensional inputs are expected to
 /// correspond to memref's indices, and its symbolic inputs if any should be
 /// provided in `symbolOperands`.
 ///
-/// `domInstFilter`, if non-null, restricts the replacement to only those
-/// operations that are dominated by the former; similarly, `postDomInstFilter`
+/// `domOpFilter`, if non-null, restricts the replacement to only those
+/// operations that are dominated by the former; similarly, `postDomOpFilter`
 /// restricts replacement to only those operations that are postdominated by it.
 ///
 /// 'allowNonDereferencingOps', if set, allows replacement of non-dereferencing
-/// uses of a memref without any requirement for access index rewrites. The
-/// default value of this flag variable is false.
+/// uses of a memref without any requirement for access index rewrites as long
+/// as the user operation has the MemRefsNormalizable trait. The default value
+/// of this flag is false.
 ///
 /// 'replaceInDeallocOp', if set, lets DeallocOp, a non-dereferencing user, to
 /// also be a candidate for replacement. The default value of this flag is
 /// false.
 ///
 /// Returns true on success and false if the replacement is not possible,
 /// whenever a memref is used as an operand in a non-dereferencing context and
 /// 'allowNonDereferencingOps' is false, except for dealloc's on the memref
 /// which are left untouched. See comments at function definition for an
 /// example.
 //
 //  Ex: to replace load %A[%i, %j] with load %Abuf[%t mod 2, %ii - %i, %j]:
 //  The SSA value corresponding to '%t mod 2' should be in 'extraIndices', and
 //  index remap will perform (%i, %j) -> (%ii - %i, %j), i.e., indexRemap = (d0,
 //  d1, d2) -> (d0 - d1, d2), and %ii will be the extra operand. Without any
 //  extra operands, note that 'indexRemap' would just be applied to existing
 //  indices (%i, %j).
 //  TODO: allow extraIndices to be added at any position.
 LogicalResult replaceAllMemRefUsesWith(
     Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices = {},
     AffineMap indexRemap = AffineMap(), ArrayRef<Value> extraOperands = {},
-    ArrayRef<Value> symbolOperands = {}, Operation *domInstFilter = nullptr,
-    Operation *postDomInstFilter = nullptr,
-    bool allowNonDereferencingOps = false, bool replaceInDeallocOp = false);
+    ArrayRef<Value> symbolOperands = {}, Operation *domOpFilter = nullptr,
+    Operation *postDomOpFilter = nullptr, bool allowNonDereferencingOps = false,
+    bool replaceInDeallocOp = false);
 
 /// Performs the same replacement as the other version above but only for the
 /// dereferencing uses of `oldMemRef` in `op`, except in cases where
 /// 'allowNonDereferencingOps' is set to true where we replace the
 /// non-dereferencing uses as well.
 LogicalResult replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                        Operation *op,
                                        ArrayRef<Value> extraIndices = {},

mlir/lib/Transforms/Utils/Utils.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TypeSwitch.h"
+
+#define DEBUG_TYPE "transforms-utils"
+
 using namespace mlir;
 
 // Perform the replacement in `op`.
 LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
                                              Operation *op,
                                              ArrayRef<Value> extraIndices,
                                              AffineMap indexRemap,
                                              ArrayRef<Value> extraOperands,
LogicalResult mlir::replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
   op->erase();
 
   return success();
 }
 
 LogicalResult mlir::replaceAllMemRefUsesWith(
     Value oldMemRef, Value newMemRef, ArrayRef<Value> extraIndices,
     AffineMap indexRemap, ArrayRef<Value> extraOperands,
-    ArrayRef<Value> symbolOperands, Operation *domInstFilter,
-    Operation *postDomInstFilter, bool allowNonDereferencingOps,
+    ArrayRef<Value> symbolOperands, Operation *domOpFilter,
+    Operation *postDomOpFilter, bool allowNonDereferencingOps,
     bool replaceInDeallocOp) {
   unsigned newMemRefRank = newMemRef.getType().cast<MemRefType>().getRank();
   (void)newMemRefRank; // unused in opt mode
   unsigned oldMemRefRank = oldMemRef.getType().cast<MemRefType>().getRank();
   (void)oldMemRefRank;
   if (indexRemap) {
     assert(indexRemap.getNumSymbols() == symbolOperands.size() &&
            "symbol operand count mismatch");
     assert(indexRemap.getNumInputs() ==
            extraOperands.size() + oldMemRefRank + symbolOperands.size());
     assert(indexRemap.getNumResults() + extraIndices.size() == newMemRefRank);
   } else {
     assert(oldMemRefRank + extraIndices.size() == newMemRefRank);
   }
 
   // Assert same elemental type.
   assert(oldMemRef.getType().cast<MemRefType>().getElementType() ==
          newMemRef.getType().cast<MemRefType>().getElementType());
 
   std::unique_ptr<DominanceInfo> domInfo;
   std::unique_ptr<PostDominanceInfo> postDomInfo;
-  if (domInstFilter)
-    domInfo = std::make_unique<DominanceInfo>(
-        domInstFilter->getParentOfType<FuncOp>());
+  if (domOpFilter)
+    domInfo =
+        std::make_unique<DominanceInfo>(domOpFilter->getParentOfType<FuncOp>());
 
-  if (postDomInstFilter)
+  if (postDomOpFilter)
     postDomInfo = std::make_unique<PostDominanceInfo>(
-        postDomInstFilter->getParentOfType<FuncOp>());
+        postDomOpFilter->getParentOfType<FuncOp>());
 
   // Walk all uses of old memref; collect ops to perform replacement. We use a
   // DenseSet since an operation could potentially have multiple uses of a
   // memref (although rare), and the replacement later is going to erase ops.
   DenseSet<Operation *> opsToReplace;
   for (auto *op : oldMemRef.getUsers()) {
-    // Skip this use if it's not dominated by domInstFilter.
-    if (domInstFilter && !domInfo->dominates(domInstFilter, op))
+    // Skip this use if it's not dominated by domOpFilter.
+    if (domOpFilter && !domInfo->dominates(domOpFilter, op))
       continue;
 
-    // Skip this use if it's not post-dominated by postDomInstFilter.
-    if (postDomInstFilter && !postDomInfo->postDominates(postDomInstFilter, op))
+    // Skip this use if it's not post-dominated by postDomOpFilter.
+    if (postDomOpFilter && !postDomInfo->postDominates(postDomOpFilter, op))
       continue;
 
     // Skip dealloc's - no replacement is necessary, and a memref replacement
     // at other uses doesn't hurt these dealloc's.
     if (isa<memref::DeallocOp>(op) && !replaceInDeallocOp)
       continue;
 
     // Check if the memref was used in a non-dereferencing context. It is fine
     // for the memref to be used in a non-dereferencing way outside of the
     // region where this replacement is happening.
     if (!isa<AffineMapAccessInterface>(*op)) {
-      if (!allowNonDereferencingOps)
+      if (!allowNonDereferencingOps) {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "Memref replacement failed: non-deferencing memref op: \n"
+                   << *op << '\n');
         return failure();
-      // Currently we support the following non-dereferencing ops to be a
-      // candidate for replacement: Dealloc, CallOp and ReturnOp.
-      // TODO: Add support for other kinds of ops.
-      if (!op->hasTrait<OpTrait::MemRefsNormalizable>())
+      }
+      // Non-dereferencing ops with the MemRefsNormalizable trait are
+      // supported for replacement.
+      if (!op->hasTrait<OpTrait::MemRefsNormalizable>()) {
+        LLVM_DEBUG(llvm::dbgs() << "Memref replacement failed: use without a "
+                                   "memrefs normalizable trait: \n"
+                                << *op << '\n');
         return failure();
       }
+    }
 
     // We'll first collect and then replace --- since replacement erases the op
     // that has the use, and that op could be postDomFilter or domFilter itself!
     opsToReplace.insert(op);
   }
 
   for (auto *op : opsToReplace) {
     if (failed(replaceAllMemRefUsesWith(
newAlloc = b.create<memref::AllocOp>(allocOp->getLoc(), newMemRefType,
                                          allocOp->alignmentAttr());
   }
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
                                       /*extraIndices=*/{},
                                       /*indexRemap=*/layoutMap,
                                       /*extraOperands=*/{},
                                       /*symbolOperands=*/symbolOperands,
-                                      /*domInstFilter=*/nullptr,
-                                      /*postDomInstFilter=*/nullptr,
+                                      /*domOpFilter=*/nullptr,
+                                      /*postDomOpFilter=*/nullptr,
                                       /*allowDereferencingOps=*/true))) {
     // If it failed (due to escapes for example), bail out.
     newAlloc.erase();
     return failure();
   }
   // Replace any uses of the original alloc op and erase it. All remaining uses
   // have to be dealloc's; RAMUW above would've failed otherwise.
   assert(llvm::all_of(oldMemRef.getUsers(), [](Operation *op) {