Add native stable CUDA sorting for bounded UTF-8 strings

sdm/tensor/string.py

@@ -16,7 +16,9 @@
 
 from sdm._warnings import warn_once
 from sdm.tensor import VarLenTensor
+from sdm.tensor._string_keys import pack
 from sdm.tensor.io import arrow_as_tensor, combine_arrow_chunks
+from sdm.tensor.var_len import _select_entries
 
 if TYPE_CHECKING:
     import cudf
@@ -523,6 +525,12 @@ def _sort(
     if inp.dim() != 1:
         raise NotImplementedError("'sort' only supports one-dimensional input")
 
+    # Bound radix passes to 12 words on CUDA and eight on CPU.
+    cuda = inp.is_cuda
+    native = cuda or (inp.is_cpu and torch.compiler.is_compiling())
+    if native and inp._max_length <= (84 if cuda else 56):
+        return _sort_native(inp, descending=descending)
+
     backend: Literal["arrow", "cudf"] = "arrow"
     if inp.is_cuda:
         if importlib.util.find_spec("cudf") is not None:
@@ -553,3 +561,27 @@ def _sort(
             perm = torch.from_dlpack(perm_ser.astype("int64").to_cupy())
 
     return cast(StringTensor, inp[perm]), perm
+
+
+def _sort_native(
+    inp: StringTensor,
+    *,
+    descending: bool,
+) -> tuple[StringTensor, Tensor]:
+    words, valid = pack(inp)
+    if inp._valid is not None:
+        # Nulls sort last in both directions, including stable duplicate nulls.
+        sentinel = -1 if descending else torch.iinfo(torch.int64).max
+        words = torch.where(valid[:, None], words, sentinel)
+
+    # Least-significant word first preserves lexicographic ordering and ties.
+    perm = words[:, -1].argsort(descending=descending, stable=True)
+    for column in range(words.size(1) - 2, -1, -1):
+        order = words[perm, column].argsort(
+            descending=descending,
+            stable=True,
+        )
+        perm = perm[order]
+
+    out = _select_entries(inp, perm)
+    return cast(StringTensor, out), perm

test/tensor/test_string.py

@@ -8,6 +8,7 @@
 import torch
 
 from sdm import StringTensor
+from sdm.tensor.string import _sort_native
 from sdm.testing import onlyCUDA, withCUDA
 
 
@@ -382,3 +383,146 @@ def test_registers_as_a_pytorch_buffer() -> None:
     module.register_buffer("value", tensor)
 
     assert tuple(module.buffers()) == (tensor,)
+
+
+@pytest.mark.parametrize("native", [False, True])
+@pytest.mark.parametrize("descending", [False, True])
+@pytest.mark.parametrize(
+    "values",
+    [
+        [],
+        [None, None],
+        ["", "", None, "\0", "\0a", "a", "a", None],
+        ["é", "e", "中", "😀", "é", "\x7f", "\x80"],
+        ["a" * 7 + "b", "a" * 8, "a" * 7, "a" * 14, None],
+        ["a" * 55 + "b", "a" * 56, "a" * 55, "a" * 56, None],
+        ["a" * 57 + "z", "a" * 58, "a" * 57, "a" * 58, None, None],
+        ["é" * 37 + "b", "é" * 37 + "a", "é" * 37, None],
+        ["a" * 83 + "b", "a" * 84, "a" * 83, "a" * 84, None],
+        ["é" * 42, "é" * 41 + "e", "é" * 41, None],
+        ["a" * 84 + "b", "a" * 85, "a" * 84, "a" * 85, None],
+    ],
+)
+@withCUDA
+def test_sort_stable(
+    values: list[str | None],
+    descending: bool,
+    native: bool,
+    device: torch.device,
+) -> None:
+    # The sort indices are relative to the view, including a storage offset.
+    storage: list[str | None] = ["unused"]
+    for value in values:
+        storage.extend([value, "unused"])
+    tensor = StringTensor.from_list(storage, device=device)[1::2]
+    indices = sorted(
+        (i for i, value in enumerate(values) if value is not None),
+        key=lambda i: cast(str, values[i]).encode("utf-8"),
+        reverse=descending,
+    )
+    indices.extend(i for i, value in enumerate(values) if value is None)
+
+    if native:
+        out, perm = _sort_native(
+            cast(StringTensor, tensor), descending=descending
+        )
+    else:
+        out, perm = torch.sort(tensor, descending=descending, stable=True)
+    assert out.tolist() == [values[i] for i in indices]
+    assert perm.tolist() == indices
+    assert out.device == device
+
+
+@pytest.mark.parametrize("descending", [False, True])
+@withCUDA
+def test_sort_native_fullgraph(
+    compilation: None,
+    descending: bool,
+    device: torch.device,
+) -> None:
+    tensor = StringTensor.from_list(
+        [
+            "unused",
+            "é",
+            "unused",
+            "aaaaaaaa",
+            "unused",
+            None,
+            "unused",
+            "aaaaaaaa",
+            "unused",
+            "",
+            "unused",
+            "中",
+        ],
+        device=device,
+    )[1::2]
+
+    def sort(inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.sort(inp, descending=descending, stable=True)
+
+    out, perm = torch.compile(sort, fullgraph=True)(tensor)
+    expected = [5, 0, 1, 3, 4, 2] if descending else [4, 1, 3, 0, 5, 2]
+    assert perm.tolist() == expected
+    values = tensor.tolist()
+    assert out.tolist() == [values[i] for i in expected]
+
+
+@pytest.mark.parametrize("width", [75, 84])
+@pytest.mark.parametrize("descending", [False, True])
+@onlyCUDA
+def test_sort_wide_fullgraph(
+    compilation: None,
+    width: int,
+    descending: bool,
+) -> None:
+    values = ["a" * (width - 1) + "b", "a" * width, "a" * width, None, ""]
+    tensor = StringTensor.from_list(values, device="cuda")
+
+    def sort(inp: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.sort(inp, descending=descending, stable=True)
+
+    out, perm = torch.compile(sort, fullgraph=True)(tensor)
+    expected = [0, 1, 2, 4, 3] if descending else [4, 1, 2, 0, 3]
+    assert perm.tolist() == expected
+    assert out.tolist() == [values[i] for i in expected]
+
+
+@withCUDA
+def test_sort_null_payload(device: torch.device) -> None:
+    # Null slots may contain arbitrary bytes and must retain their input order.
+    array = pa.Array.from_buffers(
+        pa.string(),
+        4,
+        [
+            pa.py_buffer(bytes([0b00000101])),
+            pa.py_buffer(
+                torch.tensor([0, 1, 4, 5, 8], dtype=torch.int32).numpy()
+            ),
+            pa.py_buffer(bytearray(b"bzzzazzz")),
+        ],
+    )
+    tensor = StringTensor.from_arrow(array, device=device)
+    out, perm = tensor.sort(stable=True)
+    assert out.tolist() == ["a", "b", None, None]
+    assert perm.tolist() == [2, 0, 1, 3]
+
+
+@withCUDA
+def test_sort_repeated_view(device: torch.device) -> None:
+    tensor = StringTensor.from_list(["abc"], device=device).expand(5)
+    out, perm = tensor.sort(stable=True)
+    assert out.tolist() == ["abc"] * 5
+    assert perm.tolist() == list(range(5))
+
+
+@pytest.mark.parametrize("values", [[], [None, None], ["", ""]])
+@withCUDA
+def test_sort_empty_payload(
+    values: list[str | None],
+    device: torch.device,
+) -> None:
+    tensor = StringTensor.from_list(values, device=device)
+    out, perm = tensor.sort(stable=True)
+    assert out.tolist() == values
+    assert perm.tolist() == list(range(len(values)))