[mlir][sparse] introduce sparse_tensor::unpack operation

[Peiming]

Summary

An inverse operation of sparse_tenosr::pack

[Peiming]

rebase

[Peiming]

rebase.

[Peiming]

update example code

[aartbik]

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

61–62

please rebase with the Pure and doc change before this

85–86

yeah, new example coming up ;-)

110

indice -> indices

111

an integer

114

an unpacked

114

E.g. on next line

117

I would formulate this differently:

The unpack operation ends the life time of the sparse tensor, and using this after the unpack is undefined behavior

122

use 1.1 and 5.5 or so to make it jump out

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

623

this error will be bit generic when it happens

how about just, must operate on a COO tensor

623

do we have an invalid.mlir for this?

[Peiming]

address comments.

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

623

Added

[aartbik]

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

623

must operate (no s)

[Peiming]

fix typo.

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

 }
 
 /// 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 type is a type for a COO tensor with the last
 /// dimension level type being unique.
-bool isUniqueCOOType(RankedTensorType tp);
+bool isUniqueCOOType(TensorType tp);
 
 /// Returns the starting dimension for a trailing COO region that spans across
 /// at least two dimensions. If no such COO region is found, returns the rank
 /// of the tensor.
 unsigned getCOOStart(SparseTensorEncodingAttr enc);
 
 /// Helpers to setup a COO type.
 RankedTensorType getCOOFromTypeWithOrdering(RankedTensorType src,

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

string description = [{
     sparse_tensor.new %source : !Source to tensor<1024x1024xf64, #CSR>
     ```
   }];
   let assemblyFormat = "(`expand_symmetry` $expandSymmetry^)? $source attr-dict"
                        "`:` type($source) `to` type($result)";
   let hasVerifier = 1;
 }
 
 def SparseTensor_PackOp : SparseTensor_Op<"pack", [Pure]>,
-    Arguments<(ins AnyRankedTensor:$data,
-                   AnyRankedTensor:$indices)>,
+    Arguments<(ins 1DTensorOf<[AnyType]>:$data,
+                   2DTensorOf<[AnySignlessIntegerOrIndex]>:$indices)>,
     Results<(outs AnySparseTensor: $result)> {
   let summary = "Returns a sparse tensor from the given (data, indices) pair";
 
   let description = [{
     Packs the data/indices into a COO sparse tensor. The coordinates in `indices`
     shall not exceed the dimension sizes of the returned sparse tensor. Note
     that the returned tensor must be statically shaped because it is impossible
     to infer the shape from sparse coordinates.
 
     `$indices`: stored via a 2-D tensor of integer elements with shape [N, ndims],
     which specifies the indices of the elements in the sparse tensor that contains
     non-zero values.
 
     `$data`: stored via a 1-D tensor with shape [N], that supplies the corresponding
     values for the indices.
 
     The operation can be used to materialize a sparse tensor from external sources.
     E.g., when passing from Python as two numpy arrays for data and indices.
 
     Example:
     ```mlir
     %data    = arith.constant dense<[ 1.1,   2.2,   3.3 ]> : tensor<3xf64>
     %indices = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
 
     %st = sparse_tensor.pack %data, %indices : 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|
     ```
   }];
 
   let assemblyFormat = "$data `,` $indices attr-dict `:` type($data) `,` type($indices)"
                                                    "`to` type($result)";
 
   let hasVerifier = 1;
 }
 
+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/indices
+    extracted from a COO sparse tensor. Additionally, it also returns an integer value
+    indicating the number of entries in the source tensor.
+
+    The operation can be used to return an unpacked MLIR sparse tensor to frontend.
+    E.g., returning two numpy arrays for data and indices.
+
+    The unpack operation ends the life time of the sparse tensor, and using this
+    after the unpack is undefined behavior.
+
+    Example:
+    ```mlir
+    // input 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|
+    %data, %indices, %nnz = sparse_tensor.unpack %st
+                          : tensor<3x4xf64, #COO>
+                         to tensor<2xf64>, tensor<2x2xindex>, index
+
+    // %data    = arith.constant dense<[ 1.1,   2.2,   3.3 ]> : tensor<3xf64>
+    // %indices = arith.constant dense<[[0,0], [1,2], [1,3]]> : tensor<3x2xindex>
+    // %nnz = 2
+
+    ```
+  }];
+
+  let assemblyFormat = "$tensor attr-dict `:` type($tensor) "
+                                        "`to` type($data) `,` type($indices)`,` type($nnz)";
+
+  let hasVerifier = 1;
+}
+
 def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
   [Pure, SameOperandsAndResultElementType]>,
     Arguments<(ins AnyTensor:$source)>,
     Results<(outs AnyTensor:$dest)> {
   string summary = "Converts between different tensor types";
   string description = [{
     Converts one sparse or dense tensor type to another tensor type. The rank
     of the source and destination types must match exactly, and the dimension

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

if (!isSingletonDim(enc, i))
       return false;
 
   // If isUnique is true, then make sure that the last dimension level is
   // unique, that is, rank == 1 (unique the only compressed) and rank > 1
   // (unique on the last singleton).
   return !isUnique || isUniqueDLT(getDimLevelType(enc, rank - 1));
 }
 
-bool mlir::sparse_tensor::isUniqueCOOType(RankedTensorType tp) {
+bool mlir::sparse_tensor::isUniqueCOOType(TensorType tp) {
   SparseTensorEncodingAttr enc = getSparseTensorEncoding(tp);
   return enc && isCOOType(enc, 0, /*isUnique=*/true);
 }
 
 unsigned mlir::sparse_tensor::getCOOStart(SparseTensorEncodingAttr enc) {
   const unsigned rank = enc.getDimLevelType().size();
   // We only consider COO region with at least two dimensions for the purpose
   // of AOS storage optimization.

 }
 
 LogicalResult NewOp::verify() {
   if (getExpandSymmetry() && getTypeRank(getResult()) != 2)
     return emitOpError("expand_symmetry can only be used for 2D tensors");
   return success();
 }
 
-LogicalResult PackOp::verify() {
-  TensorType dataTp = getData().getType(), idxTp = getIndices().getType();
-  TensorType retTp = getResult().getType();
-
-  if (!isUniqueCOOType(retTp.cast<RankedTensorType>()))
-    return emitError("must be packed into a COO tensor");
+static LogicalResult verifyPackUnPack(Operation *op, TensorType cooTp,
+                                      TensorType dataTp, TensorType idxTp) {
+  if (!isUniqueCOOType(cooTp))
+    return op->emitError("must operate on a COO tensor");
 
-  if (!retTp.hasStaticShape() || !dataTp.hasStaticShape() ||
-      !idxTp.hasStaticShape())
-    return emitError("all input types must be statically shaped");
-
-  if (dataTp.getRank() != 1 || idxTp.getRank() != 2) {
-    return emitError(
-        "requires rank 1 tensor for value and rank 2 tensor for indices");
-  }
-
-  auto enc = getSparseTensorEncoding(retTp);
+  auto enc = getSparseTensorEncoding(cooTp);
   if (idxTp.getElementType() != enc.getIndexType() ||
-      dataTp.getElementType() != retTp.getElementType())
-    return emitError("unmatched type between input and output");
+      dataTp.getElementType() != cooTp.getElementType())
+    return op->emitError("unmatched type between input and output");
 
   auto dNOE = dataTp.getShape()[0];
   auto iNOE = idxTp.getShape()[0];
   if (!ShapedType::isDynamic(dNOE) && !ShapedType::isDynamic(iNOE) &&
       dNOE != iNOE)
-    return emitError("unmatched number of elements in data and indices");
+    return op->emitError("unmatched number of elements in data and indices");
 
   // A tensor<?xNxi32> for indices means the input COO is rank N
   auto inRank = idxTp.getShape()[1];
-  auto ouRank = retTp.getRank();
+  auto ouRank = cooTp.getRank();
   if (!ShapedType::isDynamic(inRank) && inRank != ouRank)
-    return emitError("unmatched rank between input and output");
+    return op->emitError("unmatched rank between input and output");
 
   return success();
 }
 
+LogicalResult PackOp::verify() {
+  TensorType dataTp = getData().getType(), idxTp = getIndices().getType();
+  TensorType retTp = getResult().getType();
+
+  if (!retTp.hasStaticShape() || !dataTp.hasStaticShape() ||
+      !idxTp.hasStaticShape())
+    return emitError("all input types must be statically shaped");
+
+  return verifyPackUnPack(*this, retTp, dataTp, idxTp);
+}
+
+LogicalResult UnpackOp::verify() {
+  TensorType dataTp = getData().getType(), idxTp = getIndices().getType();
+  TensorType srcTp = getTensor().getType();
+
+  return verifyPackUnPack(*this, srcTp, dataTp, idxTp);
+}
+
 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 shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       // Accept size matches between the source and the destination type

mlir/test/Dialect/SparseTensor/invalid.mlir

func.func @non_static_pack_ret(%data: tensor<6xf64>, %index: tensor<6x1xi32>)
   // expected-error@+1 {{all input types must be statically shaped}}
   %0 = sparse_tensor.pack %data, %index : tensor<6xf64>, tensor<6x1xi32>
                                        to tensor<?xf64, #SparseVector>
   return %0 : tensor<?xf64, #SparseVector>
 }
 
 // -----
 
+#DenseVector = #sparse_tensor.encoding<{dimLevelType = ["dense"], indexBitWidth=32}>
+
+func.func @invalid_pack_dense(%data: tensor<6xf64>, %index: tensor<6x1xi32>)
+                            -> tensor<100xf64, #DenseVector> {
+  // expected-error@+1 {{must operate on a COO tensor}}
+  %0 = sparse_tensor.pack %data, %index : tensor<6xf64>, tensor<6x1xi32>
+                                       to tensor<100xf64, #DenseVector>
+  return %0 : tensor<100xf64, #DenseVector>
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
 
 func.func @invalid_pack_data(%data: tensor<6x1xf64>, %index: tensor<6x1xi32>)
                             -> tensor<100xf64, #SparseVector> {
-  // expected-error@+1 {{requires rank 1 tensor for value and rank 2 tensor for indices}}
+  // expected-error@+1 {{'sparse_tensor.pack' op operand #0 must be 1D tensor of any type values}}
   %0 = sparse_tensor.pack %data, %index : tensor<6x1xf64>, tensor<6x1xi32>
                                        to tensor<100xf64, #SparseVector>
   return %0 : tensor<100xf64, #SparseVector>
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
func.func @invalid_pack_type(%data: tensor<6xf64>, %index: tensor<6x2xi32>)
   // expected-error@+1 {{unmatched rank between input and output}}
   %0 = sparse_tensor.pack %data, %index : tensor<6xf64>, tensor<6x2xi32>
                                        to tensor<100xf64, #SparseVector>
   return %0 : tensor<100xf64, #SparseVector>
 }
 
 // -----
 
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
+
+func.func @invalid_unpack_type(%sp: tensor<100xf32, #SparseVector>)
+                            -> (tensor<6xf64>, tensor<6x1xi32>, i32) {
+  // expected-error@+1 {{unmatched type between input and output}}
+  %d, %i, %n = sparse_tensor.unpack %sp :  tensor<100xf32, #SparseVector>
+                                        to tensor<6xf64>, tensor<6x1xi32>, i32
+  return %d, %i, %n : tensor<6xf64>, tensor<6x1xi32>, i32
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
+
+func.func @invalid_unpack_type(%sp: tensor<100xf32, #SparseVector>)
+                            -> (tensor<5xf32>, tensor<6x1xi32>, i32) {
+  // expected-error@+1 {{unmatched number of elements in data and indices}}
+  %d, %i, %n = sparse_tensor.unpack %sp :  tensor<100xf32, #SparseVector>
+                                        to tensor<5xf32>, tensor<6x1xi32>, i32
+  return %d, %i, %n : tensor<5xf32>, tensor<6x1xi32>, i32
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
+
+func.func @invalid_unpack_type(%sp: tensor<100xf32, #SparseVector>)
+                            -> (tensor<6xf32>, tensor<6x2xi32>, i32) {
+  // expected-error@+1 {{unmatched rank between input and output}}
+  %d, %i, %n = sparse_tensor.unpack %sp :  tensor<100xf32, #SparseVector>
+                                        to tensor<6xf32>, tensor<6x2xi32>, i32
+  return %d, %i, %n : tensor<6xf32>, tensor<6x2xi32>, i32
+}
+
+// -----
+
 func.func @invalid_pointers_dense(%arg0: tensor<128xf64>) -> memref<?xindex> {
   // expected-error@+1 {{'sparse_tensor.pointers' op operand #0 must be sparse tensor of any type values, but got 'tensor<128xf64>'}}
   %0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<128xf64> to memref<?xindex>
   return %0 : memref<?xindex>
 }
 
 // -----
 

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>
 }
 
 // -----
 
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=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) {
+  %data, %indices, %nnz = sparse_tensor.unpack %sp : tensor<100xf64, #SparseVector>
+                                                  to tensor<6xf64>, tensor<6x1xi32>, i32
+  return %data, %indices, %nnz : tensor<6xf64>, tensor<6x1xi32>, i32
+}
+
+// -----
+
 #SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 
 // CHECK-LABEL: func @sparse_new_symmetry(
 // CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[T:.*]] = sparse_tensor.new expand_symmetry %[[A]] : !llvm.ptr<i8> to tensor<?x?xf64, #{{.*}}>
 //       CHECK: return %[[T]] : tensor<?x?xf64, #{{.*}}>
 func.func @sparse_new_symmetry(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf64, #SparseMatrix> {
   %0 = sparse_tensor.new expand_symmetry %arg0 : !llvm.ptr<i8> to tensor<?x?xf64, #SparseMatrix>