[mlir][sparse] add init sparse tensor operation

[aartbik]

Summary

This is the first step towards supporting general sparse tensors as output
of operations. The init sparse tensor is used to materialize an empty sparse
tensor of given shape and sparsity into a subsequent computation (similar to
the dense tensor init operation counterpart).

Example:

%c = sparse_tensor.init %d1, %d2 : tensor<?x?xf32, #SparseMatrix>
%0 = linalg.matmul
  ins(%a, %b: tensor<?x?xf32>, tensor<?x?xf32>)
  outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>

[bixia]

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

228–231

So we don't allow "dynamic value" in "sizes" if the corresponding dimension in the tensor type is static, something like this:

%0 = sparse_tensor.init [%d1, %d2] : tensor<?x10xf32, #SparseMatrix>. // not allow? %d1 %d2 aren't ConstantOp

%0 = sparse_tensor.init [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>. // allow

Will we have runtime checking to verify the consistency of the dimensions? I think the needed runtime checking for both of the above cases are similar and we probably should support the first case as well.

[aartbik]

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

228–231

No, we only allow runtime SSA values for the dynamic case (?), similar to e.g. the alloc operation.

I thought about not specifying the static cases at all (ie. 10x?xf32 only takes one argument), similar to alloc, but what I have now feels a bit more readable.

There could be some value in allowing dynamic SSA values on the static case and then asserting that it matches at runtime, but we have no precedent for this. For example, linalg generic operations check static sizes to make sure the shape matches, but no runtime checks are ever generated.

But I am okay to relax the requirements in follow up revisions if you see value (I also suspect we will probably merge the "init" between linalg, tensor, sparse_tensor at some point). In any form, however, I need something like this to get the sparse output started.

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

 
 //===----------------------------------------------------------------------===//
 // Operations.
 //===----------------------------------------------------------------------===//
 
 def SparseTensor_NewOp : SparseTensor_Op<"new", []>,
     Arguments<(ins AnyType:$source)>,
     Results<(outs TensorOf<[AnyType]>:$result)> {
-  string summary = "Constructs a new sparse tensor";
+  string summary = "Materializes a new sparse tensor from given source";
   string description = [{
-    Constructs a sparse tensor value with contents taken from an opaque
-    pointer provided by `source`. For targets that have access to a file
-    system, for example, this pointer may be a filename (or file) of a sparse
+    Materializes a sparse tensor with contents taken from an opaque pointer
+    provided by `source`. For targets that have access to a file system,
+    for example, this pointer may be a filename (or file) of a sparse
     tensor in a particular external storage format. The form of the operation
     is kept deliberately very general to allow for alternative implementations
     in the future, such as pointers to buffers or runnable initialization
-    code. The operation is provided as an anchor that materializes a fully
-    typed sparse tensor values into a computation.
+    code. The operation is provided as an anchor that materializes a properly
+    typed sparse tensor with inital contents into a computation.
 
     Example:
 
     ```mlir
     sparse_tensor.new %source : !Source to tensor<1024x1024xf64, #CSR>
     ```
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
 }
 
+def SparseTensor_InitOp : SparseTensor_Op<"init", []>,
+    Arguments<(ins Variadic<Index>:$sizes)>,
+    Results<(outs AnyTensor:$result)> {
+  string summary = "Materializes an empty sparse tensor";
+  string description = [{
+    Materializes an empty sparse tensor with given shape (either static or dynamic).
+    The operation is provided as an anchor that materializes a properly typed sparse
+    tensor into the output clause of a subsequent operation that yields a sparse tensor
+    as the result.
+
+    Example:
+
+    ```mlir
+    %c = sparse_tensor.init_tensor [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>
+    %0 = linalg.matmul
+      ins(%a, %b: tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
+    ```
+  }];
+  let assemblyFormat = "`[` $sizes `]` attr-dict `:` type($result)";
+}
+
 def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
   [NoSideEffect, SameOperandsAndResultType]>,
     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
     and dimensions of the source and destination types must match exactly,
def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
   let hasFolder = 1;
 }
 
 def SparseTensor_ReleaseOp : SparseTensor_Op<"release", []>,
     Arguments<(ins AnyTensor:$tensor)> {
   string description = [{
     Releases the underlying sparse storage scheme for a tensor that
-    materialized earlier through a `new` operator or a non-trivial
-    `convert` operator with an annotated tensor type as destination.
+    materialized earlier through a `new` operator, `init` operator, or a
+    non-trivial `convert` operator with an annotated tensor type as destination.
     This operation should only be called once for any materialized tensor.
     Also, after this operation, any subsequent `memref` querying operation
     on the tensor returns undefined results.
 
     Example:
 
     ```mlir
     sparse_tensor.release %tensor : tensor<1024x1024xf64, #CSR>
let description = [{
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
 }
 
 def SparseTensor_ToTensorOp : SparseTensor_Op<"tensor", [NoSideEffect]>,
     Arguments<(ins Variadic<AnyStridedMemRefOfRank<1>>:$memrefs)>,
     Results<(outs AnyTensor:$result)> {
-  let summary = "Reconstructs tensor from arrays(s)";
+  let summary = "Rematerializes tensor from arrays(s)";
   let description = [{
-    Reconstructs the sparse tensor from the sparse storage scheme array(s).
+    Rematerializes the sparse tensor from the sparse storage scheme array(s).
     This is similar to the `memref.load` operation in the sense that it
     provides a bridge between a bufferized world view and a tensor world
     view. Unlike the `memref.load` operation, however, this sparse operation
     is used only temporarily to maintain a correctly typed intermediate
     representation during progressive bufferization. Eventually the operation
     is folded away.
 
     The input arrays are defined unambigously by the sparsity annotations

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

 }
 
 static LogicalResult verify(NewOp op) {
   if (!getSparseTensorEncoding(op.result().getType()))
     return op.emitError("expected a sparse tensor result");
   return success();
 }
 
+static LogicalResult verify(InitOp op) {
+  if (!getSparseTensorEncoding(op.result().getType()))
+    return op.emitError("expected a sparse tensor result");
+  RankedTensorType ttp = op.getType().cast<RankedTensorType>();
+  unsigned rank = ttp.getRank();
+  if (rank != op.sizes().size())
+    return op.emitError("unexpected mismatch between tensor rank and sizes: ")
+           << rank << " vs. " << op.sizes().size();
+  auto shape = ttp.getShape();
+  for (unsigned i = 0; i < rank; i++) {
+    if (shape[i] == ShapedType::kDynamicSize)
+      continue;
+    auto constantOp = op.sizes()[i].getDefiningOp<ConstantOp>();
+    if (!constantOp ||
+        constantOp.getValue().cast<IntegerAttr>().getInt() != shape[i])
+      return op.emitError("unexpected mismatch with static dimension size ")
+             << shape[i];
+  }
+  return success();
+}
+
 static LogicalResult verify(ConvertOp op) {
   if (auto tp1 = op.source().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = op.dest().getType().dyn_cast<RankedTensorType>()) {
       assert(tp1.getRank() == tp2.getRank());
       auto shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++) {
         if (shape1[d] != shape2[d])
-          return op.emitError()
-                 << "unexpected conversion mismatch in dimension " << d;
+          return op.emitError("unexpected conversion mismatch in dimension ")
+                 << d;
       }
       return success();
     }
   }
   return op.emitError("unexpected type in convert");
 }
 
 OpFoldResult ConvertOp::fold(ArrayRef<Attribute> operands) {
static LogicalResult verify(ToValuesOp op) {
   MemRefType mtp = op.result().getType().cast<MemRefType>();
   if (ttp.getElementType() != mtp.getElementType())
     return op.emitError("unexpected mismatch in element types");
   return success();
 }
 
 static LogicalResult verify(ToTensorOp op) {
   if (!getSparseTensorEncoding(op.result().getType()))
-    return op.emitError("expected a sparse tensor as result");
+    return op.emitError("expected a sparse tensor result");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Methods.
 //===----------------------------------------------------------------------===//
 
 void SparseTensorDialect::initialize() {

mlir/test/Dialect/SparseTensor/invalid.mlir

 func @invalid_release_dense(%arg0: tensor<4xi32>) {
   // expected-error@+1 {{expected a sparse tensor to release}}
   sparse_tensor.release %arg0 : tensor<4xi32>
   return
 }
 
 // -----
 
+func @invalid_init_dense(%arg0: index, %arg1: index) -> tensor<?x?xf32> {
+  // expected-error@+1 {{expected a sparse tensor result}}
+  %0 = sparse_tensor.init [%arg0, %arg1] : tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
+func @invalid_init_rank(%arg0: index) -> tensor<?xf32, #SparseVector> {
+  // expected-error@+1 {{unexpected mismatch between tensor rank and sizes: 1 vs. 2}}
+  %0 = sparse_tensor.init [%arg0, %arg0] : tensor<?xf32, #SparseVector>
+  return %0 : tensor<?xf32, #SparseVector>
+}
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+func @invalid_init_size() -> tensor<?x10xf32, #SparseMatrix> {
+  %c10 = constant 10 : index
+  %c20 = constant 20 : index
+  // expected-error@+1 {{unexpected mismatch with static dimension size 10}}
+  %0 = sparse_tensor.init [%c10, %c20] : tensor<?x10xf32, #SparseMatrix>
+  return %0 : tensor<?x10xf32, #SparseMatrix>
+}
+
+// -----
+
 func @invalid_pointers_dense(%arg0: tensor<128xf64>) -> memref<?xindex> {
   %c = constant 0 : index
   // expected-error@+1 {{expected a sparse tensor to get pointers}}
   %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64> to memref<?xindex>
   return %0 : memref<?xindex>
 }
 
 // -----
func @mismatch_values_types(%arg0: tensor<?xf64, #SparseVector>) -> memref<?xf32> {
   // expected-error@+1 {{unexpected mismatch in element types}}
   %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf32>
   return %0 : memref<?xf32>
 }
 
 // -----
 
 func @sparse_to_unannotated_tensor(%arg0: memref<?xf64>) -> tensor<16x32xf64> {
-  // expected-error@+1 {{expected a sparse tensor as result}}
+  // expected-error@+1 {{expected a sparse tensor result}}
   %0 = sparse_tensor.tensor %arg0 : memref<?xf64> to tensor<16x32xf64>
   return %0 : tensor<16x32xf64>
 }
 
 // -----
 
 func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
   // expected-error@+1 {{unexpected type in convert}}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 // RUN: mlir-opt %s -split-input-file | mlir-opt | 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 @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>
 }
 
 // -----
 
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_init()
+//   CHECK-DAG: %[[C16:.*]] = constant 16 : index
+//   CHECK-DAG: %[[C32:.*]] = constant 32 : index
+//       CHECK: %[[T:.*]] = sparse_tensor.init[%[[C16]], %[[C32]]] : tensor<?x32xf64, #{{.*}}>
+//       CHECK: return %[[T]] : tensor<?x32xf64, #{{.*}}>
+func @sparse_init() -> tensor<?x32xf64, #SparseMatrix> {
+  %d1 = constant 16 : index
+  %d2 = constant 32 : index
+  %0 = sparse_tensor.init [%d1, %d2] : tensor<?x32xf64, #SparseMatrix>
+  return %0 : tensor<?x32xf64, #SparseMatrix>
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_release(
 // CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>
 //       CHECK: sparse_tensor.release %[[A]] : tensor<128xf64, #{{.*}}>
 //       CHECK: return
 func @sparse_release(%arg0: tensor<128xf64, #SparseVector>) {
   sparse_tensor.release %arg0 : tensor<128xf64, #SparseVector>