[mlir][sparse] implement lowering rules for sparse_tensor::unpack

[Peiming]

rebase.

[aartbik]

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

102

some rebasing here

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

157

Even though I know what you mean, this comment is not very clear "does not count".
How about rephrasing this a bit?

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

1193

Say in words, Converts Memref back to Tensor.

[Peiming]

address comments.

[aartbik]

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

158

memroy -> memory

[Peiming]

fix typo

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

def SparseTensor_PackOp : SparseTensor_Op<"pack">,
   }];
 
   let assemblyFormat = "$data `,` $indices attr-dict `:` type($data) `,` type($indices)"
                                                    "`to` type($result)";
 
   let hasVerifier = 1;
 }
 
-def SparseTensor_UnpackOp : SparseTensor_Op<"unpack">,
+def SparseTensor_UnpackOp : SparseTensor_Op<"unpack", [Pure]>,
     Arguments<(ins AnySparseTensor:$tensor)>,
     Results<(outs 1DTensorOf<[AnyType]>:$data,
                   2DTensorOf<[AnySignlessIntegerOrIndex]>:$indices,
                   AnySignlessIntegerOrIndex:$nnz)> {
   let summary = "Returns the (data, indices) pair unpacked from the input tensor";
 
   let description = [{
-    Unpack is the inverse operation of `sparse_tensor::pack`. It returns the data/indice        extracted from a COO sparse tensor. Additionaly, it also returns a integer value
+    Unpack is the inverse operation of `sparse_tensor::pack`. It returns the data/indice
+    extracted from a COO sparse tensor. Additionaly, it also returns a integer value
     indicating the number of entries in the source tensor.
 
     The operation can be used to return a unpacked MLIR sparse tensor to frontend. E.g.,
     returning two numpy arrays for data and indices.
 
     The sparse tensor is considered to be dead after unpacking, and it is an undefined
     behavior to continue using it after being unpacked.
 

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

struct PackOpInterface
   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 data/indices 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;
+  }
+};
+
+struct UnpackOpInterface
+    : public BufferizableOpInterface::ExternalModel<UnpackOpInterface,
+                                                    sparse_tensor::UnpackOp> {
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    // Reallocate does not count (similar to InsertOp).
+    return false;
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return false;
+  }
+
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    // Conceptually, UnpackOp equals to a list of toIndices/toValueOp
+    return {};
+  }
 };
 
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
                                                     sparse_tensor::InsertOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     return true;
registry.addExtension(+[](MLIRContext *ctx,
                             sparse_tensor::SparseTensorDialect *dialect) {
     sparse_tensor::ConcatenateOp::attachInterface<ConcatenateOpInterface>(*ctx);
     sparse_tensor::ConvertOp::attachInterface<ConvertOpInterface>(*ctx);
     sparse_tensor::LoadOp::attachInterface<LoadOpInterface>(*ctx);
     sparse_tensor::NewOp::attachInterface<NewOpInterface>(*ctx);
     sparse_tensor::InsertOp::attachInterface<InsertOpInterface>(*ctx);
     sparse_tensor::NumberOfEntriesOp::attachInterface<
         NumberOfEntriesOpInterface>(*ctx);
+    sparse_tensor::PackOp::attachInterface<PackOpInterface>(*ctx);
+    sparse_tensor::UnpackOp::attachInterface<UnpackOpInterface>(*ctx);
     sparse_tensor::ToIndicesBufferOp::attachInterface<
         ToIndicesBufferOpInterface>(*ctx);
     sparse_tensor::ToIndicesOp::attachInterface<ToIndicesOpInterface>(*ctx);
     sparse_tensor::ToPointersOp::attachInterface<ToPointersOpInterface>(*ctx);
     sparse_tensor::ToValuesOp::attachInterface<ToValuesOpInterface>(*ctx);
   });
 }

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

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;
+  LogicalResult
+  matchAndRewrite(UnpackOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    Location loc = op.getLoc();
+    int64_t rank = op.getTensor().getType().getRank();
+
+    assert(isUniqueCOOType(op.getTensor().getType()) &&
+           desc.getFields().size() == 4);
+
+    Value flatBuf = rank == 1 ? desc.getIdxMemRefOrView(rewriter, loc, 0)
+                              : desc.getAOSMemRef();
+    Value dataBuf = desc.getValMemRef();
+
+    // If frontend requests a static buffer, we reallocate the data/indices
+    // to ensure that we meet their need.
+    TensorType dataTp = op.getData().getType();
+    if (dataTp.hasStaticShape()) {
+      dataBuf = rewriter.create<memref::ReallocOp>(
+          loc, MemRefType::get(dataTp.getShape(), dataTp.getElementType()),
+          dataBuf);
+    }
+
+    TensorType indicesTp = op.getIndices().getType();
+    if (indicesTp.hasStaticShape()) {
+      auto len = indicesTp.getShape()[0] * indicesTp.getShape()[1];
+      flatBuf = rewriter.create<memref::ReallocOp>(
+          loc, MemRefType::get({len}, indicesTp.getElementType()), flatBuf);
+    }
+
+    Value idxBuf = rewriter.create<memref::ExpandShapeOp>(
+        loc, MemRefType::get(indicesTp.getShape(), indicesTp.getElementType()),
+        flatBuf, ArrayRef{ReassociationIndices{0, 1}});
+
+    // Memref => Tensor
+    Value data = rewriter.create<bufferization::ToTensorOp>(loc, dataBuf);
+    Value indices = rewriter.create<bufferization::ToTensorOp>(loc, idxBuf);
+    Value nnz = toType(rewriter, loc, desc.getValMemSize(rewriter, loc),
+                       op.getNnz().getType());
+
+    rewriter.replaceOp(op, {data, indices, nnz});
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     bool enableBufferInitialization) {
-  patterns.add<SparsePackOpConverter, SparseReturnConverter,
-               SparseCallConverter, SparseDimOpConverter, SparseCastConverter,
-               SparseTensorDeallocConverter, SparseExtractSliceCoverter,
-               SparseTensorLoadConverter, SparseExpandConverter,
-               SparseCompressConverter, SparseInsertConverter,
-               SparseToPointersConverter, SparseToIndicesConverter,
-               SparseToIndicesBufferConverter, SparseToValuesConverter,
-               SparseConvertConverter, SparseNumberOfEntriesConverter>(
-      typeConverter, patterns.getContext());
+  patterns.add<SparsePackOpConverter, SparseUnpackOpConverter,
+               SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
+               SparseCastConverter, SparseTensorDeallocConverter,
+               SparseExtractSliceCoverter, SparseTensorLoadConverter,
+               SparseExpandConverter, SparseCompressConverter,
+               SparseInsertConverter, SparseToPointersConverter,
+               SparseToIndicesConverter, SparseToIndicesBufferConverter,
+               SparseToValuesConverter, SparseConvertConverter,
+               SparseNumberOfEntriesConverter>(typeConverter,
+                                               patterns.getContext());
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

mlir/test/Dialect/SparseTensor/sparse_pack.mlir

 // CHECK:           return %[[VAL_4]], %[[VAL_7]], %[[VAL_9]], %[[VAL_22]] : memref<?xindex>, memref<?xi32>, memref<?xf64>,
 // CHECK:         }
 func.func @sparse_pack(%data: tensor<6xf64>, %index: tensor<6x2xi32>)
                     -> tensor<100x100xf64, #COO> {
   %0 = sparse_tensor.pack %data, %index : 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:.*]] = memref.realloc %[[VAL_2]] : memref<?xf64> to memref<6xf64>
+// CHECK:           %[[VAL_5:.*]] = memref.realloc %[[VAL_1]] : memref<?xi32> to memref<12xi32>
+// CHECK:           %[[VAL_6:.*]] = memref.expand_shape %[[VAL_5]] {{\[\[}}0, 1]] : memref<12xi32> into memref<6x2xi32>
+// CHECK:           %[[VAL_7:.*]] = bufferization.to_tensor %[[VAL_4]] : memref<6xf64>
+// CHECK:           %[[VAL_8:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<6x2xi32>
+// CHECK:           %[[VAL_9:.*]] = sparse_tensor.storage_specifier.get %[[VAL_3]]  val_mem_sz
+// CHECK:           %[[VAL_10:.*]] = arith.index_cast %[[VAL_9]] : i32 to index
+// CHECK:           return %[[VAL_7]], %[[VAL_8]], %[[VAL_10]] : 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

#SortedCOOI32 = #sparse_tensor.encoding<{
   indexBitWidth = 32
 }>
 
 module {
   //
   // Main driver.
   //
   func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %f0 = arith.constant 0.0 : f64
+    %i0 = arith.constant 0 : i32
     //
     // Initialize a 3-dim dense tensor.
     //
     %data = arith.constant dense<
        [  1.0,  2.0,  3.0]
     > : tensor<3xf64>
 
     %index = arith.constant dense<
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
+
     return
   }
 }