[mlir][sparse] introduce sparse_tensor.pack operation

[Peiming]

revert unintended changes

[Peiming]

fix typo

[aartbik]

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

65

use brackets for readability

.... for the given (data, indices) pair.

68

Describe data/indices in more detail

data : array wit the value of each stored element
indices: array with the indices of each stored element, in AOS format

Perhaps even give example like matrix with elements a_00, a_41 is stored as
data a_00 a_41
indices 0 0 4 1

69

shall not exceed (no s)

75

do we allow ?x?xf64 as return format

probably not, since we cannot guess size from just data, since matrix may not be filled

so we probably need to check that and document that

EDIT: I see below you enforce that, so just document that

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

453

period at end

640

ranke -> rank

671

As first glance, it did not feel right to have rewriting in IR, shouldn't this be in transforms somehow? But looking around, all other dialects define canonicaliazation in this file, so good to go ;-)

[Peiming]

address comments.

[Peiming]

only accept COO as the result encoding.

[Peiming]

update doc.

[aartbik]

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

89

COO?

[Peiming]

update doc

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

 /// dimension level type being unique.
 bool isUniqueCOOType(RankedTensorType 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,
+                                            AffineMap ordering, bool ordered);
+
+RankedTensorType getCOOFromType(RankedTensorType src, bool ordered);
+
 //
 // Dimension level types.
 //
 
 // MSVC does not allow this function to be constexpr, because
 // `SparseTensorEncodingAttr::operator bool` isn't declared constexpr.
 // And therefore all functions calling it cannot be constexpr either.
 // TODO: since Clang does allow these to be constexpr, perhaps we should

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">,
+    Arguments<(ins AnyRankedTensor:$data,
+                   AnyRankedTensor:$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     , 5     ]> : tensor<3xf64>
+    %indices = arith.constant dense<[[0,  0],[1,  2]]> : tensor<3x2xindex>
+
+    %st = sparse_tensor.pack %data, %indices : tensor<6xf64>, tensor<6x2xi32
+                                            to tensor<100x100xf64, #COO>
+    // %st = [[1, 0, 0, 0],
+    //        [0, 0, 5, 0],
+    //        [0, 0, 0, 0]]
+    ```
+  }];
+
+  let assemblyFormat = "$data `,` $indices attr-dict `:` type($data) `,` type($indices)"
+                                                   "`to` type($result)";
+
+  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

 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
unsigned mlir::sparse_tensor::getCOOStart(SparseTensorEncodingAttr enc) {
   if (rank > 1)
     for (unsigned r = 0; r < rank - 1; r++)
       if (isCOOType(enc, r, /*isUnique=*/false))
         return r;
 
   return rank;
 }
 
+// Helpers to setup a COO type.
+RankedTensorType sparse_tensor::getCOOFromTypeWithOrdering(RankedTensorType src,
+                                                           AffineMap ordering,
+                                                           bool ordered) {
+  auto *ctx = src.getContext();
+  auto rank = src.getRank();
+  SmallVector<DimLevelType> dims;
+
+  // An unordered and non-unique compressed dim at beginning.
+  // If this is also the last dimension, then it is unique.
+  dims.push_back(*getDimLevelType(LevelFormat::Compressed, ordered, rank == 1));
+  if (rank > 1) {
+    // TODO: it is actually ordered at the level for ordered input.
+    // Followed by unordered non-unique n-2 singleton levels.
+    std::fill_n(std::back_inserter(dims), rank - 2,
+                *getDimLevelType(LevelFormat::Singleton, ordered, false));
+    // Ends by a unique singleton level unless the tensor rank is 1.
+    dims.push_back(*getDimLevelType(LevelFormat::Singleton, ordered, true));
+  }
+
+  SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(src);
+  // TODO: Maybe pick the bitwidth based on input/output tensors (probably the
+  // largest one among them) in the original operation instead of using the
+  // default value.
+  unsigned pointerBitWidth = encSrc ? encSrc.getPointerBitWidth() : 0;
+  unsigned indexBitWidth = encSrc ? encSrc.getIndexBitWidth() : 0;
+  auto enc = SparseTensorEncodingAttr::get(ctx, dims, ordering, AffineMap(),
+                                           pointerBitWidth, indexBitWidth);
+  return RankedTensorType::get(src.getShape(), src.getElementType(), enc);
+}
+
+RankedTensorType sparse_tensor::getCOOFromType(RankedTensorType src,
+                                               bool ordered) {
+  return getCOOFromTypeWithOrdering(
+      src, AffineMap::getMultiDimIdentityMap(src.getRank(), src.getContext()),
+      ordered);
+}
+
 uint64_t mlir::sparse_tensor::toOrigDim(SparseTensorEncodingAttr enc,
                                         uint64_t d) {
   if (enc) {
     auto order = enc.getDimOrdering();
     if (order) {
       assert(order.isPermutation());
       return order.getDimPosition(d);
     }

 }
 
 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");
+
+  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);
+  if (idxTp.getElementType() != enc.getIndexType() ||
+      dataTp.getElementType() != retTp.getElementType())
+    return 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");
+
+  // A tensor<?xNxi32> for indices means the input COO is rank N
+  auto inRank = idxTp.getShape()[1];
+  auto ouRank = retTp.getRank();
+  if (!ShapedType::isDynamic(inRank) && inRank != ouRank)
+    return emitError("unmatched rank between input and output");
+
+  return success();
+}
+
 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
       // (e.g. 10 vs. 10, 10 vs. ?, or ? vs. ?), but reject direct mismatches or
       // matches that would need a runtime assert (e.g. 10 vs. 20 or ? vs. 10).
       for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++)
         if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamic)
           return emitError("unexpected conversion mismatch in dimension ") << d;
       return success();
     }
   }
   return emitError("unexpected type in convert");
 }
 
 OpFoldResult ConvertOp::fold(FoldAdaptor adaptor) {
   Type dstType = getType();
   // Fold trivial dense-to-dense convert and leave trivial sparse-to-sparse
   // convert for codegen to remove. This is because we use trivial
   // sparse-to-sparse convert to tell bufferization that the sparse codegen
   // will expand the tensor buffer into sparse tensor storage.
   if (!getSparseTensorEncoding(dstType) && dstType == getSource().getType())
     return getSource();

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

for (const auto &d : enumerate(stp.getShape())) {
     else
       dim = constantIndex(builder, loc, d.value());
     sizes.push_back(dim);
   }
 }
 
 // TODO: The dim level property of the COO type relies on input tensors, the
 // shape relies on the output tensor
-// Helpers to setup a COO type.
 static RankedTensorType
 getUnorderedCOOFromTypeWithOrdering(RankedTensorType src, AffineMap ordering) {
-  auto *ctx = src.getContext();
-  auto rank = src.getRank();
-  SmallVector<DimLevelType> dims;
-
-  // An unordered and non-unique compressed dim at beginning.
-  dims.push_back(DimLevelType::CompressedNuNo);
-
-  if (rank > 1) {
-    // TODO: it is actually ordered at the level for ordered input.
-    // Followed by unordered non-unique n-2 singleton levels.
-    std::fill_n(std::back_inserter(dims), rank - 2,
-                DimLevelType::SingletonNuNo);
-    // TODO: only if all the inputs (for concatentate) are unique at the last
-    // level should the COO has a unique level at the end. Ends by a unordered
-    // unique singleton level unless the tensor rank is 1.
-    dims.push_back(DimLevelType::SingletonNo);
-  }
-  SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(src);
-  // TODO: Maybe pick the bitwidth based on input/output tensors (probably the
-  // largest one among them) in the original operation instead of using the
-  // default value.
-  unsigned pointerBitWidth = encSrc ? encSrc.getPointerBitWidth() : 0;
-  unsigned indexBitWidth = encSrc ? encSrc.getIndexBitWidth() : 0;
-  auto enc = SparseTensorEncodingAttr::get(ctx, dims, ordering, AffineMap(),
-                                           pointerBitWidth, indexBitWidth);
-  return RankedTensorType::get(src.getShape(), src.getElementType(), enc);
+  return getCOOFromTypeWithOrdering(src, ordering, false);
 }
 
 static RankedTensorType getUnorderedCOOFromType(RankedTensorType src) {
-  return getUnorderedCOOFromTypeWithOrdering(
-      src, AffineMap::getMultiDimIdentityMap(src.getRank(), src.getContext()));
+  return getCOOFromType(src, false);
 }
 
 /// Collects the dynamic dimension sizes for `tp` with the assumption that
 /// `sizes` are the dimension sizes for the type. Stores the dynamic dimension
 /// sizes to dynSizes.
 static void getDynamicSizes(RankedTensorType tp,
                             const SmallVectorImpl<Value> &sizes,
                             SmallVectorImpl<Value> &dynSizes) {

mlir/test/Dialect/SparseTensor/invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 func.func @invalid_new_dense(%arg0: !llvm.ptr<i8>) -> tensor<32xf32> {
   // expected-error@+1 {{'sparse_tensor.new' op result #0 must be sparse tensor of any type values, but got 'tensor<32xf32>'}}
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<32xf32>
   return %0 : tensor<32xf32>
 }
 
 // -----
 
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
+
+func.func @non_static_pack_ret(%data: tensor<6xf64>, %index: tensor<6x1xi32>)
+                            -> tensor<?xf64, #SparseVector> {
+  // 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>
+}
+
+// -----
+
+#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}}
+  %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<6x1xi32>)
+                            -> tensor<100xf32, #SparseVector> {
+  // expected-error@+1 {{unmatched type between input and output}}
+  %0 = sparse_tensor.pack %data, %index : tensor<6xf64>, tensor<6x1xi32>
+                                       to tensor<100xf32, #SparseVector>
+  return %0 : tensor<100xf32, #SparseVector>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
+
+func.func @invalid_pack_type(%data: tensor<5xf64>, %index: tensor<6x1xi32>)
+                            -> tensor<100xf64, #SparseVector> {
+  // expected-error@+1 {{unmatched number of elements in data and indices}}
+  %0 = sparse_tensor.pack %data, %index : tensor<5xf64>, 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>)
+                            -> tensor<100xf64, #SparseVector> {
+  // 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>
+}
+
+// -----
+
 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

 // RUN: mlir-opt %s -split-input-file | mlir-opt -split-input-file | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_new(
 // CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[T:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<128xf64, #{{.*}}>
 //       CHECK: return %[[T]] : tensor<128xf64, #{{.*}}>
 func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
   return %0 : tensor<128xf64, #SparseVector>
 }
 
 // -----
 
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"], indexBitWidth=32}>
+
+// CHECK-LABEL: func @sparse_pack(
+// CHECK-SAME: %[[D:.*]]: tensor<6xf64>,
+// CHECK-SAME: %[[I:.*]]: tensor<6x1xi32>)
+//       CHECK: %[[R:.*]] = sparse_tensor.pack %[[D]], %[[I]]
+//       CHECK: return %[[R]] : tensor<100xf64, #{{.*}}>
+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>
+}
+
+// -----
+
 #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>