[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"

 
 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.
+//===----------------------------------------------------------------------===//
+
+/// Return string containing the essential bufferization information for the
+/// given value.
+static std::string printValueInfo(Value value) {
+  std::string result;
+  llvm::raw_string_ostream os(result);
+  AsmState state(value.getParentRegion()->getParentOfType<mlir::FuncOp>());
+  auto printTensorOrBufferInfo = [&](std::string prefix, Value value) {
+    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)) << "]";
+  };
+  auto *op = value.getDefiningOp();
+  if (op) {
+    os << op->getName();
+    SmallVector<Value> shapedOperands;
+    for (Value value : op->getOperands())
+      printTensorOrBufferInfo("\n\t-> ", value);
+    for (Value result : op->getOpResults())
+      printTensorOrBufferInfo("\n\t<- ", result);
+  } else {
+    os << value;
+    printTensorOrBufferInfo("\n\t - ", value);
+  }
+  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.
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())

 
 /// Postprocess the linalg.buffer_layout annotation across function boundaries.
 /// This is a purely mechanical process that may later become part of a
 /// separate pass with its own layout assignment heuristic.
 static void layoutPostProcessing(ModuleOp moduleOp) {
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   auto res = getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap);
-  (void) res;
+  (void)res;
   assert(succeeded(res) && "unexpected getFuncOpsOrderedByCalls failure");
 
   for (FuncOp funcOp : orderedFuncOps) {
     DenseMap<Operation *, SmallVector<Value>> operandsPerCaller;
     foreachCaller(callerMap, funcOp, [&](Operation *caller) {
       operandsPerCaller.try_emplace(caller, SmallVector<Value>());
     });