[mlir][sparse] extend pack operation to support packing a batched COO type

[Peiming]

fix typos

[Peiming]

fix bugs.

[aartbik]

can you also
(1) add an invalid.mlir that is specific to batching
(2) add the roundtrip for batching to ops.mlir

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

99

the operation (singular)

[Peiming]

address comments.

[Peiming]

cleanup.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 inline MemRefType getMemRefType(T t) {
   return t.getType().template cast<MemRefType>();
 }
 
 /// Convenience method to get a sparse encoding attribute from a type.
 /// Returns null-attribute for any type without an encoding.
 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
+/// Returns true iff the given sparse tensor encoding attribute has a trailing
+/// COO region starting at the given level.
+bool isCOOType(SparseTensorEncodingAttr enc, Level startLvl, bool isUnique);
+
 /// Returns true iff the given type is a COO type where the last level
 /// is unique.
 bool isUniqueCOOType(Type tp);
 
 /// Returns the starting level for a trailing COO region that spans
 /// at least two levels.  If no such COO region is found, then returns
 /// the level-rank.
 Level getCOOStart(SparseTensorEncodingAttr enc);

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

string description = [{
     ```mlir
     sparse_tensor.new %source : !Source to tensor<1024x1024xf64, #CSR>
     ```
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
 }
 
 def SparseTensor_PackOp : SparseTensor_Op<"pack", [Pure]>,
-    Arguments<(ins 1DTensorOf<[AnyType]>:$values,
-                   2DTensorOf<[AnySignlessIntegerOrIndex]>:$coordinates)>,
+    Arguments<(ins TensorOf<[AnyType]>:$values,
+                   TensorOf<[AnySignlessIntegerOrIndex]>:$coordinates,
+                   OptionalAttr<IndexAttr>:$batched_lvls)>,
     Results<(outs AnySparseTensor: $result)> {
   let summary = "Returns a sparse tensor from the given (values, coordinates) pair";
 
   let description = [{
     Packs the values/coordinates into a COO sparse tensor.  The length
     of `values` must match the outer-length of `coordinates`, since these
     two tensors are "zipped" together.  The `coordinates` argument provides
     level-coords for each value, therefore, the inner-length of `coordinates`
     must match the level-rank of the returned tensor, and each level-coords
     must be valid for the level-sizes of the returned tensor.  Note that
     the returned tensor must be statically shaped because it is impossible
     to infer the dimension-shape from level-coordinates alone.
 
     TODO: The returned tensor is allowed (in principle) to have non-identity
     dimOrdering/higherOrdering mappings.  However, the current implementation
     does not yet support them.
 
     - `coordinates : tensor<NSE x lvlRank x iType>`
       supplies the level-coords for each element in `values`.
     - `values : tensor<NSE x V>`
       supplies the corresponding values for each entry in `coordinates`.
+    - `batched_lvls : optional<index>`
+      supplies the number of leading levels that are batched.
 
     This operation can be used to materialize a sparse tensor from external
     sources; e.g., when passing two numpy arrays from Python.
 
     Example:
 
     ```mlir
     %values      = arith.constant dense<[ 1.1,   2.2,   3.3 ]> : tensor<3xf64>
     %coordinates = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
     %st = sparse_tensor.pack %values, %coordinates
         : tensor<3xf64>, tensor<3x2xindex> to tensor<3x4xf64, #COO>
     // yields 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|
     ```
+
+    If `batched_lvls` is provided, the operation materializes a batched sparse tensor.
+    Example:
+
+    ```mlir
+    %values      = arith.constant dense<[[ 1.1,   2.2,   3.3 ],
+                                         [ 1.2,   2.3,   0.0 ]]> : tensor<2x3xf64>
+    %coordinates = arith.constant dense<[[ [0],   [1],   [2] ],
+                                         [ [1],   [2],   [3] ]> : tensor<2x3x1xindex>
+    %st = sparse_tensor.pack %values, %coordinates batched_lvls=1
+        : tensor<2x3xf64>, tensor<2x3x1xindex> to tensor<2x4xf64, #BCOO>
+    // yields BCOO format |1.1, 2.2, 3.3, 0.0|
+    //      of 2x4 matrix |0.0, 1.2, 2.3, 0.0|
+    ```
+  }];
+
+  let extraClassDeclaration = [{
+    /// Returns the number of leading levels that are batched.
+    unsigned getNumBatchedLvls();
   }];
 
   let assemblyFormat =
-    "$values `,` $coordinates attr-dict"
+    "$values `,` $coordinates (`batched_lvls` `=` $batched_lvls^)? attr-dict"
     "`:` type($values) `,` type($coordinates) `to` type($result)";
 
   let hasVerifier = 1;
 }
 
 def SparseTensor_UnpackOp : SparseTensor_Op<"unpack", [Pure]>,
     Arguments<(ins AnySparseTensor:$tensor)>,
     Results<(outs 1DTensorOf<[AnyType]>:$values,

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

if (auto ttp = type.dyn_cast<RankedTensorType>())
     return ttp.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
   if (auto mdtp = type.dyn_cast<StorageSpecifierType>())
     return mdtp.getEncoding();
   return nullptr;
 }
 
 /// Returns true iff the given sparse tensor encoding attribute has a trailing
 /// COO region starting at the given level.
-static bool isCOOType(SparseTensorEncodingAttr enc, Level startLvl,
-                      bool isUnique) {
-  if (!enc || !enc.isCompressedLvl(startLvl))
+bool mlir::sparse_tensor::isCOOType(SparseTensorEncodingAttr enc,
+                                    Level startLvl, bool isUnique) {
+  if (!enc ||
+      !(enc.isCompressedLvl(startLvl) || enc.isCompressedWithHiLvl(startLvl)))
     return false;
   const Level lvlRank = enc.getLvlRank();
   for (Level l = startLvl + 1; l < lvlRank; ++l)
     if (!enc.isSingletonLvl(l))
       return false;
   // If isUnique is true, then make sure that the last level is unique,
   // that is, lvlRank == 1 (unique the only compressed) and lvlRank > 1
   // (unique on the last singleton).
if (mdKind == StorageSpecifierKind::PosMemSize && enc.isSingletonLvl(l))
           "requested position memory size on a singleton level");
   }
   return success();
 }
 
 static LogicalResult verifyPackUnPack(Operation *op, bool requiresStaticShape,
                                       SparseTensorType tensorTp,
                                       RankedTensorType valuesTp,
-                                      RankedTensorType coordinatesTp) {
+                                      RankedTensorType coordinatesTp,
+                                      IntegerAttr batchedLvls) {
+  unsigned nBatched = batchedLvls ? batchedLvls.getValue().getZExtValue() : 0;
   if (requiresStaticShape && !tensorTp.hasStaticDimShape())
     return op->emitError("the sparse-tensor must have static shape");
   if (!tensorTp.hasEncoding())
     return op->emitError("the sparse-tensor must have an encoding attribute");
   if (!tensorTp.isIdentity())
     return op->emitError("the sparse-tensor must have the identity mapping");
-  if (!isUniqueCOOType(tensorTp))
+  if (!isCOOType(tensorTp.getEncoding(), nBatched, true))
     return op->emitError("the sparse-tensor must have a COO type");
 
-  if (coordinatesTp.getRank() != 2)
-    return op->emitError("coordinates must have rank 2");
+  if (coordinatesTp.getRank() != 2 + nBatched)
+    return op->emitError("coordinates must have rank 2 + batched_lvls");
   if (requiresStaticShape && !coordinatesTp.hasStaticShape())
     return op->emitError("coordinates must have static shape");
   if (coordinatesTp.getElementType() != tensorTp.getCrdType())
     return op->emitError("input/output coordinate-types don't match");
 
-  if (valuesTp.getRank() != 1)
-    return op->emitError("values must have rank 1");
+  if (valuesTp.getRank() != 1 + nBatched)
+    return op->emitError("values must have rank 1 + batched_lvls");
   if (requiresStaticShape && !valuesTp.hasStaticShape())
     return op->emitError("values must have static shape");
   if (valuesTp.getElementType() != tensorTp.getElementType())
     return op->emitError("input/output element-types don't match");
 
-  const auto valuesNSE = valuesTp.getShape()[0];
-  const auto coordsNSE = coordinatesTp.getShape()[0];
+  for (unsigned i = 0; i < nBatched; i++) {
+    const auto valBatch = valuesTp.getShape()[i];
+    const auto crdBatch = coordinatesTp.getShape()[i];
+    if (ShapedType::isDynamic(valBatch) || ShapedType::isDynamic(crdBatch) ||
+        crdBatch != valBatch) {
+      return op->emitError(
+          "values/coordinates batched level sizes don't match statically");
+    }
+  }
+
+  const auto valuesNSE = valuesTp.getShape()[nBatched];
+  const auto coordsNSE = coordinatesTp.getShape()[nBatched];
   if (!ShapedType::isDynamic(valuesNSE) && !ShapedType::isDynamic(coordsNSE) &&
       valuesNSE != coordsNSE)
     return op->emitError("values/coordinates number-of-elements don't match");
 
   // NOTE: We use `getLvlRank` because the `coordinatesTp` is for
   // level-coordinates (cf., the op documentation).
-  const DynSize coordsRank = coordinatesTp.getShape()[1];
+  const DynSize coordsRank = coordinatesTp.getShape()[1 + nBatched];
   const Level tensorRank = tensorTp.getLvlRank();
   // FIXME: replace the `operator!=` with our backported `safelyNE`.
   if (!ShapedType::isDynamic(coordsRank) &&
-      coordsRank != static_cast<DynSize>(tensorRank))
+      coordsRank != static_cast<DynSize>(tensorRank) - nBatched)
     return op->emitError("input/output level-ranks don't match");
 
   return success();
 }
 
 LogicalResult PackOp::verify() {
   const auto valuesTp = getRankedTensorType(getValues());
   const auto coordinatesTp = getRankedTensorType(getCoordinates());
   const auto resTp = getSparseTensorType(getResult());
-  return verifyPackUnPack(*this, true, resTp, valuesTp, coordinatesTp);
+  return verifyPackUnPack(*this, true, resTp, valuesTp, coordinatesTp,
+                          getBatchedLvlsAttr());
+}
+
+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);
+  return verifyPackUnPack(*this, false, srcTp, valuesTp, coordinatesTp,
+                          nullptr);
 }
 
 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

struct PackOpInterface
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     return false;
   }
 
   AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     assert(op->getNumResults() == 1);
-    assert(isUniqueCOOType(op->getResultTypes()[0].cast<RankedTensorType>()));
     // PackOp reuses the input tensors as values/coordinates instead of
     // creating new ones when packing into a COO format.
     return {{op->getOpResult(0), BufferRelation::Equivalent}};
   }
 
   BufferRelation bufferRelation(Operation *oo, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Unknown;

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

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 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 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();
+              Value pred = builder.create<arith::CmpIOp>(
+                  loc, arith::CmpIPredicate::eq, curLen, c0);
+              auto ifOp = builder.create<scf::IfOp>(
+                  loc, TypeRange{builder.getI1Type()}, pred, true);
+              {
+                OpBuilder::InsertionGuard guard(builder);
+                // if len == 0.
+                builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+                builder.create<scf::YieldOp>(loc,
+                                             constantI1(builder, loc, false));
+                // Else branch.
+                builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+                indices.push_back(
+                    builder.create<arith::SubIOp>(loc, curLen, c1));
+                Value val = builder.create<memref::LoadOp>(loc, batV, indices);
+                indices.pop_back();
+                Value cont = builder.create<arith::CmpFOp>(
+                    loc, arith::CmpFPredicate::OEQ, val,
+                    constantZero(builder, loc, val.getType()));
+                builder.create<scf::YieldOp>(loc, cont);
+              }
+              builder.create<scf::ConditionOp>(loc, ifOp.getResults()[0], vs);
+            },
+            [c1](OpBuilder &builder, Location loc, ValueRange vs) {
+              // len --;
+              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);
+        // 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;
   LogicalResult
   matchAndRewrite(PackOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-
+    const unsigned batchedLvls = op.getNumBatchedLvls();
+    unsigned nse = op.getValues().getType().getDimSize(batchedLvls);
     const auto stt = getSparseTensorType(op.getResult());
-    assert(isUniqueCOOType(stt));
+    assert(isCOOType(stt.getEncoding(), batchedLvls, true));
+
+    unsigned batchedCount = 1;
+    SmallVector<unsigned> batchDimSzs;
+    batchDimSzs.reserve(batchedLvls);
+    for (unsigned i = 0; i < batchedLvls; i++) {
+      // Should already be guaranteed by verifier.
+      assert(!ShapedType::isDynamic(stt.getDimShape()[i]));
+      batchedCount *= stt.getDimShape()[i];
+      batchDimSzs.push_back(stt.getDimShape()[i]);
+    }
 
     SmallVector<Value> fields;
     Location loc = op.getLoc();
 
     foreachFieldAndTypeInSparseTensor(
         stt,
-        [&rewriter, &fields, &op, stt,
+        [&rewriter, &fields, &op, &batchDimSzs, nse, batchedCount, stt,
          loc](Type fType, FieldIndex fIdx, SparseTensorFieldKind fKind,
-              Level /*lvl*/, DimLevelType /*dlt*/) -> bool {
+              Level /*lvl*/, DimLevelType dlt) -> bool {
           assert(fields.size() == fIdx);
           Value field;
           switch (fKind) {
           case SparseTensorFieldKind::StorageSpec:
             field = SparseTensorSpecifier::getInitValue(rewriter, loc, stt);
             break;
           case SparseTensorFieldKind::PosMemRef: {
             // TACO-style COO starts with a PosBuffer
             // By creating a constant value for it, we avoid the complexity of
             // memory management.
             const auto posTp = stt.getPosType();
-            auto tensorType = RankedTensorType::get({2}, posTp);
-            auto memrefType = MemRefType::get(tensorType.getShape(),
-                                              tensorType.getElementType());
+            if (isCompressedDLT(dlt)) {
+              RankedTensorType tensorType;
+              SmallVector<Attribute> posAttr;
+              tensorType = RankedTensorType::get({batchedCount + 1}, posTp);
+              posAttr.push_back(IntegerAttr::get(posTp, 0));
+              for (unsigned i = 0; i < batchedCount; i++) {
+                // The postion memref will have values as
+                // [0, nse, 2 * nse, ..., batchedCount * nse]
+                posAttr.push_back(IntegerAttr::get(posTp, nse * (i + 1)));
+              }
+              MemRefType memrefType = MemRefType::get(
+                  tensorType.getShape(), tensorType.getElementType());
               auto cstPtr = rewriter.create<arith::ConstantOp>(
-                loc, tensorType,
-                DenseElementsAttr::get(
-                    tensorType,
-                    ArrayRef<Attribute>{
-                        IntegerAttr::get(posTp, 0),
-                        IntegerAttr::get(
-                            posTp, op.getValues().getType().getShape()[0])}));
-            field = rewriter.create<bufferization::ToMemrefOp>(loc, memrefType,
-                                                               cstPtr);
+                  loc, tensorType, DenseElementsAttr::get(tensorType, posAttr));
+              field = rewriter.create<bufferization::ToMemrefOp>(
+                  loc, memrefType, cstPtr);
+            } else {
+              assert(isCompressedWithHiDLT(dlt) && !batchDimSzs.empty());
+              MemRefType posMemTp = MemRefType::get({batchedCount * 2}, posTp);
+              field = rewriter.create<memref::AllocaOp>(loc, posMemTp);
+              populateCompressedWithHiPosArray(rewriter, loc, batchDimSzs,
+                                               field, nse, op);
+            }
             break;
           }
           case SparseTensorFieldKind::CrdMemRef: {
             auto tensorType = op.getCoordinates().getType();
             auto memrefType = MemRefType::get(tensorType.getShape(),
                                               tensorType.getElementType());
-            auto crdMemRef = rewriter.create<bufferization::ToMemrefOp>(
+            field = rewriter.create<bufferization::ToMemrefOp>(
                 op->getLoc(), memrefType, op.getCoordinates());
-            ReassociationIndices reassociation;
-            for (int i = 0, e = tensorType.getRank(); i < e; i++)
-              reassociation.push_back(i);
 
-            // Flattened the indices buffer to rank 1.
-            field = rewriter.create<memref::CollapseShapeOp>(
-                loc, crdMemRef, ArrayRef<ReassociationIndices>(reassociation));
             break;
           }
           case SparseTensorFieldKind::ValMemRef: {
             auto tensorType = op.getValues().getType();
             auto memrefType = MemRefType::get(tensorType.getShape(),
                                               tensorType.getElementType());
             field = rewriter.create<bufferization::ToMemrefOp>(
                 op->getLoc(), memrefType, op.getValues());
             break;
           }
           }
 
           assert(field);
+          if (auto memrefTp = field.getType().dyn_cast<MemRefType>();
+              memrefTp && memrefTp.getRank() > 1) {
+            ReassociationIndices reassociation;
+            for (int i = 0, e = memrefTp.getRank(); i < e; i++)
+              reassociation.push_back(i);
+            // Flattens the buffer to rank 1. The value buffer might need be
+            // collapsed as well due to batching.
+            field = rewriter.create<memref::CollapseShapeOp>(
+                loc, field, ArrayRef<ReassociationIndices>(reassociation));
+          }
+
           if (fType != field.getType())
             field = rewriter.create<memref::CastOp>(loc, fType, field);
           fields.push_back(field);
           // Returns true to continue the iteration.
           return true;
         });
 
     MutSparseTensorDescriptor desc(stt, fields);

mlir/test/Dialect/SparseTensor/invalid.mlir

 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], crdWidth=32}>
 
 func.func @invalid_pack_data(%values: tensor<6x1xf64>, %coordinates: tensor<6x1xi32>)
                             -> tensor<100xf64, #SparseVector> {
-  // expected-error@+1 {{'sparse_tensor.pack' op operand #0 must be 1D tensor of any type values}}
+  // expected-error@+1 {{values must have rank 1 + batched_lvls}}
   %0 = sparse_tensor.pack %values, %coordinates
      : tensor<6x1xf64>, tensor<6x1xi32> to tensor<100xf64, #SparseVector>
   return %0 : tensor<100xf64, #SparseVector>
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], crdWidth=32}>
func.func @invalid_pack_type(%values: tensor<6xf64>, %coordinates: tensor<6x2xi32>)
   // expected-error@+1 {{input/output level-ranks don't match}}
   %0 = sparse_tensor.pack %values, %coordinates
      : tensor<6xf64>, tensor<6x2xi32> to tensor<100xf64, #SparseVector>
   return %0 : tensor<100xf64, #SparseVector>
 }
 
 // -----
 
+#BCOO = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed-hi"], crdWidth=32}>
+
+func.func @invalid_pack_batched(%values: tensor<2x6xf64>, %coordinates: tensor<3x6x1xi32>)
+                              -> tensor<2x100xf64, #BCOO> {
+  // expected-error@+1 {{values/coordinates batched level sizes don't match statically}}
+  %0 = sparse_tensor.pack %values, %coordinates batched_lvls=1
+     : tensor<2x6xf64>, tensor<3x6x1xi32> to tensor<2x100xf64, #BCOO>
+  return %0 : tensor<2x100xf64, #BCOO>
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], crdWidth=32}>
 
 func.func @invalid_unpack_type(%sp: tensor<100xf32, #SparseVector>)
                             -> (tensor<6xf64>, tensor<6x1xi32>, i32) {
   // expected-error@+1 {{input/output element-types don't match}}
   %values, %coordinates, %nse = sparse_tensor.unpack %sp
      : tensor<100xf32, #SparseVector> to tensor<6xf64>, tensor<6x1xi32>, i32
   return %values, %coordinates, %nse : tensor<6xf64>, tensor<6x1xi32>, i32

mlir/test/Dialect/SparseTensor/roundtrip.mlir

func.func @sparse_pack(%data: tensor<6xf64>, %index: tensor<6x1xi32>)
                             -> tensor<100xf64, #SparseVector> {
   %0 = sparse_tensor.pack %data, %index : tensor<6xf64>, tensor<6x1xi32>
                                        to tensor<100xf64, #SparseVector>
   return %0 : tensor<100xf64, #SparseVector>
 }
 
 // -----
 
+#BCOO = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed-hi"], crdWidth=32}>
+// CHECK-LABEL: func @sparse_pack_batched(
+// CHECK-SAME: %[[D:.*]]: tensor<2x6xf64>,
+// CHECK-SAME: %[[I:.*]]: tensor<2x6x1xi32>)
+//       CHECK: %[[R:.*]] = sparse_tensor.pack %[[D]], %[[I]] batched_lvls = 1
+//       CHECK: return %[[R]] : tensor<2x100xf64, #{{.*}}>
+func.func @sparse_pack_batched(%values: tensor<2x6xf64>, %coordinates: tensor<2x6x1xi32>)
+                            -> tensor<2x100xf64, #BCOO> {
+  %0 = sparse_tensor.pack %values, %coordinates batched_lvls=1
+     : tensor<2x6xf64>, tensor<2x6x1xi32> to tensor<2x100xf64, #BCOO>
+  return %0 : tensor<2x100xf64, #BCOO>
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], crdWidth=32}>
 
 // CHECK-LABEL: func @sparse_unpack(
 //  CHECK-SAME: %[[T:.*]]: tensor<100xf64, #
 //       CHECK: %[[D:.*]], %[[I:.*]], %[[N:.*]] = sparse_tensor.unpack %[[T]]
 //       CHECK: return %[[D]], %[[I]], %[[N]]
 func.func @sparse_unpack(%sp : tensor<100xf64, #SparseVector>)
                        -> (tensor<6xf64>, tensor<6x1xi32>, i32) {