diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py index a31438a97..76d0003a8 100644 --- a/sdm/tensor/var_len.py +++ b/sdm/tensor/var_len.py @@ -1559,3 +1559,190 @@ def _materialize( stride=stride, storage_offset=0, ) + + +def _permute_entries(inp: VarLenTensor, perm: Tensor) -> VarLenTensor: + """Reorder one-dimensional entries by a permutation. + + A permutation of a view with nonzero stride selects at most the bytes of + its backing storage, so the output keeps that static capacity regardless + of individual entry lengths. Bytes after the final offset are zero. + + Eager CUDA copies in tiles, which needs one host read of entry counts; + compiled code searches each element's entry, which Inductor fuses into + one kernel. A single tiled path for both was 2.7x slower eager on 1M + URLs and used up to 2.9x the memory, so the two paths stay separate. + """ + if inp.stride(0) == 0: + return cast(VarLenTensor, inp[perm]) + + rows = perm * inp.stride(0) + inp._logical_storage_offset + valid = inp._valid[rows] if inp._valid is not None else None + capacity = inp._data.numel() + dtype = torch.int32 + if capacity + 2 * _TILE + _SCRATCH > torch.iinfo(torch.int32).max: + dtype = torch.int64 + shift = inp._offset[rows].to(dtype) + offset = inp._offset[1:][rows].to(dtype) - shift + del rows + offset = torch.cat((offset.new_zeros(1), offset.cumsum(0, dtype=dtype))) + shift -= offset[:-1] # Source minus output position of each entry. + + if inp.is_cuda and not torch.compiler.is_compiling(): + data = _gather_tiles(inp._data, shift, offset) + else: + data = _gather_elements(inp._data, shift, offset) + del shift + + return inp.__class__( + data=data, + offset=offset.to(inp._offset.dtype), + valid=valid, + size=perm.shape, + ) + + +_TILE = 32 +_SCRATCH = 1024 + + +def _gather_elements(data: Tensor, shift: Tensor, offset: Tensor) -> Tensor: + """Copy entries element by element with static shapes. + + Ends of non-empty entries strictly increase, so at most ``2 * _TILE`` of + them fall in any ``2 * _TILE``-element window. A binary search of + ``(2 * _TILE).bit_length()`` steps within the window's range finds the + entry of each output element; compiled, this is a single fused kernel. + """ + capacity, n, width = data.numel(), shift.numel(), 2 * _TILE + end = offset[1:] + nonempty = end > offset[:-1] + # Non-empty entries in output order; others land in the spare slot ``n``. + rank = (nonempty.cumsum(0, dtype=end.dtype) - 1).where(nonempty, n) + last = end.new_empty(n + 1).scatter_(0, rank, end) + shift = shift.new_empty(n + 1).scatter_(0, rank, shift) + # Slots after the last non-empty entry pad up to the capacity. + unused = torch.arange(n + 1, device=end.device) >= nonempty.sum() + last.masked_fill_(unused, capacity + width) + shift.masked_fill_(unused, -capacity) + + pos = torch.arange(capacity, dtype=end.dtype, device=end.device) + window = torch.arange( + 0, + capacity + width, + width, + dtype=end.dtype, + device=end.device, + ) + bound = torch.searchsorted( + last, + window, + right=True, + out_int32=end.dtype == torch.int32, + ) + low, high = bound[pos // width], bound[pos // width + 1] + for _ in range(width.bit_length()): + mid = (low + high) >> 1 + right = last[mid] <= pos + low, high = (mid + 1).where(right, low), high.where(right, mid) + source = pos + shift[low] + return data[source.clamp_min(0)].masked_fill(source < 0, 0) + + +def _gather_tiles(data: Tensor, shift: Tensor, offset: Tensor) -> Tensor: + """Copy entries in ``_TILE``-element windows; synchronizes once. + + Each output tile is copied from the entry owning its first element, which + is exact up to the first entry boundary inside the tile. Entries of at + least ``_TILE`` elements then rewrite their first ``_TILE`` elements, and + shorter entries are copied element by element. + """ + capacity, n = data.numel(), shift.numel() + if capacity == 0 or n == 0: + return data.new_zeros(capacity) + head, end = offset[:-1], offset[1:] + i32 = offset.dtype == torch.int32 + + # Per-entry bookkeeping before the payload is allocated. + length = end - head + short = (length > 0) & (length < _TILE) + total, n_short, short_size = torch.stack( + (end[-1:].sum(), short.sum(), (length * short).sum()) + ).tolist() + width = _TILE if total >= _TILE * n and capacity >= _TILE else 1 + if width > 1: + short = torch.nonzero_static(short, size=n_short).squeeze(1) + short = short.to(offset.dtype) + # Tiles are clamped into the payload only for the long entry whose + # source ends within the last ``width`` elements. + tail = (length >= width) & (head + shift + length > capacity - width) + tail = torch.nonzero_static(tail, size=1, fill_value=n) + del length + + out = data.new_empty(-(-capacity // width) * width + width + _SCRATCH) + src = data.unfold(0, width, 1) # [capacity - width + 1, width] + scratch = out.numel() - width - _SCRATCH + + # [tiles, width] tiles from the owner of their first element. + tiles, block = -(-total // width), 1 << (20 if width == 1 else 18) + for lo in range(0, tiles, block): + hi = min(lo + block, tiles) + row = torch.arange( + lo * width, + hi * width, + width, + dtype=offset.dtype, + device=data.device, + ) + owner = torch.searchsorted(end, row, right=True, out_int32=i32) + row += shift.index_select(0, owner) + del owner + torch.index_select( + src, + 0, + row.clamp_max_(capacity - width), + out=out[lo * width : hi * width].view(-1, width), + ) + if width == 1: + out[total:capacity].zero_() + return out[:capacity] + + # [n, width] first elements of long entries; others write scratch rows. + dst = out.unfold(0, width, 1) + step = (3 << 20) // width + for lo in range(0, n, step): + hi = min(lo + step, n) + ok = end[lo:hi] - head[lo:hi] >= width + pad = torch.arange(lo, hi, dtype=offset.dtype, device=data.device) + pad = pad.remainder_(_SCRATCH).add_(scratch) + dst.index_copy_( + 0, + head[lo:hi].where(ok, pad).long(), + src.index_select(0, (head[lo:hi] + shift[lo:hi]).where(ok, 0)), + ) + + # Rewrite the last elements of the entry whose tiles were clamped. + ok = tail < n + tail = tail.clamp_max(n - 1) + dst.index_copy_( + 0, + (end[tail] - width).where(ok, scratch).long().view(-1), + src.index_select( + 0, (end[tail] + shift[tail] - width).where(ok, 0).view(-1) + ), + ) + + # [short_size] elements of short entries, packed back to back. + first = (end - head)[short].cumsum(0, dtype=offset.dtype) + move = head[short] + move[1:] -= first[:-1] # Packed to output position. + shift = shift[short] + for lo in range(0, short_size, 1 << 17): + hi = min(lo + (1 << 17), short_size) + pos = torch.arange(lo, hi, dtype=offset.dtype, device=data.device) + k = torch.searchsorted(first, pos, right=True, out_int32=i32) + pos += move.index_select(0, k) + out[pos] = data[pos + shift.index_select(0, k)] + + out[total:capacity].zero_() + return out[:capacity] diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py index 04c6f641b..87ed983c9 100644 --- a/test/tensor/test_var_len.py +++ b/test/tensor/test_var_len.py @@ -9,7 +9,7 @@ from torch import Tensor from sdm import StringTensor, VarLenTensor -from sdm.tensor.var_len import _compact +from sdm.tensor.var_len import _compact, _permute_entries from sdm.testing import onlyCUDA, withCUDA @@ -1059,3 +1059,119 @@ def test_reshape_inference_mode() -> None: assert not torch.is_inference(view) assert torch.is_inference(copied) + + +@withCUDA +@pytest.mark.parametrize("offset_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("perm", [[2, 0, 1], [0, 1, 2]]) +@pytest.mark.parametrize("layout", ["slice", "stride", "broadcast"]) +def test_permute_entries( + device: torch.device, + offset_dtype: torch.dtype, + perm: list[int], + layout: str, +) -> None: + values = StringTensor.from_arrow( + pa.array(["unused", "é🙂", None, "", "long尾"]) + ) + assert values._valid is not None + values = StringTensor( + data=values._data.to(device), + offset=values._offset.to(device=device, dtype=offset_dtype), + valid=values._valid.to(device), + size=values.shape, + ) + if layout == "slice": + values = cast(StringTensor, values[1:4]) + elif layout == "stride": + values = cast(StringTensor, values[::2]) + else: + values = cast(StringTensor, values[1:2].expand(3)) + index = torch.tensor(perm, device=device) + expected = cast(StringTensor, values[index]) + actual = _permute_entries(values, index) + assert actual.tolist() == expected.tolist() + assert actual.device == expected.device + assert actual.data_offset[1].dtype == expected.data_offset[1].dtype + assert actual.data_offset[1].diff().equal(expected.data_offset[1].diff()) + assert not actual._data[actual._offset[-1] :].any() + + +def test_permute_entries_skewed_lengths() -> None: + # Capacity follows the stored bytes, not count times the longest entry. + values = StringTensor.from_list(["a"] * 1000 + ["b" * 100_000]) + perm = torch.arange(values.numel()).flip(0) + actual = _permute_entries(values, perm) + assert actual.tolist() == values[perm].tolist() + assert actual._data.numel() == values._data.numel() + + +@withCUDA +@pytest.mark.parametrize("perm", [[4, 0, 3, 1, 2], [0, 1, 2, 3, 4]]) +def test_permute_entries_long_and_short( + device: torch.device, + perm: list[int], +) -> None: + # Mostly long entries, mixed with short, empty and null ones; the last + # entry ends at the end of the payload. + values = StringTensor.from_arrow( + pa.array(["a" * 41, "bc", None, "", "é" * 59 + "z"]), device=device + ) + index = torch.tensor(perm, device=device) + actual = _permute_entries(values, index) + assert actual.tolist() == values[index].tolist() + assert actual._data.numel() == values._data.numel() + + +def test_permute_entries_fullgraph(compilation: None) -> None: + values = StringTensor.from_arrow(pa.array(["é🙂", None, "", "long尾"])) + perm = torch.tensor([3, 0, 2, 1]) + compiled = torch.compile(_permute_entries, fullgraph=True) + assert compiled(values, perm).tolist() == values[perm].tolist() + + +def test_permute_entries_meta() -> None: + values = StringTensor( + data=torch.empty(12, dtype=torch.uint8, device="meta"), + offset=torch.empty(4, dtype=torch.int32, device="meta"), + valid=torch.empty(3, dtype=torch.bool, device="meta"), + size=(3,), + ) + perm = torch.empty(3, dtype=torch.int64, device="meta") + actual = _permute_entries(values, perm) + assert actual.shape == perm.shape + assert actual.dtype == torch.uint8 + assert actual._offset.shape == (perm.numel() + 1,) + assert actual._offset.dtype == torch.int32 + assert actual._data.numel() == values._data.numel() + assert actual.valid is not None + assert actual.valid.shape == perm.shape + + +def test_permute_entries_payloads() -> None: + cases = [ + # Boolean payload. + VarLenTensor( + data=torch.tensor([True, False, True, True]), + offset=torch.tensor([0, 1, 4]), + valid=None, + size=(2,), + ), + # A null slot holding a byte. + StringTensor( + data=torch.tensor([97, 255, 98], dtype=torch.uint8), + offset=torch.tensor([0, 1, 2, 3], dtype=torch.int32), + valid=torch.tensor([True, False, True]), + size=(3,), + ), + # Empty strings with an empty buffer, and no entries at all. + StringTensor.from_arrow(pa.array(["", None, ""])), + StringTensor.from_arrow(pa.array([], type=pa.string())), + ] + for values in cases: + perm = torch.arange(values.numel()).roll(1) + expected = cast(VarLenTensor, values[perm]) + actual = _permute_entries(values, perm) + assert actual.dtype == values.dtype + assert actual.tolist() == expected.tolist() + assert actual.data_offset[1].equal(expected.data_offset[1])