[mlir][sparse] remove singleton dimension level type (for now)

Although we have plans to support this, and many other, dimension level type(s), currently the tag is not supported. It will be easy to add this back once support is added.

NOTE: based on discussion in https://discourse.llvm.org/t/overcoming-sparsification-limitation-on-level-types/62585

https://github.com/llvm/llvm-project/issues/51658

Reviewed By: Peiming

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

diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index 252ec68..ac2b8b6 100644 (file)
--- a/mlir/include/mlir-c/Dialect/SparseTensor.h
+++ b/mlir/include/mlir-c/Dialect/SparseTensor.h
@@ -20,12 +20,10 @@ extern "C" {
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
 
 /// Dimension level types that define sparse tensors:
-///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE      - dimension is dense, every
+///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE - dimension is dense, every
 ///   entry is stored
 ///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED - dimension is sparse,
-///   only nonzeros are stored.
-///   - MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON  - dimension contains single
-///   coordinate, no siblings.
+///   only nonzeros are stored (no duplicates).
 ///
 /// These correspond to SparseTensorEncodingAttr::DimLevelType in the C++ API.
 /// If updating, keep them in sync and update the static_assert in the impl
@@ -33,7 +31,6 @@ MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor);
 enum MlirSparseTensorDimLevelType {
   MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE,
   MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED,
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON,
 };
 
 //===----------------------------------------------------------------------===//

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index 76f408d..4fbac0d 100644 (file)
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -84,11 +84,10 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
 
   let extraClassDeclaration = [{
     // Dimension level types that define sparse tensors:
-    //   Dense      - dimension is dense, every entry is stored
-    //   Compressed - dimension is sparse, only nonzeros are stored
-    //   Singleton  - dimension contains single coordinate, no siblings
+    //   Dense               - dimension is dense, every entry is stored
+    //   Compressed (unique) - dimension is sparse, only nonzeros are stored (no duplicates)
     enum class DimLevelType {
-      Dense, Compressed, Singleton
+      Dense, Compressed
     };
   }];
 }

diff --git a/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h b/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
index db83ce0..86d36b1 100644 (file)
--- a/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h
@@ -21,7 +21,7 @@ namespace mlir {
 namespace sparse_tensor {
 
 /// Dimension level type for a tensor (undef means index does not appear).
-enum Dim { kSparse, kDense, kSingle, kUndef };
+enum Dim { kSparse, kDense, kUndef };
 
 /// Tensor expression kind.
 enum Kind {

diff --git a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
index b24d024..49b4a89 100644 (file)
--- a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
+++ b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
@@ -18,8 +18,7 @@ using namespace mlir::python::adaptors;
 static void populateDialectSparseTensorSubmodule(const py::module &m) {
   py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType", py::module_local())
       .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE)
-      .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED)
-      .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON);
+      .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED);
 
   mlir_attribute_subclass(m, "EncodingAttr",
                           mlirAttributeIsASparseTensorEncodingAttr)

diff --git a/mlir/lib/CAPI/Dialect/SparseTensor.cpp b/mlir/lib/CAPI/Dialect/SparseTensor.cpp
index f35c14a..b7b2fd5 100644 (file)
--- a/mlir/lib/CAPI/Dialect/SparseTensor.cpp
+++ b/mlir/lib/CAPI/Dialect/SparseTensor.cpp
@@ -25,9 +25,7 @@ static_assert(
             static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Dense) &&
         static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED) ==
             static_cast<int>(
-                SparseTensorEncodingAttr::DimLevelType::Compressed) &&
-        static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON) ==
-            static_cast<int>(SparseTensorEncodingAttr::DimLevelType::Singleton),
+                SparseTensorEncodingAttr::DimLevelType::Compressed),
     "MlirSparseTensorDimLevelType (C-API) and DimLevelType (C++) mismatch");
 
 bool mlirAttributeIsASparseTensorEncodingAttr(MlirAttribute attr) {

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 6bc49d5..8092aff 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -71,8 +71,6 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
           dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Dense);
         } else if (strVal == "compressed") {
           dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Compressed);
-        } else if (strVal == "singleton") {
-          dlt.push_back(SparseTensorEncodingAttr::DimLevelType::Singleton);
         } else {
           parser.emitError(parser.getNameLoc(),
                            "unexpected dimension level type: ")
@@ -126,9 +124,6 @@ void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
     case DimLevelType::Compressed:
       printer << "\"compressed\"";
       break;
-    case DimLevelType::Singleton:
-      printer << "\"singleton\"";
-      break;
     }
     if (i != e - 1)
       printer << ", ";

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
index 5533d02..d30a81d 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
@@ -148,8 +148,6 @@ DimLevelType mlir::sparse_tensor::dimLevelTypeEncoding(
     return DimLevelType::kDense;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return DimLevelType::kCompressed;
-  case SparseTensorEncodingAttr::DimLevelType::Singleton:
-    return DimLevelType::kSingleton;
   }
   llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
 }

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index b37703a..dca1c52 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -378,11 +378,6 @@ static bool canUseDirectConversion(
       if (alreadyCompressed)
         return false; // Dense after Compressed not yet supported.
       break;
-    case SparseTensorEncodingAttr::DimLevelType::Singleton:
-      // Although Singleton isn't generally supported yet, the direct
-      // conversion method doesn't have any particular problems with
-      // singleton after compressed.
-      break;
     }
   }
   return true;

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 3904835..f392e18 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -147,8 +147,6 @@ static Dim toDim(const SparseTensorEncodingAttr &enc, unsigned d) {
     SparseTensorEncodingAttr::DimLevelType tp = enc.getDimLevelType()[d];
     if (tp == SparseTensorEncodingAttr::DimLevelType::Compressed)
       return Dim::kSparse;
-    if (tp == SparseTensorEncodingAttr::DimLevelType::Singleton)
-      return Dim::kSingle;
   }
   return Dim::kDense;
 }

diff --git a/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp b/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
index 5ecd260..f844987 100644 (file)
--- a/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp
@@ -591,9 +591,6 @@ void Merger::dumpBits(const BitVector &bits) const {
       case kDense:
         llvm::dbgs() << "D";
         break;
-      case kSingle:
-        llvm::dbgs() << "T";
-        break;
       case kUndef:
         llvm::dbgs() << "U";
         break;

diff --git a/mlir/test/CAPI/sparse_tensor.c b/mlir/test/CAPI/sparse_tensor.c
index 10cd812..c7ee71f 100644 (file)
--- a/mlir/test/CAPI/sparse_tensor.c
+++ b/mlir/test/CAPI/sparse_tensor.c
@@ -25,7 +25,7 @@ static int testRoundtripEncoding(MlirContext ctx) {
   // clang-format off
   const char *originalAsm =
     "#sparse_tensor.encoding<{ "
-    "dimLevelType = [ \"dense\", \"compressed\", \"singleton\"], "
+    "dimLevelType = [ \"dense\", \"compressed\", \"compressed\"], "
     "dimOrdering = affine_map<(d0, d1, d2) -> (d0, d1, d2)>, "
     "pointerBitWidth = 32, indexBitWidth = 64 }>";
   // clang-format on
@@ -40,7 +40,7 @@ static int testRoundtripEncoding(MlirContext ctx) {
   mlirAffineMapDump(dimOrdering);
   // CHECK: level_type: 0
   // CHECK: level_type: 1
-  // CHECK: level_type: 2
+  // CHECK: level_type: 1
   int numLevelTypes = mlirSparseTensorEncodingGetNumDimLevelTypes(originalAttr);
   enum MlirSparseTensorDimLevelType *levelTypes =
       malloc(sizeof(enum MlirSparseTensorDimLevelType) * numLevelTypes);