[mlir][sparse] Introduce sparse_tensor.storage operator to create a sparse tensor storage tuple

[Peiming]

format code...

[aartbik]

Super nit pick, but important to me

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

632

nitpick: Since we are just materializing an SSA value, can we use something other than "new' here, like .make_storage(a,b,c) or, or .pack(ab,c) just .storage(a,b,c) or something like that. The term "new" is too much associated with heap allocation, imho, so anything that avoids that confusion seems a bit better.

[aartbik]

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

646

What do you think of a syntax with brackets?

%0 = sparse_tensor.storage(%1, %2) : memref<?xf64>, memref<?xf64>

[Peiming]

address comments from Aart.

[Peiming]

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

632

Okay, I chose it so that it has the same number of characters as storage_get.

I will use .storage(a, b, c)

[aartbik]

Beautiful! Thanks!

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

def SparseTensor_YieldOp : SparseTensor_Op<"yield", [NoSideEffect, Terminator]>,
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Storage Operation. These operations are used internally by
 // sparse tensor codegen to progressively lower sparse tensors.
 //===----------------------------------------------------------------------===//
 
+def SparseTensor_StorageNewOp : SparseTensor_Op<"storage", []>,
+    Arguments<(ins Variadic<AnyType>:$inputs)>,
+    Results<(outs AnyTuple:$result)> {
+  let summary = "Pack a list of value into one sparse tensor storage value";
+  let description = [{
+     Pack a list of value into one sparse tensor storage value (represented as
+     a tuple) at the given index.
+
+     The result tuple elements' type should match the corresponding type in the
+     input array.
+
+     Example:
+
+     ```mlir
+     %0 = sparse_tensor.storage(%1, %2): memref<?xf64>, memref<?xf64>
+                                to tuple<memref<?xf64>, memref<?xf64>>
+     ```
+   }];
+
+  let assemblyFormat = " attr-dict `(` $inputs `)``:` type($inputs) `to` type($result)";
+  let hasVerifier = 1;
+}
+
 def SparseTensor_StorageGetOp : SparseTensor_Op<"storage_get", []>,
     Arguments<(ins AnyTuple:$storage,
                    IndexAttr:$idx)>,
     Results<(outs AnyType:$result)> {
   let summary = "Get the data stored in the sparse tensor storage at the given index";
   let description = [{
      Get the data stored in the sparse tensor storage (represented as a tuple)
      at the given index.

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

LogicalResult YieldOp::verify() {
   return emitOpError(
       "expected parent op to be sparse_tensor unary, binary, or reduce");
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Storage Operation.
 //===----------------------------------------------------------------------===//
 
+LogicalResult StorageNewOp::verify() {
+  auto retTypes = getResult().getType().getTypes();
+  if (retTypes.size() != getInputs().size())
+    return emitError("The number of inputs is inconsistent with output tuple");
+
+  for (auto pair : llvm::zip(getInputs(), retTypes)) {
+    auto input = std::get<0>(pair);
+    auto retTy = std::get<1>(pair);
+
+    if (input.getType() != retTy)
+      return emitError(llvm::formatv("Type mismatch between input (type={0}) "
+                                     "and output tuple element (type={1})",
+                                     input.getType(), retTy));
+  }
+  return success();
+}
+
 LogicalResult StorageGetOp::verify() {
   uint64_t extractIdx = getIdx().getZExtValue();
   auto innerTypeArray = getStorage().getType().getTypes();
   if (extractIdx >= innerTypeArray.size())
     return emitError(llvm::formatv(
         "Out-of-bound access with index={0} on tuple with length={1}",
         extractIdx, innerTypeArray.size()));
 

mlir/test/Dialect/SparseTensor/invalid.mlir

%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
        : tensor<2x4xf64, #DC>,
          tensor<3x3xf64, #DC>,
          tensor<4x4xf64, #DC> to tensor<9x4xf64, #DC>
   return %0 : tensor<9x4xf64, #DC>
 }
 
 // -----
 
+func.func @sparse_storage_new(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: f64) ->
+                               tuple<memref<?xf64>, memref<?xf64>> {
+  // expected-error@+1{{The number of inputs is inconsistent with output}}
+  %0 = sparse_tensor.storage(%arg0, %arg1, %arg2)
+       : memref<?xf64>, memref<?xf64>, f64 to tuple<memref<?xf64>, memref<?xf64>>
+  return %0 : tuple<memref<?xf64>, memref<?xf64>>
+}
+
+// -----
+
+func.func @sparse_storage_new(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: f64) ->
+                               tuple<memref<?xi64>, memref<?xf64>, f64> {
+  // expected-error@+1{{Type mismatch between}}
+  %0 = sparse_tensor.storage(%arg0, %arg1, %arg2)
+       : memref<?xf64>, memref<?xf64>, f64 to tuple<memref<?xi64>, memref<?xf64>, f64>
+  return %0 : tuple<memref<?xi64>, memref<?xf64>, f64>
+}
+
+// -----
+
 func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
   // expected-error@+1{{Out-of-bound access}}
   %0 = sparse_tensor.storage_get %arg0[3]
        : tuple<memref<?xf64>, memref<?xf64>, f64> to
          memref<?xf64>
   return %0 : memref<?xf64>
 }
 

mlir/test/Dialect/SparseTensor/roundtrip.mlir

%0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
        : tensor<2x4xf64, #SparseMatrix>,
          tensor<3x4xf64, #SparseMatrix>,
          tensor<4x4xf64, #SparseMatrix> to tensor<9x4xf64, #SparseMatrix>
   return %0 : tensor<9x4xf64, #SparseMatrix>
 }
 
 // -----
 
+
+// CHECK: func @sparse_storage_new(
+//  CHECK-SAME: %[[A0:.*0]]: memref<?xf64>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xf64>,
+//  CHECK-SAME: %[[A2:.*]]: f64
+//       CHECK: %[[TMP_0:.*]] = sparse_tensor.storage(%[[A0]], %[[A1]], %[[A2]])
+//       CHECK: return %[[TMP_0]] : tuple<memref<?xf64>, memref<?xf64>, f64>
+func.func @sparse_storage_new(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: f64) ->
+                               tuple<memref<?xf64>, memref<?xf64>, f64> {
+  %0 = sparse_tensor.storage(%arg0, %arg1, %arg2)
+       : memref<?xf64>, memref<?xf64>, f64 to tuple<memref<?xf64>, memref<?xf64>, f64>
+  return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
+}
+
+// -----
+
 // CHECK-LABEL: func @sparse_storage_get(
 //  CHECK-SAME:   %[[A0:.*]]: tuple<memref<?xf64>, memref<?xf64>, f64>
 //       CHECK:   %[[TMP0:.*]] = sparse_tensor.storage_get %[[A0]][0] :
 //  CHECK-SAME:     tuple<memref<?xf64>, memref<?xf64>, f64>
 //  CHECK-SAME:     to memref<?xf64>
 //       CHECK:   return %[[TMP0]] : memref<?xf64>
 func.func @sparse_storage_get(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>) -> memref<?xf64> {
   %0 = sparse_tensor.storage_get %arg0[0]
        : tuple<memref<?xf64>, memref<?xf64>, f64> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
-// ----
+// -----
 
 // CHECK-LABEL: func @sparse_storage_set(
 //  CHECK-SAME:   %[[A0:.*]]: tuple<memref<?xf64>, memref<?xf64>, f64>,
 //  CHECK-SAME:   %[[A1:.*]]: memref<?xf64>
 //       CHECK:   %[[TMP0:.*]] = sparse_tensor.storage_set %[[A0]][0], %[[A1]] :
 //  CHECK-SAME:     tuple<memref<?xf64>, memref<?xf64>, f64>,
 //  CHECK-SAME:     memref<?xf64>
 //  CHECK-SAME:     to tuple<memref<?xf64>, memref<?xf64>, f64>
 //       CHECK:   return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
 func.func @sparse_storage_set(%arg0: tuple<memref<?xf64>, memref<?xf64>, f64>, %arg1: memref<?xf64>) -> tuple<memref<?xf64>, memref<?xf64>, f64> {
   %0 = sparse_tensor.storage_set %arg0[0], %arg1
        : tuple<memref<?xf64>, memref<?xf64>, f64>, memref<?xf64> to
          tuple<memref<?xf64>, memref<?xf64>, f64>
   return %0 : tuple<memref<?xf64>, memref<?xf64>, f64>
 }