[mlir][sparse] Make sparse_tensor::StorageLayout publicly available.

[Peiming]

rebase.

[Peiming]

revert unintended changes.

[Peiming]

update bazel file

[Peiming]

minor fixes.

[Peiming]

remove outdated files.

[wrengr]

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

14

This macro should match the name of the file

23

Since this whole section (lines 22–82) is intended as documentation, it should use "///" in lieu of "//".

24

Although this probably unchanged from before, this modifier is very strange. I'd suggest dropping it (without adding anything in its place).

85–97

Instead of doing these static-assertions, it'd be a lot simpler and clearer to define the SparseTensorFieldKind elements directly in terms of their corresponding StorageSpecifierKind elements. I've attached a change which demonstrates what I mean.

I didn't adjust the definition of SparseTensorFieldKind::StorageSpec since it didn't have a corresponding static_assert; but you can always repeat the pattern if you wanted it to match StorageSpecifierKind::LvlSize. (And if you're paranoid about the possibility of the StorageSpec clashing with the other elements, then you can always add static asserts to check that.)

121

it'd be better to use const STT& (since STT doesn't use the pimpl idiom).

122

afaict there's no benefit to marking this parameter const (since STEA uses the pimpl idiom)

154

This defn should be moved to the cpp file, rather than being defined in the class (since it's too big to benefit from being inlined)

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

46

I think it would be better to move this whole section into a separate "SparseTensorStorageLayout.cpp" file, rather than putting it here. (Fwiw, in my work on updating STEA to use the new syntax, I'm also planning on splitting all the methods for that type off into their own cpp file as well)

49

this should be static constexpr

55

this should be "lvl"

56

ditto

59

be sure to rebase this CL over D150330 (and D150822, D150830)

64

should be "level"

66

This sentence is unnecessary (since level-types are always in the same order as the levels themselves) and confusing (since it does not match the new "dimension"/"level" terminology).

70–73

I think it would be a lot cleaner to define new helper predicates: constexpr bool hasPosDLT(dlt) { return isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt); } and constexpr bool hasCrdDLT(dlt) { return isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt) || isSingletonDLT(dlt); } and then use those here instead. Especially because those predicates will also be helpful elsewhere as well.

80

You should define static constexpr Level kInvalidLevel = -1u; and use that here, instead of using the literal -1u.

84

ditto

101

Should be "specType" or "specifierType" (since we renamed the StorageSpecifierType)

145

should be "specifier", right? Also, isn't the specifier actually the first field rather than the last?

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorDescriptor.h

65–66

Since you have a StorageLayout ctor which takes the STT directly (and calls getEncoding itself), you should just use "layout(stt)" here

66

You should have the StorageLayout ctor assert that the STEA isn't null, instead of doing this assertion here.

[Peiming]

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

46

I think a better way is to keep it in SparseTensorDialect.cpp, but moving the other ops-related part into SparseTensorOps.cpp as other dialects in MLIR.

Either way, I will probably keep it this way for this revision.

70–73

I will address this in a separate patch.

145

No, the specifier is the last field.

[Peiming]

rebase and address comments.

[aartbik]

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

29

I think this is the fourth time we moved this description around ;-)

135

since you document every getter except the first one, you might as well get rid of this "section header" and add doc to L137

157

Here it seems wortwhile to add a "section header" and say that some method (inline or prototype) follow

Storage related methods.
//

or something like that

179

extra line before closing namespaces

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

46

don't use the sparse_tensor:: qualifier here to be consistent with all the other header docs in this file

46

Yeah, right now we have one SparseTensorDialect.cpp to "rule them all" and the definitely makes things simpler (getting all the deps right for example).
But the file did grow in size over time.

Something to think about for the future, but I am okay with the current placement

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorDescriptor.h

1

this top level comment is now out of date

13

so is this macro (seems the mistake goes even one generation back?)

[Peiming]

address comments.

[wrengr]

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

114–116

It's cleaner to delegate to the other ctor.

134

"Gets"

[wrengr]

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

151

Since FieldIndex isn't Level, I don't think you should use kInvalidLevel here. So I guess you need to define kInvalidField too.

[Peiming]

address comments.

[Peiming]

rebase.

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

+//===- SparseTensorStorageLayout.h ------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines utilities for the sparse memory layout.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORBUILDER_H_
+#define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORBUILDER_H_
+
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
+
+namespace mlir {
+namespace sparse_tensor {
+
+//===----------------------------------------------------------------------===//
+// The sparse tensor storage scheme for a rank-dimensional tensor is organized
+// as a single compound type with the following fields. Note that every memref
+// with ? size actually behaves as a "vector", i.e. the stored size is the
+// capacity and the used size resides in the storage_specifier struct.
+//
+// struct {
+//   ; per-level l:
+//   ;  if dense:
+//        <nothing>
+//   ;  if compresed:
+//        memref<? x pos>  positions-l   ; positions for sparse level l
+//        memref<? x crd>  coordinates-l ; coordinates for sparse level l
+//   ;  if singleton:
+//        memref<? x crd>  coordinates-l ; coordinates for singleton level l
+//
+//   memref<? x eltType> values        ; values
+//
+//   struct sparse_tensor.storage_specifier {
+//     array<rank x int> lvlSizes    ; sizes/cardinalities for each level
+//     array<n x int> memSizes;      ; sizes/lengths for each data memref
+//   }
+// };
+//
+// In addition, for a "trailing COO region", defined as a compressed level
+// followed by one or more singleton levels, the default SOA storage that
+// is inherent to the TACO format is optimized into an AOS storage where
+// all coordinates of a stored element appear consecutively.  In such cases,
+// a special operation (sparse_tensor.coordinates_buffer) must be used to
+// access the AOS coordinates array. In the code below, the method `getCOOStart`
+// is used to find the start of the "trailing COO region".
+//
+// If the sparse tensor is a slice (produced by `tensor.extract_slice`
+// operation), instead of allocating a new sparse tensor for it, it reuses the
+// same sets of MemRefs but attaching a additional set of slicing-metadata for
+// per-dimension slice offset and stride.
+//
+// Examples.
+//
+// #CSR storage of 2-dim matrix yields
+//   memref<?xindex>                           ; positions-1
+//   memref<?xindex>                           ; coordinates-1
+//   memref<?xf64>                             ; values
+//   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
+//
+// #COO storage of 2-dim matrix yields
+//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
+//   memref<?xindex>                           ; AOS coordinates storage
+//   memref<?xf64>                             ; values
+//   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
+//
+// Slice on #COO storage of 2-dim matrix yields
+//   ;; Inherited from the original sparse tensors
+//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
+//   memref<?xindex>                           ; AOS coordinates storage
+//   memref<?xf64>                             ; values
+//   struct<(array<2 x i64>, array<3 x i64>,   ; lvl0, lvl1, 3xsizes
+//   ;; Extra slicing-metadata
+//           array<2 x i64>, array<2 x i64>)>) ; dim offset, dim stride.
+//
+//===----------------------------------------------------------------------===//
+
+enum class SparseTensorFieldKind : uint32_t {
+  StorageSpec = 0,
+  PosMemRef = 1,
+  CrdMemRef = 2,
+  ValMemRef = 3
+};
+
+static_assert(static_cast<uint32_t>(SparseTensorFieldKind::PosMemRef) ==
+              static_cast<uint32_t>(StorageSpecifierKind::PosMemSize));
+static_assert(static_cast<uint32_t>(SparseTensorFieldKind::CrdMemRef) ==
+              static_cast<uint32_t>(StorageSpecifierKind::CrdMemSize));
+static_assert(static_cast<uint32_t>(SparseTensorFieldKind::ValMemRef) ==
+              static_cast<uint32_t>(StorageSpecifierKind::ValMemSize));
+
+inline StorageSpecifierKind toSpecifierKind(SparseTensorFieldKind kind) {
+  assert(kind != SparseTensorFieldKind::StorageSpec);
+  return static_cast<StorageSpecifierKind>(kind);
+}
+
+inline SparseTensorFieldKind toFieldKind(StorageSpecifierKind kind) {
+  assert(kind != StorageSpecifierKind::LvlSize);
+  return static_cast<SparseTensorFieldKind>(kind);
+}
+
+/// The type of field indices.  This alias is to help code be more
+/// self-documenting; unfortunately it is not type-checked, so it only
+/// provides documentation rather than doing anything to prevent mixups.
+using FieldIndex = unsigned;
+
+/// Provides methods to access fields of a sparse tensor with the given
+/// encoding.
+class StorageLayout {
+public:
+  // TODO: Functions/methods marked with [NUMFIELDS] might should use
+  // `FieldIndex` for their return type, via the same reasoning for why
+  // `Dimension`/`Level` are used both for identifiers and ranks.
+  explicit StorageLayout(const SparseTensorType stt) : enc(stt.getEncoding()) {}
+  explicit StorageLayout(const SparseTensorEncodingAttr enc) : enc(enc) {}
+
+  /// For each field that will be allocated for the given sparse tensor
+  /// encoding, calls the callback with the corresponding field index,
+  /// field kind, level, and level-type (the last two are only for level
+  /// memrefs).  The field index always starts with zero and increments
+  /// by one between each callback invocation.  Ideally, all other methods
+  /// should rely on this function to query a sparse tensor fields instead
+  /// of relying on ad-hoc index computation.
+  void foreachField(
+      llvm::function_ref<bool(
+          FieldIndex /*fieldIdx*/, SparseTensorFieldKind /*fieldKind*/,
+          Level /*lvl (if applicable)*/, DimLevelType /*DLT (if applicable)*/)>)
+      const;
+
+  ///
+  /// Getters: get the field index for required field.
+  ///
+
+  FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind,
+                                 std::optional<Level> lvl) const {
+    return getFieldIndexAndStride(kind, lvl).first;
+  }
+
+  /// Gets the total number of fields for the given sparse tensor encoding.
+  unsigned getNumFields() const;
+
+  /// Gets the total number of data fields (coordinate arrays, position
+  /// arrays, and a value array) for the given sparse tensor encoding.
+  unsigned getNumDataFields() const;
+
+  std::pair<FieldIndex, unsigned>
+  getFieldIndexAndStride(SparseTensorFieldKind kind,
+                         std::optional<Level> lvl) const {
+    FieldIndex fieldIdx = -1u;
+    unsigned stride = 1;
+    if (kind == SparseTensorFieldKind::CrdMemRef) {
+      assert(lvl.has_value());
+      const Level cooStart = getCOOStart(enc);
+      const Level lvlRank = enc.getLvlRank();
+      if (lvl.value() >= cooStart && lvl.value() < lvlRank) {
+        lvl = cooStart;
+        stride = lvlRank - cooStart;
+      }
+    }
+    foreachField([lvl, kind, &fieldIdx](FieldIndex fIdx,
+                                        SparseTensorFieldKind fKind, Level fLvl,
+                                        DimLevelType dlt) -> bool {
+      if ((lvl && fLvl == lvl.value() && kind == fKind) ||
+          (kind == fKind && fKind == SparseTensorFieldKind::ValMemRef)) {
+        fieldIdx = fIdx;
+        // Returns false to break the iteration.
+        return false;
+      }
+      return true;
+    });
+    assert(fieldIdx != -1u);
+    return std::pair<FieldIndex, unsigned>(fieldIdx, stride);
+  }
+
+private:
+  const SparseTensorEncodingAttr enc;
+};
+
+// TODO: See note [NUMFIELDS].
+inline unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+  return StorageLayout(enc).getNumFields();
+}
+
+// TODO: See note [NUMFIELDS].
+inline unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+  return StorageLayout(enc).getNumDataFields();
+}
+
+inline void foreachFieldInSparseTensor(
+    SparseTensorEncodingAttr enc,
+    llvm::function_ref<bool(FieldIndex, SparseTensorFieldKind, Level,
+                            DimLevelType)>
+        callback) {
+  return StorageLayout(enc).foreachField(callback);
+}
+
+void foreachFieldAndTypeInSparseTensor(
+    SparseTensorType,
+    llvm::function_ref<bool(Type, FieldIndex, SparseTensorFieldKind, Level,
+                            DimLevelType)>);
+} // namespace sparse_tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORBUILDER_H_

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

 //===- SparseTensorDialect.cpp - Sparse tensor dialect implementation -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include <utility>
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 
 #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 "mlir/IR/PatternMatch.h"

 /// Gets the dimension-rank of the type of some `T`.  (In particular
 /// this is only used for `Value` and `TypedValue<RankedTensorType>`.)
 template <typename T>
 static inline Dimension getDimRank(T t) {
   return getRankedTensorType(t).getRank();
 }
 
 //===----------------------------------------------------------------------===//
+// sparse_tensor::StorageLayout
+//===----------------------------------------------------------------------===//
+
+constexpr FieldIndex kDataFieldStartingIdx = 0;
+
+void StorageLayout::foreachField(
+    llvm::function_ref<bool(FieldIndex, SparseTensorFieldKind, Level,
+                            DimLevelType)>
+        callback) const {
+#define RETURN_ON_FALSE(fidx, kind, dim, dlt)                                  \
+  if (!(callback(fidx, kind, dim, dlt)))                                       \
+    return;
+
+  const auto lvlTypes = enc.getDimLevelType();
+  const Level lvlRank = enc.getLvlRank();
+  const Level cooStart = getCOOStart(enc);
+  const Level end = cooStart == lvlRank ? cooStart : cooStart + 1;
+  FieldIndex fieldIdx = kDataFieldStartingIdx;
+  // Per-dimension storage.
+  for (Level l = 0; l < end; l++) {
+    // Dimension level types apply in order to the reordered dimension.
+    // As a result, the compound type can be constructed directly in the given
+    // order.
+    const auto dlt = lvlTypes[l];
+    if (isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt)) {
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::PosMemRef, l, dlt);
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::CrdMemRef, l, dlt);
+    } else if (isSingletonDLT(dlt)) {
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::CrdMemRef, l, dlt);
+    } else {
+      assert(isDenseDLT(dlt)); // no fields
+    }
+  }
+  // The values array.
+  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::ValMemRef, -1u,
+                  DimLevelType::Undef);
+
+  // Put metadata at the end.
+  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::StorageSpec, -1u,
+                  DimLevelType::Undef);
+
+#undef RETURN_ON_FALSE
+}
+
+void sparse_tensor::foreachFieldAndTypeInSparseTensor(
+    SparseTensorType stt,
+    llvm::function_ref<bool(Type, FieldIndex, SparseTensorFieldKind, Level,
+                            DimLevelType)>
+        callback) {
+  assert(stt.hasEncoding());
+  // Construct the basic types.
+  const Type crdType = stt.getCrdType();
+  const Type posType = stt.getPosType();
+  const Type eltType = stt.getElementType();
+
+  const Type metaDataType = StorageSpecifierType::get(stt.getEncoding());
+  // memref<? x pos>  positions
+  const Type posMemType = MemRefType::get({ShapedType::kDynamic}, posType);
+  // memref<? x crd>  coordinates
+  const Type crdMemType = MemRefType::get({ShapedType::kDynamic}, crdType);
+  // memref<? x eltType> values
+  const Type valMemType = MemRefType::get({ShapedType::kDynamic}, eltType);
+
+  StorageLayout(stt).foreachField(
+      [metaDataType, posMemType, crdMemType, valMemType,
+       callback](FieldIndex fieldIdx, SparseTensorFieldKind fieldKind,
+                 Level lvl, DimLevelType dlt) -> bool {
+        switch (fieldKind) {
+        case SparseTensorFieldKind::StorageSpec:
+          return callback(metaDataType, fieldIdx, fieldKind, lvl, dlt);
+        case SparseTensorFieldKind::PosMemRef:
+          return callback(posMemType, fieldIdx, fieldKind, lvl, dlt);
+        case SparseTensorFieldKind::CrdMemRef:
+          return callback(crdMemType, fieldIdx, fieldKind, lvl, dlt);
+        case SparseTensorFieldKind::ValMemRef:
+          return callback(valMemType, fieldIdx, fieldKind, lvl, dlt);
+        };
+        llvm_unreachable("unrecognized field kind");
+      });
+}
+
+unsigned StorageLayout::getNumFields() const {
+  unsigned numFields = 0;
+  foreachField([&numFields](FieldIndex, SparseTensorFieldKind, Level,
+                            DimLevelType) -> bool {
+    numFields++;
+    return true;
+  });
+  return numFields;
+}
+
+unsigned StorageLayout::getNumDataFields() const {
+  unsigned numFields = 0; // one value memref
+  foreachField([&numFields](FieldIndex fidx, SparseTensorFieldKind, Level,
+                            DimLevelType) -> bool {
+    if (fidx >= kDataFieldStartingIdx)
+      numFields++;
+    return true;
+  });
+  numFields -= 1; // the last field is MetaData field
+  assert(numFields == getNumFields() - kDataFieldStartingIdx - 1);
+  return numFields;
+}
+
+//===----------------------------------------------------------------------===//
 // TensorDialect Attribute Methods.
 //===----------------------------------------------------------------------===//
 
 static bool acceptBitWidth(unsigned bitWidth) {
   switch (bitWidth) {
   case 0:
   case 8:
   case 16:
   case 32:

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

 add_mlir_dialect_library(MLIRSparseTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   CodegenEnv.cpp
   CodegenUtils.cpp
   LoopEmitter.cpp
   SparseBufferRewriting.cpp
   SparseGPUCodegen.cpp
   SparseStorageSpecifierToLLVM.cpp
   SparseTensorCodegen.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
   SparseTensorRewriting.cpp
-  SparseTensorStorageLayout.cpp
+  SparseTensorDescriptor.cpp
   SparseVectorization.cpp
   Sparsification.cpp
   SparsificationAndBufferizationPass.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS

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

 //===- SparseStorageSpecifierToLLVM.cpp - convert specifier to llvm -------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
-#include "SparseTensorStorageLayout.h"
 
+#include "mlir/Conversion/LLVMCommon/StructBuilder.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+
 #include <optional>
 
 using namespace mlir;
 using namespace sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
matchAndRewrite(SourceOp op, OpAdaptor adaptor,
     case StorageSpecifierKind::CrdMemSize:
     case StorageSpecifierKind::PosMemSize:
     case StorageSpecifierKind::ValMemSize: {
       auto enc = op.getSpecifier().getType().getEncoding();
       StorageLayout layout(enc);
       std::optional<unsigned> lvl;
       if (op.getLevel())
         lvl = (*op.getLevel());
-      unsigned idx = layout.getMemRefFieldIndex(op.getSpecifierKind(), lvl);
+      unsigned idx =
+          layout.getMemRefFieldIndex(toFieldKind(op.getSpecifierKind()), lvl);
       Value v = Base::onMemSize(rewriter, op, spec, idx);
       rewriter.replaceOp(op, v);
       return success();
     }
     }
     llvm_unreachable("unrecognized specifer kind");
   }
 };

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorDescriptor.h

 //===- SparseTensorStorageLayout.h ------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines utilities for the sparse memory layout.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_
 
-#include "mlir/Conversion/LLVMCommon/StructBuilder.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
 //===----------------------------------------------------------------------===//
 // SparseTensorDescriptor and helpers that manage the sparse tensor memory
 // layout scheme during "direct code generation" (i.e. when sparsification
 // generates the buffers as part of actual IR, in constrast with the library
 // approach where data structures are hidden behind opaque pointers).
-//
-// The sparse tensor storage scheme for a rank-dimensional tensor is organized
-// as a single compound type with the following fields. Note that every memref
-// with ? size actually behaves as a "vector", i.e. the stored size is the
-// capacity and the used size resides in the storage_specifier struct.
-//
-// struct {
-//   ; per-level l:
-//   ;  if dense:
-//        <nothing>
-//   ;  if compresed:
-//        memref<? x pos>  positions-l   ; positions for sparse level l
-//        memref<? x crd>  coordinates-l ; coordinates for sparse level l
-//   ;  if singleton:
-//        memref<? x crd>  coordinates-l ; coordinates for singleton level l
-//
-//   memref<? x eltType> values        ; values
-//
-//   struct sparse_tensor.storage_specifier {
-//     array<rank x int> lvlSizes    ; sizes/cardinalities for each level
-//     array<n x int> memSizes;      ; sizes/lengths for each data memref
-//   }
-// };
-//
-// In addition, for a "trailing COO region", defined as a compressed level
-// followed by one or more singleton levels, the default SOA storage that
-// is inherent to the TACO format is optimized into an AOS storage where
-// all coordinates of a stored element appear consecutively.  In such cases,
-// a special operation (sparse_tensor.coordinates_buffer) must be used to
-// access the AOS coordinates array. In the code below, the method `getCOOStart`
-// is used to find the start of the "trailing COO region".
-//
-// If the sparse tensor is a slice (produced by `tensor.extract_slice`
-// operation), instead of allocating a new sparse tensor for it, it reuses the
-// same sets of MemRefs but attaching a additional set of slicing-metadata for
-// per-dimension slice offset and stride.
-//
-// Examples.
-//
-// #CSR storage of 2-dim matrix yields
-//   memref<?xindex>                           ; positions-1
-//   memref<?xindex>                           ; coordinates-1
-//   memref<?xf64>                             ; values
-//   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
-//
-// #COO storage of 2-dim matrix yields
-//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
-//   memref<?xindex>                           ; AOS coordinates storage
-//   memref<?xf64>                             ; values
-//   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
-//
-// Slice on #COO storage of 2-dim matrix yields
-//   ;; Inherited from the original sparse tensors
-//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
-//   memref<?xindex>                           ; AOS coordinates storage
-//   memref<?xf64>                             ; values
-//   struct<(array<2 x i64>, array<3 x i64>,   ; lvl0, lvl1, 3xsizes
-//   ;; Extra slicing-metadata
-//           array<2 x i64>, array<2 x i64>)>) ; dim offset, dim stride.
-//
 //===----------------------------------------------------------------------===//
 
-enum class SparseTensorFieldKind : uint32_t {
-  StorageSpec = 0,
-  PosMemRef = 1,
-  CrdMemRef = 2,
-  ValMemRef = 3
-};
-
-static_assert(static_cast<uint32_t>(SparseTensorFieldKind::PosMemRef) ==
-              static_cast<uint32_t>(StorageSpecifierKind::PosMemSize));
-static_assert(static_cast<uint32_t>(SparseTensorFieldKind::CrdMemRef) ==
-              static_cast<uint32_t>(StorageSpecifierKind::CrdMemSize));
-static_assert(static_cast<uint32_t>(SparseTensorFieldKind::ValMemRef) ==
-              static_cast<uint32_t>(StorageSpecifierKind::ValMemSize));
-
-/// The type of field indices.  This alias is to help code be more
-/// self-documenting; unfortunately it is not type-checked, so it only
-/// provides documentation rather than doing anything to prevent mixups.
-using FieldIndex = unsigned;
-
-// TODO: Functions/methods marked with [NUMFIELDS] might should use
-// `FieldIndex` for their return type, via the same reasoning for why
-// `Dimension`/`Level` are used both for identifiers and ranks.
-
-/// For each field that will be allocated for the given sparse tensor
-/// encoding, calls the callback with the corresponding field index,
-/// field kind, level, and level-type (the last two are only for level
-/// memrefs).  The field index always starts with zero and increments
-/// by one between each callback invocation.  Ideally, all other methods
-/// should rely on this function to query a sparse tensor fields instead
-/// of relying on ad-hoc index computation.
-void foreachFieldInSparseTensor(
-    SparseTensorEncodingAttr,
-    llvm::function_ref<bool(
-        FieldIndex /*fieldIdx*/, SparseTensorFieldKind /*fieldKind*/,
-        Level /*lvl (if applicable)*/, DimLevelType /*DLT (if applicable)*/)>);
-
-/// Same as above, except that it also builds the Type for the corresponding
-/// field.
-void foreachFieldAndTypeInSparseTensor(
-    SparseTensorType,
-    llvm::function_ref<bool(Type /*fieldType*/, FieldIndex /*fieldIdx*/,
-                            SparseTensorFieldKind /*fieldKind*/,
-                            Level /*lvl (if applicable)*/,
-                            DimLevelType /*DLT (if applicable)*/)>);
-
-/// Gets the total number of fields for the given sparse tensor encoding.
-// TODO: See note [NUMFIELDS].
-unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc);
-
-/// Gets the total number of data fields (coordinate arrays, position
-/// arrays, and a value array) for the given sparse tensor encoding.
-// TODO: See note [NUMFIELDS].
-unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc);
-
-inline StorageSpecifierKind toSpecifierKind(SparseTensorFieldKind kind) {
-  assert(kind != SparseTensorFieldKind::StorageSpec);
-  return static_cast<StorageSpecifierKind>(kind);
-}
-
-inline SparseTensorFieldKind toFieldKind(StorageSpecifierKind kind) {
-  assert(kind != StorageSpecifierKind::LvlSize);
-  return static_cast<SparseTensorFieldKind>(kind);
-}
-
-/// Provides methods to access fields of a sparse tensor with the given
-/// encoding.
-class StorageLayout {
-public:
-  explicit StorageLayout(SparseTensorEncodingAttr enc) : enc(enc) {}
-
-  ///
-  /// Getters: get the field index for required field.
-  ///
-
-  FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind,
-                                 std::optional<Level> lvl) const {
-    return getFieldIndexAndStride(kind, lvl).first;
-  }
-
-  FieldIndex getMemRefFieldIndex(StorageSpecifierKind kind,
-                                 std::optional<Level> lvl) const {
-    return getMemRefFieldIndex(toFieldKind(kind), lvl);
-  }
-
-  // TODO: See note [NUMFIELDS].
-  static unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-    return sparse_tensor::getNumFieldsFromEncoding(enc);
-  }
-
-  static void foreachFieldInSparseTensor(
-      const SparseTensorEncodingAttr enc,
-      llvm::function_ref<bool(FieldIndex, SparseTensorFieldKind, Level,
-                              DimLevelType)>
-          callback) {
-    return sparse_tensor::foreachFieldInSparseTensor(enc, callback);
-  }
-
-  std::pair<FieldIndex, unsigned>
-  getFieldIndexAndStride(SparseTensorFieldKind kind,
-                         std::optional<Level> lvl) const {
-    FieldIndex fieldIdx = -1u;
-    unsigned stride = 1;
-    if (kind == SparseTensorFieldKind::CrdMemRef) {
-      assert(lvl.has_value());
-      const Level cooStart = getCOOStart(enc);
-      const Level lvlRank = enc.getLvlRank();
-      if (lvl.value() >= cooStart && lvl.value() < lvlRank) {
-        lvl = cooStart;
-        stride = lvlRank - cooStart;
-      }
-    }
-    foreachFieldInSparseTensor(
-        enc,
-        [lvl, kind, &fieldIdx](FieldIndex fIdx, SparseTensorFieldKind fKind,
-                               Level fLvl, DimLevelType dlt) -> bool {
-          if ((lvl && fLvl == lvl.value() && kind == fKind) ||
-              (kind == fKind && fKind == SparseTensorFieldKind::ValMemRef)) {
-            fieldIdx = fIdx;
-            // Returns false to break the iteration.
-            return false;
-          }
-          return true;
-        });
-    assert(fieldIdx != -1u);
-    return std::pair<FieldIndex, unsigned>(fieldIdx, stride);
-  }
-
-private:
-  SparseTensorEncodingAttr enc;
-};
-
 class SparseTensorSpecifier {
 public:
   explicit SparseTensorSpecifier(Value specifier)
       : specifier(cast<TypedValue<StorageSpecifierType>>(specifier)) {}
 
   // Undef value for level-sizes, all zero values for memory-sizes.
   static Value getInitValue(OpBuilder &builder, Location loc,
                             SparseTensorType stt);

 /// A helper class around an array of values that corresponds to a sparse
 /// tensor. This class provides a set of meaningful APIs to query and update
 /// a particular field in a consistent way. Users should not make assumptions
 /// on how a sparse tensor is laid out but instead rely on this class to access
 /// the right value for the right field.
 template <typename ValueArrayRef>
 class SparseTensorDescriptorImpl {
 protected:
+  // TODO: Functions/methods marked with [NUMFIELDS] might should use
+  // `FieldIndex` for their return type, via the same reasoning for why
+  // `Dimension`/`Level` are used both for identifiers and ranks.
   SparseTensorDescriptorImpl(SparseTensorType stt, ValueArrayRef fields)
-      : rType(stt), fields(fields) {
-    assert(stt.hasEncoding() &&
-           getNumFieldsFromEncoding(stt.getEncoding()) == getNumFields());
+      : rType(stt), fields(fields), layout(stt.getEncoding()) {
+    assert(stt.hasEncoding() && layout.getNumFields() == getNumFields());
     // We should make sure the class is trivially copyable (and should be small
     // enough) such that we can pass it by value.
     static_assert(std::is_trivially_copyable_v<
                   SparseTensorDescriptorImpl<ValueArrayRef>>);
   }
 
 public:
   FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind,
                                  std::optional<Level> lvl) const {
     // Delegates to storage layout.
-    StorageLayout layout(rType.getEncoding());
     return layout.getMemRefFieldIndex(kind, lvl);
   }
 
   // TODO: See note [NUMFIELDS].
   unsigned getNumFields() const { return fields.size(); }
 
   ///
   /// Getters: get the value for required field.
public:
   }
 
   ValueRange getMemRefFields() const {
     // Drop the last metadata fields.
     return fields.drop_back();
   }
 
   std::pair<FieldIndex, unsigned> getCrdMemRefIndexAndStride(Level lvl) const {
-    StorageLayout layout(rType.getEncoding());
     return layout.getFieldIndexAndStride(SparseTensorFieldKind::CrdMemRef, lvl);
   }
 
   Value getAOSMemRef() const {
     const Level cooStart = getCOOStart(rType.getEncoding());
     assert(cooStart < rType.getLvlRank());
     return getMemRefField(SparseTensorFieldKind::CrdMemRef, cooStart);
   }
 
   RankedTensorType getRankedTensorType() const { return rType; }
   ValueArrayRef getFields() const { return fields; }
 
 protected:
   SparseTensorType rType;
   ValueArrayRef fields;
+  StorageLayout layout;
 };
 
 /// Uses ValueRange for immutable descriptors.
 class SparseTensorDescriptor : public SparseTensorDescriptorImpl<ValueRange> {
 public:
   SparseTensorDescriptor(SparseTensorType stt, ValueRange buffers)
       : SparseTensorDescriptorImpl<ValueRange>(stt, buffers) {}
 

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

 //===- SparseTensorStorageLayout.cpp --------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
-#include "SparseTensorStorageLayout.h"
+#include "SparseTensorDescriptor.h"
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
Value sparse_tensor::SparseTensorDescriptor::getCrdMemRefOrView(
   Value size = getCrdMemSize(builder, loc, cooStart);
   size = builder.create<arith::DivUIOp>(loc, size, stride);
   return builder.create<memref::SubViewOp>(
       loc, getMemRefField(SparseTensorFieldKind::CrdMemRef, cooStart),
       /*offset=*/ValueRange{constantIndex(builder, loc, lvl - cooStart)},
       /*size=*/ValueRange{size},
       /*step=*/ValueRange{stride});
 }
-
-//===----------------------------------------------------------------------===//
-// Public methods.
-//===----------------------------------------------------------------------===//
-
-constexpr FieldIndex kDataFieldStartingIdx = 0;
-
-void sparse_tensor::foreachFieldInSparseTensor(
-    const SparseTensorEncodingAttr enc,
-    llvm::function_ref<bool(FieldIndex, SparseTensorFieldKind, Level,
-                            DimLevelType)>
-        callback) {
-  assert(enc);
-
-#define RETURN_ON_FALSE(fidx, kind, dim, dlt)                                  \
-  if (!(callback(fidx, kind, dim, dlt)))                                       \
-    return;
-
-  const auto lvlTypes = enc.getDimLevelType();
-  const Level lvlRank = enc.getLvlRank();
-  const Level cooStart = getCOOStart(enc);
-  const Level end = cooStart == lvlRank ? cooStart : cooStart + 1;
-  FieldIndex fieldIdx = kDataFieldStartingIdx;
-  // Per-dimension storage.
-  for (Level l = 0; l < end; l++) {
-    // Dimension level types apply in order to the reordered dimension.
-    // As a result, the compound type can be constructed directly in the given
-    // order.
-    const auto dlt = lvlTypes[l];
-    if (isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt)) {
-      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::PosMemRef, l, dlt);
-      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::CrdMemRef, l, dlt);
-    } else if (isSingletonDLT(dlt)) {
-      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::CrdMemRef, l, dlt);
-    } else {
-      assert(isDenseDLT(dlt)); // no fields
-    }
-  }
-  // The values array.
-  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::ValMemRef, -1u,
-                  DimLevelType::Undef);
-
-  // Put metadata at the end.
-  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::StorageSpec, -1u,
-                  DimLevelType::Undef);
-
-#undef RETURN_ON_FALSE
-}
-
-void sparse_tensor::foreachFieldAndTypeInSparseTensor(
-    SparseTensorType stt,
-    llvm::function_ref<bool(Type, FieldIndex, SparseTensorFieldKind, Level,
-                            DimLevelType)>
-        callback) {
-  assert(stt.hasEncoding());
-  // Construct the basic types.
-  const Type crdType = stt.getCrdType();
-  const Type posType = stt.getPosType();
-  const Type eltType = stt.getElementType();
-
-  const Type metaDataType = StorageSpecifierType::get(stt.getEncoding());
-  // memref<? x pos>  positions
-  const Type posMemType = MemRefType::get({ShapedType::kDynamic}, posType);
-  // memref<? x crd>  coordinates
-  const Type crdMemType = MemRefType::get({ShapedType::kDynamic}, crdType);
-  // memref<? x eltType> values
-  const Type valMemType = MemRefType::get({ShapedType::kDynamic}, eltType);
-
-  foreachFieldInSparseTensor(
-      stt.getEncoding(),
-      [metaDataType, posMemType, crdMemType, valMemType,
-       callback](FieldIndex fieldIdx, SparseTensorFieldKind fieldKind,
-                 Level lvl, DimLevelType dlt) -> bool {
-        switch (fieldKind) {
-        case SparseTensorFieldKind::StorageSpec:
-          return callback(metaDataType, fieldIdx, fieldKind, lvl, dlt);
-        case SparseTensorFieldKind::PosMemRef:
-          return callback(posMemType, fieldIdx, fieldKind, lvl, dlt);
-        case SparseTensorFieldKind::CrdMemRef:
-          return callback(crdMemType, fieldIdx, fieldKind, lvl, dlt);
-        case SparseTensorFieldKind::ValMemRef:
-          return callback(valMemType, fieldIdx, fieldKind, lvl, dlt);
-        };
-        llvm_unreachable("unrecognized field kind");
-      });
-}
-
-unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-  unsigned numFields = 0;
-  foreachFieldInSparseTensor(enc,
-                             [&numFields](FieldIndex, SparseTensorFieldKind,
-                                          Level, DimLevelType) -> bool {
-                               numFields++;
-                               return true;
-                             });
-  return numFields;
-}
-
-unsigned
-sparse_tensor::getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-  unsigned numFields = 0; // one value memref
-  foreachFieldInSparseTensor(enc,
-                             [&numFields](FieldIndex fidx,
-                                          SparseTensorFieldKind, Level,
-                                          DimLevelType) -> bool {
-                               if (fidx >= kDataFieldStartingIdx)
-                                 numFields++;
-                               return true;
-                             });
-  numFields -= 1; // the last field is MetaData field
-  assert(numFields ==
-         getNumFieldsFromEncoding(enc) - kDataFieldStartingIdx - 1);
-  return numFields;
-}

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

 //===- SparseTensorStorageLayout.h ------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines utilities for the sparse memory layout.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_
 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_
 
-#include "mlir/Conversion/LLVMCommon/StructBuilder.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
 //===----------------------------------------------------------------------===//
 // SparseTensorDescriptor and helpers that manage the sparse tensor memory
 // layout scheme during "direct code generation" (i.e. when sparsification
 // generates the buffers as part of actual IR, in constrast with the library
 // approach where data structures are hidden behind opaque pointers).
-//
-// The sparse tensor storage scheme for a rank-dimensional tensor is organized
-// as a single compound type with the following fields. Note that every memref
-// with ? size actually behaves as a "vector", i.e. the stored size is the
-// capacity and the used size resides in the storage_specifier struct.
-//
-// struct {
-//   ; per-level l:
-//   ;  if dense:
-//        <nothing>
-//   ;  if compresed:
-//        memref<? x pos>  positions-l   ; positions for sparse level l
-//        memref<? x crd>  coordinates-l ; coordinates for sparse level l
-//   ;  if singleton:
-//        memref<? x crd>  coordinates-l ; coordinates for singleton level l
-//
-//   memref<? x eltType> values        ; values
-//
-//   struct sparse_tensor.storage_specifier {
-//     array<rank x int> lvlSizes    ; sizes/cardinalities for each level
-//     array<n x int> memSizes;      ; sizes/lengths for each data memref
-//   }
-// };
-//
-// In addition, for a "trailing COO region", defined as a compressed level
-// followed by one or more singleton levels, the default SOA storage that
-// is inherent to the TACO format is optimized into an AOS storage where
-// all coordinates of a stored element appear consecutively.  In such cases,
-// a special operation (sparse_tensor.coordinates_buffer) must be used to
-// access the AOS coordinates array. In the code below, the method `getCOOStart`
-// is used to find the start of the "trailing COO region".
-//
-// If the sparse tensor is a slice (produced by `tensor.extract_slice`
-// operation), instead of allocating a new sparse tensor for it, it reuses the
-// same sets of MemRefs but attaching a additional set of slicing-metadata for
-// per-dimension slice offset and stride.
-//
-// Examples.
-//
-// #CSR storage of 2-dim matrix yields
-//   memref<?xindex>                           ; positions-1
-//   memref<?xindex>                           ; coordinates-1
-//   memref<?xf64>                             ; values
-//   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
-//
-// #COO storage of 2-dim matrix yields
-//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
-//   memref<?xindex>                           ; AOS coordinates storage
-//   memref<?xf64>                             ; values
-//   struct<(array<2 x i64>, array<3 x i64>)>) ; lvl0, lvl1, 3xsizes
-//
-// Slice on #COO storage of 2-dim matrix yields
-//   ;; Inherited from the original sparse tensors
-//   memref<?xindex>,                          ; positions-0, essentially [0,sz]
-//   memref<?xindex>                           ; AOS coordinates storage
-//   memref<?xf64>                             ; values
-//   struct<(array<2 x i64>, array<3 x i64>,   ; lvl0, lvl1, 3xsizes
-//   ;; Extra slicing-metadata
-//           array<2 x i64>, array<2 x i64>)>) ; dim offset, dim stride.
-//
 //===----------------------------------------------------------------------===//
 
-enum class SparseTensorFieldKind : uint32_t {
-  StorageSpec = 0,
-  PosMemRef = 1,
-  CrdMemRef = 2,
-  ValMemRef = 3
-};
-
-static_assert(static_cast<uint32_t>(SparseTensorFieldKind::PosMemRef) ==
-              static_cast<uint32_t>(StorageSpecifierKind::PosMemSize));
-static_assert(static_cast<uint32_t>(SparseTensorFieldKind::CrdMemRef) ==
-              static_cast<uint32_t>(StorageSpecifierKind::CrdMemSize));
-static_assert(static_cast<uint32_t>(SparseTensorFieldKind::ValMemRef) ==
-              static_cast<uint32_t>(StorageSpecifierKind::ValMemSize));
-
-/// The type of field indices.  This alias is to help code be more
-/// self-documenting; unfortunately it is not type-checked, so it only
-/// provides documentation rather than doing anything to prevent mixups.
-using FieldIndex = unsigned;
-
-// TODO: Functions/methods marked with [NUMFIELDS] might should use
-// `FieldIndex` for their return type, via the same reasoning for why
-// `Dimension`/`Level` are used both for identifiers and ranks.
-
-/// For each field that will be allocated for the given sparse tensor
-/// encoding, calls the callback with the corresponding field index,
-/// field kind, level, and level-type (the last two are only for level
-/// memrefs).  The field index always starts with zero and increments
-/// by one between each callback invocation.  Ideally, all other methods
-/// should rely on this function to query a sparse tensor fields instead
-/// of relying on ad-hoc index computation.
-void foreachFieldInSparseTensor(
-    SparseTensorEncodingAttr,
-    llvm::function_ref<bool(
-        FieldIndex /*fieldIdx*/, SparseTensorFieldKind /*fieldKind*/,
-        Level /*lvl (if applicable)*/, DimLevelType /*DLT (if applicable)*/)>);
-
-/// Same as above, except that it also builds the Type for the corresponding
-/// field.
-void foreachFieldAndTypeInSparseTensor(
-    SparseTensorType,
-    llvm::function_ref<bool(Type /*fieldType*/, FieldIndex /*fieldIdx*/,
-                            SparseTensorFieldKind /*fieldKind*/,
-                            Level /*lvl (if applicable)*/,
-                            DimLevelType /*DLT (if applicable)*/)>);
-
-/// Gets the total number of fields for the given sparse tensor encoding.
-// TODO: See note [NUMFIELDS].
-unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc);
-
-/// Gets the total number of data fields (coordinate arrays, position
-/// arrays, and a value array) for the given sparse tensor encoding.
-// TODO: See note [NUMFIELDS].
-unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc);
-
-inline StorageSpecifierKind toSpecifierKind(SparseTensorFieldKind kind) {
-  assert(kind != SparseTensorFieldKind::StorageSpec);
-  return static_cast<StorageSpecifierKind>(kind);
-}
-
-inline SparseTensorFieldKind toFieldKind(StorageSpecifierKind kind) {
-  assert(kind != StorageSpecifierKind::LvlSize);
-  return static_cast<SparseTensorFieldKind>(kind);
-}
-
-/// Provides methods to access fields of a sparse tensor with the given
-/// encoding.
-class StorageLayout {
-public:
-  explicit StorageLayout(SparseTensorEncodingAttr enc) : enc(enc) {}
-
-  ///
-  /// Getters: get the field index for required field.
-  ///
-
-  FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind,
-                                 std::optional<Level> lvl) const {
-    return getFieldIndexAndStride(kind, lvl).first;
-  }
-
-  FieldIndex getMemRefFieldIndex(StorageSpecifierKind kind,
-                                 std::optional<Level> lvl) const {
-    return getMemRefFieldIndex(toFieldKind(kind), lvl);
-  }
-
-  // TODO: See note [NUMFIELDS].
-  static unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-    return sparse_tensor::getNumFieldsFromEncoding(enc);
-  }
-
-  static void foreachFieldInSparseTensor(
-      const SparseTensorEncodingAttr enc,
-      llvm::function_ref<bool(FieldIndex, SparseTensorFieldKind, Level,
-                              DimLevelType)>
-          callback) {
-    return sparse_tensor::foreachFieldInSparseTensor(enc, callback);
-  }
-
-  std::pair<FieldIndex, unsigned>
-  getFieldIndexAndStride(SparseTensorFieldKind kind,
-                         std::optional<Level> lvl) const {
-    FieldIndex fieldIdx = -1u;
-    unsigned stride = 1;
-    if (kind == SparseTensorFieldKind::CrdMemRef) {
-      assert(lvl.has_value());
-      const Level cooStart = getCOOStart(enc);
-      const Level lvlRank = enc.getLvlRank();
-      if (lvl.value() >= cooStart && lvl.value() < lvlRank) {
-        lvl = cooStart;
-        stride = lvlRank - cooStart;
-      }
-    }
-    foreachFieldInSparseTensor(
-        enc,
-        [lvl, kind, &fieldIdx](FieldIndex fIdx, SparseTensorFieldKind fKind,
-                               Level fLvl, DimLevelType dlt) -> bool {
-          if ((lvl && fLvl == lvl.value() && kind == fKind) ||
-              (kind == fKind && fKind == SparseTensorFieldKind::ValMemRef)) {
-            fieldIdx = fIdx;
-            // Returns false to break the iteration.
-            return false;
-          }
-          return true;
-        });
-    assert(fieldIdx != -1u);
-    return std::pair<FieldIndex, unsigned>(fieldIdx, stride);
-  }
-
-private:
-  SparseTensorEncodingAttr enc;
-};
-
 class SparseTensorSpecifier {
 public:
   explicit SparseTensorSpecifier(Value specifier)
       : specifier(cast<TypedValue<StorageSpecifierType>>(specifier)) {}
 
   // Undef value for level-sizes, all zero values for memory-sizes.
   static Value getInitValue(OpBuilder &builder, Location loc,
                             SparseTensorType stt);

 /// A helper class around an array of values that corresponds to a sparse
 /// tensor. This class provides a set of meaningful APIs to query and update
 /// a particular field in a consistent way. Users should not make assumptions
 /// on how a sparse tensor is laid out but instead rely on this class to access
 /// the right value for the right field.
 template <typename ValueArrayRef>
 class SparseTensorDescriptorImpl {
 protected:
+  // TODO: Functions/methods marked with [NUMFIELDS] might should use
+  // `FieldIndex` for their return type, via the same reasoning for why
+  // `Dimension`/`Level` are used both for identifiers and ranks.
   SparseTensorDescriptorImpl(SparseTensorType stt, ValueArrayRef fields)
-      : rType(stt), fields(fields) {
-    assert(stt.hasEncoding() &&
-           getNumFieldsFromEncoding(stt.getEncoding()) == getNumFields());
+      : rType(stt), fields(fields), layout(stt.getEncoding()) {
+    assert(stt.hasEncoding() && layout.getNumFields() == getNumFields());
     // We should make sure the class is trivially copyable (and should be small
     // enough) such that we can pass it by value.
     static_assert(std::is_trivially_copyable_v<
                   SparseTensorDescriptorImpl<ValueArrayRef>>);
   }
 
 public:
   FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind,
                                  std::optional<Level> lvl) const {
     // Delegates to storage layout.
-    StorageLayout layout(rType.getEncoding());
     return layout.getMemRefFieldIndex(kind, lvl);
   }
 
   // TODO: See note [NUMFIELDS].
   unsigned getNumFields() const { return fields.size(); }
 
   ///
   /// Getters: get the value for required field.
public:
   }
 
   ValueRange getMemRefFields() const {
     // Drop the last metadata fields.
     return fields.drop_back();
   }
 
   std::pair<FieldIndex, unsigned> getCrdMemRefIndexAndStride(Level lvl) const {
-    StorageLayout layout(rType.getEncoding());
     return layout.getFieldIndexAndStride(SparseTensorFieldKind::CrdMemRef, lvl);
   }
 
   Value getAOSMemRef() const {
     const Level cooStart = getCOOStart(rType.getEncoding());
     assert(cooStart < rType.getLvlRank());
     return getMemRefField(SparseTensorFieldKind::CrdMemRef, cooStart);
   }
 
   RankedTensorType getRankedTensorType() const { return rType; }
   ValueArrayRef getFields() const { return fields; }
 
 protected:
   SparseTensorType rType;
   ValueArrayRef fields;
+  StorageLayout layout;
 };
 
 /// Uses ValueRange for immutable descriptors.
 class SparseTensorDescriptor : public SparseTensorDescriptorImpl<ValueRange> {
 public:
   SparseTensorDescriptor(SparseTensorType stt, ValueRange buffers)
       : SparseTensorDescriptorImpl<ValueRange>(stt, buffers) {}
 

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

cc_library(
     includes = ["include"],
 )
 
 cc_library(
     name = "SparseTensorDialect",
     srcs = ["lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp"],
     hdrs = [
         "include/mlir/Dialect/SparseTensor/IR/SparseTensor.h",
+        "include/mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h",
         "include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h",
     ],
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":SparseTensorAttrDefsIncGen",