[mlir][sparse] Move n:m printing into toMLIRString

[yinying-lisa-li]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Yinying Li (yinying-lisa-li)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/84264.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h (+4)
• (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+2-14)

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]);
   }
 }