diff --git a/sdm/tensor/string.py b/sdm/tensor/string.py index 17b8da35f..9492ef7a8 100644 --- a/sdm/tensor/string.py +++ b/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 diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py index d70c6a388..31bcbea88 100644 --- a/test/tensor/test_string.py +++ b/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)))