[mlir] Truncate/skip long strings in ViewOpGraph.cpp

[springerm]

Summary

• New pass option max-label-len: Truncate attributes/result types that have more #chars.
• New pass option print-attrs: Activate/deactivate rendering of attributes.
• New pass option printResultTypes: Activate/deactivate rendering of result types.

Depends On D106253

[springerm]

rebase

[springerm]

address comments

[springerm]

minor update

[rriddle]

mlir/include/mlir/Transforms/Passes.td

701

The name here is a bit confusing, this isn't just attributes. Can you change this to something more general?

[springerm]

rebase

[springerm]

rebase

[springerm]

rebase

mlir/include/mlir/Transforms/Passes.td

def ViewOpGraphPass : Pass<"view-op-graph", "ModuleOp"> {
   let summary = "Print Graphviz dataflow visualization of module";
   let description = [{
     This pass prints a Graphviz dataflow graph of a module.
 
     - Operations are represented as nodes;
     - Uses as edges;
     - Regions/blocks as subgraphs.
   }];
+  let options = [
+    Option<"maxAttrLen", "max-attr-len", "unsigned",
+            /*default=*/"20", "Limit attribute/type length to number of chars">,
+    Option<"printAttrs", "print-attrs", "bool",
+           /*default=*/"false", "Print attributes of operations">,
+    Option<"printResultTypes", "print-result-types", "bool",
+            /*default=*/"true", "Print result types of operations">
+  ];
   let constructor = "mlir::createPrintOpGraphPass()";
 }
 
 #endif // MLIR_TRANSFORMS_PASSES

mlir/lib/Transforms/ViewOpGraph.cpp

void emitMlirAttr(raw_ostream &os, Attribute attr) {
 
     auto array = attr.dyn_cast<ArrayAttr>();
     if (array && static_cast<int64_t>(array.size()) > largeAttrLimit) {
       os << "[...]";
       return;
     }
 
     // Print all other attributes.
-    attr.print(os);
+    std::string buf;
+    llvm::raw_string_ostream ss(buf);
+    attr.print(ss);
+    os << truncateString(ss.str());
   }
 
   /// Append an edge to the list of edges.
   /// Note: Edges are written to the output stream via `emitAllEdgeStmts`.
   void emitEdgeStmt(Node n1, Node n2, std::string label,
                     std::string style = kLineStyleDataFlow) {
     std::string buf;
     llvm::raw_string_ostream edge(buf);
private:
 
   /// Generate a label for an operation.
   std::string getLabel(Operation *op) {
     // Reuse the print output for the node labels.
     std::string buf;
     llvm::raw_string_ostream os(buf);
 
     // Print operation name and type.
-    os << op->getName() << " : (";
-    interleaveComma(op->getResultTypes(), os);
-    os << ")\n";
+    os << op->getName();
+    if (this->printResultTypes) {
+      os << " : (";
+      std::string buf;
+      llvm::raw_string_ostream ss(buf);
+      interleaveComma(op->getResultTypes(), ss);
+      os << truncateString(ss.str()) << ")";
+      os << ")";
+    }
 
     // Print attributes.
+    if (this->printAttrs) {
+      os << "\n";
       for (auto attr : op->getAttrs()) {
         os << '\n' << attr.first << ": ";
         emitMlirAttr(os, attr.second);
       }
+    }
 
     return os.str();
   }
 
   /// Generate a label for a block argument.
   std::string getLabel(BlockArgument arg) {
     return "arg" + std::to_string(arg.getArgNumber());
   }
private:
   }
 
   // Process a region.
   void processRegion(Region &region) {
     for (auto &block : region.getBlocks())
       processBlock(block);
   }
 
+  /// Truncate long strings.
+  std::string truncateString(std::string str) {
+    if (str.length() <= this->maxAttrLen)
+      return str;
+    return str.substr(0, this->maxAttrLen) + "...";
+  }
+
   // Current indentation (for writing to `os`).
   unsigned indent = 0;
   // Output stream to write DOT file to.
   raw_ostream &os;
   // A list of edges. For simplicity, should be emitted after all nodes were
   // emitted.
   std::vector<std::string> edges;
   // Mapping of SSA values to Graphviz nodes/clusters.