From c7e37c961fd8ab1e576add2214ca14391ce21760 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:59:52 -0700 Subject: [PATCH] Record string length bounds for compiled string operations Record an upper bound on every string's byte length when string columns are loaded from Arrow or cuDF, rounded up to 2**k - 1 bytes, and keep it through copies, views, concatenation, compilation and pickling. Compiled string operations can then size their buffers without reading lengths from the device. --- sdm/tensor/string.py | 31 +++++++++++++++++++++ sdm/tensor/table.py | 1 + sdm/tensor/var_len.py | 42 +++++++++++++++++++++++++--- test/tensor/test_string.py | 56 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 126 insertions(+), 4 deletions(-) diff --git a/sdm/tensor/string.py b/sdm/tensor/string.py index 6d3d4e7ac0..fdcf265d18 100644 --- a/sdm/tensor/string.py +++ b/sdm/tensor/string.py @@ -34,6 +34,9 @@ class StringTensor(VarLenTensor): size: The shape of the tensor. stride: The stride of the tensor. storage_offset: The offset into the logical ``offset`` storage. + max_length: Upper bound on the byte length of every stored string, + or :obj:`None` if unknown. Lets compiled operations size buffers + without reading lengths from the device. Must not underestimate. """ # NOTE Assume that `data` stores valid UTF-8 bytes and do not validate it. @@ -51,6 +54,7 @@ def _make_var_len_tensor( size: tuple[int, ...], stride: tuple[int, ...], storage_offset: int, + max_length: int | None, cls: type[VarLenTensor] | None = None, ) -> VarLenTensor: return VarLenTensor._make_var_len_tensor( @@ -60,6 +64,7 @@ def _make_var_len_tensor( size=size, stride=stride, storage_offset=storage_offset, + max_length=max_length, cls=StringTensor if cls is None else cls, ) @@ -131,6 +136,10 @@ def from_arrow( valid=valid, size=size, storage_offset=storage_offset, + # Inspect CPU offsets before moving buffers. + max_length=_length_bound( + offset[storage_offset : storage_offset + len(array) + 1] + ), ) @override @@ -203,6 +212,7 @@ def from_cudf( offset=torch.zeros(1, dtype=torch.int32, device=data.device), valid=None, size=size, + max_length=0, ) offsets = column.children()[0] # (base_size + 1,) INT32/INT64 values @@ -228,6 +238,10 @@ def from_cudf( valid=valid, size=size, storage_offset=storage_offset, + # Reads one value from the device, once at ingestion. + max_length=_length_bound( + offset[storage_offset : storage_offset + len(ser) + 1] + ), ) def to_cudf(self) -> cudf.Series: @@ -546,3 +560,20 @@ def _sort( perm = torch.from_dlpack(perm_ser.astype("int64").to_cupy()) return cast(StringTensor, inp[perm]), perm + + +def _length_bound(offset: Tensor) -> int: + """Return an upper bound on the byte length of every stored string. + + ``torch.compile`` treats the bound as a constant and compiles a separate + graph for each distinct value. Rounding the longest length up to the + next ``2**k - 1`` bytes (1, 3, 7, ..., 127, 255, ...) lets inputs with + similar lengths share one graph, at the cost of a bound up to about + twice the longest length. Pass only the offsets of the tensor's own + strings: views and copies of the tensor cannot reach other strings in + a shared buffer, so the bound stays valid for them. + """ + lengths = offset[1:] - offset[:-1] + longest = int(lengths.max()) if lengths.numel() > 0 else 0 + # The smallest 2**k - 1 that is at least `longest`: + return 2 ** longest.bit_length() - 1 diff --git a/sdm/tensor/table.py b/sdm/tensor/table.py index 996fc8b0e6..5a853bffda 100644 --- a/sdm/tensor/table.py +++ b/sdm/tensor/table.py @@ -306,6 +306,7 @@ def __new__( offset=torch.zeros(1, dtype=torch.int32, device=device), valid=None, size=(*size, 0), + max_length=0, ) if id is None: id = ColumnarTensor((), size=size, device=device) diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py index fae8102ef2..e515e39f00 100644 --- a/sdm/tensor/var_len.py +++ b/sdm/tensor/var_len.py @@ -61,6 +61,9 @@ class VarLenTensor(Tensor): size: The shape of the tensor. stride: The stride of the tensor. storage_offset: The offset into the logical ``offset`` storage. + max_length: Upper bound on the length of every stored entry, or + :obj:`None` if unknown. Lets compiled operations size buffers + without reading lengths from the device. Must not underestimate. """ ALLOWED_DTYPES: ClassVar[tuple[torch.dtype, ...] | None] = None @@ -71,6 +74,7 @@ class VarLenTensor(Tensor): _data: Tensor _offset: Tensor _valid: Tensor | None + _max_length: int | None _layout: Tensor # Constructors ############################################################ @@ -83,6 +87,7 @@ def __init__( size: Sequence[int], stride: Sequence[int] | None = None, storage_offset: int = 0, + max_length: int | None = None, ) -> None: pass @@ -94,6 +99,7 @@ def __new__( size: Sequence[int], stride: Sequence[int] | None = None, storage_offset: int = 0, + max_length: int | None = None, ) -> Self: r"""Create a tensor wrapper.""" size = tuple(size) @@ -221,6 +227,7 @@ def __new__( size=size, stride=stride, storage_offset=storage_offset, + max_length=max_length, cls=None if "_make_var_len_tensor" in cls.__dict__ else cls, ) @@ -234,6 +241,7 @@ def _make_var_len_tensor( size: tuple[int, ...], stride: tuple[int, ...], storage_offset: int, + max_length: int | None, cls: type[VarLenTensor] | None = None, ) -> VarLenTensor: cls = VarLenTensor if cls is None else cls @@ -249,6 +257,7 @@ def _make_var_len_tensor( out._data = data out._offset = offset out._valid = valid + out._max_length = max_length # Mirror size and stride in an empty tensor's shape so torch.compile # can make them symbolic. A tensor nested in another one, like a # categorical vocabulary, would otherwise keep its first traced size. @@ -291,6 +300,7 @@ def from_tensor( size=tensor.size(), stride=tensor.stride(), storage_offset=int(tensor.storage_offset()), + max_length=min(data.numel(), 1), # Every entry has length 1. ) @classmethod @@ -565,7 +575,7 @@ def __tensor_flatten__(self) -> tuple[list[str], tuple[Any, ...]]: attrs = ["_offset", "_data", "_layout"] if self._valid is not None: attrs.append("_valid") - return attrs, (self.__class__, self.storage_offset()) + return attrs, (self.__class__, self.storage_offset(), self._max_length) @staticmethod def __tensor_unflatten__( @@ -574,7 +584,7 @@ def __tensor_unflatten__( outer_size: tuple[int, ...], outer_stride: tuple[int, ...], ) -> VarLenTensor: - cls, storage_offset = ctx + cls, storage_offset, max_length = ctx layout = inner_tensors["_layout"] out = cls( data=inner_tensors["_data"], @@ -583,6 +593,7 @@ def __tensor_unflatten__( size=layout.size()[:-1], stride=layout.stride()[:-1], storage_offset=storage_offset, + max_length=max_length, ) out._layout = layout return out @@ -595,6 +606,7 @@ def __reduce_ex__(self, proto: SupportsIndex) -> Any: tuple(self.size()), tuple(self.stride()), int(self.storage_offset()), + self._max_length, ) return (self.__class__, args) @@ -736,6 +748,7 @@ def _alias(inp: VarLenTensor) -> VarLenTensor: size=inp.size(), stride=inp.stride(), storage_offset=int(inp.storage_offset()), + max_length=inp._max_length, ) @@ -845,6 +858,7 @@ def _to_dtype_layout( if memory_format == torch.preserve_format else None, storage_offset=0, + max_length=inp._max_length, ) @@ -954,6 +968,7 @@ def _detach(inp: VarLenTensor) -> VarLenTensor: size=inp.size(), stride=inp.stride(), storage_offset=int(inp.storage_offset()), + max_length=inp._max_length, ) @@ -972,6 +987,7 @@ def _pin_memory(inp: VarLenTensor) -> VarLenTensor: size=inp.size(), stride=inp.stride(), storage_offset=int(inp.storage_offset()), + max_length=inp._max_length, ) @@ -1366,6 +1382,10 @@ def _cat(tensors: Sequence[Tensor], dim: int = 0) -> VarLenTensor: tensors = tuple( cast(VarLenTensor, tensor.contiguous()) for tensor in tensors ) + lengths = tuple(tensor._max_length for tensor in tensors) + max_length = ( + None if None in lengths else max(cast(tuple[int, ...], lengths)) + ) data_list, offsets = zip(*(_data_offset(tensor) for tensor in tensors)) valid: Tensor | None = None if any(tensor._valid is not None for tensor in tensors): @@ -1407,7 +1427,13 @@ def _cat(tensors: Sequence[Tensor], dim: int = 0) -> VarLenTensor: if math.prod(tensors[0].size()[:dim]) == 1: # Contiguous path: offset[-1] = storage_offset - return tensor_cls(data=data, offset=offset, valid=valid, size=size) + return tensor_cls( + data=data, + offset=offset, + valid=valid, + size=size, + max_length=max_length, + ) end = torch.cat(end_views, dim=dim) start = torch.as_strided(start, size=(start.numel(),), stride=(1,)) @@ -1415,7 +1441,13 @@ def _cat(tensors: Sequence[Tensor], dim: int = 0) -> VarLenTensor: offset, index = _compact(start, end) - return tensor_cls(data=data[index], offset=offset, valid=valid, size=size) + return tensor_cls( + data=data[index], + offset=offset, + valid=valid, + size=size, + max_length=max_length, + ) @VarLenTensor.implements(aten.stack.default) @@ -1465,6 +1497,7 @@ def _from_layout_view(inp: VarLenTensor, view: Tensor) -> VarLenTensor: size=view.size(), stride=view.stride(), storage_offset=0 if view.numel() == 0 else int(storage_offset), + max_length=inp._max_length, ) @@ -1554,4 +1587,5 @@ def _materialize( size=size, stride=stride, storage_offset=0, + max_length=inp._max_length, ) diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py index 3311c440c1..6b2c9e4b37 100644 --- a/test/tensor/test_string.py +++ b/test/tensor/test_string.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pickle from typing import cast import pyarrow as pa @@ -159,6 +160,8 @@ def test_from_cudf_sliced_values() -> None: assert tensor._offset.equal( torch.tensor([0, 1, 3, 5, 5], device=tensor.device) ) + # The bound covers the slice, rounded up to 2**k - 1 bytes. + assert tensor._max_length == 3 @onlyCUDA @@ -184,6 +187,7 @@ def test_from_cudf_all_empty_values() -> None: assert tensor.tolist() == ["", ""] assert tensor._data.numel() == 0 assert tensor._offset.equal(torch.tensor([0, 0, 0], device=tensor.device)) + assert tensor._max_length == 0 @onlyCUDA @@ -198,6 +202,7 @@ def test_from_cudf_empty_values() -> None: assert tensor.tolist() == [] assert tensor._data.numel() == 0 assert tensor._offset.equal(torch.tensor([0], device=tensor.device)) + assert tensor._max_length == 0 @pytest.mark.parametrize( @@ -406,3 +411,54 @@ def test_registers_as_a_pytorch_buffer() -> None: module.register_buffer("value", tensor) assert tuple(module.buffers()) == (tensor,) + + +@pytest.mark.parametrize( + ("length", "bound"), + [(0, 0), (1, 1), (3, 3), (4, 7), (127, 127), (128, 255)], +) +def test_bound_rounds_up(length: int, bound: int) -> None: + values = StringTensor.from_list(["a" * length, ""]) + assert values._max_length == bound + + +@withCUDA +@pytest.mark.parametrize("bound", [None, 0, 15]) +def test_bound_provenance( + compilation: None, + device: torch.device, + bound: int | None, +) -> None: + texts = ["", None] if bound == 0 else ["abcdefgh", "é"] + # A non-null Arrow slice shares storage with a longer, unreachable string. + array = pa.array(texts) + if bound == 15: + array = pa.array(["x" * 40, *texts]).slice(1) + values = StringTensor.from_arrow(array, device=device) + if bound is None: + values = StringTensor( + data=values._data, + offset=values._offset, + valid=values._valid, + size=values.size(), + ) + compiled = torch.compile(lambda value: value.detach(), fullgraph=True) + for copied, expected in ( + (values[:], texts), + (values[::2], texts[::2]), + (values[torch.tensor([1, 0], device=device)], texts[::-1]), + (values.clone(), texts), + (values.detach(), texts), + (values.to(device, copy=True), texts), + (compiled(values), texts), + (pickle.loads(pickle.dumps(values)), texts), + ): + copied = cast(StringTensor, copied) + assert copied._max_length == bound + assert copied.tolist() == expected + for suffix in ("", "x" * 16): + extra = StringTensor.from_list([suffix], device=device) + joined = cast(StringTensor, torch.cat((values, extra))) + new_bound = None if bound is None else max(bound, 31 if suffix else 0) + assert joined._max_length == new_bound + assert joined.tolist() == [*texts, suffix]