add one-sweep radix sort (#694)

* add one-sweep radix sort
* implement lockless decoupled-lookback scan for sort
* support inputs other and 2d tensors
* merge several data-size-dependent grid dimension to avoid exceeding max grid size

Author: iclementine

benchmark/test_special_perf.py

@@ -118,7 +118,7 @@ def unique_input_fn(shape, dtype, device):
 def test_perf_sort():
     class SortBenchmark(GenericBenchmark2DOnly):
         def set_more_shapes(self):
-            return [(1024, 1), (1024, 512)]
+            return [(1024, 1), (1024, 512), (16, 128 * 1024), (8, 256 * 1024)]
 
     def sort_input_fn(shape, dtype, device):
         inp = generate_tensor_input(shape, dtype, device)
@@ -128,7 +128,7 @@ def sort_input_fn(shape, dtype, device):
         input_fn=sort_input_fn,
         op_name="sort",
         torch_op=torch.sort,
-        dtypes=INT_DTYPES + FLOAT_DTYPES,
+        dtypes=BOOL_DTYPES + INT_DTYPES + FLOAT_DTYPES,
     )
     bench.run()
 

src/flag_gems/ops/sort.py

@@ -1,9 +1,9 @@
 import logging
-import math
 
 import torch
 import triton
 import triton.language as tl
+from triton.language.core import _unwrap_if_constexpr
 
 from ..runtime import torch_device_fn
 from ..utils import libentry
@@ -12,6 +12,313 @@
 logger = logging.getLogger(__name__)
 
 
+@tl.constexpr
+def get_int_t(num_bits: tl.constexpr, signed: tl.constexpr) -> tl.dtype:
+    num_bits = _unwrap_if_constexpr(num_bits)
+    signed = _unwrap_if_constexpr(signed)
+    return tl.core.get_int_dtype(num_bits, signed)
+
+
+@tl.constexpr
+def one_zeros(num_bits: tl.constexpr) -> int:
+    num_bits = _unwrap_if_constexpr(num_bits)
+    return 1 << (num_bits - 1)
+
+
+@tl.constexpr
+def zero_ones(num_bits: tl.constexpr) -> int:
+    num_bits = _unwrap_if_constexpr(num_bits)
+    return (1 << (num_bits - 1)) - 1
+
+
+@triton.jit
+def uint_to_uint(x, descending: tl.constexpr = False):
+    out = ~x if descending else x
+    return out
+
+
+@triton.jit
+def int_to_uint(x, descending: tl.constexpr = False):
+    num_bits: tl.constexpr = x.dtype.primitive_bitwidth
+    udtype = get_int_t(num_bits, False)
+    ux = tl.cast(x, udtype, bitcast=True)
+    if descending:
+        # 0111111....1
+        bit_mask: tl.constexpr = zero_ones(num_bits)
+        out = ux ^ bit_mask
+    else:
+        # 1000000...0
+        sign_bit_mask: tl.constexpr = one_zeros(num_bits)
+        out = ux ^ sign_bit_mask
+    return out
+
+
+@triton.jit
+def floating_to_uint(x, descending: tl.constexpr = False):
+    num_bits: tl.constexpr = x.dtype.primitive_bitwidth
+    sdtype = get_int_t(num_bits, True)
+    udtype = get_int_t(num_bits, False)
+    sx = x.to(sdtype, bitcast=True)
+    ux = x.to(udtype, bitcast=True)
+
+    sign_bit_mask: tl.constexpr = one_zeros(num_bits)
+    # mind the dtype, right_shift for signed is arithmetic right shift
+    mask = sign_bit_mask | (sx >> (num_bits - 1)).to(udtype, bitcast=True)
+    # 1000000000...0 for positive
+    # 1111111111...1 for negative
+    if descending:
+        out = ux ^ (~mask)
+    else:
+        out = ux ^ mask
+    return out.to(udtype, bitcast=True)
+
+
+@triton.jit
+def convert_to_uint_preverse_order(x: tl.tensor, descending: tl.constexpr = False):
+    if x.dtype.is_floating():
+        out = floating_to_uint(x, descending)
+    elif x.dtype.is_int_signed():
+        out = int_to_uint(x, descending)
+    elif x.dtype.is_int_unsigned():
+        out = uint_to_uint(x, descending)
+    return out
+
+
+@triton.jit
+def compute_global_hist_kernel(
+    arr_ptr,
+    out_ptr,
+    num_passes,
+    m,
+    n,
+    tiles_n_per_cta,
+    TILE_N: tl.constexpr,
+    TILE_R: tl.constexpr,
+    num_bits_per_pass: tl.constexpr,
+    descending: tl.constexpr,
+):
+    # arr_ptr: (m, n)
+    # out_ptr: (m, n_passes, r), where r = 2 ** k_bits is the number of bins
+    pid = tl.program_id(0)
+    pid_n = pid // m
+    pid_m = pid % m
+
+    r: tl.constexpr = 2**num_bits_per_pass
+    bfe_mask: tl.constexpr = (1 << num_bits_per_pass) - 1  # a.k.a. 2 ** k_bits - 1
+    CTA_TILE_N: tl.constexpr = TILE_N * tiles_n_per_cta
+    cta_n_start = CTA_TILE_N * pid_n
+    cta_n_end = tl.minimum(cta_n_start + CTA_TILE_N, n)
+
+    for p in range(0, num_passes):  # parallel
+        bit_offset = p * num_bits_per_pass
+        for r_start in range(0, r, TILE_R):  # parallel
+            bin_indices = r_start + tl.arange(0, TILE_R)
+            acc = tl.zeros((TILE_R, TILE_N), dtype=tl.int64)
+            for n_start in range(cta_n_start, cta_n_end, TILE_N):  # sequantial
+                n_offsets = n_start + tl.arange(0, TILE_N)  # (TILE_N, )
+                mask = n_offsets < cta_n_end
+                arr = tl.load(arr_ptr + pid_m * n + n_offsets, mask=mask)
+                arr = convert_to_uint_preverse_order(arr, descending)
+                key = (arr >> bit_offset) & bfe_mask  # (TILE_N, )
+                matches = tl.where(
+                    mask, (bin_indices[:, None] == key), False
+                )  # (TILE_R, TILE_N)
+                acc += matches
+            local_sum = tl.sum(acc, axis=1)
+            tl.atomic_add(
+                out_ptr + pid_m * num_passes * r + p * r + bin_indices,
+                local_sum,
+                sem="relaxed",
+            )
+
+
+@triton.jit
+def sweep(
+    arr_ptr,
+    associate_arr_ptr,  # inputs: (key & value)
+    out_ptr,
+    associate_out_ptr,  # outputs: (key & value)
+    excumsum_bins_ptr,
+    status_ptr,  # aux input and status
+    n_passes,
+    pass_id,
+    bit_offset,
+    m,
+    N,
+    OUT_N,
+    TILE_N: tl.constexpr,
+    TILE_R: tl.constexpr,
+    k_bits: tl.constexpr,
+    descending: tl.constexpr,
+):
+    # r: num_bins = 2 ** k_bits
+    # OUT_N: grid_n = cdiv(N, )
+
+    # arr_ptr: (m, N)
+    # out_ptr: (m, N)
+    # excumsum_bins_ptr: (m, n_passes, r)
+    # flag_ptr: (m, r, OUT_N)
+
+    # grid: (m, grid_r, grid_n)
+
+    # load data
+    pid = tl.program_id(0)
+    pid_m = pid % m
+    pid_n = pid // m
+    pid_r = tl.program_id(1)
+
+    # bit masks
+    aggregate_mask: tl.constexpr = 1 << 30
+    inclusive_prefix_mask: tl.constexpr = 1 << 31
+    v_mask: tl.constexpr = (1 << 30) - 1
+    bfe_mask: tl.constexpr = (1 << k_bits) - 1  # a.k.a. 2 ** k_bits - 1
+
+    # initialize flag to zero-local sum is not ready
+    r: tl.constexpr = 2**k_bits
+    cta_r_start = pid_r * TILE_R
+    cta_r_end = tl.minimum(cta_r_start + TILE_R, r)
+
+    # cumsum for a bin_index
+    n_offsets = pid_n * TILE_N + tl.arange(0, TILE_N)  # (TILE_N, )
+    mask = n_offsets < N
+    arr = tl.load(arr_ptr + pid_m * N + n_offsets, mask=mask)
+    arr_u = convert_to_uint_preverse_order(arr, descending)
+    key = (arr_u >> bit_offset) & bfe_mask  # (TILE_N, )
+
+    # since triton can only use scalar as condition, loop by bin_index
+    # status must be pre zero-initialized, or else we have to initialize it
+    for bin_index in range(cta_r_start, cta_r_end):
+        matches = tl.where(mask, key == bin_index, False)  # (TILE_N, ) bool
+        # cta level cumsum per bin
+        # CAUTION: tl.sum in triton 3.2 does not promote type
+        local_sum = tl.sum(matches.to(tl.uint32), axis=0)
+        pack0 = aggregate_mask | local_sum
+        status_offset = pid_m * (r * OUT_N) + bin_index * OUT_N + pid_n
+        tl.store(status_ptr + status_offset, pack0, cache_modifier=".cg")
+
+        # decoupled lookback
+        exclusive_prefix = tl.zeros((), dtype=tl.uint32)
+        i_lookback = pid_n - 1
+        while i_lookback >= 0:
+            flag_offset_i = pid_m * (r * OUT_N) + bin_index * OUT_N + i_lookback
+            pack1 = tl.load(status_ptr + flag_offset_i, volatile=True)  # uin32
+            while pack1 == 0:
+                pack1 = tl.load(status_ptr + flag_offset_i, volatile=True)
+            exclusive_prefix += pack1 & v_mask
+            if (pack1 & aggregate_mask) == aggregate_mask:
+                i_lookback -= 1
+            else:
+                i_lookback = -1
+        pack2 = inclusive_prefix_mask | (exclusive_prefix + local_sum)
+        tl.store(status_ptr + status_offset, pack2, cache_modifier=".cg")
+
+        local_ex_cumsum = (
+            tl.cumsum(matches.to(tl.uint32), axis=0) - matches
+        )  # (TILE_N, )
+        ex_cumsum_in_bin = (
+            exclusive_prefix + local_ex_cumsum
+        )  # global ex_cumsum_in_bin (TILE_N, )
+
+        # ex_cumsum_bins (m, n_passes, r)
+        ex_cumsum_bins = tl.load(
+            excumsum_bins_ptr + pid_m * (n_passes * r) + pass_id * r + bin_index
+        )  # scalar
+        pos = ex_cumsum_bins + ex_cumsum_in_bin  # (TILE_N, )
+
+        # scatter
+        tl.store(out_ptr + pid_m * N + pos, arr, mask=matches)
+        if associate_arr_ptr is not None:
+            associate_arr = tl.load(
+                associate_arr_ptr + pid_m * N + n_offsets, mask=mask
+            )
+            tl.store(associate_out_ptr + pid_m * N + pos, associate_arr, mask=matches)
+
+
+def radix_sort(arr, k_bits=8, descending=False):
+    n = arr.shape[-1]
+    m = arr.numel() // n
+    assert n < (1 << 30), "we have not implemented 2**30 per launch"
+    dtype = arr.dtype
+    num_bits = 1 if dtype == torch.bool else (arr.itemsize * 8)
+
+    TILE_N = 1024
+    tiles_n_per_cta = 8
+    CTA_TILE_N = tiles_n_per_cta * TILE_N
+
+    num_bins = 2**k_bits
+    n_passes = triton.cdiv(num_bits, k_bits)
+    TILE_R = 16
+
+    grid_n = triton.cdiv(n, CTA_TILE_N)
+    grid_for_global_hist = (m * grid_n, 1, 1)
+
+    with torch_device_fn.device(arr.device):
+        global_hist = torch.zeros(
+            (m, n_passes, num_bins), device=arr.device, dtype=torch.int32
+        )
+        compute_global_hist_kernel[grid_for_global_hist](
+            arr,
+            global_hist,
+            n_passes,
+            m,
+            n,
+            tiles_n_per_cta,
+            TILE_N,
+            TILE_R,
+            k_bits,
+            descending,
+        )
+        ex_cumsum_bins = torch.cumsum(global_hist, -1) - global_hist
+        ex_cumsum_bins = ex_cumsum_bins.to(torch.uint32)
+
+        # sort
+        arr_in = torch.clone(arr)
+        indices_in = (
+            torch.arange(0, n, dtype=torch.int64, device=arr_in.device)
+            .broadcast_to(arr.shape)
+            .contiguous()
+        )
+        arr_out = torch.empty_like(arr)
+        indices_out = torch.empty_like(indices_in)
+
+        TILE_R = 8
+        grid_r = triton.cdiv(num_bins, TILE_R)
+        TILE_N = 2048
+        grid_n = triton.cdiv(n, TILE_N)
+        grid_for_sweep = (m * grid_n, grid_r)
+
+        status = torch.empty(
+            (m, num_bins, grid_n), device=arr.device, dtype=torch.uint32
+        )
+
+        for i in range(0, n_passes):
+            bit_offset = i * k_bits
+            status.zero_()
+            sweep[grid_for_sweep](
+                arr_in,
+                indices_in,
+                arr_out,
+                indices_out,
+                ex_cumsum_bins,
+                status,
+                n_passes,
+                i,
+                bit_offset,
+                m,
+                n,
+                grid_n,
+                TILE_N,
+                TILE_R,
+                k_bits,
+                descending,
+            )
+            # print(f"< sorted last {bit_offset + k_bits:>2d} bits: {arr_out}")
+            arr_in, arr_out = arr_out, arr_in
+            indices_in, indices_out = indices_out, indices_in
+
+    return arr_in, indices_in
+
+
 @libentry()
 @triton.jit()
 def sort_kernel(
@@ -51,32 +358,17 @@ def sort(inp, dim=-1, descending=False):
     sort_elem_cnt = inp.shape[dim]
     if sort_elem_cnt == 1:
         return inp, torch.zeros_like(inp, dtype=torch.int64)
-    elif sort_elem_cnt > 512:  # TODO: Optimize implementation for large cases.
-        return torch.sort(inp, stable=False, dim=dim, descending=descending)
-    block_size = triton.next_power_of_2(sort_elem_cnt)
 
     if dim < 0:
         dim = dim + inp.ndim
     if dim != inp.ndim - 1:
         inp = torch.movedim(inp, dim, -1).contiguous()
     else:
         inp = inp.contiguous()
-    batch_size = math.prod(inp.shape) // sort_elem_cnt
-
-    out = torch.empty_like(inp)
-    out_index = torch.empty_like(inp, dtype=torch.int64)
-
-    with torch_device_fn.device(inp.device):
-        sort_kernel[batch_size,](
-            inp,
-            out,
-            out_index,
-            N=sort_elem_cnt,
-            BLOCK_SIZE=block_size,
-            DESCENDING=descending,
-            IS_FLOAT=inp.is_floating_point(),
-            num_warps=4,
-        )
+
+    dtype = inp.dtype
+    num_bits_per_pass = 1 if dtype == torch.bool else 4
+    out, out_index = radix_sort(inp, num_bits_per_pass, descending)
 
     if dim != inp.ndim - 1:
         out = torch.movedim(out, -1, dim)