Add descending flag to topk (#9355)

Author: denix56

python/test/unit/language/test_standard.py

@@ -43,7 +43,7 @@ def sort_kernel(X, stride_xm, Z, stride_zm, M: tl.constexpr, N: tl.constexpr, k:
         if k is None or x.numel < k:
             z = tl.sort(x, descending=descending)
         else:
-            z = tl.topk(x, k)
+            z = tl.topk(x, k, descending=descending)
         offs_z = offs_m[:, None] * stride_zm + offs_z_n[None, :]
         tl.store(Z + offs_z, z)
 
@@ -54,7 +54,7 @@ def sort_kernel(X, stride_xm, Z, stride_zm, M: tl.constexpr, N: tl.constexpr, k:
     if k is None or x.numel() < k:
         y = torch.sort(x, descending=descending)[0]
     else:
-        y = torch.topk(x, k=k).values
+        y = torch.topk(x, k=k, largest=descending).values
     sort_kernel[(1, )](x, x.stride(0), z, z.stride(0), M, N, k, descending, num_warps=8)
     assert (y == z).all(), (y, z)
 

python/triton/language/standard.py

@@ -466,9 +466,9 @@ def sort(x, dim: core.constexpr = None, descending: core.constexpr = core.CONSTE
 
 
 @jit
-def topk(x, k: core.constexpr, dim: core.constexpr = None):
+def topk(x, k: core.constexpr, dim: core.constexpr = None, descending: core.constexpr = True):
     """
-    Returns the k largest elements of the input tensor along the specified dimension.
+    Returns the k largest (or smallest) elements of the input tensor along the specified dimension.
 
     The elements are returned in sorted order (largest first).
 
@@ -479,6 +479,8 @@ def topk(x, k: core.constexpr, dim: core.constexpr = None):
     :param dim: The dimension along which to find the top k elements.
                 If None, uses the last dimension. Currently only the last dimension is supported.
     :type dim: int, optional
+    :param descending: If set to True, returns k largest elements. If set to False, returns k smallest elements.
+    :type descending: bool, optional
     :return: A tensor containing the k largest elements along the specified dimension.
     :rtype: Tensor
 
@@ -488,7 +490,7 @@ def topk(x, k: core.constexpr, dim: core.constexpr = None):
         x = tl.arange(0, 16)
         top4 = tl.topk(x, 4)  # Returns [15, 14, 13, 12]
     """
-    return sort_impl(x, k=k, dim=dim, descending=True)
+    return sort_impl(x, k=k, dim=dim, descending=descending)
 
 
 @jit