[Torch] emit aten.argsort and decompose it to aten.sort (#4027)

Author: qingyunqu

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -15176,6 +15176,31 @@ def Torch_AtenSortOp : Torch_Op<"aten.sort", [
   let hasFolder = 1;
 }
 
+def Torch_AtenArgsortOp : Torch_Op<"aten.argsort", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::argsort : (Tensor, int, bool) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$dim,
+    Torch_BoolType:$descending
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenArgsortOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 3, 1);
+    }
+    void AtenArgsortOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 3, 1);
+    }
+  }];
+}
+
 def Torch_AtenSplitTensorOp : Torch_Op<"aten.split.Tensor", [
     AllowsTypeRefinement,
     ReadOnly

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -10295,6 +10295,13 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %1 = torch.prim.TupleConstruct %0#1, %int4 : !torch.int, !torch.int -> !torch.tuple<int, int>\n"
 "    return %1 : !torch.tuple<int, int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.argsort\"(%arg0: !torch.list<int>, %arg1: !torch.int, %arg2: !torch.bool) -> !torch.list<int> {\n"
+"    return %arg0 : !torch.list<int>\n"
+"  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.argsort\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int, %arg2: !torch.bool) -> !torch.int {\n"
+"    %int4 = torch.constant.int 4\n"
+"    return %int4 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.narrow\"(%arg0: !torch.list<int>, %arg1: !torch.int, %arg2: !torch.int, %arg3: !torch.int) -> !torch.list<int> {\n"
 "    %int1 = torch.constant.int 1\n"
 "    %0 = torch.aten.add.int %arg2, %arg3 : !torch.int, !torch.int -> !torch.int\n"

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -10214,6 +10214,31 @@ class DecomposeAtenTopkOp : public OpRewritePattern<AtenTopkOp> {
 };
 } // namespace
 
+namespace {
+// decompose aten.argsort to aten.sort
+class DecomposeAtenArgsortOp : public OpRewritePattern<AtenArgsortOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenArgsortOp op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    auto context = op.getContext();
+
+    Value self = op.getSelf();
+    Value dim = op.getDim();
+    Value descending = op.getDescending();
+    auto selfType = cast<BaseTensorType>(self.getType());
+    auto sortIndicesType = selfType.getWithSizesAndDtype(
+        selfType.getOptionalSizes(),
+        IntegerType::get(context, 64, IntegerType::Signed));
+    auto sortOpResult = rewriter.create<AtenSortOp>(
+        loc, self.getType(), sortIndicesType, self, dim, descending);
+    rewriter.replaceOp(op, sortOpResult->getResult(1));
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 
 /// Creates coefficients based on DFT definition, see
@@ -11781,6 +11806,7 @@ class DecomposeComplexOpsPass
         patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenVarMeanDimOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenTopkOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenArgsortOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenFftRfftOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenHannWindowPeriodicOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenScalarTensor>(patterns);

lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

@@ -557,6 +557,7 @@ static void markDecomposedOpsAsIllegal(MLIRContext *context,
   target.addIllegalOp<AtenCrossEntropyLossOp>();
   target.addIllegalOp<AtenVarMeanDimOp>();
   target.addIllegalOp<AtenTopkOp>();
+  target.addIllegalOp<AtenArgsortOp>();
   target.addIllegalOp<AtenHannWindowPeriodicOp>();
   target.addIllegalOp<AtenScalarTensorOp>();
   target.addIllegalOp<AtenScatterValueOp>();

projects/pt1/e2e_testing/xfail_sets.py

@@ -531,6 +531,8 @@
 }
 
 FX_IMPORTER_STABLEHLO_XFAIL_SET = {
+    "ArgsortTensor_basic",
+    "ArgsortTensorInteger_basic",
     "AddFloatIntModule_basic",
     "AtenKthvalueDynamicDimsModule_basic",
     "AtenKthvalueFloat64DynamicDimsModule_basic",
@@ -3348,6 +3350,8 @@
 }
 
 FX_IMPORTER_TOSA_XFAIL_SET = {
+    "ArgsortTensor_basic",
+    "ArgsortTensorInteger_basic",
     "AtenSymConstrainRangeForSize_basic",
     "AtenSymConstrainRange_basic",
     "Aten_AssertScalar_basic",

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

@@ -1988,6 +1988,12 @@ def aten〇sort〡dtype(self_rank_dtype: Tuple[int, int], dim: int = -1, descend
     _, input_dtype = self_rank_dtype
     return input_dtype, torch.long
 
+def aten〇argsort〡shape(self: List[int], dim: int = -1, descending: bool = False) -> List[int]:
+    return self
+
+def aten〇argsort〡dtype(self_rank_dtype: Tuple[int, int], dim: int = -1, descending: bool = False) -> int:
+    return torch.long
+
 def aten〇narrow〡shape(self: List[int], dim: int, start: int, length: int) -> List[int]:
     return upstream_shape_functions.slice(self, dim, start, start + length, 1)
 

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -1081,6 +1081,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::any.bool : (bool[]) -> (bool)", has_folder=True)
     emit("aten::sort.int : (int[], bool) -> ()", has_canonicalizer=True)
     emit("aten::sort : (Tensor, int, bool) -> (Tensor, Tensor)", has_folder=True)
+    emit("aten::argsort : (Tensor, int, bool) -> (Tensor)")
     emit("aten::split.Tensor : (Tensor, int, int) -> (Tensor[])")
     emit("aten::split_with_sizes : (Tensor, int[], int) -> (Tensor[])")
     emit(

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py

@@ -5357,6 +5357,36 @@ def SortTensorNegativeDimension_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4, 5))
 
 
+class ArgsortTensor(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([-1, -1, -1], torch.float32, True)])
+    def forward(self, input):
+        return torch.argsort(input)
+
+
+@register_test_case(module_factory=lambda: ArgsortTensor())
+def ArgsortTensor_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4, 5))
+
+
+class ArgsortTensorInteger(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([-1, -1], torch.int64, True)])
+    def forward(self, input):
+        return torch.argsort(input)
+
+
+@register_test_case(module_factory=lambda: ArgsortTensorInteger())
+def ArgsortTensorInteger_basic(module, tu: TestUtils):
+    module.forward(tu.randint(2, 3))
+
+
 # ==============================================================================