[mlir][sparse] Add AOS optimization.

[bixia]

Summary

Use an array of structures to represent the indices for the tailing COO region
of a sparse tensor.

[aartbik]

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

266 ↗

(On Diff #485918)

this is the only change in the file, and one I don't want ;-)
Can you please revert this file from the revision

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

10 ↗

(On Diff #485918)

why this change? was lint complaining?

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

125

can't we pass an mutable as an immutable? if so, then it seems this is not a very pure distinction?

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

85

why the selective field/buffer renaming

also note that comments and method names are out of sync now at a lot of places

we use "field" here in the tuple sense, so I think I prefer keeping it that way, unless you have a very strong argument for the renaming?

[bixia]

Removed an unnecessary include; reverted field/buffer renaming.

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

266 ↗

(On Diff #485918)

This line and the routine after is added from the previous PR. There must be an issue with rebasing the source. Fixed.

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

10 ↗

(On Diff #485918)

Removed.

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

125

We can't pass a mutable as an immutable or the other way.
Mutable keeps the linear buffers, immutable needs to explicitly create views of the linear buffers.

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

85

Revert the field/buffer change here and other places.

[aartbik]

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

75

here and below, I would even remove the "Disp"
does not add much information and only makes it longer

116

this class needs a comment now, describing what the template var means

307

either after the statement, or otherwise as a sentence Drop ... fields.

319

why is this so moved around? makes it hard to review?

can we do a prep revision with this, if needed ,first?

338

this changes the simplicity of the original setup quite a bit.

Can't we use template specialization to define the methods?
https://en.cppreference.com/w/cpp/language/template_specialization

[bixia]

Split some change to https://reviews.llvm.org/D141002.
Address review comments.

[bixia]

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

75

Fix here and another place

116

Added description for the class. PTAL.

319

move change to https://reviews.llvm.org/D141002.

338

MutSparseTensorDescriptor now contains the getter/setter that are only meaningful to a mutable descriptor. PTAL.

[bixia]

Rebase.

[bixia]

Clang-format.

[aartbik]

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

304–305

why was there commented out code?

[bixia]

Remove commented out code.

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 layout(enc);
+      StorageLayout<true> 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 std::optional<Value> sizeFromTensorAtDim(OpBuilder &builder,
   // Any other query can consult the dimSizes array at field DimSizesIdx,
   // accounting for the reordering applied to the sparse storage.
   return desc.getDimSize(builder, loc, toStoredDim(rtp, dim));
 }
 
 // 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,
-                      SparseTensorDescriptor desc, unsigned d) {
+                      MutSparseTensorDescriptor 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 desc.getDimSize(builder, loc, d);
 }
static Value genCompressed(OpBuilder &builder, Location loc,
                            MutSparseTensorDescriptor desc,
                            SmallVectorImpl<Value> &indices, Value value,
                            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 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 msz = desc.getIdxMemSize(builder, loc, d);
-  // Value msz = desc.getMemSize(builder, loc, getFieldMemSizesIndex(idxIndex));
-
+  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);
   // 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, desc.getMemRefField(idxIndex), phim1);
+  Value crd = genLoad(
+      builder, loc, desc.getMemRefField(idxIndex),
+      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);
   builder.create<scf::YieldOp>(loc, constantI1(builder, loc, false));
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();
 
-    auto desc = getDescriptorFromTensorTuple(adaptor.getSource());
-    auto sz = sizeFromTensorAtDim(rewriter, op.getLoc(), desc, *index);
+    Location loc = op.getLoc();
+    auto desc =
+        getDescriptorFromTensorTuple(rewriter, loc, adaptor.getSource());
+    auto sz = sizeFromTensorAtDim(rewriter, loc, desc, *index);
 
     if (!sz)
       return failure();
 
     rewriter.replaceOp(op, *sz);
     return success();
   }
 };
public:
   matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto enc = getSparseTensorEncoding(op.getTensor().getType());
     if (!enc)
       return failure();
 
     // Replace the sparse tensor deallocation with field deallocations.
     Location loc = op.getLoc();
-    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    SmallVector<Value> fields;
+    auto desc = getMutDescriptorFromTensorTuple(adaptor.getTensor(), fields);
     for (auto input : desc.getMemRefFields())
       // Deallocate every buffer used to store the sparse tensor handler.
       rewriter.create<memref::DeallocOp>(loc, input);
 
     rewriter.eraseOp(op);
     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(adaptor.getTensor());
+    auto desc =
+        getDescriptorFromTensorTuple(rewriter, loc, 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(adaptor.getTensor());
+    auto desc = getDescriptorFromTensorTuple(rewriter, op.getLoc(),
+                                             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(adaptor.getTensor());
+    auto desc = getDescriptorFromTensorTuple(rewriter, op.getLoc(),
+                                             adaptor.getTensor());
     uint64_t dim = op.getDimension().getZExtValue();
     Value field = desc.getIdxMemRef(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())
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(adaptor.getTensor());
+    auto desc = getDescriptorFromTensorTuple(rewriter, op.getLoc(),
+                                             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

 
 /// For each field 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 on ad-hoc index computation.
 void foreachFieldInSparseTensor(
     SparseTensorEncodingAttr,
     llvm::function_ref<bool(unsigned /*fieldIdx*/,
                             SparseTensorFieldKind /*fieldKind*/,
                             unsigned /*dim (if applicable)*/,
-                            DimLevelType /*DLT (if applicable)*/)>);
+                            DimLevelType /*DLT (if applicable)*/)>,
+    bool isBuffer = false);
 
 /// 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);
+unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc, bool isBuffer);
 
 /// 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>
 class StorageLayout {
 public:
   explicit StorageLayout(SparseTensorEncodingAttr enc) : enc(enc) {}
 
   ///
   /// Getters: get the field index for required field.
   ///
+
   unsigned getMemRefFieldIndex(SparseTensorFieldKind kind,
-                               Optional<unsigned> dim) const;
+                               std::optional<unsigned> dim) const {
+    return getFieldIndexAndStride(kind, dim).first;
+  }
 
   unsigned getMemRefFieldIndex(StorageSpecifierKind kind,
-                               Optional<unsigned> dim) const;
+                               std::optional<unsigned> dim) const {
+    return getMemRefFieldIndex(toFieldKind(kind), dim);
+  }
+
+  static unsigned getNumFieldsFromEncoding(SparseTensorEncodingAttr enc) {
+    return sparse_tensor::getNumFieldsFromEncoding(enc, isBuffer);
+  }
+
+  static void foreachFieldInSparseTensor(
+      const SparseTensorEncodingAttr enc,
+      llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
+                              DimLevelType)>
+          callback) {
+    return sparse_tensor::foreachFieldInSparseTensor(enc, callback, isBuffer);
+  }
+
+  std::pair<unsigned, unsigned>
+  getFieldIndexAndStride(SparseTensorFieldKind kind,
+                         std::optional<unsigned> dim) const {
+    unsigned fieldIdx = -1u;
+    unsigned stride = 1;
+    if (isBuffer && 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;
+      }
+    }
+    foreachFieldInSparseTensor(
+        enc,
+        [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
+                               unsigned fDim, DimLevelType dlt) -> bool {
+          if ((dim && fDim == dim.value() && kind == fKind) ||
+              (kind == fKind && fKind == SparseTensorFieldKind::ValMemRef)) {
+            fieldIdx = fIdx;
+            // Returns false to break the iteration.
+            return false;
+          }
+          return true;
+        });
+    assert(fieldIdx != -1u);
+    return std::pair<unsigned, unsigned>(fieldIdx, stride);
+  }
 
 private:
-  unsigned getFieldIndex(unsigned dim, SparseTensorFieldKind kind) const;
   SparseTensorEncodingAttr enc;
 };
 
 class SparseTensorSpecifier {
 public:
   explicit SparseTensorSpecifier(Value specifier) : specifier(specifier) {}
 
   // Undef value for dimension sizes, all zero value for memory sizes.
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;
 
-public:
+  SparseTensorDescriptorImpl(Type tp)
+      : rType(tp.cast<RankedTensorType>()), fields() {}
+
   SparseTensorDescriptorImpl(Type tp, ValueArrayRef fields)
       : rType(tp.cast<RankedTensorType>()), fields(fields) {
-    assert(getSparseTensorEncoding(tp) &&
-           getNumFieldsFromEncoding(getSparseTensorEncoding(tp)) ==
-               fields.size());
+    sanityCheck();
+  }
+
+  void sanityCheck() {
+    assert(getSparseTensorEncoding(rType) &&
+           StorageLayout<mut>::getNumFieldsFromEncoding(
+               getSparseTensorEncoding(rType)) == 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 (and cheap) type conversion from MutSparseTensorDescriptor to
-  // SparseTensorDescriptor.
-  template <typename T = SparseTensorDescriptorImpl<true>>
-  /*implicit*/ SparseTensorDescriptorImpl(std::enable_if_t<!mut, T> &mDesc)
-      : rType(mDesc.getTensorType()), fields(mDesc.getFields()) {}
-
+public:
   unsigned getMemRefFieldIndex(SparseTensorFieldKind kind,
                                Optional<unsigned> dim) const {
     // Delegates to storage layout.
-    StorageLayout layout(getSparseTensorEncoding(rType));
+    StorageLayout<mut> layout(getSparseTensorEncoding(rType));
     return layout.getMemRefFieldIndex(kind, dim);
   }
 
   unsigned getPtrMemRefIndex(unsigned ptrDim) const {
     return getMemRefFieldIndex(SparseTensorFieldKind::PtrMemRef, ptrDim);
   }
 
   unsigned getIdxMemRefIndex(unsigned idxDim) const {
public:
 
   Value getField(unsigned fidx) const {
     assert(fidx < fields.size());
     return fields[fidx];
   }
 
   ValueRange getMemRefFields() const {
     ValueRange ret = fields;
     // Drop the last metadata fields.
     return ret.slice(0, fields.size() - 1);
   }
 
   Type getMemRefElementType(SparseTensorFieldKind kind,
                             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 getValMemSize(OpBuilder &builder, Location loc) const {
void setSpecifierField(OpBuilder &builder, Location loc,
     md.setSpecifierField(builder, loc, v, kind, dim);
     fields.back() = md;
   }
 
   template <typename T = Value>
   void setDimSize(OpBuilder &builder, Location loc, unsigned dim, T v) {
     setSpecifierField(builder, loc, StorageSpecifierKind::DimSize, dim, v);
   }
+
+  std::pair<unsigned, unsigned>
+  getIdxMemRefIndexAndStride(unsigned idxDim) const {
+    StorageLayout<true> layout(getSparseTensorEncoding(rType));
+    return layout.getFieldIndexAndStride(SparseTensorFieldKind::IdxMemRef,
+                                         idxDim);
+  }
 };
 
-using SparseTensorDescriptor = SparseTensorDescriptorImpl<false>;
+class SparseTensorDescriptor : public SparseTensorDescriptorImpl<false> {
+public:
+  SparseTensorDescriptor(OpBuilder &builder, Location loc, Type tp,
+                         ValueArrayRef buffers);
+
+private:
+  // Store the fields passed to SparseTensorDescriptorImpl when the tensor has
+  // a COO region.
+  SmallVector<Value> expandedFields;
+};
 
 /// 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,
-                      SparseTensorDescriptor desc) {
+                      MutSparseTensorDescriptor desc) {
   return genTuple(builder, loc, desc.getTensorType(), desc.getFields());
 }
 
-inline SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor) {
+inline SparseTensorDescriptor
+getDescriptorFromTensorTuple(OpBuilder &builder, Location loc, Value tensor) {
   auto tuple = getTuple(tensor);
-  return SparseTensorDescriptor(tuple.getResultTypes()[0], tuple.getInputs());
+  return SparseTensorDescriptor(builder, loc, 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

 //===- SparseTensorStorageLayout.cpp --------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "SparseTensorStorageLayout.h"
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace sparse_tensor;
 
 //===----------------------------------------------------------------------===//
if (!getSparseTensorEncoding(tp))
       // Not a sparse tensor.
       return std::nullopt;
     // Sparse compiler knows how to cancel out these casts.
     return genTuple(builder, loc, tp, inputs);
   });
 }
 
 //===----------------------------------------------------------------------===//
-// StorageLayout methods.
-//===----------------------------------------------------------------------===//
-
-unsigned StorageLayout::getMemRefFieldIndex(SparseTensorFieldKind kind,
-                                            std::optional<unsigned> dim) const {
-  unsigned fieldIdx = -1u;
-  foreachFieldInSparseTensor(
-      enc,
-      [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
-                             unsigned fDim, DimLevelType dlt) -> bool {
-        if ((dim && fDim == dim.value() && kind == fKind) ||
-            (kind == fKind && fKind == SparseTensorFieldKind::ValMemRef)) {
-          fieldIdx = fIdx;
-          // Returns false to break the iteration.
-          return false;
-        }
-        return true;
-      });
-  assert(fieldIdx != -1u);
-  return fieldIdx;
-}
-
-unsigned StorageLayout::getMemRefFieldIndex(StorageSpecifierKind kind,
-                                            std::optional<unsigned> dim) const {
-  return getMemRefFieldIndex(toFieldKind(kind), dim);
-}
-
-//===----------------------------------------------------------------------===//
 // StorageTensorSpecifier methods.
 //===----------------------------------------------------------------------===//
 
 Value SparseTensorSpecifier::getInitValue(OpBuilder &builder, Location loc,
                                           RankedTensorType rtp) {
   return builder.create<StorageSpecifierInitOp>(
       loc, StorageSpecifierType::get(getSparseTensorEncoding(rtp)));
 }
void SparseTensorSpecifier::setSpecifierField(OpBuilder &builder, Location loc,
                                               StorageSpecifierKind kind,
                                               std::optional<unsigned> dim) {
   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();
+  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];
+
+    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();
+}
+
+//===----------------------------------------------------------------------===//
 // Public methods.
 //===----------------------------------------------------------------------===//
 
 constexpr uint64_t kDataFieldStartingIdx = 0;
 
 void sparse_tensor::foreachFieldInSparseTensor(
     const SparseTensorEncodingAttr enc,
     llvm::function_ref<bool(unsigned, SparseTensorFieldKind, unsigned,
                             DimLevelType)>
-        callback) {
+        callback,
+    bool isBuffer) {
   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;
   static_assert(kDataFieldStartingIdx == 0);
   unsigned fieldIdx = kDataFieldStartingIdx;
   // Per-dimension storage.
-  for (unsigned r = 0, rank = enc.getDimLevelType().size(); r < rank; r++) {
+  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.
     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) {
+unsigned sparse_tensor::getNumFieldsFromEncoding(SparseTensorEncodingAttr enc,
+                                                 bool isBuffer) {
   unsigned numFields = 0;
-  foreachFieldInSparseTensor(enc,
-                             [&numFields](unsigned, SparseTensorFieldKind,
-                                          unsigned, DimLevelType) -> bool {
-                               numFields++;
-                               return true;
-                             });
+  foreachFieldInSparseTensor(
+      enc,
+      [&numFields](unsigned, SparseTensorFieldKind, unsigned,
+                   DimLevelType) -> bool {
+      },
+      isBuffer);
   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 >= kDataFieldStartingIdx)
-                                 numFields++;
-                               return true;
-                             });
+  foreachFieldInSparseTensor(
+      enc,
+      [&numFields](unsigned fidx, SparseTensorFieldKind, unsigned,
+                   DimLevelType) -> bool {
+      },
+      /*isBuffer=*/true);
   numFields -= 1; // the last field is MetaData field
-  assert(numFields ==
-         getNumFieldsFromEncoding(enc) - kDataFieldStartingIdx - 1);
+  assert(numFields == getNumFieldsFromEncoding(enc, /*isBuffer=*/true) -
+                          kDataFieldStartingIdx - 1);
   return numFields;
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

#Dense3D = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense", "dense" ],
   dimOrdering = affine_map<(i, j, k) -> (k, i, j)>
 }>
 
 #Coo = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ]
 }>
 
+#ccoo = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed-nu", "singleton" ]
+}>
+
 // CHECK-LABEL: func @sparse_nop(
 //  CHECK-SAME: %[[A0:.*]]: memref<?xi32>,
 //  CHECK-SAME: %[[A1:.*]]: memref<?xi64>,
 //  CHECK-SAME: %[[A2:.*]]: memref<?xf64>,
 //  CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] :
 //  CHECK-SAME:   memref<?xi32>, memref<?xi64>, memref<?xf64>, !sparse_tensor.storage_specifier
 func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {

 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
 //       CHECK: return %[[A4]] : memref<?xf64>
 func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
+// CHECK-LABEL: func.func @sparse_values_coo(
+//  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
+//       CHECK: return %[[A4]] : memref<?xf64>
+func.func @sparse_values_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xf64> {
+  %0 = sparse_tensor.values %arg0 : tensor<?x?x?xf64, #ccoo> to memref<?xf64>
+  return %0 : memref<?xf64>
+}
+
+
+// CHECK-LABEL:   func.func @sparse_indices_coo(
+//  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
+//       CHECK: %[[C2:.*]] = arith.constant 2 : index
+//       CHECK: %[[S0:.*]] = sparse_tensor.storage_specifier.get %[[A5]]  idx_mem_sz at 1
+//       CHECK: %[[S1:.*]] = arith.index_cast %[[S0]]
+//       CHECK: %[[S2:.*]] = arith.divui %[[S1]], %[[C2]] : index
+//       CHECK: %[[R1:.*]] = memref.subview %[[A3]][0] {{\[}}%[[S2]]] [2] : memref<?xindex> to memref<?xindex, strided<[2]>>
+//       CHECK: %[[R2:.*]] = memref.cast %[[R1]] : memref<?xindex, strided<[2]>> to memref<?xindex, strided<[?], offset: ?>>
+//       CHECK: return %[[R2]] : memref<?xindex, strided<[?], offset: ?>>
+func.func @sparse_indices_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xindex, strided<[?], offset: ?>> {
+  %0 = sparse_tensor.indices  %arg0 { dimension = 1 : index } : tensor<?x?x?xf64, #ccoo> to memref<?xindex, strided<[?], offset: ?>>
+  return %0 : memref<?xindex, strided<[?], offset: ?>>
+}
+
 // CHECK-LABEL: func @sparse_noe(
 //  CHECK-SAME: %[[A0:.*]]: memref<?xi32>,
 //  CHECK-SAME: %[[A1:.*]]: memref<?xi64>,
 //  CHECK-SAME: %[[A2:.*]]: memref<?xf64>,
 //  CHECK-SAME: %[[A3:.*]]: !sparse_tensor.storage_specifier
 //       CHECK: %[[A4:.*]] = sparse_tensor.storage_specifier.get %[[A3]] val_mem_sz
 //       CHECK: %[[NOE:.*]] = arith.index_cast %[[A4]] : i64 to index
 //       CHECK: return %[[NOE]] : index

 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3
 func.func @sparse_insert_typed(%arg0: tensor<128xf64, #SparseVector>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SparseVector>
   %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SparseVector>
   return %1 : tensor<128xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func.func private @"_insert_compressed-nu_singleton_5_6_f64_0_0"(
-//  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
-//  CHECK-SAME: %[[A4:.*4]]: !sparse_tensor.storage_specifier
+//  CHECK-SAME: %[[A1:.*0]]: memref<?xindex>,
+//  CHECK-SAME: %[[A2:.*1]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*2]]: memref<?xf64>,
+//  CHECK-SAME: %[[A4:.*3]]: !sparse_tensor.storage_specifier
+//  CHECK-SAME: %[[A5:.*4]]: index,
 //  CHECK-SAME: %[[A5:.*5]]: index,
-//  CHECK-SAME: %[[A5:.*6]]: index,
-//  CHECK-SAME: %[[A7:.*7]]: f64)
+//  CHECK-SAME: %[[A7:.*6]]: f64)
 //
 // CHECK-LABEL: func.func @sparse_insert_coo(
 //  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
 //  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
-//  CHECK-SAME: %[[A4:.*4]]: !sparse_tensor.storage_specifier
-//  CHECK-SAME: %[[A5:.*5]]: index,
-//  CHECK-SAME: %[[A6:.*6]]: f64)
-//       CHECK: %[[R:.*]]:5 = call @"_insert_compressed-nu_singleton_5_6_f64_0_0"(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A5]], %[[A6]])
-//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>,
+//  CHECK-SAME: %[[A3:.*3]]: !sparse_tensor.storage_specifier
+//  CHECK-SAME: %[[A4:.*4]]: index,
+//  CHECK-SAME: %[[A5:.*5]]: f64)
+//       CHECK: %[[R:.*]]:4 = call @"_insert_compressed-nu_singleton_5_6_f64_0_0"(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A4]], %[[A5]])
+//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3
 func.func @sparse_insert_coo(%arg0: tensor<5x6xf64, #Coo>, %arg1: index, %arg2: f64) -> tensor<5x6xf64, #Coo> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1, %arg1] : tensor<5x6xf64, #Coo>
   %1 = sparse_tensor.load %0 hasInserts : tensor<5x6xf64, #Coo>
   return %1 : tensor<5x6xf64, #Coo>
 }
 
 // CHECK-LABEL: func.func @sparse_nop_convert(
 //  CHECK-SAME: %[[A1:.*]]: memref<?xi32>,

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir

func.func @dump_dense(%arg0: tensor<4x3xf64, #Dense>) {
     return
   }
 
   func.func @dump_coo(%arg0: tensor<4x3xf64, #SortedCOO>) {
     %c0 = arith.constant 0 : index
     %cu = arith.constant -1 : index
     %fu = arith.constant 99.0 : f64
     %p0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex>
-    %i0 = sparse_tensor.indices  %arg0 { dimension = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex>
-    %i1 = sparse_tensor.indices  %arg0 { dimension = 1 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex>
+    %i0 = sparse_tensor.indices  %arg0 { dimension = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
+    %i1 = sparse_tensor.indices  %arg0 { dimension = 1 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %v = sparse_tensor.values %arg0 : tensor<4x3xf64, #SortedCOO> to memref<?xf64>
     %vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
     vector.print %vp0 : vector<2xindex>
-    %vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex>, vector<4xindex>
+    %vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex, strided<[?], offset: ?>>, vector<4xindex>
     vector.print %vi0 : vector<4xindex>
-    %vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex>, vector<4xindex>
+    %vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex, strided<[?], offset: ?>>, vector<4xindex>
     vector.print %vi1 : vector<4xindex>
     %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
     vector.print %vv : vector<4xf64>
     return
   }
 
   func.func @dump_csr(%arg0: tensor<4x3xf64, #CSR>) {
     %c0 = arith.constant 0 : index

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir

 #TensorCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "dense", "compressed" ]
 }>
 
 #TensorRow = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "dense" ]
 }>
 
+#CCoo = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed-nu", "singleton" ]
+}>
+
+#DCoo = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed-nu", "singleton" ]
+}>
+
+
 module {
 
   func.func @dump(%arg0: tensor<5x4x3xf64, #TensorCSR>) {
     %c0 = arith.constant 0 : index
     %fu = arith.constant 99.0 : f64
     %p0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
     %i0 = sparse_tensor.indices  %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
     %p2 = sparse_tensor.pointers %arg0 { dimension = 2 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
func.func @dump_row(%arg0: tensor<5x4x3xf64, #TensorRow>) {
     vector.print %vp1 : vector<3xindex>
     %vi1 = vector.transfer_read %i1[%c0], %c0: memref<?xindex>, vector<4xindex>
     vector.print %vi1 : vector<4xindex>
     %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<12xf64>
     vector.print %vv : vector<12xf64>
     return
   }
 
+func.func @dump_ccoo(%arg0: tensor<5x4x3xf64, #CCoo>) {
+    %c0 = arith.constant 0 : index
+    %fu = arith.constant 99.0 : f64
+    %p0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #CCoo> to memref<?xindex>
+    %i0 = sparse_tensor.indices  %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #CCoo> to memref<?xindex>
+    %p1 = sparse_tensor.pointers %arg0 { dimension = 1 : index } : tensor<5x4x3xf64, #CCoo> to memref<?xindex>
+    %i1 = sparse_tensor.indices  %arg0 { dimension = 1 : index } : tensor<5x4x3xf64, #CCoo> to memref<?xindex>
+    %i2 = sparse_tensor.indices  %arg0 { dimension = 2 : index } : tensor<5x4x3xf64, #CCoo> to memref<?xindex>
+    %v = sparse_tensor.values %arg0 : tensor<5x4x3xf64, #CCoo> to memref<?xf64>
+    %vp0 = vector.transfer_read %p0[%c0], %c0: memref<?xindex>, vector<2xindex>
+    vector.print %vp0 : vector<2xindex>
+    %vi0 = vector.transfer_read %i0[%c0], %c0: memref<?xindex>, vector<2xindex>
+    vector.print %vi0 : vector<2xindex>
+    %vp1 = vector.transfer_read %p1[%c0], %c0: memref<?xindex>, vector<3xindex>
+    vector.print %vp1 : vector<3xindex>
+    %vi1 = vector.transfer_read %i1[%c0], %c0: memref<?xindex>, vector<5xindex>
+    vector.print %vi1 : vector<5xindex>
+    %vi2 = vector.transfer_read %i2[%c0], %c0: memref<?xindex>, vector<5xindex>
+    vector.print %vi2 : vector<5xindex>
+    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<5xf64>
+    vector.print %vv : vector<5xf64>
+    return
+  }
+
+func.func @dump_dcoo(%arg0: tensor<5x4x3xf64, #DCoo>) {
+    %c0 = arith.constant 0 : index
+    %fu = arith.constant 99.0 : f64
+    %p1 = sparse_tensor.pointers %arg0 { dimension = 1 : index } : tensor<5x4x3xf64, #DCoo> to memref<?xindex>
+    %i1 = sparse_tensor.indices  %arg0 { dimension = 1 : index } : tensor<5x4x3xf64, #DCoo> to memref<?xindex>
+    %i2 = sparse_tensor.indices  %arg0 { dimension = 2 : index } : tensor<5x4x3xf64, #DCoo> to memref<?xindex>
+    %v = sparse_tensor.values %arg0 : tensor<5x4x3xf64, #DCoo> to memref<?xf64>
+    %vp1 = vector.transfer_read %p1[%c0], %c0: memref<?xindex>, vector<6xindex>
+    vector.print %vp1 : vector<6xindex>
+    %vi1 = vector.transfer_read %i1[%c0], %c0: memref<?xindex>, vector<5xindex>
+    vector.print %vi1 : vector<5xindex>
+    %vi2 = vector.transfer_read %i2[%c0], %c0: memref<?xindex>, vector<5xindex>
+    vector.print %vi2 : vector<5xindex>
+    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<5xf64>
+    vector.print %vv : vector<5xf64>
+    return
+}
+
   //
   // Main driver.
   //
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
     %c2 = arith.constant 2 : index
     %c3 = arith.constant 3 : index
func.func @entry() {
     %row1 = sparse_tensor.insert %f1 into %rowa[%c3, %c0, %c1] : tensor<5x4x3xf64, #TensorRow>
     %row2 = sparse_tensor.insert %f2 into %row1[%c3, %c0, %c2] : tensor<5x4x3xf64, #TensorRow>
     %row3 = sparse_tensor.insert %f3 into %row2[%c3, %c3, %c1] : tensor<5x4x3xf64, #TensorRow>
     %row4 = sparse_tensor.insert %f4 into %row3[%c4, %c2, %c2] : tensor<5x4x3xf64, #TensorRow>
     %row5 = sparse_tensor.insert %f5 into %row4[%c4, %c3, %c2] : tensor<5x4x3xf64, #TensorRow>
     %rowm = sparse_tensor.load %row5 hasInserts : tensor<5x4x3xf64, #TensorRow>
     call @dump_row(%rowm) : (tensor<5x4x3xf64, #TensorRow>) -> ()
 
+    //
+    // CHECK: ( 0, 2 )
+    // CHECK-NEXT: ( 3, 4 )
+    // CHECK-NEXT: ( 0, 3, 5 )
+    // CHECK-NEXT: ( 0, 0, 3, 2, 3 )
+    // CHECK-NEXT: ( 1, 2, 1, 2, 2 )
+    // CHECK-NEXT: ( 1.1, 2.2, 3.3, 4.4, 5.5 )
+    %ccoo = bufferization.alloc_tensor() : tensor<5x4x3xf64, #CCoo>
+    %ccoo1 = sparse_tensor.insert %f1 into %ccoo[%c3, %c0, %c1] : tensor<5x4x3xf64, #CCoo>
+    %ccoo2 = sparse_tensor.insert %f2 into %ccoo1[%c3, %c0, %c2] : tensor<5x4x3xf64, #CCoo>
+    %ccoo3 = sparse_tensor.insert %f3 into %ccoo2[%c3, %c3, %c1] : tensor<5x4x3xf64, #CCoo>
+    %ccoo4 = sparse_tensor.insert %f4 into %ccoo3[%c4, %c2, %c2] : tensor<5x4x3xf64, #CCoo>
+    %ccoo5 = sparse_tensor.insert %f5 into %ccoo4[%c4, %c3, %c2] : tensor<5x4x3xf64, #CCoo>
+    %ccoom = sparse_tensor.load %ccoo5 hasInserts : tensor<5x4x3xf64, #CCoo>
+    call @dump_ccoo(%ccoom) : (tensor<5x4x3xf64, #CCoo>) -> ()
+
+    //
+    // CHECK-NEXT: ( 0, 0, 0, 0, 3, 5 )
+    // CHECK-NEXT: ( 0, 0, 3, 2, 3 )
+    // CHECK-NEXT: ( 1, 2, 1, 2, 2 )
+    // CHECK-NEXT: ( 1.1, 2.2, 3.3, 4.4, 5.5 )
+    %dcoo = bufferization.alloc_tensor() : tensor<5x4x3xf64, #DCoo>
+    %dcoo1 = sparse_tensor.insert %f1 into %dcoo[%c3, %c0, %c1] : tensor<5x4x3xf64, #DCoo>
+    %dcoo2 = sparse_tensor.insert %f2 into %dcoo1[%c3, %c0, %c2] : tensor<5x4x3xf64, #DCoo>
+    %dcoo3 = sparse_tensor.insert %f3 into %dcoo2[%c3, %c3, %c1] : tensor<5x4x3xf64, #DCoo>
+    %dcoo4 = sparse_tensor.insert %f4 into %dcoo3[%c4, %c2, %c2] : tensor<5x4x3xf64, #DCoo>
+    %dcoo5 = sparse_tensor.insert %f5 into %dcoo4[%c4, %c3, %c2] : tensor<5x4x3xf64, #DCoo>
+    %dcoom = sparse_tensor.load %dcoo5 hasInserts : tensor<5x4x3xf64, #DCoo>
+    call @dump_dcoo(%dcoom) : (tensor<5x4x3xf64, #DCoo>) -> ()
+
     // NOE sanity check.
     //
     // CHECK-NEXT: 5
     // CHECK-NEXT: 12
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 5
     //
     %noe1 = sparse_tensor.number_of_entries %tensorm : tensor<5x4x3xf64, #TensorCSR>
     vector.print %noe1 : index
     %noe2 = sparse_tensor.number_of_entries %rowm : tensor<5x4x3xf64, #TensorRow>
     vector.print %noe2 : index
+    %noe3 = sparse_tensor.number_of_entries %ccoom : tensor<5x4x3xf64, #CCoo>
+    vector.print %noe3 : index
+    %noe4 = sparse_tensor.number_of_entries %dcoom : tensor<5x4x3xf64, #DCoo>
+    vector.print %noe4 : index
 
     // Release resources.
     bufferization.dealloc_tensor %tensorm : tensor<5x4x3xf64, #TensorCSR>
     bufferization.dealloc_tensor %rowm : tensor<5x4x3xf64, #TensorRow>
+    bufferization.dealloc_tensor %ccoom : tensor<5x4x3xf64, #CCoo>
+    bufferization.dealloc_tensor %dcoom : tensor<5x4x3xf64, #DCoo>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sorted_coo.mlir

module {
 
   func.func @dumpi(%arg0: memref<?xindex>) {
     %c0 = arith.constant 0 : index
     %v = vector.transfer_read %arg0[%c0], %c0: memref<?xindex>, vector<20xindex>
     vector.print %v : vector<20xindex>
     return
   }
 
+  func.func @dumpsi(%arg0: memref<?xindex, strided<[?], offset: ?>>) {
+    %c0 = arith.constant 0 : index
+    %v = vector.transfer_read %arg0[%c0], %c0: memref<?xindex, strided<[?], offset: ?>>, vector<20xindex>
+    vector.print %v : vector<20xindex>
+    return
+  }
+
   func.func @dumpf(%arg0: memref<?xf64>) {
     %c0 = arith.constant 0 : index
     %nan = arith.constant 0x0 : f64
     %v = vector.transfer_read %arg0[%c0], %nan: memref<?xf64>, vector<20xf64>
     vector.print %v : vector<20xf64>
     return
   }
 
func.func @entry() {
     // CHECK:      ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 0, 0, 1, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 126, 127, 254, 1, 253, 2, 0, 1, 3, 98, 126, 127, 128, 249, 253, 255, 0, 0, 0 )
     // CHECK-NEXT: ( -1, 2, -3, 4, -5, 6, -7, 8, -9, 10, -11, 12, -13, 14, -15, 16, -17, 0, 0, 0 )
     //
     %p0 = sparse_tensor.pointers %0 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i00 = sparse_tensor.indices %0 { dimension = 0 : index }
-      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %i01 = sparse_tensor.indices %0 { dimension = 1 : index }
-      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %v0 = sparse_tensor.values %0
       : tensor<?x?xf64, #SortedCOO> to memref<?xf64>
     call @dumpi(%p0)  : (memref<?xindex>) -> ()
-    call @dumpi(%i00) : (memref<?xindex>) -> ()
-    call @dumpi(%i01) : (memref<?xindex>) -> ()
+    call @dumpsi(%i00) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i01) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
     call @dumpf(%v0)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 1, 1, 2, 3, 98, 126, 126, 127, 127, 128, 249, 253, 253, 254, 255, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 3, 1, 3, 2, 3, 3, 0, 3, 0, 3, 3, 3, 1, 3, 0, 3, 0, 0, 0 )
     // CHECK-NEXT: ( -1, 8, -5, -9, -7, 10, -11, 2, 12, -3, -13, 14, -15, 6, 16, 4, -17, 0, 0, 0 )
     //
     %p1 = sparse_tensor.pointers %1 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
     %i10 = sparse_tensor.indices %1 { dimension = 0 : index }
-      : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex, strided<[?], offset: ?>>
     %i11 = sparse_tensor.indices %1 { dimension = 1 : index }
-      : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex, strided<[?], offset: ?>>
     %v1 = sparse_tensor.values %1
       : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xf64>
     call @dumpi(%p1)  : (memref<?xindex>) -> ()
-    call @dumpi(%i10) : (memref<?xindex>) -> ()
-    call @dumpi(%i11) : (memref<?xindex>) -> ()
+    call @dumpsi(%i10) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i11) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
     call @dumpf(%v1)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 1, 1, 2, 2, 2, 2, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 1, 1, 2, 2, 2, 2, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 0 )
     // CHECK-NEXT: ( 3, 63, 11, 100, 66, 61, 13, 43, 77, 10, 46, 61, 53, 3, 75, 22, 18, 0, 0, 0 )
     //
     %p2 = sparse_tensor.pointers %2 { dimension = 0 : index }
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
     %i20 = sparse_tensor.indices %2 { dimension = 0 : index }
-      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
+      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex, strided<[?], offset: ?>>
     %i21 = sparse_tensor.indices %2 { dimension = 1 : index }
-      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
+      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex, strided<[?], offset: ?>>
     %i22 = sparse_tensor.indices %2 { dimension = 2 : index }
-      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
+      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex, strided<[?], offset: ?>>
     %v2 = sparse_tensor.values %2
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xf64>
     call @dumpi(%p2)  : (memref<?xindex>) -> ()
-    call @dumpi(%i20) : (memref<?xindex>) -> ()
-    call @dumpi(%i21) : (memref<?xindex>) -> ()
-    call @dumpi(%i21) : (memref<?xindex>) -> ()
+    call @dumpsi(%i20) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i21) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i21) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
     call @dumpf(%v2)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0 )
     // CHECK-NEXT: ( 66, 77, 61, 11, 61, 53, 22, 3, 100, 13, 10, 3, 18, 63, 43, 46, 75, 0, 0, 0 )
     //
     %p3 = sparse_tensor.pointers %3 { dimension = 0 : index }
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
     %i30 = sparse_tensor.indices %3 { dimension = 0 : index }
-      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
+      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex, strided<[?], offset: ?>>
     %i31 = sparse_tensor.indices %3 { dimension = 1 : index }
-      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
+      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex, strided<[?], offset: ?>>
     %i32 = sparse_tensor.indices %3 { dimension = 2 : index }
-      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
+      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex, strided<[?], offset: ?>>
     %v3 = sparse_tensor.values %3
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xf64>
     call @dumpi(%p3)  : (memref<?xindex>) -> ()
-    call @dumpi(%i30) : (memref<?xindex>) -> ()
-    call @dumpi(%i31) : (memref<?xindex>) -> ()
-    call @dumpi(%i31) : (memref<?xindex>) -> ()
+    call @dumpsi(%i30) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i31) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i31) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
     call @dumpf(%v3)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 2, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 3, 0, 3, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
    // CHECK-NEXT: ( 6, 5, 4, 3, 2, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     //
     %p4 = sparse_tensor.pointers %4 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i40 = sparse_tensor.indices %4 { dimension = 0 : index }
-      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %i41 = sparse_tensor.indices %4 { dimension = 1 : index }
-      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %v4 = sparse_tensor.values %4
       : tensor<?x?xf64, #SortedCOO> to memref<?xf64>
     call @dumpi(%p4)  : (memref<?xindex>) -> ()
-    call @dumpi(%i40) : (memref<?xindex>) -> ()
-    call @dumpi(%i41) : (memref<?xindex>) -> ()
+    call @dumpsi(%i40) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i41) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
     call @dumpf(%v4)  : (memref<?xf64>) -> ()
 
     // And last but not least, an actual operation applied to COO.
     // Note that this performs the operation "in place".
     %5 = call @sparse_scale(%4) : (tensor<?x?xf64, #SortedCOO>) -> tensor<?x?xf64, #SortedCOO>
 
     //
     // CHECK-NEXT: ( 0, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 2, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 3, 0, 3, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 12, 10, 8, 6, 4, 22, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     //
     %p5 = sparse_tensor.pointers %5 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i50 = sparse_tensor.indices %5 { dimension = 0 : index }
-      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %i51 = sparse_tensor.indices %5 { dimension = 1 : index }
-      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
     %v5 = sparse_tensor.values %5
       : tensor<?x?xf64, #SortedCOO> to memref<?xf64>
     call @dumpi(%p5)  : (memref<?xindex>) -> ()
-    call @dumpi(%i50) : (memref<?xindex>) -> ()
-    call @dumpi(%i51) : (memref<?xindex>) -> ()
+    call @dumpsi(%i50) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
+    call @dumpsi(%i51) : (memref<?xindex, strided<[?], offset: ?>>) -> ()
     call @dumpf(%v5)  : (memref<?xf64>) -> ()
 
     // Release the resources.
     bufferization.dealloc_tensor %0 : tensor<?x?xf64, #SortedCOO>
     bufferization.dealloc_tensor %1 : tensor<?x?xf64, #SortedCOOPermuted>
     bufferization.dealloc_tensor %2 : tensor<?x?x?xf64, #SortedCOO3D>
     bufferization.dealloc_tensor %3 : tensor<?x?x?xf64, #SortedCOO3DPermuted>
     bufferization.dealloc_tensor %4 : tensor<?x?xf64, #SortedCOO>
 
     return
   }
 }