[Mlir-commits] [mlir] [mlir][sparse] Move n:m printing into toMLIRString (PR #84264)

[Yinying Li]

https://github.com/yinying-lisa-li created https://github.com/llvm/llvm-project/pull/84264

None

>From f0ab9bd85ed49d6d1ce518615ed4d33d1d7a1a80 Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Wed, 6 Mar 2024 23:37:49 +0000
Subject: [PATCH] [mlir][sparse] Move n:m printing into toMLIRString

---
 .../include/mlir/Dialect/SparseTensor/IR/Enums.h |  4 ++++
 .../SparseTensor/IR/SparseTensorDialect.cpp      | 16 ++--------------
 2 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
index 5563cb907e9353..33f613a46bad84 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
@@ -360,6 +360,10 @@ struct LevelType {
   std::string toMLIRString() const {
     std::string lvlStr = toFormatString(getLvlFmt());
     std::string propStr = "";
+    if (isa<LevelFormat::NOutOfM>()) {
+      lvlStr +=
+          "[" + std::to_string(getN()) + ", " + std::to_string(getM()) + "]";
+    }
     if (isa<LevelPropNonDefault::Nonunique>())
       propStr += toPropString(LevelPropNonDefault::Nonunique);
 
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 6ba8b46370b038..c19907a945d3bb 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -646,28 +646,16 @@ void SparseTensorEncodingAttr::printDimensions(
   }
 }
 
-std::string getNOutOfMString(LevelType lt) {
-  if (isNOutOfMLT(lt)) {
-    unsigned n = getN(lt);
-    unsigned m = getM(lt);
-    auto output = "[" + std::to_string(n) + ", " + std::to_string(m) + "]";
-    return output;
-  }
-  return "";
-}
-
 void SparseTensorEncodingAttr::printLevels(AffineMap &map, AsmPrinter &printer,
                                            ArrayRef<LevelType> lvlTypes) const {
   for (unsigned i = 0, n = map.getNumResults() - 1; i < n; i++) {
     map.getResult(i).print(printer.getStream());
-    printer << " : " << toMLIRString(lvlTypes[i])
-            << getNOutOfMString(lvlTypes[i]) << ", ";
+    printer << " : " << toMLIRString(lvlTypes[i]) << ", ";
   }
   if (map.getNumResults() >= 1) {
     auto lastIndex = map.getNumResults() - 1;
     map.getResult(lastIndex).print(printer.getStream());
-    printer << " : " << toMLIRString(lvlTypes[lastIndex])
-            << getNOutOfMString(lvlTypes[lastIndex]);
+    printer << " : " << toMLIRString(lvlTypes[lastIndex]);
   }
 }