[TorchToStablehlo] Add support for lowering torch.aten.sort to stablehlo.sort (#4633)

## Description
This PR implements the lowering support for `torch.aten.sort` in the
`convert-torch-to-stablehlo` pass, which is a prerequisite for compiling
operators like `topk` (since `topk` decomposes into `sort` + `slice`).
  See semantics in https://openxla.org/stablehlo/spec#sort.
  Fix #4337.

## Details
- Implemented `ConvertAtenSortOp` to lower `torch.aten.sort` to
`stablehlo.sort`.
- Used `stablehlo.iota` to generate initial indices along the sorting
dimension.
- Created the comparator region inside `stablehlo.sort` using
`stablehlo.compare` (GT for descending, LT for ascending).

## Testing
- Added regression test case in
`test/Conversion/TorchToStablehlo/reduction.mlir`.
 - Verified all unit tests pass with `check-torch-mlir`.
- the output stablehlo mlir is compiled and executed by iree, the result
is expected.

Signed-off-by: hsqStephenZhang <stephenzhang666666@gmail.com>

Author: hsqStephenZhang

lib/Conversion/TorchToStablehlo/Basic.cpp

@@ -2204,6 +2204,111 @@ LogicalResult ConvertAtenOp<AtenIsfiniteOp>::matchAndRewrite(
   return success();
 }
 
+// AtenSortOp
+template <>
+LogicalResult ConvertAtenOp<AtenSortOp>::matchAndRewrite(
+    AtenSortOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  Value self = adaptor.getSelf();
+  auto selfType = dyn_cast<RankedTensorType>(self.getType());
+  if (!selfType)
+    return rewriter.notifyMatchFailure(op, "expected ranked tensor for self");
+
+  // Values and indices return types
+  RankedTensorType valuesType = dyn_cast<RankedTensorType>(
+      getTypeConverter()->convertType(op.getResult(0).getType()));
+  RankedTensorType indicesType = dyn_cast<RankedTensorType>(
+      getTypeConverter()->convertType(op.getResult(1).getType()));
+  if (!valuesType || !indicesType)
+    return rewriter.notifyMatchFailure(op,
+                                       "expected ranked tensor output types");
+
+  Location loc = op.getLoc();
+  Value dimVal = op.getDim();
+  Value descendingVal = op.getDescending();
+
+  int64_t dim;
+  if (auto constantOp = dimVal.getDefiningOp<ConstantIntOp>()) {
+    dim = constantOp.getValueAttr().getInt();
+  } else {
+    return rewriter.notifyMatchFailure(op, "non-constant dim parameter");
+  }
+  int64_t rank = selfType.getRank();
+  if (dim < -rank || dim >= rank) {
+    return rewriter.notifyMatchFailure(op, "dimension out of range");
+  }
+  if (dim < 0) {
+    dim += rank;
+  }
+
+  bool descending;
+  if (auto constantOp = descendingVal.getDefiningOp<ConstantBoolOp>()) {
+    descending = constantOp.getValue();
+  } else {
+    return rewriter.notifyMatchFailure(op, "non-constant descending parameter");
+  }
+
+  // 1. Generate indices tensor using stablehlo.iota
+  Value indices = stablehlo::IotaOp::create(rewriter, loc, indicesType,
+                                            rewriter.getI64IntegerAttr(dim));
+
+  // 2. Create stablehlo.sort op
+  auto sortOp = stablehlo::SortOp::create(
+      rewriter, loc, TypeRange{valuesType, indicesType},
+      ValueRange{self, indices}, rewriter.getI64IntegerAttr(dim),
+      rewriter.getBoolAttr(false));
+
+  // 3. Build comparator block
+  Block &block = sortOp.getComparator().emplaceBlock();
+
+  auto blockValArgumentType =
+      RankedTensorType::get({}, valuesType.getElementType());
+  auto blockIdxArgumentType =
+      RankedTensorType::get({}, indicesType.getElementType());
+
+  block.addArgument(blockValArgumentType, loc);
+  block.addArgument(blockValArgumentType, loc);
+  block.addArgument(blockIdxArgumentType, loc);
+  block.addArgument(blockIdxArgumentType, loc);
+
+  auto *firstValArg = block.args_begin();
+  auto *secondValArg = std::next(firstValArg);
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(&block);
+
+  auto compareDirectionAttr = stablehlo::ComparisonDirectionAttr::get(
+      rewriter.getContext(), descending ? stablehlo::ComparisonDirection::GT
+                                        : stablehlo::ComparisonDirection::LT);
+
+  stablehlo::ComparisonTypeAttr compareTypeAttr;
+  Type elemTy = valuesType.getElementType();
+
+  if (isa<mlir::FloatType>(valuesType.getElementType())) {
+    compareTypeAttr = stablehlo::ComparisonTypeAttr::get(
+        rewriter.getContext(), stablehlo::ComparisonType::FLOAT);
+  } else if (isa<mlir::IntegerType>(valuesType.getElementType())) {
+    if (elemTy.isInteger(1)) {
+      compareTypeAttr = stablehlo::ComparisonTypeAttr::get(
+          rewriter.getContext(), stablehlo::ComparisonType::UNSIGNED);
+    } else {
+      compareTypeAttr = stablehlo::ComparisonTypeAttr::get(
+          rewriter.getContext(), stablehlo::ComparisonType::SIGNED);
+    }
+  }
+
+  auto compareResultType = RankedTensorType::get({}, rewriter.getI1Type());
+
+  Value compareResult = stablehlo::CompareOp::create(
+      rewriter, loc, compareResultType, *firstValArg, *secondValArg,
+      compareDirectionAttr, compareTypeAttr);
+
+  stablehlo::ReturnOp::create(rewriter, loc, compareResult);
+
+  rewriter.replaceOp(op, sortOp.getResults());
+  return success();
+}
+
 void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     ConversionTarget &target, const TorchToStablehloOptions &options) {
@@ -2380,6 +2485,8 @@ void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
 
   INSERT_ATENOP_PATTERN(AtenTrilOp);
   INSERT_ATENOP_PATTERN(AtenIsfiniteOp);
+  INSERT_ATENOP_PATTERN(AtenSortOp);
+
 #undef INSERT_ATENOP_PATTERN
 
 #define INSERT_BINARY_BROADCAST_PATTERN(AtenOp, StablehloOp)                   \

test/Conversion/TorchToStablehlo/basic.mlir

@@ -339,3 +339,24 @@ func.func @torch.aten.tril(%arg0: !torch.vtensor<[2,3,5],f32>, %arg1: !torch.int
   %0 = torch.aten.tril %arg0, %arg1:!torch.vtensor<[2,3,5],f32>, !torch.int -> !torch.vtensor<[2,3,5],f32>
   return %0 : !torch.vtensor<[2,3,5],f32>
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @torch.aten.sort(
+// CHECK-SAME:                  %[[ARG_0:.*]]: !torch.vtensor<[2,3],f32>) -> (!torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>) {
+// CHECK:           %[[VAL_0:.*]] = torch_c.to_builtin_tensor %[[ARG0]] : !torch.vtensor<[2,3],f32> -> tensor<2x3xf32>
+// CHECK:           %[[VAL_1:.*]] = stablehlo.iota dim = 1 : tensor<2x3xi64>
+// CHECK:           %[[VAL_2:.*]]:2 = "stablehlo.sort"(%[[VAL_0]], %[[VAL_1]]) <{dimension = 1 : i64, is_stable = false}> ({
+// CHECK:           ^bb0(%[[ARG1:.*]]: tensor<f32>, %[[ARG2:.*]]: tensor<f32>, %[[ARG3:.*]]: tensor<i64>, %[[ARG4:.*]]: tensor<i64>):
+// CHECK:             %[[VAL_3:.*]] = stablehlo.compare GT, %[[ARG1]], %[[ARG2]], FLOAT : (tensor<f32>, tensor<f32>) -> tensor<i1>
+// CHECK:             stablehlo.return %[[VAL_3]] : tensor<i1>
+// CHECK:           }) : (tensor<2x3xf32>, tensor<2x3xi64>) -> (tensor<2x3xf32>, tensor<2x3xi64>)
+// CHECK:           %[[VAL_4:.*]] = torch_c.from_builtin_tensor %[[VAL_2]]#1 : tensor<2x3xi64> -> !torch.vtensor<[2,3],si64>
+// CHECK:           %[[VAL_5:.*]] = torch_c.from_builtin_tensor %[[VAL_2]]#0 : tensor<2x3xf32> -> !torch.vtensor<[2,3],f32>
+// CHECK:           return %[[VAL_5]], %[[VAL_4]] : !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>
+func.func @torch.aten.sort(%arg0: !torch.vtensor<[2,3],f32>) -> (!torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>) {
+  %int-1 = torch.constant.int -1
+  %true = torch.constant.bool true
+  %values, %indices = torch.aten.sort %arg0, %int-1, %true : !torch.vtensor<[2,3],f32>, !torch.int, !torch.bool -> !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>
+  return %values, %indices : !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],si64>
+}