From 03ddabe57732b54a02f6d336115c00071fefb4c5 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:38:25 -0700 Subject: [PATCH 1/4] Add bounded native variable-length payload selection --- sdm/tensor/var_len.py | 36 ++++++++++++++++++++ test/tensor/test_var_len.py | 68 ++++++++++++++++++++++++++++++++++++- 2 files changed, 103 insertions(+), 1 deletion(-) diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py index 8557d3845..308c6f880 100644 --- a/sdm/tensor/var_len.py +++ b/sdm/tensor/var_len.py @@ -1577,3 +1577,39 @@ def _materialize( ) out._max_length = inp._max_length return out + + +def _select_entries(inp: VarLenTensor, index: Tensor) -> VarLenTensor: + """Compact selected entries into bounded backing storage. + + The input and index are one-dimensional. Padding after the final offset + keeps the output capacity static without truncating selected values. + """ + index = index.where(index >= 0, index + inp.numel()) + positions = index * inp.stride(0) + inp._logical_storage_offset + start = inp._offset[positions] + end = inp._offset[positions + 1] + max_length = inp._max_length + capacity = index.numel() * max_length + if capacity > torch.iinfo(inp._offset.dtype).max: + return cast(VarLenTensor, inp[index]) + offset = torch.cat((start.new_zeros(1), (end - start).cumsum(0))) + offset = offset.to(inp._offset.dtype) + if index.numel() == 0 or inp._data.numel() == 0: + data = inp._data.new_empty(0) + else: + byte = torch.arange(capacity, device=inp.device, dtype=offset.dtype) + owner = torch.searchsorted(offset[1:], byte, right=True) + owner = owner.clamp_max(index.numel() - 1) + source = start[owner] + byte - offset[owner] + used = byte < offset[-1] + source = source.where(used, 0) + data = inp._data[source].where(used, 0) + out = inp.__class__( + data=data, + offset=offset, + valid=inp._valid[positions] if inp._valid is not None else None, + size=index.shape, + ) + out._max_length = max_length + return out diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py index 864a4e085..eb3ba4775 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, _select_entries from sdm.testing import onlyCUDA, withCUDA @@ -1034,3 +1034,69 @@ 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("indices", [[], [2, 0, 2, 1], [0, 1, 2], [-1, 0]]) +@pytest.mark.parametrize("layout", ["slice", "stride", "broadcast"]) +def test_native_string_selection( + device: torch.device, + offset_dtype: torch.dtype, + indices: 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(indices, device=device, dtype=torch.int64) + expected = cast(StringTensor, values[index]) + actual = _select_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()) + + +def test_native_string_selection_fullgraph() -> None: + values = StringTensor.from_arrow(pa.array(["é🙂", None, "", "long尾"])) + index = torch.tensor([3, 0, 3, 1, 2]) + compiled = torch.compile(_select_entries, fullgraph=True) + actual = compiled(values, index) + assert actual.tolist() == values[index].tolist() + + +@pytest.mark.parametrize("empty", [False, True]) +def test_native_string_selection_empty_payload(empty: bool) -> None: + values = StringTensor.from_arrow( + pa.array([], type=pa.string()) if empty else pa.array(["", None, ""]) + ) + index = torch.tensor([] if empty else [2, 0, 1, 2], dtype=torch.int64) + actual = _select_entries(values, index) + assert actual.tolist() == values[index].tolist() + + +def test_native_string_selection_null_payload() -> None: + values = 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,), + ) + index = torch.tensor([1, 2, 0, 1]) + actual = _select_entries(values, index) + assert actual.tolist() == values[index].tolist() + assert actual.data_offset[1].diff().tolist() == [1, 1, 1, 1] From c2bc40c28dffa805d87c373b8c841e3c81b33edd Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:41:32 -0700 Subject: [PATCH 2/4] Use compilation fixture for native string selection test --- test/tensor/test_var_len.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py index eb3ba4775..aa24a5f2b 100644 --- a/test/tensor/test_var_len.py +++ b/test/tensor/test_var_len.py @@ -1071,7 +1071,7 @@ def test_native_string_selection( assert actual.data_offset[1].diff().equal(expected.data_offset[1].diff()) -def test_native_string_selection_fullgraph() -> None: +def test_native_string_selection_fullgraph(compilation: None) -> None: values = StringTensor.from_arrow(pa.array(["é🙂", None, "", "long尾"])) index = torch.tensor([3, 0, 3, 1, 2]) compiled = torch.compile(_select_entries, fullgraph=True) From 10d6b3b8dc25b84aa1a3ad8edba0812cc93f6b0f Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:53:14 -0700 Subject: [PATCH 3/4] Verify native string selection metadata inference --- test/tensor/test_var_len.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/test/tensor/test_var_len.py b/test/tensor/test_var_len.py index aa24a5f2b..f1be5030e 100644 --- a/test/tensor/test_var_len.py +++ b/test/tensor/test_var_len.py @@ -1079,6 +1079,25 @@ def test_native_string_selection_fullgraph(compilation: None) -> None: assert actual.tolist() == values[index].tolist() +def test_native_string_selection_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,), + ) + index = torch.empty(4, dtype=torch.int64, device="meta") + actual = _select_entries(values, index) + assert actual.shape == index.shape + assert actual.dtype == torch.uint8 + assert actual.device == index.device + assert actual._offset.shape == (index.numel() + 1,) + assert actual._offset.dtype == torch.int32 + assert actual._data.numel() == index.numel() * 12 + assert actual.valid is not None + assert actual.valid.shape == index.shape + + @pytest.mark.parametrize("empty", [False, True]) def test_native_string_selection_empty_payload(empty: bool) -> None: values = StringTensor.from_arrow( From 3600d42ce2f87201106fef1907829bdc90eaf861 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:58:18 -0700 Subject: [PATCH 4/4] Accelerate CPU payload selection while preserving the CUDA path --- sdm/tensor/var_len.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/sdm/tensor/var_len.py b/sdm/tensor/var_len.py index 308c6f880..39902348c 100644 --- a/sdm/tensor/var_len.py +++ b/sdm/tensor/var_len.py @@ -1597,6 +1597,16 @@ def _select_entries(inp: VarLenTensor, index: Tensor) -> VarLenTensor: offset = offset.to(inp._offset.dtype) if index.numel() == 0 or inp._data.numel() == 0: data = inp._data.new_empty(0) + elif inp._data.device.type == "cpu": + byte = torch.arange(max_length, device=inp.device, dtype=offset.dtype) + used = byte < (end - start)[:, None] + source = (start[:, None] + byte).clamp_max(inp._data.numel() - 1) + target = torch.where(used, offset[:-1, None] + byte, capacity) + values = inp._data[source].where(used, 0) + # Padding writes only to the extra byte, outside the logical payload. + data = inp._data.new_zeros(capacity + 1).scatter_( + 0, target.flatten().long(), values.flatten() + )[:-1] else: byte = torch.arange(capacity, device=inp.device, dtype=offset.dtype) owner = torch.searchsorted(offset[1:], byte, right=True)