use heap memory for position buffer allocated for PackOp.

[aartbik]

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

1345–1346

this comment seems out of date now

1448

typo

1485

const

[Peiming]

address comments.

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", [Pure]>,
+def SparseTensor_UnpackOp : SparseTensor_Op<"unpack">,
     Arguments<(ins AnySparseTensor:$tensor)>,
     Results<(outs 1DTensorOf<[AnyType]>:$values,
                   2DTensorOf<[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

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

if (createDeallocs) {
         // Deallocate every buffer used to store the sparse tensor handler.
         rewriter.create<memref::DeallocOp>(loc, input);
     }
     rewriter.eraseOp(op);
     return success();
   }
 
 private:
-  bool createDeallocs;
+  const bool createDeallocs;
 };
 
 /// Sparse codegen rule for tensor rematerialization.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LoadOp op, OpAdaptor adaptor,
foreachFieldAndTypeInSparseTensor(
               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();
             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++) {
+              auto memrefType = MemRefType::get({batchedCount + 1}, posTp);
+              field = rewriter.create<memref::AllocOp>(loc, memrefType);
+              Value c0 = constantIndex(rewriter, loc, 0);
+              genStore(rewriter, loc, c0, field, c0);
+              for (unsigned i = 1; 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)));
+                Value idx = constantIndex(rewriter, loc, i);
+                Value val = constantIndex(rewriter, loc, nse * i);
+                genStore(rewriter, loc, val, field, idx);
               }
-              MemRefType memrefType = MemRefType::get(
-                  tensorType.getShape(), tensorType.getElementType());
-              auto cstPtr = rewriter.create<arith::ConstantOp>(
-                  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);
+              field = rewriter.create<memref::AllocOp>(loc, posMemTp);
               populateCompressedWithHiPosArray(rewriter, loc, batchDimSzs,
                                                field, nse, op);
             }
             break;
           }
           case SparseTensorFieldKind::CrdMemRef: {
             auto tensorType = op.getCoordinates().getType();
             auto memrefType = MemRefType::get(tensorType.getShape(),
matchAndRewrite(PackOp op, OpAdaptor adaptor,
 
     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);
+    }
 
     // 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);
     }
matchAndRewrite(UnpackOp op, OpAdaptor adaptor,
     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();
   }
+
+private:
+  const bool createDeallocs;
 };
 
 struct SparseNewOpConverter : public OpConversionPattern<NewOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NewOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();

 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     bool createSparseDeallocs, bool enableBufferInitialization) {
-  patterns.add<SparsePackOpConverter, SparseUnpackOpConverter,
-               SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
-               SparseCastConverter, SparseExtractSliceConverter,
-               SparseTensorLoadConverter, SparseExpandConverter,
-               SparseCompressConverter, SparseInsertConverter,
+  patterns.add<SparsePackOpConverter, SparseReturnConverter,
+               SparseCallConverter, SparseDimOpConverter, SparseCastConverter,
+               SparseExtractSliceConverter, SparseTensorLoadConverter,
+               SparseExpandConverter, SparseCompressConverter,
+               SparseInsertConverter,
                SparseSliceGetterOpConverter<ToSliceOffsetOp,
                                             StorageSpecifierKind::DimOffset>,
                SparseSliceGetterOpConverter<ToSliceStrideOp,
                                             StorageSpecifierKind::DimStride>,
                SparseToPositionsConverter, SparseToCoordinatesConverter,
                SparseToCoordinatesBufferConverter, SparseToValuesConverter,
                SparseConvertConverter, SparseNewOpConverter,
                SparseNumberOfEntriesConverter>(typeConverter,
                                                patterns.getContext());
-  patterns.add<SparseTensorDeallocConverter>(
+  patterns.add<SparseTensorDeallocConverter, SparseUnpackOpConverter>(
       typeConverter, patterns.getContext(), createSparseDeallocs);
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

mlir/test/Dialect/SparseTensor/sparse_pack.mlir

 // RUN: mlir-opt %s --canonicalize --post-sparsification-rewrite="enable-runtime-library=false" --sparse-tensor-codegen -cse | FileCheck %s
 
 #COO = #sparse_tensor.encoding<{
   dimLevelType = ["compressed-nu", "singleton"],
   crdWidth=32
 }>
 
 // CHECK-LABEL:   func.func @sparse_pack(
 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<6xf64>,
-// CHECK-SAME:      %[[VAL_1:.*]]: tensor<6x2xi32>) -> (memref<?xindex>, memref<?xi32>, memref<?xf64>,
-// CHECK:           %[[VAL_2:.*]] = arith.constant dense<[0, 6]> : tensor<2xindex>
-// CHECK:           %[[VAL_3:.*]] = bufferization.to_memref %[[VAL_2]] : memref<2xindex>
-// CHECK:           %[[VAL_4:.*]] = memref.cast %[[VAL_3]] : memref<2xindex> to memref<?xindex>
-// CHECK:           %[[VAL_5:.*]] = bufferization.to_memref %[[VAL_1]] : memref<6x2xi32>
-// CHECK:           %[[VAL_6:.*]] = memref.collapse_shape %[[VAL_5]] {{\[\[}}0, 1]] : memref<6x2xi32> into memref<12xi32>
-// CHECK:           %[[VAL_7:.*]] = memref.cast %[[VAL_6]] : memref<12xi32> to memref<?xi32>
-// CHECK:           %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_0]] : memref<6xf64>
-// CHECK:           %[[VAL_9:.*]] = memref.cast %[[VAL_8]] : memref<6xf64> to memref<?xf64>
-// CHECK:           %[[VAL_10:.*]] = sparse_tensor.storage_specifier.init :
-// CHECK:           %[[VAL_11:.*]] = arith.constant 6 : index
-// CHECK:           %[[VAL_12:.*]] = arith.constant 100 : index
-// CHECK:           %[[VAL_14:.*]] = sparse_tensor.storage_specifier.set %[[VAL_10]]  lvl_sz at 0 with %[[VAL_12]]
+// CHECK-SAME:      %[[VAL_1:.*]]: tensor<6x2xi32>)
+// CHECK-DAG:       %[[VAL_2:.*]] = memref.alloc() : memref<2xindex>
+// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK-DAG:       memref.store %[[VAL_3]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<2xindex>
+// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 6 : index
+// CHECK-DAG:       memref.store %[[VAL_5]], %[[VAL_2]]{{\[}}%[[VAL_4]]] : memref<2xindex>
+// CHECK:           %[[VAL_6:.*]] = memref.cast %[[VAL_2]] : memref<2xindex> to memref<?xindex>
+// CHECK:           %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_1]] : memref<6x2xi32>
+// CHECK:           %[[VAL_8:.*]] = memref.collapse_shape %[[VAL_7]] {{\[\[}}0, 1]] : memref<6x2xi32> into memref<12xi32>
+// CHECK:           %[[VAL_9:.*]] = memref.cast %[[VAL_8]] : memref<12xi32> to memref<?xi32>
+// CHECK:           %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_0]] : memref<6xf64>
+// CHECK:           %[[VAL_11:.*]] = memref.cast %[[VAL_10]] : memref<6xf64> to memref<?xf64>
+// CHECK:           %[[VAL_12:.*]] = sparse_tensor.storage_specifier.init
+// CHECK:           %[[VAL_13:.*]] = arith.constant 100 : index
+// CHECK:           %[[VAL_14:.*]] = sparse_tensor.storage_specifier.set %[[VAL_12]]  lvl_sz at 0 with %[[VAL_13]]
 // CHECK:           %[[VAL_15:.*]] = arith.constant 2 : index
-// CHECK:           %[[VAL_17:.*]] = sparse_tensor.storage_specifier.set %[[VAL_14]]  pos_mem_sz at 0 with %[[VAL_15]]
-// CHECK:           %[[VAL_19:.*]] = sparse_tensor.storage_specifier.set %[[VAL_17]]  crd_mem_sz at 0 with %[[VAL_11]]
-// CHECK:           %[[VAL_20:.*]] = sparse_tensor.storage_specifier.set %[[VAL_19]]  lvl_sz at 1 with %[[VAL_12]]
-// CHECK:           %[[VAL_21:.*]] = sparse_tensor.storage_specifier.set %[[VAL_20]]  crd_mem_sz at 1 with %[[VAL_11]]
-// CHECK:           %[[VAL_22:.*]] = sparse_tensor.storage_specifier.set %[[VAL_21]]  val_mem_sz with %[[VAL_11]]
-// CHECK:           return %[[VAL_4]], %[[VAL_7]], %[[VAL_9]], %[[VAL_22]] : memref<?xindex>, memref<?xi32>, memref<?xf64>,
+// CHECK:           %[[VAL_16:.*]] = sparse_tensor.storage_specifier.set %[[VAL_14]]  pos_mem_sz at 0 with %[[VAL_15]]
+// CHECK:           %[[VAL_17:.*]] = sparse_tensor.storage_specifier.set %[[VAL_16]]  crd_mem_sz at 0 with %[[VAL_5]]
+// CHECK:           %[[VAL_18:.*]] = sparse_tensor.storage_specifier.set %[[VAL_17]]  lvl_sz at 1 with %[[VAL_13]]
+// CHECK:           %[[VAL_19:.*]] = sparse_tensor.storage_specifier.set %[[VAL_18]]  crd_mem_sz at 1 with %[[VAL_5]]
+// CHECK:           %[[VAL_20:.*]] = sparse_tensor.storage_specifier.set %[[VAL_19]]  val_mem_sz with %[[VAL_5]]
+// CHECK:           return %[[VAL_6]], %[[VAL_9]], %[[VAL_11]], %[[VAL_20]]
 // CHECK:         }
 func.func @sparse_pack(%values: tensor<6xf64>, %coordinates: tensor<6x2xi32>)
                     -> tensor<100x100xf64, #COO> {
   %0 = sparse_tensor.pack %values, %coordinates
      : tensor<6xf64>, tensor<6x2xi32> to tensor<100x100xf64, #COO>
   return %0 : tensor<100x100xf64, #COO>
 }
 
 // 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:.*]]: !sparse_tensor.storage_specifier
-// CHECK:           %[[VAL_4:.*]] = arith.constant 6 : index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 0 : index
+// 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>

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

func.func @entry() {
     // 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
      }
 
-
     %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
 
+    %d1, %i1, %n1 = sparse_tensor.unpack %s4 : tensor<10x10xf64, #SortedCOO>
+                                         to tensor<3xf64>, tensor<3x2xindex>, index
     return
   }
 }