[mlir][sparse] unifies sparse_tensor.sort_coo/sort into one operation. (#66722)

The use cases of the two operations are largely overlapped, let's
simplify it and only use one of them.

1 files changed, 142 insertions, 172 deletions

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp
index 029ecb070894..3181395a474c 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseBufferRewriting.cpp
@@ -45,46 +45,43 @@ 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)>;
+using FuncGeneratorType = function_ref<void(OpBuilder &, ModuleOp, func::FuncOp,
+                                            AffineMap, uint64_t, 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
+///   <namePrefix><xPerm>_<x type>_<y0 type>..._<yn type> for sort
+///   <namePrefix><xPerm>_<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) {
-  nameOstream << namePrefix << nx << "_"
-              << getMemRefType(operands[xStartIdx]).getElementType();
+                                         StringRef namePrefix, AffineMap xPerm,
+                                         uint64_t ny, ValueRange operands) {
+  nameOstream << namePrefix;
+  for (auto res : xPerm.getResults())
+    nameOstream << res.cast<AffineDimExpr>().getPosition() << "_";
 
-  if (isCoo)
-    nameOstream << "_coo_" << ny;
+  nameOstream << getMemRefType(operands[xStartIdx]).getElementType();
+  nameOstream << "_coo_" << ny;
 
-  uint64_t yBufferOffset = isCoo ? 1 : nx;
+  constexpr uint64_t yBufferOffset = 1;
   for (Value v : operands.drop_front(xStartIdx + yBufferOffset))
     nameOstream << "_" << getMemRefType(v).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.
+/// parameters `xPerm` and `ny` tell the number of x and y values provided
+/// by the buffer in xStartIdx.
 //
 // 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,
-                         uint32_t nTrailingP = 0) {
+static FlatSymbolRefAttr getMangledSortHelperFunc(
+    OpBuilder &builder, func::FuncOp insertPoint, TypeRange resultTypes,
+    StringRef namePrefix, AffineMap xPerm, uint64_t ny, ValueRange operands,
+    FuncGeneratorType createFunc, uint32_t nTrailingP = 0) {
   SmallString<32> nameBuffer;
   llvm::raw_svector_ostream nameOstream(nameBuffer);
-  getMangledSortHelperFuncName(nameOstream, namePrefix, nx, ny, isCoo,
+  getMangledSortHelperFuncName(nameOstream, namePrefix, xPerm, ny,
                                operands.drop_back(nTrailingP));
 
   ModuleOp module = insertPoint->getParentOfType<ModuleOp>();
@@ -101,7 +98,7 @@ getMangledSortHelperFunc(OpBuilder &builder, func::FuncOp insertPoint,
         loc, nameOstream.str(),
         FunctionType::get(context, operands.getTypes(), resultTypes));
     func.setPrivate();
-    createFunc(builder, module, func, nx, ny, isCoo, nTrailingP);
+    createFunc(builder, module, func, xPerm, ny, nTrailingP);
   }
 
   return result;
@@ -110,27 +107,19 @@ getMangledSortHelperFunc(OpBuilder &builder, func::FuncOp insertPoint,
 /// 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(
-    OpBuilder &builder, Location loc, ValueRange args, uint64_t nx, uint64_t ny,
-    bool isCoo, function_ref<void(uint64_t, Value, Value, Value)> bodyBuilder) {
-  Value iOffset, jOffset;
-  if (isCoo) {
-    Value cstep = constantIndex(builder, loc, nx + ny);
-    iOffset = builder.create<arith::MulIOp>(loc, args[0], cstep);
-    jOffset = builder.create<arith::MulIOp>(loc, args[1], cstep);
-  }
-  for (uint64_t k = 0; k < nx; k++) {
-    scf::IfOp ifOp;
-    Value i, j, buffer;
-    if (isCoo) {
-      Value ck = constantIndex(builder, loc, k);
-      i = builder.create<arith::AddIOp>(loc, ck, iOffset);
-      j = builder.create<arith::AddIOp>(loc, ck, jOffset);
-      buffer = args[xStartIdx];
-    } else {
-      i = args[0];
-      j = args[1];
-      buffer = args[xStartIdx + k];
-    }
+    OpBuilder &builder, Location loc, ValueRange args, AffineMap xPerm,
+    uint64_t ny,
+    function_ref<void(uint64_t, Value, Value, Value)> bodyBuilder) {
+  Value cstep = constantIndex(builder, loc, xPerm.getNumResults() + ny);
+  Value iOffset = builder.create<arith::MulIOp>(loc, args[0], cstep);
+  Value jOffset = builder.create<arith::MulIOp>(loc, args[1], cstep);
+  for (unsigned k = 0, e = xPerm.getNumResults(); k < e; k++) {
+    unsigned actualK = xPerm.getResult(k).cast<AffineDimExpr>().getPosition();
+    Value ak = constantIndex(builder, loc, actualK);
+    Value i = builder.create<arith::AddIOp>(loc, ak, iOffset);
+    Value j = builder.create<arith::AddIOp>(loc, ak, jOffset);
+    Value buffer = args[xStartIdx];
+
     bodyBuilder(k, i, j, buffer);
   }
 }
@@ -138,21 +127,28 @@ static void forEachIJPairInXs(
 /// Creates a code block to process each pair of (xys[i], xys[j]) for sorting.
 /// The code to process the value pairs is generated by `bodyBuilder`.
 static void forEachIJPairInAllBuffers(
-    OpBuilder &builder, Location loc, ValueRange args, uint64_t nx, uint64_t ny,
-    bool isCoo, function_ref<void(uint64_t, Value, Value, Value)> bodyBuilder) {
-
-  // Create code for the first (nx + ny) buffers. When isCoo==true, these
-  // logical buffers are all from the xy buffer of the sort_coo operator.
-  forEachIJPairInXs(builder, loc, args, nx + ny, 0, isCoo, bodyBuilder);
+    OpBuilder &builder, Location loc, ValueRange args, AffineMap xPerm,
+    uint64_t ny,
+    function_ref<void(uint64_t, Value, Value, Value)> bodyBuilder) {
+
+  // Create code for the first (xPerm + ny) buffers.
+  SmallVector<AffineExpr> exps(xPerm.getResults().begin(),
+                               xPerm.getResults().end());
+  for (unsigned y = 0; y < ny; y++) {
+    exps.push_back(builder.getAffineDimExpr(y + xPerm.getNumResults()));
+  }
+  AffineMap xyPerm = AffineMap::get(exps.size(), 0, exps, builder.getContext());
+  assert(xyPerm.isPermutation());
 
-  uint64_t numHandledBuffers = isCoo ? 1 : nx + ny;
+  forEachIJPairInXs(builder, loc, args, xyPerm, 0, bodyBuilder);
 
+  constexpr uint64_t numHandledBuffers = 1;
   // Create code for the remaining buffers.
   Value i = args[0];
   Value j = args[1];
   for (const auto &arg :
        llvm::enumerate(args.drop_front(xStartIdx + numHandledBuffers))) {
-    bodyBuilder(arg.index() + nx + ny, i, j, arg.value());
+    bodyBuilder(arg.index() + xPerm.getNumResults() + ny, i, j, arg.value());
   }
 }
 
@@ -168,7 +164,7 @@ static void forEachIJPairInAllBuffers(
 //     ...
 //     swap(yn[i], yn[j]);
 static void createSwap(OpBuilder &builder, Location loc, ValueRange args,
-                       uint64_t nx, uint64_t ny, bool isCoo) {
+                       AffineMap xPerm, uint64_t ny) {
   auto swapOnePair = [&](uint64_t unused, Value i, Value j, Value buffer) {
     Value vi = builder.create<memref::LoadOp>(loc, buffer, i);
     Value vj = builder.create<memref::LoadOp>(loc, buffer, j);
@@ -176,20 +172,20 @@ static void createSwap(OpBuilder &builder, Location loc, ValueRange args,
     builder.create<memref::StoreOp>(loc, vi, buffer, j);
   };
 
-  forEachIJPairInAllBuffers(builder, loc, args, nx, ny, isCoo, swapOnePair);
+  forEachIJPairInAllBuffers(builder, loc, args, xPerm, ny, swapOnePair);
 }
 
 /// Creates code to compare all the (xs[i], xs[j]) pairs. The method to compare
 /// each pair is create via `compareBuilder`.
 static Value createInlinedCompareImplementation(
-    OpBuilder &builder, Location loc, ValueRange args, uint64_t nx, uint64_t ny,
-    bool isCoo,
+    OpBuilder &builder, Location loc, ValueRange args, AffineMap xPerm,
+    uint64_t ny,
     function_ref<Value(OpBuilder &, Location, Value, Value, Value, bool, bool)>
         compareBuilder) {
   Value result;
   auto bodyBuilder = [&](uint64_t k, Value i, Value j, Value buffer) {
     bool isFirstDim = (k == 0);
-    bool isLastDim = (k == nx - 1);
+    bool isLastDim = (k == xPerm.getNumResults() - 1);
     Value val =
         compareBuilder(builder, loc, i, j, buffer, isFirstDim, isLastDim);
     if (isFirstDim) {
@@ -202,7 +198,7 @@ static Value createInlinedCompareImplementation(
     }
   };
 
-  forEachIJPairInXs(builder, loc, args, nx, ny, isCoo, bodyBuilder);
+  forEachIJPairInXs(builder, loc, args, xPerm, ny, bodyBuilder);
 
   builder.setInsertionPointAfterValue(result);
   return result;
@@ -252,12 +248,12 @@ static Value createEqCompare(OpBuilder &builder, Location loc, Value i, Value j,
 //     else if (x2[2] != x2[j]))
 //       and so on ...
 static Value createInlinedEqCompare(OpBuilder &builder, Location loc,
-                                    ValueRange args, uint64_t nx, uint64_t ny,
-                                    bool isCoo, uint32_t nTrailingP = 0) {
+                                    ValueRange args, AffineMap xPerm,
+                                    uint64_t ny, uint32_t nTrailingP = 0) {
   // Compare functions don't use trailing parameters.
   (void)nTrailingP;
   assert(nTrailingP == 0);
-  return createInlinedCompareImplementation(builder, loc, args, nx, ny, isCoo,
+  return createInlinedCompareImplementation(builder, loc, args, xPerm, ny,
                                             createEqCompare);
 }
 
@@ -306,12 +302,12 @@ static Value createLessThanCompare(OpBuilder &builder, Location loc, Value i,
 //   else
 //       and so on ...
 static Value createInlinedLessThan(OpBuilder &builder, Location loc,
-                                   ValueRange args, uint64_t nx, uint64_t ny,
-                                   bool isCoo, uint32_t nTrailingP = 0) {
+                                   ValueRange args, AffineMap xPerm,
+                                   uint64_t ny, uint32_t nTrailingP = 0) {
   // Compare functions don't use trailing parameters.
   (void)nTrailingP;
   assert(nTrailingP == 0);
-  return createInlinedCompareImplementation(builder, loc, args, nx, ny, isCoo,
+  return createInlinedCompareImplementation(builder, loc, args, xPerm, ny,
                                             createLessThanCompare);
 }
 
@@ -329,8 +325,8 @@ static Value createInlinedLessThan(OpBuilder &builder, Location loc,
 //   return lo;
 //
 static void createBinarySearchFunc(OpBuilder &builder, ModuleOp module,
-                                   func::FuncOp func, uint64_t nx, uint64_t ny,
-                                   bool isCoo, uint32_t nTrailingP = 0) {
+                                   func::FuncOp func, AffineMap xPerm,
+                                   uint64_t ny, uint32_t nTrailingP = 0) {
   // Binary search doesn't use trailing parameters.
   (void)nTrailingP;
   assert(nTrailingP == 0);
@@ -368,11 +364,10 @@ static void createBinarySearchFunc(OpBuilder &builder, ModuleOp module,
 
   // Compare xs[p] < xs[mid].
   SmallVector<Value> compareOperands{p, mid};
-  uint64_t numXBuffers = isCoo ? 1 : nx;
+  constexpr uint64_t numXBuffers = 1;
   compareOperands.append(args.begin() + xStartIdx,
                          args.begin() + xStartIdx + numXBuffers);
-  Value cond2 =
-      createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
+  Value cond2 = createInlinedLessThan(builder, loc, compareOperands, xPerm, ny);
   // Update lo and hi for the WhileOp as follows:
   //   if (xs[p] < xs[mid]))
   //     hi = mid;
@@ -392,10 +387,11 @@ static void createBinarySearchFunc(OpBuilder &builder, ModuleOp module,
 ///   while (xs[i] > xs[p]) i += step (step < 0)
 /// The routine returns i as well as a boolean value to indicate whether
 /// xs[i] == xs[p].
-static std::pair<Value, Value>
-createScanLoop(OpBuilder &builder, ModuleOp module, func::FuncOp func,
-               ValueRange xs, Value i, Value p, uint64_t nx, uint64_t ny,
-               bool isCoo, int step) {
+static std::pair<Value, Value> createScanLoop(OpBuilder &builder,
+                                              ModuleOp module,
+                                              func::FuncOp func, ValueRange xs,
+                                              Value i, Value p, AffineMap xPerm,
+                                              uint64_t ny, int step) {
   Location loc = func.getLoc();
   scf::WhileOp whileOp =
       builder.create<scf::WhileOp>(loc, TypeRange{i.getType()}, ValueRange{i});
@@ -413,8 +409,7 @@ createScanLoop(OpBuilder &builder, ModuleOp module, func::FuncOp func,
     compareOperands.push_back(before->getArgument(0));
   }
   compareOperands.append(xs.begin(), xs.end());
-  Value cond =
-      createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
+  Value cond = createInlinedLessThan(builder, loc, compareOperands, xPerm, ny);
   builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
 
   Block *after =
@@ -429,7 +424,7 @@ createScanLoop(OpBuilder &builder, ModuleOp module, func::FuncOp func,
   compareOperands[0] = i;
   compareOperands[1] = p;
   Value compareEq =
-      createInlinedEqCompare(builder, loc, compareOperands, nx, ny, isCoo);
+      createInlinedEqCompare(builder, loc, compareOperands, xPerm, ny);
 
   return std::make_pair(whileOp.getResult(0), compareEq);
 }
@@ -438,67 +433,63 @@ createScanLoop(OpBuilder &builder, ModuleOp module, func::FuncOp func,
 /// if compareFunc(data[b], data[a]) returns true. The new insertion point is
 /// right after the swap instructions.
 static scf::IfOp createCompareThenSwap(OpBuilder &builder, Location loc,
-                                       uint64_t nx, uint64_t ny, bool isCoo,
+                                       AffineMap xPerm, uint64_t ny,
                                        SmallVectorImpl<Value> &swapOperands,
                                        SmallVectorImpl<Value> &compareOperands,
                                        Value a, Value b) {
   // Compare(data[b], data[a]).
   compareOperands[0] = b;
   compareOperands[1] = a;
-  Value cond =
-      createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
+  Value cond = createInlinedLessThan(builder, loc, compareOperands, xPerm, ny);
   scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else=*/false);
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   swapOperands[0] = b;
   swapOperands[1] = a;
-  createSwap(builder, loc, swapOperands, nx, ny, isCoo);
+  createSwap(builder, loc, swapOperands, xPerm, ny);
   return ifOp;
 }
 
 /// Creates code to insert the 3rd element to a list of two sorted elements.
-static void createInsert3rd(OpBuilder &builder, Location loc, uint64_t nx,
-                            uint64_t ny, bool isCoo,
-                            SmallVectorImpl<Value> &swapOperands,
+static void createInsert3rd(OpBuilder &builder, Location loc, AffineMap xPerm,
+                            uint64_t ny, SmallVectorImpl<Value> &swapOperands,
                             SmallVectorImpl<Value> &compareOperands, Value v0,
                             Value v1, Value v2) {
-  scf::IfOp ifOp = createCompareThenSwap(builder, loc, nx, ny, isCoo,
-                                         swapOperands, compareOperands, v1, v2);
-  createCompareThenSwap(builder, loc, nx, ny, isCoo, swapOperands,
-                        compareOperands, v0, v1);
+  scf::IfOp ifOp = createCompareThenSwap(builder, loc, xPerm, ny, swapOperands,
+                                         compareOperands, v1, v2);
+  createCompareThenSwap(builder, loc, xPerm, ny, swapOperands, compareOperands,
+                        v0, v1);
   builder.setInsertionPointAfter(ifOp);
 }
 
 /// Creates code to sort 3 elements.
-static void createSort3(OpBuilder &builder, Location loc, uint64_t nx,
-                        uint64_t ny, bool isCoo,
-                        SmallVectorImpl<Value> &swapOperands,
+static void createSort3(OpBuilder &builder, Location loc, AffineMap xPerm,
+                        uint64_t ny, SmallVectorImpl<Value> &swapOperands,
                         SmallVectorImpl<Value> &compareOperands, Value v0,
                         Value v1, Value v2) {
   // Sort the first 2 elements.
-  scf::IfOp ifOp1 = createCompareThenSwap(
-      builder, loc, nx, ny, isCoo, swapOperands, compareOperands, v0, v1);
+  scf::IfOp ifOp1 = createCompareThenSwap(builder, loc, xPerm, ny, swapOperands,
+                                          compareOperands, v0, v1);
   builder.setInsertionPointAfter(ifOp1);
 
   // Insert the 3th element.
-  createInsert3rd(builder, loc, nx, ny, isCoo, swapOperands, compareOperands,
-                  v0, v1, v2);
+  createInsert3rd(builder, loc, xPerm, ny, swapOperands, compareOperands, v0,
+                  v1, v2);
 }
 
 /// Creates code to sort 5 elements.
-static void createSort5(OpBuilder &builder, Location loc, uint64_t nx,
-                        uint64_t ny, bool isCoo,
-                        SmallVectorImpl<Value> &swapOperands,
+static void createSort5(OpBuilder &builder, Location loc, AffineMap xPerm,
+                        uint64_t ny, SmallVectorImpl<Value> &swapOperands,
                         SmallVectorImpl<Value> &compareOperands, Value v0,
                         Value v1, Value v2, Value v3, Value v4) {
   // Sort the first 3 elements.
-  createSort3(builder, loc, nx, ny, isCoo, swapOperands, compareOperands, v0,
-              v1, v2);
+  createSort3(builder, loc, xPerm, ny, swapOperands, compareOperands, v0, v1,
+              v2);
 
   auto insert4th = [&]() {
     scf::IfOp ifOp = createCompareThenSwap(
-        builder, loc, nx, ny, isCoo, swapOperands, compareOperands, v2, v3);
-    createInsert3rd(builder, loc, nx, ny, isCoo, swapOperands, compareOperands,
-                    v0, v1, v2);
+        builder, loc, xPerm, ny, swapOperands, compareOperands, v2, v3);
+    createInsert3rd(builder, loc, xPerm, ny, swapOperands, compareOperands, v0,
+                    v1, v2);
     builder.setInsertionPointAfter(ifOp);
   };
 
@@ -506,8 +497,8 @@ static void createSort5(OpBuilder &builder, Location loc, uint64_t nx,
   insert4th();
 
   // Insert the 5th element.
-  scf::IfOp ifOp = createCompareThenSwap(builder, loc, nx, ny, isCoo,
-                                         swapOperands, compareOperands, v3, v4);
+  scf::IfOp ifOp = createCompareThenSwap(builder, loc, xPerm, ny, swapOperands,
+                                         compareOperands, v3, v4);
   insert4th();
   builder.setInsertionPointAfter(ifOp);
 }
@@ -517,11 +508,10 @@ static void createSort5(OpBuilder &builder, Location loc, uint64_t nx,
 /// the number of values in range [lo, hi) is more than a threshold, we also
 /// include the middle of [lo, mi) and [mi, hi) and sort a total of five values.
 static void createChoosePivot(OpBuilder &builder, ModuleOp module,
-                              func::FuncOp func, uint64_t nx, uint64_t ny,
-                              bool isCoo, Value lo, Value hi, Value mi,
-                              ValueRange args) {
+                              func::FuncOp func, AffineMap xPerm, uint64_t ny,
+                              Value lo, Value hi, Value mi, ValueRange args) {
   SmallVector<Value> compareOperands{mi, lo};
-  uint64_t numXBuffers = isCoo ? 1 : nx;
+  constexpr uint64_t numXBuffers = 1;
   compareOperands.append(args.begin() + xStartIdx,
                          args.begin() + xStartIdx + numXBuffers);
   SmallVector<Value> swapOperands{mi, lo};
@@ -537,8 +527,8 @@ static void createChoosePivot(OpBuilder &builder, ModuleOp module,
 
   // When len < 1000, choose pivot from median of 3 values.
   builder.setInsertionPointToStart(&lenIf.getThenRegion().front());
-  createSort3(builder, loc, nx, ny, isCoo, swapOperands, compareOperands, lo,
-              mi, hi);
+  createSort3(builder, loc, xPerm, ny, swapOperands, compareOperands, lo, mi,
+              hi);
 
   // When len >= 1000, choose pivot from median of 5 values.
   builder.setInsertionPointToStart(&lenIf.getElseRegion().front());
@@ -549,8 +539,8 @@ static void createChoosePivot(OpBuilder &builder, ModuleOp module,
   Value b = builder.create<arith::AddIOp>(loc, mi, hiP1);
   // Value b is the middle between [mi, hi].
   b = builder.create<arith::ShRUIOp>(loc, b, c1);
-  createSort5(builder, loc, nx, ny, isCoo, swapOperands, compareOperands, lo, a,
-              mi, b, hi);
+  createSort5(builder, loc, xPerm, ny, swapOperands, compareOperands, lo, a, mi,
+              b, hi);
 
   builder.setInsertionPointAfter(lenIf);
 }
@@ -586,8 +576,8 @@ static void createChoosePivot(OpBuilder &builder, ModuleOp module,
 //   }
 // }
 static void createPartitionFunc(OpBuilder &builder, ModuleOp module,
-                                func::FuncOp func, uint64_t nx, uint64_t ny,
-                                bool isCoo, uint32_t nTrailingP = 0) {
+                                func::FuncOp func, AffineMap xPerm, uint64_t ny,
+                                uint32_t nTrailingP = 0) {
   // Quick sort partition doesn't use trailing parameters.
   (void)nTrailingP;
   assert(nTrailingP == 0);
@@ -606,7 +596,7 @@ static void createPartitionFunc(OpBuilder &builder, ModuleOp module,
 
   Value i = lo;
   Value j = builder.create<arith::SubIOp>(loc, hi, c1);
-  createChoosePivot(builder, module, func, nx, ny, isCoo, i, j, p, args);
+  createChoosePivot(builder, module, func, xPerm, ny, i, j, p, args);
   Value trueVal = constantI1(builder, loc, true); // The value for while (true)
   SmallVector<Value, 4> operands{i, j, p, trueVal}; // Exactly four values.
   SmallVector<Type, 4> types{i.getType(), j.getType(), p.getType(),
@@ -628,14 +618,14 @@ static void createPartitionFunc(OpBuilder &builder, ModuleOp module,
   j = after->getArgument(1);
   p = after->getArgument(2);
 
-  uint64_t numXBuffers = isCoo ? 1 : nx;
+  constexpr uint64_t numXBuffers = 1;
   auto [iresult, iCompareEq] =
       createScanLoop(builder, module, func, args.slice(xStartIdx, numXBuffers),
-                     i, p, nx, ny, isCoo, 1);
+                     i, p, xPerm, ny, 1);
   i = iresult;
   auto [jresult, jCompareEq] =
       createScanLoop(builder, module, func, args.slice(xStartIdx, numXBuffers),
-                     j, p, nx, ny, isCoo, -1);
+                     j, p, xPerm, ny, -1);
   j = jresult;
 
   // If i < j:
@@ -645,7 +635,7 @@ static void createPartitionFunc(OpBuilder &builder, ModuleOp module,
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   SmallVector<Value> swapOperands{i, j};
   swapOperands.append(args.begin() + xStartIdx, args.end());
-  createSwap(builder, loc, swapOperands, nx, ny, isCoo);
+  createSwap(builder, loc, swapOperands, xPerm, ny);
   // If the pivot is moved, update p with the new pivot.
   Value icond =
       builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, i, p);
@@ -737,8 +727,8 @@ static Value createSubTwoDividedByTwo(OpBuilder &builder, Location loc,
 // }
 //
 static void createShiftDownFunc(OpBuilder &builder, ModuleOp module,
-                                func::FuncOp func, uint64_t nx, uint64_t ny,
-                                bool isCoo, uint32_t nTrailingP) {
+                                func::FuncOp func, AffineMap xPerm, uint64_t ny,
+                                uint32_t nTrailingP) {
   // The value n is passed in as a trailing parameter.
   assert(nTrailingP == 1);
   OpBuilder::InsertionGuard insertionGuard(builder);
@@ -768,7 +758,7 @@ static void createShiftDownFunc(OpBuilder &builder, ModuleOp module,
   builder.setInsertionPointToStart(&ifNc.getThenRegion().front());
   Value c1 = constantIndex(builder, loc, 1);
   SmallVector<Value> compareOperands{start, start};
-  uint64_t numXBuffers = isCoo ? 1 : nx;
+  constexpr uint64_t numXBuffers = 1;
   compareOperands.append(args.begin() + xStartIdx,
                          args.begin() + xStartIdx + numXBuffers);
 
@@ -794,7 +784,7 @@ static void createShiftDownFunc(OpBuilder &builder, ModuleOp module,
     compareOperands[0] = lChildIdx;
     compareOperands[1] = rChildIdx;
     Value cond2 =
-        createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
+        createInlinedLessThan(builder, loc, compareOperands, xPerm, ny);
     scf::IfOp if2 =
         builder.create<scf::IfOp>(loc, ifTypes, cond2, /*else=*/true);
     builder.setInsertionPointToStart(&if2.getThenRegion().front());
@@ -825,8 +815,7 @@ static void createShiftDownFunc(OpBuilder &builder, ModuleOp module,
   childIdx = before->getArgument(2);
   compareOperands[0] = start;
   compareOperands[1] = childIdx;
-  Value cond =
-      createInlinedLessThan(builder, loc, compareOperands, nx, ny, isCoo);
+  Value cond = createInlinedLessThan(builder, loc, compareOperands, xPerm, ny);
   builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
 
   // The after-region of the WhileOp.
@@ -836,7 +825,7 @@ static void createShiftDownFunc(OpBuilder &builder, ModuleOp module,
   childIdx = after->getArgument(2);
   SmallVector<Value> swapOperands{start, childIdx};
   swapOperands.append(args.begin() + xStartIdx, args.end());
-  createSwap(builder, loc, swapOperands, nx, ny, isCoo);
+  createSwap(builder, loc, swapOperands, xPerm, ny);
   start = childIdx;
   Value cond2 =
       builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge, t, child);
@@ -869,8 +858,8 @@ static void createShiftDownFunc(OpBuilder &builder, ModuleOp module,
 //      shiftdown(lo, lo, l-1)
 // }
 static void createHeapSortFunc(OpBuilder &builder, ModuleOp module,
-                               func::FuncOp func, uint64_t nx, uint64_t ny,
-                               bool isCoo, uint32_t nTrailingP) {
+                               func::FuncOp func, AffineMap xPerm, uint64_t ny,
+                               uint32_t nTrailingP) {
   // Heap sort function doesn't have trailing parameters.
   (void)nTrailingP;
   assert(nTrailingP == 0);
@@ -897,7 +886,7 @@ static void createHeapSortFunc(OpBuilder &builder, ModuleOp module,
   shiftDownOperands.append(args.begin() + xStartIdx, args.end());
   shiftDownOperands.push_back(n);
   FlatSymbolRefAttr shiftDownFunc = getMangledSortHelperFunc(
-      builder, func, TypeRange(), kShiftDownFuncNamePrefix, nx, ny, isCoo,
+      builder, func, TypeRange(), kShiftDownFuncNamePrefix, xPerm, ny,
       shiftDownOperands, createShiftDownFunc, /*nTrailingP=*/1);
   builder.create<func::CallOp>(loc, shiftDownFunc, TypeRange(),
                                shiftDownOperands);
@@ -912,7 +901,7 @@ static void createHeapSortFunc(OpBuilder &builder, ModuleOp module,
   loplm1 = builder.create<arith::SubIOp>(loc, loplm1, c1);
   SmallVector<Value> swapOperands{lo, loplm1};
   swapOperands.append(args.begin() + xStartIdx, args.end());
-  createSwap(builder, loc, swapOperands, nx, ny, isCoo);
+  createSwap(builder, loc, swapOperands, xPerm, ny);
   shiftDownOperands[1] = lo;
   shiftDownOperands[shiftDownOperands.size() - 1] =
       builder.create<arith::SubIOp>(loc, l, c1);
@@ -928,7 +917,7 @@ static void createHeapSortFunc(OpBuilder &builder, ModuleOp module,
 /// the bigger partition to be processed by the enclosed while-loop.
 static std::pair<Value, Value>
 createQuickSort(OpBuilder &builder, ModuleOp module, func::FuncOp func,
-                ValueRange args, uint64_t nx, uint64_t ny, bool isCoo,
+                ValueRange args, AffineMap xPerm, uint64_t ny,
                 uint32_t nTrailingP) {
   MLIRContext *context = module.getContext();
   Location loc = func.getLoc();
@@ -937,8 +926,8 @@ createQuickSort(OpBuilder &builder, ModuleOp module, func::FuncOp func,
   SmallVector<Type, 2> types(2, lo.getType()); // Only two types.
 
   FlatSymbolRefAttr partitionFunc = getMangledSortHelperFunc(
-      builder, func, {IndexType::get(context)}, kPartitionFuncNamePrefix, nx,
-      ny, isCoo, args.drop_back(nTrailingP), createPartitionFunc);
+      builder, func, {IndexType::get(context)}, kPartitionFuncNamePrefix, xPerm,
+      ny, args.drop_back(nTrailingP), createPartitionFunc);
   Value p = builder
                 .create<func::CallOp>(loc, partitionFunc,
                                       TypeRange{IndexType::get(context)},
@@ -1008,8 +997,8 @@ createQuickSort(OpBuilder &builder, ModuleOp module, func::FuncOp func,
 //   }
 // }
 static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
-                                 func::FuncOp func, uint64_t nx, uint64_t ny,
-                                 bool isCoo, uint32_t nTrailingP) {
+                                 func::FuncOp func, AffineMap xPerm,
+                                 uint64_t ny, uint32_t nTrailingP) {
   // Stable sort function doesn't use trailing parameters.
   (void)nTrailingP;
   assert(nTrailingP == 0);
@@ -1034,8 +1023,8 @@ static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
   SmallVector<Value> operands{lo, i};
   operands.append(args.begin() + xStartIdx, args.end());
   FlatSymbolRefAttr searchFunc = getMangledSortHelperFunc(
-      builder, func, {IndexType::get(context)}, kBinarySearchFuncNamePrefix, nx,
-      ny, isCoo, operands, createBinarySearchFunc);
+      builder, func, {IndexType::get(context)}, kBinarySearchFuncNamePrefix,
+      xPerm, ny, operands, createBinarySearchFunc);
   Value p = builder
                 .create<func::CallOp>(loc, searchFunc, TypeRange{c1.getType()},
                                       operands)
@@ -1045,7 +1034,7 @@ static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
   operands[0] = operands[1] = i;
   SmallVector<Value> d;
   forEachIJPairInAllBuffers(
-      builder, loc, operands, nx, ny, isCoo,
+      builder, loc, operands, xPerm, ny,
       [&](uint64_t unused, Value i, Value unused2, Value buffer) {
         d.push_back(builder.create<memref::LoadOp>(loc, buffer, i));
       });
@@ -1061,7 +1050,7 @@ static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
   operands[1] = imj;
   operands[0] = builder.create<arith::SubIOp>(loc, imj, c1);
   forEachIJPairInAllBuffers(
-      builder, loc, operands, nx, ny, isCoo,
+      builder, loc, operands, xPerm, ny,
       [&](uint64_t unused, Value imjm1, Value imj, Value buffer) {
         Value t = builder.create<memref::LoadOp>(loc, buffer, imjm1);
         builder.create<memref::StoreOp>(loc, t, buffer, imj);
@@ -1071,7 +1060,7 @@ static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
   builder.setInsertionPointAfter(forOpJ);
   operands[0] = operands[1] = p;
   forEachIJPairInAllBuffers(
-      builder, loc, operands, nx, ny, isCoo,
+      builder, loc, operands, xPerm, ny,
       [&](uint64_t k, Value p, Value usused, Value buffer) {
         builder.create<memref::StoreOp>(loc, d[k], buffer, p);
       });
@@ -1123,8 +1112,8 @@ static void createSortStableFunc(OpBuilder &builder, ModuleOp module,
 // }
 //
 static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
-                                func::FuncOp func, uint64_t nx, uint64_t ny,
-                                bool isCoo, uint32_t nTrailingP) {
+                                func::FuncOp func, AffineMap xPerm, uint64_t ny,
+                                uint32_t nTrailingP) {
   assert(nTrailingP == 1 || nTrailingP == 0);
   bool isHybrid = (nTrailingP == 1);
   OpBuilder::InsertionGuard insertionGuard(builder);
@@ -1173,7 +1162,7 @@ static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
     // When len <= limit.
     builder.setInsertionPointToStart(&lenIf.getThenRegion().front());
     FlatSymbolRefAttr insertionSortFunc = getMangledSortHelperFunc(
-        builder, func, TypeRange(), kSortStableFuncNamePrefix, nx, ny, isCoo,
+        builder, func, TypeRange(), kSortStableFuncNamePrefix, xPerm, ny,
         ValueRange(args).drop_back(nTrailingP), createSortStableFunc);
     builder.create<func::CallOp>(loc, insertionSortFunc, TypeRange(),
                                  ValueRange(args).drop_back(nTrailingP));
@@ -1193,7 +1182,7 @@ static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
     // When depth exceeds limit.
     builder.setInsertionPointToStart(&depthIf.getThenRegion().front());
     FlatSymbolRefAttr heapSortFunc = getMangledSortHelperFunc(
-        builder, func, TypeRange(), kHeapSortFuncNamePrefix, nx, ny, isCoo,
+        builder, func, TypeRange(), kHeapSortFuncNamePrefix, xPerm, ny,
         ValueRange(args).drop_back(nTrailingP), createHeapSortFunc);
     builder.create<func::CallOp>(loc, heapSortFunc, TypeRange(),
                                  ValueRange(args).drop_back(nTrailingP));
@@ -1203,7 +1192,7 @@ static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
     builder.setInsertionPointToStart(&depthIf.getElseRegion().front());
     args.back() = depthLimit;
     std::tie(lo, hi) =
-        createQuickSort(builder, module, func, args, nx, ny, isCoo, nTrailingP);
+        createQuickSort(builder, module, func, args, xPerm, ny, nTrailingP);
     builder.create<scf::YieldOp>(loc, ValueRange{lo, hi});
 
     builder.setInsertionPointAfter(depthIf);
@@ -1216,7 +1205,7 @@ static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
     hi = lenIf.getResult(1);
   } else {
     std::tie(lo, hi) =
-        createQuickSort(builder, module, func, args, nx, ny, isCoo, nTrailingP);
+        createQuickSort(builder, module, func, args, xPerm, ny, nTrailingP);
   }
 
   // New [lo, hi) for the next while-loop iteration.
@@ -1229,9 +1218,8 @@ static void createQuickSortFunc(OpBuilder &builder, ModuleOp module,
 
 /// 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) {
+LogicalResult matchAndRewriteSortOp(OpTy op, ValueRange xys, AffineMap xPerm,
+                                    uint64_t ny, PatternRewriter &rewriter) {
   Location loc = op.getLoc();
   SmallVector<Value> operands{constantIndex(rewriter, loc, 0), op.getN()};
 
@@ -1285,8 +1273,8 @@ LogicalResult matchAndRewriteSortOp(OpTy op, ValueRange xys, uint64_t nx,
   }
 
   FlatSymbolRefAttr func =
-      getMangledSortHelperFunc(rewriter, insertPoint, TypeRange(), funcName, nx,
-                               ny, isCoo, operands, funcGenerator, nTrailingP);
+      getMangledSortHelperFunc(rewriter, insertPoint, TypeRange(), funcName,
+                               xPerm, ny, operands, funcGenerator, nTrailingP);
   rewriter.replaceOpWithNewOp<func::CallOp>(op, func, TypeRange(), operands);
   return success();
 }
@@ -1296,7 +1284,6 @@ LogicalResult matchAndRewriteSortOp(OpTy op, ValueRange xys, uint64_t nx,
 //===---------------------------------------------------------------------===//
 
 namespace {
-
 /// Sparse rewriting rule for the push_back operator.
 struct PushBackRewriter : OpRewritePattern<PushBackOp> {
 public:
@@ -1410,20 +1397,6 @@ private:
   bool enableBufferInitialization;
 };
 
-/// Sparse rewriting rule for the sort operator.
-struct SortRewriter : public OpRewritePattern<SortOp> {
-public:
-  using OpRewritePattern<SortOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(SortOp op,
-                                PatternRewriter &rewriter) const override {
-    SmallVector<Value> xys(op.getXs());
-    xys.append(op.getYs().begin(), op.getYs().end());
-    return matchAndRewriteSortOp(op, xys, op.getXs().size(), /*ny=*/0,
-                                 /*isCoo=*/false, rewriter);
-  }
-};
-
 /// Sparse rewriting rule for the sort_coo operator.
 struct SortCooRewriter : public OpRewritePattern<SortCooOp> {
 public:
@@ -1434,16 +1407,13 @@ public:
     SmallVector<Value> xys;
     xys.push_back(op.getXy());
     xys.append(op.getYs().begin(), op.getYs().end());
-    uint64_t nx = 1;
-    if (auto nxAttr = op.getNxAttr())
-      nx = nxAttr.getInt();
 
+    auto xPerm = op.getPermMap();
     uint64_t ny = 0;
     if (auto nyAttr = op.getNyAttr())
       ny = nyAttr.getInt();
 
-    return matchAndRewriteSortOp(op, xys, nx, ny,
-                                 /*isCoo=*/true, rewriter);
+    return matchAndRewriteSortOp(op, xys, xPerm, ny, rewriter);
   }
 };
 
@@ -1457,5 +1427,5 @@ void mlir::populateSparseBufferRewriting(RewritePatternSet &patterns,
                                          bool enableBufferInitialization) {
   patterns.add<PushBackRewriter>(patterns.getContext(),
                                  enableBufferInitialization);
-  patterns.add<SortRewriter, SortCooRewriter>(patterns.getContext());
+  patterns.add<SortCooRewriter>(patterns.getContext());
 }