[mlir][bufferize][NFC] Better debug output for One-Shot Analysis

[springerm]

Summary

Run mlir-opt with -debug-only="one-shot-analysis" for detailed debug output.

[springerm]

rebase

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
+// Run mlir-opt with `-debug-only="one-shot-analysis"` for detailed debug
+// output.
+#define DEBUG_TYPE "one-shot-analysis"
+
 using namespace mlir;
 using namespace mlir::bufferization;
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 // These are for testing and debugging only. Bufferization information is

 static bool hasReadAfterWriteInterference(
     const DenseSet<OpOperand *> &usesRead,
     const DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo,
     AnalysisState &state, const BufferizationAliasInfo &aliasInfo) {
   const BufferizationOptions &options = state.getOptions();
 
   // Check if op dominance can be used to rule out read-after-write conflicts.
   bool useDominance = canUseOpDominance(usesRead, usesWrite, state);
+  LLVM_DEBUG(llvm::dbgs() << "\n- useDominance = " << useDominance << "\n");
 
   for (OpOperand *uRead : usesRead) {
     Operation *readingOp = uRead->getOwner();
 
     // Find most recent writes of uRead by following the SSA use-def chain.
     // E.g.:
     //
     // %0 = "writing_op"(%t) : tensor<?x32> -> tensor<?xf32>
     // %1 = "aliasing_op"(%0) : tensor<?x32> -> tensor<?xf32>
     // %2 = "reading_op"(%1) : : tensor<?x32> -> not_a_tensor_type
     //
     // In the above example, if uRead is the OpOperand of reading_op, lastWrite
     // is %0. Note that operations that create an alias but do not write (such
     // as ExtractSliceOp) are skipped.
     SetVector<Value> lastWrites = state.findLastPrecedingWrite(uRead->get());
 
     // Look for conflicting memory writes. Potential conflicts are writes to an
     // alias that have been decided to bufferize inplace.
     for (OpOperand *uConflictingWrite : usesWrite) {
+      LLVM_DEBUG(llvm::dbgs() << "\n- check conflict:\n");
+      LLVM_DEBUG(llvm::dbgs()
+                 << "  uRead = operand " << uRead->getOperandNumber() << " of "
+                 << *uRead->getOwner() << "\n");
+      LLVM_DEBUG(llvm::dbgs() << "  unConflictingWrite = operand "
+                              << uConflictingWrite->getOperandNumber() << " of "
+                              << *uConflictingWrite->getOwner() << "\n");
+
       // Throughout this loop, check for multiple requirements that have to be
       // met for uConflictingWrite to be an actual conflict.
       Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
       // Inside of repetitive regions, ops may be executed multiple times and op
       // dominance cannot be used to rule out conflicts.
       if (useDominance) {
         // No conflict if the readingOp dominates conflictingWritingOp, i.e.,
         // the write is not visible when reading.
         //
         // Note: If ops are executed multiple times (e.g., because they are
         //       inside a loop), there may be no meaningful `happensBefore`
         //       relationship.
-        if (happensBefore(readingOp, conflictingWritingOp, domInfo))
+        if (happensBefore(readingOp, conflictingWritingOp, domInfo)) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "  no conflict: read happens before write\n");
           continue;
+        }
 
         // No conflict if the reading use equals the use of the conflicting
         // write. A use cannot conflict with itself.
         //
         // Note: Just being the same op is not enough. It has to be the same
         //       use.
         // Note: If the op is executed multiple times (e.g., because it is
         //       inside a loop), it may be conflicting with itself.
-        if (uConflictingWrite == uRead)
+        if (uConflictingWrite == uRead) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "  no conflict: read and write are same use\n");
           continue;
+        }
 
         // Ops are not conflicting if they are in mutually exclusive regions.
         //
         // Note: If ops are executed multiple times (e.g., because they are
         //       inside a loop), mutually exclusive regions may be executed
         //       multiple times.
-        if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
+        if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp)) {
+          LLVM_DEBUG(llvm::dbgs() << "  no conflict: read and write are in "
+                                     "mutually exclusive regions\n");
           continue;
         }
+      }
 
       // No conflict if the op interface says so.
-      if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp))
-        if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state))
+      if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp)) {
+        if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state)) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "  no conflict: op interace of reading op says 'no'\n");
           continue;
+        }
+      }
 
-      if (conflictingWritingOp != readingOp)
+      if (conflictingWritingOp != readingOp) {
         if (auto bufferizableOp =
-                options.dynCastBufferizableOp(conflictingWritingOp))
-          if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state))
+                options.dynCastBufferizableOp(conflictingWritingOp)) {
+          if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite,
+                                              state)) {
+            LLVM_DEBUG(
+                llvm::dbgs()
+                << "  no conflict: op interace of writing op says 'no'\n");
             continue;
+          }
+        }
+      }
 
       // Check all possible last writes.
       for (Value lastWrite : lastWrites) {
+        LLVM_DEBUG(llvm::dbgs() << "  * lastWrite = " << lastWrite << "\n");
+
         // No conflict if the conflicting write happens before the last
         // write.
         if (Operation *writingOp = lastWrite.getDefiningOp()) {
-          if (happensBefore(conflictingWritingOp, writingOp, domInfo))
+          if (happensBefore(conflictingWritingOp, writingOp, domInfo)) {
             // conflictingWritingOp happens before writingOp. No conflict.
+            LLVM_DEBUG(llvm::dbgs()
+                       << "    no conflict: write happens before last write\n");
             continue;
+          }
           // No conflict if conflictingWritingOp is contained in writingOp.
-          if (writingOp->isProperAncestor(conflictingWritingOp))
+          if (writingOp->isProperAncestor(conflictingWritingOp)) {
+            LLVM_DEBUG(
+                llvm::dbgs()
+                << "    no conflict: write is contained in last write\n");
             continue;
+          }
         } else {
           auto bbArg = lastWrite.cast<BlockArgument>();
           Block *block = bbArg.getOwner();
-          if (!block->findAncestorOpInBlock(*conflictingWritingOp))
+          if (!block->findAncestorOpInBlock(*conflictingWritingOp)) {
+            LLVM_DEBUG(llvm::dbgs() << "    no conflict: last write is bbArg "
+                                       "and write happens outside of block\n");
             // conflictingWritingOp happens outside of the block. No
             // conflict.
             continue;
           }
+        }
 
         // No conflict if the conflicting write and the last write are the same
         // use.
         SmallVector<OpResult> aliasingOpResult =
             state.getAliasingOpResult(*uConflictingWrite);
-        if (aliasingOpResult.size() == 1 && aliasingOpResult[0] == lastWrite)
+        if (aliasingOpResult.size() == 1 && aliasingOpResult[0] == lastWrite) {
+          LLVM_DEBUG(llvm::dbgs()
+                     << "    no conflict: last write and write are same\n");
           continue;
+        }
 
         // All requirements are met. Conflict found!
 
         if (options.printConflicts)
           annotateConflict(uRead, uConflictingWrite, lastWrite);
-
+        LLVM_DEBUG(llvm::dbgs() << "  => RaW CONFLICT FOUND\n");
         return true;
       }
     }
   }
 
   return false;
 }
 
for (OpResult result : state.getAliasingOpResult(operand)) {
     getAliasingInplaceWrites(usesWrite, result, aliasInfo, state);
   }
   if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
   // Assuming that `operand` bufferizes in-place: For each write (to each
   // alias), check if there is a non-writable tensor in the reverse SSA use-def
   // chain.
-  for (OpOperand *uWrite : usesWrite)
+  for (OpOperand *uWrite : usesWrite) {
     if (hasPrecedingAliasingNonWritableTensor(uWrite->get(), &operand,
-                                              aliasInfo, state))
+                                              aliasInfo, state)) {
+      LLVM_DEBUG(llvm::dbgs() << "=> NOT WRITABLE\n");
       return true;
+    }
+  }
 
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place.
 static LogicalResult bufferizableInPlaceAnalysisImpl(
     OpOperand &operand, BufferizationAliasInfo &aliasInfo,
     OneShotAnalysisState &state, const DominanceInfo &domInfo) {
+  LLVM_DEBUG(
+      llvm::dbgs() << "//===-------------------------------------------===//\n"
+                   << "Analyzing operand #" << operand.getOperandNumber()
+                   << " of " << *operand.getOwner() << "\n");
+
   bool foundInterference =
       wouldCreateWriteToNonWritableBuffer(operand, aliasInfo, state) ||
       wouldCreateReadAfterWriteInterference(operand, domInfo, state, aliasInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(operand);
   else
     aliasInfo.bufferizeInPlace(operand, state);
 
+  LLVM_DEBUG(llvm::dbgs()
+             << "//===-------------------------------------------===//\n");
   return success();
 }
 
 /// Analyze the `ops` to determine which OpOperands are inplaceable. Walk ops in
 /// reverse and bufferize ops greedily. This is a good starter heuristic.
 ///
 /// Even if an op does not read or write, it may still create an alias when
 /// bufferized in-place. An example of such ops is tensor.extract_slice.