[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

519

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

[Peiming]

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

524

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

524

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

528

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

mlir/test/Dialect/SparseTensor/invalid.mlir

472

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

519

Thanks!

524

Add comment and description in the td definition.

528

Change to allow >= n.

mlir/test/Dialect/SparseTensor/invalid.mlir

472

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

514

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

524

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

524

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

[NathanielMcVicar]

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

524

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", []>,
+    Arguments<(ins Index:$n,
+               StridedMemRefRankOf<[AnyType], [1]>:$values,
+               Variadic<StridedMemRefRankOf<[AnyInteger, Index], [1]>>:$coordinates)>  {
+  string summary = "Sorts COO data with values and coordinates";
+  string description = [{
+    Sort the first `n` values in `values` by non-decreasing coordinate values
+    provided by `coordinates`. This operator modifies `values` and `coordinates`
+    with the result of the sort.
+
+    Example:
+
+    ```mlir
+    sparse_tensor.sort %n, %values, %coordinate
+      : memref<?xf32>, memref<?xindex>
+    ```
+  }];
+  let assemblyFormat = "$n `,` $values (`,` $coordinates^)? attr-dict "
+                       "`:` type($values) (`,` type($coordinates)^)?";
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Custom Linalg.Generic 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";
   let description = [{
       Defines a computation within a `linalg.generic` operation that takes two
       operands and executes one of the regions depending on whether both operands
       or either operand is nonzero (i.e. stored explicitly in the sparse storage

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

if (i == concatDim) {
         prev = d;
       }
     }
   }
 
   return success();
 }
 
+LogicalResult SortOp::verify() {
+  if (getCoordinates().empty())
+    return emitError("Need at least one coordinate buffer.");
+
+  int64_t dim = getValues().getType().cast<MemRefType>().getShape()[0];
+  Type etp =
+      getCoordinates().front().getType().cast<MemRefType>().getElementType();
+  for (size_t i = 0; i < getCoordinates().size(); i++) {
+    Value coord = getCoordinates()[i];
+    MemRefType mtp = coord.getType().cast<MemRefType>();
+    uint64_t cur_dim = mtp.getShape()[0];
+    if (dim == ShapedType::kDynamicSize || cur_dim == ShapedType::kDynamicSize)
+      dim = ShapedType::kDynamicSize;
+    else if (dim != cur_dim)
+      return emitError(llvm::formatv("Buffers for values and coordinates need "
+                                     "to have matching dimension: {0} != {1}",
+                                     dim, cur_dim));
+    if (etp != mtp.getElementType())
+      return emitError("Mismatch in index array element types");
+  }
+
+  return success();
+}
+
 LogicalResult ReduceOp::verify() {
   Type inputType = getX().getType();
   LogicalResult regionResult = success();
 
   // Check correct number of block arguments and return type.
   Region &formula = getRegion();
   regionResult = verifyNumBlockArgs(this, formula, "reduce",
                                     TypeRange{inputType, inputType}, inputType);
LogicalResult SelectOp::verify() {
 
   return success();
 }
 
 LogicalResult YieldOp::verify() {
   // Check for compatible parent.
   auto *parentOp = (*this)->getParentOp();
   if (isa<BinaryOp>(parentOp) || isa<UnaryOp>(parentOp) ||
       isa<ReduceOp>(parentOp) || isa<SelectOp>(parentOp))
     return success();
 
   return emitOpError("expected parent op to be sparse_tensor unary, binary, "
                      "reduce, or select");
 }
 
 //===----------------------------------------------------------------------===//
 // TensorDialect Methods.
 //===----------------------------------------------------------------------===//
 
 void SparseTensorDialect::initialize() {
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.cpp.inc"
       >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.cpp.inc"
       >();
 }
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.cpp.inc"
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.cpp.inc"

mlir/test/Dialect/SparseTensor/invalid.mlir

func.func @invalid_concat_size_mismatch(%arg0: tensor<2x4xf64, #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>
 }
+
+// -----
+
+func.func @sparse_sort_coordinate_type(%arg0: index, %arg1: memref<?xf32>,%arg2: memref<?xf32>) {
+  // expected-error@+1 {{op operand #2 must be 1D memref of integer or index values}}
+  sparse_tensor.sort %arg0, %arg1, %arg2 : memref<?xf32>, memref<?xf32>
+}
+
+// -----
+
+func.func @sparse_sort_no_coordinate(%arg0: index, %arg1: memref<20xf32>) {
+  // expected-error@+1 {{Need at least one coordinate buffer}}
+  sparse_tensor.sort %arg0, %arg1 : memref<20xf32>
+  return
+}
+
+// -----
+
+func.func @sparse_sort_inconsistent(%arg0: index, %arg1: memref<20xf32>, %arg2: memref<10xindex>) {
+  // expected-error@+1 {{Buffers for values and coordinates need to have matching dimension: 20 != 10}}
+  sparse_tensor.sort %arg0, %arg1, %arg2 : memref<20xf32>, memref<10xindex>
+  return
+}
+
+// -----
+
+func.func @sparse_sort_indices(%arg0: index, %arg1: memref<10xf64>, %arg2: memref<10xindex>, %arg3: memref<10xi8>)
+  -> (memref<10xf64>, memref<10xindex>, memref<10xi8>) {
+  // expected-error@+1 {{Mismatch in index array element types}}
+  sparse_tensor.sort %arg0, %arg1, %arg2, %arg3 : memref<10xf64>, memref<10xindex>, memref<10xi8>
+  return %arg1, %arg2, %arg3 : memref<10xf64>, memref<10xindex>, memref<10xi8>
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

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>
 }
+
+// -----
+
+// CHECK-LABEL: func @sparse_sort_1d(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<?xf32>,
+//  CHECK-SAME: %[[C:.*]]: memref<?xindex>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]], %[[C]] : memref<?xf32>, memref<?xindex>
+//       CHECK: return %[[B]], %[[C]]
+func.func @sparse_sort_1d(%arg0: index, %arg1: memref<?xf32>, %arg2: memref<?xindex>) -> (memref<?xf32>, memref<?xindex>) {
+  sparse_tensor.sort %arg0, %arg1, %arg2 : memref<?xf32>, memref<?xindex>
+  return %arg1, %arg2 : memref<?xf32>, memref<?xindex>
+}
+
+// -----
+
+// CHECK-LABEL: func @sparse_sort_2d(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<10xf64>,
+//  CHECK-SAME: %[[C:.*]]: memref<?xi8>,
+//  CHECK-SAME: %[[D:.*]]: memref<10xi8>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]], %[[C]], %[[D]] : memref<10xf64>, memref<?xi8>, memref<10xi8>
+//       CHECK: return %[[B]], %[[C]], %[[D]]
+func.func @sparse_sort_2d(%arg0: index, %arg1: memref<10xf64>, %arg2: memref<?xi8>, %arg3: memref<10xi8>) -> (memref<10xf64>, memref<?xi8>, memref<10xi8>) {
+  sparse_tensor.sort %arg0, %arg1, %arg2, %arg3 : memref<10xf64>, memref<?xi8>, memref<10xi8>
+  return %arg1, %arg2, %arg3 : memref<10xf64>, memref<?xi8>, memref<10xi8>
+}