[Mlir-commits] [mlir] 9c9b47c - [mlir][sparse] add dim level type toString convenience method

[Aart Bik]

Author: Aart Bik
Date: 2022-12-15T16:11:30-08:00
New Revision: 9c9b47c976238100add391f27cd7fb8bb20b5371

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

LOG: [mlir][sparse] add dim level type toString convenience method

Reviewed By: wrengr, bixia

Differential Revision: https://reviews.llvm.org/D140165

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
    mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
index 62118130a3fab..d46e325433a5d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
@@ -157,6 +157,34 @@ enum class DimLevelType : uint8_t {
   SingletonNuNo = 19,  // 0b100_11
 };
 
+/// Returns string representation of the given dimension level type.
+inline std::string toMLIRString(DimLevelType dlt) {
+  switch (dlt) {
+  // TODO: should probably raise an error instead of printing it...
+  case DimLevelType::Undef:
+    return "\"undef\"";
+  case DimLevelType::Dense:
+    return "\"dense\"";
+  case DimLevelType::Compressed:
+    return "\"compressed\"";
+  case DimLevelType::CompressedNu:
+    return "\"compressed-nu\"";
+  case DimLevelType::CompressedNo:
+    return "\"compressed-no\"";
+  case DimLevelType::CompressedNuNo:
+    return "\"compressed-nu-no\"";
+  case DimLevelType::Singleton:
+    return "\"singleton\"";
+  case DimLevelType::SingletonNu:
+    return "\"singleton-nu\"";
+  case DimLevelType::SingletonNo:
+    return "\"singleton-no\"";
+  case DimLevelType::SingletonNuNo:
+    return "\"singleton-nu-no\"";
+  }
+  return "";
+}
+
 /// Check that the `DimLevelType` contains a valid (possibly undefined) value.
 constexpr bool isValidDLT(DimLevelType dlt) {
   const uint8_t formatBits = static_cast<uint8_t>(dlt) >> 2;

diff  --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 628b391bdaaa5..06a451fc78ffe 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -160,39 +160,7 @@ void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
   // Print the struct-like storage in dictionary fashion.
   printer << "<{ dimLevelType = [ ";
   for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
-    switch (getDimLevelType()[i]) {
-    case DimLevelType::Undef:
-      // TODO: should probably raise an error instead of printing it...
-      printer << "\"undef\"";
-      break;
-    case DimLevelType::Dense:
-      printer << "\"dense\"";
-      break;
-    case DimLevelType::Compressed:
-      printer << "\"compressed\"";
-      break;
-    case DimLevelType::CompressedNu:
-      printer << "\"compressed-nu\"";
-      break;
-    case DimLevelType::CompressedNo:
-      printer << "\"compressed-no\"";
-      break;
-    case DimLevelType::CompressedNuNo:
-      printer << "\"compressed-nu-no\"";
-      break;
-    case DimLevelType::Singleton:
-      printer << "\"singleton\"";
-      break;
-    case DimLevelType::SingletonNu:
-      printer << "\"singleton-nu\"";
-      break;
-    case DimLevelType::SingletonNo:
-      printer << "\"singleton-no\"";
-      break;
-    case DimLevelType::SingletonNuNo:
-      printer << "\"singleton-nu-no\"";
-      break;
-    }
+    printer << toMLIRString(getDimLevelType()[i]);
     if (i != e - 1)
       printer << ", ";
   }