[mlir][sparse] Refactoring: abstract sparse tensor memory scheme into a SparseTensorDescriptor class.

[Peiming]

Summary

This patch abstracts sparse tensor memory scheme into a SparseTensorDescriptor class. Previously, the field accesses are performed in a relatively error-prone way, this patch hides the hairy details behind a SparseTensorDescriptor class to allow users access sparse tensor fields in a more cohesive way.

[Peiming]

minor improvement.

[Peiming]

minor fixes.

[Peiming]

code cleanup.

[Peiming]

replaces more places using SparseTensorDescriptor.

[Peiming]

add some comments.

[Peiming]

replace all SmallVectorImpl to SparseTensorDescriptor

[Peiming]

minor fixes.

[aartbik]

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

99

This description describes the layout and thus drives all the code.
Can you please move this into the header, inside the class documentation.

[Peiming]

add some comments.

[Peiming]

cleanup.

[Peiming]

cleanup.

[Peiming]

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

168

A better way to create PushBackOp can be PushBackOp(builder, loc, descriptor, fidx, value), which is more clear and can be read as "push a value into the sparse tensor descriptor at the given field index). But it would require us to put SparseTensorDescriptor to a publicly available place (and I am not sure whether it is wanted).

[Peiming]

rebase

[Peiming]

cleanup.

[Peiming]

code cleanup.

[Peiming]

rebase

[aartbik]

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

571

I think this line should not be here? Copied from somewhere else?

592

relying *on* ad-hoc

612

, and a
(Oxford comma ;-)

622

an array

625

layouted -> laid out? (I think that is the proper english, not sure though...)

647

to restore (no d)

672

it feels like we have many more getters than we really need?
am I right, is there a way to condense the API a bit?

676

Although this solution is much better at information hiding than the original, we now need to scan each time for every field query. Although this is never very deep (rank bounded), it is still a bit wasteful. Can we avoid this by precomputing offsets per dimenison?

683

this feels very convoluted? why do we need a dim for a value query?

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

252

note that this works for now, but we actually planned to implement a heuristic here, which will need to scan the ranks + level types again

274

continue

387

yeah, agreed, this made more sense in the original

but the alternative is to copy and create a fully new array....

[Peiming]

remove unused APIs + fix typos.

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

676

I also think about it, we could, but storing the map will make the class expensive to copy.

How about I do it in next revision and change all the descriptor to reference?

683

deleted.

[Peiming]

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

252

Then, you can have another foreach before this to compute the heuristic.

[Peiming]

rebase.

[PeteSteinfeld]

@Peiming, this patch breaks the flang build.

[vzakhari]

In D138627#3972470, @PeteSteinfeld wrote:

@Peiming, this patch breaks the flang build.

It is not flang specific issue, though. The following gcc buildbot fails: https://lab.llvm.org/buildbot/#/builders/160/builds/13724

[stella.stamenova]

In D138627#3972495, @vzakhari wrote:

In D138627#3972470, @PeteSteinfeld wrote:

@Peiming, this patch breaks the flang build.

It is not flang specific issue, though. The following gcc buildbot fails: https://lab.llvm.org/buildbot/#/builders/160/builds/13724

This also broke the windows mlir buildbot. Please address the issues or revert.

[Peiming]

In D138627#3972566, @stella.stamenova wrote:

In D138627#3972495, @vzakhari wrote:

In D138627#3972470, @PeteSteinfeld wrote:

@Peiming, this patch breaks the flang build.

It is not flang specific issue, though. The following gcc buildbot fails: https://lab.llvm.org/buildbot/#/builders/160/builds/13724

This also broke the windows mlir buildbot. Please address the issues or revert.

Okay, Investigating

[Peiming]

The failure seems unrelated to this patch...

[Peiming]

In D138627#3972566, @stella.stamenova wrote:

In D138627#3972495, @vzakhari wrote:

In D138627#3972470, @PeteSteinfeld wrote:

@Peiming, this patch breaks the flang build.

It is not flang specific issue, though. The following gcc buildbot fails: https://lab.llvm.org/buildbot/#/builders/160/builds/13724

This also broke the windows mlir buildbot. Please address the issues or revert.

The windows failure should be fixed by https://reviews.llvm.org/D139383

[vzakhari]

In D138627#3972777, @Peiming wrote:

The failure seems unrelated to this patch...

The build log seems to directly point to your code (https://lab.llvm.org/buildbot/#/builders/160/builds/13724/steps/5/logs/stdio):

In file included from ../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp:14:
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:396:13: error: explicit specialization in non-namespace scope ‘class mlir::sparse_tensor::SparseTensorDescriptorImpl<mut>’
  396 |   template <>
      |             ^
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:397:10: error: too few template-parameter-lists
  397 |   struct ArrayStorage<false> {
      |          ^~~~~~~~~~~~~~~~~~~
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:401:13: error: explicit specialization in non-namespace scope ‘class mlir::sparse_tensor::SparseTensorDescriptorImpl<mut>’
  401 |   template <>
      |             ^
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:402:10: error: too few template-parameter-lists
  402 |   struct ArrayStorage<true> {
      |          ^~~~~~~~~~~~~~~~~~

[Peiming]

In D138627#3972803, @vzakhari wrote:

In D138627#3972777, @Peiming wrote:

The failure seems unrelated to this patch...

The build log seems to directly point to your code (https://lab.llvm.org/buildbot/#/builders/160/builds/13724/steps/5/logs/stdio):

In file included from ../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp:14:
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:396:13: error: explicit specialization in non-namespace scope ‘class mlir::sparse_tensor::SparseTensorDescriptorImpl<mut>’
  396 |   template <>
      |             ^
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:397:10: error: too few template-parameter-lists
  397 |   struct ArrayStorage<false> {
      |          ^~~~~~~~~~~~~~~~~~~
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:401:13: error: explicit specialization in non-namespace scope ‘class mlir::sparse_tensor::SparseTensorDescriptorImpl<mut>’
  401 |   template <>
      |             ^
../llvm-project/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h:402:10: error: too few template-parameter-lists
  402 |   struct ArrayStorage<true> {
      |          ^~~~~~~~~~~~~~~~~~

Okay, I did not see this. I will try with GCC

[stella.stamenova]

@Peiming : I reverted your change because it broke multiple buildbots - windows and gcc included. I see that you've re-committed it. Did you make sure to address all three build breaks that were reported including the one @vzakhari pointed out? In cases like this, please make sure NOT to recommit your changes without making sure that all build breaks are addressed as it is disruptive to have buildbots that are broken especially when a change is identified as the source.

[Peiming]

In D138627#3972826, @stella.stamenova wrote:

@Peiming : I reverted your change because it broke multiple buildbots - windows and gcc included. I see that you've re-committed it. Did you make sure to address all three build breaks that were reported including the one @vzakhari pointed out? In cases like this, please make sure NOT to recommit your changes without making sure that all build breaks are addressed as it is disruptive to have buildbots that are broken especially when a change is identified as the source.

Yes, of course. I haven't push the code, waiting the pre-merge windows build to finish.

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

inline bool isCompressedDim(RankedTensorType type, uint64_t d) {
   return isCompressedDLT(getDimLevelType(type, d));
 }
 
 /// Convenience function to test for singleton dimension (0 <= d < rank).
 inline bool isSingletonDim(RankedTensorType type, uint64_t d) {
   return isSingletonDLT(getDimLevelType(type, d));
 }
 
+/// Convenience function to test for dense dimension (0 <= d < rank).
+inline bool isDenseDim(SparseTensorEncodingAttr enc, uint64_t d) {
+  return isDenseDLT(getDimLevelType(enc, d));
+}
+
+/// Convenience function to test for compressed dimension (0 <= d < rank).
+inline bool isCompressedDim(SparseTensorEncodingAttr enc, uint64_t d) {
+  return isCompressedDLT(getDimLevelType(enc, d));
+}
+
+/// Convenience function to test for singleton dimension (0 <= d < rank).
+inline bool isSingletonDim(SparseTensorEncodingAttr enc, uint64_t d) {
+  return isSingletonDLT(getDimLevelType(enc, d));
+}
+
 //
 // Dimension level properties.
 //
 
 /// Convenience function to test for ordered property in the
 /// given dimension (0 <= d < rank).
 inline bool isOrderedDim(RankedTensorType type, uint64_t d) {
   return isOrderedDLT(getDimLevelType(type, d));

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

inline Value constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
   return constantI8(builder, loc, static_cast<uint8_t>(dlt));
 }
 
 inline bool isZeroRankedTensorOrScalar(Type type) {
   auto rtp = type.dyn_cast<RankedTensorType>();
   return !rtp || rtp.getRank() == 0;
 }
 
+enum class SparseTensorFieldKind {
+  DimSizes,
+  MemSizes,
+  PtrMemRef,
+  IdxMemRef,
+  ValMemRef
+};
+
+//===----------------------------------------------------------------------===//
+// SparseTensorDescriptor and helpers, manage the sparse tensor memory layout
+// scheme.
+//===----------------------------------------------------------------------===//
+
+constexpr uint64_t dimSizesIdx = 0;
+constexpr uint64_t memSizesIdx = dimSizesIdx + 1;
+constexpr uint64_t dataFieldIdx = memSizesIdx + 1;
+
+/// Foreach fields that will be allocated for the given sparse tensor encoding,
+/// calls the callback with the corresponding field index, field kind, dimension
+/// (for sparse tensor level memrefs) and dimlevelType.
+/// The field index always starts with zero and increments by one between two
+/// callback invocations.
+/// Ideally, all other methods should rely on this function to query a sparse
+/// tensor fields instead of relying ad-hoc index computation to allow better
+/// future extension.
+void foreachFieldInSparseTensor(
+    SparseTensorEncodingAttr,
+    llvm::function_ref<bool(unsigned /*fieldIdx*/,
+                            SparseTensorFieldKind /*fieldKind*/,
+                            unsigned /*dim (if applicable)*/,
+                            DimLevelType /*DLT (if applicable)*/)>);
+
+/// Same as above, except it also build the Type (for now, all MemRefTypes) for
+/// the corresponding field.
+void foreachFieldAndTypeInSparseTensor(
+    RankedTensorType,
+    llvm::function_ref<bool(Type /*fieldType*/, unsigned /*fieldIdx*/,
+                            SparseTensorFieldKind /*fieldKind*/,
+                            unsigned /*dim (if applicable)*/,
+                            DimLevelType /*DLT (if applicable)*/)>);
+
+unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc);
+unsigned getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc);
+
+inline unsigned getFieldMemSizesIndex(unsigned fid) {
+  assert(fid >= dataFieldIdx);
+  return fid - dataFieldIdx;
+}
+
+/// A helper class around a array of values that corresponding to a sparse
+/// tensor, provides a set of meaningful APIs to query and update a particular
+/// field in a consistent way.
+/// Users should not make assumption on how a sparse tensor is layouted but
+/// instead relies on this class to access the right value for the right field.
+template <bool mut>
+class SparseTensorDescriptorImpl {
+private:
+  template <bool>
+  struct ArrayStorage;
+
+  template <>
+  struct ArrayStorage<false> {
+    using ValueArray = ValueRange;
+  };
+
+  template <>
+  struct ArrayStorage<true> {
+    using ValueArray = SmallVectorImpl<Value> &;
+  };
+
+  // Uses ValueRange for immuatable descriptor; uses SmallVectorImpl<Value> &
+  // for mutable descriptor.
+  // Using SmallVector for mutable descriptor allows users to reuse it as a tmp
+  // buffers to append value for some special cases, though users should be
+  // responsible to restored the buffer to legal states after their use. It is
+  // probably not a clean way, but it is the most efficient way to avoid copying
+  // the fields into another SmallVector. If a more clear way is wanted, we
+  // should change it to MutableArrayRef instead.
+  using Storage = typename ArrayStorage<mut>::ValueArray;
+
+public:
+  SparseTensorDescriptorImpl(Type tp, Storage fields)
+      : rType(tp.cast<RankedTensorType>()), fields(fields) {
+    assert(getSparseTensorEncoding(tp) &&
+           getNumFieldsFromEncoding(getSparseTensorEncoding(tp)) ==
+               fields.size());
+    // 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<mut>>);
+  }
+
+  // implicit type conversion
+  template <typename T = SparseTensorDescriptorImpl<true>>
+  /*implicit*/ SparseTensorDescriptorImpl(std::enable_if_t<!mut, T> &mDesc)
+      : rType(mDesc.getTensorType()), fields(mDesc.getFields()) {}
+
+  unsigned getFieldIndex(unsigned dim, SparseTensorFieldKind kind) const {
+    unsigned fieldIdx = -1u;
+    foreachFieldInSparseTensor(
+        getSparseTensorEncoding(rType),
+        [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
+                               unsigned fDim, DimLevelType dlt) -> bool {
+          if (fDim == dim && kind == fKind) {
+            fieldIdx = fIdx;
+            // Returns false to break the iteration.
+            return false;
+          }
+          return true;
+        });
+    assert(fieldIdx != -1u);
+    return fieldIdx;
+  }
+
+  unsigned getPtrMemRefIndex(unsigned ptrDim) const {
+    return getFieldIndex(ptrDim, SparseTensorFieldKind::PtrMemRef);
+  }
+
+  unsigned getIdxMemRefIndex(unsigned idxDim) const {
+    return getFieldIndex(idxDim, SparseTensorFieldKind::IdxMemRef);
+  }
+
+  unsigned getDataFieldIndex(unsigned dim, SparseTensorFieldKind kind) const {
+    if (kind == SparseTensorFieldKind::ValMemRef)
+      return getValMemRefIndex();
+    return getFieldIndex(dim, kind);
+  }
+
+  unsigned getValMemRefIndex() const { return fields.size() - 1; }
+
+  unsigned getPtrMemSizesIndex(unsigned dim) const {
+    return getPtrMemRefIndex(dim) - dataFieldIdx;
+  }
+
+  unsigned getIdxMemSizesIndex(unsigned dim) const {
+    return getIdxMemRefIndex(dim) - dataFieldIdx;
+  }
+
+  unsigned getValMemSizesIndex() const {
+    return getValMemRefIndex() - dataFieldIdx;
+  }
+
+  unsigned getFieldMemSizesIndex(unsigned dim,
+                                 SparseTensorFieldKind kind) const {
+    return getFieldMemSizesIndex(getDataFieldIndex(dim, kind));
+  }
+
+  unsigned getNumFields() const { return fields.size(); }
+
+  Value getDimSizesMemRef() const { return fields[dimSizesIdx]; }
+  Value getMemSizesMemRef() const { return fields[memSizesIdx]; }
+
+  Value getPtrMemRef(unsigned ptrDim) const {
+    return fields[getPtrMemRefIndex(ptrDim)];
+  }
+
+  Value getIdxMemRef(unsigned idxDim) const {
+    return fields[getIdxMemRefIndex(idxDim)];
+  }
+
+  Value getValMemRef() const { return fields[getValMemRefIndex()]; }
+
+  Value getField(unsigned fid) const {
+    assert(fid < fields.size());
+    return fields[fid];
+  }
+
+  template <typename T = Value>
+  void setDimSizesMemRef(std::enable_if_t<mut, T> v) {
+    fields[dimSizesIdx] = v;
+  }
+
+  template <typename T = Value>
+  void setMemSizesMemRef(std::enable_if_t<mut, T> v) {
+    fields[memSizesIdx] = v;
+  }
+
+  template <typename T = Value>
+  void setPtrMemRef(unsigned ptrDim, std::enable_if_t<mut, T> v) {
+    fields[getPtrMemRefIndex(ptrDim)] = v;
+  }
+
+  template <typename T = Value>
+  void setIdxMemRef(unsigned idxDim, std::enable_if_t<mut, T> v) {
+    fields[getIdxMemRefIndex(idxDim)] = v;
+  }
+
+  template <typename T = Value>
+  void setLvlMemRef(unsigned dim, SparseTensorFieldKind kind,
+                    std::enable_if_t<mut, T> v) {
+    fields[getDataFieldIndex(dim, kind)] = v;
+  }
+
+  template <typename T = Value>
+  void setValMemRef(std::enable_if_t<mut, T> v) {
+    fields[getValMemRefIndex()] = v;
+  }
+
+  template <typename T = Value>
+  void setField(unsigned fid, std::enable_if_t<mut, T> v) {
+    assert(fid < fields.size());
+    fields[fid] = v;
+  }
+
+  Type getElementType(unsigned fidx) const {
+    return fields[fidx].getType().template cast<MemRefType>().getElementType();
+  }
+
+  RankedTensorType getTensorType() const { return rType; }
+  Storage getFields() const { return fields; }
+
+  // TODO: a better places for these functions should be in
+  // SparseTensorEncodingAttr.
+  Type getPtrElementType() const {
+    auto *ctx = rType.getContext();
+    unsigned ptrWidth = getSparseTensorEncoding(rType).getPointerBitWidth();
+    Type indexType = IndexType::get(ctx);
+    return ptrWidth ? IntegerType::get(ctx, ptrWidth) : indexType;
+  }
+
+  Type getIdxElementType() const {
+    auto *ctx = rType.getContext();
+    unsigned idxWidth = getSparseTensorEncoding(rType).getIndexBitWidth();
+    Type indexType = IndexType::get(ctx);
+    return idxWidth ? IntegerType::get(ctx, idxWidth) : indexType;
+  }
+
+private:
+  RankedTensorType rType;
+  Storage fields;
+};
+
+using SparseTensorDescriptor = SparseTensorDescriptorImpl<false>;
+using MutSparseTensorDescriptor = SparseTensorDescriptorImpl<true>;
+
 //===----------------------------------------------------------------------===//
-// SparseTensorLoopEmiter class, manages sparse tensors and helps to generate
-// loop structure to (co)-iterate sparse tensors.
+// SparseTensorLoopEmiter class, manages sparse tensors and helps to
+// generate loop structure to (co)-iterate sparse tensors.
 //
 // An example usage:
 // To generate the following loops over T1<?x?> and T2<?x?>
 //
 // for i in TENSOR_1_0 {
 //   for j : TENSOR_2_0 {
 //     for k : TENSOR_1_1 {}
 //     for k : TENSOR_2_1 {}
 //   }
 // }
 //
 // One can use
 //
 // SparseTensorLoopEmiter loopEmiter({T1, T1});
 // loopEmiter.initializeLoopEmit();
 // loopEmiter.enterLoopOverTensorAtDim(T1, 0);
 // loopEmiter.enterLoopOverTensorAtDim(T2, 0);
 // loopEmiter.enterLoopOverTensorAtDim(T1, 1);
 // loopEmiter.exitCurrentLoop();
 // loopEmiter.enterLoopOverTensorAtDim(T2, 1);
 // loopEmiter.exitCurrentLoop(); // exit k
 // loopEmiter.exitCurrentLoop(); // exit j
 // loopEmiter.exitCurrentLoop(); // exit i
 //===----------------------------------------------------------------------===//
 
 class SparseTensorLoopEmitter {
 public:
   /// Optional callback function to setup dense output tensors when
   /// initializing the loop emitter (e.g., to fill a dense output with zeros).
   using OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc,
                                            Value memref, Value tensor)>;
 
-  /// Constructor: take an array of tensors inputs, on which the generated loops
-  /// will iterate on. The index of the tensor in the array is also the
+  /// Constructor: take an array of tensors inputs, on which the generated
+  /// loops will iterate on. The index of the tensor in the array is also the
   /// tensor id (tid) used in related functions.
   /// If isSparseOut is set, loop emitter assume that the sparse output tensor
   /// is empty, and will always generate loops on it based on the dim sizes.
   /// An optional array could be provided (by sparsification) to indicate the
   /// loop id sequence that will be generated. It is used to establish the
-  /// mapping between affineDimExpr to the corresponding loop index in the loop
-  /// stack that are maintained by the loop emitter.
+  /// mapping between affineDimExpr to the corresponding loop index in the
+  /// loop stack that are maintained by the loop emitter.
   explicit SparseTensorLoopEmitter(ValueRange tensors,
                                    StringAttr loopTag = nullptr,
                                    bool hasOutput = false,
                                    bool isSparseOut = false,
                                    ArrayRef<unsigned> topSort = {});
 
   /// Starts a loop emitting session by generating all the buffers needed to
   /// iterate tensors.
   void initializeLoopEmit(OpBuilder &builder, Location loc,
                           OutputUpdater updater = nullptr);
 
   /// Generates a list of operations to compute the affine expression.
   Value genAffine(OpBuilder &builder, AffineExpr a, Location loc);
 
-  /// Enters a new loop sequence, the loops within the same sequence starts from
-  /// the break points of previous loop instead of starting over from 0.
+  /// Enters a new loop sequence, the loops within the same sequence starts
+  /// from the break points of previous loop instead of starting over from 0.
   /// e.g.,
   /// {
   ///   // loop sequence start.
   ///   p0 = while(xxx)
   ///     ...
   ///     break p0
   ///
   ///   // Starts loop from p0
   ///   for (i = p0; i < end; i++)
   ///     ...
   ///   // loop sequence end.
   /// }
   void enterNewLoopSeq(OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
                        ArrayRef<size_t> dims);
 
   // exit the current loop sequence, this will reset universal index to 0.
   void exitCurrentLoopSeq() {
     assert(loopSeqStack.size() == loopStack.size() + 1);
     loopSeqStack.pop_back();
   }
 
   // TODO: Gets rid of `dim` in the argument list? Track the dimension we
   // are currently at internally. Then it would be enterNextDimForTensor.
   // Still need a way to specify the dim for non annoated dense tensor though,
   // as it can be accessed out of order.
   /// Emits loop over tensor_tid_dim, it assumes that loops between
   /// tensor_tid_[0, dim - 1] have already been generated.
   /// The function will also perform in-place update on the `reduc` vector to
   /// return the reduction variable used inside the generated loop.
   Operation *enterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
                                       size_t tid, size_t dim,
                                       MutableArrayRef<Value> reduc = {},
                                       bool isParallel = false,
                                       ArrayRef<size_t> extraTids = {},
                                       ArrayRef<size_t> extraDims = {});
 
   Operation *enterFilterLoopOverTensorAtDim(OpBuilder &builder, Location loc,
                                             size_t tid, size_t dim,
                                             AffineExpr affine,
                                             MutableArrayRef<Value> reduc = {});
 
   void genDenseAffineAddressAtCurLevel(OpBuilder &builder, Location loc,
                                        size_t tid, size_t dim,
                                        AffineExpr affine);
 
   /// Emits a co-iteration loop over a set of tensors.
   Operation *enterCoIterationOverTensorsAtDims(
       OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
       ArrayRef<size_t> dims, bool needsUniv, MutableArrayRef<Value> reduc = {},
       ArrayRef<size_t> extraTids = {}, ArrayRef<size_t> extraDims = {});
 
   void exitCurrentLoop(RewriterBase &rewriter, Location loc,
                        MutableArrayRef<Value> reduc = {});
 
   /// Returns the array of coordinate for all the loop generated till now.
   void getCoordinateArray(SmallVectorImpl<Value> &coords) const {
     for (auto &l : loopStack)
       coords.push_back(l.iv);
   }
 
   /// Gets loop induction variable at the given level.
   unsigned getCurrentDepth() const { return loopStack.size(); }
 
   /// Gets loop induction variable at the given level.
   Value getLoopIV(size_t level) const {
     if (level < loopStack.size())
       return loopStack[level].iv;
     return nullptr;
   }
 
   ///
   /// Getters.
   ///
   const std::vector<std::vector<Value>> &getPidxs() const { return pidxs; };
   const std::vector<std::vector<Value>> &getCoord() const { return coord; };
   const std::vector<std::vector<Value>> &getHighs() const { return highs; };
   const std::vector<std::vector<Value>> &getPtrBuffer() const {
     return ptrBuffer;
   };
   const std::vector<std::vector<Value>> &getIdxBuffer() const {
     return idxBuffer;
   };
   const std::vector<Value> &getValBuffer() const { return valBuffer; };
 
private:
   /// will be transformed into
   /// %ret = parallel () init(%args) {
   ///   ...
   ///   scf.reduce(%c) bb0(%0, %1){
   ///     %val = op %0, %1
   ///     scf.reduce.return %val
   ///   }
   /// }
-  /// NOTE: only one instruction will be moved into reduce block, transformation
-  /// will fail if multiple instructions are used to compute the reduction
-  /// value.
-  /// Return %ret to user, while %val is provided by users (`reduc`).
+  /// NOTE: only one instruction will be moved into reduce block,
+  /// transformation will fail if multiple instructions are used to compute
+  /// the reduction value. Return %ret to user, while %val is provided by
+  /// users (`reduc`).
   void exitForLoop(RewriterBase &rewriter, Location loc,
                    MutableArrayRef<Value> reduc);
 
   /// Exits a while loop, returns the reduction results.
   void exitCoIterationLoop(OpBuilder &builder, Location loc,
                            MutableArrayRef<Value> reduc);
 
-  /// A optional string attribute that should be attached to the loop generated
-  /// by loop emitter, it might help following passes to identify loops that
-  /// operates on sparse tensors more easily.
+  /// A optional string attribute that should be attached to the loop
+  /// generated by loop emitter, it might help following passes to identify
+  /// loops that operates on sparse tensors more easily.
   StringAttr loopTag;
   /// Whether the loop emitter needs to treat the last tensor as the output
   /// tensor.
   bool hasOutput;
   bool isSparseOut;
   /// Input and (optional) output tensors.
   std::vector<Value> tensors;
   /// The dim type array for each tensor.
   std::vector<std::vector<DimLevelType>> dimTypes;
   /// Sparse iteration information (by tensor and dim). These arrays
   /// are updated to remain current within the current loop.
   std::vector<std::vector<Value>> pidxs;
   std::vector<std::vector<Value>> coord;
   std::vector<std::vector<Value>> highs;
   std::vector<std::vector<Value>> ptrBuffer; // to_pointers
   std::vector<std::vector<Value>> idxBuffer; // to_indices
   std::vector<Value> valBuffer;              // to_value
 
-  // Loop Stack, stores the information of all the nested loops that are alive.
+  // Loop Stack, stores the information of all the nested loops that are
+  // alive.
   std::vector<LoopLevelInfo> loopStack;
 
   // Loop Sequence Stack, stores the unversial index for the current loop
   // sequence.
   std::vector<Value> loopSeqStack;
 
   // Maps AffineDimExpr to the index of the loop in loopStack.
   // TODO: We should probably use a callback function here to make it more

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

static Value genIndexAndValueForDense(OpBuilder &builder, Location loc,
                                       Value tensor,
                                       SmallVectorImpl<Value> &indicesArray,
                                       ValueRange ivs) {
   Value val = genValueForDense(builder, loc, tensor, ivs);
   indicesArray.append(ivs.begin(), ivs.end());
   return val;
 }
 
+void sparse_tensor::foreachFieldInSparseTensor(
+    const SparseTensorEncodingAttr enc,
+    llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
+                            DimLevelType)>
+        callback) {
+  //
+  // Sparse tensor storage scheme for 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 memSizes array.
+  //
+  // struct {
+  //   memref<rank x index> dimSizes     ; size in each dimension
+  //   memref<n x index> memSizes        ; sizes of ptrs/inds/values
+  //   ; per-dimension d:
+  //   ;  if dense:
+  //        <nothing>
+  //   ;  if compresed:
+  //        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
+  //        memref<? x idx>  indices-d   ; indices for sparse dim d
+  //   ;  if singleton:
+  //        memref<? x idx>  indices-d   ; indices for singleton dim d
+  //   memref<? x eltType> values        ; values
+  // };
+  //
+  // The dimSizes array and memSizes array.
+  assert(enc);
+
+#define RETURN_ON_FALSE(idx, kind, dim, dlt)                                   \
+  if (!(callback(idx, kind, dim, dlt)))                                        \
+    return;
+
+  RETURN_ON_FALSE(dimSizesIdx, SparseTensorFieldKind::DimSizes, -1u,
+                  DimLevelType::Undef);
+  RETURN_ON_FALSE(memSizesIdx, SparseTensorFieldKind::MemSizes, -1u,
+                  DimLevelType::Undef);
+
+  static_assert(dataFieldIdx == memSizesIdx + 1);
+  unsigned fieldIdx = dataFieldIdx;
+  // Per-dimension storage.
+  for (unsigned r = 0, rank = enc.getDimLevelType().size(); r < rank; r++) {
+    // Dimension level types apply in order to the reordered dimension.
+    // As a result, the compound type can be constructed directly in the given
+    // order.
+    auto dlt = getDimLevelType(enc, r);
+    if (isCompressedDLT(dlt)) {
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::PtrMemRef, r, dlt);
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::IdxMemRef, r, dlt);
+    } else if (isSingletonDLT(dlt)) {
+      RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::IdxMemRef, r, dlt);
+    } else {
+      assert(isDenseDLT(dlt)); // no fields
+    }
+  }
+  // The values array.
+  RETURN_ON_FALSE(fieldIdx++, SparseTensorFieldKind::ValMemRef, -1u,
+                  DimLevelType::Undef);
+
+#undef RETURN_ON_FALSE
+}
+
+void sparse_tensor::foreachFieldAndTypeInSparseTensor(
+    RankedTensorType rType,
+    llvm::function_ref<bool(Type, unsigned, SparseTensorFieldKind, unsigned,
+                            DimLevelType)>
+        callback) {
+  auto enc = getSparseTensorEncoding(rType);
+  assert(enc);
+  // Construct the basic types.
+  auto *context = rType.getContext();
+  unsigned idxWidth = enc.getIndexBitWidth();
+  unsigned ptrWidth = enc.getPointerBitWidth();
+  Type indexType = IndexType::get(context);
+  Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
+  Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
+  Type eltType = rType.getElementType();
+  unsigned rank = rType.getShape().size();
+  // memref<rank x index> dimSizes
+  Type dimSizeType = MemRefType::get({rank}, indexType);
+  // memref<n x index> memSizes
+  Type memSizeType =
+      MemRefType::get({getNumDataFieldsFromEncoding(enc)}, indexType);
+  // memref<? x ptr>  pointers
+  Type ptrMemType = MemRefType::get({ShapedType::kDynamic}, ptrType);
+  // memref<? x idx>  indices
+  Type idxMemType = MemRefType::get({ShapedType::kDynamic}, idxType);
+  // memref<? x eltType> values
+  Type valMemType = MemRefType::get({ShapedType::kDynamic}, eltType);
+
+  foreachFieldInSparseTensor(
+      enc,
+      [dimSizeType, memSizeType, ptrMemType, idxMemType, valMemType,
+       callback](unsigned fieldIdx, SparseTensorFieldKind fieldKind,
+                 unsigned dim, DimLevelType dlt) -> bool {
+        switch (fieldKind) {
+        case SparseTensorFieldKind::DimSizes:
+          return callback(dimSizeType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::MemSizes:
+          return callback(memSizeType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::PtrMemRef:
+          return callback(ptrMemType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::IdxMemRef:
+          return callback(idxMemType, fieldIdx, fieldKind, dim, dlt);
+        case SparseTensorFieldKind::ValMemRef:
+          return callback(valMemType, fieldIdx, fieldKind, dim, dlt);
+        };
+      });
+}
+
+unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+  unsigned numFields = 0;
+  foreachFieldInSparseTensor(enc,
+                             [&numFields](unsigned, SparseTensorFieldKind,
+                                          unsigned, DimLevelType) -> bool {
+                               numFields++;
+                               return true;
+                             });
+  return numFields;
+}
+
+unsigned
+sparse_tensor::getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+  unsigned numFields = 0; // one value memref
+  foreachFieldInSparseTensor(enc,
+                             [&numFields](unsigned fidx, SparseTensorFieldKind,
+                                          unsigned, DimLevelType) -> bool {
+                               if (fidx >= dataFieldIdx)
+                                 numFields++;
+                               return true;
+                             });
+  assert(numFields == getNumFieldsFromEncoding(enc) - dataFieldIdx);
+  return numFields;
+}
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
                                                  StringAttr loopTag,
                                                  bool hasOutput,
                                                  bool isSparseOut,

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

 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
-
-static constexpr uint64_t dimSizesIdx = 0;
-static constexpr uint64_t memSizesIdx = 1;
-static constexpr uint64_t fieldsIdx = 2;
-
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
 /// Returns the "tuple" value of the adapted tensor.
 static UnrealizedConversionCastOp getTuple(Value tensor) {
   return llvm::cast<UnrealizedConversionCastOp>(tensor.getDefiningOp());
 }
 
+static SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor) {
+  auto tuple = getTuple(tensor);
+  return SparseTensorDescriptor(tuple.getResultTypes()[0], tuple.getInputs());
+}
+
 /// Packs the given values as a "tuple" value.
 static Value genTuple(OpBuilder &builder, Location loc, Type tp,
                       ValueRange values) {
   return builder.create<UnrealizedConversionCastOp>(loc, TypeRange(tp), values)
       .getResult(0);
 }
 
 /// Flatten a list of operands that may contain sparse tensors.
static void genStore(OpBuilder &builder, Location loc, Value val, Value mem,
   idx = toType(builder, loc, idx, builder.getIndexType());
   val = toType(builder, loc, val,
                mem.getType().cast<ShapedType>().getElementType());
   builder.create<memref::StoreOp>(loc, val, mem, idx);
 }
 
 /// Creates a straightforward counting for-loop.
 static scf::ForOp createFor(OpBuilder &builder, Location loc, Value upper,
-                            SmallVectorImpl<Value> &fields,
+                            MutableArrayRef<Value> fields,
                             Value lower = Value()) {
   Type indexType = builder.getIndexType();
   if (!lower)
     lower = constantZero(builder, loc, indexType);
   Value one = constantOne(builder, loc, indexType);
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lower, upper, one, fields);
   for (unsigned i = 0, e = fields.size(); i < e; i++)
     fields[i] = forOp.getRegionIterArg(i);
static Optional<Value> sizeFromTensorAtDim(OpBuilder &builder, Location loc,
   // Access into static dimension can query original type directly.
   // Note that this is typically already done by DimOp's folding.
   auto shape = tensorTp.getShape();
   if (!ShapedType::isDynamic(shape[dim]))
     return constantIndex(builder, loc, shape[dim]);
 
   // Any other query can consult the dimSizes array at field DimSizesIdx,
   // accounting for the reordering applied to the sparse storage.
-  auto tuple = getTuple(adaptedValue);
+  auto desc = getDescriptorFromTensorTuple(adaptedValue);
   Value idx = constantIndex(builder, loc, toStoredDim(tensorTp, dim));
-  return builder
-      .create<memref::LoadOp>(loc, tuple.getInputs()[dimSizesIdx], idx)
+  return builder.create<memref::LoadOp>(loc, desc.getDimSizesMemRef(), idx)
       .getResult();
 }
 
+//// Gets the dimension size at the given stored dimension 'd', either as a
+//// constant for a static size, or otherwise dynamically descizes.
+// Value sizeAtStoredDim(OpBuilder &builder, Location loc, Rankdesc rtp,
+//                      SmallVectorImpl<Value> &fieldesc d) {
+//  unsigned dim = tdesc, d);
+//  auto shape descpe();
+//  if (!ShapedType::isDyndescim]))
+//    return constantIndex(builder, descim]);
+//  return genLoad(builder, loc, fiedescIdx],
+//                 constandescdesc;
+//}
+
 // Gets the dimension size at the given stored dimension 'd', either as a
 // constant for a static size, or otherwise dynamically through memSizes.
-Value sizeAtStoredDim(OpBuilder &builder, Location loc, RankedTensorType rtp,
-                      SmallVectorImpl<Value> &fields, unsigned d) {
+Value sizeAtStoredDim(OpBuilder &builder, Location loc,
+                      SparseTensorDescriptor desc, unsigned d) {
+  RankedTensorType rtp = desc.getTensorType();
   unsigned dim = toOrigDim(rtp, d);
   auto shape = rtp.getShape();
   if (!ShapedType::isDynamic(shape[dim]))
     return constantIndex(builder, loc, shape[dim]);
-  return genLoad(builder, loc, fields[dimSizesIdx],
-                 constantIndex(builder, loc, d));
-}
 
-/// Translates field index to memSizes index.
-static unsigned getMemSizesIndex(unsigned field) {
-  assert(fieldsIdx <= field);
-  return field - fieldsIdx;
+  return genLoad(builder, loc, desc.getDimSizesMemRef(),
+                 constantIndex(builder, loc, d));
 }
 
-/// Creates a pushback op for given field and updates the fields array
-/// accordingly. This operation also updates the memSizes contents.
 static void createPushback(OpBuilder &builder, Location loc,
-                           SmallVectorImpl<Value> &fields, unsigned field,
+                           MutSparseTensorDescriptor desc, unsigned fidx,
                            Value value, Value repeat = Value()) {
-  assert(fieldsIdx <= field && field < fields.size());
-  Type etp = fields[field].getType().cast<ShapedType>().getElementType();
-  fields[field] = builder.create<PushBackOp>(
-      loc, fields[field].getType(), fields[memSizesIdx], fields[field],
-      toType(builder, loc, value, etp), APInt(64, getMemSizesIndex(field)),
+  Type etp = desc.getElementType(fidx);
+  Value field = desc.getField(fidx);
+  Value newField = builder.create<PushBackOp>(
+      loc, field.getType(), desc.getMemSizesMemRef(), field,
+      toType(builder, loc, value, etp), APInt(64, getFieldMemSizesIndex(fidx)),
       repeat);
-}
-
-/// Returns field index of sparse tensor type for pointers/indices, when set.
-static unsigned getFieldIndex(Type type, unsigned ptrDim, unsigned idxDim) {
-  assert(getSparseTensorEncoding(type));
-  RankedTensorType rType = type.cast<RankedTensorType>();
-  unsigned field = fieldsIdx; // start past header
-  for (unsigned r = 0, rank = rType.getShape().size(); r < rank; r++) {
-    if (isCompressedDim(rType, r)) {
-      if (r == ptrDim)
-        return field;
-      field++;
-      if (r == idxDim)
-        return field;
-      field++;
-    } else if (isSingletonDim(rType, r)) {
-      if (r == idxDim)
-        return field;
-      field++;
-    } else {
-      assert(isDenseDim(rType, r)); // no fields
-    }
-  }
-  assert(ptrDim == -1u && idxDim == -1u);
-  return field + 1; // return values field index
+  desc.setField(fidx, newField);
 }
 
 /// Maps a sparse tensor type to the appropriate compounded buffers.
 static Optional<LogicalResult>
 convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
   auto enc = getSparseTensorEncoding(type);
   if (!enc)
     return llvm::None;
-  // Construct the basic types.
-  auto *context = type.getContext();
-  unsigned idxWidth = enc.getIndexBitWidth();
-  unsigned ptrWidth = enc.getPointerBitWidth();
+
   RankedTensorType rType = type.cast<RankedTensorType>();
-  Type indexType = IndexType::get(context);
-  Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
-  Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
-  Type eltType = rType.getElementType();
-  //
-  // Sparse tensor storage scheme for 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 memSizes array.
-  //
-  // struct {
-  //   memref<rank x index> dimSizes     ; size in each dimension
-  //   memref<n x index> memSizes        ; sizes of ptrs/inds/values
-  //   ; per-dimension d:
-  //   ;  if dense:
-  //        <nothing>
-  //   ;  if compresed:
-  //        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
-  //        memref<? x idx>  indices-d   ; indices for sparse dim d
-  //   ;  if singleton:
-  //        memref<? x idx>  indices-d   ; indices for singleton dim d
-  //   memref<? x eltType> values        ; values
-  // };
-  //
-  unsigned rank = rType.getShape().size();
-  unsigned lastField = getFieldIndex(type, -1u, -1u);
-  // The dimSizes array and memSizes array.
-  fields.push_back(MemRefType::get({rank}, indexType));
-  fields.push_back(MemRefType::get({getMemSizesIndex(lastField)}, indexType));
-  // Per-dimension storage.
-  for (unsigned r = 0; r < rank; r++) {
-    // Dimension level types apply in order to the reordered dimension.
-    // As a result, the compound type can be constructed directly in the given
-    // order. Clients of this type know what field is what from the sparse
-    // tensor type.
-    if (isCompressedDim(rType, r)) {
-      fields.push_back(MemRefType::get({ShapedType::kDynamic}, ptrType));
-      fields.push_back(MemRefType::get({ShapedType::kDynamic}, idxType));
-    } else if (isSingletonDim(rType, r)) {
-      fields.push_back(MemRefType::get({ShapedType::kDynamic}, idxType));
-    } else {
-      assert(isDenseDim(rType, r)); // no fields
-    }
-  }
-  // The values array.
-  fields.push_back(MemRefType::get({ShapedType::kDynamic}, eltType));
-  assert(fields.size() == lastField);
+  foreachFieldAndTypeInSparseTensor(
+      rType,
+      [&fields](Type fieldType, unsigned fieldIdx,
+                SparseTensorFieldKind /*fieldKind*/, unsigned /*dim*/,
+                DimLevelType /*dlt*/) -> bool {
+        assert(fieldIdx == fields.size());
+        fields.push_back(fieldType);
+        return true;
+      });
   return success();
 }
 
 /// Generates code that allocates a sparse storage scheme for given rank.
 static void allocSchemeForRank(OpBuilder &builder, Location loc,
-                               RankedTensorType rtp,
-                               SmallVectorImpl<Value> &fields, unsigned field,
-                               unsigned r0) {
+                               MutSparseTensorDescriptor desc, unsigned r0) {
+  RankedTensorType rtp = desc.getTensorType();
   unsigned rank = rtp.getShape().size();
   Value linear = constantIndex(builder, loc, 1);
   for (unsigned r = r0; r < rank; r++) {
     if (isCompressedDim(rtp, r)) {
       // Append linear x pointers, initialized to zero. Since each compressed
       // dimension initially already has a single zero entry, this maintains
       // the desired "linear + 1" length property at all times.
-      unsigned ptrWidth = getSparseTensorEncoding(rtp).getPointerBitWidth();
-      Type indexType = builder.getIndexType();
-      Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
+      Type ptrType = desc.getPtrElementType();
       Value ptrZero = constantZero(builder, loc, ptrType);
-      createPushback(builder, loc, fields, field, ptrZero, linear);
+      unsigned fidx = desc.getPtrMemRefIndex(r);
+      createPushback(builder, loc, desc, fidx, ptrZero, linear);
       return;
     }
     if (isSingletonDim(rtp, r)) {
       return; // nothing to do
     }         // Keep compounding the size, but nothing needs to be initialized
-      // at this level. We will eventually reach a compressed level or
-      // otherwise the values array for the from-here "all-dense" case.
-      assert(isDenseDim(rtp, r));
-      Value size = sizeAtStoredDim(builder, loc, rtp, fields, r);
-      linear = builder.create<arith::MulIOp>(loc, linear, size);
+    Value size = sizeAtStoredDim(builder, loc, desc, r);
   }
   // Reached values array so prepare for an insertion.
   Value valZero = constantZero(builder, loc, rtp.getElementType());
-  createPushback(builder, loc, fields, field, valZero, linear);
-  assert(fields.size() == ++field);
+  createPushback(builder, loc, desc, desc.getValMemRefIndex(), valZero, linear);
 }
 
 /// Creates allocation operation.
-static Value createAllocation(OpBuilder &builder, Location loc, Type type,
-                              Value sz, bool enableInit) {
-  auto memType = MemRefType::get({ShapedType::kDynamic}, type);
-  Value buffer = builder.create<memref::AllocOp>(loc, memType, sz);
+static Value createAllocation(OpBuilder &builder, Location loc,
+                              MemRefType memRefType, Value sz,
+                              bool enableInit) {
+  Value buffer = builder.create<memref::AllocOp>(loc, memRefType, sz);
+  Type elemType = memRefType.getElementType();
   if (enableInit) {
-    Value fillValue =
-        builder.create<arith::ConstantOp>(loc, type, builder.getZeroAttr(type));
+    Value fillValue = builder.create<arith::ConstantOp>(
+        loc, elemType, builder.getZeroAttr(elemType));
     builder.create<linalg::FillOp>(loc, fillValue, buffer);
   }
   return buffer;
 }
 
 /// Creates allocation for each field in sparse tensor type. Note that
 /// for all dynamic memrefs, the memory size is really the capacity of
 /// the "vector", while the actual size resides in the sizes array.
 ///
 /// TODO: for efficiency, we will need heuristis to make educated guesses
 ///       on the required capacities (see heuristic variable).
 ///
 static void createAllocFields(OpBuilder &builder, Location loc, Type type,
                               ValueRange dynSizes, bool enableInit,
                               SmallVectorImpl<Value> &fields) {
-  auto enc = getSparseTensorEncoding(type);
-  assert(enc);
-  // Construct the basic types.
-  unsigned idxWidth = enc.getIndexBitWidth();
-  unsigned ptrWidth = enc.getPointerBitWidth();
   RankedTensorType rtp = type.cast<RankedTensorType>();
-  Type indexType = builder.getIndexType();
-  Type idxType = idxWidth ? builder.getIntegerType(idxWidth) : indexType;
-  Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
-  Type eltType = rtp.getElementType();
-  auto shape = rtp.getShape();
-  unsigned rank = shape.size();
   Value heuristic = constantIndex(builder, loc, 16);
+
+  foreachFieldAndTypeInSparseTensor(
+      rtp,
+      [&](Type fType, unsigned fIdx, SparseTensorFieldKind fKind,
+          unsigned /*dim*/, DimLevelType /*dlt*/) -> bool {
+        assert(fields.size() == fIdx);
+        auto memRefTp = fType.cast<MemRefType>();
+        Value field;
+        switch (fKind) {
+        case SparseTensorFieldKind::DimSizes:
+        case SparseTensorFieldKind::MemSizes:
+          field = builder.create<memref::AllocOp>(loc, memRefTp);
+          break;
+        case SparseTensorFieldKind::PtrMemRef:
+        case SparseTensorFieldKind::IdxMemRef:
+        case SparseTensorFieldKind::ValMemRef:
+          field =
+              createAllocation(builder, loc, memRefTp, heuristic, enableInit);
+          break;
+        }
+        assert(field);
+        fields.push_back(field);
+        // Returns true to ontinue the iteration.
+        return true;
+      });
+
+  MutSparseTensorDescriptor desc(rtp, fields);
+
   // Build original sizes.
   SmallVector<Value> sizes;
+  auto shape = rtp.getShape();
+  unsigned rank = shape.size();
   for (unsigned r = 0, o = 0; r < rank; r++) {
     if (ShapedType::isDynamic(shape[r]))
       sizes.push_back(dynSizes[o++]);
     else
       sizes.push_back(constantIndex(builder, loc, shape[r]));
   }
-  // The dimSizes array and memSizes array.
-  unsigned lastField = getFieldIndex(type, -1u, -1u);
-  Value dimSizes =
-      builder.create<memref::AllocOp>(loc, MemRefType::get({rank}, indexType));
-  Value memSizes = builder.create<memref::AllocOp>(
-      loc, MemRefType::get({getMemSizesIndex(lastField)}, indexType));
-  fields.push_back(dimSizes);
-  fields.push_back(memSizes);
-  // Per-dimension storage.
-  for (unsigned r = 0; r < rank; r++) {
-    if (isCompressedDim(rtp, r)) {
-      fields.push_back(
-          createAllocation(builder, loc, ptrType, heuristic, enableInit));
-      fields.push_back(
-          createAllocation(builder, loc, idxType, heuristic, enableInit));
-    } else if (isSingletonDim(rtp, r)) {
-      fields.push_back(
-          createAllocation(builder, loc, idxType, heuristic, enableInit));
-    } else {
-      assert(isDenseDim(rtp, r)); // no fields
-    }
-  }
-  // The values array.
-  fields.push_back(
-      createAllocation(builder, loc, eltType, heuristic, enableInit));
-  assert(fields.size() == lastField);
   // Initialize the storage scheme to an empty tensor. Initialized memSizes
   // to all zeros, sets the dimSizes to known values and gives all pointer
   // fields an initial zero entry, so that it is easier to maintain the
   // "linear + 1" length property.
   builder.create<linalg::FillOp>(
-      loc, ValueRange{constantZero(builder, loc, indexType)},
-      ValueRange{memSizes}); // zero memSizes
-  Value ptrZero = constantZero(builder, loc, ptrType);
-  for (unsigned r = 0, field = fieldsIdx; r < rank; r++) {
+      loc, constantZero(builder, loc, builder.getIndexType()),
+      desc.getMemSizesMemRef()); // zero memSizes
+
+  Value ptrZero = constantZero(builder, loc, desc.getPtrElementType());
+  for (unsigned r = 0; r < rank; r++) {
     unsigned ro = toOrigDim(rtp, r);
-    genStore(builder, loc, sizes[ro], dimSizes, constantIndex(builder, loc, r));
+    // Fills dim sizes array.
+    genStore(builder, loc, sizes[ro], desc.getDimSizesMemRef(),
+             constantIndex(builder, loc, r));
+
+    // Pushes a leading zero to pointers memref.
     if (isCompressedDim(rtp, r)) {
-      createPushback(builder, loc, fields, field, ptrZero);
-      field += 2;
-    } else if (isSingletonDim(rtp, r)) {
-      field += 1;
+      unsigned fidx =
+          desc.getDataFieldIndex(r, SparseTensorFieldKind::PtrMemRef);
+      createPushback(builder, loc, desc, fidx, ptrZero);
     }
   }
-  allocSchemeForRank(builder, loc, rtp, fields, fieldsIdx, /*rank=*/0);
+  allocSchemeForRank(builder, loc, desc, /*rank=*/0);
 }
 
 /// Helper method that generates block specific to compressed case:
 ///
 ///  plo = pointers[d][pos[d-1]]
 ///  phi = pointers[d][pos[d-1]+1]
 ///  msz = indices[d].size()
 ///  if (plo < phi) {
 ///    present = indices[d][phi-1] == i[d]
 ///  } else { // first insertion
 ///    present = false
 ///    pointers[d][pos[d-1]] = msz
 ///  }
 ///  if (present) { // index already present
 ///    next = phi-1
 ///  } else {
 ///    indices[d].push_back(i[d])
 ///    pointers[d][pos[d-1]+1] = msz+1
 ///    next = msz
 ///    <prepare dimension d + 1>
 ///  }
 ///  pos[d] = next
 static Value genCompressed(OpBuilder &builder, Location loc,
-                           RankedTensorType rtp, SmallVectorImpl<Value> &fields,
+                           MutSparseTensorDescriptor desc,
                            SmallVectorImpl<Value> &indices, Value value,
-                           Value pos, unsigned field, unsigned d) {
+                           Value pos, unsigned d) {
+  RankedTensorType rtp = desc.getTensorType();
   unsigned rank = rtp.getShape().size();
   SmallVector<Type> types;
   Type indexType = builder.getIndexType();
   Type boolType = builder.getIntegerType(1);
+  unsigned idxIndex = desc.getIdxMemRefIndex(d);
+  unsigned ptrIndex = desc.getPtrMemRefIndex(d);
   Value one = constantIndex(builder, loc, 1);
   Value pp1 = builder.create<arith::AddIOp>(loc, pos, one);
-  Value plo = genLoad(builder, loc, fields[field], pos);
-  Value phi = genLoad(builder, loc, fields[field], pp1);
-  Value psz = constantIndex(builder, loc, getMemSizesIndex(field + 1));
-  Value msz = genLoad(builder, loc, fields[memSizesIdx], psz);
+  Value plo = genLoad(builder, loc, desc.getField(ptrIndex), pos);
+  Value phi = genLoad(builder, loc, desc.getField(ptrIndex), pp1);
+  Value psz = constantIndex(builder, loc, getFieldMemSizesIndex(idxIndex));
+  Value msz = genLoad(builder, loc, desc.getMemSizesMemRef(), psz);
   Value phim1 = builder.create<arith::SubIOp>(
       loc, toType(builder, loc, phi, indexType), one);
   // Conditional expression.
   Value lt =
       builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, plo, phi);
   types.push_back(boolType);
   scf::IfOp ifOp1 = builder.create<scf::IfOp>(loc, types, lt, /*else*/ true);
   types.pop_back();
   builder.setInsertionPointToStart(&ifOp1.getThenRegion().front());
-  Value crd = genLoad(builder, loc, fields[field + 1], phim1);
+  Value crd = genLoad(builder, loc, desc.getField(idxIndex), phim1);
   Value eq = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                            toType(builder, loc, crd, indexType),
                                            indices[d]);
   builder.create<scf::YieldOp>(loc, eq);
   builder.setInsertionPointToStart(&ifOp1.getElseRegion().front());
   if (d > 0)
-    genStore(builder, loc, msz, fields[field], pos);
+    genStore(builder, loc, msz, desc.getField(ptrIndex), pos);
   builder.create<scf::YieldOp>(loc, constantI1(builder, loc, false));
   builder.setInsertionPointAfter(ifOp1);
   Value p = ifOp1.getResult(0);
-  // If present construct. Note that for a non-unique dimension level, we simply
-  // set the condition to false and rely on CSE/DCE to clean up the IR.
+  // If present construct. Note that for a non-unique dimension level, we
+  // simply set the condition to false and rely on CSE/DCE to clean up the IR.
   //
   // TODO: generate less temporary IR?
   //
-  for (unsigned i = 0, e = fields.size(); i < e; i++)
-    types.push_back(fields[i].getType());
+  for (unsigned i = 0, e = desc.getNumFields(); i < e; i++)
+    types.push_back(desc.getField(i).getType());
   types.push_back(indexType);
   if (!isUniqueDim(rtp, d))
     p = constantI1(builder, loc, false);
   scf::IfOp ifOp2 = builder.create<scf::IfOp>(loc, types, p, /*else*/ true);
   // If present (fields unaffected, update next to phim1).
   builder.setInsertionPointToStart(&ifOp2.getThenRegion().front());
-  fields.push_back(phim1);
-  builder.create<scf::YieldOp>(loc, fields);
-  fields.pop_back();
+
+  // FIXME: This does not looks like a clean way, but probably the most
+  // efficient way.
+  desc.getFields().push_back(phim1);
+  builder.create<scf::YieldOp>(loc, desc.getFields());
+  desc.getFields().pop_back();
+
   // If !present (changes fields, update next).
   builder.setInsertionPointToStart(&ifOp2.getElseRegion().front());
   Value mszp1 = builder.create<arith::AddIOp>(loc, msz, one);
-  genStore(builder, loc, mszp1, fields[field], pp1);
-  createPushback(builder, loc, fields, field + 1, indices[d]);
+  genStore(builder, loc, mszp1, desc.getField(ptrIndex), pp1);
+  createPushback(builder, loc, desc, idxIndex, indices[d]);
   // Prepare the next dimension "as needed".
   if ((d + 1) < rank)
-    allocSchemeForRank(builder, loc, rtp, fields, field + 2, d + 1);
-  fields.push_back(msz);
-  builder.create<scf::YieldOp>(loc, fields);
-  fields.pop_back();
+    allocSchemeForRank(builder, loc, desc, d + 1);
+
+  desc.getFields().push_back(msz);
+  builder.create<scf::YieldOp>(loc, desc.getFields());
+  desc.getFields().pop_back();
+
   // Update fields and return next pos.
   builder.setInsertionPointAfter(ifOp2);
   unsigned o = 0;
-  for (unsigned i = 0, e = fields.size(); i < e; i++)
-    fields[i] = ifOp2.getResult(o++);
+  for (unsigned i = 0, e = desc.getNumFields(); i < e; i++)
+    desc.setField(i, ifOp2.getResult(o++));
   return ifOp2.getResult(o);
 }
 
 /// Generates code along an insertion path without the need for a "cursor".
 /// This current insertion strategy comes at the expense of some testing
 /// overhead for each insertion. The strategy will be optimized later for
 /// common insertion patterns. The current insertion strategy also assumes
 /// insertions occur in "a reasonable order" that enables building the
 /// storage scheme in an appending/inserting kind of fashion (i.e. no
 /// in-between insertions that need data movement). The implementation
 /// relies on CSE/DCE to clean up all bookkeeping that is not needed.
 ///
 /// TODO: better unord/not-unique; also generalize, optimize, specialize!
 ///
 static void genInsert(OpBuilder &builder, Location loc, RankedTensorType rtp,
                       SmallVectorImpl<Value> &fields,
                       SmallVectorImpl<Value> &indices, Value value) {
   unsigned rank = rtp.getShape().size();
   assert(rank == indices.size());
-  unsigned field = fieldsIdx; // start past header
+  MutSparseTensorDescriptor desc(rtp, fields);
   Value pos = constantZero(builder, loc, builder.getIndexType());
   // Generate code for every dimension.
   for (unsigned d = 0; d < rank; d++) {
     if (isCompressedDim(rtp, d)) {
       // Create:
       //   if (!present) {
       //     indices[d].push_back(i[d])
       //     <update pointers and prepare dimension d + 1>
       //   }
       //   pos[d] = indices.size() - 1
       //   <insert @ pos[d] at next dimension d + 1>
-      pos = genCompressed(builder, loc, rtp, fields, indices, value, pos, field,
-                          d);
-      field += 2;
+      pos = genCompressed(builder, loc, desc, indices, value, pos, d);
     } else if (isSingletonDim(rtp, d)) {
       // Create:
       //   indices[d].push_back(i[d])
       //   pos[d] = pos[d-1]
       //   <insert @ pos[d] at next dimension d + 1>
-      createPushback(builder, loc, fields, field, indices[d]);
-      field += 1;
+      unsigned fidx = desc.getIdxMemRefIndex(d);
+      createPushback(builder, loc, desc, fidx, indices[d]);
     } else {
       assert(isDenseDim(rtp, d));
       // Construct the new position as:
       //   pos[d] = size * pos[d-1] + i[d]
       //   <insert @ pos[d] at next dimension d + 1>
-      Value size = sizeAtStoredDim(builder, loc, rtp, fields, d);
+      Value size = sizeAtStoredDim(builder, loc, desc, d);
       Value mult = builder.create<arith::MulIOp>(loc, size, pos);
       pos = builder.create<arith::AddIOp>(loc, mult, indices[d]);
     }
   }
   // Reached the actual value append/insert.
-  if (!isDenseDim(rtp, rank - 1))
-    createPushback(builder, loc, fields, field++, value);
-  else
-    genStore(builder, loc, value, fields[field++], pos);
-  assert(fields.size() == field);
+  unsigned valIdx = desc.getValMemRefIndex();
+  if (!isDenseDim(rtp, rank - 1)) {
+    createPushback(builder, loc, desc, valIdx, value);
+  } else
+    genStore(builder, loc, value, desc.getValMemRef(), pos);
 }
 
 /// Generations insertion finalization code.
 static void genEndInsert(OpBuilder &builder, Location loc, RankedTensorType rtp,
                          SmallVectorImpl<Value> &fields) {
   unsigned rank = rtp.getShape().size();
-  unsigned field = fieldsIdx; // start past header
+  unsigned field = dataFieldIdx; // start past header
   for (unsigned d = 0; d < rank; d++) {
     if (isCompressedDim(rtp, d)) {
       // Compressed dimensions need a pointer cleanup for all entries
       // that were not visited during the insertion pass.
       //
-      // TODO: avoid cleanup and keep compressed scheme consistent at all times?
+      // TODO: avoid cleanup and keep compressed scheme consistent at all
+      // times?
       //
       if (d > 0) {
         unsigned ptrWidth = getSparseTensorEncoding(rtp).getPointerBitWidth();
         Type indexType = builder.getIndexType();
         Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
-        Value mz = constantIndex(builder, loc, getMemSizesIndex(field));
+        Value mz = constantIndex(builder, loc, getFieldMemSizesIndex(field));
         Value hi = genLoad(builder, loc, fields[memSizesIdx], mz);
         Value zero = constantIndex(builder, loc, 0);
         Value one = constantIndex(builder, loc, 1);
         // Vector of only one, but needed by createFor's prototype.
         SmallVector<Value, 1> inits{genLoad(builder, loc, fields[field], zero)};
         scf::ForOp loop = createFor(builder, loc, hi, inits, one);
         Value i = loop.getInductionVar();
         Value oldv = loop.getRegionIterArg(0);
public:
   using OpAdaptor = typename SourceOp::Adaptor;
   using OpConversionPattern<SourceOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Replace the requested pointer access with corresponding field.
     // The cast_op is inserted by type converter to intermix 1:N type
     // conversion.
-    auto tuple = getTuple(adaptor.getTensor());
-    unsigned idx = Base::getIndexForOp(tuple, op);
-    auto fields = tuple.getInputs();
-    assert(idx < fields.size());
-    rewriter.replaceOp(op, fields[idx]);
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    Value field = Base::getFieldForOp(desc, op);
+    rewriter.replaceOp(op, field);
     return success();
   }
 };
 
 /// Sparse codegen rule for pointer accesses.
 class SparseToPointersConverter
     : public SparseGetterOpConverter<ToPointersOp, SparseToPointersConverter> {
 public:
   using SparseGetterOpConverter::SparseGetterOpConverter;
   // Callback for SparseGetterOpConverter.
-  static unsigned getIndexForOp(UnrealizedConversionCastOp /*tuple*/,
-                                ToPointersOp op) {
+  static Value getFieldForOp(const SparseTensorDescriptor &desc,
     uint64_t dim = op.getDimension().getZExtValue();
-    return getFieldIndex(op.getTensor().getType(), /*ptrDim=*/dim, -1u);
+    return desc.getPtrMemRef(dim);
   }
 };
 
 /// Sparse codegen rule for index accesses.
 class SparseToIndicesConverter
     : public SparseGetterOpConverter<ToIndicesOp, SparseToIndicesConverter> {
 public:
   using SparseGetterOpConverter::SparseGetterOpConverter;
   // Callback for SparseGetterOpConverter.
-  static unsigned getIndexForOp(UnrealizedConversionCastOp /*tuple*/,
-                                ToIndicesOp op) {
+  static Value getFieldForOp(const SparseTensorDescriptor &desc,
     uint64_t dim = op.getDimension().getZExtValue();
-    return getFieldIndex(op.getTensor().getType(), -1u, /*idxDim=*/dim);
+    return desc.getIdxMemRef(dim);
   }
 };
 
 /// Sparse codegen rule for value accesses.
 class SparseToValuesConverter
     : public SparseGetterOpConverter<ToValuesOp, SparseToValuesConverter> {
 public:
   using SparseGetterOpConverter::SparseGetterOpConverter;
   // Callback for SparseGetterOpConverter.
-  static unsigned getIndexForOp(UnrealizedConversionCastOp tuple,
-                                ToValuesOp /*op*/) {
-    // The last field holds the value buffer.
-    return tuple.getInputs().size() - 1;
+  static Value getFieldForOp(const SparseTensorDescriptor &desc,
+    return desc.getValMemRef();
   }
 };
 
 /// Sparse codegen rule for the convert operator.
 class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult

 class SparseNumberOfEntriesConverter
     : public OpConversionPattern<NumberOfEntriesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Query memSizes for the actually stored values size.
-    auto tuple = getTuple(adaptor.getTensor());
-    auto fields = tuple.getInputs();
-    unsigned lastField = fields.size() - 1;
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     Value field =
-        constantIndex(rewriter, op.getLoc(), getMemSizesIndex(lastField));
-    rewriter.replaceOpWithNewOp<memref::LoadOp>(op, fields[memSizesIdx], field);
+        constantIndex(rewriter, op.getLoc(), desc.getValMemSizesIndex());
+    rewriter.replaceOpWithNewOp<memref::LoadOp>(op, desc.getMemSizesMemRef(),
+                                                field);
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into an actual buffer.