[mlir][sparse] Extend sorting function generator to support operand beyond (lo, hi, xs, ys).

[bixia]

Summary

This is to prepare for implementing a hybrid quick sort, which switches to heap
sort when the recursive depth exceeds certain limits.

[aartbik]

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

254

I find this a very distracting idiom to force use of an otherwise unused parameter. Is there no better way to achieve this?

[aartbik]

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

284

not really used in this revision?

[wrengr]

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

254

Assuming the issue here is warnings about unused variables, then iirc you should be able to rephrase the prototype to (..., uint32_t = 0) —i.e., simply drop the name of the variable.

If there are more complicated problems at hand, then I'm thinking you should rephrase the typedef to something like template <typename ...Args> using FuncGeneratorType = function_ref<void(blah blah blah, Args&&...)>; so that this particular function instantiates that template at the empty Args

[bixia]

Removed irrelevant changes.

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

284

Reversed this change.

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

254

Please refer to discussion (void)unused_var; is in LLVM style guide, when we have c++ 17 we can use [[maybe_unused]] attribute.
I am keeping (void)nTrailingP here and in other places, it also allows "assert" in debug mode.

[bixia]

Rebase.

[bixia]

Rebase.

[bixia]

Rebase.

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

 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 kSortStableFuncNamePrefix[] =
     "_sparse_sort_stable_";
 
-using FuncGeneratorType = function_ref<void(OpBuilder &, ModuleOp, func::FuncOp,
-                                            uint64_t, uint64_t, bool)>;
+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,
                                          StringRef namePrefix, uint64_t nx,
                                          uint64_t ny, bool isCoo,
                                          ValueRange operands) {
for (Value v : operands.drop_front(xStartIdx + yBufferOffset))
     nameOstream << "_" << v.getType().cast<MemRefType>().getElementType();
 }
 
 /// Looks up a function that is appropriate for the given operands being
 /// sorted, and creates such a function if it doesn't exist yet. The
 /// parameters `nx` and `ny` tell the number of x and y values provided
 /// by the buffer in xStartIdx, and `isCoo` indicates whether the instruction
 /// being processed is a sparse_tensor.sort or sparse_tensor.sort_coo.
+//
+// All sorting function generators take (lo, hi, xs, ys) in `operands` as
+// parameters for the sorting functions. Other parameters, such as the recursive
+// call depth, are appended to the end of the parameter list as
+// "trailing parameters".
 static FlatSymbolRefAttr
 getMangledSortHelperFunc(OpBuilder &builder, func::FuncOp insertPoint,
                          TypeRange resultTypes, StringRef namePrefix,
                          uint64_t nx, uint64_t ny, bool isCoo,
-                         ValueRange operands, FuncGeneratorType createFunc) {
+                         ValueRange operands, FuncGeneratorType createFunc,
+                         uint32_t nTrailingP = 0) {
   SmallString<32> nameBuffer;
   llvm::raw_svector_ostream nameOstream(nameBuffer);
   getMangledSortHelperFuncName(nameOstream, namePrefix, nx, ny, isCoo,
-                               operands);
+                               operands.drop_back(nTrailingP));
 
   ModuleOp module = insertPoint->getParentOfType<ModuleOp>();
   MLIRContext *context = module.getContext();
   auto result = SymbolRefAttr::get(context, nameOstream.str());
   auto func = module.lookupSymbol<func::FuncOp>(result.getAttr());
 
   if (!func) {
     // Create the function.
     OpBuilder::InsertionGuard insertionGuard(builder);
     builder.setInsertionPoint(insertPoint);
     Location loc = insertPoint.getLoc();
     func = builder.create<func::FuncOp>(
         loc, nameOstream.str(),
         FunctionType::get(context, operands.getTypes(), resultTypes));
     func.setPrivate();
-    createFunc(builder, module, func, nx, ny, isCoo);
+    createFunc(builder, module, func, nx, ny, isCoo, nTrailingP);
   }
 
   return result;
 }
 
 /// Creates a code block to process each pair of (xs[i], xs[j]) for sorting.
 /// The code to process the value pairs is generated by `bodyBuilder`.
 static void forEachIJPairInXs(

 //     return false;
 //   else
 //     if (x1[i] != x1[j])
 //       return false;
 //     else if (x2[2] != x2[j]))
 //       and so on ...
 static void createEqCompareFunc(OpBuilder &builder, ModuleOp unused,
                                 func::FuncOp func, uint64_t nx, uint64_t ny,
-                                bool isCoo) {
+                                bool isCoo, uint32_t nTrailingP = 0) {
+  // Compare functions don't use trailing parameters.
+  (void)nTrailingP;
+  assert(nTrailingP == 0);
   createCompareFuncImplementation(builder, unused, func, nx, ny, isCoo,
                                   createEqCompare);
 }
 
 /// Generates an if-statement to compare whether x[i] is less than x[j].
 static scf::IfOp createLessThanCompare(OpBuilder &builder, Location loc,
                                        Value i, Value j, Value x,
                                        bool isLastDim) {

 //     return false;
 //   else
 //     if (x1[i] < x1[j])
 //       return true;
 //     else if (x1[j] < x1[i]))
 //       and so on ...
 static void createLessThanFunc(OpBuilder &builder, ModuleOp unused,
                                func::FuncOp func, uint64_t nx, uint64_t ny,
-                               bool isCoo) {
+                               bool isCoo, uint32_t nTrailingP = 0) {
+  // Compare functions don't use trailing parameters.
+  (void)nTrailingP;
+  assert(nTrailingP == 0);
   createCompareFuncImplementation(builder, unused, func, nx, ny, isCoo,
                                   createLessThanCompare);
 }
 
 /// Creates a function to use a binary search to find the insertion point for
 /// inserting xs[hi] to the sorted values xs[lo..hi).
 //
 // The generate IR corresponds to this C like algorithm:
 //   p = hi
 //   while (lo < hi)
 //      mid = (lo + hi) >> 1
 //      if (xs[p] < xs[mid])
 //        hi = mid
 //      else
 //        lo = mid - 1
 //   return lo;
 //
 static void createBinarySearchFunc(OpBuilder &builder, ModuleOp module,
                                    func::FuncOp func, uint64_t nx, uint64_t ny,
-                                   bool isCoo) {
+                                   bool isCoo, uint32_t nTrailingP = 0) {
+  // Binary search doesn't use trailing parameters.
+  (void)nTrailingP;
+  assert(nTrailingP == 0);
   OpBuilder::InsertionGuard insertionGuard(builder);
   Block *entryBlock = func.addEntryBlock();
   builder.setInsertionPointToStart(entryBlock);
 
   Location loc = func.getLoc();
   ValueRange args = entryBlock->getArguments();
   Value p = args[hiIdx];
-  SmallVector<Type, 2> types(2, p.getType()); // only two
+  SmallVector<Type, 2> types(2, p.getType()); // Only two types.
   scf::WhileOp whileOp = builder.create<scf::WhileOp>(
       loc, types, SmallVector<Value, 2>{args[loIdx], args[hiIdx]});
 
   // The before-region of the WhileOp.
   Block *before =
       builder.createBlock(&whileOp.getBefore(), {}, types, {loc, loc});
   builder.setInsertionPointToEnd(before);
   Value cond1 = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
static void createBinarySearchFunc(OpBuilder &builder, ModuleOp module,
   // Compare xs[p] < xs[mid].
   SmallVector<Value> compareOperands{p, mid};
   uint64_t numXBuffers = isCoo ? 1 : nx;
   compareOperands.append(args.begin() + xStartIdx,
                          args.begin() + xStartIdx + numXBuffers);
   Type i1Type = IntegerType::get(module.getContext(), 1, IntegerType::Signless);
   FlatSymbolRefAttr lessThanFunc = getMangledSortHelperFunc(
       builder, func, {i1Type}, kLessThanFuncNamePrefix, nx, ny, isCoo,
-      compareOperands, createLessThanFunc);
+      compareOperands, createLessThanFunc, nTrailingP);
   Value cond2 = builder
                     .create<func::CallOp>(loc, lessThanFunc, TypeRange{i1Type},
                                           compareOperands)
                     .getResult(0);
 
   // Update lo and hi for the WhileOp as follows:
   //   if (xs[p] < xs[mid]))
   //     hi = mid;

 //         --j;
 //       }
 //     }
 //   }
 //   return p
 //   }
 static void createPartitionFunc(OpBuilder &builder, ModuleOp module,
                                 func::FuncOp func, uint64_t nx, uint64_t ny,
-                                bool isCoo) {
+                                bool isCoo, uint32_t nTrailingP = 0) {
+  // Quick sort partition doesn't use trailing parameters.
+  (void)nTrailingP;
+  assert(nTrailingP == 0);
   OpBuilder::InsertionGuard insertionGuard(builder);
 
   Block *entryBlock = func.addEntryBlock();
   builder.setInsertionPointToStart(entryBlock);
 
   Location loc = func.getLoc();
   ValueRange args = entryBlock->getArguments();
   Value lo = args[loIdx];

 //   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) {
+                                    bool isCoo, uint32_t nTrailingP) {
+  (void)nTrailingP;
   OpBuilder::InsertionGuard insertionGuard(builder);
   Block *entryBlock = func.addEntryBlock();
   builder.setInsertionPointToStart(entryBlock);
 
   MLIRContext *context = module.getContext();
   Location loc = func.getLoc();
   ValueRange args = entryBlock->getArguments();
   Value lo = args[loIdx];

 //      p = binarySearch(lo, i-1, data)
 //      for (j = 0; j > i - p; j++)
 //        data[i-j] = data[i-j-1]
 //      data[p] = d
 //   }
 // }
 static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
                                  func::FuncOp func, uint64_t nx, uint64_t ny,
-                                 bool isCoo) {
+                                 bool isCoo, uint32_t nTrailingP) {
+  // Stable sort function doesn't use trailing parameters.
+  (void)nTrailingP;
+  assert(nTrailingP == 0);
   OpBuilder::InsertionGuard insertionGuard(builder);
   Block *entryBlock = func.addEntryBlock();
   builder.setInsertionPointToStart(entryBlock);
 
   MLIRContext *context = module.getContext();
   Location loc = func.getLoc();
   ValueRange args = entryBlock->getArguments();
   Value c1 = constantIndex(builder, loc, 1);
for (Value v : xys) {
     }
     operands.push_back(v);
   }
   auto insertPoint = op->template getParentOfType<func::FuncOp>();
   SmallString<32> funcName(op.getStable() ? kSortStableFuncNamePrefix
                                           : kSortNonstableFuncNamePrefix);
   FuncGeneratorType funcGenerator =
       op.getStable() ? createSortStableFunc : createSortNonstableFunc;
+  uint32_t nTrailingP = 0;
   FlatSymbolRefAttr func =
       getMangledSortHelperFunc(rewriter, insertPoint, TypeRange(), funcName, nx,
-                               ny, isCoo, operands, funcGenerator);
+                               ny, isCoo, operands, funcGenerator, nTrailingP);
   rewriter.replaceOpWithNewOp<func::CallOp>(op, func, TypeRange(), operands);
   return success();
 }
 
 //===---------------------------------------------------------------------===//
 // The actual sparse buffer rewriting rules.
 //===---------------------------------------------------------------------===//