[Mlir-commits] [mlir] [MLIR][Presburger] Make printing aligned to assist in debugging (PR #107648)

[Amy Wang]

https://github.com/kaitingwang updated https://github.com/llvm/llvm-project/pull/107648

>From 1dba1bb0c2e558abc3ac15f9fcb5b333c1e08cef Mon Sep 17 00:00:00 2001
From: Amy Wang <kai.ting.wang at huawei.com>
Date: Fri, 6 Sep 2024 17:51:11 -0400
Subject: [PATCH] [MLIR][Presburger] Make printing aligned to assist in
 debugging

---
 mlir/include/mlir/Analysis/Presburger/Utils.h | 44 +++++++++++++++++++
 .../Analysis/Presburger/IntegerRelation.cpp   | 24 +++++++---
 mlir/lib/Analysis/Presburger/Matrix.cpp       | 14 ++++--
 mlir/lib/Analysis/Presburger/Simplex.cpp      | 10 ++++-
 4 files changed, 82 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/Utils.h b/mlir/include/mlir/Analysis/Presburger/Utils.h
index d3c0802c240bc1..1f471c79cb3b24 100644
--- a/mlir/include/mlir/Analysis/Presburger/Utils.h
+++ b/mlir/include/mlir/Analysis/Presburger/Utils.h
@@ -17,7 +17,9 @@
 #include "llvm/ADT/DynamicAPInt.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Support/raw_ostream.h"
 #include <optional>
+#include <string>
 
 namespace mlir {
 namespace presburger {
@@ -292,6 +294,48 @@ std::vector<Fraction> multiplyPolynomials(ArrayRef<Fraction> a,
 
 bool isRangeZero(ArrayRef<Fraction> arr);
 
+struct PrintTableMetrics {
+  // If unknown, set to 0 and pass the struct into updatePrintMetrics.
+  unsigned maxPreIndent;
+  unsigned maxPostIndent;
+  // The substring expected prior to alignment.
+  std::string preAlign;
+};
+
+// Updates 'm' given a table entry val. Iterate over each val in the table
+// with .maxPreIndent and .maxPostIndent initialized to 0.
+// class T is any type that can be handled by llvm::raw_string_ostream.
+template <class T>
+void updatePrintMetrics(T val, PrintTableMetrics &m) {
+  std::string str;
+  llvm::raw_string_ostream(str) << val;
+  if (str.length() == 0)
+    return;
+  unsigned int preIndent = str.find(m.preAlign);
+  preIndent = (preIndent != std::string::npos) ? preIndent + 1 : 0;
+  m.maxPreIndent = std::max(m.maxPreIndent, preIndent);
+  m.maxPostIndent =
+      std::max(m.maxPostIndent, (unsigned int)(str.length() - preIndent));
+}
+
+// Print val in the table with metrics specified in 'm'.
+template <class T>
+void printWithPrintMetrics(raw_ostream &os, T val, unsigned minSpacing,
+                           const PrintTableMetrics &m) {
+  std::string str;
+  llvm::raw_string_ostream(str) << val;
+  unsigned preIndent;
+  if (str.length() != 0) {
+    preIndent = str.find(m.preAlign);
+    preIndent = (preIndent != std::string::npos) ? preIndent + 1 : 0;
+  } else
+    preIndent = 0;
+  for (unsigned i = 0; i < (minSpacing + m.maxPreIndent - preIndent); ++i)
+    os << " ";
+  os << str;
+  for (unsigned i = 0; i < m.maxPostIndent - (str.length() - preIndent); ++i)
+    os << " ";
+}
 } // namespace presburger
 } // namespace mlir
 
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 94af81f955e5a5..182232996eb529 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -32,7 +32,10 @@
 #include <cassert>
 #include <functional>
 #include <memory>
+#include <numeric>
 #include <optional>
+#include <sstream>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -2589,19 +2592,28 @@ void IntegerRelation::mergeAndCompose(const IntegerRelation &other) {
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());
   printSpace(os);
+  PrintTableMetrics ptm = {0, 0, "-"};
+  for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
+    for (unsigned j = 0, f = getNumCols(); j < f; ++j)
+      updatePrintMetrics<DynamicAPInt>(atEq(i, j), ptm);
+  }
+  for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
+    for (unsigned j = 0, f = getNumCols(); j < f; ++j)
+      updatePrintMetrics<DynamicAPInt>(atIneq(i, j), ptm);
+  }
+  // Print using PrintMetrics.
+  unsigned MIN_SPACING = 1;
   for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
-    os << " ";
     for (unsigned j = 0, f = getNumCols(); j < f; ++j) {
-      os << atEq(i, j) << "\t";
+      printWithPrintMetrics<DynamicAPInt>(os, atEq(i, j), MIN_SPACING, ptm);
     }
-    os << "= 0\n";
+    os << "  = 0\n";
   }
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
-    os << " ";
     for (unsigned j = 0, f = getNumCols(); j < f; ++j) {
-      os << atIneq(i, j) << "\t";
+      printWithPrintMetrics<DynamicAPInt>(os, atIneq(i, j), MIN_SPACING, ptm);
     }
-    os << ">= 0\n";
+    os << " >= 0\n";
   }
   os << '\n';
 }
diff --git a/mlir/lib/Analysis/Presburger/Matrix.cpp b/mlir/lib/Analysis/Presburger/Matrix.cpp
index 110c5df1af37c0..8f3a22a4c18d01 100644
--- a/mlir/lib/Analysis/Presburger/Matrix.cpp
+++ b/mlir/lib/Analysis/Presburger/Matrix.cpp
@@ -398,10 +398,18 @@ Matrix<T> Matrix<T>::getSubMatrix(unsigned fromRow, unsigned toRow,
 
 template <typename T>
 void Matrix<T>::print(raw_ostream &os) const {
+  PrintTableMetrics ptm = {0, 0, "-"};
   for (unsigned row = 0; row < nRows; ++row) {
-    for (unsigned column = 0; column < nColumns; ++column)
-      os << at(row, column) << ' ';
-    os << '\n';
+    for (unsigned column = 0; column < nColumns; ++column) {
+      updatePrintMetrics<T>(at(row, column), ptm);
+    }
+  }
+  unsigned MIN_SPACING = 1;
+  for (unsigned row = 0; row < nRows; ++row) {
+    for (unsigned column = 0; column < nColumns; ++column) {
+      printWithPrintMetrics<T>(os, at(row, column), MIN_SPACING, ptm);
+    }
+    os << "\n";
   }
 }
 
diff --git a/mlir/lib/Analysis/Presburger/Simplex.cpp b/mlir/lib/Analysis/Presburger/Simplex.cpp
index c78a0723a6c0fa..77c8ea7b94b2ad 100644
--- a/mlir/lib/Analysis/Presburger/Simplex.cpp
+++ b/mlir/lib/Analysis/Presburger/Simplex.cpp
@@ -2153,9 +2153,17 @@ void SimplexBase::print(raw_ostream &os) const {
   for (unsigned col = 2, e = getNumColumns(); col < e; ++col)
     os << ", c" << col << ": " << colUnknown[col];
   os << '\n';
+  PrintTableMetrics ptm = {0, 0, "-"};
   for (unsigned row = 0, numRows = getNumRows(); row < numRows; ++row) {
     for (unsigned col = 0, numCols = getNumColumns(); col < numCols; ++col)
-      os << tableau(row, col) << '\t';
+      updatePrintMetrics<DynamicAPInt>(tableau(row, col), ptm);
+  }
+  unsigned MIN_SPACING = 1;
+  for (unsigned row = 0, numRows = getNumRows(); row < numRows; ++row) {
+    for (unsigned col = 0, numCols = getNumColumns(); col < numCols; ++col) {
+      printWithPrintMetrics<DynamicAPInt>(os, tableau(row, col), MIN_SPACING,
+                                          ptm);
+    }
     os << '\n';
   }
   os << '\n';