[mlir][sparse] use SparseStorageEncodingAttr for sparse_tensor.storage_specifier

[Peiming]

Summary

SparseTensor::StorageSpecifierType used to take SparseTensorEncodingAttr to calculate the memory layout. Yet, SparseTensorEncodingAttr holds extra information that are not needed by storage specifier (e.g., dimension level properties, dimension ordering, etc).

[Peiming]

fix typo.

[Peiming]

revert unintended change.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

160–161

I think it's clearer to say that the enum defines all the "storage formats/representations themselves, without any level properties."

163

Fwiw, now that I've introduced the "dim"-vs-"lvl" terminology and gotten the bulk of the codebase rewritten to make that distinction, I'm planning to rename DimLevelType to just LevelType in the near future. And the name I've had in mind for the propertyless-DLT enum has been LevelFormat (mnemonic: "type = format + properties").

So if you want to minimize future code churn, you should name this LevelFormat. Though if you'd rather wait until I do the s/DimLevelType/LevelType/g change, then I think StorageLevelType is a fine name for the interim.

171

Could you add some static_asserts that the DLT and SLT versions of this function agree? Or file a ticket to do so, assign it to me (so I don't forget), and add a TODO with the ticket number?

260

This name is too generic, since we could in principle construct a bunch of different things out of a DLT. I'd suggest something more like getStorageLevelType(/getLevelFormat).

266–267

I think it'd be better to just static_assert(std::is_same_v<std::underlying_type_t<StorageLevelType>, std::underlying_type_t<DimLevelType>>); That's enough to make sure their underlying types agree, but minimizes the number of places we'd need to change things if we decide we need to change the underlying type to something wider.

You should also add a sentence to the documentation explaining that this is a simplification to help ensure efficiency or whatever. (E.g., in the long run if we end up expanding the LevelType out to a larger struct, or add a ton of new properties, then it'd no longer make sense to require the LevelFormat have the same bitwidth.) My main concern here is that this is a rather different sort of assertion than the others, which are strictly about ensuring that the bitbashing is correct.

271–276

I'm not sure I like these assertions... The next block of assertions that fromDLT does what we expect, those are great. But these assertions that fromDLT is bitwise identity on the ordered+unique DLTs, these smell fishy to me.

For example, we could instead define SLT fromDLT(DLT x) { return static_cast<SLT>(static_cast<uint8_t>(x) >> 2); } (and update the enum class SLT accordingly). Why should the assertions rule that out? If there's any code that relies on fromDLT being bitwise identity on ordered+unique arguments, I expect that code is going to be depending on the underlying bit patterns of DLT in ways that it shouldn't.

In fact, I'd go so far as to say that the x >> 2 implementation should be the preferred one, precisely because it will catch out client code that buggily depends on the underlying bit patterns.

280–289

I'm a bit surprised that this compiles as written, since you haven't defined operator== on StorageLevelType. On a quick look I couldn't find anything in the spec that defines how default operator== works on enum class, so I'm not sure how portable this'll be; though it was just a quick look, so I may well have missed it.

(Personally, as a language designer, I'd think the default operator== should be fine here since it doesn't do anything to break the encapsulation that enum class provides. But just because I think it's reasonable doesn't mean the C++ standards committee does ;)

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

171

I think string comparison can not be a constexpr until C++20

[wrengr]

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

31

"to encode" or "for encoding"

31

"about"

31

Since the driving motivation for factoring out StorageLevelType and SparseStorageEncodingAttr is to remove the "properties" of DLT, you should avoid using that word here since it's liable to cause confusion.

So perhaps something more like, "An attribute for encoding TACO-style information about a storage format for represent ing sparse tensors, without the additional semantic properties encoded by SparseTensorEncodingAttr."

40

should be "storage level type" (or "level format")

40

should be "level"

43

"per level storage type" (or "per level storage format")

44

This should be plural.

53–54

Both of those should be "SparseTensorEncodingAttr" (since just saying "TensorEncodingAttr" suggests that it might be some more general thing that dense tensors could also have)

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

48

Is there any way you can get this to share more code with SparseTensorEncodingAttr::parse? It'll be a nightmare to try to keep two copies of almost identical parsers in sync

58

the name dlt doesn't match StorageLevelType

77

For DLT the undef value is purely for internal use (e.g., when we need to preinitialize an array of DLTs), and hence "undef" is intentionally not a valid MLIR token, since MLIR code should never construct the undef value.

Do you really mean for SLT::Undef to be something that MLIR code can construct?

77–84

You can reduce maintenance burden by instead having something more like: for (auto x : allSLTs) { if (strVal == toMLIRString(x)) { dlt.push_back(x); break; }}. That way we don't have to worry about keeping two sets of string literals in sync.

488

I think this is a perfectly good place to use auto. Also the variable name dlts is no longer appropriate with the new type

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

271–276

Okay, I will delete it. It is somehow overlapped with the tests below too.

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

77

This is a good question, I don't think undef should be in STEA in the first place (You can use the DLT internally, but it can not and should not be parsed in). Since you raised the question, I would suggest we fix it when parsing STEA as well (And verify fails when there is an undef DLT in the STEA)

[Peiming]

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

77

Oh, shot! I did not realize undef is not a branch when parsing STEA... But verification failure is still something we can add ;-)

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

171

Ah, yes, quite right.

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

77

I think it's better to fail during the parsing step rather than during the subsequent validation. Failing earlier typically makes it easier to give more relevant/targeted error messages, plus it reduces the cost/need for subsequent validation.

77–84

This implementation also makes it easy to fail-to-parse-undef, just don't have SLT::Undef in the allSLTs list you're looping over :)

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

162

I would mention the fact that this focuses on the "format" and not the "properties" (see comment above)

269

Ensure would be more consistent with others (I actually prefer the Ensures, but looking at other style in the file, I would go with majority)

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

26

, -> .

34

Read more in ....

Or,

For more information, see ....

194

Can we derive from the storage class to avoid duplicating these fields?
Or is that too much OO in tablegen. Also, it probably would require factoring out the "format" part....

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

79

aren't we duplicating a lot of very similar code now for the printing/parsing?

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

enum class DimLevelType : uint8_t {
   CompressedNo = 10,   // 0b010_10
   CompressedNuNo = 11, // 0b010_11
   Singleton = 16,      // 0b100_00
   SingletonNu = 17,    // 0b100_01
   SingletonNo = 18,    // 0b100_10
   SingletonNuNo = 19,  // 0b100_11
 };
 
+/// This enum defines all the sparse representations supportable by
+/// the SparseTensor dialect. Unlike DimLevelType, it does not encode level
+/// properties, which are irrelevant to sparse tensor storage scheme.
+enum class StorageLevelType : uint8_t {
+  Undef = 0,      // 0b000_00
+  Dense = 4,      // 0b001_00
+  Compressed = 8, // 0b010_00
+  Singleton = 16, // 0b100_00
+};
+
+/// Returns string representation of the given storage level type.
+inline std::string toMLIRString(StorageLevelType dlt) {
+  switch (dlt) {
+  case StorageLevelType::Undef:
+    return "\"undef\"";
+  case StorageLevelType::Dense:
+    return "\"dense\"";
+  case StorageLevelType::Compressed:
+    return "\"compressed\"";
+  case StorageLevelType::Singleton:
+    return "\"singleton\"";
+  }
+  return "";
+}
+
 /// 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\"";
constexpr bool isOrderedDLT(DimLevelType dlt) {
   return !(static_cast<uint8_t>(dlt) & 2);
 }
 
 /// Check if the `DimLevelType` is unique (regardless of storage format).
 constexpr bool isUniqueDLT(DimLevelType dlt) {
   return !(static_cast<uint8_t>(dlt) & 1);
 }
 
+/// Convert a DimLevelType to its corresponding StorageLevelType.
+constexpr StorageLevelType fromDLT(DimLevelType dlt) {
+  return static_cast<StorageLevelType>(static_cast<uint8_t>(dlt) & ~3);
+}
+
+/// Ensure both StorageLevelType and DimLevelType have the same underlying type.
+static_assert(
+    std::is_same_v<std::underlying_type_t<StorageLevelType>, uint8_t>);
+static_assert(std::is_same_v<std::underlying_type_t<DimLevelType>, uint8_t>);
+
+/// Ensures both StorageLevelType and DimLevelType use the same value to
+/// encoding the same level type.
+static_assert(static_cast<uint8_t>(StorageLevelType::Dense) ==
+              static_cast<uint8_t>(DimLevelType::Dense));
+static_assert(static_cast<uint8_t>(StorageLevelType::Compressed) ==
+              static_cast<uint8_t>(DimLevelType::Compressed));
+static_assert(static_cast<uint8_t>(StorageLevelType::Singleton) ==
+              static_cast<uint8_t>(DimLevelType::Singleton));
+
+/// Ensure the above conversion works as intended.
+static_assert(
+    (fromDLT(DimLevelType::Undef) == StorageLevelType::Undef &&
+     fromDLT(DimLevelType::Dense) == StorageLevelType::Dense &&
+     fromDLT(DimLevelType::Compressed) == StorageLevelType::Compressed &&
+     fromDLT(DimLevelType::CompressedNu) == StorageLevelType::Compressed &&
+     fromDLT(DimLevelType::CompressedNo) == StorageLevelType::Compressed &&
+     fromDLT(DimLevelType::CompressedNuNo) == StorageLevelType::Compressed &&
+     fromDLT(DimLevelType::Singleton) == StorageLevelType::Singleton &&
+     fromDLT(DimLevelType::SingletonNu) == StorageLevelType::Singleton &&
+     fromDLT(DimLevelType::SingletonNo) == StorageLevelType::Singleton &&
+     fromDLT(DimLevelType::SingletonNuNo) == StorageLevelType::Singleton),
+    "fromDLT conversion is broken");
+
 // Ensure the above predicates work as intended.
 static_assert((isValidDLT(DimLevelType::Undef) &&
                isValidDLT(DimLevelType::Dense) &&
                isValidDLT(DimLevelType::Compressed) &&
                isValidDLT(DimLevelType::CompressedNu) &&
                isValidDLT(DimLevelType::CompressedNo) &&
                isValidDLT(DimLevelType::CompressedNuNo) &&
                isValidDLT(DimLevelType::Singleton) &&

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

 // Sparse Tensor Type Encoding Attribute
 //===----------------------------------------------------------------------===//
 
 // All of the Tensor attributes will extend this class.
 class SparseTensor_Attr<string name,
                         list<Trait> traits = []>
     : AttrDef<SparseTensor_Dialect, name, traits>;
 
+// Sparse tensor storage attribute,
+def SparseStorageEncodingAttr : SparseTensor_Attr<"SparseStorageEncoding",
+         [ DeclareAttrInterfaceMethods<VerifiableTensorEncoding> ] > {
+  let mnemonic = "storage";
+  let description = [{
+    An attribute to encoding TACO-style information on sparsity properties
+    of tensors that are related to sparse tensor storage layout.
+
+    Read more from `SparseTensorEncodingAttr`.
+  }];
+
+  // Data in sparse tensor encoding.
+  let parameters = (
+    ins
+    // A dimension level type for each dimension of the tensor type.
+    ArrayRefParameter<
+      "::mlir::sparse_tensor::StorageLevelType",
+      "per dimension level type"
+      >: $storageLevelType,
+
+    // The required bit width for pointer storage.
+    "unsigned":$pointerBitWidth,
+    // The required bit width for index storage.
+    "unsigned":$indexBitWidth
+  );
+
+  let extraClassDeclaration = [{
+     // Build a StorageEncodingAttr from TensorEncodingAttr
+     static SparseStorageEncodingAttr fromTensorEncodingAttr(SparseTensorEncodingAttr attr);
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
 // Sparse tensor encoding attribute.
 def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
          [ DeclareAttrInterfaceMethods<VerifiableTensorEncoding> ] > {
   let mnemonic = "encoding";
 
   let description = [{
     An attribute to encode TACO-style information on sparsity properties
     of tensors. The encoding is eventually used by a **sparse compiler**
let parameters = (
     ArrayRefParameter<
       "::mlir::sparse_tensor::DimLevelType",
       "per dimension level type"
       >: $dimLevelType,
     // A dimension order on the indices of this tensor type.
     "AffineMap":$dimOrdering,
     // A mapping between the original and higher-ordering index space.
     "AffineMap":$higherOrdering,
     // The required bit width for pointer storage.
     "unsigned":$pointerBitWidth,
     // The required bit width for index storage.
     "unsigned":$indexBitWidth
   );
 
   let extraClassDeclaration = [{
     /// Returns the type for pointer storage based on pointerBitWidth
     Type getPointerType() const;

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

let description = [{
     Examples:
 
     ```mlir
     // A storage specifier that can be used to store storage scheme metadata from CSR matrix.
     !storage_specifier<#CSR>
     ```
   }];
 
-  let parameters = (ins SparseTensorEncodingAttr : $encoding);
+  let parameters = (ins SparseStorageEncodingAttr : $encoding);
   let builders = [
-    TypeBuilderWithInferredContext<(ins "SparseTensorEncodingAttr":$encoding), [{
+    TypeBuilderWithInferredContext<(ins "SparseStorageEncodingAttr":$encoding), [{
       assert(encoding && "sparse tensor encoding should not be null");
       return $_get(encoding.getContext(), encoding);
     }]>,
+    TypeBuilderWithInferredContext<(ins "SparseTensorEncodingAttr":$encoding), [{
+      return get(SparseStorageEncodingAttr::fromTensorEncodingAttr(encoding));
+    }]>,
     TypeBuilderWithInferredContext<(ins "Type":$type), [{
       return get(getSparseTensorEncoding(type));
     }]>,
     TypeBuilderWithInferredContext<(ins "Value":$tensor), [{
       return get(tensor.getType());
     }]>
   ];
 

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

static bool acceptBitWidth(unsigned bitWidth) {
   case 32:
   case 64:
     return true;
   default:
     return false;
   }
 }
 
+Attribute SparseStorageEncodingAttr::parse(AsmParser &parser, Type type) {
+  if (failed(parser.parseLess()))
+    return {};
+  // Parse the data as a dictionary.
+  DictionaryAttr dict;
+  if (failed(parser.parseAttribute(dict)))
+    return {};
+  if (failed(parser.parseGreater()))
+    return {};
+  // Process the data from the parsed dictionary value into struct-like data.
+  SmallVector<StorageLevelType> dlt;
+  unsigned ptr = 0;
+  unsigned ind = 0;
+  for (const NamedAttribute &attr : dict) {
+    if (attr.getName() == "storageLevelType") {
+      auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
+      if (!arrayAttr) {
+        parser.emitError(parser.getNameLoc(),
+                         "expected an array for dimension level types");
+        return {};
+      }
+      for (auto i : arrayAttr) {
+        auto strAttr = i.dyn_cast<StringAttr>();
+        if (!strAttr) {
+          parser.emitError(parser.getNameLoc(),
+                           "expected a string value in dimension level types");
+          return {};
+        }
+        auto strVal = strAttr.getValue();
+        if (strVal == "undef") {
+          dlt.push_back(StorageLevelType::Undef);
+        } else if (strVal == "dense") {
+          dlt.push_back(StorageLevelType::Dense);
+        } else if (strVal == "compressed") {
+          dlt.push_back(StorageLevelType::Compressed);
+        } else if (strVal == "singleton") {
+          dlt.push_back(StorageLevelType::Singleton);
+        } else {
+          parser.emitError(parser.getNameLoc(),
+                           "unexpected dimension level type: ")
+              << strVal;
+          return {};
+        }
+      }
+    } else if (attr.getName() == "pointerBitWidth") {
+      auto intAttr = attr.getValue().dyn_cast<IntegerAttr>();
+      if (!intAttr) {
+        parser.emitError(parser.getNameLoc(),
+                         "expected an integral pointer bitwidth");
+        return {};
+      }
+      ptr = intAttr.getInt();
+    } else if (attr.getName() == "indexBitWidth") {
+      auto intAttr = attr.getValue().dyn_cast<IntegerAttr>();
+      if (!intAttr) {
+        parser.emitError(parser.getNameLoc(),
+                         "expected an integral index bitwidth");
+        return {};
+      }
+      ind = intAttr.getInt();
+    } else {
+      parser.emitError(parser.getNameLoc(), "unexpected key: ")
+          << attr.getName().strref();
+      return {};
+    }
+  }
+  // Construct struct-like storage for attribute.
+  return parser.getChecked<SparseStorageEncodingAttr>(parser.getContext(), dlt,
+                                                      ptr, ind);
+}
+
+void SparseStorageEncodingAttr::print(AsmPrinter &printer) const {
+  // Print the struct-like storage in dictionary fashion.
+  printer << "<{ storageLevelType = [ ";
+  for (unsigned i = 0, e = getStorageLevelType().size(); i < e; i++) {
+    printer << toMLIRString(getStorageLevelType()[i]);
+    if (i != e - 1)
+      printer << ", ";
+  }
+  printer << " ]";
+  if (getPointerBitWidth())
+    printer << ", pointerBitWidth = " << getPointerBitWidth();
+  if (getIndexBitWidth())
+    printer << ", indexBitWidth = " << getIndexBitWidth();
+  printer << " }>";
+}
+
+LogicalResult SparseStorageEncodingAttr::verifyEncoding(
+    ArrayRef<int64_t> shape, Type elementType,
+    ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError) const {
+  if (acceptBitWidth(getPointerBitWidth()) &&
+      acceptBitWidth(getIndexBitWidth()))
+    return success();
+
+  return failure();
+}
+
+SparseStorageEncodingAttr SparseStorageEncodingAttr::fromTensorEncodingAttr(
+    SparseTensorEncodingAttr attr) {
+  SmallVector<StorageLevelType> slt;
+  for (auto dlt : attr.getDimLevelType())
+    slt.push_back(fromDLT(dlt));
+
+  return SparseStorageEncodingAttr::get(attr.getContext(), slt,
+                                        attr.getPointerBitWidth(),
+                                        attr.getIndexBitWidth());
+}
+
 Type SparseTensorEncodingAttr::getPointerType() const {
   unsigned ptrWidth = getPointerBitWidth();
   Type indexType = IndexType::get(getContext());
   return ptrWidth ? IntegerType::get(getContext(), ptrWidth) : indexType;
 }
 
 Type SparseTensorEncodingAttr::getIndexType() const {
   unsigned idxWidth = getIndexBitWidth();

 //===----------------------------------------------------------------------===//
 // Convenience Methods.
 //===----------------------------------------------------------------------===//
 
 SparseTensorEncodingAttr
 mlir::sparse_tensor::getSparseTensorEncoding(Type type) {
   if (auto ttp = type.dyn_cast<RankedTensorType>())
     return ttp.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
-  if (auto mdtp = type.dyn_cast<StorageSpecifierType>())
-    return mdtp.getEncoding();
   return nullptr;
 }
 
 bool mlir::sparse_tensor::isUniqueCOOType(RankedTensorType tp) {
   SparseTensorEncodingAttr enc = getSparseTensorEncoding(tp);
 
   if (!enc)
     return false;
verifySparsifierGetterSetter(StorageSpecifierKind mdKind, Optional<APInt> dim,
                              TypedValue<StorageSpecifierType> md,
                              Operation *op) {
   if (mdKind == StorageSpecifierKind::ValMemSize && dim) {
     return op->emitError(
         "redundant dimension argument for querying value memory size");
   }
 
   auto enc = md.getType().getEncoding();
-  ArrayRef<DimLevelType> dlts = enc.getDimLevelType();
+  ArrayRef<StorageLevelType> dlts = enc.getStorageLevelType();
   unsigned rank = dlts.size();
 
   if (mdKind != StorageSpecifierKind::ValMemSize) {
     if (!dim)
       return op->emitError("missing dimension argument");
 
     unsigned d = dim.value().getZExtValue();
     if (d >= rank)
       return op->emitError("requested dimension out of bound");
 
-    if (mdKind == StorageSpecifierKind::PtrMemSize && isSingletonDLT(dlts[d]))
+    if (mdKind == StorageSpecifierKind::PtrMemSize &&
+        dlts[d] == StorageLevelType::Singleton)
       return op->emitError(
           "requested pointer memory size on a singleton level");
   }
   return success();
 }
 
 LogicalResult NewOp::verify() {
   if (getExpandSymmetry() &&
       getResult().getType().cast<RankedTensorType>().getRank() != 2)
     return emitOpError("expand_symmetry can only be used for 2D tensors");
   return success();
 }
 
 LogicalResult ConvertOp::verify() {
   if (auto tp1 = getSource().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
       if (tp1.getRank() != tp2.getRank())
         return emitError("unexpected conversion mismatch in rank");
       auto shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       // Accept size matches between the source and the destination type
-      // (e.g. 10 vs. 10, 10 vs. ?, or ? vs. ?), but reject direct mismatches or
-      // matches that would need a runtime assert (e.g. 10 vs. 20 or ? vs. 10).
+      // (e.g. 10 vs. 10, 10 vs. ?, or ? vs. ?), but reject direct mismatches
+      // or matches that would need a runtime assert (e.g. 10 vs. 20 or ? vs.
+      // 10).
       for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++)
         if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamic)
           return emitError("unexpected conversion mismatch in dimension ") << d;
       return success();
     }
   }
   return emitError("unexpected type in convert");
 }

mlir/test/Dialect/SparseTensor/fold.mlir

func.func @sparse_concat_dce(%arg0: tensor<2xf64, #SparseVector>,
                              %arg2: tensor<4xf64, #SparseVector>) {
   %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
        : tensor<2xf64, #SparseVector>,
          tensor<3xf64, #SparseVector>,
          tensor<4xf64, #SparseVector> to tensor<9xf64, #SparseVector>
   return
 }
 
+#SparseVectorStorage = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
+
 // CHECK-LABEL: func @sparse_get_specifier_dce_fold(
 //  CHECK-SAME:  %[[A0:.*]]: !sparse_tensor.storage_specifier
 //  CHECK-SAME:  %[[A1:.*]]: i64,
 //  CHECK-SAME:  %[[A2:.*]]: i64)
 //   CHECK-NOT:  sparse_tensor.storage_specifier.set
 //   CHECK-NOT:  sparse_tensor.storage_specifier.get
 //       CHECK:  return %[[A1]]
-func.func @sparse_get_specifier_dce_fold(%arg0: !sparse_tensor.storage_specifier<#SparseVector>, %arg1: i64, %arg2: i64) -> i64 {
+func.func @sparse_get_specifier_dce_fold(%arg0: !sparse_tensor.storage_specifier<#SparseVectorStorage>, %arg1: i64, %arg2: i64) -> i64 {
   %0 = sparse_tensor.storage_specifier.set %arg0 dim_sz at 0 with %arg1
-       : i64, !sparse_tensor.storage_specifier<#SparseVector>
+       : i64, !sparse_tensor.storage_specifier<#SparseVectorStorage>
   %1 = sparse_tensor.storage_specifier.set %0 ptr_mem_sz at 0 with %arg2
-       : i64, !sparse_tensor.storage_specifier<#SparseVector>
+       : i64, !sparse_tensor.storage_specifier<#SparseVectorStorage>
   %2 = sparse_tensor.storage_specifier.get %1 dim_sz at 0
-       : !sparse_tensor.storage_specifier<#SparseVector> to i64
+       : !sparse_tensor.storage_specifier<#SparseVectorStorage> to i64
   return %2 : i64
 }

mlir/test/Dialect/SparseTensor/invalid.mlir

 func.func @mismatch_values_types(%arg0: tensor<?xf64, #SparseVector>) -> memref<?xf32> {
   // expected-error@+1 {{unexpected mismatch in element types}}
   %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf32>
   return %0 : memref<?xf32>
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   // expected-error@+1 {{redundant dimension argument for querying value memory size}}
   %0 = sparse_tensor.storage_specifier.get %arg0 val_mem_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   // expected-error@+1 {{missing dimension argument}}
   %0 = sparse_tensor.storage_specifier.get %arg0 idx_mem_sz
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   // expected-error@+1 {{requested dimension out of bound}}
   %0 = sparse_tensor.storage_specifier.get %arg0 dim_sz at 1
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }
 
 // -----
 
-#COO = #sparse_tensor.encoding<{dimLevelType = ["compressed-nu", "singleton"]}>
+#COO = #sparse_tensor.storage<{storageLevelType = ["compressed", "singleton"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#COO>) -> i64 {
   // expected-error@+1 {{requested pointer memory size on a singleton level}}
   %0 = sparse_tensor.storage_specifier.get %arg0 ptr_mem_sz at 1
        : !sparse_tensor.storage_specifier<#COO> to i64
   return %0 : i64
 }
 
 // -----
 
-#COO = #sparse_tensor.encoding<{dimLevelType = ["compressed-nu", "singleton"]}>
+#COO = #sparse_tensor.storage<{storageLevelType = ["compressed", "singleton"]}>
 
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#COO>) -> i64 {
   // expected-error@+1 {{type mismatch between requested }}
   %0 = sparse_tensor.storage_specifier.get %arg0 ptr_mem_sz at 0
        : !sparse_tensor.storage_specifier<#COO> to i32
   return %0 : i32
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 func.func @sparse_set_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>,
                          %arg1: i32)
           -> !sparse_tensor.storage_specifier<#SparseVector> {
   // expected-error@+1 {{type mismatch between requested }}
   %0 = sparse_tensor.storage_specifier.set %arg0 dim_sz at 0 with %arg1
        : i32, !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : !sparse_tensor.storage_specifier<#SparseVector>

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: return %[[T]] : memref<?xf64>
 func.func @sparse_values(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_metadata_init(
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#{{.*}}>
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_metadata_init() -> !sparse_tensor.storage_specifier<#SparseVector> {
   %0 = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : !sparse_tensor.storage_specifier<#SparseVector>
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_get_md(
 //  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.get %[[A]] dim_sz at 0
 //       CHECK: return %[[T]] : i64
 func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>) -> i64 {
   %0 = sparse_tensor.storage_specifier.get %arg0 dim_sz at 0
        : !sparse_tensor.storage_specifier<#SparseVector> to i64
   return %0 : i64
 }
 
 // -----
 
-#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#SparseVector = #sparse_tensor.storage<{storageLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_set_md(
 //  CHECK-SAME: %[[A:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>,
 //  CHECK-SAME: %[[I:.*]]: i64)
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.set %[[A]] dim_sz at 0 with %[[I]]
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_set_md(%arg0: !sparse_tensor.storage_specifier<#SparseVector>, %arg1: i64)
           -> !sparse_tensor.storage_specifier<#SparseVector> {