[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

458

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

463

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

458

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

518

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

[Peiming]

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

523

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

523

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

527

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

518

Thanks!

523

Add comment and description in the td definition.

527

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

513

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

523

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

523

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

[NathanielMcVicar]

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

523

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]>>:$xs,
+               Variadic<StridedMemRefRankOf<[AnyType], [1]>>:$ys)>  {
+  string summary = "Sorts the arrays in xs and ys lexicographically on the "
+                   "integral values found in the xs list";
+  string description = [{
+    Lexicographically sort the first `n` values in `xs` along with the values in
+    `ys`. Values in `ys` needed to be sorted along with values in `xs` but don't
+    affect the lexicographical order. Thus, the order in which arrays appear in
+    `xs` matter. This operator updates `xs` and `ys` in place with the result of
+    the sort.
+
+    For example, assume x1=[4, 3], x2=[1, 2], y1=[10, 5], then the output of
+    "sort 2, x1, x2 jointly y1" are x1=[3, 4], x2=[2, 1], y1=[5, 10] while the
+    output of "sort 2, x2, x1, jointly y1" are x2=[1, 2], x1=[4, 3], y1=[10, 5].
+
+    Buffers in `xs` needs to have the same element type while buffers in `ys`
+    can have different element types. All buffers in `xs` and `ys` should have
+    a dimension not less than `n`. The operator requires at least one buffer in
+    `xs` while `ys` can be empty.
+
+    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, %x1, %x2 jointly y1, %y2
+      : memref<?xindex>, memref<?xindex> jointly memref<?xindex>, memref<?xf32>
+    ```
+  }];
+  let assemblyFormat = "$n `,` $xs (`jointly` $ys^)? attr-dict"
+                       "`:` type($xs) (`jointly` type($ys)^)?";
+
+  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/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorDialect
   SparseTensorDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   DEPENDS
   MLIRSparseTensorAttrDefsIncGen
   MLIRSparseTensorOpsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRArithmeticDialect
   MLIRDialect
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSupport
   )

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

 //===- SparseTensorDialect.cpp - Sparse tensor dialect implementation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
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 (getXs().empty())
+    return emitError("need at least one xs buffer.");
+
+  auto n = getN().getDefiningOp<arith::ConstantIndexOp>();
+
+  auto mtp = getXs().front().getType().cast<MemRefType>();
+  Type etp = mtp.getElementType();
+  auto checkTypes = [&](ValueRange operands,
+                        bool checkEleType = true) -> LogicalResult {
+    for (Value opnd : operands) {
+      MemRefType cur_mtp = opnd.getType().cast<MemRefType>();
+      uint64_t cur_dim = cur_mtp.getShape()[0];
+      // We can't check the size of dynamic dimension at compile-time, but the
+      // operator requires all xs and ys have a dimension not less than n.
+      if (n && cur_dim != ShapedType::kDynamicSize && cur_dim < n.value())
+        return emitError(llvm::formatv("xs and ys need to have a dimension >= n"
+                                       ": {0} < {1}",
+                                       cur_dim, n.value()));
+
+      if (checkEleType && etp != cur_mtp.getElementType())
+        return emitError("mismatch xs element types");
+    }
+    return success();
+  };
+
+  LogicalResult result = checkTypes(getXs());
+  if (failed(result))
+    return result;
+
+  if (n)
+    return checkTypes(getYs(), false);
+
+  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) ||
       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
 }
+
+// -----
+
+// TODO: a test case with empty xs doesn't work due to some parser issues.
+
+func.func @sparse_sort_x_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_dim_too_small(%arg0: memref<10xindex>) {
+  %i20 = arith.constant 20 : index
+  // expected-error@+1 {{xs and ys need to have a dimension >= n: 10 < 20}}
+  sparse_tensor.sort %i20, %arg0 : memref<10xindex>
+  return
+}
+
+// -----
+
+func.func @sparse_sort_mismatch_x_type(%arg0: index, %arg1: memref<10xindex>, %arg2: memref<10xi8>) {
+  // expected-error@+1 {{mismatch xs 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_1d0v(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<?xindex>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]] : memref<?xindex>
+//       CHECK: return %[[B]]
+func.func @sparse_sort_1d0v(%arg0: index, %arg1: memref<?xindex>) -> (memref<?xindex>) {
+  sparse_tensor.sort %arg0, %arg1 : memref<?xindex>
+  return %arg1 : memref<?xindex>
+}
+
+// -----
+
+// CHECK-LABEL: func @sparse_sort_1d2v(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<20xindex>,
+//  CHECK-SAME: %[[C:.*]]: memref<10xindex>,
+//  CHECK-SAME: %[[D:.*]]: memref<?xf32>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]] jointly %[[C]], %[[D]] : memref<20xindex> jointly memref<10xindex>, memref<?xf32>
+//       CHECK: return %[[B]], %[[C]], %[[D]]
+func.func @sparse_sort_1d2v(%arg0: index, %arg1: memref<20xindex>, %arg2: memref<10xindex>, %arg3: memref<?xf32>) -> (memref<20xindex>, memref<10xindex>, memref<?xf32>) {
+  sparse_tensor.sort %arg0, %arg1 jointly %arg2, %arg3 : memref<20xindex> jointly memref<10xindex>, memref<?xf32>
+  return %arg1, %arg2, %arg3 : memref<20xindex>, memref<10xindex>, memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @sparse_sort_2d1v(
+//  CHECK-SAME: %[[A:.*]]: index,
+//  CHECK-SAME: %[[B:.*]]: memref<10xi8>,
+//  CHECK-SAME: %[[C:.*]]: memref<20xi8>,
+//  CHECK-SAME: %[[D:.*]]: memref<10xf64>)
+//       CHECK: sparse_tensor.sort %[[A]], %[[B]], %[[C]] jointly %[[D]] : memref<10xi8>, memref<20xi8> jointly memref<10xf64>
+//       CHECK: return %[[B]], %[[C]], %[[D]]
+func.func @sparse_sort_2d1v(%arg0: index, %arg1: memref<10xi8>, %arg2: memref<20xi8>, %arg3: memref<10xf64>) -> (memref<10xi8>, memref<20xi8>, memref<10xf64>) {
+  sparse_tensor.sort %arg0, %arg1, %arg2 jointly %arg3 : memref<10xi8>, memref<20xi8> jointly memref<10xf64>
+  return %arg1, %arg2, %arg3 : memref<10xi8>, memref<20xi8>, memref<10xf64>
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 )
 
 cc_library(
     name = "SparseTensorDialect",
     srcs = ["lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp"],
     hdrs = ["include/mlir/Dialect/SparseTensor/IR/SparseTensor.h"],
     includes = ["include"],
     deps = [
+        ":ArithmeticDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":SparseTensorAttrDefsIncGen",
         ":SparseTensorOpsIncGen",
         "//llvm:Support",
     ],
 )