[mlir][sparse] Add sparse_tensor.sort operator.

[aartbik]

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

395

I would not present this as a COO case per se.

What it really does is sorting parallel arrays, with one array of one type, and all other arrays of index type.
But what is missing here is a criteria on how to compare elements. Is it always lexicographic order in coordinates?

For example, in compress, we will need to sort a *single* array of type index, so that can only be expressed as

sparse_tensor.sort %n, %values

with index typed %values. But will that work too?

400

Please add the usual disclaimer (note that we have two sections, pure ops, and ops that are (for now) defined with side effects.
This is currently the latter

Note that this operation is "impure" in the sense that its behavior
 is solely defined by side-effects and not SSA values. The semantics
 may be refined over time as our sparse abstractions evolve.

404

coordinates

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

457

ironically, our errors do not follow our comment rule,
so start without capital letter, no period at end

462

Please use

for (size_t i = 0, e = getCoordinates().size(); i < e; i++)

idiom

[bixia]

Revise the operator format, fix comments to address review comments.

[bixia]

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

395

Revise the op per offline discussion.

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

457

Fix here and other places.

[bixia]

Revise the operator syntax per offline discussion.

[Peiming]

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

398

It is ambiguous, how about "so that the resulting coordinates are in lexicographical order when putting together". (or something better)

for example:
[1, 3, 2]
[1, 2, 3]

are sorted into
[1, 2, 3]
[1, 3, 2] (note that this array is not in non-decreasing order)

not
[1, 2, 3]
[1, 2, 3]
right?

403

Maybe change it to "This operator updates coordinates and values in place"?

416

Is it the trick to omit jointly when value is empty?

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

517

why not for (Value t : operands)? since you are not using i other than operands[i]

[Peiming]

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

522

This reminds me of the debate I had with Wren about whether or not support dynamic dimension for concat.

While I believe here dynamic support is crucial, but can we add a disclaimer that it is UB when dynamic sized array have unmatched dimensions?

[bixia]

Use xs and ys instead of coordinates and values per offline discussion.
Address review comment.

[aartbik]

Yeah, looking good. A few nits on doc and verification

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

395

Please update short summary as well, something like

Sorts the arrays in xs and ys lexicographically on the integral values found in the xs list

399

Please make a note that, thus, the order in which arrays appear in the xs list matter, before going into the example

408

Don't we want to relax this to, all arrays should have length >= n
?

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

522

I think it is safe to state that all arrays should have length >= n or otherwise UB results.
We can check the static sizes, but skip the dynamic sizes

526

as above, do we really want same length, or just >= n?

mlir/test/Dialect/SparseTensor/invalid.mlir

505

I don't think all custom verification errors are covered yet?
e.g. need at least one xs buffer?

[bixia]

Relax the requirement of dimension.

[bixia]

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

398

Revise this part per offline discussion and also add an example. PTAL.

408

Yes, we can.

416

Yes.

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

517

Thanks!

522

Add comment and description in the td definition.

526

Change to allow >= n.

mlir/test/Dialect/SparseTensor/invalid.mlir

505

The test doesn't work possible due to some parse issue. Add a TODO.

[wrengr]

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

391

Can you check to see if there's a way to use the SameOperandsShape and SameOperandsElementType traits but restricting them to the memrefs in $xs (rather than for all operands to the op)? If we can, that'd save a lot of code in the verifier. Though last time I checked there wasn't a clean way to do this

393

As mentioned in chat, you should use Variadic<...> {let minSize = 1;} for this. I don't think you can do the let inline, so you can just float out the definition class NonemptyVariadic<Type type> : Variadic<type> { let minSize = 1; } and then use NonemptyVariadic<...>:$xs

397–405

It might help clarify things to say that the relationship between memrefs is as if we did zip(zip(xs), zip(ys)) to obtain a list of pairs of lists, and then lexicographically sort the outer list on the first component of the pair.

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

512

Minor nit, but I'd rather this be named xsType (or xsTp or xtp)

[bixia]

Address review comments.

[bixia]

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

391

We now allow different dims for xs and ys. SameOperandsElementType applies to all operands, but we only need it for xs. So we can't use this trait here.

393

I add a TODO for adding this in tablegen and improve here after.

[aartbik]

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

403

put `` on the zip part, so that it looks like code in our online doc as well

415

perhaps mention UB if this condition is not met (since we don't have runtime checks for that)

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

522

Nit: if n == 0, you can skip the whole loop that tests this

[bixia]

Fix comments to address review feedback.

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

522

Per offline discussion, I will keep the current loop which does two checks.

[NathanielMcVicar]

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

522

I believe MSVC is unhappy with the signed comparison here. https://lab.llvm.org/buildbot/#/builders/13/builds/26316

C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\IR\SparseTensorDialect.cpp(523): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\IR\SparseTensorDialect.cpp(523): warning C4018: '<': signed/unsigned mismatch

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

string description = [{
 
     ```mlir
     sparse_tensor.out %t, %dest : tensor<1024x1024xf64, #CSR>, !Dest
     ```
   }];
   let assemblyFormat = "$tensor `,` $dest attr-dict `:` type($tensor) `,` type($dest)";
 }
 
+def SparseTensor_SortOp : SparseTensor_Op<"sort", [AttrSizedOperandSegments]>,
+    Arguments<(ins Index:$n,
+               Variadic<StridedMemRefRankOf<[AnyInteger, Index], [1]>>:$coordinates,
+               Variadic<StridedMemRefRankOf<[AnyType], [1]>>:$values)>  {
+  string summary = "Sorts values based on coordinates";
+  string description = [{
+    Sort the first `n` values in buffers provided by `coordinates` and `values`
+    so that the resulting `coordinates` have non-decreasing values. The
+    `coordinates` needs to have the same element type while the `values` can
+    have different element types. All buffers in `coordinates` and `values`
+    should have the same dimensions. The operator requires at least one buffer
+    for `coordinates` while `values` can be empty. This operator modifies
+    `coordinates` and `values` to reflect the result of the sorting.
+
+    Note that this operation is "impure" in the sense that its behavior is
+    solely defined by side-effects and not SSA values. The semantics may be
+    refined over time as our sparse abstractions evolve.
+
+    Example:
+
+    ```mlir
+    sparse_tensor.sort $n, %coords1, %coords2 (%values1, %values2)
+      : memref<?xindex>, memref<?xindex> (memref<?xindex>, memref<?xf32>)
+    ```
+  }];
+  let assemblyFormat = "$n `,` $coordinates (`(`$values^`)`)? attr-dict "
+                       "`:` type($coordinates) (`(`type($values)^`)`)?";
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Syntax Operations.
 //===----------------------------------------------------------------------===//
 
 def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [NoSideEffect]>,
     Arguments<(ins AnyType:$x, AnyType:$y, UnitAttr:$left_identity, UnitAttr:$right_identity)>,
     Results<(outs AnyType:$output)> {
   let summary = "Binary set operation utilized within linalg.generic";

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

for (unsigned i = 0; i < rank; i++) {
     }
   }
 
   return success();
 }
 
 LogicalResult ForeachOp::verify() {
   auto t = getTensor().getType().cast<RankedTensorType>();
   auto args = getBody()->getArguments();
 
   if (static_cast<size_t>(t.getRank()) + 1 != args.size())
     return emitError("Unmatched number of arguments in the block");
 
   for (int64_t i = 0, e = t.getRank(); i < e; i++)
     if (args[i].getType() != IndexType::get(getContext()))
       emitError(
           llvm::formatv("Expecting Index type for argument at index {0}", i));
 
   auto elemTp = t.getElementType();
   auto valueTp = args.back().getType();
   if (elemTp != valueTp)
     emitError(llvm::formatv("Unmatched element type between input tensor and "
LogicalResult SelectOp::verify() {
   regionResult = verifyNumBlockArgs(this, formula, "select",
                                     TypeRange{inputType}, boolType);
   if (failed(regionResult))
     return regionResult;
 
   return success();
 }
 
+LogicalResult SortOp::verify() {
+  if (getCoordinates().empty())
+    return emitError("need at least one coordinate buffer.");
+
+  auto mtp = getCoordinates().front().getType().cast<MemRefType>();
+  Type etp = mtp.getElementType();
+  int64_t dim = mtp.getShape()[0];
+  auto checkTypes = [&](ValueRange operands,
+                        bool checkEleType = true) -> LogicalResult {
+    for (size_t i = 0, e = operands.size(); i < e; i++) {
+      Value opnd = operands[i];
+      MemRefType cur_mtp = opnd.getType().cast<MemRefType>();
+      uint64_t cur_dim = cur_mtp.getShape()[0];
+      if (dim == ShapedType::kDynamicSize &&
+          cur_dim != ShapedType::kDynamicSize)
+        dim = cur_dim;
+      if (dim != cur_dim && cur_dim != ShapedType::kDynamicSize)
+        return emitError(
+            llvm::formatv("buffers for values and coordinates need "
+                          "to have matching dimension: {0} != {1}",
+                          dim, cur_dim));
+      if (checkEleType && etp != cur_mtp.getElementType())
+        return emitError("mismatch in index array element types");
+    }
+    return success();
+  };
+
+  LogicalResult result = checkTypes(getCoordinates().drop_front());
+  if (failed(result))
+    return result;
+
+  return checkTypes(getValues(), false);
+}
+
 LogicalResult YieldOp::verify() {
   // Check for compatible parent.
   auto *parentOp = (*this)->getParentOp();
   if (isa<BinaryOp>(parentOp) || isa<UnaryOp>(parentOp) ||
       isa<ReduceOp>(parentOp) || isa<SelectOp>(parentOp) ||
       isa<ForeachOp>(parentOp))
     return success();
 

mlir/test/Dialect/SparseTensor/invalid.mlir

 #DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
   // expected-error@+1 {{Unmatched element type between input tensor and block argument}}
   sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
     ^bb0(%1: index, %2: index, %v: f32) :
   }
   return
 }
+
+// -----
+
+func.func @sparse_sort_coordinate_type( %arg0: index, %arg1: memref<?xf32>) {
+  // expected-error@+1 {{operand #1 must be 1D memref of integer or index values}}
+  sparse_tensor.sort %arg0, %arg1: memref<?xf32>
+}
+
+// -----
+
+func.func @sparse_sort_inconsistent(%arg0: index, %arg1: memref<10xindex>, %arg2: memref<20xf32>) {
+  // expected-error@+1 {{buffers for values and coordinates need to have matching dimension: 10 != 20}}
+  sparse_tensor.sort %arg0, %arg1 (%arg2) : memref<10xindex> (memref<20xf32>)
+  return
+}
+
+// -----
+
+func.func @sparse_sort_indices(%arg0: index, %arg1: memref<10xindex>, %arg2: memref<10xi8>) {
+  // expected-error@+1 {{mismatch in index array element types}}
+  sparse_tensor.sort %arg0, %arg1, %arg2 : memref<10xindex>, memref<10xi8>
+  return
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: sparse_tensor.foreach in %[[A0]] : 
 //       CHECK:  ^bb0(%arg1: index, %arg2: index, %arg3: f64):
 func.func @sparse_tensor_foreach(%arg0: tensor<2x4xf64, #DCSR>) -> () {
   sparse_tensor.foreach in %arg0 : tensor<2x4xf64, #DCSR> do {
     ^bb0(%1: index, %2: index, %v: f64) :
   }
   return
 }
+
+// ----
+
+// CHECK-LABEL: func @sparse_sort_1dv0(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<?xindex>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]] : memref<?xindex>
+//       CHECK: return %[[B]]
+func.func @sparse_sort_1dv0(%arg0: index, %arg1: memref<?xindex>) -> (memref<?xindex>) {
+  sparse_tensor.sort %arg0, %arg1 : memref<?xindex>
+  return %arg1 : memref<?xindex>
+}
+
+// -----
+
+// CHECK-LABEL: func @sparse_sort_1dv2(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<?xindex>,
+//  CHECK-SAME: %[[C:.*]]: memref<10xindex>,
+//  CHECK-SAME: %[[D:.*]]: memref<?xf32>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]](%[[C]], %[[D]]) : memref<?xindex>(memref<10xindex>, memref<?xf32>)
+//       CHECK: return %[[B]], %[[C]], %[[D]]
+func.func @sparse_sort_1dv2(%arg0: index, %arg1: memref<?xindex>, %arg2: memref<10xindex>, %arg3: memref<?xf32>) -> (memref<?xindex>, memref<10xindex>, memref<?xf32>) {
+  sparse_tensor.sort %arg0, %arg1 (%arg2, %arg3) : memref<?xindex> (memref<10xindex>, memref<?xf32>)
+  return %arg1, %arg2, %arg3 : memref<?xindex>, memref<10xindex>, memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @sparse_sort_2d(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<10xi8>,
+//  CHECK-SAME: %[[C:.*]]: memref<?xi8>,
+//  CHECK-SAME: %[[D:.*]]: memref<10xf64>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]], %[[C]](%[[D]]) : memref<10xi8>, memref<?xi8>(memref<10xf64>)
+//       CHECK: return %[[B]], %[[C]], %[[D]]
+func.func @sparse_sort_2d(%arg0: index, %arg1: memref<10xi8>, %arg2: memref<?xi8>, %arg3: memref<10xf64>) -> (memref<10xi8>, memref<?xi8>, memref<10xf64>) {
+  sparse_tensor.sort %arg0, %arg1, %arg2 (%arg3) : memref<10xi8>, memref<?xi8> (memref<10xf64>)
+  return %arg1, %arg2, %arg3 : memref<10xi8>, memref<?xi8>, memref<10xf64>
+}