[mlir][sparse] implement singleton dimension level type

[aartbik]

Summary

This is a first step towards fully implementing the new dimension
level types and properties, illustrating with a fully functional
sorted COO of any dimension. Note that the sparsification part is
pretty complete. The required parts in the runtime support library
have been kept to a minimum, to avoid huge conflicts with Wren's
ongoing refactoring. The missing parts will be filled in later.

[Peiming]

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

1390–1391

How about having a merger.isSparseLevel utils?

[Peiming]

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

1327–1328

What if you have lo = 0, hi = 1 here for singleton? (to allow the loop unroll opt kicking in later?)

1369

And you do not overwrite the value here. (the pidxs for singleton will become loop invariant, but it is okay because there is no loop actually).

[aartbik]

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

1327–1328

Yeah that would be a good way to get the current loop unroller working.

But I also feel we would break the regularity of codegen in sparsification just for the sake of working around a limitation in unrolling. Besides, I have put one of my best people on the task of improving unrolling ;-)

1369

Note that the "loop" could co-iterates with another, and then it remains. I have a CHECK example for that ready in the next revision.

1390–1391

Yeah, good idea, I can also use that for the scan one.
(do you insist of having it in this revision?)

[Peiming]

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

1221

I am wondering which one of the followings should we use?

(d0) -> (d0 + 1)
or
()[s0]->(s0 + 1)

[aartbik]

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

1221

Waiting to see what feedback you get on https://reviews.llvm.org/D134204

[Peiming]

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

1390–1391

No, not really.

[aartbik]

rebased with Peiming's unrolling: single loop for COO now!

[aartbik]

edits

[aartbik]

Holding off landing this revision until Wren's chain of refactoring lands, and I will rebase this then before submitting.

[aartbik]

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

1221

Note that I reverted this again to addi over affine apply

[aartbik]

rebased with Peiming's recent change (affine apply -> addi)

[aartbik]

addi

[aartbik]

rebase with main

[aartbik]

rebased with main

[aartbik]

rebased with main

[aartbik]

rebased against main

[aartbik]

rebased with main
(note that all changes in support lib for singleton
disappeared since Wren already prefetched them into
the refactored support lib revisions)

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

for (unsigned d = 1, rank = map.getNumResults(); d < rank; d++) {
       addAffineOrderings(adjM, inDegree, f, t, 0);
     }
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     if (mask & SortMask::kIncludeUndef) {
       unsigned tensor = t->getOperandNumber();
       for (unsigned i = 0; i < n; i++)
         if (merger.isDimLevelType(tensor, i, DimLvlType::kCompressed) ||
-            merger.isDimLevelType(tensor, i, DimLvlType::kSingleton))
+            merger.isDimLevelType(tensor, i, DimLvlType::kSingleton)) {
           for (unsigned j = 0; j < n; j++)
             if (merger.isDimLevelType(tensor, j, DimLvlType::kUndef)) {
               adjM[i][j] = true;
               inDegree[j]++;
             }
+        } else {
+          assert(merger.isDimLevelType(tensor, i, DimLvlType::kDense) ||
+                 merger.isDimLevelType(tensor, i, DimLvlType::kUndef));
+        }
     }
   }
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
   topSort.clear();
   topSort.reserve(n);
   return topSortOptimal(n, iteratorTypes, topSort, inDegree, adjM);
 }
static bool isAdmissableTensorExp(Merger &merger, linalg::GenericOp op,
   bool allDense = true;
   auto iteratorTypes = op.iterator_types().getValue();
   unsigned numLoops = iteratorTypes.size();
   for (unsigned i = 0; i < numLoops; i++)
     if (merger.isDimLevelType(tensor, i, DimLvlType::kCompressed) ||
         merger.isDimLevelType(tensor, i, DimLvlType::kSingleton)) {
       allDense = false;
       break;
+    } else {
+      assert(merger.isDimLevelType(tensor, i, DimLvlType::kDense) ||
+             merger.isDimLevelType(tensor, i, DimLvlType::kUndef));
     }
   if (allDense)
     return true;
   // A tensor expression with a sparse output tensor that changes its values
   // but not its nonzero structure, an operation called "simply dynamic" in
   // [Bik96,Ch9], is also admissable without special codegen.
   if (merger.isSingleCondition(tensor, exp))
     return true;
static Value genOutputBuffer(CodeGen &codegen, OpBuilder &builder,
   if (!isInit) {
     Value zero = constantZero(builder, loc, denseTp.getElementType());
     builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{init});
   }
   return init;
 }
 
 /// Local bufferization of all dense and sparse data structures.
-/// This code enables testing the first prototype sparse compiler.
-// TODO: replace this with a proliferated bufferization strategy
 static void genBuffers(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                        linalg::GenericOp op) {
   Location loc = op.getLoc();
   assert(op.getNumInputsAndOutputs() == op.getNumInputs() + 1);
   // For every tensor, find lower and upper bound on dimensions, set the
   // same bounds on loop indices, and obtain dense or sparse buffer(s).
+  auto dynShape = {ShapedType::kDynamicSize};
   SmallVector<Value, 4> args;
   for (OpOperand *t : op.getInputAndOutputOperands()) {
     unsigned tensor = t->getOperandNumber();
     auto shape = op.getShape(t);
     auto map = op.getTiedIndexingMap(t);
     auto enc = getSparseTensorEncoding(t->get().getType());
     // Scan all dimensions of current tensor.
     args.clear();
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
       if (a.getKind() != AffineExprKind::DimId)
         continue; // compound
       unsigned idx = a.cast<AffineDimExpr>().getPosition();
-      // Handle sparse storage schemes.
+      // Handle the different storage schemes.
       if (merger.isDimLevelType(tensor, idx, DimLvlType::kCompressed)) {
-        auto dynShape = {ShapedType::kDynamicSize};
+        // Compressed dimension, fetch pointer and indices.
         auto ptrTp =
             MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
-        // Generate sparse primitives to obtains pointer and indices.
         codegen.pointers[tensor][idx] =
             builder.create<ToPointersOp>(loc, ptrTp, t->get(), dim);
         codegen.indices[tensor][idx] =
             builder.create<ToIndicesOp>(loc, indTp, t->get(), dim);
       } else if (merger.isDimLevelType(tensor, idx, DimLvlType::kSingleton)) {
-        llvm_unreachable("TODO: not implemented yet");
+        // Singleton dimension, fetch indices.
+        auto indTp =
+            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
+        auto dim = builder.getIndexAttr(d);
+        codegen.indices[tensor][idx] =
+            builder.create<ToIndicesOp>(loc, indTp, t->get(), dim);
+      } else {
+        // Dense dimension, nothing to fetch.
+        assert(merger.isDimLevelType(tensor, idx, DimLvlType::kDense));
       }
       // Find upper bound in current dimension.
       unsigned p = perm(enc, d);
       Value up = linalg::createOrFoldDimOp(builder, loc, t->get(), p);
       if (ShapedType::isDynamic(shape[p]))
         args.push_back(up);
       assert(codegen.highs[tensor][idx] == nullptr);
       codegen.sizes[idx] = codegen.highs[tensor][idx] = up;
if (!enc) {
         codegen.buffers[tensor] =
             builder.create<bufferization::ToMemrefOp>(loc, denseTp, t->get());
       else
         codegen.buffers[tensor] =
             genOutputBuffer(codegen, builder, op, denseTp, args);
     } else if (t == codegen.sparseOut) {
       // True sparse output needs a lexIdx array.
       Value rank = constantIndex(builder, loc, op.getRank(t));
-      auto dynShape = {ShapedType::kDynamicSize};
       auto memTp = MemRefType::get(dynShape, builder.getIndexType());
       codegen.lexIdx = builder.create<memref::AllocaOp>(loc, memTp, rank);
       codegen.lexVal = builder.create<memref::AllocaOp>(
           loc, MemRefType::get({}, elementType));
     } else {
       // Annotated sparse tensors.
-      auto dynShape = {ShapedType::kDynamicSize};
       auto sparseTp = MemRefType::get(dynShape, elementType);
       codegen.buffers[tensor] =
           builder.create<ToValuesOp>(loc, sparseTp, t->get());
     }
   }
 }
 
 /// Constructs vector type.
static bool genInit(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned at, BitVector &inits) {
   bool needsUniv = false;
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse positions.
   for (unsigned b = 0, be = inits.size(); b < be; b++) {
-    if (inits[b]) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      if (merger.isDimLevelType(b, DimLvlType::kCompressed)) {
-        // Initialize sparse index.
-        unsigned pat = at;
-        for (; pat != 0; pat--) {
-          if (codegen.pidxs[tensor][topSort[pat - 1]])
-            break;
-        }
-        Value ptr = codegen.pointers[tensor][idx];
-        Value one = constantIndex(builder, loc, 1);
-        Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
-                              : codegen.pidxs[tensor][topSort[pat - 1]];
-        codegen.pidxs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p0);
-        Value p1 = builder.create<arith::AddIOp>(loc, p0, one);
-        codegen.highs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p1);
-      } else if (merger.isDimLevelType(b, DimLvlType::kSingleton)) {
-        llvm_unreachable("TODO: not implemented yet");
-      } else {
-        // Dense index still in play.
-        needsUniv = true;
-      }
-    }
-  }
+    if (!inits[b])
+      continue;
+      // Initialize sparse index that will implement the iteration:
+      //   for pidx_idx = pointers(pidx_idx-1), pointers(1+pidx_idx-1)
+      // Initialize sparse index that will implement the "iteration":
+      //   for pidx_idx = pidx_idx-1, 1+pidx_idx-1
+      // We rely on subsequent loop unrolling to get rid of the loop
+      // if it is not involved in co-iteration with anything else.
+      unsigned pat = at;
+      for (; pat != 0; pat--) {
+        if (codegen.pidxs[tensor][topSort[pat - 1]])
+          break;
+      }
+      Value one = constantIndex(builder, loc, 1);
+      Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
+                            : codegen.pidxs[tensor][topSort[pat - 1]];
+      codegen.pidxs[tensor][idx] = p0;
+      codegen.highs[tensor][idx] = builder.create<arith::AddIOp>(loc, p0, one);
+      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
+             merger.isDimLevelType(b, DimLvlType::kUndef));
 
   // Initialize the universal dense index.
   codegen.loops[idx] = constantIndex(builder, loc, 0);
   return needsUniv;
 }
 
 /// Returns vectorization strategy. Any implicit inner loop in the Linalg
 /// operation is a candidate. Whether it is actually converted to SIMD code

 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
                          unsigned idx, BitVector &indices) {
   unsigned fb = indices.find_first();
   unsigned tensor = merger.tensor(fb);
   assert(idx == merger.index(fb));
   auto iteratorTypes = op.iterator_types().getValue();
   bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
-  bool isSparse = merger.isDimLevelType(fb, DimLvlType::kCompressed);
+  bool isSparse = merger.isDimLevelType(fb, DimLvlType::kCompressed) ||
+                  merger.isDimLevelType(fb, DimLvlType::kSingleton);
   bool isVector = isVectorFor(codegen, isInner, isReduction, isSparse) &&
                   denseUnitStrides(merger, op, idx);
   bool isParallel =
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
-  assert(!merger.isDimLevelType(fb, DimLvlType::kSingleton) &&
-         "TODO: implement");
-
   // Prepare vector length.
   if (isVector)
     codegen.curVecLength = codegen.options.vectorLength;
 
   // Loop bounds and increment.
   Location loc = op.getLoc();
   Value lo = isSparse ? codegen.pidxs[tensor][idx] : codegen.loops[idx];
   Value hi = isSparse ? codegen.highs[tensor][idx] : codegen.sizes[idx];
   Value step = constantIndex(builder, loc, codegen.curVecLength);
   if (isVector && codegen.options.enableVLAVectorization) {
     Value vscale = builder.create<vector::VectorScaleOp>(
         loc, IndexType::get(builder.getContext()));
     step = builder.create<arith::MulIOp>(loc, vscale, step);
   }
 
   // Emit a parallel loop.
static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, operands);
   if (codegen.redVal)
     updateReduc(merger, codegen, forOp.getRegionIterArgs().front());
   if (codegen.expValues)
     codegen.expCount = forOp.getRegionIterArgs().back();
   // Assign induction variable to sparse or dense index.
   Value iv = forOp.getInductionVar();
   if (isSparse)
     codegen.pidxs[tensor][idx] = iv;
   else
     codegen.loops[idx] = iv;
   builder.setInsertionPointToStart(forOp.getBody());
   // Share vector iteration mask between all subsequent loads/stores.
   if (isVector)
     codegen.curVecMask = genVectorMask(codegen, builder, iv, lo, hi, step);
   return forOp;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
 static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                            linalg::GenericOp op, unsigned idx, bool needsUniv,
                            BitVector &indices) {
   SmallVector<Type, 4> types;
   SmallVector<Value, 4> operands;
   // Construct the while-loop with a parameter for each index.
   Type indexType = builder.getIndexType();
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
-    if (indices[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
+    if (!indices[b])
+      continue;
+    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
+        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       types.push_back(indexType);
       operands.push_back(codegen.pidxs[tensor][idx]);
+    } else {
+      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
+             merger.isDimLevelType(b, DimLvlType::kUndef));
     }
   }
   if (codegen.redVal) {
     types.push_back(codegen.redVal.getType());
     operands.push_back(codegen.redVal);
   }
   if (codegen.expValues) {
     types.push_back(indexType);
static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
-    // TODO: singleton
-    if (indices[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
+    if (!indices[b])
+      continue;
+    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
+        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = before->getArgument(o);
       Value op2 = codegen.highs[tensor][idx];
       Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
                                                 op1, op2);
       cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
       codegen.pidxs[tensor][idx] = after->getArgument(o++);
+    } else {
+      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
+             merger.isDimLevelType(b, DimLvlType::kUndef));
     }
   }
   if (codegen.redVal)
     updateReduc(merger, codegen, after->getArgument(o++));
   if (codegen.expValues)
     codegen.expCount = after->getArgument(o++);
   if (needsUniv)
     codegen.loops[idx] = after->getArgument(o++);
static void genLocals(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                       linalg::GenericOp op, std::vector<unsigned> &topSort,
                       unsigned at, bool needsUniv, BitVector &locals) {
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse indices.
   Value min;
   for (unsigned b = 0, be = locals.size(); b < be; b++) {
-    // TODO: singleton
-    if (locals[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
+    if (!locals[b])
+      continue;
+    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
+        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value ptr = codegen.indices[tensor][idx];
       Value s = codegen.pidxs[tensor][idx];
       Value load = genLoad(codegen, builder, loc, ptr, s);
       codegen.idxs[tensor][idx] = load;
       if (!needsUniv) {
         if (min) {
           Value cmp = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::ult, load, min);
           min = builder.create<arith::SelectOp>(loc, cmp, load, min);
         } else {
           min = load;
         }
       }
+    } else {
+      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
+             merger.isDimLevelType(b, DimLvlType::kUndef));
     }
   }
 
   // Merge dense universal index over minimum.
   if (min) {
     assert(!needsUniv);
     codegen.loops[idx] = min;
   }
static void genWhileInduction(Merger &merger, CodeGen &codegen,
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
   SmallVector<Value, 4> operands;
   Value one = constantIndex(builder, loc, 1);
   for (unsigned b = 0, be = induction.size(); b < be; b++) {
-    // TODO: singleton
-    if (induction[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
+    if (!induction[b])
+      continue;
+    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
+        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = codegen.idxs[tensor][idx];
       Value op2 = codegen.loops[idx];
       Value op3 = codegen.pidxs[tensor][idx];
       Value cmp = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                 op1, op2);
       Value add = builder.create<arith::AddIOp>(loc, op3, one);
       operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, op3));
       codegen.pidxs[tensor][idx] = whileOp->getResult(o++);
+    } else {
+      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
+             merger.isDimLevelType(b, DimLvlType::kUndef));
     }
   }
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, whileOp->getResult(o++));
   }
   if (codegen.expValues) {
     operands.push_back(codegen.expCount);

 /// Generates a single if-statement within a while-loop.
 static scf::IfOp genIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                        linalg::GenericOp op, unsigned idx,
                        BitVector &conditions) {
   Location loc = op.getLoc();
   SmallVector<Type, 4> types;
   Value cond;
   for (unsigned b = 0, be = conditions.size(); b < be; b++) {
-    if (conditions[b]) {
-      unsigned tensor = merger.tensor(b);
-      assert(idx == merger.index(b));
-      Value clause;
-      // TODO: singleton
-      if (merger.isDimLevelType(b, DimLvlType::kCompressed)) {
-        Value op1 = codegen.idxs[tensor][idx];
-        Value op2 = codegen.loops[idx];
-        clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
-                                               op1, op2);
-      } else {
-        clause = constantI1(builder, loc, true);
-      }
-      cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
-    }
-  }
+    if (!conditions[b])
+      continue;
+    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
+        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+      clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, op1,
+                                             op2);
+      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
+             merger.isDimLevelType(b, DimLvlType::kUndef));
   if (codegen.redVal)
     types.push_back(codegen.redVal.getType());
   if (codegen.expValues)
     types.push_back(builder.getIndexType());
   scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, cond, /*else=*/true);
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   return ifOp;
 }

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

SparseTensorStorageBase(const std::vector<uint64_t> &dimSizes,
       : dimSizes(dimSizes), rev(getRank()),
         dimTypes(sparsity, sparsity + getRank()) {
     assert(perm && sparsity);
     const uint64_t rank = getRank();
     // Validate parameters.
     assert(rank > 0 && "Trivial shape is unsupported");
     for (uint64_t r = 0; r < rank; r++) {
       assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
-      assert((dimTypes[r] == DimLevelType::kDense ||
-              dimTypes[r] == DimLevelType::kCompressed) &&
+      assert((isDenseDim(r) || isCompressedDim(r) || isSingletonDim(r)) &&
              "Unsupported DimLevelType");
     }
     // Construct the "reverse" (i.e., inverse) permutation.
     for (uint64_t r = 0; r < rank; r++)
       rev[perm[r]] = r;
   }
 
   virtual ~SparseTensorStorageBase() = default;
public:
 
   /// Getter for the "reverse" permutation, which maps this object's
   /// storage-order to the tensor's semantic-order.
   const std::vector<uint64_t> &getRev() const { return rev; }
 
   /// Getter for the dimension-types array, in storage-order.
   const std::vector<DimLevelType> &getDimTypes() const { return dimTypes; }
 
+  /// Safely check if the (storage-order) dimension uses dense storage.
+  bool isDenseDim(uint64_t d) const {
+    assert(d < getRank());
+    return dimTypes[d] == DimLevelType::kDense;
+  }
+
   /// Safely check if the (storage-order) dimension uses compressed storage.
   bool isCompressedDim(uint64_t d) const {
     assert(d < getRank());
-    return (dimTypes[d] == DimLevelType::kCompressed);
+    switch (dimTypes[d]) {
+    case DimLevelType::kCompressed:
+    case DimLevelType::kCompressedNu:
+    case DimLevelType::kCompressedNo:
+    case DimLevelType::kCompressedNuNo:
+      return true;
+    default:
+      return false;
+    }
+  }
+
+  /// Safely check if the (storage-order) dimension uses singleton storage.
+  bool isSingletonDim(uint64_t d) const {
+    assert(d < getRank());
+    switch (dimTypes[d]) {
+    case DimLevelType::kSingleton:
+    case DimLevelType::kSingletonNu:
+    case DimLevelType::kSingletonNo:
+    case DimLevelType::kSingletonNuNo:
+      return true;
+    default:
+      return false;
+    }
+  }
+
+  /// Safely check if the (storage-order) dimension is ordered.
+  bool isOrderedDim(uint64_t d) const {
+    assert(d < getRank());
+    switch (dimTypes[d]) {
+    case DimLevelType::kCompressedNo:
+    case DimLevelType::kCompressedNuNo:
+    case DimLevelType::kSingletonNo:
+    case DimLevelType::kSingletonNuNo:
+      return false;
+    default:
+      return true;
+    }
+  }
+
+  /// Safely check if the (storage-order) dimension is unique.
+  bool isUniqueDim(uint64_t d) const {
+    assert(d < getRank());
+    switch (dimTypes[d]) {
+    case DimLevelType::kCompressedNu:
+    case DimLevelType::kCompressedNuNo:
+    case DimLevelType::kSingletonNu:
+    case DimLevelType::kSingletonNuNo:
+      return false;
+    default:
+      return true;
+    }
   }
 
   /// Allocate a new enumerator.
 #define DECL_NEWENUMERATOR(VNAME, V)                                           \
   virtual void newEnumerator(SparseTensorEnumeratorBase<V> **, uint64_t,       \
                              const uint64_t *) const {                         \
     FATAL_PIV("newEnumerator" #VNAME);                                         \
   }
for (uint64_t r = 0, rank = getRank(); r < rank; r++) {
       if (isCompressedDim(r)) {
         // TODO: Take a parameter between 1 and `dimSizes[r]`, and multiply
         // `sz` by that before reserving. (For now we just use 1.)
         pointers[r].reserve(sz + 1);
         pointers[r].push_back(0);
         indices[r].reserve(sz);
         sz = 1;
         allDense = false;
+      } else if (isSingletonDim(r)) {
+        indices[r].reserve(sz);
+        sz = 1;
+        allDense = false;
       } else { // Dense dimension.
+        assert(isDenseDim(r));
         sz = checkedMul(sz, getDimSizes()[r]);
       }
     }
     // Then assign contents from coordinate scheme tensor if provided.
     if (coo) {
       // Ensure both preconditions of `fromCOO`.
       assert(coo->getDimSizes() == getDimSizes() && "Tensor size mismatch");
       coo->sort();
private:
   /// `indices[d]` array, checking that `i` is representable in the `I`
   /// type; however, we do not verify other semantic requirements (e.g.,
   /// that `i` is in bounds for `dimSizes[d]`, and not previously occurring
   /// in the same segment).  For dense dimensions, this method instead
   /// appends the appropriate number of zeros to the `values` array,
   /// where `full` is the number of "entries" already written to `values`
   /// for this segment (aka one after the highest index previously appended).
   void appendIndex(uint64_t d, uint64_t full, uint64_t i) {
-    if (isCompressedDim(d)) {
+    if (isCompressedDim(d) || isSingletonDim(d)) {
       assert(i <= std::numeric_limits<I>::max() &&
              "Index value is too large for the I-type");
       indices[d].push_back(static_cast<I>(i));
     } else { // Dense dimension.
+      assert(isDenseDim(d));
       assert(i >= full && "Index was already filled");
       if (i == full)
         return; // Short-circuit, since it'll be a nop.
       if (d + 1 == getRank())
         values.insert(values.end(), i - full, 0);
       else
         finalizeSegment(d + 1, 0, i - full);
     }
if (d == rank) {
       return;
     }
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) { // If `hi` is unchanged, then `lo < elements.size()`.
       // Find segment in interval with same index elements in this dimension.
       uint64_t i = elements[lo].indices[d];
       uint64_t seg = lo + 1;
-      while (seg < hi && elements[seg].indices[d] == i)
+      bool merge = isUniqueDim(d);
+      while (merge && seg < hi && elements[seg].indices[d] == i)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
       appendIndex(d, full, i);
       full = i + 1;
       fromCOO(elements, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
     finalizeSegment(d, full);
   }
 
   /// Finalize the sparse pointer structure at this dimension.
   void finalizeSegment(uint64_t d, uint64_t full = 0, uint64_t count = 1) {
     if (count == 0)
       return; // Short-circuit, since it'll be a nop.
     if (isCompressedDim(d)) {
       appendPointer(d, indices[d].size(), count);
+    } else if (isSingletonDim(d)) {
+      return;
     } else { // Dense dimension.
+      assert(isDenseDim(d));
       const uint64_t sz = getDimSizes()[d];
       assert(sz >= full && "Segment is overfull");
       count = checkedMul(count, sz - full);
       // For dense storage we must enumerate all the remaining coordinates
       // in this dimension (i.e., coordinates after the last non-zero
       // element), and either fill in their zero values or else recurse
       // to finalize some deeper dimension.
       if (d + 1 == getRank())
if (d == Base::getRank()) {
       // Loop-invariant code for looking up the `d`-level coordinates/indices.
       const std::vector<I> &indicesD = src.indices[d];
       assert(pstop <= indicesD.size() && "Index position is out of bounds");
       uint64_t &cursorReordD = this->cursor[this->reord[d]];
       for (uint64_t pos = pstart; pos < pstop; pos++) {
         cursorReordD = static_cast<uint64_t>(indicesD[pos]);
         forallElements(yield, pos, d + 1);
       }
+    } else if (src.isSingletonDim(d)) {
+      FATAL("unsupported dimension level type");
     } else { // Dense dimension.
+      assert(src.isDenseDim(d));
       const uint64_t sz = src.getDimSizes()[d];
       const uint64_t pstart = parentPos * sz;
       uint64_t &cursorReordD = this->cursor[this->reord[d]];
       for (uint64_t i = 0; i < sz; i++) {
         cursorReordD = i;
         forallElements(yield, pstart + i, d + 1);
       }
     }
for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
         const uint64_t currentPos = pointers[r][parentPos];
         // This increment won't overflow the `P` type, since it can't
         // exceed the original value of `pointers[r][parentPos+1]`
         // which was already verified to be within bounds for `P`
         // when it was written to the array.
         pointers[r][parentPos]++;
         writeIndex(r, currentPos, ind[r]);
         parentPos = currentPos;
+      } else if (isSingletonDim(r)) {
       } else { // Dense dimension.
+        assert(isDenseDim(r));
         parentPos = parentPos * getDimSizes()[r] + ind[r];
       }
       parentSz = assembledSize(parentSz, r);
     }
     assert(parentPos < values.size() && "Value position is out of bounds");
     values[parentPos] = val;
   });
   // No longer need the enumerator, so we'll delete it ASAP.

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

+// RUN: mlir-opt %s -sparsification | FileCheck %s
+
+#SortedCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ]
+}>
+
+#trait_scale = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>  // X (out)
+  ],
+  iterator_types = ["parallel", "parallel"],
+  doc = "X(i,j) = X(i,j) * 2.0"
+}
+
+#trait_matvec = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (j)>,    // b
+    affine_map<(i,j) -> (i)>     // x (out)
+  ],
+  iterator_types = ["parallel","reduction"],
+  doc = "x(i) += A(i,j) * b(j)"
+}
+
+//
+// Two kernels that operate on SortedCOO format.
+//
+
+// CHECK-LABEL: func.func @sparse_scale(
+// CHECK-SAME:    %[[VAL_0:.*]]: tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> {
+// CHECK-DAG:     %[[VAL_1:.*]] = arith.constant 0 : index
+// CHECK-DAG:     %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 2.000000e+00 : f32
+// CHECK-DAG:     %[[VAL_4:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf32>
+// CHECK:         %[[VAL_6:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_1]]] : memref<?xindex>
+// CHECK:         %[[VAL_7:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:         scf.for %[[VAL_8:.*]] = %[[VAL_6]] to %[[VAL_7]] step %[[VAL_2]] {
+// CHECK:           %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_8]]] : memref<?xf32>
+// CHECK:           %[[VAL_10:.*]] = arith.mulf %[[VAL_9]], %[[VAL_3]] : f32
+// CHECK:           memref.store %[[VAL_10]], %[[VAL_5]]{{\[}}%[[VAL_8]]] : memref<?xf32>
+// CHECK:         }
+// CHECK:         %[[VAL_11:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
+// CHECK:         return %[[VAL_11]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
+// CHECK:       }
+func.func @sparse_scale(%argx: tensor<?x?xf32, #SortedCOO>) -> tensor<?x?xf32, #SortedCOO> {
+  %c = arith.constant 2.0 : f32
+  %0 = linalg.generic #trait_scale
+    outs(%argx: tensor<?x?xf32, #SortedCOO>) {
+      ^bb(%x: f32):
+        %1 = arith.mulf %x, %c : f32
+        linalg.yield %1 : f32
+  } -> tensor<?x?xf32, #SortedCOO>
+  return %0 : tensor<?x?xf32, #SortedCOO>
+}
+
+// CHECK-LABEL: func.func @matvec(
+// CHECK-SAME:    %[[VAL_0:.*]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>,
+// CHECK-SAME:    %[[VAL_1:.*]]: tensor<64xf64>,
+// CHECK-SAME:    %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
+// CHECK-DAG:     %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK-DAG:     %[[VAL_4:.*]] = arith.constant 1 : index
+// CHECK-DAG:     %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
+// CHECK-DAG:     %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
+// CHECK:         %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<64xf64>
+// CHECK:         %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xf64>
+// CHECK:         %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:         %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK:         scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
+// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK:           %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf64>
+// CHECK:           %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK:           %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xf64>
+// CHECK:           %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_16]]] : memref<64xf64>
+// CHECK:           %[[VAL_19:.*]] = arith.mulf %[[VAL_17]], %[[VAL_18]] : f64
+// CHECK:           %[[VAL_20:.*]] = arith.addf %[[VAL_15]], %[[VAL_19]] : f64
+// CHECK:           memref.store %[[VAL_20]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf64>
+// CHECK:         }
+// CHECK:         %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf64>
+// CHECK:         return %[[VAL_21]] : tensor<32xf64>
+// CHECK:       }
+func.func @matvec(%arga: tensor<32x64xf64, #SortedCOO>,
+                  %argb: tensor<64xf64>,
+                  %argx: tensor<32xf64>) -> tensor<32xf64> {
+  %0 = linalg.generic #trait_matvec
+      ins(%arga, %argb : tensor<32x64xf64, #SortedCOO>, tensor<64xf64>)
+      outs(%argx: tensor<32xf64>) {
+    ^bb(%A: f64, %b: f64, %x: f64):
+      %0 = arith.mulf %A, %b : f64
+      %1 = arith.addf %x, %0 : f64
+      linalg.yield %1 : f64
+  } -> tensor<32xf64>
+  return %0 : tensor<32xf64>
+}

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

+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
+// RUN: TENSOR1="%mlir_src_dir/test/Integration/data/mttkrp_b.tns" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+!Filename = !llvm.ptr<i8>
+
+#SortedCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ]
+}>
+
+#SortedCOOPermuted = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+#SortedCOO3D = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton-nu", "singleton" ]
+}>
+
+#SortedCOO3DPermuted = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton-nu", "singleton" ],
+  dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
+}>
+
+#trait_scale = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>  // X (out)
+  ],
+  iterator_types = ["parallel", "parallel"],
+  doc = "X(i,j) = X(i,j) * 2.0"
+}
+
+//
+// Tests reading in matrix/tensor from file into Sorted COO formats
+// as well as applying various operations to this format.
+//
+module {
+
+  func.func private @getTensorFilename(index) -> (!Filename)
+
+  //
+  // A kernel that scales a sparse matrix A by a factor of 2.0.
+  //
+  func.func @sparse_scale(%argx: tensor<?x?xf64, #SortedCOO>)
+                              -> tensor<?x?xf64, #SortedCOO> {
+    %c = arith.constant 2.0 : f64
+    %0 = linalg.generic #trait_scale
+      outs(%argx: tensor<?x?xf64, #SortedCOO>) {
+        ^bb(%x: f64):
+          %1 = arith.mulf %x, %c : f64
+          linalg.yield %1 : f64
+    } -> tensor<?x?xf64, #SortedCOO>
+    return %0 : tensor<?x?xf64, #SortedCOO>
+  }
+
+  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 @dumpf(%arg0: memref<?xf64>) {
+    %c0 = arith.constant 0 : index
+    %nan = arith.constant 0x7FF0000001000000 : f64
+    %v = vector.transfer_read %arg0[%c0], %nan: memref<?xf64>, vector<20xf64>
+    vector.print %v : vector<20xf64>
+    return
+  }
+
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+
+    %fileName0 = call @getTensorFilename(%c0) : (index) -> (!Filename)
+    %fileName1 = call @getTensorFilename(%c1) : (index) -> (!Filename)
+
+    // Read the sparse tensors from file, construct sparse storage.
+    %0 = sparse_tensor.new %fileName0 : !Filename to tensor<?x?xf64, #SortedCOO>
+    %1 = sparse_tensor.new %fileName0 : !Filename to tensor<?x?xf64, #SortedCOOPermuted>
+    %2 = sparse_tensor.new %fileName1 : !Filename to tensor<?x?x?xf64, #SortedCOO3D>
+    %3 = sparse_tensor.new %fileName1 : !Filename to tensor<?x?x?xf64, #SortedCOO3DPermuted>
+
+    // Conversion from literal.
+    %m = arith.constant sparse<
+       [ [0,0], [1,3], [2,0], [2,3], [3,1], [4,1] ],
+         [6.0, 5.0, 4.0, 3.0, 2.0, 11.0 ]
+    > : tensor<5x4xf64>
+    %4 = sparse_tensor.convert %m : tensor<5x4xf64> to tensor<?x?xf64, #SortedCOO>
+
+    //
+    // 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, nan, nan, nan )
+    //
+    %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>
+    %i01 = sparse_tensor.indices %0 { dimension = 1 : index }
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+    %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 @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, nan, nan, nan )
+    //
+    %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>
+    %i11 = sparse_tensor.indices %1 { dimension = 1 : index }
+      : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
+    %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 @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, nan, nan, nan )
+    //
+    %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>
+    %i21 = sparse_tensor.indices %2 { dimension = 1 : index }
+      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
+    %i22 = sparse_tensor.indices %2 { dimension = 2 : index }
+      : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
+    %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 @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, nan, nan, nan )
+    //
+    %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>
+    %i31 = sparse_tensor.indices %3 { dimension = 1 : index }
+      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
+    %i32 = sparse_tensor.indices %3 { dimension = 2 : index }
+      : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
+    %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 @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, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan )
+    //
+    %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>
+    %i41 = sparse_tensor.indices %4 { dimension = 1 : index }
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+    %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 @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, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan )
+    //
+    %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>
+    %i51 = sparse_tensor.indices %5 { dimension = 1 : index }
+      : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
+    %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 @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
+  }
+}