[mlir][sparse] Add new concatente operator to sparse tensor

[Peiming]

Summary

See https://www.tensorflow.org/xla/operation_semantics#concatenate for the operator semantics

[aartbik]

Much easier to review smaller revision! Thanks.

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

171

Concatenates a list of tensor into a single tensor.

173

Please mention that all operands have same rank, and that 0 <= dimension < rank

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

359

check

0 <= concatDim < rank

?

364

You could refine the error by emitting the argument number and mismatched rank
(optional, but it is usually better to have precise error messages)

368

Can you add some comment here on what you allow for dynamic sizes? The logic for static sizes is clear (all the same for != concatDim and sum of all if == concatDim, but it it not so clear what cases you allow when some or dyanmic). Would it make sense to simply require them all to be static for the first version? Did you observe the need to support dynamic?

[Peiming]

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

359

The concatDim is of type index, it should guarantee the 0 <= concatDim.

I will add check concatDim < rank

368

I will require them all to be static then

[Peiming]

addressed comments

[Peiming]

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

364

See L375

368

See SparseTensorOps.td@L179

[aartbik]

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

174

typo: negaive -> negative

also, although you are right that (unsigned) indices are always non-negative, I would , at least here, in the doc, simply say

0 <= dimension < rank

for readability

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

361

no { } needed in single stmt if

also, we can decided to allow for size 1 and simply fold it into itself
(this has something nice properties when progressive lowering stuff into lower ranked parts)

367

no { }

373

same, no { } needed

382

and here, no { } , also below

[Peiming]

Addressed comments

[Peiming]

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

361

Should I do it now or in the future (when it is needed)

Also, probably need to lowering it to sparse_tensor.convert in the case that input and output tensor have different encodings

[aartbik]

Two last nits, but good to go

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

361

We can do that later. I was just thinking "out loud" here.

404

it works but I would find code that assign the first dimension size here and then checks if they are all equal without the "prev" trick a bit cleaner (and you have just verified there are sufficient inputs for this)

[wrengr]

Some post-landing comments since I was out on vacation

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

173

Note that there's an AllRanksMatch trait to ensure that ranks match between operands (and/or results iiuc). I'll post a CL to add that. Don't know that there's a premade trait for the other constraint, but it should be easy enough to write up

179–180

As I mentioned in D130287, I think it'd be good to have this op not introduce any new dynamic sizes but rather force there to be an explicit cast doing so (like memref::CollapseShapeOp).

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

368

@aartbik fwiw, we had a big long discussion over at D130287. IMO the only valid thing would be to allow dynamic sizes for the dimension being concatenated along (and if any input has such, then the corresponding output dimension would also be dynamic). Since anything else would require introducing runtime assertions (or a sufficiently powerful analysis to statically determine that those assertions must always succeed).

[Peiming]

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

173

Good to know! Thank you!

179–180

Either way is okay for me.

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

let description = [{
     ```mlir
     %1 = sparse_tensor.values %0 : tensor<64x64xf64, #CSR> to memref<?xf64>
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
   let hasVerifier = 1;
 }
 
+def SparseTensor_ConcatenateOp : SparseTensor_Op<"concatenate", []>,
+    Arguments<(ins Variadic<AnyRankedTensor>:$inputs,
+                   IndexAttr:$dimension)>,
+    Results<(outs AnyRankedTensor:$result)> {
+
+  let summary = "Concatenates a list of concatenate op";
+  let description = [{
+     The concatenation happens on the specified `dimension`. The resulting `dimension`
+     size is the sum of all the input dimension sizes, while all the other dimensions
+     should have the same size in the input and output tensors.
+
+
+     Example:
+
+     ```mlir
+     %0 = sparse_tensor.concatenate %1, %2 { dimension = 0 : index }
+       : tensor<64x64xf64, #CSR>, tensor<64x64xf64, #CSR> to tensor<128x64xf64, #CSR>
+     ```
+   }];
+
+  let assemblyFormat = "$inputs attr-dict `:` type($inputs) `to` type($result)";
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Management Operations. These operations are "impure" in the
 // sense that they do not properly operate on SSA values. Instead, the behavior
 // is solely defined by side-effects. These operations provide a bridge between
 // the code generator and the support library. The semantics of these operations
 // may be refined over time as our sparse abstractions evolve.
 //===----------------------------------------------------------------------===//
 

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

regionResult =
         verifyNumBlockArgs(this, absent, "absent", TypeRange{}, outputType);
     if (failed(regionResult))
       return regionResult;
   }
 
   return success();
 }
 
+LogicalResult ConcatenateOp::verify() {
+  auto dstTp = getType().cast<RankedTensorType>();
+  uint64_t concatDim = getDimension().getZExtValue();
+  unsigned rank = dstTp.getRank();
+
+  for (auto input : getInputs()) {
+    auto inputRank = input.getType().cast<RankedTensorType>().getRank();
+    if (inputRank != rank) {
+      return emitError(
+          "All input tensors and output tensor should have the same rank.");
+    }
+  }
+
+  for (unsigned i = 0; i < rank; i++) {
+    auto dstDim = dstTp.getShape()[i];
+    if (i == concatDim) {
+      if (dstDim != ShapedType::kDynamicSize) {
+        unsigned sumDim = 0;
+        for (auto src : getInputs()) {
+          auto d = src.getType().cast<RankedTensorType>().getShape()[i];
+          if (d == ShapedType::kDynamicSize)
+            // If the output tensor has static dimension, yet there is an input
+            // tensor with dynamic dimension.
+            return emitError(
+                "Failed to verify the shaping rules with dynamically "
+                "shaped inputs");
+          sumDim += d;
+        }
+        // If all dimension are statically known, the sum of all the input
+        // dimensions should be equal to the output dimension.
+        if (sumDim != dstDim) {
+          return emitError(
+              "The concatenation dimension of the output tensor should be the "
+              "sum of all the concatenation dimensions of the input tensors.");
+        }
+      }
+    } else {
+      int prev = dstDim;
+      for (auto src : getInputs()) {
+        auto d = src.getType().cast<RankedTensorType>().getShape()[i];
+        if (d != ShapedType::kDynamicSize) {
+          if (prev != ShapedType::kDynamicSize && d != prev) {
+            return emitError(
+                "All dimensions (expect for the concatenating one) "
+                "should be equal.");
+          }
+          prev = d;
+        }
+      }
+    }
+  }
+
+  return success();
+}
+
 LogicalResult ReduceOp::verify() {
   Type inputType = getX().getType();
   LogicalResult regionResult = success();
 
   // Check correct number of block arguments and return type.
   Region &formula = getRegion();
   if (!formula.empty()) {
     regionResult = verifyNumBlockArgs(

mlir/test/Dialect/SparseTensor/invalid.mlir

func.func @invalid_reduce_wrong_yield(%arg0: f64, %arg1: f64) -> f64 {
   // expected-error@+1 {{reduce region must end with sparse_tensor.yield}}
   %r = sparse_tensor.reduce %arg0, %arg1, %cf1 : f64 {
       ^bb0(%x: f64, %y: f64):
         %cst = arith.constant 2 : i64
         tensor.yield %cst : i64
     }
   return %r : f64
 }
+
+// -----
+
+#C = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#DC = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+#DCC = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed", "compressed"]}>
+func.func @invalid_concat_rank_mismatch(%arg0: tensor<2xf64, #C>,
+                                        %arg1: tensor<3x4xf64, #DC>,
+                                        %arg2: tensor<4x4x4xf64, #DCC>) -> tensor<9x4xf64, #DC> {
+  // expected-error@+1 {{All input tensors and output tensor should have the same rank}}
+  %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
+       : tensor<2xf64, #C>,
+         tensor<3x4xf64, #DC>,
+         tensor<4x4x4xf64, #DCC> to tensor<9x4xf64, #DC>
+  return %0 : tensor<9x4xf64, #DC>
+}
+
+// -----
+
+#DC = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+func.func @invalid_concat_size_mismatch_dyn(%arg0: tensor<?x4xf64, #DC>,
+                                            %arg1: tensor<5x4xf64, #DC>,
+                                            %arg2: tensor<4x4xf64, #DC>) -> tensor<9x4xf64, #DC> {
+  // expected-error@+1 {{Failed to verify the shaping rules with dynamically shaped inputs}}
+  %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
+       : tensor<?x4xf64, #DC>,
+         tensor<5x4xf64, #DC>,
+         tensor<4x4xf64, #DC> to tensor<9x4xf64, #DC>
+  return %0 : tensor<9x4xf64, #DC>
+}
+
+// -----
+
+#DC = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+func.func @invalid_concat_size_mismatch(%arg0: tensor<3x4xf64, #DC>,
+                                        %arg1: tensor<5x4xf64, #DC>,
+                                        %arg2: tensor<4x4xf64, #DC>) -> tensor<9x4xf64, #DC> {
+  // expected-error@+1 {{The concatenation dimension of the output tensor should be the sum of}}
+  %0 = sparse_tensor.concatenate %arg0, %arg1, %arg2 {dimension = 0 : index}
+       : tensor<3x4xf64, #DC>,
+         tensor<5x4xf64, #DC>,
+         tensor<4x4xf64, #DC> to tensor<9x4xf64, #DC>
+  return %0 : tensor<9x4xf64, #DC>
+}
+
+// -----
+
+#DC = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+func.func @invalid_concat_size_mismatch(%arg0: tensor<2x4xf64, #DC>,
+                                        %arg1: tensor<3x3xf64, #DC>,
+                                        %arg2: tensor<4x4xf64, #DC>) -> tensor<9x4xf64, #DC> {
+  // expected-error@+1 {{All dimensions (expect for the concatenating one) should be equal}}
+  %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>
+}
+

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: }
 func.func @sparse_reduce_2d_to_1d(%arg0: f64, %arg1: f64) -> f64 {
   %cf0 = arith.constant 0.0 : f64
   %r = sparse_tensor.reduce %arg0, %arg1, %cf0 : f64 {
       ^bb0(%x: f64, %y: f64):
         sparse_tensor.yield %x : f64
     }
   return %r : f64
-}
+}
+
+// -----
+
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @concat_sparse_sparse(
+//  CHECK-SAME:   %[[A0:.*]]: tensor<2x4xf64
+//  CHECK-SAME:   %[[A1:.*]]: tensor<3x4xf64
+//  CHECK-SAME:   %[[A2:.*]]: tensor<4x4xf64
+//       CHECK:   %[[TMP0:.*]] = sparse_tensor.concatenate %[[A0]], %[[A1]], %[[A2]] {dimension = 0 : index} :
+//  CHECK-SAME:   tensor<2x4xf64
+//  CHECK-SAME:   tensor<3x4xf64
+//  CHECK-SAME:   tensor<4x4xf64
+//  CHECK-SAME:   tensor<9x4xf64
+//       CHECK:   return %[[TMP0]] : tensor<9x4xf64
+func.func @concat_sparse_sparse(%arg0: tensor<2x4xf64, #SparseMatrix>,
+                                %arg1: tensor<3x4xf64, #SparseMatrix>,
+                                %arg2: tensor<4x4xf64, #SparseMatrix>) -> tensor<9x4xf64, #SparseMatrix> {
+  %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>
+}