Record string length bounds for compiled string operations

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

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)

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,15 +1427,27 @@ 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,))
     end = torch.as_strided(end, size=(end.numel(),), stride=(1,))
 
     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,
     )

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]