[llvm-branch-commits] [mlir] eb1aacd - [MLIR-LAIR] Custom printing for multi-dim affine index maps

[Uday Bondhugula via llvm-branch-commits]

Author: Uday Bondhugula
Date: 2021-09-23T08:19:53+05:30
New Revision: eb1aacd495d84384e2d2e4e40c893aebfbd19837

URL: https://github.com/llvm/llvm-project/commit/eb1aacd495d84384e2d2e4e40c893aebfbd19837
DIFF: https://github.com/llvm/llvm-project/commit/eb1aacd495d84384e2d2e4e40c893aebfbd19837.diff

LOG: [MLIR-LAIR] Custom printing for multi-dim affine index maps

Custom printing for AffineMaps with named SSA values - to aid printing
NGDL accesses.

Added: 
    

Modified: 
    mlir/include/mlir/IR/AffineMap.h
    mlir/lib/IR/AsmPrinter.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index f1f00acd8022..4ab75ad3ea64 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -534,6 +534,11 @@ inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
   map.print(os);
   return os;
 }
+
+/// Prints an affine map index in the form [expr_1][expr_2] with the supplied
+/// names used for the inputs.
+void printAffineMapAccess(AffineMap map, ArrayRef<std::string> names,
+                          raw_ostream &os);
 } // end namespace mlir
 
 namespace llvm {

diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index b6d072f9f60a..9c717eef111c 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -2241,6 +2241,18 @@ void ModulePrinter::printAffineConstraint(AffineExpr expr, bool isEq) {
   isEq ? os << " == 0" : os << " >= 0";
 }
 
+void mlir::printAffineMapAccess(AffineMap map, ArrayRef<std::string> names,
+                                raw_ostream &os) {
+  assert(map.getNumInputs() == names.size() && "too few/many names");
+  ModulePrinter p(os);
+  auto printValueName = [&](unsigned pos, bool isSymbol) { os << names[pos]; };
+  for (auto expr : map.getResults()) {
+    os << '[';
+    p.printAffineExpr(expr, printValueName);
+    os << ']';
+  }
+}
+
 void ModulePrinter::printAffineMap(AffineMap map) {
   // Dimension identifiers.
   os << '(';