[mlir][sparse] supports sparse_tensor.pack on libgen path

[Peiming]

add gpu tests.

[aartbik]

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

153

we could start using 1 again ;-)

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

654

strange edit?

1204

typo: coordinate

1217

yeah, a "swap" would be so awesome here but... maybe one day when we control runtime of external source too ;-)

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

665

I was about to comment that you can now remove the static util in SparseTensorCodegen.cpp
but you already did that ;-)

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

216

that point ... (since pointers is plural)

1314

strange edit?

[Peiming]

fix asan.

[Peiming]

revert unintended changes.

[Peiming]

address comments.

[DavidSpickett]

Since this change we have a test failing with a double free on our AArch64 bots: https://lab.llvm.org/buildbot/#/builders/179/builds/7122

This also occurs on our SVE bots, but they were failing to build for a time which masked the issue. The difference is that on the SVE bots the test runs natively, it's using QEMU otherwise (https://lab.llvm.org/buildbot/#/builders/184/builds/5631). Which means the failure isn't down to QEMU.

[DavidSpickett]

The 2 stage SVE bots aren't failing, but that's because they don't run the tests. You need to pass -DMLIR_INCLUDE_INTEGRATION_TESTS=True -DMLIR_RUN_ARM_SVE_TESTS=True to CMake to enable them.

[DavidSpickett]

ASAN results:

=================================================================
==3831480==ERROR: AddressSanitizer: attempting double-free on 0xffff956192a0 in thread T0:
    #0 0xaaaac7242ec4 in __interceptor_free /home/tcwg-buildslave/workspace/tcwg-llvm-release/tcwg-jade-03/final/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:52:3
    #1 0xffff9990d314  (<unknown module>)
    #2 0xffff9990d384  (<unknown module>)
    #3 0xaaaac7f2efc8 in compileAndExecute((anonymous namespace)::Options&, mlir::Operation*, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig, void**, std::unique_ptr<llvm::TargetMachine, std::default_delete<llvm::TargetMachine>>) /home/david.spickett/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:215:3
    #4 0xaaaac7f28754 in compileAndExecuteVoidFunction((anonymous namespace)::Options&, mlir::Operation*, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig, std::unique_ptr<llvm::TargetMachine, std::default_delete<llvm::TargetMachine>>) /home/david.spickett/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:234:10
    #5 0xaaaac7f23c28 in mlir::JitRunnerMain(int, char**, mlir::DialectRegistry const&, mlir::JitRunnerConfig) /home/david.spickett/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:391:23
    #6 0xaaaac72727ec in main /home/david.spickett/llvm-project/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp:33:10
    #7 0xffff9b140e0c in __libc_start_main /build/glibc-RIFKjK/glibc-2.31/csu/../csu/libc-start.c:308:16
    #8 0xaaaac71cc414 in _start (/home/david.spickett/build-llvm-aarch64/bin/mlir-cpu-runner+0x1afe414)

0xffff956192a0 is located 0 bytes inside of 88-byte region [0xffff956192a0,0xffff956192f8)
freed by thread T0 here:
    #0 0xaaaac7242ec4 in __interceptor_free /home/tcwg-buildslave/workspace/tcwg-llvm-release/tcwg-jade-03/final/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:52:3
    #1 0xffff9990d2c4  (<unknown module>)

previously allocated by thread T0 here:
    #0 0xaaaac7243158 in malloc /home/tcwg-buildslave/workspace/tcwg-llvm-release/tcwg-jade-03/final/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:69:3
    #1 0xffff9990d024  (<unknown module>)
    #2 0xaaaac7f28754 in compileAndExecuteVoidFunction((anonymous namespace)::Options&, mlir::Operation*, llvm::StringRef, (anonymous namespace)::CompileAndExecuteConfig, std::unique_ptr<llvm::TargetMachine, std::default_delete<llvm::TargetMachine>>) /home/david.spickett/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:234:10
    #3 0xaaaac7f23c28 in mlir::JitRunnerMain(int, char**, mlir::DialectRegistry const&, mlir::JitRunnerConfig) /home/david.spickett/llvm-project/mlir/lib/ExecutionEngine/JitRunner.cpp:391:23
    #4 0xaaaac72727ec in main /home/david.spickett/llvm-project/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp:33:10
    #5 0xffff9b140e0c in __libc_start_main /build/glibc-RIFKjK/glibc-2.31/csu/../csu/libc-start.c:308:16
    #6 0xaaaac71cc414 in _start (/home/david.spickett/build-llvm-aarch64/bin/mlir-cpu-runner+0x1afe414)

SUMMARY: AddressSanitizer: double-free /home/tcwg-buildslave/workspace/tcwg-llvm-release/tcwg-jade-03/final/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:52:3 in __interceptor_free
==3831480==ABORTING
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /home/david.spickett/build-llvm-aarch64/bin/FileCheck /home/david.spickett/llvm-project/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_libgen.mlir

--

[awarzynski]

@DavidSpickett , thanks for reporting!

To avoid a full a revert: https://github.com/llvm/llvm-project/commit/51eaee3b42f057fc2cc3fae7ec71d44fa1b82057
Issue reported here: https://github.com/llvm/llvm-project/issues/64727

@Peiming and @aartbik - do you see where the "double free" could be coming from?

-Andrzej

[Peiming]

In D158012#4591331, @awarzynski wrote:

@DavidSpickett , thanks for reporting!

To avoid a full a revert: https://github.com/llvm/llvm-project/commit/51eaee3b42f057fc2cc3fae7ec71d44fa1b82057
Issue reported here: https://github.com/llvm/llvm-project/issues/64727

@Peiming and @aartbik - do you see where the "double free" could be coming from?

-Andrzej

Thanks for reporting! I think you partial revert is the right fix: I forget to change the flag for VLA test and my local test environment can not cover that path either.

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

enum class Action : uint32_t {
   kEmpty = 0,
   // newSparseTensor no longer handles `kFromFile=1`, so we leave this
   // number reserved to help catch any code that still needs updating.
   kFromCOO = 2,
   kSparseToSparse = 3,
   kEmptyCOO = 4,
   kToCOO = 5,
   kToIterator = 6,
+  kPack = 7,
 };
 
 /// This enum defines all the sparse representations supportable by
 /// the SparseTensor dialect.  We use a lightweight encoding to encode
 /// both the "format" per se (dense, compressed, singleton) as well as
 /// the "properties" (ordered, unique).  The encoding is chosen for
 /// performance of the runtime library, and thus may change in future
 /// versions; consequently, client code should use the predicate functions

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

public:
   ///
   /// Asserts:
   /// * `lvlRank == lvlEnumerator.getTrgRank()`.
   SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
                       uint64_t lvlRank, const DimLevelType *lvlTypes,
                       const uint64_t *lvl2dim,
                       SparseTensorEnumeratorBase<V> &lvlEnumerator);
 
+  /// Constructs a sparse tensor with the given encoding, and initializes
+  /// the contents from the level buffers.  This ctor allocates exactly
+  /// the required amount of overhead storage, not using any heuristics.
+  /// It assumes that the data provided by `lvlBufs` can be directly used to
+  /// interpret the result sparse tensor and performs *NO* integrity test on the
+  /// input data. It also assume that the trailing COO coordinate buffer is
+  /// passed in as a single AoS memory.
+  SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
+                      uint64_t lvlRank, const uint64_t *lvlSizes,
+                      const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
+                      const intptr_t *lvlBufs);
+
   /// Allocates a new empty sparse tensor.  The preconditions/assertions
   /// are as per the `SparseTensorStorageBase` ctor; which is to say,
   /// the `dimSizes` and `lvlSizes` must both be "sizes" not "shapes",
   /// since there's nowhere to reconstruct dynamic sizes from.
   static SparseTensorStorage<P, C, V> *
   newEmpty(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
            const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
            const uint64_t *lvl2dim) {
public:
   // implementation).
   static SparseTensorStorage<P, C, V> *
   newFromSparseTensor(uint64_t dimRank, const uint64_t *dimShape,
                       uint64_t lvlRank, const uint64_t *lvlSizes,
                       const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
                       uint64_t srcRank, const uint64_t *src2lvl,
                       const SparseTensorStorageBase &source);
 
+  /// Allocates a new sparse tensor and initialize it with the data stored level
+  /// buffers directly.
+  ///
+  /// Precondition:
+  /// * as per the `SparseTensorStorageBase` ctor.
+  /// * the data integrity stored in `buffers` is guaranteed by users already.
+  static SparseTensorStorage<P, C, V> *
+  packFromLvlBuffers(uint64_t dimRank, const uint64_t *dimShape,
+                     uint64_t lvlRank, const uint64_t *lvlSizes,
+                     const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
+                     uint64_t srcRank, const uint64_t *src2lvl,
+                     const intptr_t *buffers);
+
   ~SparseTensorStorage() final = default;
 
   /// Partially specialize these getter methods based on template types.
   void getPositions(std::vector<P> **out, uint64_t lvl) final {
     assert(out && "Received nullptr for out parameter");
     ASSERT_VALID_LVL(lvl);
     *out = &positions[lvl];
   }
void fromCOO(const std::vector<Element<V>> &lvlElements, uint64_t lo,
     }
     // Finalize the sparse position structure at this level.
     finalizeSegment(l, full);
   }
 
   /// Finalizes the sparse position structure at this level.
   void finalizeSegment(uint64_t l, uint64_t full = 0, uint64_t count = 1) {
     if (count == 0)
-      return; // Short-circuit, since it'll be a nop.
+      return;                       // Short-circuit, since it'll be a nop.
     const auto dlt = getLvlType(l); // Avoid redundant bounds checking.
     if (isCompressedDLT(dlt)) {
       appendPos(l, coordinates[l].size(), count);
     } else if (isSingletonDLT(dlt)) {
       return; // Nothing to finalize.
     } else {  // Dense dimension.
       ASSERT_DENSE_DLT(dlt);
       const uint64_t sz = getLvlSizes()[l];
#endif
   source.newEnumerator(&lvlEnumerator, lvlRank, lvlSizes, srcRank, src2lvl);
   auto *tensor = new SparseTensorStorage<P, C, V>(
       dimRank, dimSizes.data(), lvlRank, lvlTypes, lvl2dim, *lvlEnumerator);
   delete lvlEnumerator;
   return tensor;
 }
 
 template <typename P, typename C, typename V>
+SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::packFromLvlBuffers(
+    uint64_t dimRank, const uint64_t *dimShape, uint64_t lvlRank,
+    const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
+    const uint64_t *lvl2dim, uint64_t srcRank, const uint64_t *src2lvl,
+    const intptr_t *buffers) {
+  assert(dimShape && "Got nullptr for dimension shape");
+  auto *tensor = new SparseTensorStorage<P, C, V>(
+      dimRank, dimShape, lvlRank, lvlSizes, lvlTypes, lvl2dim, buffers);
+  return tensor;
+}
+
+template <typename P, typename C, typename V>
 SparseTensorStorage<P, C, V>::SparseTensorStorage(
     uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
     const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
     const uint64_t *lvl2dim, bool initializeValuesIfAllDense)
     : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
                           lvl2dim) {
   // Provide hints on capacity of positions and coordinates.
   // TODO: needs much fine-tuning based on actual sparsity; currently
if (isCompressedDLT(dlt)) {
       // This assertion is for future-proofing.
       assert((isDenseDLT(dlt) || isSingletonDLT(dlt)) &&
              "Level is neither dense nor singleton");
     }
     parentSz = assembledSize(parentSz, l);
   }
 }
 
+template <typename P, typename C, typename V>
+SparseTensorStorage<P, C, V>::SparseTensorStorage(
+    uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
+    const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
+    const uint64_t *lvl2dim, const intptr_t *lvlBufs)
+    : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
+                          lvl2dim) {
+  uint64_t trailCOOLen = 0, parentSz = 1, bufIdx = 0;
+  for (uint64_t l = 0; l < lvlRank; l++) {
+    if (!isUniqueLvl(l) && isCompressedLvl(l)) {
+      // A `compressed-nu` level marks the start of trailing COO start level.
+      // Since the coordinate buffer used for trailing COO are passed in as AoS
+      // scheme, and SparseTensorStorage uses a SoA scheme, we can not simply
+      // copy the value from the provided buffers.
+      trailCOOLen = lvlRank - l;
+      break;
+    }
+    assert(!isSingletonLvl(l) &&
+           "Singleton level not following a compressed-nu level");
+    if (isCompressedLvl(l)) {
+      P *posPtr = reinterpret_cast<P *>(lvlBufs[bufIdx++]);
+      C *crdPtr = reinterpret_cast<C *>(lvlBufs[bufIdx++]);
+      // Copies the lvlBuf into the vectors. The buffer can not be simply reused
+      // because the memory passed from users is not necessarily allocated on
+      // heap.
+      positions[l].assign(posPtr, posPtr + parentSz + 1);
+      coordinates[l].assign(crdPtr, crdPtr + positions[l][parentSz]);
+    } else {
+      assert(isDenseLvl(l) && "Level is not dense");
+    }
+    parentSz = assembledSize(parentSz, l);
+  }
+
+  if (trailCOOLen != 0) {
+    uint64_t cooStartLvl = lvlRank - trailCOOLen;
+    assert(!isUniqueLvl(cooStartLvl) && isCompressedLvl(cooStartLvl));
+    P *posPtr = reinterpret_cast<P *>(lvlBufs[bufIdx++]);
+    C *aosCrdPtr = reinterpret_cast<C *>(lvlBufs[bufIdx++]);
+    positions[cooStartLvl].assign(posPtr, posPtr + parentSz + 1);
+    P crdLen = positions[cooStartLvl][parentSz];
+    for (uint64_t l = cooStartLvl; l < lvlRank; l++) {
+      coordinates[l].resize(crdLen);
+      for (uint64_t n = 0; n < crdLen; n++) {
+        coordinates[l][n] = *(aosCrdPtr + (l - cooStartLvl) + n * trailCOOLen);
+      }
+    }
+    parentSz = assembledSize(parentSz, cooStartLvl);
+  }
+
+  V *valPtr = reinterpret_cast<V *>(lvlBufs[bufIdx]);
+  values.assign(valPtr, valPtr + parentSz);
+}
+
 #undef ASSERT_DENSE_DLT
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H

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

inline Value constantDimLevelTypeEncoding(OpBuilder &builder, Location loc,
   return constantI8(builder, loc, static_cast<uint8_t>(dlt));
 }
 
 inline bool isZeroRankedTensorOrScalar(Type type) {
   auto rtp = dyn_cast<RankedTensorType>(type);
   return !rtp || rtp.getRank() == 0;
 }
 
+// Generates code to cast a tensor to a memref.
+TypedValue<BaseMemRefType> genToMemref(OpBuilder &builder, Location loc,
+                                       Value tensor);
+
 /// Infers the result type and generates `ToPositionsOp`.
 Value genToPositions(OpBuilder &builder, Location loc, Value tensor, Level lvl);
 
 /// Infers the result type and generates `ToCoordinatesOp`.  If the
 /// level is within a COO region, the result type is a memref with unknown
 /// stride and offset.  Otherwise, the result type is a memref without
 /// any specified layout.
 Value genToCoordinates(OpBuilder &builder, Location loc, Value tensor,

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

Value sparse_tensor::reshapeValuesToLevels(OpBuilder &builder, Location loc,
   lvlCoords = builder.create<memref::CastOp>(loc, lvlSizesTp, lvlCoords);
   // Finally, create the ReshapeOp.
   const SmallVector<DynSize> resShape(lvlRank, ShapedType::kDynamic);
   const Type elemTp = getMemRefType(valuesBuffer).getElementType();
   const auto resTp = MemRefType::get(resShape, elemTp);
   return builder.create<memref::ReshapeOp>(loc, resTp, valuesBuffer, lvlCoords);
 }
 
+TypedValue<BaseMemRefType>
+sparse_tensor::genToMemref(OpBuilder &builder, Location loc, Value tensor) {
+  auto tTp = llvm::cast<TensorType>(tensor.getType());
+  auto mTp = MemRefType::get(tTp.getShape(), tTp.getElementType());
+  return builder.create<bufferization::ToMemrefOp>(loc, mTp, tensor)
+      .getResult();
+}
+
 Value sparse_tensor::genToPositions(OpBuilder &builder, Location loc,
                                     Value tensor, Level lvl) {
   const auto srcTp = getSparseTensorType(tensor);
   const Type posTp = srcTp.getPosType();
   const Type memTp = get1DMemRefType(posTp, /*withLayout=*/false);
   return builder.create<ToPositionsOp>(loc, memTp, tensor,
                                        builder.getIndexAttr(lvl));
 }

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

if (isCompressedDLT(dlt)) {
         builder.setInsertionPointAfter(loop);
       }
     } else {
       assert(isDenseDLT(dlt) || isSingletonDLT(dlt));
     }
   }
 }
 
-static TypedValue<BaseMemRefType> genToMemref(OpBuilder &builder, Location loc,
-                                              Value tensor) {
-  auto tTp = llvm::cast<TensorType>(tensor.getType());
-  auto mTp = MemRefType::get(tTp.getShape(), tTp.getElementType());
-  return builder.create<bufferization::ToMemrefOp>(loc, mTp, tensor)
-      .getResult();
-}
-
-Value genSliceToSize(OpBuilder &builder, Location loc, Value mem, Value sz) {
+static Value genSliceToSize(OpBuilder &builder, Location loc, Value mem,
+                            Value sz) {
   auto elemTp = llvm::cast<MemRefType>(mem.getType()).getElementType();
   return builder
       .create<memref::SubViewOp>(
           loc, MemRefType::get({ShapedType::kDynamic}, elemTp), mem,
           ValueRange{}, ValueRange{sz}, ValueRange{},
           ArrayRef<int64_t>{0},                    // static offset
           ArrayRef<int64_t>{ShapedType::kDynamic}, // dynamic size
           ArrayRef<int64_t>{1})                    // static stride

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

static void fillDimSizes(OpBuilder &builder, Location loc, SparseTensorType stt,
   const Dimension dimRank = stt.getDimRank();
   out.clear();
   out.reserve(dimRank);
   for (Dimension d = 0; d < dimRank; d++)
     out.push_back(createOrFoldDimCall(builder, loc, stt, tensor, d));
 }
 
 /// Returns an array with the dimension-sizes of the given tensor.
+/// If the *tensor* parameters is null, the tensor type is assumed to have a
+/// static shape.
 static SmallVector<Value> getDimSizes(OpBuilder &builder, Location loc,
-                                      SparseTensorType stt, Value tensor) {
+                                      SparseTensorType stt,
+                                      Value tensor = Value()) {
   SmallVector<Value> out;
   fillDimSizes(builder, loc, stt, tensor, out);
   return out;
 }
 
 /// Populates the array with the dimension-shape of the given
 /// `SparseTensorType`, where dynamic sizes are represented by zero.
 static void fillDimShape(OpBuilder &builder, Location loc, SparseTensorType stt,
static Value genLvlTypesBuffer(OpBuilder &builder, Location loc,
                                SparseTensorType stt) {
   SmallVector<Value> lvlTypes;
   lvlTypes.reserve(stt.getLvlRank());
   for (const auto dlt : stt.getEncoding().getLvlTypes())
     lvlTypes.push_back(constantDimLevelTypeEncoding(builder, loc, dlt));
   return allocaBuffer(builder, loc, lvlTypes);
 }
 
+/// Extracts the bare (aligned) pointers that point to the tensor.
+static Value extractBarePtrFromTensor(OpBuilder &builder, Location loc,
+                                      Value tensor) {
+  auto buf = genToMemref(builder, loc, tensor);
+  return builder.create<memref::ExtractAlignedPointerAsIndexOp>(loc, buf);
+}
+
+/// Generates a temporary buffer for the level-types of the given encoding.
+static Value genLvlPtrsBuffers(OpBuilder &builder, Location loc,
+                               ValueRange lvlTensors, Value valTensor) {
+  SmallVector<Value> lvlBarePtrs;
+  lvlBarePtrs.reserve(lvlTensors.size() + 1);
+  // Passing in lvl buffer pointers.
+  for (const auto lvl : lvlTensors)
+    lvlBarePtrs.push_back(extractBarePtrFromTensor(builder, loc, lvl));
+
+  // Passing in value buffer pointers.
+  lvlBarePtrs.push_back(extractBarePtrFromTensor(builder, loc, valTensor));
+  Value idxPtr = builder.create<memref::ExtractAlignedPointerAsIndexOp>(
+      loc, allocaBuffer(builder, loc, lvlBarePtrs));
+  Value idxCast =
+      builder.create<arith::IndexCastOp>(loc, builder.getI64Type(), idxPtr);
+  return builder.create<LLVM::IntToPtrOp>(loc, getOpaquePointerType(builder),
+                                          idxCast);
+}
+
 /// This class abstracts over the API of `_mlir_ciface_newSparseTensor`:
 /// the "swiss army knife" method of the sparse runtime support library
 /// for materializing sparse tensors into the computation.  This abstraction
 /// reduces the need to make modifications to client code whenever that
 /// API changes.
 class NewCallParams final {
 public:
   /// Allocates the `ValueRange` for the `func::CallOp` parameters,
matchAndRewrite(ConcatenateOp op, OpAdaptor adaptor,
     //    return a
     Location loc = op.getLoc();
     const auto dstTp = getSparseTensorType(op);
     const auto dstEnc = dstTp.getEncoding();
     const Type elemTp = dstTp.getElementType();
     const Dimension concatDim = op.getDimension();
     const Dimension dimRank = dstTp.getDimRank();
 
-    Value dst;     // destination tensor
+    Value dst;         // destination tensor
     Value dstDimToLvl; // destination tensor permutation (if sparse out)
     // A pointer to the value being inserted (if dense => sparse)
     Value elemPtr;
     // Memory that holds the dim-coords for destination tensor (if sparse out)
     Value dstDimCoords;
     // The offset applied to the dimension to be concated (starting from 0)
     Value offset = constantIndex(rewriter, loc, 0);
 
matchAndRewrite(OutOp op, OpAdaptor adaptor,
     SmallString<18> name{"outSparseTensor", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {}, outParams, EmitCInterface::Off);
     genDelCOOCall(rewriter, loc, elemTp, coo);
     rewriter.eraseOp(op);
     return success();
   }
 };
 
+/// Sparse conversion rule for the sparse_tensor.pack operator.
+class SparseTensorPackConverter : public OpConversionPattern<PackOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(PackOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    const Location loc = op->getLoc();
+    const auto dstTp = getSparseTensorType(op.getResult());
+    // PackOps always returns a static shaped tensor result.
+    assert(dstTp.hasStaticDimShape());
+    SmallVector<Value> dimSizes = getDimSizes(rewriter, loc, dstTp);
+    Value dst =
+        NewCallParams(rewriter, loc)
+            .genBuffers(dstTp.withoutDimToLvl(), dimSizes)
+            .genNewCall(Action::kPack,
+                        genLvlPtrsBuffers(rewriter, loc, adaptor.getLevels(),
+                                          adaptor.getValues()));
+    rewriter.replaceOp(op, dst);
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into opaque pointer.
 //===----------------------------------------------------------------------===//
 
 mlir::SparseTensorTypeToPtrConverter::SparseTensorTypeToPtrConverter() {
   addConversion([](Type type) { return type; });
   addConversion(convertSparseTensorTypes);
 }
 
 //===----------------------------------------------------------------------===//
 // 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,
     const SparseTensorConversionOptions &options) {
-  patterns
-      .add<SparseReturnConverter, SparseTensorToDimSizeConverter,
+  patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
                SparseCastConverter, SparseTensorNewConverter,
                SparseReshapeConverter<tensor::ExpandShapeOp>,
                SparseReshapeConverter<tensor::CollapseShapeOp>,
                SparseTensorConcatConverter, SparseTensorAllocConverter,
                SparseTensorDeallocConverter, SparseTensorToPositionsConverter,
-           SparseTensorToCoordinatesConverter, SparseTensorToValuesConverter,
-           SparseNumberOfEntriesConverter, SparseTensorLoadConverter,
-           SparseTensorInsertConverter, SparseTensorExpandConverter,
-           SparseTensorCompressConverter, SparseTensorOutConverter>(
-          typeConverter, patterns.getContext());
+               SparseTensorToCoordinatesConverter,
+               SparseTensorToValuesConverter, SparseNumberOfEntriesConverter,
+               SparseTensorLoadConverter, SparseTensorInsertConverter,
+               SparseTensorExpandConverter, SparseTensorCompressConverter,
+               SparseTensorOutConverter, SparseTensorPackConverter>(
 
   patterns.add<SparseTensorConvertConverter>(typeConverter,
                                              patterns.getContext(), options);
 }

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

case Action::kToCOO: {                                                     \
       return tensor.toCOO(lvlRank, lvlSizes, dimRank, dim2lvl);                \
     }                                                                          \
     case Action::kToIterator: {                                                \
       assert(ptr && "Received nullptr for SparseTensorStorage object");        \
       auto &tensor = *static_cast<SparseTensorStorage<P, C, V> *>(ptr);        \
       auto *coo = tensor.toCOO(lvlRank, lvlSizes, dimRank, dim2lvl);           \
       return new SparseTensorIterator<V>(coo);                                 \
     }                                                                          \
+    case Action::kPack: {                                                      \
+      assert(ptr && "Received nullptr for SparseTensorStorage object");        \
+      intptr_t *buffers = static_cast<intptr_t *>(ptr);                        \
+      return SparseTensorStorage<P, C, V>::packFromLvlBuffers(                 \
+          dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, lvl2dim, dimRank,    \
+          dim2lvl, buffers);                                                   \
+    }                                                                          \
     }                                                                          \
     MLIR_SPARSETENSOR_FATAL("unknown action: %d\n",                            \
                             static_cast<uint32_t>(action));                    \
   }
 
 #define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
 
 // Assume index_type is in fact uint64_t, so that _mlir_ciface_newSparseTensor

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

 
 // REDEFINE: %{sparse_compiler_opts} = enable-runtime-library=false
 // RUN: %{compile} | %{run} | FileCheck %s
 //
 // Do the same run, but now with VLA vectorization.
 // REDEFINE: %{sparse_compiler_opts} = enable-runtime-library=false vl=4
 // RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}
 
-// TODO: Pack only support CodeGen Path
+// TODO: support sparse_tensor.unpack on libgen path.
 
 #SortedCOO = #sparse_tensor.encoding<{
   lvlTypes = [ "compressed-nu", "singleton" ]
 }>
 
 #SortedCOOI32 = #sparse_tensor.encoding<{
   lvlTypes = [ "compressed-nu", "singleton" ],
   posWidth = 32,

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

 // DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
 // DEFINE: %{run_opts} = -e entry -entry-point-result=void
 // DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
 // DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
 //
 // DEFINE: %{env} =
 //--------------------------------------------------------------------------------------------------
 
-// REDEFINE: %{sparse_compiler_opts} = enable-runtime-library=false
 // RUN: %{compile} | %{run} | FileCheck %s
 //
 // Do the same run, but now with VLA vectorization.
 // REDEFINE: %{sparse_compiler_opts} = enable-runtime-library=false vl=4
 // RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}
 
-// TODO: Pack only support CodeGen Path
+// TODO: This is considered to be a short-living tests and should be merged with sparse_pack.mlir
+// after sparse_tensor.unpack is supported on libgen path.
 
 #SortedCOO = #sparse_tensor.encoding<{
   lvlTypes = [ "compressed-nu", "singleton" ]
 }>
 
 #SortedCOOI32 = #sparse_tensor.encoding<{
   lvlTypes = [ "compressed-nu", "singleton" ],
   posWidth = 32,
   crdWidth = 32
 }>
 
 #CSR = #sparse_tensor.encoding<{
   lvlTypes = [ "dense", "compressed" ],
   posWidth = 32,
   crdWidth = 32
 }>
 
-#BCOO = #sparse_tensor.encoding<{
-  lvlTypes = [ "dense", "compressed-hi-nu", "singleton" ]
-}>
+// TODO: "compressed-hi" is not supported by libgen path.
+// #BCOO = #sparse_tensor.encoding<{
+//   lvlTypes = [ "dense", "compressed-hi-nu", "singleton" ]
+//}>
 
 module {
   //
   // Main driver.
   //
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %f0 = arith.constant 0.0 : f64
func.func @entry() {
     > : tensor<3xi32>
 
     %csr_index32 = arith.constant dense<
        [1, 0, 1]
     > : tensor<3xi32>
     %csr= sparse_tensor.pack %csr_data, %csr_pos32, %csr_index32 : tensor<4xf64>, tensor<3xi32>, tensor<3xi32>
                                            to tensor<2x2xf64, #CSR>
 
-    %bdata = arith.constant dense<
-       [  1.0,  2.0,  3.0,  4.0,  5.0,  0.0]
-    > : tensor<6xf64>
-
-    %bpos = arith.constant dense<
-       [0, 3, 3, 5]
-    > : tensor<4xindex>
-
-    %bindex = arith.constant dense<
-      [[  1,  2],
-       [  5,  6],
-       [  7,  8],
-       [  2,  3],
-       [  4,  2],
-       [ 10, 10]]
-    > : tensor<6x2xindex>
-    %bs = sparse_tensor.pack %bdata, %bpos, %bindex :
-          tensor<6xf64>, tensor<4xindex>,  tensor<6x2xindex> to tensor<2x10x10xf64, #BCOO>
-
     // CHECK:1
     // CHECK-NEXT:2
     // CHECK-NEXT:1
     //
     // CHECK-NEXT:5
     // CHECK-NEXT:6
     // CHECK-NEXT:2
     //
func.func @entry() {
     // CHECK-NEXT:3
     sparse_tensor.foreach in %csr : tensor<2x2xf64, #CSR> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
      }
 
-    %d_csr = tensor.empty() : tensor<4xf64>
-    %p_csr = tensor.empty() : tensor<3xi32>
-    %i_csr = tensor.empty() : tensor<3xi32>
-    %rd_csr, %rp_csr, %ri_csr, %ld_csr, %lp_csr, %li_csr = sparse_tensor.unpack %csr : tensor<2x2xf64, #CSR>
-                 outs(%d_csr, %p_csr, %i_csr : tensor<4xf64>, tensor<3xi32>, tensor<3xi32>)
-                 -> tensor<4xf64>, (tensor<3xi32>, tensor<3xi32>), index, (index, index)
-
-    // CHECK-NEXT: ( 1, 2, 3, {{.*}} )
-    %vd_csr = vector.transfer_read %rd_csr[%c0], %f0 : tensor<4xf64>, vector<4xf64>
-    vector.print %vd_csr : vector<4xf64>
-
-    // 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
-     }
-
-    %od = tensor.empty() : tensor<3xf64>
-    %op = tensor.empty() : tensor<2xi32>
-    %oi = tensor.empty() : tensor<3x2xi32>
-    %d, %p, %i, %dl, %pl, %il = sparse_tensor.unpack %s5 : tensor<10x10xf64, #SortedCOOI32>
-                 outs(%od, %op, %oi : tensor<3xf64>, tensor<2xi32>, tensor<3x2xi32>)
-                 -> tensor<3xf64>, (tensor<2xi32>, tensor<3x2xi32>), index, (index, index)
-
-    // 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>
-
-
-    %bod = tensor.empty() : tensor<6xf64>
-    %bop = tensor.empty() : tensor<4xindex>
-    %boi = tensor.empty() : tensor<6x2xindex>
-    %bd, %bp, %bi, %ld, %lp, %li = sparse_tensor.unpack %bs : tensor<2x10x10xf64, #BCOO>
-                    outs(%bod, %bop, %boi : tensor<6xf64>, tensor<4xindex>, tensor<6x2xindex>)
-                    -> tensor<6xf64>, (tensor<4xindex>, tensor<6x2xindex>), index, (index, index)
-
-    // CHECK-NEXT: ( 1, 2, 3, 4, 5, {{.*}} )
-    %vbd = vector.transfer_read %bd[%c0], %f0 : tensor<6xf64>, vector<6xf64>
-    vector.print %vbd : vector<6xf64>
-    // CHECK-NEXT: 5
-    vector.print %ld : index
 
-    // CHECK-NEXT: ( ( 1, 2 ), ( 5, 6 ), ( 7, 8 ), ( 2, 3 ), ( 4, 2 ), ( {{.*}}, {{.*}} ) )
-    %vbi = vector.transfer_read %bi[%c0, %c0], %c0 : tensor<6x2xindex>, vector<6x2xindex>
-    vector.print %vbi : vector<6x2xindex>
-    // CHECK-NEXT: 10
-    vector.print %li : index
+    bufferization.dealloc_tensor %s4  : tensor<10x10xf64, #SortedCOO>
+    bufferization.dealloc_tensor %s5  : tensor<10x10xf64, #SortedCOOI32>
+    bufferization.dealloc_tensor %csr  : tensor<2x2xf64, #CSR>
 
     return
   }
 }