[mlir][sparse] introduce sparse_tensor::StorageSpecifierType and related operations on it

[Peiming]

update bazel file.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/CMakeLists.txt

7–8

Since we already have the MLIRSparseTensorEnums library, you should rename these to something else to avoid confusion. Elsewhere in MLIR they often just append "enums" to the name of the file containing the definitions (e.g., SparseTensorAttrDefsEnums.{h,cpp}.inc; though I'd shorten that to SparseTensorAttrEnums.{h,cpp}.inc since that's more legible)

[Peiming]

fix bazel

[wrengr]

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

181–184

Elsewhere we've been moving away from the "k" prefix on enum ctors. Do we really want the "k" here or no?

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

186–187

perhaps rephrase this to something more like "holds the sizes for dimensions, pointer arrays, index arrays, value arrays. This includes the buffer memory for storing those sizes, but not the buffer memory for the tensor itself." Though I'm not sure if the second sentence is really necessary (since imo it's implied by the preceding sentence).

[wrengr]

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

208

grammatically that should be "held by". Though I think it'd be better to say "Returns the requested field of the metadata."

212

"size of the"

212

This should be "level" rather than "dimension", and ditto for most (if not all) of the other places in this CL.

The one place I can think of where "dimension" is correct is for some of the uses of kDimSize. However, we will actually want to have both kDimSize and kLvlSize. See D139165 for more details about why we want both.

234

Grammatically that should be "Sets the requested field of". Though semantically it'd be better to rephrase the sentence as "Sets the field of the metadata to the given value."

239

"level".

239

"size of the"

[wrengr]

Also it looks like you forgot to commit the contents of mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.h

[Peiming]

address comments.

[Peiming]

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

181–184

I personally have no preference

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

186–187

Agree, I did not include the second sentence.

[Peiming]

remove leading k in enum class

[Peiming]

replace "sparse_tensor.storage_specifier" with "sparse_tensor.storage_specifier.init"

[Peiming]

update doc.

[aartbik]

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

23

if you prefer, you can avoid all of this by simply repeating the define

#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/SparseTensor/IR/SparseTensorEnums.h.inc"

#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"

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

16–20

Introduce a new section here, and below for the new stuff
(and probably for the traits implementation)

just so we split the file a bit into parts

===----------------------------------------------------------------------===
Sparse Tensor Type Encoding Attribute
===----------------------------------------------------------------------===//

178

period at end

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

33

is this still in the making?

85

newline?

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

344

can we move this to the top with the other #define/include parts above?
I realize we do this "mid file" trick at the end near L810, but that has its own reason

that way we have less includes spread out

[Peiming]

address comments from Aart.

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

23

I actually learned the trick for IREE through code search. But yeah, your way seems much more clear. Updated!

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

33

Overlooked it... Now complete

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

344

Okay. make sense!

[aartbik]

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

182

still missin period after kind ;-)

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

184

Return an initial ....

Also, I think initial can use some explanation here

185

index arrays, and

(Oxford comma)

208

In all other ops, we either just say Example, or add explanation to the "header".
So

Example of querying the size ....

would be more consistent?

235

same suggestion, combine with Example header

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

25

wait, this is not Sparse Tensor Types, but Storage Specifiers right? Different header?

Although we have a bit of a misnomer filename (SparseTensorTypes, should that be StorageSpecifierTypes?
(It is getting silly, I know ;-)

[aartbik]

Approving pending addressing the comments.

[Peiming]

address comments.

[Peiming]

fix typo.

mlir/include/mlir/Dialect/SparseTensor/IR/CMakeLists.txt

 add_mlir_dialect(SparseTensorOps sparse_tensor)
 add_mlir_doc(SparseTensorOps SparseTensorOps Dialects/ -gen-dialect-doc)
 
 set(LLVM_TARGET_DEFINITIONS SparseTensorAttrDefs.td)
 mlir_tablegen(SparseTensorAttrDefs.h.inc -gen-attrdef-decls)
 mlir_tablegen(SparseTensorAttrDefs.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(SparseTensorAttrEnums.h.inc -gen-enum-decls)
+mlir_tablegen(SparseTensorAttrEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRSparseTensorAttrDefsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS SparseTensorTypes.td)
+mlir_tablegen(SparseTensorTypes.h.inc -gen-typedef-decls)
+mlir_tablegen(SparseTensorTypes.cpp.inc -gen-typedef-defs)
+add_public_tablegen_target(MLIRSparseTensorTypesIncGen)
+

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

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TensorEncoding.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+// We must include Enums.h.inc before AttrDefs.h.inc due to dependency between
+// StorageSpecifierKindAttr and StorageSpeciferKind Enum.
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.h.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"
 
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.h.inc"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
 
 namespace mlir {
 namespace sparse_tensor {
 

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

 //===-- SparseTensorAttrDefs.td - attributes definitions ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPARSETENSOR_ATTRDEFS
 #define SPARSETENSOR_ATTRDEFS
 
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
 include "mlir/IR/TensorEncoding.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 encoding attribute.
 def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
          [ DeclareAttrInterfaceMethods<VerifiableTensorEncoding> ] > {
let extraClassDeclaration = [{
     /// Returns the type for index storage based on indexBitWidth
     Type getIndexType() const;
   }];
 
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Sparse Tensor Storage Specifier Enum Attribute
+//===----------------------------------------------------------------------===//
+
+// The C++ enum for Storage Specifier kind.
+def SparseTensorStorageSpecifierKindEnum
+    : I32EnumAttr<"StorageSpecifierKind", "sparse tensor storage specifier kind", [
+        I32EnumAttrCase<"DimSize",    0, "dim_sz">,
+        I32EnumAttrCase<"PtrMemSize", 1, "ptr_mem_sz">,
+        I32EnumAttrCase<"IdxMemSize", 2, "idx_mem_sz">,
+        I32EnumAttrCase<"ValMemSize", 3, "val_mem_sz">,
+      ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = SparseTensor_Dialect.cppNamespace;
+}
+
+// Define the enum StorageSpecifier kind attribute.
+def SparseTensorStorageSpecifierKindAttr
+    : EnumAttr<SparseTensor_Dialect, SparseTensorStorageSpecifierKindEnum,
+               "SparseTensorStorageSpecifierKind"> {
+   let mnemonic = "kind";
+}
+
+//===----------------------------------------------------------------------===//
+// Sparse Tensor Traits
+//===----------------------------------------------------------------------===//
+
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,
 // `RankedTensorOf`, `AnyRankedTensor`.
 
 class SparseTensorOf<list<Type> allowedTypes>
   : TensorOf<allowedTypes, [IsSparseTensorPred], "sparse tensor">;

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

let description = [{
     and Saman Amarasinghe. Tensor Algebra Compilation with Workspaces,
     Proceedings of the IEEE/ACM International Symposium on Code Generation
     and Optimization, 2019.
     * [Kjolstad20] Fredrik Berg Kjolstad. Sparse Tensor Algebra Compilation.
     PhD thesis, MIT, February, 2020.
   }];
 
   let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
 }
 
 #endif // SPARSETENSOR_BASE

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

 //===- SparseTensorOps.td - Sparse tensor dialect ops ------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPARSETENSOR_OPS
 #define SPARSETENSOR_OPS
 
 include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
 include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
+include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // Base class.
 //===----------------------------------------------------------------------===//
 
 class SparseTensor_Op<string mnemonic, list<Trait> traits = []>
let description = [{
     ```mlir
     %1 = sparse_tensor.values %0 : tensor<64x64xf64, #CSR> to memref<?xf64>
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
   let hasVerifier = 1;
 }
 
+def SparseTensor_StorageSpecifierInitOp : SparseTensor_Op<"storage_specifier.init", [Pure]>,
+    Results<(outs SparseTensorStorageSpecifier:$result)> {
+  let summary = "";
+  let description = [{
+    Returns an initial storage specifier value. A storage specifier value holds
+    the sizes for tensor dimensions, pointer arrays, index arrays, and the value array.
+
+    Example:
+
+    ```mlir
+    %0 = sparse_tensor.storage_specifier.init :  !sparse_tensor.storage_specifier<#CSR>
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict `:` qualified(type($result))";
+}
+
+def SparseTensor_GetStorageSpecifierOp : SparseTensor_Op<"storage_specifier.get", [Pure]>,
+    Arguments<(ins SparseTensorStorageSpecifier:$specifier,
+                   SparseTensorStorageSpecifierKindAttr:$specifierKind,
+                   OptionalAttr<IndexAttr>:$dim)>,
+    Results<(outs AnyType:$result)> {
+  let summary = "";
+  let description = [{
+    Returns the requested field of the given storage_specifier.
+
+    Example of querying the size of the index array for level 0:
+
+    ```mlir
+    %0 = sparse_tensor.storage_specifier.get %arg0 idx_mem_sz at 0
+         : !sparse_tensor.storage_specifier<#COO> to i64
+    ```
+  }];
+
+  let assemblyFormat = "$specifier $specifierKind (`at` $dim^)? attr-dict `:` "
+                       "qualified(type($specifier)) `to` type($result)";
+  let hasVerifier = 1;
+}
+
+def SparseTensor_SetStorageSpecifierOp : SparseTensor_Op<"storage_specifier.set",
+    [Pure, AllTypesMatch<["result", "specifier"]>]>,
+    Arguments<(ins SparseTensorStorageSpecifier:$specifier,
+                   SparseTensorStorageSpecifierKindAttr:$specifierKind,
+                   OptionalAttr<IndexAttr>:$dim,
+                   AnyType:$value)>,
+    Results<(outs SparseTensorStorageSpecifier:$result)> {
+  let summary = "";
+  let description = [{
+    Set the field of the storage specifier to the given input value. Returns
+    the updated storage_specifier as a new SSA value.
+
+    Example of updating the sizes of the index array for level 0:
+
+    ```mlir
+    %0 = sparse_tensor.storage_specifier.set %arg0 idx_mem_sz at 0 with %new_sz
+       : i32, !sparse_tensor.storage_specifier<#COO>
+
+    ```
+  }];
+  let assemblyFormat = "$specifier $specifierKind (`at` $dim^)? `with` $value attr-dict `:` "
+                       "type($value) `,` qualified(type($result))";
+  let hasVerifier = 1;
+}
+
 def SparseTensor_NumberOfEntriesOp : SparseTensor_Op<"number_of_entries", [Pure]>,
     Arguments<(ins AnySparseTensor:$tensor)>,
     Results<(outs Index:$result)> {
   let summary = "Returns the number of entries that are stored in the tensor.";
   let description = [{
     Returns the number of entries that are stored in the given sparse tensor.
     Note that this is typically the number of nonzero elements in the tensor,
     but since explicit zeros may appear in the storage formats, the more

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

+//===- SparseTensorTypes.td - Sparse tensor dialect types ------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef SPARSETENSOR_TYPES
+#define SPARSETENSOR_TYPES
+
+include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
+include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
+
+//===----------------------------------------------------------------------===//
+// Base class.
+//===----------------------------------------------------------------------===//
+
+// Base class for Builtin dialect types.
+class SparseTensor_Type<string name, list<Trait> traits = [],
+                   string baseCppClass = "::mlir::Type">
+    : TypeDef<SparseTensor_Dialect, name, traits, baseCppClass> {}
+
+//===----------------------------------------------------------------------===//
+// Sparse Tensor Dialect Types.
+//===----------------------------------------------------------------------===//
+
+def SparseTensor_StorageSpecifier : SparseTensor_Type<"StorageSpecifier"> {
+  let mnemonic = "storage_specifier";
+  let summary = "Structured metadata for sparse tensor low-level storage scheme";
+  
+  let description = [{
+    Syntax:
+
+    ```
+    storage_specifier-type ::= `!storage_specifier` `<` encoding `>`
+    encoding ::= attribute-value
+    ```
+
+    Values with storage_specifier types represent aggregated storage scheme metadata
+    for the given sparse tensor encoding.
+    It currently holds a set of values for sizes of sparse tensor dimension, index array,
+    pointer array and value array.
+    Note that the type is not yet stable and subject to change in the near future.
+    
+    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 builders = [
+    TypeBuilderWithInferredContext<(ins "SparseTensorEncodingAttr":$encoding), [{
+      assert(encoding && "sparse tensor encoding should not be null");
+      return $_get(encoding.getContext(), encoding);
+    }]>,
+    TypeBuilderWithInferredContext<(ins "Type":$type), [{
+      return get(getSparseTensorEncoding(type));
+    }]>,
+    TypeBuilderWithInferredContext<(ins "Value":$tensor), [{
+      return get(tensor.getType());
+    }]>
+  ];
+  
+  let extraClassDeclaration = [{
+    // Get the integer type used to store memory and dimension sizes.
+    IntegerType getSizesType() const;
+    Type getFieldType(StorageSpecifierKind kind, Optional<unsigned> dim) const;
+    Type getFieldType(StorageSpecifierKind kind, Optional<APInt> dim) const;
+  }];
+
+  let assemblyFormat="`<` qualified($encoding) `>`";
+}
+
+def IsSparseTensorStorageSpecifierTypePred
+    : CPred<"$_self.isa<::mlir::sparse_tensor::StorageSpecifierType>()">;
+    
+def SparseTensorStorageSpecifier
+    : Type<CPred<"$_self.isa<::mlir::sparse_tensor::StorageSpecifierType>()">, "metadata",
+          "::mlir::sparse_tensor::StorageSpecifierType">;
+    
+#endif // SPARSETENSOR_TYPES

mlir/lib/Dialect/SparseTensor/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRSparseTensorDialect
   SparseTensorDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS
   MLIRSparseTensorAttrDefsIncGen
   MLIRSparseTensorOpsIncGen
+  MLIRSparseTensorTypesIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialect
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSupport
   MLIRSparseTensorEnums
   )

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

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.cpp.inc"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.cpp.inc"
+
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Attribute Methods.
 //===----------------------------------------------------------------------===//
 
-#define GET_ATTRDEF_CLASSES
-#include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
-
 static bool acceptBitWidth(unsigned bitWidth) {
   switch (bitWidth) {
   case 0:
   case 8:
   case 16:
   case 32:
   case 64:
     return true;

 //===----------------------------------------------------------------------===//
 // 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;

 }
 
 uint64_t mlir::sparse_tensor::toStoredDim(RankedTensorType type, uint64_t d) {
   assert(d < static_cast<uint64_t>(type.getRank()));
   return toStoredDim(getSparseTensorEncoding(type), d);
 }
 
 //===----------------------------------------------------------------------===//
-// TensorDialect Operations.
+// SparseTensorDialect Types.
+//===----------------------------------------------------------------------===//
+
+IntegerType StorageSpecifierType::getSizesType() const {
+  unsigned idxBitWidth =
+      getEncoding().getIndexBitWidth() ? getEncoding().getIndexBitWidth() : 64u;
+  unsigned ptrBitWidth =
+      getEncoding().getIndexBitWidth() ? getEncoding().getIndexBitWidth() : 64u;
+
+  return IntegerType::get(getContext(), std::max(idxBitWidth, ptrBitWidth));
+}
+
+Type StorageSpecifierType::getFieldType(StorageSpecifierKind kind,
+                                        Optional<unsigned> dim) const {
+  if (kind != StorageSpecifierKind::ValMemSize)
+    assert(dim);
+
+  // Right now, we store every sizes metadata using the same size type.
+  // TODO: the field size type can be defined dimensional wise after sparse
+  // tensor encoding supports per dimension index/pointer bitwidth.
+  return getSizesType();
+}
+
+Type StorageSpecifierType::getFieldType(StorageSpecifierKind kind,
+                                        Optional<APInt> dim) const {
+  Optional<unsigned> intDim = std::nullopt;
+  if (dim)
+    intDim = dim.value().getZExtValue();
+  return getFieldType(kind, intDim);
+}
+
+//===----------------------------------------------------------------------===//
+// SparseTensorDialect Operations.
 //===----------------------------------------------------------------------===//
 
 static LogicalResult isInBounds(uint64_t dim, Value tensor) {
   uint64_t rank = tensor.getType().cast<RankedTensorType>().getRank();
   if (dim >= rank)
     return failure();
   return success(); // in bounds
 }
 
 static LogicalResult isMatchingWidth(Value result, unsigned width) {
   Type etp = result.getType().cast<MemRefType>().getElementType();
   if ((width == 0 && etp.isIndex()) || (width > 0 && etp.isInteger(width)))
     return success();
   return failure();
 }
 
+static LogicalResult
+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();
+  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]))
+      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() {

 LogicalResult ToValuesOp::verify() {
   RankedTensorType ttp = getTensor().getType().cast<RankedTensorType>();
   MemRefType mtp = getResult().getType().cast<MemRefType>();
   if (ttp.getElementType() != mtp.getElementType())
     return emitError("unexpected mismatch in element types");
   return success();
 }
 
+LogicalResult GetStorageSpecifierOp::verify() {
+  if (failed(verifySparsifierGetterSetter(getSpecifierKind(), getDim(),
+                                          getSpecifier(), getOperation()))) {
+    return failure();
+  }
+
+  // Checks the result type
+  if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=
+      getResult().getType()) {
+    return emitError(
+        "type mismatch between requested specifier field and result value");
+  }
+  return success();
+}
+
+LogicalResult SetStorageSpecifierOp::verify() {
+  if (failed(verifySparsifierGetterSetter(getSpecifierKind(), getDim(),
+                                          getSpecifier(), getOperation()))) {
+    return failure();
+  }
+
+  // Checks the input type
+  if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=
+      getValue().getType()) {
+    return emitError(
+        "type mismatch between requested specifier field and input value");
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TensorDialect Linalg.Generic Operations.
 //===----------------------------------------------------------------------===//
 
 template <class T>
 static LogicalResult verifyNumBlockArgs(T *op, Region &region,
                                         const char *regionName,
                                         TypeRange inputTypes, Type outputType) {

 // TensorDialect Methods.
 //===----------------------------------------------------------------------===//
 
 void SparseTensorDialect::initialize() {
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
       >();
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.cpp.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.cpp.inc"
       >();
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.cpp.inc"
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.cpp.inc"

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"]}>
+
+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"]}>
+
+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"]}>
+
+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"]}>
+
+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"]}>
+
+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"]}>
+
+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>
+}
+
+// -----
+
 func.func @sparse_unannotated_load(%arg0: tensor<16x32xf64>) -> tensor<16x32xf64> {
   // expected-error@+1 {{'sparse_tensor.load' op operand #0 must be sparse tensor of any type values, but got 'tensor<16x32xf64>'}}
   %0 = sparse_tensor.load %arg0 : tensor<16x32xf64>
   return %0 : tensor<16x32xf64>
 }
 
 // -----
 

mlir/test/Dialect/SparseTensor/roundtrip.mlir

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"]}>
 
+// 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"]}>
+
+// 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"]}>
+
+// 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> {
+  %0 = sparse_tensor.storage_specifier.set %arg0 dim_sz at 0 with %arg1
+       : i64, !sparse_tensor.storage_specifier<#SparseVector>
+  return %0 : !sparse_tensor.storage_specifier<#SparseVector>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
 // CHECK-LABEL: func @sparse_noe(
 //  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>)
 //       CHECK: %[[T:.*]] = sparse_tensor.number_of_entries %[[A]] : tensor<128xf64, #{{.*}}>
 //       CHECK: return %[[T]] : index
 func.func @sparse_noe(%arg0: tensor<128xf64, #SparseVector>) -> index {
   %0 = sparse_tensor.number_of_entries %arg0 : tensor<128xf64, #SparseVector>
   return %0 : index
 }
func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>, %arg1: f32) -> () {
   %ret = sparse_tensor.foreach in %arg0 init(%arg1): tensor<2x4xf64, #DCSR>, f32 -> f32
   do {
     ^bb0(%1: index, %2: index, %v: f64, %r: f32) :
       sparse_tensor.yield %r : f32
   }
   return
 }
 
-// ----
+// -----
 
 // CHECK-LABEL: func @sparse_sort_1d0v(
 //  CHECK-SAME: %[[A:.*]]: index,
 //  CHECK-SAME: %[[B:.*]]: memref<?xindex>)
 //       CHECK: sparse_tensor.sort %[[A]], %[[B]] : memref<?xindex>
 //       CHECK: return %[[B]]
 func.func @sparse_sort_1d0v(%arg0: index, %arg1: memref<?xindex>) -> (memref<?xindex>) {
   sparse_tensor.sort %arg0, %arg1 : memref<?xindex>

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 ##---------------------------------------------------------------------------##
 
 td_library(
     name = "SparseTensorTdFiles",
     srcs = [
         "include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td",
         "include/mlir/Dialect/SparseTensor/IR/SparseTensorBase.td",
         "include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td",
+        "include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.td",
     ],
     includes = ["include"],
     deps = [
         ":InferTypeOpInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
     ],
 )
 
 gentbl_cc_library(
     name = "SparseTensorAttrDefsIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             ["--gen-attrdef-decls"],
             "include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc",
         ),
         (
             ["--gen-attrdef-defs"],
             "include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc",
         ),
+        (
+            ["--gen-enum-decls"],
+            "include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.h.inc",
+        ),
+        (
+            ["--gen-enum-defs"],
+            "include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.cpp.inc",
+        ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td",
     deps = [":SparseTensorTdFiles"],
 )
 
 gentbl_cc_library(
     name = "SparseTensorOpsIncGen",
tbl_outs = [
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td",
     deps = [":SparseTensorTdFiles"],
 )
 
 gentbl_cc_library(
+    name = "SparseTensorTypesIncGen",
+    strip_include_prefix = "include",
+    tbl_outs = [
+        (
+            ["--gen-typedef-decls"],
+            "include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.h.inc",
+        ),
+        (
+            ["--gen-typedef-defs"],
+            "include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.cpp.inc",
+        ),
+    ],
+    tblgen = ":mlir-tblgen",
+    td_file = "include/mlir/Dialect/SparseTensor/IR/SparseTensorTypes.td",
+    deps = [":SparseTensorTdFiles"],
+)
+
+gentbl_cc_library(
     name = "SparseTensorPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             [
                 "-gen-pass-decls",
                 "-name=SparseTensor",
             ],
cc_library(
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":SparseTensorAttrDefsIncGen",
         ":SparseTensorEnums",
         ":SparseTensorOpsIncGen",
+        ":SparseTensorTypesIncGen",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorUtils",
     srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),