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

[aartbik]

Summary

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

mlir/include/mlir-c/Dialect/SparseTensor.h

 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 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
 /// file.
 enum MlirSparseTensorDimLevelType {
   MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE,
   MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED,
-  MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON,
 };
 
 //===----------------------------------------------------------------------===//
 // SparseTensorEncodingAttr
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given attribute is a sparse_tensor.encoding attribute.
 MLIR_CAPI_EXPORTED bool

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

let parameters = (
     "unsigned":$indexBitWidth
   );
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 
   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
     };
   }];
 }
 
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/BitVector.h"
 
 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 {
   // Leaf.
   kTensor = 0,
   kInvariant,
   kIndex,
   // Unary operations.

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

 namespace py = pybind11;
 using namespace llvm;
 using namespace mlir;
 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)
       .def_classmethod(
           "get",
           [](py::object cls,
              std::vector<MlirSparseTensorDimLevelType> dimLevelTypes,
              llvm::Optional<MlirAffineMap> dimOrdering, int pointerBitWidth,

mlir/lib/CAPI/Dialect/SparseTensor.cpp

MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(SparseTensor, sparse_tensor,
                                       mlir::sparse_tensor::SparseTensorDialect)
 
 // Ensure the C-API enums are int-castable to C++ equivalents.
 static_assert(
     static_cast<int>(MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE) ==
             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) {
   return unwrap(attr).isa<SparseTensorEncodingAttr>();
 }
 
 MlirAttribute mlirSparseTensorEncodingAttrGet(
     MlirContext ctx, intptr_t numDimLevelTypes,

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

if (attr.getName() == "dimLevelType") {
                            "expected a string value in dimension level types");
           return {};
         }
         auto strVal = strAttr.getValue();
         if (strVal == "dense") {
           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: ")
               << strVal;
           return {};
         }
       }
     } else if (attr.getName() == "dimOrdering") {
void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
   for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
     switch (getDimLevelType()[i]) {
     case DimLevelType::Dense:
       printer << "\"dense\"";
       break;
     case DimLevelType::Compressed:
       printer << "\"compressed\"";
       break;
-    case DimLevelType::Singleton:
-      printer << "\"singleton\"";
-      break;
     }
     if (i != e - 1)
       printer << ", ";
   }
   printer << " ]";
   if (getDimOrdering())
     printer << ", dimOrdering = affine_map<" << getDimOrdering() << ">";
   printer << ", pointerBitWidth = " << getPointerBitWidth()

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

 
 DimLevelType mlir::sparse_tensor::dimLevelTypeEncoding(
     SparseTensorEncodingAttr::DimLevelType dlt) {
   switch (dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Dense:
     return DimLevelType::kDense;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return DimLevelType::kCompressed;
-  case SparseTensorEncodingAttr::DimLevelType::Singleton:
-    return DimLevelType::kSingleton;
   }
   llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
 }
 
 //===----------------------------------------------------------------------===//
 // Misc code generators.
 //===----------------------------------------------------------------------===//
 

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

case SparseTensorEncodingAttr::DimLevelType::Compressed:
       if (alreadyCompressed)
         return false; // Multiple compressed dimensions not yet supported.
       alreadyCompressed = true;
       break;
     case SparseTensorEncodingAttr::DimLevelType::Dense:
       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;
 }
 
 /// Helper method to translate indices during a reshaping operation.
 /// TODO: provide as general utility to MLIR at large?
 static void translateIndices(Location loc, ConversionPatternRewriter &rewriter,

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

 }
 
 /// Helper method to translate dim level type to internal representation.
 static Dim toDim(const SparseTensorEncodingAttr &enc, unsigned d) {
   if (enc) {
     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;
 }
 
 /// Helper method to inspect affine expressions. Rejects cases where the
 /// same index is used more than once. Also rejects affine expressions
 /// that are not a direct index for annotated tensors.
 // TODO: accept more affine cases for sparse tensors

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

if (bits[b]) {
       llvm::dbgs() << " i_" << t << "_" << i << "_";
       switch (dims[t][i]) {
       case kSparse:
         llvm::dbgs() << "S";
         break;
       case kDense:
         llvm::dbgs() << "D";
         break;
-      case kSingle:
-        llvm::dbgs() << "T";
-        break;
       case kUndef:
         llvm::dbgs() << "U";
         break;
       }
     }
   }
 }
 

mlir/test/CAPI/sparse_tensor.c

 #include <string.h>
 
 // CHECK-LABEL: testRoundtripEncoding()
 static int testRoundtripEncoding(MlirContext ctx) {
   fprintf(stderr, "testRoundtripEncoding()\n");
   // 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
   MlirAttribute originalAttr =
       mlirAttributeParseGet(ctx, mlirStringRefCreateFromCString(originalAsm));
   // CHECK: isa: 1
   fprintf(stderr, "isa: %d\n",
           mlirAttributeIsASparseTensorEncodingAttr(originalAttr));
   MlirAffineMap dimOrdering =
       mlirSparseTensorEncodingAttrGetDimOrdering(originalAttr);
   // CHECK: (d0, d1, d2) -> (d0, d1, d2)
   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);
   for (int i = 0; i < numLevelTypes; ++i) {
     levelTypes[i] =
         mlirSparseTensorEncodingAttrGetDimLevelType(originalAttr, i);
     fprintf(stderr, "level_type: %d\n", levelTypes[i]);
   }