[mlir][sparse] Generate AOS subviews on-demand.

[bixia]

Summary

Previously, we generate AOS subviews for indices buffers when constructing an
immutable sparse tensor descriptor. We now only generate such subviews when
getIdxMemRefOrView is requested.

[Peiming]

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

122

I am not sure whether a mutDesc is really needed here.

[Peiming]

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

313–314

You can move the getter to the base class (shared by mut/immut)

[Peiming]

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

123

Seems that you only need the idxMemSize from mutDesc, but these getters can be safely moved to the base class (they are not updating the descriptor anyway, see my previous comments), then you can get rid of this mutDesc here.

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

339–340

Did you rebase? if not, these template need to be removed as well.

[bixia]

Move some getters to the Impl class, remove the need for mutDesc for a case.

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

123

Removed mutDesc here.

[Peiming]

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

393

These two parameter seems unecessary now

[bixia]

Removed unused parameters.

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

matchAndRewrite(SourceOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     SpecifierStructBuilder spec(adaptor.getSpecifier());
     Value v;
     if (op.getSpecifierKind() == StorageSpecifierKind::DimSize) {
       v = Base::onDimSize(rewriter, op, spec,
                           op.getDim().value().getZExtValue());
     } else {
       auto enc = op.getSpecifier().getType().getEncoding();
-      StorageLayout<true> layout(enc);
+      StorageLayout layout(enc);
       Optional<unsigned> dim = std::nullopt;
       if (op.getDim())
         dim = op.getDim().value().getZExtValue();
       unsigned idx = layout.getMemRefFieldIndex(op.getSpecifierKind(), dim);
       v = Base::onMemSize(rewriter, op, spec, idx);
     }
 
     rewriter.replaceOp(op, v);

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

static Value genCompressed(OpBuilder &builder, Location loc,
   RankedTensorType rtp = desc.getTensorType();
   unsigned rank = rtp.getShape().size();
   SmallVector<Type> types;
   Type indexType = builder.getIndexType();
   Type boolType = builder.getIntegerType(1);
   unsigned idxIndex;
   unsigned idxStride;
   std::tie(idxIndex, idxStride) = desc.getIdxMemRefIndexAndStride(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, desc.getMemRefField(ptrIndex), pos);
-  Value phi = genLoad(builder, loc, desc.getMemRefField(ptrIndex), pp1);
+  Value plo = genLoad(builder, loc, desc.getPtrMemRef(d), pos);
+  Value phi = genLoad(builder, loc, desc.getPtrMemRef(d), pp1);
   Value msz = desc.getIdxMemSize(builder, loc, d);
   Value idxStrideC;
   if (idxStride > 1) {
     idxStrideC = constantIndex(builder, loc, idxStride);
     msz = builder.create<arith::DivUIOp>(loc, msz, idxStrideC);
   }
   Value phim1 = builder.create<arith::SubIOp>(
       loc, toType(builder, loc, phi, indexType), one);
Value crd = genLoad(
       idxStride > 1 ? builder.create<arith::MulIOp>(loc, phim1, idxStrideC)
                     : 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, desc.getMemRefField(ptrIndex), pos);
+    genStore(builder, loc, msz, desc.getPtrMemRef(d), 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.
   //
   // TODO: generate less temporary IR?
   //
static Value genCompressed(OpBuilder &builder, Location loc,
   // 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, desc.getMemRefField(ptrIndex), pp1);
+  genStore(builder, loc, mszp1, desc.getPtrMemRef(d), pp1);
   createPushback(builder, loc, desc, SparseTensorFieldKind::IdxMemRef, d,
                  indices[d]);
   // Prepare the next dimension "as needed".
   if ((d + 1) < rank)
     allocSchemeForRank(builder, loc, desc, d + 1);
 
   desc.getFields().push_back(msz);
   builder.create<scf::YieldOp>(loc, desc.getFields());
public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     std::optional<int64_t> index = op.getConstantIndex();
     if (!index || !getSparseTensorEncoding(adaptor.getSource().getType()))
       return failure();
 
-    Location loc = op.getLoc();
-    auto desc =
-        getDescriptorFromTensorTuple(rewriter, loc, adaptor.getSource());
-    auto sz = sizeFromTensorAtDim(rewriter, loc, desc, *index);
+    auto desc = getDescriptorFromTensorTuple(adaptor.getSource());
+    auto sz = sizeFromTensorAtDim(rewriter, op.getLoc(), desc, *index);
 
     if (!sz)
       return failure();
 
     rewriter.replaceOp(op, *sz);
     return success();
   }
 };

 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (!getSparseTensorEncoding(op.getTensor().getType()))
       return failure();
     Location loc = op->getLoc();
-    auto desc =
-        getDescriptorFromTensorTuple(rewriter, loc, adaptor.getTensor());
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     RankedTensorType srcType =
         op.getTensor().getType().cast<RankedTensorType>();
     Type eltType = srcType.getElementType();
     Type boolType = rewriter.getIntegerType(1);
     Type idxType = rewriter.getIndexType();
     // All initialization should be done on entry of the loop nest.
     rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
     // Determine the size for access expansion (always the innermost stored
public:
   using OpAdaptor = typename ToPointersOp::Adaptor;
   using OpConversionPattern<ToPointersOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToPointersOp 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 desc = getDescriptorFromTensorTuple(rewriter, op.getLoc(),
-                                             adaptor.getTensor());
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     uint64_t dim = op.getDimension().getZExtValue();
     rewriter.replaceOp(op, desc.getPtrMemRef(dim));
     return success();
   }
 };
 
 /// Sparse codegen rule for index accesses.
 class SparseToIndicesConverter : public OpConversionPattern<ToIndicesOp> {
 public:
   using OpAdaptor = typename ToIndicesOp::Adaptor;
   using OpConversionPattern<ToIndicesOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToIndicesOp 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 desc = getDescriptorFromTensorTuple(rewriter, op.getLoc(),
-                                             adaptor.getTensor());
+    Location loc = op.getLoc();
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     uint64_t dim = op.getDimension().getZExtValue();
-    Value field = desc.getIdxMemRef(dim);
+    Value field = desc.getIdxMemRefOrView(rewriter, loc, dim);
 
     // Insert a cast to bridge the actual type to the user expected type. If the
     // actual type and the user expected type aren't compatible, the compiler or
     // the runtime will issue an error.
     Type resType = op.getResult().getType();
     if (resType != field.getType())
-      field = rewriter.create<memref::CastOp>(op.getLoc(), resType, field);
+      field = rewriter.create<memref::CastOp>(loc, resType, field);
     rewriter.replaceOp(op, field);
 
     return success();
   }
 };
 
 /// Sparse codegen rule for accessing the linear indices buffer.
 class SparseToIndicesBufferConverter
public:
   using OpAdaptor = typename ToValuesOp::Adaptor;
   using OpConversionPattern<ToValuesOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToValuesOp 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 desc = getDescriptorFromTensorTuple(rewriter, op.getLoc(),
-                                             adaptor.getTensor());
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     rewriter.replaceOp(op, desc.getValMemRef());
     return success();
   }
 };
 
 /// Sparse codegen rule for the convert operator.
 class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
 public:

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

 /// callback invocations.
 /// 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(unsigned /*fieldIdx*/,
                             SparseTensorFieldKind /*fieldKind*/,
                             unsigned /*dim (if applicable)*/,
-                            DimLevelType /*DLT (if applicable)*/)>,
-    bool isBuffer = false);
+                            DimLevelType /*DLT (if applicable)*/)>);
 
 /// Same as above, except that it also builds the Type 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)*/)>);
 
 /// Gets the total number of fields for the given sparse tensor encoding.
-unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc, bool isBuffer);
+unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc);
 
 /// Gets the total number of data fields (index arrays, pointer arrays, and a
 /// value array) for the given sparse tensor encoding.
 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::DimSize);
   return static_cast<SparseTensorFieldKind>(kind);
 }
 
 /// Provides methods to access fields of a sparse tensor with the given
-/// encoding. When isBuffer is true, the fields are the actual buffers of the
-/// sparse tensor storage. In particular, when a linear buffer is used to
-/// store the COO data as an array-of-structures, the fields include the
-/// linear buffer (isBuffer=true) or includes the subviews of the buffer for the
-/// indices (isBuffer=false).
-template <bool isBuffer>
+/// encoding.
 class StorageLayout {
 public:
   explicit StorageLayout(SparseTensorEncodingAttr enc) : enc(enc) {}
 
   ///
   /// Getters: get the field index for required field.
   ///
 
   unsigned getMemRefFieldIndex(SparseTensorFieldKind kind,
                                std::optional<unsigned> dim) const {
     return getFieldIndexAndStride(kind, dim).first;
   }
 
   unsigned getMemRefFieldIndex(StorageSpecifierKind kind,
                                std::optional<unsigned> dim) const {
     return getMemRefFieldIndex(toFieldKind(kind), dim);
   }
 
   static unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
-    return sparse_tensor::getNumFieldsFromEncoding(enc, isBuffer);
+    return sparse_tensor::getNumFieldsFromEncoding(enc);
   }
 
   static void foreachFieldInSparseTensor(
       const SparseTensorEncodingAttr enc,
       llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
                               DimLevelType)>
           callback) {
-    return sparse_tensor::foreachFieldInSparseTensor(enc, callback, isBuffer);
+    return sparse_tensor::foreachFieldInSparseTensor(enc, callback);
   }
 
   std::pair<unsigned, unsigned>
   getFieldIndexAndStride(SparseTensorFieldKind kind,
                          std::optional<unsigned> dim) const {
     unsigned fieldIdx = -1u;
     unsigned stride = 1;
-    if (isBuffer && kind == SparseTensorFieldKind::IdxMemRef) {
+    if (kind == SparseTensorFieldKind::IdxMemRef) {
       assert(dim.has_value());
       unsigned cooStart = getCOOStart(enc);
       unsigned rank = enc.getDimLevelType().size();
       if (dim.value() >= cooStart && dim.value() < rank) {
         dim = cooStart;
         stride = rank - cooStart;
       }
     }
protected:
   // buffers to append value for some special cases, though users should be
   // responsible to restore 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 ValueArrayRef = typename std::conditional<mut, SmallVectorImpl<Value> &,
                                                   ValueRange>::type;
 
-  SparseTensorDescriptorImpl(Type tp)
-      : rType(tp.cast<RankedTensorType>()), fields() {}
-
   SparseTensorDescriptorImpl(Type tp, ValueArrayRef fields)
       : rType(tp.cast<RankedTensorType>()), fields(fields) {
-    sanityCheck();
-  }
-
-  void sanityCheck() {
-    assert(getSparseTensorEncoding(rType) &&
-           StorageLayout<mut>::getNumFieldsFromEncoding(
-               getSparseTensorEncoding(rType)) == fields.size());
+    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>>);
   }
 
 public:
   unsigned getMemRefFieldIndex(SparseTensorFieldKind kind,
                                std::optional<unsigned> dim) const {
     // Delegates to storage layout.
-    StorageLayout<mut> layout(getSparseTensorEncoding(rType));
+    StorageLayout layout(getSparseTensorEncoding(rType));
     return layout.getMemRefFieldIndex(kind, dim);
   }
 
-  unsigned getPtrMemRefIndex(unsigned ptrDim) const {
-    return getMemRefFieldIndex(SparseTensorFieldKind::PtrMemRef, ptrDim);
-  }
-
-  unsigned getIdxMemRefIndex(unsigned idxDim) const {
-    return getMemRefFieldIndex(SparseTensorFieldKind::IdxMemRef, idxDim);
-  }
-
-  unsigned getValMemRefIndex() const {
-    return getMemRefFieldIndex(SparseTensorFieldKind::ValMemRef, std::nullopt);
-  }
-
   unsigned getNumFields() const { return fields.size(); }
 
   ///
   /// Getters: get the value for required field.
   ///
 
   Value getSpecifierField(OpBuilder &builder, Location loc,
                           StorageSpecifierKind kind,
                           std::optional<unsigned> dim) const {
     SparseTensorSpecifier md(fields.back());
     return md.getSpecifierField(builder, loc, kind, dim);
   }
 
   Value getDimSize(OpBuilder &builder, Location loc, unsigned dim) const {
     return getSpecifierField(builder, loc, StorageSpecifierKind::DimSize, dim);
   }
 
   Value getPtrMemRef(unsigned ptrDim) const {
     return getMemRefField(SparseTensorFieldKind::PtrMemRef, ptrDim);
   }
 
-  Value getIdxMemRef(unsigned idxDim) const {
-    return getMemRefField(SparseTensorFieldKind::IdxMemRef, idxDim);
-  }
-
   Value getValMemRef() const {
     return getMemRefField(SparseTensorFieldKind::ValMemRef, std::nullopt);
   }
 
   Value getMemRefField(SparseTensorFieldKind kind,
                        std::optional<unsigned> dim) const {
     return fields[getMemRefFieldIndex(kind, dim)];
   }
 
   Value getMemRefField(unsigned fidx) const {
     assert(fidx < fields.size() - 1);
     return fields[fidx];
   }
 
-  Value getField(unsigned fidx) const {
-    assert(fidx < fields.size());
-    return fields[fidx];
+  Value getPtrMemSize(OpBuilder &builder, Location loc, unsigned dim) const {
+    return getSpecifierField(builder, loc, StorageSpecifierKind::PtrMemSize,
+                             dim);
   }
 
-  ValueRange getMemRefFields() const {
-    ValueRange ret = fields;
-    // Drop the last metadata fields.
-    return ret.slice(0, fields.size() - 1);
+  Value getIdxMemSize(OpBuilder &builder, Location loc, unsigned dim) const {
+    return getSpecifierField(builder, loc, StorageSpecifierKind::IdxMemSize,
+                             dim);
+  }
+
+  Value getValMemSize(OpBuilder &builder, Location loc) const {
+    return getSpecifierField(builder, loc, StorageSpecifierKind::ValMemSize,
+                             std::nullopt);
   }
 
   Type getMemRefElementType(SparseTensorFieldKind kind,
                             std::optional<unsigned> dim) const {
     return getMemRefField(kind, dim)
         .getType()
         .template cast<MemRefType>()
         .getElementType();
   }
 
   RankedTensorType getTensorType() const { return rType; }
   ValueArrayRef getFields() const { return fields; }
 
 protected:
   RankedTensorType rType;
   ValueArrayRef fields;
 };
 
 class MutSparseTensorDescriptor : public SparseTensorDescriptorImpl<true> {
 public:
   MutSparseTensorDescriptor(Type tp, ValueArrayRef buffers)
       : SparseTensorDescriptorImpl<true>(tp, buffers) {}
 
-  ///
-  /// Getters: get the value for required field.
-  ///
-
-  Value getPtrMemSize(OpBuilder &builder, Location loc, unsigned dim) const {
-    return getSpecifierField(builder, loc, StorageSpecifierKind::PtrMemSize,
-                             dim);
-  }
-
-  Value getIdxMemSize(OpBuilder &builder, Location loc, unsigned dim) const {
-    return getSpecifierField(builder, loc, StorageSpecifierKind::IdxMemSize,
-                             dim);
+  Value getField(unsigned fidx) const {
+    assert(fidx < fields.size());
+    return fields[fidx];
   }
 
-  Value getValMemSize(OpBuilder &builder, Location loc) const {
-    return getSpecifierField(builder, loc, StorageSpecifierKind::ValMemSize,
-                             std::nullopt);
+  ValueRange getMemRefFields() const {
+    ValueRange ret = fields;
+    // Drop the last metadata fields.
+    return ret.slice(0, fields.size() - 1);
   }
 
   ///
   /// Setters: update the value for required field (only enabled for
   /// MutSparseTensorDescriptor).
   ///
 
   void setMemRefField(SparseTensorFieldKind kind, std::optional<unsigned> dim,
                       Value v) {
     fields[getMemRefFieldIndex(kind, dim)] = v;
   }
 
   void setMemRefField(unsigned fidx, Value v) {
     assert(fidx < fields.size() - 1);
     fields[fidx] = v;
   }
 
   void setField(unsigned fidx, Value v) {
     assert(fidx < fields.size());
     fields[fidx] = v;
   }
 
   void setSpecifierField(OpBuilder &builder, Location loc,
                          StorageSpecifierKind kind, std::optional<unsigned> dim,
                          Value v) {
     SparseTensorSpecifier md(fields.back());
     md.setSpecifierField(builder, loc, v, kind, dim);
     fields.back() = md;
   }
 
   void setDimSize(OpBuilder &builder, Location loc, unsigned dim, Value v) {
     setSpecifierField(builder, loc, StorageSpecifierKind::DimSize, dim, v);
   }
 
   std::pair<unsigned, unsigned>
   getIdxMemRefIndexAndStride(unsigned idxDim) const {
-    StorageLayout<true> layout(getSparseTensorEncoding(rType));
+    StorageLayout layout(getSparseTensorEncoding(rType));
     return layout.getFieldIndexAndStride(SparseTensorFieldKind::IdxMemRef,
                                          idxDim);
   }
 
   Value getAOSMemRef() const {
     auto enc = getSparseTensorEncoding(rType);
     unsigned cooStart = getCOOStart(enc);
     assert(cooStart < enc.getDimLevelType().size());
-    return getIdxMemRef(cooStart);
+    return getMemRefField(SparseTensorFieldKind::IdxMemRef, cooStart);
   }
 };
 
 class SparseTensorDescriptor : public SparseTensorDescriptorImpl<false> {
 public:
-  SparseTensorDescriptor(OpBuilder &builder, Location loc, Type tp,
-                         ValueArrayRef buffers);
+  SparseTensorDescriptor(Type tp, ValueArrayRef buffers)
+      : SparseTensorDescriptorImpl<false>(tp, buffers) {}
 
-private:
-  // Store the fields passed to SparseTensorDescriptorImpl when the tensor has
-  // a COO region.
-  SmallVector<Value> expandedFields;
+  Value getIdxMemRefOrView(OpBuilder &builder, Location loc,
+                           unsigned idxDim) const;
 };
 
 /// Returns the "tuple" value of the adapted tensor.
 inline UnrealizedConversionCastOp getTuple(Value tensor) {
   return llvm::cast<UnrealizedConversionCastOp>(tensor.getDefiningOp());
 }
 
 /// Packs the given values as a "tuple" value.
 inline Value genTuple(OpBuilder &builder, Location loc, Type tp,
                       ValueRange values) {
   return builder.create<UnrealizedConversionCastOp>(loc, TypeRange(tp), values)
       .getResult(0);
 }
 
 inline Value genTuple(OpBuilder &builder, Location loc,
                       MutSparseTensorDescriptor desc) {
   return genTuple(builder, loc, desc.getTensorType(), desc.getFields());
 }
 
-inline SparseTensorDescriptor
-getDescriptorFromTensorTuple(OpBuilder &builder, Location loc, Value tensor) {
+inline SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor) {
   auto tuple = getTuple(tensor);
-  return SparseTensorDescriptor(builder, loc, tuple.getResultTypes()[0],
-                                tuple.getInputs());
+  return SparseTensorDescriptor(tuple.getResultTypes()[0], tuple.getInputs());
 }
 
 inline MutSparseTensorDescriptor
 getMutDescriptorFromTensorTuple(Value tensor, SmallVectorImpl<Value> &fields) {
   auto tuple = getTuple(tensor);
   fields.assign(tuple.getInputs().begin(), tuple.getInputs().end());
   return MutSparseTensorDescriptor(tuple.getResultTypes()[0], fields);
 }
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_SPARSETENSORBUILDER_H_

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

specifier = builder.create<SetStorageSpecifierOp>(
       loc, specifier, kind, fromOptionalInt(specifier.getContext(), dim),
       createIndexCast(builder, loc, v, getFieldType(kind, dim)));
 }
 
 //===----------------------------------------------------------------------===//
 // SparseTensorDescriptor methods.
 //===----------------------------------------------------------------------===//
 
-sparse_tensor::SparseTensorDescriptor::SparseTensorDescriptor(
-    OpBuilder &builder, Location loc, Type tp, ValueArrayRef buffers)
-    : SparseTensorDescriptorImpl<false>(tp), expandedFields() {
-  SparseTensorEncodingAttr enc = getSparseTensorEncoding(tp);
-  unsigned rank = enc.getDimLevelType().size();
+Value sparse_tensor::SparseTensorDescriptor::getIdxMemRefOrView(
+    OpBuilder &builder, Location loc, unsigned idxDim) const {
+  auto enc = getSparseTensorEncoding(rType);
   unsigned cooStart = getCOOStart(enc);
-  if (cooStart < rank) {
-    ValueRange beforeFields = buffers.drop_back(3);
-    expandedFields.append(beforeFields.begin(), beforeFields.end());
-    Value buffer = buffers[buffers.size() - 3];
-
+  unsigned idx = idxDim >= cooStart ? cooStart : idxDim;
+  Value buffer = getMemRefField(SparseTensorFieldKind::IdxMemRef, idx);
+  if (idxDim >= cooStart) {
+    unsigned rank = enc.getDimLevelType().size();
     Value stride = constantIndex(builder, loc, rank - cooStart);
-    SmallVector<Value> buffersArray(buffers.begin(), buffers.end());
-    MutSparseTensorDescriptor mutDesc(tp, buffersArray);
-    // Calculate subbuffer size as memSizes[idx] / (stride).
-    Value subBufferSize = mutDesc.getIdxMemSize(builder, loc, cooStart);
-    subBufferSize = builder.create<arith::DivUIOp>(loc, subBufferSize, stride);
-
-    // Create views of the linear idx buffer for the COO indices.
-    for (unsigned i = cooStart; i < rank; i++) {
-      Value subBuffer = builder.create<memref::SubViewOp>(
-          loc, buffer,
-          /*offset=*/ValueRange{constantIndex(builder, loc, i - cooStart)},
-          /*size=*/ValueRange{subBufferSize},
-          /*step=*/ValueRange{stride});
-      expandedFields.push_back(subBuffer);
-    }
-    expandedFields.push_back(buffers[buffers.size() - 2]); // The Values memref.
-    expandedFields.push_back(buffers.back());              // The specifier.
-    fields = expandedFields;
-  } else {
-    fields = buffers;
-  }
-
-  sanityCheck();
+    Value size = getIdxMemSize(builder, loc, cooStart);
+    size = builder.create<arith::DivUIOp>(loc, size, stride);
+    buffer = builder.create<memref::SubViewOp>(
+        /*offset=*/ValueRange{constantIndex(builder, loc, idxDim - cooStart)},
+        /*size=*/ValueRange{size},
+  return buffer;
 }
 
 //===----------------------------------------------------------------------===//
 // Public methods.
 //===----------------------------------------------------------------------===//
 
 constexpr uint64_t kDataFieldStartingIdx = 0;
 
 void sparse_tensor::foreachFieldInSparseTensor(
     const SparseTensorEncodingAttr enc,
     llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
                             DimLevelType)>
-        callback,
-    bool isBuffer) {
+        callback) {
   assert(enc);
 
 #define RETURN_ON_FALSE(idx, kind, dim, dlt)                                   \
   if (!(callback(idx, kind, dim, dlt)))                                        \
     return;
 
   unsigned rank = enc.getDimLevelType().size();
-  unsigned cooStart = isBuffer ? getCOOStart(enc) : rank;
+  unsigned end = getCOOStart(enc);
+  if (end != rank)
+    end += 1;
   static_assert(kDataFieldStartingIdx == 0);
   unsigned fieldIdx = kDataFieldStartingIdx;
   // Per-dimension storage.
-  for (unsigned r = 0; r < rank; r++) {
+  for (unsigned r = 0; r < end; 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)) {
-      if (r < cooStart)
-        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);
 
foreachFieldInSparseTensor(
         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);
         };
         llvm_unreachable("unrecognized field kind");
-      },
-      /*isBuffer=*/true);
+      });
 }
 
-unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc,
-                                                 bool isBuffer) {
+unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
   unsigned numFields = 0;
-  foreachFieldInSparseTensor(
-      enc,
-      [&numFields](unsigned, SparseTensorFieldKind, unsigned,
-                   DimLevelType) -> bool {
+  foreachFieldInSparseTensor(enc,
+                             [&numFields](unsigned, SparseTensorFieldKind,
+                                          unsigned, DimLevelType) -> bool {
                                numFields++;
                                return true;
-      },
-      isBuffer);
+                             });
   return numFields;
 }
 
 unsigned
 sparse_tensor::getNumDataFieldsFromEncoding(SparseTensorEncodingAttr enc) {
   unsigned numFields = 0; // one value memref
-  foreachFieldInSparseTensor(
-      enc,
-      [&numFields](unsigned fidx, SparseTensorFieldKind, unsigned,
-                   DimLevelType) -> bool {
+  foreachFieldInSparseTensor(enc,
+                             [&numFields](unsigned fidx, SparseTensorFieldKind,
+                                          unsigned, DimLevelType) -> bool {
                                if (fidx >= kDataFieldStartingIdx)
                                  numFields++;
                                return true;
-      },
-      /*isBuffer=*/true);
+                             });
   numFields -= 1; // the last field is MetaData field
-  assert(numFields == getNumFieldsFromEncoding(enc, /*isBuffer=*/true) -
-                          kDataFieldStartingIdx - 1);
+  assert(numFields ==
+         getNumFieldsFromEncoding(enc) - kDataFieldStartingIdx - 1);
   return numFields;
 }