[mlir][sparse] Add an attribute to the sort operator for stable sorting.

[aartbik]

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

441

indicates that .... -> indicates whether ....

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

680

build(... , /*stable=*/false);

[bixia]

Address review comments.

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

 }
 
 def SparseTensor_SortOp : SparseTensor_Op<"sort", [AttrSizedOperandSegments]>,
     // TODO: May want to extend tablegen with
     // class NonemptyVariadic<Type type> : Variadic<type> { let minSize = 1; }
     // and then use NonemptyVariadic<...>:$xs here.
     Arguments<(ins Index:$n,
                Variadic<StridedMemRefRankOf<[AnyInteger, Index], [1]>>:$xs,
-               Variadic<StridedMemRefRankOf<[AnyType], [1]>>:$ys)>  {
+               Variadic<StridedMemRefRankOf<[AnyType], [1]>>:$ys,
+               UnitAttr:$stable)>  {
   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`. Conceptually, the values being sorted are tuples produced by
     `zip(zip(xs), zip(ys))`. In particular, values in `ys` needed to be sorted
     along with values in `xs`, but values in `ys` don't affect the
     lexicographical order. The order in which arrays appear in `xs` affects the
     sorting result. The operator updates `xs` and `ys` in place with the result
     of the sorting.
 
     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 integral element type while buffers
     in `ys` can have different numeric element types. All buffers in `xs` and
     `ys` should have a dimension not less than `n`. The behavior of the operator
     is undefined if this condition is not met. The operator requires at least
     one buffer in `xs` while `ys` can be empty.
 
+    The `stable` attribute indicates whether a stable sorting algorithm should
+    be used to implement the operator.
+
     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>
     ```
+
+    ```mlir
+    sparse_tensor.sort stable %n, %x1, %x2 jointly y1, %y2
+      : memref<?xindex>, memref<?xindex> jointly memref<?xindex>, memref<?xf32>
+    ```
   }];
-  let assemblyFormat = "$n `,` $xs (`jointly` $ys^)? attr-dict"
+  let assemblyFormat = "(`stable` $stable^)? $n"
+                       "`,`$xs (`jointly` $ys^)? attr-dict"
                        "`:` type($xs) (`jointly` type($ys)^)?";
-
+  let builders = [
+        OpBuilder<(ins "Value":$n, "ValueRange":$xs, "ValueRange":$ys)>
+  ];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Syntax Operations.
 //===----------------------------------------------------------------------===//
 
 def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [NoSideEffect]>,

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

LogicalResult SelectOp::verify() {
   regionResult = verifyNumBlockArgs(this, formula, "select",
                                     TypeRange{inputType}, boolType);
   if (failed(regionResult))
     return regionResult;
 
   return success();
 }
 
+void SortOp::build(OpBuilder &odsBuilder, OperationState &odsState, Value n,
+                   ValueRange xs, ValueRange ys) {
+  build(odsBuilder, odsState, n, xs, ys, /*stable=*/false);
+}
+
 LogicalResult SortOp::verify() {
   if (getXs().empty())
     return emitError("need at least one xs buffer.");
 
   auto n = getN().getDefiningOp<arith::ConstantIndexOp>();
 
   Type xtp = getXs().front().getType().cast<MemRefType>().getElementType();
   auto checkTypes = [&](ValueRange operands,

mlir/test/Dialect/SparseTensor/roundtrip.mlir

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