[mlir][sparse] extend unpack operation to support unpacking a batched COO type

[Peiming]

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

157

@springerm Do you know why this does not work and I still need to manually deallocate returned tensors?

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

692

This is a unrelated bug (but I think it should be trivial enough to be reviewed in the same patch). Let me know if you want it in a separate patch! @aartbik

[Peiming]

remove debugging code.

[Peiming]

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

167

@springerm See here

[aartbik]

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

130

new syntax and verification needs roundtrip.mlir and invalid.mlir coverage

150

this needs a new example

[Peiming]

add roundtrip/invalid test + examples.

[Peiming]

fix typo.

[aartbik]

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

156–157

(1) We allocate

(2) Can you keep the comment for == 0, (The "Similar to InsertOp ...."

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

221

This "as explained above" has become a bit of a misleading link after all the code movement

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

692

Fine for now, but indeed, in the future just send such one liners out as separate revision

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

1233

this should really be in the "helper method" section above (we try to keep the rewriting rules together, and the helper methods together)

1245

same here, in principle, but since that was already here, we can leave that for a later cleanup

1434

I am okay keeping the genUnBatchedUnpackOp/genBatchedUnpackOp methods here, since they are really just the conversion rules

1477

this is a helper, so move up

[Peiming]

address comments.

[Peiming]

provide more build() implementation.

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

def SparseTensor_PackOp : SparseTensor_Op<"pack", [Pure]>,
   let assemblyFormat =
     "$values `,` $coordinates (`batched_lvls` `=` $batched_lvls^)? attr-dict"
     "`:` type($values) `,` type($coordinates) `to` type($result)";
 
   let hasVerifier = 1;
 }
 
 def SparseTensor_UnpackOp : SparseTensor_Op<"unpack">,
-    Arguments<(ins AnySparseTensor:$tensor)>,
-    Results<(outs 1DTensorOf<[AnyType]>:$values,
-                  2DTensorOf<[AnySignlessIntegerOrIndex]>:$coordinates,
+    Arguments<(ins AnySparseTensor:$tensor,
+                   OptionalAttr<IndexAttr>:$batched_lvls)>,
+    Results<(outs TensorOf<[AnyType]>:$values,
+                  TensorOf<[AnySignlessIntegerOrIndex]>:$coordinates,
                   AnySignlessIntegerOrIndex:$nse)> {
   let summary = "Returns the (values, coordinates) pair unpacked from the input tensor";
 
   let description = [{
     The unpack operation is the inverse of `sparse_tensor::pack`.  It returns
     the values, level-coordinates, and number-of-stored-entries extracted
     from the sparse tensor.  The source tensor is allowed (in principle)
     to have non-identity dimOrdering/higherOrdering mappings.  Regardless
     of the mappings, the returned `coordinates` are always level-coordinates,
     because this is what we mean by "unpacking" as opposed to other forms
     of exposing sparse tensors to external clients.  This operation can be
     used for returning an unpacked MLIR sparse tensor to frontend; e.g.,
     returning two numpy arrays to Python.
 
     TODO: the current implementation does not yet support non-identity mappings.
 
     This operation ends the lifetime of the sparse tensor, and using
     the tensor after the unpack is undefined behavior.
 
     Example:
 
     ```mlir
     // input COO format |1.1, 0.0, 0.0, 0.0|
     //    of 3x4 matrix |0.0, 0.0, 2.2, 3.3|
     //                  |0.0, 0.0, 0.0, 0.0|
     %values, %coordinates, %nse
       = sparse_tensor.unpack %st
       : tensor<3x4xf64, #COO> to tensor<2xf64>, tensor<2x2xindex>, index
     // %values      = arith.constant dense<[ 1.1,   2.2,   3.3 ]> : tensor<3xf64>
     // %coordinates = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
     // %nse = 3
     ```
   }];
 
+  let extraClassDeclaration = [{
+    /// Returns the number of leading levels that are batched.
+    unsigned getNumBatchedLvls();
+  }];
+
   let assemblyFormat =
-    "$tensor attr-dict `:` type($tensor)"
-    "`to` type($values) `,` type($coordinates) `,` type($nse)";
+    "$tensor (`batched_lvls` `=` $batched_lvls^)? attr-dict `:`"
+    "type($tensor) `to` type($values) `,` type($coordinates) `,` type($nse)";
 
   let hasVerifier = 1;
 }
 
 def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
   [Pure]>,
     Arguments<(ins AnyTensor:$source)>,
     Results<(outs AnyTensor:$dest)> {

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

unsigned PackOp::getNumBatchedLvls() {
   return getBatchedLvls().has_value() ? getBatchedLvls()->getZExtValue() : 0;
 }
 
 LogicalResult UnpackOp::verify() {
   const auto valuesTp = getRankedTensorType(getValues());
   const auto coordinatesTp = getRankedTensorType(getCoordinates());
   const auto srcTp = getSparseTensorType(getTensor());
   return verifyPackUnPack(*this, false, srcTp, valuesTp, coordinatesTp,
-                          nullptr);
+                          getBatchedLvlsAttr());
+}
+
+unsigned UnpackOp::getNumBatchedLvls() {
+  return getBatchedLvls().has_value() ? getBatchedLvls()->getZExtValue() : 0;
 }
 
 LogicalResult ConvertOp::verify() {
   if (auto tp1 = getSource().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
       if (tp1.getRank() != tp2.getRank())
         return emitError("unexpected conversion mismatch in rank");
       auto dstEnc =

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

BufferRelation bufferRelation(Operation *oo, OpResult opResult,
     return BufferRelation::Unknown;
   }
 };
 
 struct UnpackOpInterface
     : public BufferizableOpInterface::ExternalModel<UnpackOpInterface,
                                                     sparse_tensor::UnpackOp> {
   bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
-    // Similar to InsertOp, reallocation is not considered to allocate a new
-    // piece of memory.
-    return false;
+    // We allocates and returns unpacked memory if this is a batched unpack.
+    return llvm::cast<UnpackOp>(op).getNumBatchedLvls() != 0;
   }
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,

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

 
 //===----------------------------------------------------------------------===//
 // Misc code generators and utilities.
 //===----------------------------------------------------------------------===//
 
 /// Add type casting between arith and index types when needed.
 Value genCast(OpBuilder &builder, Location loc, Value value, Type dstTy);
 
+/// Generates a pointer/index load from the sparse storage scheme. Narrower
+/// data types need to be zero extended before casting the value into the
+/// index type used for looping and indexing.
+Value genIndexLoad(OpBuilder &builder, Location loc, Value mem, Value s);
+
 /// Generates a 1-valued attribute of the given type.  This supports
 /// all the same types as `getZeroAttr`; however, unlike `getZeroAttr`,
 /// for unsupported types we raise `llvm_unreachable` rather than
 /// returning a null attribute.
 TypedAttr getOneAttr(Builder &builder, Type tp);
 
 /// Generates the comparison `v != 0` where `v` is of numeric type.
 /// For floating types, we use the "unordered" comparator (i.e., returns

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

Value sparse_tensor::genCast(OpBuilder &builder, Location loc, Value value,
   if (srcTp.isa<IndexType>() || dstTp.isa<IndexType>())
     return builder.create<arith::IndexCastOp>(loc, dstTp, value);
 
   const auto srcIntTp = srcTp.dyn_cast_or_null<IntegerType>();
   const bool isUnsignedCast = srcIntTp ? srcIntTp.isUnsigned() : false;
   return mlir::convertScalarToDtype(builder, loc, value, dstTp, isUnsignedCast);
 }
 
+Value sparse_tensor::genIndexLoad(OpBuilder &builder, Location loc, Value mem,
+                                  Value s) {
+  // For the scalar case, we simply zero extend narrower indices into 64-bit
+  // values before casting to index without a performance penalty. Here too,
+  // however, indices that already are 64-bit, in theory, cannot express the
+  // full range as explained above.
+  Value load = builder.create<memref::LoadOp>(loc, mem, s);
+  if (!load.getType().isa<IndexType>()) {
+    if (load.getType().getIntOrFloatBitWidth() < 64)
+      load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
+    load =
+        builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
+  }
+  return load;
+}
+
 mlir::TypedAttr mlir::sparse_tensor::getOneAttr(Builder &builder, Type tp) {
   if (tp.isa<FloatType>())
     return builder.getFloatAttr(tp, 1.0);
   if (tp.isa<IndexType>())
     return builder.getIndexAttr(1);
   if (auto intTp = tp.dyn_cast<IntegerType>())
     return builder.getIntegerAttr(tp, APInt(intTp.getWidth(), 1));
   if (tp.isa<RankedTensorType, VectorType>()) {

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

 #define SUBI(lhs, rhs) (builder.create<arith::SubIOp>(loc, (lhs), (rhs)))
 #define MULI(lhs, rhs) (builder.create<arith::MulIOp>(loc, (lhs), (rhs)))
 #define SELECT(c, l, r) (builder.create<arith::SelectOp>(loc, (c), (l), (r)))
 
 //===----------------------------------------------------------------------===//
 // File local helper functions.
 //===----------------------------------------------------------------------===//
 
-/// Generates a pointer/index load from the sparse storage scheme. Narrower
-/// data types need to be zero extended before casting the value into the
-/// index type used for looping and indexing.
-static Value genIndexLoad(OpBuilder &builder, Location loc, Value mem,
-                          Value s) {
-  // For the scalar case, we simply zero extend narrower indices into 64-bit
-  // values before casting to index without a performance penalty. Here too,
-  // however, indices that already are 64-bit, in theory, cannot express the
-  // full range as explained above.
-  Value load = builder.create<memref::LoadOp>(loc, mem, s);
-  if (!load.getType().isa<IndexType>()) {
-    if (load.getType().getIntOrFloatBitWidth() < 64)
-      load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
-    load =
-        builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
-  }
-  return load;
-}
-
 static Value genSliceOffset(OpBuilder &builder, Location loc, Value tensor,
                             Level lvl) {
   auto enc = getSparseTensorEncoding(tensor.getType());
   // FIXME: `toOrigDim` is deprecated
   return createOrFoldSliceOffsetOp(builder, loc, tensor, toOrigDim(enc, lvl));
 }
 
 static Value genSliceStride(OpBuilder &builder, Location loc, Value tensor,
for (auto [t, l] : llvm::zip(tids, lvls)) {
     if (!dependentLvlMap[t][l].empty() && !isDenseDLT(lvlType)) {
       assert(!isSparseSliceCond && !isSparseCond);
       isSparseSliceCond = true;
       tid = t;
       lvl = l;
       continue;
     }
 
-    bool isSparse = isCompressedDLT(lvlType) || isSingletonDLT(lvlType);
+    bool isSparse = isCompressedDLT(lvlType) || isSingletonDLT(lvlType) ||
+                    isCompressedWithHiDLT(lvlType);
     // We can at most have one sparse input, otherwise, a while loop is
     // required to co-iterate multiple sparse tensors.
     assert(!isSparseCond || !isSparse);
     assert(!isSparseSliceCond || !isSparseCond);
     if (isSparse) {
       tid = t;
       lvl = l;
     }

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

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Query memSizes for the actually stored values.
     rewriter.replaceOp(
         op, genValMemSize(rewriter, op.getLoc(), adaptor.getTensor()));
     return success();
   }
 };
 
+static Value linearize(OpBuilder &builder, Location loc, ValueRange ivs,
+                       ValueRange bounds) {
+  assert(ivs.size() == bounds.size());
+  Value crd = constantIndex(builder, loc, 0);
+  for (unsigned i = 0, e = ivs.size(); i < e; i++) {
+    crd = builder.create<arith::AddIOp>(loc, crd, ivs[i]);
+    if (i != ivs.size() - 1)
+      crd = builder.create<arith::MulIOp>(loc, crd, bounds[i + 1]);
+  }
+  return crd;
+}
+
 static void populateCompressedWithHiPosArray(OpBuilder &builder, Location loc,
                                              ArrayRef<unsigned> batchDimSzs,
                                              Value posMemRef, unsigned nse,
                                              PackOp op) {
   SmallVector<Value> lbs, ubs, steps;
   Value c0 = constantIndex(builder, loc, 0);
   Value c1 = constantIndex(builder, loc, 1);
   Value c2 = constantIndex(builder, loc, 2);
   for (unsigned dimSz : batchDimSzs) {
     lbs.push_back(c0);
     ubs.push_back(constantIndex(builder, loc, dimSz));
     steps.push_back(c1);
   }
   auto tensorType = op.getValues().getType();
   auto memrefType =
       MemRefType::get(tensorType.getShape(), tensorType.getElementType());
   Value batV = builder.create<bufferization::ToMemrefOp>(loc, memrefType,
                                                          op.getValues());
   scf::buildLoopNest(
       builder, loc, lbs, ubs, steps,
       [&ubs, c0, c1, c2, nse, batV, posMemRef](OpBuilder &builder, Location loc,
                                                ValueRange ivs) {
         // Linearize index variables
-        Value crd = constantIndex(builder, loc, 0);
-        for (unsigned i = 0, e = ivs.size(); i < e; i++) {
-          crd = builder.create<arith::AddIOp>(loc, crd, ivs[i]);
-          if (i != ivs.size() - 1)
-            crd = builder.create<arith::MulIOp>(loc, crd, ubs[i + 1]);
-        }
+        Value crd = linearize(builder, loc, ivs, ubs);
         Value len = constantIndex(builder, loc, nse);
         Value pLo = builder.create<arith::MulIOp>(loc, crd, len);
         SmallVector<Value> indices(ivs.begin(), ivs.end());
         auto whileOp = builder.create<scf::WhileOp>(
             loc, TypeRange{builder.getIndexType()}, ValueRange{len},
             [&indices, c0, c1, batV](OpBuilder &builder, Location loc,
                                      ValueRange vs) {
               Value curLen = vs.front();
scf::buildLoopNest(
               Value nxLen = builder.create<arith::SubIOp>(loc, vs.front(), c1);
               builder.create<scf::YieldOp>(loc, nxLen);
             });
         len = whileOp.getResults()[0];
         Value pHi = builder.create<arith::AddIOp>(loc, pLo, len);
         // Stores position lower bound.
         Value idx = builder.create<arith::MulIOp>(loc, crd, c2);
         genStore(builder, loc, pLo, posMemRef, idx);
+        // builder.create<vector::PrintOp>(loc, idx);
+        // builder.create<vector::PrintOp>(loc, pLo);
+        // builder.create<vector::PrintOp>(loc, pHi);
         // Stores position upper bound.
         idx = builder.create<arith::AddIOp>(loc, idx, c1);
         genStore(builder, loc, pHi, posMemRef, idx);
       });
 }
 
 struct SparsePackOpConverter : public OpConversionPattern<PackOp> {
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(PackOp op, OpAdaptor adaptor,
     }
     desc.setValMemSize(rewriter, loc, noe);
 
     rewriter.replaceOp(op, genTuple(rewriter, loc, desc));
     return success();
   }
 };
 
-struct SparseUnpackOpConverter : public OpConversionPattern<UnpackOp> {
-  using OpConversionPattern::OpConversionPattern;
-  SparseUnpackOpConverter(TypeConverter &typeConverter, MLIRContext *context,
-                          bool createDeallocs)
-      : OpConversionPattern(typeConverter, context),
-        createDeallocs(createDeallocs) {}
-
-  LogicalResult
-  matchAndRewrite(UnpackOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
-    Location loc = op.getLoc();
-    const auto srcTp = getSparseTensorType(op.getTensor());
-    const Level lvlRank = srcTp.getLvlRank();
-
-    assert(isUniqueCOOType(srcTp) && desc.getFields().size() == 4);
-
-    Value flatBuf = lvlRank == 1 ? desc.getCrdMemRefOrView(rewriter, loc, 0)
-                                 : desc.getAOSMemRef();
-    Value valuesBuf = desc.getValMemRef();
-    Value posBuf = desc.getPosMemRef(0);
-    if (createDeallocs) {
-      // Unpack ends the lifetime of the sparse tensor. While the value array
-      // and coordinate array are unpacked and returned, the position array
-      // becomes useless and need to be freed (if user requests).
-      rewriter.create<memref::DeallocOp>(loc, posBuf);
-    }
+static LogicalResult genUnBatchedUnpackOp(UnpackOp op,
+                                          SparseTensorDescriptor desc,
+                                          ConversionPatternRewriter &rewriter) {
 
-    // If frontend requests a static buffer, we reallocate the
-    // values/coordinates to ensure that we meet their need.
-    const auto valuesTp = getRankedTensorType(op.getValues());
-    if (valuesTp.hasStaticShape()) {
-      valuesBuf =
-          reallocOrSubView(rewriter, loc, valuesTp.getShape()[0], valuesBuf);
-    }
+    // FIXME: Reallocation is not always safe! E.g., if we are unpacking a
+    // tensor that is packed from constants.
 
-    const auto coordinatesTp = getRankedTensorType(op.getCoordinates());
-    if (coordinatesTp.hasStaticShape()) {
-      auto len = coordinatesTp.getShape()[0] * coordinatesTp.getShape()[1];
-      flatBuf = reallocOrSubView(rewriter, loc, len, flatBuf);
-    }
+    // FIXME: Reallocation is not always safe! E.g., if we are unpacking a
+    // tensor that is packed from constants.
 
-    Value coordinatesBuf = rewriter.create<memref::ExpandShapeOp>(
-        loc,
-        MemRefType::get(coordinatesTp.getShape(),
-                        coordinatesTp.getElementType()),
-        flatBuf, ArrayRef{ReassociationIndices{0, 1}});
+      MemRefType::get(coordinatesTp.getShape(), coordinatesTp.getElementType()),
 
-    // Converts MemRefs back to Tensors.
-    Value values = rewriter.create<bufferization::ToTensorOp>(loc, valuesBuf);
-    Value coordinates =
-        rewriter.create<bufferization::ToTensorOp>(loc, coordinatesBuf);
-    Value nse = genCast(rewriter, loc, desc.getValMemSize(rewriter, loc),
-                        op.getNse().getType());
 
-    rewriter.replaceOp(op, {values, coordinates, nse});
-    return success();
-  }
 
+ReassociationIndices getReassociationForFlattening(ShapedType srcTp) {
+  ReassociationIndices reassociation;
+  for (int i = 0, e = srcTp.getRank(); i < e; i++)
+    reassociation.push_back(i);
+  return reassociation;
+}
+
+static LogicalResult genBatchedUnpackOp(UnpackOp op, unsigned nBatched,
+                                        SparseTensorDescriptor desc,
+                                        ConversionPatternRewriter &rewriter) {
+  assert(nBatched != 0);
+  Location loc = op.getLoc();
+  Value c0 = constantIndex(rewriter, loc, 0);
+  Value c1 = constantIndex(rewriter, loc, 1);
+  Value c2 = constantIndex(rewriter, loc, 2);
+
+  auto genZeroedAlloc = [loc,
+                         &rewriter](TensorType tt) -> TypedValue<MemRefType> {
+    auto mem = rewriter
+                   .create<memref::AllocOp>(
+                       loc, MemRefType::get(tt.getShape(), tt.getElementType()))
+                   .getMemref();
+    // TODO: Instead of filling the entire buffer, we can only fill the
+    // trailing zeros.
+    rewriter.create<linalg::FillOp>(
+        loc, ValueRange{constantZero(rewriter, loc, tt.getElementType())}, mem);
+    return mem;
+  };
+  SparseTensorType stt = getSparseTensorType(op.getTensor());
+  TensorType valTensorTp = op.getValues().getType();
+  TensorType crdTensorTp = op.getCoordinates().getType();
+  TypedValue<MemRefType> valMemref = genZeroedAlloc(valTensorTp);
+  TypedValue<MemRefType> crdMemref = genZeroedAlloc(crdTensorTp);
+  assert(valTensorTp.hasStaticShape() && crdTensorTp.hasStaticShape());
+
+  SmallVector<Value> lbs(nBatched, c0), steps(nBatched, c1);
+  SmallVector<Value> ubs;
+  for (unsigned i = 0; i < nBatched; i++) {
+    assert(!ShapedType::isDynamic(stt.getDimShape()[i]));
+    ubs.push_back(constantIndex(rewriter, loc, stt.getDimShape()[i]));
+  }
+
+  DimLevelType dlt = stt.getLvlType(nBatched);
+  assert(isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt));
+  Value posStep = isCompressedDLT(dlt) ? c1  // forward position index by 1
+                                       : c2; // forward position index by 2
+  auto loopNest = scf::buildLoopNest(
+      rewriter, loc, lbs, ubs, steps, {c0 /*maximum nse*/},
+      [&ubs, c0, c1, posStep, desc, nBatched, &valMemref,
+       &crdMemref](OpBuilder &builder, Location loc, ValueRange ivs,
+                   ValueRange args) -> scf::ValueVector {
+        // crdMemref has shape: <... x nse x rank>
+        unsigned unBatchedRank = crdMemref.getType().getShape().back();
+        Value values = desc.getValMemRef();
+        Value flatCrds = unBatchedRank == 1
+                             ? desc.getCrdMemRefOrView(builder, loc, 0)
+                             : desc.getAOSMemRef();
+
+        Value positions = desc.getPosMemRef(nBatched);
+        Value positLo = builder.create<arith::MulIOp>(
+            loc, linearize(builder, loc, ivs, ubs), posStep);
+        Value positHi = builder.create<arith::AddIOp>(loc, positLo, c1);
+
+        Value pLo = genIndexLoad(builder, loc, positions, positLo);
+        Value pHi = genIndexLoad(builder, loc, positions, positHi);
+        Value nse = builder.create<arith::SubIOp>(loc, pHi, pLo);
+
+        Value crdLo = builder.create<arith::MulIOp>(
+            loc, pLo, constantIndex(builder, loc, unBatchedRank));
+        Value nCrd = builder.create<arith::MulIOp>(
+            loc, nse, constantIndex(builder, loc, unBatchedRank));
+
+        SmallVector<Value> offsets, sizes, strides;
+        for (unsigned i = 0; i < nBatched; i++) {
+          offsets.push_back(ivs[i]);
+          sizes.push_back(c1);
+          strides.push_back(c1);
+        }
+        // [0, nse, 1].
+        offsets.push_back(c0);
+        sizes.push_back(nse);
+        strides.push_back(c1);
+
+        auto valView = builder.create<memref::SubViewOp>(
+            loc, valMemref, offsets, sizes, strides);
+        auto valReass = getReassociationForFlattening(valView.getType());
+        Value valDst =
+            builder.create<memref::CollapseShapeOp>(loc, valView, valReass);
+        Value valSrc =
+            builder.create<memref::SubViewOp>(loc, values, pLo, nse, c1);
+        builder.create<memref::CopyOp>(loc, valSrc, valDst);
+
+        // [0, rank, 1].
+        offsets.push_back(c0);
+        sizes.push_back(constantIndex(builder, loc, unBatchedRank));
+        strides.push_back(c1);
+
+        auto crdView = builder.create<memref::SubViewOp>(
+            loc, crdMemref, offsets, sizes, strides);
+        auto crdReass = getReassociationForFlattening(crdView.getType());
+        Value crdDst =
+            builder.create<memref::CollapseShapeOp>(loc, crdView, crdReass);
+        Value crdSrc =
+            builder.create<memref::SubViewOp>(loc, flatCrds, crdLo, nCrd, c1);
+        builder.create<memref::CopyOp>(loc, crdSrc, crdDst);
+
+        Value pred = builder.create<arith::CmpIOp>(
+            loc, arith::CmpIPredicate::ugt, nse, args[0]);
+        // Choose the larger NSE
+        return {builder.create<arith::SelectOp>(loc, pred, nse, args[0])};
+      });
+
+  // Converts MemRefs back to Tensors.
+  Value values = rewriter.create<bufferization::ToTensorOp>(loc, valMemref);
+  Value coordinates =
+      rewriter.create<bufferization::ToTensorOp>(loc, crdMemref);
+  Value nse =
+      genCast(rewriter, loc, loopNest.results.front(), op.getNse().getType());
+
+  rewriter.replaceOp(op, {values, coordinates, nse});
+  return success();
+}
+
+struct SparseUnpackOpConverter : public OpConversionPattern<UnpackOp> {
+  using OpConversionPattern::OpConversionPattern;
+  SparseUnpackOpConverter(TypeConverter &typeConverter, MLIRContext *context,
+                          bool createDeallocs)
+      : OpConversionPattern(typeConverter, context),
+        createDeallocs(createDeallocs) {}
+
+  LogicalResult
+  matchAndRewrite(UnpackOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    const auto srcTp = getSparseTensorType(op.getTensor());
+    const unsigned nBatched = op.getNumBatchedLvls();
+    assert(isCOOType(srcTp.getEncoding(), nBatched, true) &&
+           desc.getFields().size() == 4); // specifier + pos + crds + values
+    auto logicRes = nBatched == 0
+                        ? genUnBatchedUnpackOp(op, desc, rewriter)
+                        : genBatchedUnpackOp(op, nBatched, desc, rewriter);
+    Value posBuf = desc.getPosMemRef(nBatched);
+
+    if (createDeallocs) {
+      // Unpack ends the lifetime of the sparse tensor. While the value array
+      // and coordinate array are unpacked and returned, the position array
+      // becomes useless and need to be freed (if user requests).
+      // FIXME: Depending on whether the tensor being unpacked is created by
+      // PackOp or not, we may or may not need to free other memref fields of
+      // the sparse tensor too (PackOp borrows value/coordinate buffer).
+      rewriter.create<memref::DeallocOp>(op.getLoc(), posBuf);
+    }
+
+    return logicRes;
+  }
+
 private:
   const bool createDeallocs;
 };
 
 struct SparseNewOpConverter : public OpConversionPattern<NewOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NewOp op, OpAdaptor adaptor,

mlir/test/Dialect/SparseTensor/sparse_2d.mlir

func.func @add_ss_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32, #Tss>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
   } -> tensor<32x16xf32>
   return %0 : tensor<32x16xf32>
 }
 
 #BatchedVector = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed-hi" ],
 }>
 // CHECK-LABEL:   func.func @sub_ss_batched(
-// CHECK-SAME:      %[[VAL_0:.*]]: tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>,
-// CHECK-SAME:      %[[VAL_1:.*]]: tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>) -> tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> {
+// CHECK-SAME:      %[[VAL_0:.*]]: tensor<2x3xf64, #{{.*}}>>,
+// CHECK-SAME:      %[[VAL_1:.*]]: tensor<2x3xf64, #{{.*}}>>) -> tensor<2x3xf64, #{{.*}}>> {
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 2 : index
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
-// CHECK-DAG:       %[[VAL_5:.*]] = bufferization.alloc_tensor() : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 1 : index} : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
-// CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 1 : index} : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.coordinates %[[VAL_1]] {level = 1 : index} : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
-// CHECK:           %[[VAL_12:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_3]] to %[[VAL_2]] step %[[VAL_4]] iter_args(%[[VAL_14:.*]] = %[[VAL_5]]) -> (tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK-DAG:       %[[VAL_5:.*]] = bufferization.alloc_tensor() : tensor<2x3xf64, #{{.*}}>>
+// CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<2x3xf64, #{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_7:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 1 : index} : tensor<2x3xf64, #{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<2x3xf64, #{{.*}}>> to memref<?xf64>
+// CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 1 : index} : tensor<2x3xf64, #{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.coordinates %[[VAL_1]] {level = 1 : index} : tensor<2x3xf64, #{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<2x3xf64, #{{.*}}>> to memref<?xf64>
+// CHECK:           %[[VAL_12:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_3]] to %[[VAL_2]] step %[[VAL_4]] iter_args(%[[VAL_14:.*]] = %[[VAL_5]]) -> (tensor<2x3xf64, #{{.*}}>>) {
 // CHECK:             %[[VAL_15:.*]] = arith.muli %[[VAL_13]], %[[VAL_2]] : index
 // CHECK:             %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_15]]] : memref<?xindex>
 // CHECK:             %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_4]] : index
 // CHECK:             %[[VAL_18:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_17]]] : memref<?xindex>
 // CHECK:             %[[VAL_19:.*]] = arith.muli %[[VAL_13]], %[[VAL_2]] : index
 // CHECK:             %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex>
 // CHECK:             %[[VAL_21:.*]] = arith.addi %[[VAL_19]], %[[VAL_4]] : index
 // CHECK:             %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xindex>
 // CHECK:             %[[VAL_23:.*]]:3 = scf.while (%[[VAL_24:.*]] = %[[VAL_16]], %[[VAL_25:.*]] = %[[VAL_20]], %[[VAL_26:.*]] = %[[VAL_14]])
 // CHECK:               %[[VAL_27:.*]] = arith.cmpi ult, %[[VAL_24]], %[[VAL_18]] : index
 // CHECK:               %[[VAL_28:.*]] = arith.cmpi ult, %[[VAL_25]], %[[VAL_22]] : index
 // CHECK:               %[[VAL_29:.*]] = arith.andi %[[VAL_27]], %[[VAL_28]] : i1
-// CHECK:               scf.condition(%[[VAL_29]]) %[[VAL_24]], %[[VAL_25]], %[[VAL_26]] : index, index, tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:               scf.condition(%[[VAL_29]]) %[[VAL_24]], %[[VAL_25]], %[[VAL_26]] : index, index, tensor<2x3xf64, #{{.*}}>>
 // CHECK:             } do {
-// CHECK:             ^bb0(%[[VAL_30:.*]]: index, %[[VAL_31:.*]]: index, %[[VAL_32:.*]]: tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>):
+// CHECK:             ^bb0(%[[VAL_30:.*]]: index, %[[VAL_31:.*]]: index, %[[VAL_32:.*]]: tensor<2x3xf64, #{{.*}}>>):
 // CHECK:               %[[VAL_33:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_30]]] : memref<?xindex>
 // CHECK:               %[[VAL_34:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_31]]] : memref<?xindex>
 // CHECK:               %[[VAL_35:.*]] = arith.cmpi ult, %[[VAL_34]], %[[VAL_33]] : index
 // CHECK:               %[[VAL_36:.*]] = arith.select %[[VAL_35]], %[[VAL_34]], %[[VAL_33]] : index
 // CHECK:               %[[VAL_37:.*]] = arith.cmpi eq, %[[VAL_33]], %[[VAL_36]] : index
 // CHECK:               %[[VAL_38:.*]] = arith.cmpi eq, %[[VAL_34]], %[[VAL_36]] : index
 // CHECK:               %[[VAL_39:.*]] = arith.andi %[[VAL_37]], %[[VAL_38]] : i1
-// CHECK:               %[[VAL_40:.*]] = scf.if %[[VAL_39]] -> (tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK:               %[[VAL_40:.*]] = scf.if %[[VAL_39]] -> (tensor<2x3xf64, #{{.*}}>>) {
 // CHECK:                 %[[VAL_41:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_30]]] : memref<?xf64>
 // CHECK:                 %[[VAL_42:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_31]]] : memref<?xf64>
 // CHECK:                 %[[VAL_43:.*]] = arith.subf %[[VAL_41]], %[[VAL_42]] : f64
-// CHECK:                 %[[VAL_44:.*]] = sparse_tensor.insert %[[VAL_43]] into %[[VAL_32]]{{\[}}%[[VAL_13]], %[[VAL_36]]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK:                 scf.yield %[[VAL_44]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:                 %[[VAL_44:.*]] = sparse_tensor.insert %[[VAL_43]] into %[[VAL_32]]{{\[}}%[[VAL_13]], %[[VAL_36]]] : tensor<2x3xf64, #{{.*}}>>
+// CHECK:                 scf.yield %[[VAL_44]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:               } else {
 // CHECK:                 %[[VAL_45:.*]] = arith.cmpi eq, %[[VAL_33]], %[[VAL_36]] : index
-// CHECK:                 %[[VAL_46:.*]] = scf.if %[[VAL_45]] -> (tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK:                 %[[VAL_46:.*]] = scf.if %[[VAL_45]] -> (tensor<2x3xf64, #{{.*}}>>) {
 // CHECK:                   %[[VAL_47:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_30]]] : memref<?xf64>
-// CHECK:                   %[[VAL_48:.*]] = sparse_tensor.insert %[[VAL_47]] into %[[VAL_32]]{{\[}}%[[VAL_13]], %[[VAL_36]]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK:                   scf.yield %[[VAL_48]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:                   %[[VAL_48:.*]] = sparse_tensor.insert %[[VAL_47]] into %[[VAL_32]]{{\[}}%[[VAL_13]], %[[VAL_36]]] : tensor<2x3xf64, #{{.*}}>>
+// CHECK:                   scf.yield %[[VAL_48]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:                 } else {
 // CHECK:                   %[[VAL_49:.*]] = arith.cmpi eq, %[[VAL_34]], %[[VAL_36]] : index
-// CHECK:                   %[[VAL_50:.*]] = scf.if %[[VAL_49]] -> (tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK:                   %[[VAL_50:.*]] = scf.if %[[VAL_49]] -> (tensor<2x3xf64, #{{.*}}>>) {
 // CHECK:                     %[[VAL_51:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_31]]] : memref<?xf64>
 // CHECK:                     %[[VAL_52:.*]] = arith.negf %[[VAL_51]] : f64
-// CHECK:                     %[[VAL_53:.*]] = sparse_tensor.insert %[[VAL_52]] into %[[VAL_32]]{{\[}}%[[VAL_13]], %[[VAL_36]]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK:                     scf.yield %[[VAL_53]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:                     %[[VAL_53:.*]] = sparse_tensor.insert %[[VAL_52]] into %[[VAL_32]]{{\[}}%[[VAL_13]], %[[VAL_36]]] : tensor<2x3xf64, #{{.*}}>>
+// CHECK:                     scf.yield %[[VAL_53]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:                   } else {
-// CHECK:                     scf.yield %[[VAL_32]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:                     scf.yield %[[VAL_32]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:                   }
-// CHECK:                   scf.yield %[[VAL_54:.*]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:                   scf.yield %[[VAL_54:.*]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:                 }
-// CHECK:                 scf.yield %[[VAL_55:.*]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:                 scf.yield %[[VAL_55:.*]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:               }
 // CHECK:               %[[VAL_56:.*]] = arith.cmpi eq, %[[VAL_33]], %[[VAL_36]] : index
 // CHECK:               %[[VAL_57:.*]] = arith.addi %[[VAL_30]], %[[VAL_4]] : index
 // CHECK:               %[[VAL_58:.*]] = arith.select %[[VAL_56]], %[[VAL_57]], %[[VAL_30]] : index
 // CHECK:               %[[VAL_59:.*]] = arith.cmpi eq, %[[VAL_34]], %[[VAL_36]] : index
 // CHECK:               %[[VAL_60:.*]] = arith.addi %[[VAL_31]], %[[VAL_4]] : index
 // CHECK:               %[[VAL_61:.*]] = arith.select %[[VAL_59]], %[[VAL_60]], %[[VAL_31]] : index
-// CHECK:               scf.yield %[[VAL_58]], %[[VAL_61]], %[[VAL_62:.*]] : index, index, tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:               scf.yield %[[VAL_58]], %[[VAL_61]], %[[VAL_62:.*]] : index, index, tensor<2x3xf64, #{{.*}}>>
 // CHECK:             } attributes {"Emitted from" = "linalg.generic"}
-// CHECK:             %[[VAL_63:.*]] = scf.for %[[VAL_64:.*]] = %[[VAL_3]] to %[[VAL_18]] step %[[VAL_4]] iter_args(%[[VAL_65:.*]] = %[[VAL_66:.*]]#2)
+// CHECK:             %[[VAL_63:.*]] = scf.for %[[VAL_64:.*]] = %[[VAL_65:.*]]#0 to %[[VAL_18]] step %[[VAL_4]] iter_args(%[[VAL_66:.*]] = %[[VAL_65]]#2)
 // CHECK:               %[[VAL_67:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_64]]] : memref<?xindex>
 // CHECK:               %[[VAL_68:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_64]]] : memref<?xf64>
-// CHECK:               %[[VAL_69:.*]] = sparse_tensor.insert %[[VAL_68]] into %[[VAL_65]]{{\[}}%[[VAL_13]], %[[VAL_67]]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK:               scf.yield %[[VAL_69]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:               %[[VAL_69:.*]] = sparse_tensor.insert %[[VAL_68]] into %[[VAL_66]]{{\[}}%[[VAL_13]], %[[VAL_67]]] : tensor<2x3xf64, #{{.*}}>>
+// CHECK:               scf.yield %[[VAL_69]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:             } {"Emitted from" = "linalg.generic"}
-// CHECK:             %[[VAL_70:.*]] = scf.for %[[VAL_71:.*]] = %[[VAL_3]] to %[[VAL_22]] step %[[VAL_4]] iter_args(%[[VAL_72:.*]] = %[[VAL_73:.*]])
-// CHECK:               %[[VAL_74:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_71]]] : memref<?xindex>
-// CHECK:               %[[VAL_75:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_71]]] : memref<?xf64>
-// CHECK:               %[[VAL_76:.*]] = arith.negf %[[VAL_75]] : f64
-// CHECK:               %[[VAL_77:.*]] = sparse_tensor.insert %[[VAL_76]] into %[[VAL_72]]{{\[}}%[[VAL_13]], %[[VAL_74]]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK:               scf.yield %[[VAL_77]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:             %[[VAL_70:.*]] = scf.for %[[VAL_71:.*]] = %[[VAL_72:.*]]#1 to %[[VAL_22]] step %[[VAL_4]] iter_args(%[[VAL_73:.*]] = %[[VAL_74:.*]])
+// CHECK:               %[[VAL_75:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_71]]] : memref<?xindex>
+// CHECK:               %[[VAL_76:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_71]]] : memref<?xf64>
+// CHECK:               %[[VAL_77:.*]] = arith.negf %[[VAL_76]] : f64
+// CHECK:               %[[VAL_78:.*]] = sparse_tensor.insert %[[VAL_77]] into %[[VAL_73]]{{\[}}%[[VAL_13]], %[[VAL_75]]] : tensor<2x3xf64, #{{.*}}>>
+// CHECK:               scf.yield %[[VAL_78]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:             } {"Emitted from" = "linalg.generic"}
-// CHECK:             scf.yield %[[VAL_78:.*]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:             scf.yield %[[VAL_79:.*]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:           } {"Emitted from" = "linalg.generic"}
-// CHECK:           %[[VAL_79:.*]] = sparse_tensor.load %[[VAL_80:.*]] hasInserts : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
-// CHECK:           return %[[VAL_79]] : tensor<2x3xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:           %[[VAL_80:.*]] = sparse_tensor.load %[[VAL_81:.*]] hasInserts : tensor<2x3xf64, #{{.*}}>>
+// CHECK:           return %[[VAL_80]] : tensor<2x3xf64, #{{.*}}>>
 // CHECK:         }
 func.func @sub_ss_batched(%0: tensor<2x3xf64, #BatchedVector>, %1: tensor<2x3xf64, #BatchedVector>)
                            -> tensor<2x3xf64, #BatchedVector> {
   %2 = bufferization.alloc_tensor() : tensor<2x3xf64, #BatchedVector>
   %3 = linalg.generic #trait2
        ins(%0, %1 : tensor<2x3xf64, #BatchedVector>, tensor<2x3xf64, #BatchedVector>)
        outs(%2 : tensor<2x3xf64, #BatchedVector>) {
     ^bb0(%in: f64, %in_0: f64, %out: f64):

mlir/test/Dialect/SparseTensor/sparse_foreach.mlir

func.func @foreach_print_slice(%A: tensor<4x4xf64, #CSR_SLICE>) {
   return
 }
 
 #BCOO = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed-hi-nu", "singleton" ],
 }>
 
 // CHECK-LABEL:   func.func @foreach_bcoo(
-// CHECK-SAME:      %[[VAL_0:.*]]: tensor<4x4x4xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK-SAME:      %[[VAL_0:.*]]: tensor<4x4x4xf64, #{{.*}}>>) {
 // CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 4 : index
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 2 : index
-// CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<4x4x4xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<4x4x4xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
+// CHECK-DAG:       %[[VAL_5:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<4x4x4xf64, #{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<4x4x4xf64, #{{.*}}>> to memref<?xf64>
 // CHECK:           scf.for %[[VAL_7:.*]] = %[[VAL_2]] to %[[VAL_1]] step %[[VAL_3]] {
 // CHECK:             %[[VAL_8:.*]] = arith.muli %[[VAL_7]], %[[VAL_4]] : index
-// CHECK:             %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_3]] : index
-// CHECK:             %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_9]]] : memref<?xindex>
-// CHECK:             scf.for %[[VAL_11:.*]] = %[[VAL_2]] to %[[VAL_10]] step %[[VAL_3]] {
-// CHECK:               %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf64>
-// CHECK:               "test.use"(%[[VAL_12]]) : (f64) -> ()
-// CHECK:             }
-// CHECK:           }
+// CHECK:             %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_8]]] : memref<?xindex>
+// CHECK:             %[[VAL_10:.*]] = arith.addi %[[VAL_8]], %[[VAL_3]] : index
+// CHECK:             %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_10]]] : memref<?xindex>
+// CHECK:             scf.for %[[VAL_12:.*]] = %[[VAL_9]] to %[[VAL_11]] step %[[VAL_3]] {
+// CHECK:               %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xf64>
+// CHECK:               "test.use"(%[[VAL_13]]) : (f64) -> ()
+// CHECK:             } {"Emitted from" = "sparse_tensor.foreach"}
+// CHECK:           } {"Emitted from" = "sparse_tensor.foreach"}
 // CHECK:           return
+// CHECK:         }
 func.func @foreach_bcoo(%A: tensor<4x4x4xf64, #BCOO>) {
   sparse_tensor.foreach in %A : tensor<4x4x4xf64, #BCOO> do {
   ^bb0(%1: index, %2: index, %3: index,  %v: f64) :
     "test.use" (%v) : (f64) -> ()
   }
   return
 }

mlir/test/Dialect/SparseTensor/sparse_pack.mlir

 
 // CHECK-LABEL:   func.func @sparse_unpack(
 // CHECK-SAME:      %[[VAL_0:.*]]: memref<?xindex>,
 // CHECK-SAME:      %[[VAL_1:.*]]: memref<?xi32>,
 // CHECK-SAME:      %[[VAL_2:.*]]: memref<?xf64>,
 // CHECK-SAME:      %[[VAL_3:.*]]
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 6 : index
 // CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0 : index
-// CHECK-DAG:       memref.dealloc %[[VAL_0]] : memref<?xindex>
 // CHECK:           %[[VAL_6:.*]] = memref.dim %[[VAL_2]], %[[VAL_5]] : memref<?xf64>
 // CHECK:           %[[VAL_7:.*]] = arith.cmpi ugt, %[[VAL_4]], %[[VAL_6]] : index
 // CHECK:           %[[VAL_8:.*]] = scf.if %[[VAL_7]] -> (memref<6xf64>) {
 // CHECK:             %[[VAL_9:.*]] = memref.realloc %[[VAL_2]] : memref<?xf64> to memref<6xf64>
 // CHECK:             scf.yield %[[VAL_9]] : memref<6xf64>
 // CHECK:           } else {
 // CHECK:             %[[VAL_10:.*]] = memref.subview %[[VAL_2]][0] [6] [1] : memref<?xf64> to memref<6xf64>
 // CHECK:             scf.yield %[[VAL_10]] : memref<6xf64>
 // CHECK:           }
 // CHECK:           %[[VAL_11:.*]] = arith.constant 12 : index
 // CHECK:           %[[VAL_12:.*]] = memref.dim %[[VAL_1]], %[[VAL_5]] : memref<?xi32>
 // CHECK:           %[[VAL_13:.*]] = arith.cmpi ugt, %[[VAL_11]], %[[VAL_12]] : index
 // CHECK:           %[[VAL_14:.*]] = scf.if %[[VAL_13]] -> (memref<12xi32>) {
 // CHECK:             %[[VAL_15:.*]] = memref.realloc %[[VAL_1]] : memref<?xi32> to memref<12xi32>
 // CHECK:             scf.yield %[[VAL_15]] : memref<12xi32>
 // CHECK:           } else {
 // CHECK:             %[[VAL_16:.*]] = memref.subview %[[VAL_1]][0] [12] [1] : memref<?xi32> to memref<12xi32>
 // CHECK:             scf.yield %[[VAL_16]] : memref<12xi32>
 // CHECK:           }
 // CHECK:           %[[VAL_17:.*]] = memref.expand_shape %[[VAL_18:.*]] {{\[\[}}0, 1]] : memref<12xi32> into memref<6x2xi32>
 // CHECK:           %[[VAL_19:.*]] = bufferization.to_tensor %[[VAL_20:.*]] : memref<6xf64>
 // CHECK:           %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_17]] : memref<6x2xi32>
 // CHECK:           %[[VAL_22:.*]] = sparse_tensor.storage_specifier
+// CHECK:       memref.dealloc %[[VAL_0]] : memref<?xindex>
 // CHECK:           return %[[VAL_19]], %[[VAL_21]], %[[VAL_22]] : tensor<6xf64>, tensor<6x2xi32>, index
 // CHECK:         }
 func.func @sparse_unpack(%sp: tensor<100x100xf64, #COO>) -> (tensor<6xf64>, tensor<6x2xi32>, index) {
   %d, %i, %nnz = sparse_tensor.unpack %sp : tensor<100x100xf64, #COO>
                                          to tensor<6xf64>, tensor<6x2xi32>, index
   return %d, %i, %nnz : tensor<6xf64>, tensor<6x2xi32>, index
 }

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

 // DEFINE: %{option} = enable-runtime-library=false
 // DEFINE: %{compile} = mlir-opt %s --sparse-compiler=%{option}
 // DEFINE: %{run} = mlir-cpu-runner \
 // DEFINE:  -e entry -entry-point-result=void  \
-// DEFINE:  -shared-libs=%mlir_c_runner_utils | \
+// DEFINE:  -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils | \
 // DEFINE: FileCheck %s
 //
 // RUN: %{compile} | %{run}
 //
 
 // Do the same run, but now with direct IR generation and, if available, VLA
 // vectorization.
 // REDEFINE: %{option} = "enable-runtime-library=false vl=4 enable-arm-sve=%ENABLE_VLA"

 }>
 
 #SortedCOOI32 = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ],
   posWidth = 32,
   crdWidth = 32
 }>
 
+#BCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed-hi-nu", "singleton" ]
+}>
+
 module {
   //
   // Main driver.
   //
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %f0 = arith.constant 0.0 : f64
     %i0 = arith.constant 0 : i32
func.func @entry() {
     %index32 = arith.constant dense<
        [[  1,  2],
         [  5,  6],
         [  7,  8]]
     > : tensor<3x2xi32>
 
     %s4 = sparse_tensor.pack %data, %index : tensor<3xf64>, tensor<3x2xindex>
                                           to tensor<10x10xf64, #SortedCOO>
+    %s5= sparse_tensor.pack %data, %index32 : tensor<3xf64>, tensor<3x2xi32>
+                                           to tensor<10x10xf64, #SortedCOOI32>
+
+    %bdata = arith.constant dense<
+       [[  1.0,  2.0,  3.0],
+        [  4.0,  5.0,  0.0]]
+    > : tensor<2x3xf64>
+
+    %bindex = arith.constant dense<
+      [[[  1,  2],
+        [  5,  6],
+        [  7,  8]],
+       [[  2,  3],
+        [  4,  2],
+        [ 10, 10]]]
+    > : tensor<2x3x2xindex>
+    %bs = sparse_tensor.pack %bdata, %bindex batched_lvls = 1 :
+          tensor<2x3xf64>, tensor<2x3x2xindex> to tensor<2x10x10xf64, #BCOO>
+
     // CHECK:1
     // CHECK-NEXT:2
     // CHECK-NEXT:1
     //
     // CHECK-NEXT:5
     // CHECK-NEXT:6
     // CHECK-NEXT:2
     //
     // CHECK-NEXT:7
     // CHECK-NEXT:8
     // CHECK-NEXT:3
     sparse_tensor.foreach in %s4 : tensor<10x10xf64, #SortedCOO> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
      }
 
-    %s5= sparse_tensor.pack %data, %index32 : tensor<3xf64>, tensor<3x2xi32>
-                                          to tensor<10x10xf64, #SortedCOOI32>
     // CHECK-NEXT:1
     // CHECK-NEXT:2
     // CHECK-NEXT:1
     //
     // CHECK-NEXT:5
     // CHECK-NEXT:6
     // CHECK-NEXT:2
     //
     // CHECK-NEXT:7
     // CHECK-NEXT:8
     // CHECK-NEXT:3
     sparse_tensor.foreach in %s5 : tensor<10x10xf64, #SortedCOOI32> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
      }
 
+    // CHECK-NEXT:1
+    // CHECK-NEXT:2
+    // CHECK-NEXT:3
+    //
+    // CHECK-NEXT:4
+    // CHECK-NEXT:5
+    //
+    // Make sure the trailing zeros are not traversed.
+    // CHECK-NOT: 0
+    sparse_tensor.foreach in %bs : tensor<2x10x10xf64, #BCOO> do {
+      ^bb0(%0: index, %1: index, %2: index, %v: f64) :
+        vector.print %v: f64
+     }
+
     %d, %i, %n = sparse_tensor.unpack %s5 : tensor<10x10xf64, #SortedCOOI32>
                                          to tensor<3xf64>, tensor<3x2xi32>, i32
 
-
-
     // CHECK-NEXT: ( 1, 2, 3 )
     %vd = vector.transfer_read %d[%c0], %f0 : tensor<3xf64>, vector<3xf64>
     vector.print %vd : vector<3xf64>
 
     // CHECK-NEXT: ( ( 1, 2 ), ( 5, 6 ), ( 7, 8 ) )
     %vi = vector.transfer_read %i[%c0, %c0], %i0 : tensor<3x2xi32>, vector<3x2xi32>
     vector.print %vi : vector<3x2xi32>
 
     // CHECK-NEXT: 3
     vector.print %n : i32
 
+
+    %bd, %bi, %bn = sparse_tensor.unpack %bs batched_lvls=1 :
+       tensor<2x10x10xf64, #BCOO> to tensor<2x3xf64>, tensor<2x3x2xindex>, i32
+
+    // CHECK-NEXT: ( ( 1, 2, 3 ), ( 4, 5, 0 ) )
+    %vbd = vector.transfer_read %bd[%c0, %c0], %f0 : tensor<2x3xf64>, vector<2x3xf64>
+    vector.print %vbd : vector<2x3xf64>
+
+    // CHECK-NEXT: ( ( ( 1, 2 ), ( 5, 6 ), ( 7, 8 ) ), ( ( 2, 3 ), ( 4, 2 ), ( 0, 0 ) ) )
+    %vbi = vector.transfer_read %bi[%c0, %c0, %c0], %c0 : tensor<2x3x2xindex>, vector<2x3x2xindex>
+    vector.print %vbi : vector<2x3x2xindex>
+
+    // CHECK-NEXT: 3
+    vector.print %bn : i32
+
     %d1, %i1, %n1 = sparse_tensor.unpack %s4 : tensor<10x10xf64, #SortedCOO>
                                          to tensor<3xf64>, tensor<3x2xindex>, index
+    // FIXME: This should be freed by one-shot-bufferization.
+    bufferization.dealloc_tensor %bd : tensor<2x3xf64>
+    bufferization.dealloc_tensor %bi : tensor<2x3x2xindex>
     return
   }
 }