[mlir][sparse] Implement hybrid quick sort for sparse_tensor.sort.

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

1037

isn't the nTrailingP a boolean here, deciding between two versions?
The int seems to imply more semantics than there are?

You even start with

assert(nTrailingP == 1 || nTrailingP == 0);

bool isHybrid = (nTrailingP == 1);

so why not just passing in isHybrid?

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

48

Why not name this "_sparse_qsort_" so that it matches the "_sparse_hybrid_qsort_" name?

[bixia]

Change _sparse_quick_sort to _sparse_qsort. Rebase.

[bixia]

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

1037

This is because FuncGeneratorType requires it to be an uint32_t. And the reason that FuncGeneratorType uses uint32_t not boolean is we can potentially have more trailing data, such as for collecting the number of swaps performed.

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

inline Value constantOne(OpBuilder &builder, Location loc, Type tp) {
   return builder.create<arith::ConstantOp>(loc, tp, getOneAttr(builder, tp));
 }
 
 /// Generates a constant of `index` type.
 inline Value constantIndex(OpBuilder &builder, Location loc, int64_t i) {
   return builder.create<arith::ConstantIndexOp>(loc, i);
 }
 
+/// Generates a constant of `i64` type.
+inline Value constantI64(OpBuilder &builder, Location loc, int64_t i) {
+  return builder.create<arith::ConstantIntOp>(loc, i, 64);
+}
+
 /// Generates a constant of `i32` type.
 inline Value constantI32(OpBuilder &builder, Location loc, int32_t i) {
   return builder.create<arith::ConstantIntOp>(loc, i, 32);
 }
 
 /// Generates a constant of `i16` type.
 inline Value constantI16(OpBuilder &builder, Location loc, int16_t i) {
   return builder.create<arith::ConstantIntOp>(loc, i, 16);

mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===---------------------------------------------------------------------===//
 // Helper methods for the actual rewriting rules.
 //===---------------------------------------------------------------------===//
 
 static constexpr uint64_t loIdx = 0;
 static constexpr uint64_t hiIdx = 1;
 static constexpr uint64_t xStartIdx = 2;
 
 static constexpr const char kLessThanFuncNamePrefix[] = "_sparse_less_than_";
 static constexpr const char kCompareEqFuncNamePrefix[] = "_sparse_compare_eq_";
 static constexpr const char kPartitionFuncNamePrefix[] = "_sparse_partition_";
 static constexpr const char kBinarySearchFuncNamePrefix[] =
     "_sparse_binary_search_";
-static constexpr const char kSortNonstableFuncNamePrefix[] =
-    "_sparse_sort_nonstable_";
+static constexpr const char kHybridQuickSortFuncNamePrefix[] =
+    "_sparse_hybrid_qsort_";
 static constexpr const char kSortStableFuncNamePrefix[] =
     "_sparse_sort_stable_";
 static constexpr const char kShiftDownFuncNamePrefix[] = "_sparse_shift_down_";
 static constexpr const char kHeapSortFuncNamePrefix[] = "_sparse_heap_sort_";
+static constexpr const char kQuickSortFuncNamePrefix[] = "_sparse_qsort_";
 
 using FuncGeneratorType = function_ref<void(
     OpBuilder &, ModuleOp, func::FuncOp, uint64_t, uint64_t, bool, uint32_t)>;
 
 /// Constructs a function name with this format to facilitate quick sort:
 ///   <namePrefix><nx>_<x type>_<y0 type>..._<yn type> for sort
 ///   <namePrefix><nx>_<x type>_coo_<ny>_<y0 type>..._<yn type> for sort_coo
 static void getMangledSortHelperFuncName(llvm::raw_svector_ostream &nameOstream,
shiftDownOperands[shiftDownOperands.size() - 1] =
       builder.create<arith::SubIOp>(loc, l, c1);
   builder.create<func::CallOp>(loc, shiftDownFunc, TypeRange(),
                                shiftDownOperands);
 
   builder.setInsertionPointAfter(forL);
   builder.create<func::ReturnOp>(loc);
 }
 
-/// Creates a function to perform quick sort on the value in the range of
-/// index [lo, hi).
-//
-// The generate IR corresponds to this C like algorithm:
-// void quickSort(lo, hi, data) {
-//   if (lo < hi) {
-//        p = partition(low, high, data);
-//        quickSort(lo, p, data);
-//        quickSort(p + 1, hi, data);
-//   }
-// }
-static void createSortNonstableFunc(OpBuilder &builder, ModuleOp module,
-                                    func::FuncOp func, uint64_t nx, uint64_t ny,
-                                    bool isCoo, uint32_t nTrailingP) {
-  (void)nTrailingP;
-  OpBuilder::InsertionGuard insertionGuard(builder);
-  Block *entryBlock = func.addEntryBlock();
-  builder.setInsertionPointToStart(entryBlock);
-
+static void createQuickSort(OpBuilder &builder, ModuleOp module,
+                            func::FuncOp func, ValueRange args, uint64_t nx,
+                            uint64_t ny, bool isCoo, uint32_t nTrailingP) {
   MLIRContext *context = module.getContext();
   Location loc = func.getLoc();
-  ValueRange args = entryBlock->getArguments();
   Value lo = args[loIdx];
   Value hi = args[hiIdx];
-  Value cond =
-      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, lo, hi);
-  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else=*/false);
-
-  // The if-stmt true branch.
-  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   FlatSymbolRefAttr partitionFunc = getMangledSortHelperFunc(
       builder, func, {IndexType::get(context)}, kPartitionFuncNamePrefix, nx,
-      ny, isCoo, args, createPartitionFunc);
-  auto p = builder.create<func::CallOp>(
-      loc, partitionFunc, TypeRange{IndexType::get(context)}, ValueRange(args));
+      ny, isCoo, args.drop_back(nTrailingP), createPartitionFunc);
+  auto p = builder.create<func::CallOp>(loc, partitionFunc,
+                                        TypeRange{IndexType::get(context)},
+                                        args.drop_back(nTrailingP));
 
   SmallVector<Value> lowOperands{lo, p.getResult(0)};
   lowOperands.append(args.begin() + xStartIdx, args.end());
   builder.create<func::CallOp>(loc, func, lowOperands);
 
   SmallVector<Value> highOperands{
       builder.create<arith::AddIOp>(loc, p.getResult(0),
                                     constantIndex(builder, loc, 1)),
       hi};
   highOperands.append(args.begin() + xStartIdx, args.end());
   builder.create<func::CallOp>(loc, func, highOperands);
-
-  // After the if-stmt.
-  builder.setInsertionPointAfter(ifOp);
-  builder.create<func::ReturnOp>(loc);
 }
 
 /// Creates a function to perform insertion sort on the values in the range of
 /// index [lo, hi).
 //
 // The generate IR corresponds to this C like algorithm:
 // void insertionSort(lo, hi, data) {
 //   for (i = lo+1; i < hi; i++) {
forEachIJPairInAllBuffers(
       [&](uint64_t k, Value p, Value usused, Value buffer) {
         builder.create<memref::StoreOp>(loc, d[k], buffer, p);
       });
 
   builder.setInsertionPointAfter(forOpI);
   builder.create<func::ReturnOp>(loc);
 }
 
+/// Creates a function to perform quick sort or a hybrid quick sort on the
+/// values in the range of index [lo, hi).
+//
+//
+// When nTrailingP == 0, the generated IR corresponds to this C like algorithm:
+// void quickSort(lo, hi, data) {
+//   if (lo + 1 < hi) {
+//        p = partition(low, high, data);
+//        quickSort(lo, p, data);
+//        quickSort(p + 1, hi, data);
+//   }
+// }
+//
+// When nTrailingP == 1, the generated IR corresponds to this C like algorithm:
+// void hybridQuickSort(lo, hi, data, depthLimit) {
+//   if (lo + 1 < hi) {
+//     len = hi - lo;
+//     if (len <= limit) {
+//       insertionSort(lo, hi, data);
+//     } else {
+//       depthLimit --;
+//       if (depthLimit <= 0) {
+//         heapSort(lo, hi, data);
+//       } else {
+//          p = partition(low, high, data);
+//          quickSort(lo, p, data);
+//          quickSort(p + 1, hi, data);
+//       }
+//       depthLimit ++;
+//     }
+//   }
+// }
+//
+static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
+                                func::FuncOp func, uint64_t nx, uint64_t ny,
+                                bool isCoo, uint32_t nTrailingP) {
+  assert(nTrailingP == 1 || nTrailingP == 0);
+  bool isHybrid = (nTrailingP == 1);
+  OpBuilder::InsertionGuard insertionGuard(builder);
+  Block *entryBlock = func.addEntryBlock();
+  builder.setInsertionPointToStart(entryBlock);
+
+  Location loc = func.getLoc();
+  ValueRange args = entryBlock->getArguments();
+  Value lo = args[loIdx];
+  Value hi = args[hiIdx];
+  Value loCmp =
+      builder.create<arith::AddIOp>(loc, lo, constantIndex(builder, loc, 1));
+  Value cond =
+      builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, loCmp, hi);
+  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else=*/false);
+
+  // The if-stmt true branch.
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  Value pDepthLimit;
+  Value savedDepthLimit;
+  scf::IfOp depthIf;
+
+  if (isHybrid) {
+    Value len = builder.create<arith::SubIOp>(loc, hi, lo);
+    Value lenLimit = constantIndex(builder, loc, 30);
+    Value lenCond = builder.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::ule, len, lenLimit);
+    scf::IfOp lenIf = builder.create<scf::IfOp>(loc, lenCond, /*else=*/true);
+
+    // When len <= limit.
+    builder.setInsertionPointToStart(&lenIf.getThenRegion().front());
+    FlatSymbolRefAttr insertionSortFunc = getMangledSortHelperFunc(
+        builder, func, TypeRange(), kSortStableFuncNamePrefix, nx, ny, isCoo,
+        args.drop_back(nTrailingP), createSortStableFunc);
+    builder.create<func::CallOp>(loc, insertionSortFunc, TypeRange(),
+                                 ValueRange(args.drop_back(nTrailingP)));
+
+    // When len > limit.
+    builder.setInsertionPointToStart(&lenIf.getElseRegion().front());
+    pDepthLimit = args.back();
+    savedDepthLimit = builder.create<memref::LoadOp>(loc, pDepthLimit);
+    Value depthLimit = builder.create<arith::SubIOp>(
+        loc, savedDepthLimit, constantI64(builder, loc, 1));
+    builder.create<memref::StoreOp>(loc, depthLimit, pDepthLimit);
+    Value depthCond =
+        builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ule,
+                                      depthLimit, constantI64(builder, loc, 0));
+    depthIf = builder.create<scf::IfOp>(loc, depthCond, /*else=*/true);
+
+    // When depth exceeds limit.
+    builder.setInsertionPointToStart(&depthIf.getThenRegion().front());
+    FlatSymbolRefAttr heapSortFunc = getMangledSortHelperFunc(
+        builder, func, TypeRange(), kHeapSortFuncNamePrefix, nx, ny, isCoo,
+        args.drop_back(nTrailingP), createHeapSortFunc);
+    builder.create<func::CallOp>(loc, heapSortFunc, TypeRange(),
+                                 ValueRange(args.drop_back(nTrailingP)));
+
+    // When depth doesn't exceed limit.
+    builder.setInsertionPointToStart(&depthIf.getElseRegion().front());
+  }
+
+  createQuickSort(builder, module, func, args, nx, ny, isCoo, nTrailingP);
+
+  if (isHybrid) {
+    // Restore depthLimit.
+    builder.setInsertionPointAfter(depthIf);
+    builder.create<memref::StoreOp>(loc, savedDepthLimit, pDepthLimit);
+  }
+
+  // After the if-stmt.
+  builder.setInsertionPointAfter(ifOp);
+  builder.create<func::ReturnOp>(loc);
+}
+
 /// Implements the rewriting for operator sort and sort_coo.
 template <typename OpTy>
 LogicalResult matchAndRewriteSortOp(OpTy op, ValueRange xys, uint64_t nx,
                                     uint64_t ny, bool isCoo,
                                     PatternRewriter &rewriter) {
   Location loc = op.getLoc();
   SmallVector<Value> operands{constantIndex(rewriter, loc, 0), op.getN()};
 
   // Convert `values` to have dynamic shape and append them to `operands`.
   for (Value v : xys) {
     auto mtp = getMemRefType(v);
     if (!mtp.isDynamicDim(0)) {
       auto newMtp =
           MemRefType::get({ShapedType::kDynamic}, mtp.getElementType());
       v = rewriter.create<memref::CastOp>(loc, newMtp, v);
     }
     operands.push_back(v);
   }
 
   auto insertPoint = op->template getParentOfType<func::FuncOp>();
   SmallString<32> funcName;
   FuncGeneratorType funcGenerator;
   uint32_t nTrailingP = 0;
   switch (op.getAlgorithm()) {
-  case SparseTensorSortKind::HybridQuickSort:
+  case SparseTensorSortKind::HybridQuickSort: {
+    funcName = kHybridQuickSortFuncNamePrefix;
+    funcGenerator = createQuickSortFunc;
+    nTrailingP = 1;
+    Value pDepthLimit = rewriter.create<memref::AllocaOp>(
+        loc, MemRefType::get({}, rewriter.getI64Type()));
+    operands.push_back(pDepthLimit);
+    // As a heuristics, set depthLimit = 2 * log2(n).
+    Value lo = operands[loIdx];
+    Value hi = operands[hiIdx];
+    Value len = rewriter.create<arith::IndexCastOp>(
+        loc, rewriter.getI64Type(),
+        rewriter.create<arith::SubIOp>(loc, hi, lo));
+    Value depthLimit = rewriter.create<arith::SubIOp>(
+        loc, constantI64(rewriter, loc, 64),
+        rewriter.create<math::CountLeadingZerosOp>(loc, len));
+    depthLimit = rewriter.create<arith::ShLIOp>(loc, depthLimit,
+                                                constantI64(rewriter, loc, 1));
+    rewriter.create<memref::StoreOp>(loc, depthLimit, pDepthLimit);
+    break;
+  }
   case SparseTensorSortKind::QuickSort:
-    funcName = kSortNonstableFuncNamePrefix;
-    funcGenerator = createSortNonstableFunc;
+    funcName = kQuickSortFuncNamePrefix;
+    funcGenerator = createQuickSortFunc;
     break;
   case SparseTensorSortKind::InsertionSortStable:
     funcName = kSortStableFuncNamePrefix;
     funcGenerator = createSortStableFunc;
     break;
   case SparseTensorSortKind::HeapSort:
     funcName = kHeapSortFuncNamePrefix;
     funcGenerator = createHeapSortFunc;

mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir

 // CHECK:             } else {
 // CHECK:               scf.yield %[[I2]], %[[J2]], %[[Pa]]
 // CHECK:             }
 // CHECK:             scf.yield %[[If:.*]]#0, %[[If]]#1, %[[If]]#2
 // CHECK:           }
 // CHECK:           return %[[W:.*]]#2
 // CHECK:         }
 
-// CHECK-LABEL:   func.func private @_sparse_sort_nonstable_1_i8_f32_index(
+// CHECK-LABEL:   func.func private @_sparse_qsort_1_i8_f32_index(
 // CHECK-SAME:                                                   %[[L:arg0]]: index,
 // CHECK-SAME:                                                   %[[H:.*]]: index,
 // CHECK-SAME:                                                   %[[X0:.*]]: memref<?xi8>,
 // CHECK-SAME:                                                   %[[Y0:.*]]: memref<?xf32>,
 // CHECK-SAME:                                                   %[[Y1:.*]]: memref<?xindex>) {
 // CHECK:           %[[C1:.*]] = arith.constant 1
-// CHECK:           %[[COND:.*]] = arith.cmpi ult, %[[L]], %[[H]]
+// CHECK:           %[[Lb:.*]] = arith.addi %[[L]], %[[C1]]
+// CHECK:           %[[COND:.*]] = arith.cmpi ult, %[[Lb]], %[[H]]
 // CHECK:           scf.if %[[COND]] {
 // CHECK:             %[[P:.*]] = func.call @_sparse_partition_1_i8_f32_index(%[[L]], %[[H]], %[[X0]], %[[Y0]], %[[Y1]])
-// CHECK:             func.call @_sparse_sort_nonstable_1_i8_f32_index(%[[L]], %[[P]], %[[X0]], %[[Y0]], %[[Y1]])
+// CHECK:             func.call @_sparse_qsort_1_i8_f32_index(%[[L]], %[[P]], %[[X0]], %[[Y0]], %[[Y1]])
 // CHECK:             %[[P2:.*]] = arith.addi %[[P]], %[[C1]] : index
-// CHECK:             func.call @_sparse_sort_nonstable_1_i8_f32_index(%[[P2]], %[[H]], %[[X0]], %[[Y0]], %[[Y1]])
+// CHECK:             func.call @_sparse_qsort_1_i8_f32_index(%[[P2]], %[[H]], %[[X0]], %[[Y0]], %[[Y1]])
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 
-// CHECK-LABEL:   func.func @sparse_sort_1d2v(
+// CHECK-LABEL:   func.func @sparse_sort_1d2v_quick(
 // CHECK-SAME:                                %[[N:.*]]: index,
 // CHECK-SAME:                                %[[X0:.*]]: memref<10xi8>,
 // CHECK-SAME:                                %[[Y0:.*]]: memref<?xf32>,
 // CHECK-SAME:                                %[[Y1:.*]]: memref<10xindex>) -> (memref<10xi8>, memref<?xf32>, memref<10xindex>) {
 // CHECK:           %[[C0:.*]] = arith.constant 0
 // CHECK:           %[[DX0:.*]] = memref.cast %[[X0]] : memref<10xi8> to memref<?xi8>
 // CHECK:           %[[DY1:.*]] = memref.cast %[[Y1]] : memref<10xindex> to memref<?xindex>
-// CHECK:           call @_sparse_sort_nonstable_1_i8_f32_index(%[[C0]], %[[N]], %[[DX0]], %[[Y0]], %[[DY1]])
+// CHECK:           call @_sparse_qsort_1_i8_f32_index(%[[C0]], %[[N]], %[[DX0]], %[[Y0]], %[[DY1]])
 // CHECK:           return %[[X0]], %[[Y0]], %[[Y1]]
 // CHECK:         }
-func.func @sparse_sort_1d2v(%arg0: index, %arg1: memref<10xi8>, %arg2: memref<?xf32>, %arg3: memref<10xindex>)
+func.func @sparse_sort_1d2v_quick(%arg0: index, %arg1: memref<10xi8>, %arg2: memref<?xf32>, %arg3: memref<10xindex>)
    -> (memref<10xi8>, memref<?xf32>, memref<10xindex>) {
-  sparse_tensor.sort hybrid_quick_sort %arg0, %arg1 jointly %arg2, %arg3 : memref<10xi8> jointly memref<?xf32>, memref<10xindex>
+  sparse_tensor.sort quick_sort %arg0, %arg1 jointly %arg2, %arg3 : memref<10xi8> jointly memref<?xf32>, memref<10xindex>
   return %arg1, %arg2, %arg3 : memref<10xi8>, memref<?xf32>, memref<10xindex>
 }
 
 // -----
 
 // Only check the generated supporting function now. We have integration test
 // to verify correctness of the generated code.
 //
 // CHECK-DAG:     func.func private @_sparse_less_than_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> i1 {
 // CHECK-DAG:     func.func private @_sparse_compare_eq_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> i1 {
 // CHECK-DAG:     func.func private @_sparse_partition_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> index {
-// CHECK-DAG:     func.func private @_sparse_sort_nonstable_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) {
-// CHECK-LABEL:   func.func @sparse_sort_3d
-func.func @sparse_sort_3d(%arg0: index, %arg1: memref<10xindex>, %arg2: memref<?xindex>, %arg3: memref<10xindex>) -> (memref<10xindex>, memref<?xindex>, memref<10xindex>) {
+// CHECK-DAG:     func.func private @_sparse_qsort_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) {
+// CHECK-LABEL:   func.func @sparse_sort_3d_quick
+func.func @sparse_sort_3d_quick(%arg0: index, %arg1: memref<10xindex>, %arg2: memref<?xindex>, %arg3: memref<10xindex>) -> (memref<10xindex>, memref<?xindex>, memref<10xindex>) {
+  sparse_tensor.sort quick_sort %arg0, %arg1, %arg2, %arg3 : memref<10xindex>, memref<?xindex>, memref<10xindex>
+  return %arg1, %arg2, %arg3 : memref<10xindex>, memref<?xindex>, memref<10xindex>
+}
+
+// -----
+
+// Only check the generated supporting function now. We have integration test
+// to verify correctness of the generated code.
+//
+// CHECK-DAG:     func.func private @_sparse_less_than_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> i1 {
+// CHECK-DAG:     func.func private @_sparse_binary_search_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> index {
+// CHECK-DAG:     func.func private @_sparse_sort_stable_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) {
+// CHECK-DAG:     func.func private @_sparse_shift_down_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>, %arg5: index) {
+// CHECK-DAG:     func.func private @_sparse_heap_sort_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) {
+// CHECK-DAG:     func.func private @_sparse_compare_eq_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> i1 {
+// CHECK-DAG:     func.func private @_sparse_partition_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>) -> index {
+// CHECK-DAG:     func.func private @_sparse_hybrid_qsort_3_index(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xindex>, %arg4: memref<?xindex>, %arg5: memref<i64>) {
+// CHECK-LABEL:   func.func @sparse_sort_3d_hybrid
+func.func @sparse_sort_3d_hybrid(%arg0: index, %arg1: memref<10xindex>, %arg2: memref<?xindex>, %arg3: memref<10xindex>) -> (memref<10xindex>, memref<?xindex>, memref<10xindex>) {
   sparse_tensor.sort hybrid_quick_sort %arg0, %arg1, %arg2, %arg3 : memref<10xindex>, memref<?xindex>, memref<10xindex>
   return %arg1, %arg2, %arg3 : memref<10xindex>, memref<?xindex>, memref<10xindex>
 }
 
 // -----
 
 // Only check the generated supporting functions. We have integration test to
 // verify correctness of the generated code.

 // -----
 
 // Only check the generated supporting functions. We have integration test to
 // verify correctness of the generated code.
 //
 // CHECK-DAG:     func.func private @_sparse_less_than_2_index_coo_1(%arg0: index, %arg1: index, %arg2: memref<?xindex>) -> i1 {
 // CHECK-DAG:     func.func private @_sparse_compare_eq_2_index_coo_1(%arg0: index, %arg1: index, %arg2: memref<?xindex>) -> i1 {
 // CHECK-DAG:     func.func private @_sparse_partition_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index {
-// CHECK-DAG:     func.func private @_sparse_sort_nonstable_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
-// CHECK-LABEL:   func.func @sparse_sort_coo
-func.func @sparse_sort_coo(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) {
+// CHECK-DAG:     func.func private @_sparse_qsort_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
+// CHECK-LABEL:   func.func @sparse_sort_coo_quick
+func.func @sparse_sort_coo_quick(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) {
+  sparse_tensor.sort_coo quick_sort %arg0, %arg1 jointly %arg2, %arg3 {nx = 2 : index, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32>
+  return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32>
+}
+
+// -----
+
+// Only check the generated supporting functions. We have integration test to
+// verify correctness of the generated code.
+//
+// CHECK-DAG:     func.func private @_sparse_less_than_2_index_coo_1(%arg0: index, %arg1: index, %arg2: memref<?xindex>) -> i1 {
+// CHECK-DAG:     func.func private @_sparse_binary_search_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index {
+// CHECK-DAG:     func.func private @_sparse_sort_stable_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
+// CHECK-DAG:     func.func private @_sparse_shift_down_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: index) {
+// CHECK-DAG:     func.func private @_sparse_heap_sort_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
+// CHECK-DAG:     func.func private @_sparse_compare_eq_2_index_coo_1(%arg0: index, %arg1: index, %arg2: memref<?xindex>) -> i1 {
+// CHECK-DAG:     func.func private @_sparse_partition_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) -> index {
+// CHECK-DAG:     func.func private @_sparse_hybrid_qsort_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: memref<i64>) {
+// CHECK-LABEL:   func.func @sparse_sort_coo_hybrid
+func.func @sparse_sort_coo_hybrid(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) {
   sparse_tensor.sort_coo hybrid_quick_sort %arg0, %arg1 jointly %arg2, %arg3 {nx = 2 : index, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32>
   return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32>
 }
 
 // -----
 
 // Only check the generated supporting functions. We have integration test to
 // verify correctness of the generated code.

 // CHECK-DAG:     func.func private @_sparse_less_than_2_index_coo_1(%arg0: index, %arg1: index, %arg2: memref<?xindex>) -> i1 {
 // CHECK-DAG:     func.func private @_sparse_shift_down_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: index) {
 // CHECK-DAG:     func.func private @_sparse_heap_sort_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
 // CHECK-LABEL:   func.func @sparse_sort_coo_heap
 func.func @sparse_sort_coo_heap(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) {
   sparse_tensor.sort_coo heap_sort %arg0, %arg1 jointly %arg2, %arg3 {nx = 2 : index, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32>
   return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32>
 }
 No newline at end of file

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort.mlir

func.func @entry() {
 
     // Prepare a buffer.
     %x0s = memref.alloc() : memref<5xi32>
     %x0 = memref.cast %x0s : memref<5xi32> to memref<?xi32>
     call @storeValuesTo(%x0, %c10, %c2, %c0, %c5, %c1)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
 
     // Sort 0 elements.
+    // Quick sort.
     // CHECK: [10,  2,  0,  5,  1]
-    sparse_tensor.sort hybrid_quick_sort %i0, %x0 : memref<?xi32>
+    sparse_tensor.sort quick_sort %i0, %x0 : memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
     // Stable sort.
     // CHECK: [10,  2,  0,  5,  1]
     sparse_tensor.sort insertion_sort_stable %i0, %x0 : memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
     // Heap sort.
     // CHECK: [10,  2,  0,  5,  1]
     sparse_tensor.sort heap_sort %i0, %x0 : memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
+    // Hybrid sort.
+    // CHECK: [10,  2,  0,  5,  1]
+    sparse_tensor.sort hybrid_quick_sort %i0, %x0 : memref<?xi32>
+    call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
 
     // Sort the first 4 elements, with the last valid value untouched.
+    // Quick sort.
     // CHECK: [0,  2,  5, 10,  1]
-    sparse_tensor.sort hybrid_quick_sort %i4, %x0 : memref<?xi32>
+    sparse_tensor.sort quick_sort %i4, %x0 : memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
     // Stable sort.
     // CHECK: [0,  2,  5,  10,  1]
     call @storeValuesTo(%x0, %c10, %c2, %c0, %c5, %c1)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
     sparse_tensor.sort insertion_sort_stable %i4, %x0 : memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
     // Heap sort.
     // CHECK: [0,  2,  5,  10,  1]
     call @storeValuesTo(%x0, %c10, %c2, %c0, %c5, %c1)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
     sparse_tensor.sort heap_sort %i4, %x0 : memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
+    // Hybrid sort.
+    // CHECK: [0,  2,  5, 10,  1]
+    sparse_tensor.sort hybrid_quick_sort %i4, %x0 : memref<?xi32>
+    call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
 
     // Prepare more buffers of different dimensions.
     %x1s = memref.alloc() : memref<10xi32>
     %x1 = memref.cast %x1s : memref<10xi32> to memref<?xi32>
     %x2s = memref.alloc() : memref<6xi32>
     %x2 = memref.cast %x2s : memref<6xi32> to memref<?xi32>
     %y0s = memref.alloc() : memref<7xi32>
     %y0 = memref.cast %y0s : memref<7xi32> to memref<?xi32>
 
     // Sort "parallel arrays".
     // CHECK: [1,  1,  2,  5,  10]
     // CHECK: [3,  3,  1,  10,  1
     // CHECK: [9,  9,  4,  7,  2
     // CHECK: [7,  8,  10,  9,  6
     call @storeValuesTo(%x0, %c10, %c2, %c1, %c5, %c1)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
     call @storeValuesTo(%x1, %c1, %c1, %c3, %c10, %c3)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
     call @storeValuesTo(%x2, %c2, %c4, %c9, %c7, %c9)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
     call @storeValuesTo(%y0, %c6, %c10, %c8, %c9, %c7)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
-    sparse_tensor.sort hybrid_quick_sort %i5, %x0, %x1, %x2 jointly %y0
+    sparse_tensor.sort quick_sort %i5, %x0, %x1, %x2 jointly %y0
       : memref<?xi32>, memref<?xi32>, memref<?xi32> jointly memref<?xi32>
     call @printMemref1dI32(%x0) : (memref<?xi32>) -> ()
     call @printMemref1dI32(%x1) : (memref<?xi32>) -> ()
     call @printMemref1dI32(%x2) : (memref<?xi32>) -> ()
     call @printMemref1dI32(%y0) : (memref<?xi32>) -> ()
     // Stable sort.
     // CHECK: [1,  1,  2,  5,  10]
     // CHECK: [3,  3,  1,  10,  1

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_rewrite_sort_coo.mlir

func.func @entry() {
     call @storeValuesToStrided(%x1, %c1, %c1, %c3, %c10, %c3)
       : (memref<?xi32, strided<[4], offset: ?>>, i32, i32, i32, i32, i32) -> ()
     call @storeValuesToStrided(%x2, %c2, %c4, %c9, %c7, %c9)
       : (memref<?xi32, strided<[4], offset: ?>>, i32, i32, i32, i32, i32) -> ()
     call @storeValuesToStrided(%y0, %c6, %c10, %c8, %c9, %c7)
       : (memref<?xi32, strided<[4], offset: ?>>, i32, i32, i32, i32, i32) -> ()
     call @storeValuesTo(%y1, %c5, %c7, %c4, %c9, %c7)
       : (memref<?xi32>, i32, i32, i32, i32, i32) -> ()
-    sparse_tensor.sort_coo hybrid_quick_sort %i5, %xy jointly %y1 {nx = 3 : index, ny = 1 : index}
+    sparse_tensor.sort_coo quick_sort %i5, %xy jointly %y1 {nx = 3 : index, ny = 1 : index}
       : memref<?xi32> jointly memref<?xi32>
     %x0v = vector.transfer_read %x0[%i0], %c100: memref<?xi32, strided<[4], offset: ?>>, vector<5xi32>
     vector.print %x0v : vector<5xi32>
     %x1v = vector.transfer_read %x1[%i0], %c100: memref<?xi32, strided<[4], offset: ?>>, vector<5xi32>
     vector.print %x1v : vector<5xi32>
     %x2v = vector.transfer_read %x2[%i0], %c100: memref<?xi32, strided<[4], offset: ?>>, vector<5xi32>
     vector.print %x2v : vector<5xi32>
     %y0v = vector.transfer_read %y0[%i0], %c100: memref<?xi32, strided<[4], offset: ?>>, vector<5xi32>