Add native variable-length entry permutation

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]

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])