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

[Peiming]

rebase.

[aartbik]

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

102 ↗

(On Diff #496626)

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/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 {
+    // Similar to InsertOp, reallocation is not considered to allocate a new
+    // piece of memory.
+    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}});
+
+    // Converts MemRefs back to Tensors.
+    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
   }
 }