[mlir][linalg] Format bufferization debug print outs (NFC).

[gysit]

Summary

Change the formatting of the debug print outs to elide unnecessary information.

[gysit]

@nicolasvasilache an easier solution maybe to use the normal value printer and drop anything except for the first line. However, readability is still somewhat limited.

[gysit]

• Add method to print operation info
• Use printing methods more widely

[nicolasvasilache]

Nice, thanks for the QoL improvements!

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

 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/BufferUtils.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "comprehensive-func-bufferize"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
+// Forward declarations.
+static std::string printOperationInfo(Operation *);
+static std::string printValueInfo(Value);
+
 //===----------------------------------------------------------------------===//
 // Generic helpers.
 //===----------------------------------------------------------------------===//
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// Return the FuncOp called by `callOp`.
 static FuncOp getCalledFunction(CallOpInterface callOp) {

 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 /// Wrapper for better debugging.
 static void map(BlockAndValueMapping &bvm, ValueRange keys, ValueRange values) {
   assert(!keys.empty() && "Unexpected empty keys");
-  LDBG("Map: " << keys.front() << " to " << values.front() << '\n');
+  LDBG("\n\tMap: " << printValueInfo(keys.front())
+                   << "\n\tto: " << printValueInfo(values.front()) << '\n');
   return bvm.map(keys, values);
 }
 
 /// Wrapper for better debugging.
 static void map(BlockAndValueMapping &bvm, Value key, Value value) {
-  LDBG("Map: " << key << " to " << value << '\n');
+  LDBG("\n\tMap: " << printValueInfo(key) << "\n\tto: " << printValueInfo(value)
+                   << '\n');
   return bvm.map(key, value);
 }
 
 /// Wrapper for better debugging.
 static Value lookup(const BlockAndValueMapping &bvm, Value key) {
   // TODO: if key comes from bbArg, forward.
   assert(key.getType().isa<TensorType>());
   Value v = bvm.lookupOrNull(key);
static void setInPlaceOpResult(OpResult opResult,
   Operation *op = opResult.getOwner();
   auto attr =
       op->getAttr(kInPlaceResultsAttrName).dyn_cast_or_null<ArrayAttr>();
   SmallVector<StringRef> inPlaceVector =
       attr ? SmallVector<StringRef>(
                  llvm::to_vector<4>(attr.getAsValueRange<StringAttr>()))
            : SmallVector<StringRef>(op->getNumResults(),
                                     stringify(InPlaceSpec::None));
-  LDBG("->set inPlace=" << stringify(inPlace) << ": " << *op
-                        << " @idx=" << opResult.getResultNumber() << '\n');
+  LDBG("->set inPlace=" << stringify(inPlace) << " <- #"
+                        << opResult.getResultNumber() << ": "
+                        << printOperationInfo(op) << "\n");
   inPlaceVector[opResult.getResultNumber()] = stringify(inPlace);
   op->setAttr(kInPlaceResultsAttrName,
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
 /// Get the InPlaceSpec attribute entry `kInPlaceResultsAttrName` for
 /// `opResult`. If the result is `InPlaceSpec::True`, the use-def chain analysis
 /// has guaranteed that no subsequent read of the tensor value occurs and the

 
 LLVM_ATTRIBUTE_UNUSED static InPlaceSpec getInPlace(Value v) {
   if (auto bbArg = v.dyn_cast<BlockArgument>())
     return getInPlace(bbArg);
   return getInPlace(v.cast<OpResult>());
 }
 
 //===----------------------------------------------------------------------===//
+// Printing helpers.
+//===----------------------------------------------------------------------===//
+
+/// Helper method printing the bufferization information of a buffer / tensor.
+static void printTensorOrBufferInfo(std::string prefix, Value value,
+                                    AsmState &state, llvm::raw_ostream &os) {
+  if (!value.getType().isa<ShapedType>())
+    return;
+  os << prefix;
+  value.printAsOperand(os, state);
+  os << " : " << value.getType();
+  if (getInPlace(value) == InPlaceSpec::None)
+    return;
+  os << " [InPlace=" << stringify(getInPlace(value)) << "]";
+}
+
+/// Print the operation name and bufferization information.
+static std::string printOperationInfo(Operation *op) {
+  std::string result;
+  llvm::raw_string_ostream os(result);
+  AsmState state(op->getParentOfType<mlir::FuncOp>());
+  os << op->getName();
+  SmallVector<Value> shapedOperands;
+  for (OpOperand &opOperand : op->getOpOperands()) {
+    std::string prefix =
+        llvm::formatv("\n\t-> #{0} ", opOperand.getOperandNumber());
+    printTensorOrBufferInfo(prefix, opOperand.get(), state, os);
+  }
+  for (OpResult opResult : op->getOpResults()) {
+    std::string prefix =
+        llvm::formatv("\n\t<- #{0} ", opResult.getResultNumber());
+    printTensorOrBufferInfo(prefix, opResult, state, os);
+  }
+  return result;
+}
+
+/// Print the bufferization information for the defining op or block argument.
+static std::string printValueInfo(Value value) {
+  auto *op = value.getDefiningOp();
+  if (op)
+    return printOperationInfo(op);
+  // Print the block argument bufferization information.
+  std::string result;
+  llvm::raw_string_ostream os(result);
+  AsmState state(value.getParentRegion()->getParentOfType<mlir::FuncOp>());
+  os << value;
+  printTensorOrBufferInfo("\n\t - ", value, state, os);
+  return result;
+}
+
+//===----------------------------------------------------------------------===//
 // Op-specific semantics helper to retrieve matching inplaceable result.
 // These should become proper interfaces interfaces when the time is right.
 // Modulo better naming, these helpers / interfaces comprise information on:
 //   1. Whether an op has a known bufferization behavior (i.e. an instance of
 //      BufferizableOpInterface).
 //   2. Whether an op, when bufferized inplace, can guarantee an
 //      (OpOperand, OpResult) pair bufferizes to equivalent (i.e. the same)
 //      buffers in memory.

 }
 
 /// Return true if the buffer to which `operand` would bufferize aliases a
 /// buffer that is known to not be writeable. This implies that the matching
 /// OpResult cannot be bufferized inplace.
 bool BufferizationAliasInfo::aliasesNonWriteableBuffer(
     OpOperand &operand) const {
   LDBG("----Start aliasesNonWriteableBuffer\n");
-  LDBG("-------for operand #" << operand.getOperandNumber() << ": "
-                              << *(operand.getOwner()) << '\n');
+  LDBG("-------for -> #" << operand.getOperandNumber() << ": "
+                         << printOperationInfo(operand.getOwner()) << '\n');
   for (Value v : getAliasInfoRef(operand.get())) {
-    LDBG("-----------examine: " << v << '\n');
+    LDBG("-----------examine: " << printValueInfo(v) << '\n');
     if (auto bbArg = v.dyn_cast<BlockArgument>()) {
       if (getInPlace(bbArg) == InPlaceSpec::True) {
-        LDBG("-----------bbArg is writeable -> skip: " << bbArg << '\n');
+        LDBG("-----------bbArg is writeable -> skip: " << printValueInfo(bbArg)
+                                                       << '\n');
         continue;
       }
       LDBG("-----------notWriteable\n");
       return true;
     }
 
     if (Operation *op = v.getDefiningOp()) {
       if (isa<ConstantOp>(op) || !hasKnownBufferizationAliasingBehavior(op)) {
         LDBG("-----------notWriteable\n");
         return true;
       }
     }
   }
   LDBG("---->operand is writeable\n");
   return false;
 }
 
 /// Return true if the buffer to which `operand` would bufferize is equivalent
 /// to some buffer write.
 bool BufferizationAliasInfo::aliasesInPlaceWrite(Value value) const {
   LDBG("----Start aliasesInPlaceWrite\n");
-  LDBG("-------for : " << value << '\n');
+  LDBG("-------for : " << printValueInfo(value) << '\n');
   for (Value v : getAliasInfoRef(value)) {
     for (auto &use : v.getUses()) {
       if (bufferizesToMemoryWrite(use, InPlaceSpec::True)) {
         LDBG("-----------wants to bufferize to inPlace write: "
-             << *use.getOwner() << '\n');
+             << printOperationInfo(use.getOwner()) << '\n');
         return true;
       }
     }
   }
   LDBG("----------->does not alias an inplace write\n");
   return false;
 }
 
bool BufferizationAliasInfo::wouldCreateReadAfterWriteInterference(
     OpResult result, const DominanceInfo &domInfo) const {
   Optional<OpOperand *> maybeAliasingOperand = getAliasingOpOperand(result);
   if (!maybeAliasingOperand)
     return false;
 
   Operation *opToBufferize = result.getDefiningOp();
   Value root = (*maybeAliasingOperand)->get();
   LDBG("----Start wouldCreateReadAfterWriteInterference\n");
-  LDBG("--------rootValue: " << root << "\n");
+  LDBG("--------rootValue: " << printValueInfo(root) << "\n");
 
   // Collect:
   //   1. all the inplace write uses of some alias of `root`.
   //   2. all the write uses that belong to `opToBufferize`.
   // opToBufferize is not yet inplace, we want to determine if it can be inplace
   // so we also consider all its write uses, not just the inplace ones.
   DenseSet<OpOperand *> usesWrite;
   for (Value vWrite : getAliasInfoRef(root)) {
for (auto &uRead : vRead.getUses())
         usesRead.insert(&uRead);
   for (Value vRead : getAliasInfoRef(result))
     for (auto &uRead : vRead.getUses())
       if (bufferizesToMemoryRead(uRead))
         usesRead.insert(&uRead);
 
   for (OpOperand *uRead : usesRead) {
     Operation *aliasingReadOp = uRead->getOwner();
-    LDBG("----++++aliasRead  #" << uRead->getOperandNumber()
-                                << " in: " << *aliasingReadOp << '\n');
+    LDBG("----++++aliasRead -> #"
+         << uRead->getOperandNumber()
+         << " in: " << printOperationInfo(aliasingReadOp) << '\n');
     for (OpOperand *uWrite : usesWrite) {
       // Don't consider self-use of the same operand for interference.
       // Multiple different uses within the same op is fair game though.
       if (uWrite == uRead)
         continue;
 
       Operation *aliasingWriteOp = uWrite->getOwner();
-      LDBG("----    aliasWrite #" << uWrite->getOperandNumber()
-                                  << " in: " << *aliasingWriteOp << '\n');
+      LDBG("----    aliasWrite -> #"
+           << uWrite->getOperandNumber()
+           << " in: " << printOperationInfo(aliasingWriteOp) << '\n');
       // If the candidate write is the one that produces the read value (in the
       // SSA def-use sense), this is not considered an interference.
       if (getInplaceableOpResult(*uWrite) == uRead->get())
         continue;
       // If aliasingReadOp properly dominates aliasingWriteOp, the read cannot
       // be affected by the write: there is no interference.
       if (domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp))
         continue;
       // At this point, aliasingWriteOp properly dominates aliasingReadOp or
       // there is no clear dominance and we need to be conservative.
       LDBG("---->found RaW interference\n");
-      LDBG("     Interfering  read (op #" << uRead->getOperandNumber()
-                                          << "): " << *aliasingReadOp << '\n');
-      LDBG("     Interfering write (op #" << uWrite->getOperandNumber()
-                                          << "): " << *aliasingWriteOp << '\n');
+      LDBG("     Interfering  read -> #" << uRead->getOperandNumber() << ":\n"
+                                         << printOperationInfo(aliasingReadOp)
+                                         << '\n');
+      LDBG("     Interfering write -> #" << uWrite->getOperandNumber() << ":\n"
+                                         << printOperationInfo(aliasingWriteOp)
+                                         << '\n');
       LDBG("---->opportunity to clobber RaW interference\n");
       if (isClobberedWriteBeforeRead(opToBufferize, *uRead, *uWrite, domInfo)) {
         LDBG("---->clobbered! -> skip\n");
         continue;
       }
 
       LDBG("---->not clobbered -> found an interference\n");
       return true;
for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
     fun(mit->v);
   }
 }
 
 void BufferizationAliasInfo::print(raw_ostream &os) const {
   os << "\n/========================== AliasInfo "
         "==========================\n";
   for (auto it : aliasInfo) {
-    os << "|\n| -- source: " << it.getFirst() << '\n';
+    os << "|\n| -- source: " << printValueInfo(it.getFirst()) << '\n';
     for (auto v : it.getSecond())
-      os << "| ---- target: " << v << '\n';
+      os << "| ---- target: " << printValueInfo(v) << '\n';
   }
   os << "|\n\\====================== End AliasInfo "
         "======================\n\n";
   os << "\n/********************* Equivalent Buffers *********************\n";
   for (auto it = equivalentInfo.begin(), eit = equivalentInfo.end(); it != eit;
        ++it) {
     if (!it->isLeader())
       continue;
     Value leader = it->getData();
-    os << "|\n| -- leader: " << leader << '\n';
+    os << "|\n| -- leader: " << printValueInfo(leader) << '\n';
     for (auto mit = equivalentInfo.member_begin(it),
               meit = equivalentInfo.member_end();
          mit != meit; ++mit) {
       Value v = static_cast<Value>(*mit);
-      os << "| ---- equivalent member: " << v << '\n';
+      os << "| ---- equivalent member: " << printValueInfo(v) << '\n';
     }
   }
   os << "|\n\\***************** End Equivalent Buffers *****************\n\n";
 }
 
 DenseSet<Value> &BufferizationAliasInfo::getAliasInfoRef(Value v) {
   auto it = aliasInfo.find(v);
   if (it == aliasInfo.end())
for (Value valueToClobber : {aliasingRead.get(), aliasingWrite.get()}) {
     for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
          ++mit) {
       /// Note: the "would write to memory after bufferization" condition is
       /// verified by `candidateOp` since it would produce a value that
       /// bufferizes to an equivalent buffer.
       Operation *candidateOp = mit->v.getDefiningOp();
       if (!candidateOp)
         continue;
-      LDBG("---->clobbering candidate: " << *candidateOp << '\n');
+      LDBG("---->clobbering candidate: " << printOperationInfo(candidateOp)
+                                         << '\n');
       if (domInfo.properlyDominates(aliasingWriteOp, candidateOp) &&
           domInfo.properlyDominates(candidateOp, aliasingReadOp))
         return true;
     }
   }
   return false;
 }
 

 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult
 bufferizableInPlaceAnalysis(ExtractSliceOp extractSliceOp,
                             BufferizationAliasInfo &aliasInfo,
                             const DominanceInfo &domInfo) {
   LDBG('\n');
-  LDBG("Inplace analysis for extract_slice: " << *extractSliceOp << '\n');
+  LDBG("Inplace analysis for extract_slice: "
+       << printOperationInfo(extractSliceOp) << '\n');
 
   // If `extractSliceOp` were to be bufferized inplace, it cannot end up
   // aliasing a write into a non-writeable buffer.
   bool wouldCreateAliasingWriteToNonWriteableBuffer =
       aliasInfo.aliasesInPlaceWrite(extractSliceOp.result()) &&
       aliasInfo.aliasesNonWriteableBuffer(extractSliceOp->getOpOperand(0));
 
   if (wouldCreateAliasingWriteToNonWriteableBuffer)
bufferizableInPlaceAnalysis(OpOperand &operand, OpResult result,
   Operation *op = result.getDefiningOp();
   assert(result && !isa<ExtractSliceOp>(op) &&
          "expected OpResult not coming from a ExtractSliceOp");
   (void)op;
 
   int64_t resultNumber = result.getResultNumber();
   (void)resultNumber;
   LDBG('\n');
-  LDBG("Inplace analysis for result #" << resultNumber << " (operand #"
-                                       << operand.getOperandNumber() << ") in "
-                                       << result << '\n');
+  LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
+                                   << operand.getOperandNumber() << " in "
+                                   << printValueInfo(result) << '\n');
 
   // `result` must bufferize to a writeable buffer to be a candidate.
   // This means the operand  must not alias either:
   //   1. a function bbArg that is not inplaceable or
   //   2. a constant op.
   // to be considered for inplace bufferization
   bool wouldCreateAliasingWriteToNonWriteableBuffer =
       aliasInfo.aliasesNonWriteableBuffer(operand);