[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 required sparse tensor. The coordinates in
+    `indices` shall not exceeds the dimension size of the returned sparse
+    tensor.
+
+    Example:
+    ```mlir
+    %st = sparse_tensor.pack %data, %indices : tensor<6xf64>, tensor<6x2xi32
+                                            to tensor<100x100xf64, #CSR>
+    ```
+  }];
+
+  let assemblyFormat = "$data `,` $indices attr-dict `:` type($data) `,` type($indices)"
+                                                   "`to` type($result)";
+
+  let hasVerifier = 1;
+  let hasCanonicalizer = 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
     sizes must either match exactly or relax from a static to a dynamic size.
     The sparse encoding of the two types can obviously be completely different.
     The name `convert` was preferred over `cast`, since the operation may incur
     a non-trivial cost.

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 (!retTp.hasStaticShape())
+    return emitError("returned type 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 ranke 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();
+}
+
+static LogicalResult convertNonCOOPack(PackOp packOp,
+                                       PatternRewriter &rewriter) {
+  // TODO: We assume the input indices are all unique. Extends the PackOp to
+  // allow user to specifier whether the input indices is unique/ordered or not
+  auto rtp = packOp.getResult().getType().cast<RankedTensorType>();
+  auto enc = getSparseTensorEncoding(rtp);
+  if (!isCOOType(enc, 0, /*isUnique*/ true)) {
+    // Pack data, indices -> CSR
+    // =>
+    // Pack data, indices -> COO
+    // convert COO -> CSR
+    Location loc = packOp.getLoc();
+    Value data = packOp.getData(), indices = packOp.getIndices();
+
+    auto tp = getCOOFromType(rtp, true);
+    auto coo = rewriter.create<PackOp>(loc, tp, data, indices);
+    rewriter.replaceOpWithNewOp<ConvertOp>(packOp, rtp, coo);
+  }
+
+  return success();
+}
+
+void PackOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.add(&convertNonCOOPack);
+}
+
 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/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 {{returned type 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>