[mlir][sparse] finalize sparse output in the presence of reductions

[aartbik]

Summary

This revision implements sparse outputs (from scratch) in all cases where
the loops can be reordered with all but one parallel loops outer. If the
inner parallel loop appears inside one or more reductions loops, then an
access pattern expansion is required (aka. workspaces in TACO speak).

[aartbik]

rebased with main to keep it closer to recent mods

[bixia]

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

774

This doesn't generate code for the "False branch". Maybe we want to move it to after "Value assignment" at the end in this case?

[aartbik]

addressed comment, rebased again with main

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

string description = [{
     sparse_tensor.lex_insert %tensor, %indices, %val
       : tensor<1024x1024xf64, #CSR>, memref<?xindex>, f64
     ```
   }];
   let assemblyFormat = "$tensor `,` $indices `,` $value attr-dict `:`"
                        " type($tensor) `,` type($indices) `,` type($value)";
 }
 
+def SparseTensor_ExpandOp : SparseTensor_Op<"expand", []>,
+    Arguments<(ins AnyTensor:$tensor)>,
+    Results<(outs AnyStridedMemRefOfRank<1>:$values,
+                  StridedMemRefRankOf<[I1],[1]>:$filled,
+                  StridedMemRefRankOf<[Index],[1]>:$added,
+                  Index:$count)> {
+  string summary = "Expands an access pattern for insertion";
+  string description = [{
+    Performs an access pattern expansion for the innermost dimensions of the
+    given tensor. This operation is useful to implement kernels in which a
+    sparse tensor appears as output. This technique is known under several
+    different names and using several alternative implementations,
+    for example, phase counter [Gustavson71], expanded or switch array
+    [Pissanetzky84], in phase scan [Duff90], access pattern expansion [Bik96],
+    and workspaces [Kjolstad2018].
+
+    The values and filled array have a size the suffices for a *dense* innermost
+    dimension (e.g. a full row for matrices). The added array and count are used
+    to store new indices when a false value is encountered in the filled array.
+    All arrays should be allocated before the loop (possibly even shared between
+    loops in a future optimization) so that their *dense* intitialization can be
+    amortized over many iterations. Setting and resetting the dense arrays in
+    the loop nest itself is kept *sparse* by only iterating over set elements
+    through an indirection using the added array, so that the operations are
+    kept proportional to the number of nonzeros.
+
+    Note that this operation is "impure" in the sense that its behavior
+    is solely defined by side-effects and not SSA values. The semantics
+    may be refined over time as our sparse abstractions evolve.
+
+    Example:
+
+    ```mlir
+    %values, %filled, %added, %count = sparse_tensor.expand %0
+      : tensor<4x4xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+    ```
+  }];
+  let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($values)"
+                       " `,` type($filled) `,` type($added) `,` type($count)";
+}
+
+def SparseTensor_CompressOp : SparseTensor_Op<"compress", []>,
+    Arguments<(ins AnyTensor:$tensor,
+                   StridedMemRefRankOf<[Index],[1]>:$indices,
+                   AnyStridedMemRefOfRank<1>:$values,
+                   StridedMemRefRankOf<[I1],[1]>:$filled,
+                   StridedMemRefRankOf<[Index],[1]>:$added,
+                   Index:$count)> {
+  string summary = "Compressed an access pattern for insertion";
+  string description = [{
+    Finishes a single access pattern by moving the inserted elements
+    into the sparse storage scheme. The values and filled array are reset
+    in a *sparse* fashion by only iterating over set elements through an
+    indirection using the added array, so that the operations are kept
+    proportional to the number of nonzeros. See the 'expand' operation
+    for more details.
+
+    Note that this operation is "impure" in the sense that its behavior
+    is solely defined by side-effects and not SSA values. The semantics
+    may be refined over time as our sparse abstractions evolve.
+
+    Example:
+
+    ```mlir
+    sparse_tensor.compress %0, %1, %values, %filled, %added, %2
+        : tensor<4x4xf64, #CSR>, memref<?xindex>, memref<?xf64>,
+    memref<?xi1>, memref<?xindex>, index
+    ```
+  }];
+  let assemblyFormat = "$tensor `,` $indices `,` $values `,` $filled `,`"
+                        " $added `,` $count attr-dict `:` type($tensor) `,`"
+      " type($indices) `,` type($values) `,` type($filled) `,`"
+      " type($added) `,` type($count)";
+}
+
 def SparseTensor_LoadOp : SparseTensor_Op<"load", [SameOperandsAndResultType]>,
     Arguments<(ins AnyTensor:$tensor, UnitAttr:$hasInserts)>,
     Results<(outs AnyTensor:$result)> {
   let summary =
     "Rematerializes tensor from underlying sparse storage format";
   let description = [{
     Rematerializes a tensor from the underlying sparse storage format of the
     given tensor. This is similar to the `bufferization.to_tensor` operation

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

 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(LexInsertOp op) {
   if (!getSparseTensorEncoding(op.tensor().getType()))
     return op.emitError("expected a sparse tensor for insertion");
   return success();
 }
 
+static LogicalResult verify(ExpandOp op) {
+  if (!getSparseTensorEncoding(op.tensor().getType()))
+    return op.emitError("expected a sparse tensor for expansion");
+  return success();
+}
+
+static LogicalResult verify(CompressOp op) {
+  if (!getSparseTensorEncoding(op.tensor().getType()))
+    return op.emitError("expected a sparse tensor for compression");
+  return success();
+}
+
 static LogicalResult verify(LoadOp op) {
   if (!getSparseTensorEncoding(op.tensor().getType()))
     return op.emitError("expected a sparse tensor to materialize");
   return success();
 }
 
 static LogicalResult verify(ReleaseOp op) {
   if (!getSparseTensorEncoding(op.tensor().getType()))

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

for (unsigned i = 0, rank = stp.getRank(); i < rank; i++)
     else
       sizes.push_back(constantIndex(rewriter, loc, shape[i]));
 }
 
 /// Generates an uninitialized temporary buffer of the given size and
 /// type, but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`).
 static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
-                       unsigned sz, Type tp) {
+                       Value sz, Type tp) {
   auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
-  Value a = constantIndex(rewriter, loc, sz);
-  return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{a});
+  return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
+}
+
+/// Generates an uninitialized temporary buffer of the given size and
+/// type, but returns it as type `memref<? x $tp>` (rather than as type
+/// `memref<$sz x $tp>`).
+static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
+                       unsigned sz, Type tp) {
+  return genAlloca(rewriter, loc, constantIndex(rewriter, loc, sz), tp);
 }
 
 /// Generates an uninitialized temporary buffer with room for one value
 /// of the given type, and returns the `memref<$tp>`.
 static Value genAllocaScalar(ConversionPatternRewriter &rewriter, Location loc,
                              Type tp) {
   return rewriter.create<memref::AllocaOp>(loc, MemRefType::get({}, tp));
 }
else
       llvm_unreachable("Unknown element type");
     TypeRange noTp;
     replaceOpWithFuncCall(rewriter, op, name, noTp, adaptor.getOperands(),
                           EmitCInterface::On);
     return success();
   }
 };
 
+class SparseTensorExpandConverter : public OpConversionPattern<ExpandOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+    ShapedType srcType = op.tensor().getType().cast<ShapedType>();
+    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.tensor().getDefiningOp());
+    // Determine the size for access expansion.
+    auto enc = getSparseTensorEncoding(srcType);
+    Value src = adaptor.getOperands()[0];
+    Value sz = genDimSizeCall(rewriter, op, enc, src, srcType.getRank() - 1);
+    // Allocate temporary stack buffers for values, filled-switch, and indices.
+    Value values = genAlloca(rewriter, loc, sz, eltType);
+    Value filled = genAlloca(rewriter, loc, sz, boolType);
+    Value indices = genAlloca(rewriter, loc, sz, idxType);
+    Value zero = constantZero(rewriter, loc, idxType);
+    // Reset the values/filled-switch to all-zero/false. Note that this
+    // introduces an O(N) operation into the computation, but this reset
+    // operation is amortized over the innermost loops for the access
+    // pattern expansion.
+    rewriter.create<linalg::FillOp>(loc, constantZero(rewriter, loc, eltType),
+                                    values);
+    rewriter.create<linalg::FillOp>(loc, constantZero(rewriter, loc, boolType),
+                                    filled);
+    // Replace expansion op with these buffers and initial index.
+    assert(op.getNumResults() == 4);
+    rewriter.replaceOp(op, {values, filled, indices, zero});
+    return success();
+  }
+};
+
+class SparseTensorCompressConverter : public OpConversionPattern<CompressOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(CompressOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Note that this method call resets the values/filled-switch back to
+    // all-zero/false by only iterating over the set elements, so the
+    // complexity remains proportional to the sparsity of the expanded
+    // access pattern.
+    Type srcType = op.tensor().getType();
+    Type eltType = srcType.cast<ShapedType>().getElementType();
+    StringRef name;
+    if (eltType.isF64())
+      name = "expInsertF64";
+    else if (eltType.isF32())
+      name = "expInsertF32";
+    else if (eltType.isInteger(64))
+      name = "expInsertI64";
+    else if (eltType.isInteger(32))
+      name = "expInsertI32";
+    else if (eltType.isInteger(16))
+      name = "expInsertI16";
+    else if (eltType.isInteger(8))
+      name = "expInsertI8";
+    else
+      return failure();
+    TypeRange noTp;
+    auto fn =
+        getFunc(op, name, noTp, adaptor.getOperands(), /*emitCInterface=*/true);
+    rewriter.replaceOpWithNewOp<CallOp>(op, noTp, fn, adaptor.getOperands());
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
                                                   RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
                SparseCastConverter, SparseTensorNewConverter,
                SparseTensorInitConverter, SparseTensorConvertConverter,
                SparseTensorReleaseConverter, SparseTensorToPointersConverter,
                SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
-               SparseTensorLoadConverter, SparseTensorLexInsertConverter>(
+               SparseTensorLoadConverter, SparseTensorLexInsertConverter,
+               SparseTensorExpandConverter, SparseTensorCompressConverter>(
       typeConverter, patterns.getContext());
 }

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

CodeGen(SparsificationOptions o, unsigned numTensors, unsigned numLoops,
           OpOperand *op, unsigned nest)
       : options(o), loops(numLoops), sizes(numLoops), buffers(numTensors),
         pointers(numTensors, std::vector<Value>(numLoops)),
         indices(numTensors, std::vector<Value>(numLoops)),
         highs(numTensors, std::vector<Value>(numLoops)),
         pidxs(numTensors, std::vector<Value>(numLoops)),
         idxs(numTensors, std::vector<Value>(numLoops)), redExp(-1u), redVal(),
         redKind(kNoReduc), sparseOut(op), outerParNest(nest), lexIdx(),
-        curVecLength(1), curVecMask() {}
+        expValues(), expFilled(), expAdded(), expCount(), curVecLength(1),
+        curVecMask() {}
   /// Sparsification options.
   SparsificationOptions options;
   /// Universal dense indices and upper bounds (by index). The loops array
   /// is updated with the value of the universal dense index in the current
   /// loop. The sizes array is set once with the inferred dimension sizes.
   std::vector<Value> loops;
   std::vector<Value> sizes;
   /// Buffers for storing dense and sparse numerical values (by tensor).
struct CodeGen {
   std::vector<std::vector<Value>> idxs;
   /// Current reduction, updated during code generation. When indices of a
   /// reduction are exhausted, all inner loops can use a scalarized reduction.
   unsigned redExp;
   Value redVal;
   Reduction redKind;
   // Sparse tensor as output. Implemented either through direct injective
   // insertion in lexicographic index order (where indices are updated
-  // in the temporary array `lexIdx`) or TODO: access pattern expansion
+  // in the temporary array `lexIdx`) or through access pattern expansion
+  // in the innermost loop nest (`expValues` through `expCount`).
   OpOperand *sparseOut;
   unsigned outerParNest;
   Value lexIdx;
+  Value expValues;
+  Value expFilled;
+  Value expAdded;
+  Value expCount;
   // Current vector length and mask.
   unsigned curVecLength;
   Value curVecMask;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
if (isMaterializing(lhs->get())) {
     unsigned nest = 0;
     for (unsigned i = 0; i < numLoops; i++) {
       if (isReductionIterator(iteratorTypes[topSort[i]]))
         break; // terminate at first reduction
       nest++;
     }
     // Determine admissable dynamic insertion situations:
     // (1) fully injective, since there are no reductions,
-    // (2) admissable 1-d expansion in innermost dimension. TODO: accept
-    if (nest == op.getRank(lhs)) {
+    // (2) admissable 1-d expansion in innermost dimension.
+    if (nest >= op.getRank(lhs) - 1) {
       *sparseOut = lhs;
       outerParNest = nest;
       return true;
     }
   }
   return false;
 }
 
case AffineExprKind::Constant: {
     int64_t c = a.cast<AffineConstantExpr>().getValue();
     return rewriter.create<arith::ConstantIndexOp>(loc, c);
   }
   default:
     llvm_unreachable("unexpected affine subscript");
   }
 }
 
+/// Generates index for load/store on sparse tensor.
+static Value genIndex(CodeGen &codegen, linalg::GenericOp op, OpOperand *t) {
+  auto map = op.getTiedIndexingMap(t);
+  auto enc = getSparseTensorEncoding(t->get().getType());
+  AffineExpr a = map.getResult(perm(enc, map.getNumResults() - 1));
+  assert(a.getKind() == AffineExprKind::DimId);
+  unsigned idx = a.cast<AffineDimExpr>().getPosition();
+  return codegen.loops[idx];
+}
+
 /// Generates subscript for load/store on a dense or sparse tensor.
 static Value genSubscript(CodeGen &codegen, PatternRewriter &rewriter,
                           linalg::GenericOp op, OpOperand *t,
                           SmallVector<Value, 4> &args) {
   unsigned tensor = t->getOperandNumber();
   auto map = op.getTiedIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   unsigned rank = map.getNumResults();
if (enc) {
     for (unsigned d = 0; d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
       args.push_back(genAffine(codegen, rewriter, a, op.getLoc()));
     }
   }
   return codegen.buffers[tensor];
 }
 
+/// Generates insertion code to implement dynamic tensor load.
+static Value genInsertionLoad(CodeGen &codegen, PatternRewriter &rewriter,
+                              linalg::GenericOp op, OpOperand *t) {
+  Location loc = op.getLoc();
+  // Direct lexicographic index order, tensor loads as zero.
+  if (!codegen.expValues) {
+    Type tp = getElementTypeOrSelf(t->get().getType());
+    return rewriter.create<arith::ConstantOp>(loc, tp,
+                                              rewriter.getZeroAttr(tp));
+  }
+  // Load from expanded access pattern.
+  Value index = genIndex(codegen, op, t);
+  return rewriter.create<memref::LoadOp>(loc, codegen.expValues, index);
+}
+
+/// Generates insertion code to implement dynamic tensor store.
+static void genInsertionStore(CodeGen &codegen, PatternRewriter &rewriter,
+                              linalg::GenericOp op, OpOperand *t, Value rhs) {
+  Location loc = op.getLoc();
+  // Direct insertion in lexicographic index order.
+  if (!codegen.expValues) {
+    rewriter.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, rhs);
+    return;
+  }
+  // Generates insertion code along expanded access pattern.
+  //   if (!expFilled[i]) then
+  //     expFilled[i] = true
+  //     expAdded[inserts++] = i
+  //   endif
+  //   values[i] = rhs
+  Value index = genIndex(codegen, op, t);
+  Value fval = rewriter.create<arith::ConstantIntOp>(loc, 0, 1); // false
+  Value tval = rewriter.create<arith::ConstantIntOp>(loc, 1, 1); // true
+  // If statement.
+  Value filled = rewriter.create<memref::LoadOp>(loc, codegen.expFilled, index);
+  Value cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+                                              filled, fval);
+  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, rewriter.getIndexType(),
+                                              cond, /*else=*/true);
+  // True branch.
+  rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
+  rewriter.create<memref::StoreOp>(loc, tval, codegen.expFilled, index);
+  rewriter.create<memref::StoreOp>(loc, index, codegen.expAdded,
+                                   codegen.expCount);
+  Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
+  Value add = rewriter.create<arith::AddIOp>(loc, codegen.expCount, one);
+  rewriter.create<scf::YieldOp>(loc, add);
+  // False branch.
+  rewriter.setInsertionPointToStart(&ifOp.elseRegion().front());
+  rewriter.create<scf::YieldOp>(loc, codegen.expCount);
+  rewriter.setInsertionPointAfter(ifOp);
+  // Value assignment.
+  codegen.expCount = ifOp.getResult(0);
+  rewriter.create<memref::StoreOp>(loc, rhs, codegen.expValues, index);
+}
+
 /// Generates a load on a dense or sparse tensor.
 static Value genTensorLoad(Merger &merger, CodeGen &codegen,
                            PatternRewriter &rewriter, linalg::GenericOp op,
                            unsigned exp) {
   // Test if the load was hoisted to a higher loop nest.
   Value val = merger.exp(exp).val;
   if (val) {
     if (codegen.curVecLength > 1 && !val.getType().isa<VectorType>())
       return genVectorInvariantValue(codegen, rewriter, val);
     return val;
   }
-  // Insertion (a sparse tensor output "loads" as zero).
+  // Load during insertion.
   OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
-  if (t == codegen.sparseOut) {
-    Type tp = getElementTypeOrSelf(t->get().getType());
-    return rewriter.create<arith::ConstantOp>(op.getLoc(), tp,
-                                              rewriter.getZeroAttr(tp));
-  }
+  if (t == codegen.sparseOut)
+    return genInsertionLoad(codegen, rewriter, op, t);
   // Actual load.
   SmallVector<Value, 4> args;
   Value ptr = genSubscript(codegen, rewriter, op, t, args);
   if (codegen.curVecLength > 1)
     return genVectorLoad(codegen, rewriter, ptr, args);
   return rewriter.create<memref::LoadOp>(op.getLoc(), ptr, args);
 }
 
 /// Generates a store on a dense or sparse tensor.
 static void genTensorStore(Merger &merger, CodeGen &codegen,
                            PatternRewriter &rewriter, linalg::GenericOp op,
                            Value rhs) {
   Location loc = op.getLoc();
   // Test if this is a scalarized reduction.
   if (codegen.redVal) {
     if (codegen.curVecLength > 1)
       rhs = rewriter.create<SelectOp>(loc, codegen.curVecMask, rhs,
                                       codegen.redVal);
     updateReduc(merger, codegen, rhs);
     return;
   }
-  // Insertion.
+  // Store during insertion.
   OpOperand *t = op.getOutputOperand(0);
   if (t == codegen.sparseOut) {
-    rewriter.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, rhs);
+    genInsertionStore(codegen, rewriter, op, t, rhs);
     return;
   }
   // Actual store.
   SmallVector<Value, 4> args;
   Value ptr = genSubscript(codegen, rewriter, op, t, args);
   if (codegen.curVecLength > 1)
     genVectorStore(codegen, rewriter, rhs, ptr, args);
   else
if (merger.exp(exp).kind == Kind::kTensor) {
     Kind last = merger.exp(exp).kind;
     unsigned e0 = merger.exp(exp).children.e0;
     unsigned e1 = merger.exp(exp).children.e1;
     genInvariants(merger, codegen, rewriter, op, e0, ldx, atStart, last);
     genInvariants(merger, codegen, rewriter, op, e1, ldx, atStart, last);
   }
 }
 
+/// Generates an expanded access pattern in innermost dimension.
+static void genExpansion(Merger &merger, CodeGen &codegen,
+                         PatternRewriter &rewriter, linalg::GenericOp op,
+                         unsigned at, bool atStart) {
+  OpOperand *lhs = codegen.sparseOut;
+  if (!lhs || codegen.outerParNest != op.getRank(lhs) - 1 ||
+      at != codegen.outerParNest)
+    return; // not needed at this level
+  // Generate start or end of an expanded access pattern.
+  Value tensor = lhs->get();
+  Location loc = op.getLoc();
+  if (atStart) {
+    auto dynShape = {ShapedType::kDynamicSize};
+    Type etp = tensor.getType().cast<ShapedType>().getElementType();
+    Type t1 = MemRefType::get(dynShape, etp);
+    Type t2 = MemRefType::get(dynShape, genIntType(rewriter, 1));
+    Type t3 = MemRefType::get(dynShape, genIntType(rewriter, 0));
+    Type t4 = rewriter.getIndexType();
+    auto res =
+        rewriter.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
+    assert(res.getNumResults() == 4);
+    assert(!codegen.expValues);
+    codegen.expValues = res.getResult(0);
+    codegen.expFilled = res.getResult(1);
+    codegen.expAdded = res.getResult(2);
+    codegen.expCount = res.getResult(3);
+  } else {
+    assert(codegen.expValues);
+    rewriter.create<CompressOp>(loc, tensor, codegen.lexIdx, codegen.expValues,
+                                codegen.expFilled, codegen.expAdded,
+                                codegen.expCount);
+    codegen.expValues = codegen.expFilled = codegen.expAdded =
+        codegen.expCount = Value();
+  }
+}
+
 /// Generates initialization code for the subsequent loop sequence at
 /// current index level. Returns true if the loop sequence needs to
 /// maintain the universal index.
 static bool genInit(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned at, llvm::BitVector &inits) {
   bool needsUniv = false;
   Location loc = op.getLoc();
if (codegen.redVal) {
     // In a vector loop, bring reduction into SIMD form, if not already.
     if (isVector && !codegen.redVal.getType().isa<VectorType>()) {
       VectorType vtp = vectorType(codegen, codegen.redVal.getType());
       Value vred = genVectorReducInit(codegen, rewriter, loc, vtp);
       updateReduc(merger, codegen, vred);
     }
     operands.push_back(codegen.redVal);
   }
+  if (codegen.expValues)
+    operands.push_back(codegen.expCount);
   scf::ForOp forOp = rewriter.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;
   rewriter.setInsertionPointToStart(forOp.getBody());
   // Share vector iteration mask between all subsequent loads/stores.
if (indices[b] && merger.isDim(b, Dim::kSparse)) {
       types.push_back(indexType);
       operands.push_back(codegen.pidxs[tensor][idx]);
     }
   }
   if (codegen.redVal) {
     types.push_back(codegen.redVal.getType());
     operands.push_back(codegen.redVal);
   }
+  if (codegen.expValues) {
+    types.push_back(indexType);
+    operands.push_back(codegen.expCount);
+  }
   if (needsUniv) {
     types.push_back(indexType);
     operands.push_back(codegen.loops[idx]);
   }
   assert(types.size() == operands.size());
   Location loc = op.getLoc();
   scf::WhileOp whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
   Block *before = rewriter.createBlock(&whileOp.before(), {}, types);
if (indices[b] && merger.isDim(b, Dim::kSparse)) {
       Value opc = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
                                                  op1, op2);
       cond = cond ? rewriter.create<arith::AndIOp>(loc, cond, opc) : opc;
       codegen.pidxs[tensor][idx] = after->getArgument(o++);
     }
   }
   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++);
   assert(o == operands.size());
   rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
   rewriter.setInsertionPointToStart(&whileOp.after().front());
   return whileOp;
 }
 
if ((locals[b] || merger.isOutTensor(b, idx)) &&
           break;
       Value p = (pat == 0) ? rewriter.create<arith::ConstantIndexOp>(loc, 0)
                            : codegen.pidxs[tensor][topSort[pat - 1]];
       codegen.pidxs[tensor][idx] = genAddress(
           codegen, rewriter, loc, codegen.sizes[idx], p, codegen.loops[idx]);
     }
   }
 
-  // Move the insertion indices in lexicographic index order.
-  if (codegen.sparseOut) {
+  // Move the insertion indices in lexicographic index order. During access
+  // pattern expansion, we can skip setting the innermost dimension.
+  if (codegen.sparseOut && !codegen.expValues) {
     Value pos = rewriter.create<arith::ConstantIndexOp>(loc, at);
     rewriter.create<memref::StoreOp>(loc, codegen.loops[idx], codegen.lexIdx,
                                      pos);
   }
 }
 
 /// Generates the induction structure for a while-loop.
 static void genWhileInduction(Merger &merger, CodeGen &codegen,
                               PatternRewriter &rewriter, linalg::GenericOp op,
                               unsigned idx, bool needsUniv,
                               llvm::BitVector &induction,
                               scf::WhileOp whileOp) {
   Location loc = op.getLoc();
   // Finalize each else branch of all if statements.
-  if (codegen.redVal) {
+  if (codegen.redVal || codegen.expValues) {
     while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
                rewriter.getInsertionBlock()->getParentOp())) {
-      rewriter.create<scf::YieldOp>(loc, codegen.redVal);
-      updateReduc(merger, codegen, ifOp.getResult(0));
+      unsigned y = 0;
+      SmallVector<Value, 4> yields;
+      if (codegen.redVal) {
+        yields.push_back(codegen.redVal);
+        updateReduc(merger, codegen, ifOp.getResult(y++));
+      }
+      if (codegen.expValues) {
+        yields.push_back(codegen.expCount);
+        codegen.expCount = ifOp->getResult(y++);
+      }
+      assert(y == yields.size());
+      rewriter.create<scf::YieldOp>(loc, yields);
       rewriter.setInsertionPointAfter(ifOp);
     }
   }
   rewriter.setInsertionPointToEnd(&whileOp.after().front());
   // Finalize the induction. Note that the induction could be performed
   // 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
if (induction[b] && merger.isDim(b, Dim::kSparse)) {
       operands.push_back(rewriter.create<SelectOp>(loc, cmp, add, op3));
       codegen.pidxs[tensor][idx] = whileOp->getResult(o++);
     }
   }
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, whileOp->getResult(o++));
   }
+  if (codegen.expValues) {
+    operands.push_back(codegen.expCount);
+    codegen.expCount = whileOp->getResult(o++);
+  }
   if (needsUniv) {
     operands.push_back(
         rewriter.create<arith::AddIOp>(loc, codegen.loops[idx], one));
     codegen.loops[idx] = whileOp->getResult(o++);
   }
   assert(o == operands.size());
   rewriter.create<scf::YieldOp>(loc, operands);
   rewriter.setInsertionPointAfter(whileOp);
 }
 
 /// Generates the induction structure for a for-loop.
 static void genForInduction(Merger &merger, CodeGen &codegen,
                             PatternRewriter &rewriter, linalg::GenericOp op,
                             Operation *loop) {
   Location loc = op.getLoc();
   unsigned o = 0;
   SmallVector<Value, 4> operands;
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, loop->getResult(o++));
   }
+  if (codegen.expValues) {
+    operands.push_back(codegen.expCount);
+    codegen.expCount = loop->getResult(o++);
+  }
   assert(o == operands.size());
   if (o > 0)
     rewriter.create<scf::YieldOp>(loc, operands);
   rewriter.setInsertionPointAfter(loop);
 }
 
 /// Generates a single if-statement within a while-loop.
 static scf::IfOp genIf(Merger &merger, CodeGen &codegen,
if (conditions[b]) {
       } else {
         clause = rewriter.create<arith::ConstantIntOp>(loc, 1, 1); // true
       }
       cond = cond ? rewriter.create<arith::AndIOp>(loc, cond, clause) : clause;
     }
   }
   if (codegen.redVal)
     types.push_back(codegen.redVal.getType());
+  if (codegen.expValues)
+    types.push_back(rewriter.getIndexType());
   scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, types, cond, /*else=*/true);
   rewriter.setInsertionPointToStart(&ifOp.thenRegion().front());
   return ifOp;
 }
 
 /// Generates end of true branch of if-statement within a while-loop.
 static void endIf(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
-                  linalg::GenericOp op, scf::IfOp ifOp, Value ifInput) {
+                  linalg::GenericOp op, scf::IfOp ifOp, Operation *loop,
+                  Value redInput, Value cntInput) {
+  SmallVector<Value, 4> operands;
   if (codegen.redVal) {
-    rewriter.create<scf::YieldOp>(op.getLoc(), codegen.redVal);
-    updateReduc(merger, codegen, ifInput);
+    operands.push_back(codegen.redVal);
+    updateReduc(merger, codegen, redInput);
+  }
+  if (codegen.expValues) {
+    operands.push_back(codegen.expCount);
+    codegen.expCount = cntInput;
   }
+  if (!operands.empty())
+    rewriter.create<scf::YieldOp>(op.getLoc(), operands);
   rewriter.setInsertionPointToStart(&ifOp.elseRegion().front());
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (loop sequence).
 //===----------------------------------------------------------------------===//
 
 /// Starts a loop sequence at given level. Returns true if
 /// the universal loop index must be maintained at this level.
 static bool startLoopSeq(Merger &merger, CodeGen &codegen,
                          PatternRewriter &rewriter, linalg::GenericOp op,
                          std::vector<unsigned> &topSort, unsigned exp,
                          unsigned at, unsigned idx, unsigned ldx,
                          unsigned lts) {
   assert(codegen.curVecLength == 1);
   assert(!codegen.loops[idx]);
   // Emit invariants at this loop sequence level.
   genInvariants(merger, codegen, rewriter, op, exp, ldx, /*atStart=*/true);
+  // Emit access pattern expansion for sparse tensor output.
+  genExpansion(merger, codegen, rewriter, op, at, /*atStart=*/true);
   // Emit further intitialization at this loop sequence level.
   unsigned l0 = merger.set(lts)[0];
   bool needsUniv =
       genInit(merger, codegen, rewriter, op, topSort, at, merger.lat(l0).bits);
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     unsigned lsize = merger.set(lts).size();
static bool endLoop(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
   // End a for-loop.
   genForInduction(merger, codegen, rewriter, op, loop);
   return false;
 }
 
 /// Ends a loop sequence at given level.
 static void endLoopSeq(Merger &merger, CodeGen &codegen,
                        PatternRewriter &rewriter, linalg::GenericOp op,
-                       unsigned exp, unsigned idx, unsigned ldx) {
+                       unsigned exp, unsigned at, unsigned idx, unsigned ldx) {
   assert(codegen.curVecLength == 1);
   codegen.loops[idx] = Value();
   // Bring a pending reduction back from SIMD form when sequence ends.
   if (codegen.redVal)
     if (auto vtp = codegen.redVal.getType().dyn_cast<VectorType>())
       updateReduc(merger, codegen,
                   genVectorReducEnd(codegen, rewriter, op.getLoc(), vtp));
   // Unmark bookkeeping of invariants and loop index.
   genInvariants(merger, codegen, rewriter, op, exp, ldx, /*atStart=*/false);
+  // Finalize access pattern expansion for sparse tensor output.
+  genExpansion(merger, codegen, rewriter, op, at, /*atStart=*/false);
 }
 
 /// Recursively generates code while computing iteration lattices in order
 /// to manage the complexity of implementing co-iteration over unions
 /// and intersections of sparse iterations spaces.
 static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned exp, unsigned at) {
static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
   for (unsigned i = 0; i < lsize; i++) {
     // Start a loop.
     unsigned li = merger.set(lts)[i];
     Operation *loop =
         startLoop(merger, codegen, rewriter, op, topSort, at, li, needsUniv);
 
     // Visit all lattices points with Li >= Lj to generate the
     // loop-body, possibly with if statements for coiteration.
-    Value ifInput = codegen.redVal;
+    Value redInput = codegen.redVal;
+    Value cntInput = codegen.expCount;
     bool isWhile = dyn_cast<scf::WhileOp>(loop) != nullptr;
     for (unsigned j = 0; j < lsize; j++) {
       unsigned lj = merger.set(lts)[j];
       unsigned ej = merger.lat(lj).exp;
       if (li == lj || merger.latGT(li, lj)) {
         // Recurse into body of each branch.
         if (isWhile) {
           scf::IfOp ifOp =
               genIf(merger, codegen, rewriter, op, idx, merger.lat(lj).simple);
           genStmt(merger, codegen, rewriter, op, topSort, ej, at + 1);
-          endIf(merger, codegen, rewriter, op, ifOp, ifInput);
+          endIf(merger, codegen, rewriter, op, ifOp, loop, redInput, cntInput);
         } else {
           genStmt(merger, codegen, rewriter, op, topSort, ej, at + 1);
         }
       }
     }
 
     // End a loop.
     needsUniv =
         endLoop(merger, codegen, rewriter, op, loop, idx, li, needsUniv);
   }
 
   // End a loop sequence.
-  endLoopSeq(merger, codegen, rewriter, op, exp, idx, ldx);
+  endLoopSeq(merger, codegen, rewriter, op, exp, at, idx, ldx);
 }
 
 /// Converts the result computed by the sparse kernel into the required form.
 static void genResult(Merger &merger, CodeGen &codegen,
                       PatternRewriter &rewriter, linalg::GenericOp op) {
   OpOperand *lhs = op.getOutputOperand(0);
   Type resType = lhs->get().getType();
   Value result;

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

private:
   bool iteratorLocked;
   unsigned iteratorPos;
 };
 
 /// Abstract base class of sparse tensor storage. Note that we use
 /// function overloading to implement "partial" method specialization.
 class SparseTensorStorageBase {
 public:
-  // Dimension size query.
+  /// Dimension size query.
   virtual uint64_t getDimSize(uint64_t) = 0;
 
-  // Overhead storage.
+  /// Overhead storage.
   virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
   virtual void getPointers(std::vector<uint32_t> **, uint64_t) { fatal("p32"); }
   virtual void getPointers(std::vector<uint16_t> **, uint64_t) { fatal("p16"); }
   virtual void getPointers(std::vector<uint8_t> **, uint64_t) { fatal("p8"); }
   virtual void getIndices(std::vector<uint64_t> **, uint64_t) { fatal("i64"); }
   virtual void getIndices(std::vector<uint32_t> **, uint64_t) { fatal("i32"); }
   virtual void getIndices(std::vector<uint16_t> **, uint64_t) { fatal("i16"); }
   virtual void getIndices(std::vector<uint8_t> **, uint64_t) { fatal("i8"); }
 
-  // Primary storage.
+  /// Primary storage.
   virtual void getValues(std::vector<double> **) { fatal("valf64"); }
   virtual void getValues(std::vector<float> **) { fatal("valf32"); }
   virtual void getValues(std::vector<int64_t> **) { fatal("vali64"); }
   virtual void getValues(std::vector<int32_t> **) { fatal("vali32"); }
   virtual void getValues(std::vector<int16_t> **) { fatal("vali16"); }
   virtual void getValues(std::vector<int8_t> **) { fatal("vali8"); }
 
-  // Element-wise insertion in lexicographic index order.
+  /// Element-wise insertion in lexicographic index order.
   virtual void lexInsert(uint64_t *, double) { fatal("insf64"); }
   virtual void lexInsert(uint64_t *, float) { fatal("insf32"); }
   virtual void lexInsert(uint64_t *, int64_t) { fatal("insi64"); }
   virtual void lexInsert(uint64_t *, int32_t) { fatal("insi32"); }
   virtual void lexInsert(uint64_t *, int16_t) { fatal("ins16"); }
   virtual void lexInsert(uint64_t *, int8_t) { fatal("insi8"); }
+
+  /// Expanded insertion.
+  virtual void expInsert(uint64_t *, double *, bool *, uint64_t *, uint64_t) {
+    fatal("expf64");
+  }
+  virtual void expInsert(uint64_t *, float *, bool *, uint64_t *, uint64_t) {
+    fatal("expf32");
+  }
+  virtual void expInsert(uint64_t *, int64_t *, bool *, uint64_t *, uint64_t) {
+    fatal("expi64");
+  }
+  virtual void expInsert(uint64_t *, int32_t *, bool *, uint64_t *, uint64_t) {
+    fatal("expi32");
+  }
+  virtual void expInsert(uint64_t *, int16_t *, bool *, uint64_t *, uint64_t) {
+    fatal("expi16");
+  }
+  virtual void expInsert(uint64_t *, int8_t *, bool *, uint64_t *, uint64_t) {
+    fatal("expi8");
+  }
+
+  /// Finishes insertion.
   virtual void endInsert() = 0;
 
   virtual ~SparseTensorStorageBase() {}
 
 private:
   void fatal(const char *tp) {
     fprintf(stderr, "unsupported %s\n", tp);
     exit(1);
if (!values.empty()) {
       diff = lexDiff(cursor);
       endPath(diff + 1);
       top = idx[diff] + 1;
     }
     // Then continue with insertion path.
     insPath(cursor, diff, top, val);
   }
 
+  /// Partially specialize expanded insertions based on template types.
+  /// Note that this method resets the values/filled-switch array back
+  /// to all-zero/false while only iterating over the nonzero elements.
+  void expInsert(uint64_t *cursor, V *values, bool *filled, uint64_t *added,
+                 uint64_t count) override {
+    if (count == 0)
+      return;
+    // Sort.
+    std::sort(added, added + count);
+    // Restore insertion path for first insert.
+    uint64_t rank = getRank();
+    uint64_t index = added[0];
+    cursor[rank - 1] = index;
+    lexInsert(cursor, values[index]);
+    assert(filled[index]);
+    values[index] = 0;
+    filled[index] = false;
+    // Subsequent insertions are quick.
+    for (uint64_t i = 1; i < count; i++) {
+      assert(index < added[i] && "non-lexicographic insertion");
+      index = added[i];
+      cursor[rank - 1] = index;
+      insPath(cursor, rank - 1, added[i - 1] + 1, values[index]);
+      assert(filled[index]);
+      values[index] = 0.0;
+      filled[index] = false;
+    }
+  }
+
   /// Finalizes lexicographic insertions.
   void endInsert() override {
     if (values.empty())
       endDim(0);
     else
       endPath(0);
   }
 
if (action <= Action::kFromCOO) {                                          \
       return tensor;                                                           \
     }                                                                          \
   }
 
 #define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
 
 #define IMPL_SPARSEVALUES(NAME, TYPE, LIB)                                     \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor) {    \
-    assert(ref);                                                               \
-    assert(tensor);                                                            \
+    assert(ref &&tensor);                                                      \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v);                   \
     ref->basePtr = ref->data = v->data();                                      \
     ref->offset = 0;                                                           \
     ref->sizes[0] = v->size();                                                 \
     ref->strides[0] = 1;                                                       \
   }
 
 #define IMPL_GETOVERHEAD(NAME, TYPE, LIB)                                      \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor,      \
                            index_t d) {                                        \
-    assert(ref);                                                               \
-    assert(tensor);                                                            \
+    assert(ref &&tensor);                                                      \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
     ref->basePtr = ref->data = v->data();                                      \
     ref->offset = 0;                                                           \
     ref->sizes[0] = v->size();                                                 \
     ref->strides[0] = 1;                                                       \
   }
 
 #define IMPL_ADDELT(NAME, TYPE)                                                \
   void *_mlir_ciface_##NAME(void *tensor, TYPE value,                          \
                             StridedMemRefType<index_t, 1> *iref,               \
                             StridedMemRefType<index_t, 1> *pref) {             \
-    assert(tensor);                                                            \
-    assert(iref);                                                              \
-    assert(pref);                                                              \
+    assert(tensor &&iref &&pref);                                              \
     assert(iref->strides[0] == 1 && pref->strides[0] == 1);                    \
     assert(iref->sizes[0] == pref->sizes[0]);                                  \
     const index_t *indx = iref->data + iref->offset;                           \
     const index_t *perm = pref->data + pref->offset;                           \
     uint64_t isize = iref->sizes[0];                                           \
     std::vector<index_t> indices(isize);                                       \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indices[perm[r]] = indx[r];                                              \
     static_cast<SparseTensorCOO<TYPE> *>(tensor)->add(indices, value);         \
     return tensor;                                                             \
   }
 
 #define IMPL_GETNEXT(NAME, V)                                                  \
-  bool _mlir_ciface_##NAME(void *tensor, StridedMemRefType<uint64_t, 1> *iref, \
+  bool _mlir_ciface_##NAME(void *tensor, StridedMemRefType<index_t, 1> *iref,  \
                            StridedMemRefType<V, 0> *vref) {                    \
+    assert(tensor &&iref &&vref);                                              \
     assert(iref->strides[0] == 1);                                             \
-    uint64_t *indx = iref->data + iref->offset;                                \
+    index_t *indx = iref->data + iref->offset;                                 \
     V *value = vref->data + vref->offset;                                      \
     const uint64_t isize = iref->sizes[0];                                     \
     auto iter = static_cast<SparseTensorCOO<V> *>(tensor);                     \
     const Element<V> *elem = iter->getNext();                                  \
     if (elem == nullptr) {                                                     \
       delete iter;                                                             \
       return false;                                                            \
     }                                                                          \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indx[r] = elem->indices[r];                                              \
     *value = elem->value;                                                      \
     return true;                                                               \
   }
 
 #define IMPL_LEXINSERT(NAME, V)                                                \
   void _mlir_ciface_##NAME(void *tensor, StridedMemRefType<index_t, 1> *cref,  \
                            V val) {                                            \
+    assert(tensor &&cref);                                                     \
     assert(cref->strides[0] == 1);                                             \
-    uint64_t *cursor = cref->data + cref->offset;                              \
+    index_t *cursor = cref->data + cref->offset;                               \
     assert(cursor);                                                            \
     static_cast<SparseTensorStorageBase *>(tensor)->lexInsert(cursor, val);    \
   }
 
+#define IMPL_EXPINSERT(NAME, V)                                                \
+  void _mlir_ciface_##NAME(                                                    \
+      void *tensor, StridedMemRefType<index_t, 1> *cref,                       \
+      StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref,         \
+      StridedMemRefType<index_t, 1> *aref, index_t count) {                    \
+    assert(tensor &&cref &&vref &&fref &&aref);                                \
+    assert(cref->strides[0] == 1);                                             \
+    assert(vref->strides[0] == 1);                                             \
+    assert(fref->strides[0] == 1);                                             \
+    assert(aref->strides[0] == 1);                                             \
+    assert(vref->sizes[0] == fref->sizes[0]);                                  \
+    index_t *cursor = cref->data + cref->offset;                               \
+    V *values = vref->data + vref->offset;                                     \
+    bool *filled = fref->data + fref->offset;                                  \
+    index_t *added = aref->data + aref->offset;                                \
+    static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
+        cursor, values, filled, added, count);                                 \
+  }
+
 /// Constructs a new sparse tensor. This is the "swiss army knife"
 /// method for materializing sparse tensors into the computation.
 ///
 /// Action:
 /// kEmpty = returns empty storage to fill later
 /// kFromFile = returns storage, where ptr contains filename to read
 /// kFromCOO = returns storage, where ptr contains coordinate scheme to assign
 /// kEmptyCOO = returns empty coordinate scheme to fill and use with kFromCOO

 /// Helper to insert elements in lexicograph index order, one per value type.
 IMPL_LEXINSERT(lexInsertF64, double)
 IMPL_LEXINSERT(lexInsertF32, float)
 IMPL_LEXINSERT(lexInsertI64, int64_t)
 IMPL_LEXINSERT(lexInsertI32, int32_t)
 IMPL_LEXINSERT(lexInsertI16, int16_t)
 IMPL_LEXINSERT(lexInsertI8, int8_t)
 
+/// Helper to insert using expansion, one per value type.
+IMPL_EXPINSERT(expInsertF64, double)
+IMPL_EXPINSERT(expInsertF32, float)
+IMPL_EXPINSERT(expInsertI64, int64_t)
+IMPL_EXPINSERT(expInsertI32, int32_t)
+IMPL_EXPINSERT(expInsertI16, int16_t)
+IMPL_EXPINSERT(expInsertI8, int8_t)
+
 #undef CASE
 #undef IMPL_SPARSEVALUES
 #undef IMPL_GETOVERHEAD
 #undef IMPL_ADDELT
 #undef IMPL_GETNEXT
-#undef IMPL_INSERTLEX
+#undef IMPL_LEXINSERT
+#undef IMPL_EXPINSERT
 
 //===----------------------------------------------------------------------===//
 //
 // Public API with methods that accept C-style data structures to interact
 // with sparse tensors, which are only visible as opaque pointers externally.
 // These methods can be used both by MLIR compiler-generated code as well as by
 // an external runtime that wants to interact with MLIR compiler-generated code.
 //

mlir/test/Dialect/SparseTensor/conversion.mlir

 //       CHECK: return
 func @sparse_insert(%arg0: tensor<128xf32, #SparseVector>,
                     %arg1: memref<?xindex>,
                     %arg2: f32) {
   sparse_tensor.lex_insert %arg0, %arg1, %arg2 : tensor<128xf32, #SparseVector>, memref<?xindex>, f32
   return
 }
 
+// CHECK-LABEL: func @sparse_expansion()
+//    %[[S:.*]] = call @sparseDimSize
+//    %[[V:.*]] = memref.alloca(%[[S]]) : memref<?xf64>
+//    %[[F:.*]] = memref.alloca(%[[S]]) : memref<?xi1>
+//    %[[A:.*]] = memref.alloca(%[[S]]) : memref<?xindex>
+//    linalg.fill(%{{.*}}, %[[V]]) : f64, memref<?xf64>
+//    linalg.fill(%{{.*}}, %[[F]]) : i1, memref<?xi1>
+//       CHECK: return
+func @sparse_expansion() {
+  %c = arith.constant 8 : index
+  %0 = sparse_tensor.init [%c, %c] : tensor<8x8xf64, #SparseMatrix>
+  %values, %filled, %added, %count = sparse_tensor.expand %0
+    : tensor<8x8xf64, #SparseMatrix> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return
+}
+
+// CHECK-LABEL: func @sparse_compression(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>,
+//       CHECK: call @expInsertF64(%[[A]],
+//       CHECK: return
+func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
+                         %arg1: memref<?xindex>, %arg2: memref<?xf64>, %arg3: memref<?xi1>,
+                         %arg4: memref<?xindex>, %arg5: index) {
+  sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
+    : tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return
+}

mlir/test/Dialect/SparseTensor/invalid.mlir

 func @sparse_unannotated_insert(%arg0: tensor<128xf64>, %arg1: memref<?xindex>, %arg2: f64) {
   // expected-error@+1 {{expected a sparse tensor for insertion}}
   sparse_tensor.lex_insert %arg0, %arg1, %arg2 : tensor<128xf64>, memref<?xindex>, f64
   return
 }
 
 // -----
 
+func @sparse_unannotated_expansion(%arg0: tensor<128xf64>) {
+  // expected-error@+1 {{expected a sparse tensor for expansion}}
+  %values, %filled, %added, %count = sparse_tensor.expand %arg0
+    : tensor<128xf64> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return
+}
+
+// -----
+
+func @sparse_unannotated_compression(%arg0: tensor<128xf64>, %arg1: memref<?xindex>,
+                                     %arg2: memref<?xf64>, %arg3: memref<?xi1>,
+             %arg4: memref<?xindex>, %arg5: index) {
+  // expected-error@+1 {{expected a sparse tensor for compression}}
+  sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
+    : tensor<128xf64>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+}
+
+// -----
+
 func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
   // expected-error@+1 {{unexpected type in convert}}
   %0 = sparse_tensor.convert %arg0 : tensor<*xf32> to tensor<10xf32>
   return %0 : tensor<10xf32>
 }
 
 // -----
 

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //  CHECK-SAME: %[[B:.*]]: memref<?xindex>,
 //  CHECK-SAME: %[[C:.*]]: f64) {
 //       CHECK: sparse_tensor.lex_insert %[[A]], %[[B]], %[[C]] : tensor<128xf64, #{{.*}}>, memref<?xindex>, f64
 //       CHECK: return
 func @sparse_insert(%arg0: tensor<128xf64, #SparseVector>, %arg1: memref<?xindex>, %arg2: f64) {
   sparse_tensor.lex_insert %arg0, %arg1, %arg2 : tensor<128xf64, #SparseVector>, memref<?xindex>, f64
   return
 }
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_expansion(
+//  CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>)
+//       CHECK: sparse_tensor.expand %[[A]]
+//       CHECK: return
+func @sparse_expansion(%arg0: tensor<8x8xf64, #SparseMatrix>) {
+  %values, %filled, %added, %count = sparse_tensor.expand %arg0
+    : tensor<8x8xf64, #SparseMatrix> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return
+}
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_compression(
+//  CHECK-SAME: %[[A:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
+//       CHECK: sparse_tensor.compress %[[A]]
+//       CHECK: return
+func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
+                         %arg1: memref<?xindex>, %arg2: memref<?xf64>, %arg3: memref<?xi1>,
+                         %arg4: memref<?xindex>, %arg5: index) {
+  sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
+    : tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return
+}

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

+// RUN: mlir-opt %s \
+// RUN:   --linalg-generalize-named-ops --linalg-fuse-elementwise-ops \
+// RUN:   --sparsification --sparse-tensor-conversion \
+// RUN:   --linalg-bufferize --convert-linalg-to-loops \
+// RUN:   --convert-vector-to-scf --convert-scf-to-std \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
+// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm \
+// RUN:   --convert-std-to-llvm --reconcile-unrealized-casts | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+//
+
+#CSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (i,j)>
+}>
+
+#DCSR = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (i,j)>
+}>
+
+module {
+  //
+  // Computes C = A x B with all matrices dense.
+  //
+  func @matmul1(%A: tensor<4x8xf64>,
+                %B: tensor<8x4xf64>) -> tensor<4x4xf64> {
+    %C = arith.constant dense<0.0> : tensor<4x4xf64>
+    %D = linalg.matmul
+      ins(%A, %B: tensor<4x8xf64>, tensor<8x4xf64>)
+         outs(%C: tensor<4x4xf64>) -> tensor<4x4xf64>
+    return %D: tensor<4x4xf64>
+  }
+
+  //
+  // Computes C = A x B with all matrices sparse (SpMSpM) in CSR.
+  //
+  func @matmul2(%A: tensor<4x8xf64, #CSR>,
+                %B: tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> {
+    %c4 = arith.constant 4 : index
+    %C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #CSR>
+    %D = linalg.matmul
+      ins(%A, %B: tensor<4x8xf64, #CSR>, tensor<8x4xf64, #CSR>)
+         outs(%C: tensor<4x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
+    return %D: tensor<4x4xf64, #CSR>
+  }
+
+  //
+  // Computes C = A x B with all matrices sparse (SpMSpM) in DCSR.
+  //
+  func @matmul3(%A: tensor<4x8xf64, #DCSR>,
+                %B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
+    %c4 = arith.constant 4 : index
+    %C = sparse_tensor.init [%c4, %c4] : tensor<4x4xf64, #DCSR>
+    %D = linalg.matmul
+      ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>)
+         outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
+    return %D: tensor<4x4xf64, #DCSR>
+  }
+
+  //
+  // Main driver.
+  //
+  func @entry() {
+    %c0 = arith.constant 0 : index
+    %d1 = arith.constant -1.0 : f64
+
+    // Initialize various matrices, dense for stress testing,
+    // and sparse to verify correct nonzero structure.
+    %da = arith.constant dense<[
+        [ 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1 ],
+        [ 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2 ],
+        [ 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3 ],
+        [ 1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4 ]
+    ]> : tensor<4x8xf64>
+    %db = arith.constant dense<[
+        [ 10.1, 11.1, 12.1, 13.1 ],
+        [ 10.2, 11.2, 12.2, 13.2 ],
+        [ 10.3, 11.3, 12.3, 13.3 ],
+        [ 10.4, 11.4, 12.4, 13.4 ],
+        [ 10.5, 11.5, 12.5, 13.5 ],
+        [ 10.6, 11.6, 12.6, 13.6 ],
+        [ 10.7, 11.7, 12.7, 13.7 ],
+        [ 10.8, 11.8, 12.8, 13.8 ]
+    ]> : tensor<8x4xf64>
+    %sa = arith.constant dense<[
+        [ 0.0, 2.1, 0.0, 0.0, 0.0, 6.1, 0.0, 0.0 ],
+        [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
+        [ 0.0, 2.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ],
+        [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 ]
+    ]> : tensor<4x8xf64>
+    %sb = arith.constant dense<[
+        [ 0.0, 0.0, 0.0, 1.0 ],
+        [ 0.0, 0.0, 2.0, 0.0 ],
+        [ 0.0, 3.0, 0.0, 0.0 ],
+        [ 4.0, 0.0, 0.0, 0.0 ],
+        [ 0.0, 0.0, 0.0, 0.0 ],
+        [ 0.0, 5.0, 0.0, 0.0 ],
+        [ 0.0, 0.0, 6.0, 0.0 ],
+        [ 0.0, 0.0, 7.0, 8.0 ]
+    ]> : tensor<8x4xf64>
+
+    // Convert all these matrices to sparse format.
+    %a1 = sparse_tensor.convert %da : tensor<4x8xf64> to tensor<4x8xf64, #CSR>
+    %a2 = sparse_tensor.convert %da : tensor<4x8xf64> to tensor<4x8xf64, #DCSR>
+    %a3 = sparse_tensor.convert %sa : tensor<4x8xf64> to tensor<4x8xf64, #CSR>
+    %a4 = sparse_tensor.convert %sa : tensor<4x8xf64> to tensor<4x8xf64, #DCSR>
+    %b1 = sparse_tensor.convert %db : tensor<8x4xf64> to tensor<8x4xf64, #CSR>
+    %b2 = sparse_tensor.convert %db : tensor<8x4xf64> to tensor<8x4xf64, #DCSR>
+    %b3 = sparse_tensor.convert %sb : tensor<8x4xf64> to tensor<8x4xf64, #CSR>
+    %b4 = sparse_tensor.convert %sb : tensor<8x4xf64> to tensor<8x4xf64, #DCSR>
+
+    // Call kernels with dense.
+    %0 = call @matmul1(%da, %db)
+       : (tensor<4x8xf64>, tensor<8x4xf64>) -> tensor<4x4xf64>
+    %1 = call @matmul2(%a1, %b1)
+       : (tensor<4x8xf64, #CSR>,
+          tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
+    %2 = call @matmul3(%a2, %b2)
+       : (tensor<4x8xf64, #DCSR>,
+          tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
+
+    // Call kernels with one sparse.
+    %3 = call @matmul1(%sa, %db)
+       : (tensor<4x8xf64>, tensor<8x4xf64>) -> tensor<4x4xf64>
+    %4 = call @matmul2(%a3, %b1)
+       : (tensor<4x8xf64, #CSR>,
+          tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
+    %5 = call @matmul3(%a4, %b2)
+       : (tensor<4x8xf64, #DCSR>,
+          tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
+
+    // Call kernels with sparse.
+    %6 = call @matmul1(%sa, %sb)
+       : (tensor<4x8xf64>, tensor<8x4xf64>) -> tensor<4x4xf64>
+    %7 = call @matmul2(%a3, %b3)
+       : (tensor<4x8xf64, #CSR>,
+          tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
+    %8 = call @matmul3(%a4, %b4)
+       : (tensor<4x8xf64, #DCSR>,
+          tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
+
+    //
+    // CHECK:    ( ( 388.76, 425.56, 462.36, 499.16 ),
+    // CHECK-SAME: ( 397.12, 434.72, 472.32, 509.92 ),
+    // CHECK-SAME: ( 405.48, 443.88, 482.28, 520.68 ),
+    // CHECK-SAME: ( 413.84, 453.04, 492.24, 531.44 ) )
+    //
+    %m0 = bufferization.to_memref %0 : memref<4x4xf64>
+    %v0 = vector.transfer_read %m0[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v0 : vector<4x4xf64>
+
+    //
+    // CHECK:    ( ( 388.76, 425.56, 462.36, 499.16 ),
+    // CHECK-SAME: ( 397.12, 434.72, 472.32, 509.92 ),
+    // CHECK-SAME: ( 405.48, 443.88, 482.28, 520.68 ),
+    // CHECK-SAME: ( 413.84, 453.04, 492.24, 531.44 ) )
+    //
+    %c1 = sparse_tensor.convert %1 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
+    %m1 = bufferization.to_memref %c1 : memref<4x4xf64>
+    %v1 = vector.transfer_read %m1[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v1 : vector<4x4xf64>
+
+    //
+    // CHECK:    ( ( 388.76, 425.56, 462.36, 499.16 ),
+    // CHECK-SAME: ( 397.12, 434.72, 472.32, 509.92 ),
+    // CHECK-SAME: ( 405.48, 443.88, 482.28, 520.68 ),
+    // CHECK-SAME: ( 413.84, 453.04, 492.24, 531.44 ) )
+    //
+    %c2 = sparse_tensor.convert %2 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
+    %m2 = bufferization.to_memref %c2 : memref<4x4xf64>
+    %v2 = vector.transfer_read %m2[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v2 : vector<4x4xf64>
+
+    //
+    // CHECK:    ( ( 86.08, 94.28, 102.48, 110.68 ),
+    // CHECK-SAME: ( 0, 0, 0, 0 ),
+    // CHECK-SAME: ( 23.46, 25.76, 28.06, 30.36 ),
+    // CHECK-SAME: ( 10.8, 11.8, 12.8, 13.8 ) )
+    //
+    %m3 = bufferization.to_memref %3 : memref<4x4xf64>
+    %v3 = vector.transfer_read %m3[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v3 : vector<4x4xf64>
+
+    //
+    // CHECK:    ( ( 86.08, 94.28, 102.48, 110.68 ),
+    // CHECK-SAME: ( 0, 0, 0, 0 ),
+    // CHECK-SAME: ( 23.46, 25.76, 28.06, 30.36 ),
+    // CHECK-SAME: ( 10.8, 11.8, 12.8, 13.8 ) )
+    //
+    %c4 = sparse_tensor.convert %4 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
+    %m4 = bufferization.to_memref %c4 : memref<4x4xf64>
+    %v4 = vector.transfer_read %m4[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v4 : vector<4x4xf64>
+
+    //
+    // CHECK:    ( ( 86.08, 94.28, 102.48, 110.68 ),
+    // CHECK-SAME: ( 0, 0, 0, 0 ),
+    // CHECK-SAME: ( 23.46, 25.76, 28.06, 30.36 ),
+    // CHECK-SAME: ( 10.8, 11.8, 12.8, 13.8 ) )
+    //
+    %c5 = sparse_tensor.convert %5 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
+    %m5 = bufferization.to_memref %c5 : memref<4x4xf64>
+    %v5 = vector.transfer_read %m5[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v5 : vector<4x4xf64>
+
+    //
+    // CHECK: ( ( 0, 30.5, 4.2, 0 ), ( 0, 0, 0, 0 ), ( 0, 0, 4.6, 0 ), ( 0, 0, 7, 8 ) )
+    //
+    %m6 = bufferization.to_memref %6 : memref<4x4xf64>
+    %v6 = vector.transfer_read %m6[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v6 : vector<4x4xf64>
+
+    //
+    // CHECK: ( ( 0, 30.5, 4.2, 0 ), ( 0, 0, 0, 0 ), ( 0, 0, 4.6, 0 ), ( 0, 0, 7, 8 ) )
+    //
+    %c7 = sparse_tensor.convert %7 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
+    %m7 = bufferization.to_memref %c7 : memref<4x4xf64>
+    %v7 = vector.transfer_read %m7[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v7 : vector<4x4xf64>
+
+    //
+    // CHECK: ( ( 0, 30.5, 4.2, 0 ), ( 0, 0, 0, 0 ), ( 0, 0, 4.6, 0 ), ( 0, 0, 7, 8 ) )
+    //
+    %c8 = sparse_tensor.convert %8 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
+    %m8 = bufferization.to_memref %c8 : memref<4x4xf64>
+    %v8 = vector.transfer_read %m8[%c0, %c0], %d1 : memref<4x4xf64>, vector<4x4xf64>
+    vector.print %v8 : vector<4x4xf64>
+
+    //
+    // Sanity check on nonzeros.
+    //
+    // CHECK: ( 30.5, 4.2, 4.6, 7, 8, -1, -1, -1 )
+    // CHECK: ( 30.5, 4.2, 4.6, 7, 8, -1, -1, -1 )
+    //
+    %val7 = sparse_tensor.values %7 : tensor<4x4xf64, #CSR> to memref<?xf64>
+    %val8 = sparse_tensor.values %8 : tensor<4x4xf64, #DCSR> to memref<?xf64>
+    %nz7 = vector.transfer_read %val7[%c0], %d1 : memref<?xf64>, vector<8xf64>
+    %nz8 = vector.transfer_read %val8[%c0], %d1 : memref<?xf64>, vector<8xf64>
+    vector.print %nz7 : vector<8xf64>
+    vector.print %nz8 : vector<8xf64>
+
+    // Release the resources.
+    sparse_tensor.release %a1 : tensor<4x8xf64, #CSR>
+    sparse_tensor.release %a2 : tensor<4x8xf64, #DCSR>
+    sparse_tensor.release %a3 : tensor<4x8xf64, #CSR>
+    sparse_tensor.release %a4 : tensor<4x8xf64, #DCSR>
+    sparse_tensor.release %b1 : tensor<8x4xf64, #CSR>
+    sparse_tensor.release %b2 : tensor<8x4xf64, #DCSR>
+    sparse_tensor.release %b3 : tensor<8x4xf64, #CSR>
+    sparse_tensor.release %b4 : tensor<8x4xf64, #DCSR>
+    sparse_tensor.release %1 : tensor<4x4xf64, #CSR>
+    sparse_tensor.release %2 : tensor<4x4xf64, #DCSR>
+    sparse_tensor.release %4 : tensor<4x4xf64, #CSR>
+    sparse_tensor.release %5 : tensor<4x4xf64, #DCSR>
+    sparse_tensor.release %7 : tensor<4x4xf64, #CSR>
+    sparse_tensor.release %8 : tensor<4x4xf64, #DCSR>
+    memref.dealloc %m0 : memref<4x4xf64>
+    memref.dealloc %m1 : memref<4x4xf64>
+    memref.dealloc %m2 : memref<4x4xf64>
+    memref.dealloc %m3 : memref<4x4xf64>
+    memref.dealloc %m4 : memref<4x4xf64>
+    memref.dealloc %m5 : memref<4x4xf64>
+    memref.dealloc %m6 : memref<4x4xf64>
+    memref.dealloc %m7 : memref<4x4xf64>
+    memref.dealloc %m8 : memref<4x4xf64>
+
+    return
+  }
+}