feat: add argsort operator implementation, tests and benchmark

Author: factnn

benchmark/test_special_perf.py

@@ -344,6 +344,25 @@ def sort_input_fn(shape, dtype, device):
     bench.run()
 
 
+@pytest.mark.argsort
+def test_perf_argsort():
+    class ArgsortBenchmark(GenericBenchmark2DOnly):
+        def set_more_shapes(self):
+            return [(1024, 1), (1024, 512), (16, 128 * 1024), (8, 256 * 1024)]
+
+    def argsort_input_fn(shape, dtype, device):
+        inp = generate_tensor_input(shape, dtype, device)
+        yield inp, {"dim": -1, "descending": False},
+
+    bench = ArgsortBenchmark(
+        input_fn=argsort_input_fn,
+        op_name="argsort",
+        torch_op=torch.argsort,
+        dtypes=INT_DTYPES + FLOAT_DTYPES,
+    )
+    bench.run()
+
+
 @pytest.mark.multinomial
 def test_multinomial_with_replacement():
     def multinomial_input_fn(shape, dtype, device):

conf/operators.yaml

@@ -617,6 +617,16 @@ ops:
       - LinearAlg
     stages:
       - stable: '2.2'
+  - name: argsort
+    description: Returns the indices that sort a tensor along a given dimension in ascending order by value.
+    for:
+      - argsort
+    labels:
+      - aten
+    kind:
+      - Data
+    stages:
+      - stable: '3.0'
   - name: asinh_
     description: Computes the inverse hyperbolic sine for each element of a tensor in-place.
     for:

src/flag_gems/__init__.py

@@ -90,6 +90,7 @@ def torch_ge(v):
     ("arcsinh_", arcsinh_),
     ("argmax", argmax),
     ("argmin", argmin),
+    ("argsort", argsort),
     ("asinh_", asinh_),
     ("atan", atan),
     ("atan_", atan_),

src/flag_gems/ops/__init__.py

@@ -25,8 +25,8 @@
 from flag_gems.ops.arctanh_ import arctanh_
 from flag_gems.ops.argmax import argmax
 from flag_gems.ops.argmin import argmin
+from flag_gems.ops.argsort import argsort
 from flag_gems.ops.asinh_ import asinh_
-from flag_gems.ops.assert_async import _assert_async
 from flag_gems.ops.atan import atan, atan_
 from flag_gems.ops.atan2 import atan2, atan2_out
 from flag_gems.ops.attention import (
@@ -371,6 +371,7 @@
     "arcsinh_",
     "argmax",
     "argmin",
+    "argsort",
     "asinh_",
     "atan",
     "atan_",

src/flag_gems/ops/argsort.py

@@ -0,0 +1,70 @@
+import logging
+
+import torch
+import triton
+
+from flag_gems.ops.sort import sort_kernel, sort_stable
+
+logger = logging.getLogger(__name__)
+
+
+def argsort(inp, dim=-1, descending=False):
+    """Returns the indices that sort a tensor along a given dimension.
+
+    This is equivalent to calling torch.sort and returning only the indices.
+
+    Performance Notes:
+    - For small N (≤4096): Uses bitonic sort (single kernel launch)
+      Speedup: ~2.5-3x for N=64-256
+    - For large N (>4096): Uses radix sort (multiple passes)
+      Current performance: 0.1-0.3x slower than PyTorch for N≥1024
+      TODO: Optimize radix sort for large N scenarios
+      Benchmark results (N=1024-262144):
+        - [1024, 1024]: 0.4-0.7x speedup
+        - [4096, 4096]: 0.1-0.2x speedup
+        - [1024, 65536]: 0.2-0.3x speedup
+    """
+    logger.debug("GEMS ARGSORT")
+
+    # For small N, use bitonic sort (single kernel launch)
+    # For large N, use radix sort (multiple passes)
+    if dim < 0:
+        dim = dim + inp.ndim
+
+    sort_elem_cnt = inp.shape[dim]
+
+    # Use bitonic sort for small sizes (faster single-kernel approach)
+    if sort_elem_cnt <= 4096:
+        if dim != inp.ndim - 1:
+            inp = torch.movedim(inp, dim, -1).contiguous()
+        else:
+            inp = inp.contiguous()
+
+        N = inp.shape[-1]
+        M = inp.numel() // N
+
+        out = torch.empty_like(inp)
+        out_index = torch.empty(inp.shape, dtype=torch.int64, device=inp.device)
+
+        BLOCK_SIZE = triton.next_power_of_2(N)
+        IS_FLOAT = inp.dtype.is_floating_point
+
+        grid = lambda meta: (M,)
+        sort_kernel[grid](
+            inp,
+            out,
+            out_index,
+            N,
+            BLOCK_SIZE=BLOCK_SIZE,
+            DESCENDING=descending,
+            IS_FLOAT=IS_FLOAT,
+        )
+
+        if dim != inp.ndim - 1:
+            out_index = torch.movedim(out_index, -1, dim)
+
+        return out_index
+    else:
+        # Use radix sort for large sizes
+        _, indices = sort_stable(inp, stable=True, dim=dim, descending=descending)
+        return indices

tests/test_argsort.py

@@ -0,0 +1,41 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+from .conftest import QUICK_MODE
+
+ARGSORT_HIDDENSIZE = (
+    [1, 256, 2048, 9333]
+    if QUICK_MODE
+    else [1, 256, 2048, 9333, 65536, 32768, 128 * 1024, 256 * 1024]
+)
+
+
+@pytest.mark.argsort
+@pytest.mark.parametrize("batch_size", [4, 8])
+@pytest.mark.parametrize("hiddensize", ARGSORT_HIDDENSIZE)
+@pytest.mark.parametrize("descending", [True, False])
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.INT_DTYPES)
+@pytest.mark.parametrize("dim", [0, -1])
+def test_accuracy_argsort(batch_size, hiddensize, descending, dtype, dim):
+    if dtype in utils.BOOL_TYPES:
+        y = torch.randint(
+            0, 2, (batch_size, hiddensize), dtype=dtype, device=flag_gems.device
+        )
+    elif dtype in utils.ALL_INT_DTYPES:
+        min_v, max_v = torch.iinfo(dtype).min, torch.iinfo(dtype).max
+        y = torch.randint(
+            min_v, max_v, (batch_size, hiddensize), dtype=dtype, device="cpu"
+        ).to(flag_gems.device)
+    else:
+        y = torch.randn((batch_size, hiddensize), dtype=dtype, device=flag_gems.device)
+
+    ref_y = utils.to_reference(y)
+    ref_index = torch.argsort(ref_y, dim=dim, stable=True, descending=descending)
+
+    with flag_gems.use_gems():
+        res_index = torch.argsort(y, dim=dim, stable=True, descending=descending)
+
+    utils.gems_assert_equal(res_index, ref_index)