From 051106054caa2af07b441df1453353aaefbbb681 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:40:42 -0700 Subject: [PATCH 1/5] Use tensor keys for categorical string lookup --- sdm/processing/categorical/align.py | 54 +++++++++++ test/processing/categorical/test_align.py | 105 +++++++++++++++++++--- 2 files changed, 149 insertions(+), 10 deletions(-) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 3a4cbcf924..c662a39c04 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -6,6 +6,7 @@ import torch from torch import Tensor +from torch.fx.experimental.symbolic_shapes import statically_known_true from sdm import ( CategoricalTensor, @@ -18,6 +19,7 @@ from sdm.nn._buffer import BufferList from sdm.processing import EnsembleProcessor from sdm.relational.join import join_index +from sdm.tensor._string_keys import pack from sdm.tensor.categorical import _iter_categories from sdm.tensor.var_len import _cat, _contiguous @@ -340,6 +342,21 @@ def _string_category_lookups( ) -> tuple[Tensor, ...]: if not input_categories: return () + # Bound packed-key work to eight words per category. + if all( + max(left._max_length, cast(StringTensor, right)._max_length) <= 56 + for left, right in zip( + input_categories, fitted_categories, strict=True + ) + ): + return tuple( + _native_string_lookup( + left, cast(StringTensor, right), codes.dtype + ) + for left, right in zip( + input_categories, fitted_categories, strict=True + ) + ) # StringTensor does not handle this custom op. Pass its underlying # tensors and layout without copying the string data. inputs, input_layout = _flatten_strings(input_categories) @@ -501,6 +518,43 @@ def __repr__(self, *, indent: int = 0) -> str: ) +def _native_string_lookup( + left: StringTensor, right: StringTensor, dtype: torch.dtype +) -> Tensor: + width = max(left._max_length, right._max_length) + left_keys, _ = pack(left, width) + right_keys, _ = pack(right, width) + right_ids = torch.arange(right.numel(), dtype=dtype, device=right.device) + # Limit the pairwise temporary; larger vocabularies use sorting. + if statically_known_true(left_keys.numel() * right.numel() <= 65536): + # [L, R]: category vectors contain no nulls. + equal = (left_keys[:, None, :] == right_keys[None, :, :]).all(-1) + return torch.where(equal, right_ids, -1).amax(-1) + + keys = torch.cat((left_keys, right_keys)) + order = torch.arange(keys.size(0), device=keys.device) + for column in range(keys.size(1) - 1, -1, -1): + order = order[keys[order, column].argsort(stable=True)] + sorted_keys = keys[order] + starts = torch.cat( + ( + torch.ones(1, dtype=torch.bool, device=keys.device), + (sorted_keys[1:] != sorted_keys[:-1]).any(-1), + ) + ) + groups = starts.long().cumsum(0) - 1 + values = torch.cat( + ( + right_ids.new_full((left.numel(),), -1), + right_ids, + ) + ) + matches = right_ids.new_full((keys.size(0),), -1) + matches.scatter_reduce_(0, groups, values[order], reduce="amax") + result = torch.empty_like(matches).scatter_(0, order, matches[groups]) + return result[: left.numel()] + + def _flatten_strings( categories: tuple[Tensor, ...], ) -> tuple[list[Tensor | None], list[int]]: diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index 52cf0a6b95..cafc469c98 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -506,10 +506,15 @@ def test_align_categories_ensemble_matches_member_fits( assert fitted_query_output[member_id].equal(expected_query) -def test_string_category_lookups_compile(compilation: None) -> None: +@pytest.mark.parametrize("count", [0, 300]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_string_category_lookups_compile( + compilation: None, count: int, dynamic: bool +) -> None: + extra = [f"long-prefix-{index:04}" for index in range(count)] fitted = ( - StringTensor.from_list(["a", "b"]), - StringTensor.from_list(["a", "x"]), + StringTensor.from_list(["a", "b", *extra]), + StringTensor.from_list(["a", "x", *extra]), ) codes = torch.empty((), dtype=torch.int32) @@ -520,23 +525,103 @@ def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: codes=codes, ) - compiled = torch.compile(lookups, backend="inductor", fullgraph=True) - first, second = compiled(StringTensor.from_list(["b", "z", "a"])) + graph = torch._dynamo.explain(lookups)( + StringTensor.from_list(["b", "z", "a", *extra[::-1]]) + ) + assert graph.graph_break_count == 0 + assert all( + "string_category_lookup" not in str(node.target) + for captured in graph.graphs + for node in captured.graph.nodes + ) + compiled = torch.compile( + lookups, backend="inductor", fullgraph=True, dynamic=dynamic + ) + first, second = compiled( + StringTensor.from_list(["b", "z", "a", *extra[::-1]]) + ) torch.testing.assert_close( actual=first, - expected=torch.tensor([1, -1, 0], dtype=codes.dtype), + expected=torch.tensor( + [1, -1, 0, *range(count + 1, 1, -1)], dtype=codes.dtype + ), ) torch.testing.assert_close( actual=second, - expected=torch.tensor([-1, -1, 0], dtype=codes.dtype), + expected=torch.tensor( + [-1, -1, 0, *range(count + 1, 1, -1)], dtype=codes.dtype + ), ) - first, second = compiled(StringTensor.from_list(["x", "z", "q"])) + first, second = compiled(StringTensor.from_list(["x", "z", "q", *extra])) torch.testing.assert_close( actual=first, - expected=torch.tensor([-1, -1, -1], dtype=codes.dtype), + expected=torch.tensor( + [-1, -1, -1, *range(2, count + 2)], dtype=codes.dtype + ), ) torch.testing.assert_close( actual=second, - expected=torch.tensor([1, -1, -1], dtype=codes.dtype), + expected=torch.tensor( + [1, -1, -1, *range(2, count + 2)], dtype=codes.dtype + ), + ) + + +@withCUDA +@pytest.mark.parametrize("count", [8, 300]) +def test_align_categories_utf8_views(count: int, device: torch.device) -> None: + values = [ + f"é\x00long-prefix-{index:04}" + ("\x00" if index % 2 else "") + for index in range(count) + ] + ["a", "a\x00", ""] + backing = StringTensor.from_list( + [ + "skip", + *[ + value + for key in [*values, "unknown"] + for value in (key, "skip") + ], + ], + device=device, + ) + context = _table( + [[index] for index in range(len(values))], + categories=(tuple(values[::-1]),), + device=device, + ) + query = TableTensor( + categorical=CategoricalTensor( + code=torch.arange( + -1, len(values) + 1, dtype=torch.int32, device=device + )[:, None], + categories=(backing[1::2],), + ), ) + output = AlignCategories().fit(context).transform(query) + assert output.categorical.code.squeeze(-1).tolist() == [ + -1, + *range(len(values) - 1, -1, -1), + -1, + ] + + +@withCUDA +@pytest.mark.parametrize("width", [7, 55, 56, 257]) +def test_align_categories_distinguishes_long_common_prefixes( + width: int, device: torch.device +) -> None: + prefix = "a" * width + context = _table( + [[0], [1]], + categories=((prefix + "x", prefix + "y"),), + device=device, + ) + query = _table( + [[0], [1], [2], [-1]], + categories=((prefix + "y", prefix + "z", prefix + "x"),), + device=device, + ) + output = AlignCategories().fit(context).transform(query) + assert output.categorical.code.squeeze(-1).tolist() == [1, -1, 0, -1] From 658a051f6900a660c5c5a964cb4bc62a91baf94d Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:35:52 -0700 Subject: [PATCH 2/5] Keep long-string category joins eager without custom registration --- sdm/processing/categorical/align.py | 73 ++--------------------- test/processing/categorical/test_align.py | 61 +++++++++++++++++-- 2 files changed, 62 insertions(+), 72 deletions(-) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index c662a39c04..3bb073cead 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -357,15 +357,11 @@ def _string_category_lookups( input_categories, fitted_categories, strict=True ) ) - # StringTensor does not handle this custom op. Pass its underlying - # tensors and layout without copying the string data. - inputs, input_layout = _flatten_strings(input_categories) - fitted, fitted_layout = _flatten_strings(fitted_categories) lookup = _string_category_lookup( - inputs=inputs, - input_layout=input_layout, - fitted=fitted, - fitted_layout=fitted_layout, + input_categories=input_categories, + fitted_categories=cast( + tuple[StringTensor, ...], fitted_categories + ), codes=codes, ) return lookup.split( @@ -555,45 +551,7 @@ def _native_string_lookup( return result[: left.numel()] -def _flatten_strings( - categories: tuple[Tensor, ...], -) -> tuple[list[Tensor | None], list[int]]: - """Pack vocabulary tensors and their (length, stride, offset) metadata.""" - tensors: list[Tensor | None] = [] - layout: list[int] = [] - for category in categories: - category = cast(StringTensor, category) - tensors.extend((category._data, category._offset, category._valid)) - layout.extend( - ( - category.numel(), - category.stride(0), - category._logical_storage_offset, - ) - ) - return tensors, layout - - -def _unflatten_strings( - tensors: list[Tensor | None], layout: list[int] -) -> tuple[StringTensor, ...]: - """Rebuild string views over the same tensors for the join.""" - return tuple( - StringTensor( - data=cast(Tensor, tensors[3 * i]), - offset=cast(Tensor, tensors[3 * i + 1]), - valid=tensors[3 * i + 2], - size=(layout[3 * i],), - stride=(layout[3 * i + 1],), - storage_offset=layout[3 * i + 2], - ) - for i in range(len(layout) // 3) - ) - - def _concatenate_strings(categories: tuple[StringTensor, ...]) -> StringTensor: - # Calling torch.cat here fails on StringTensor during compiled execution. - # Call SDM's concatenation implementation directly. contiguous = tuple( category if category.is_contiguous() else _contiguous(category) for category in categories @@ -601,17 +559,11 @@ def _concatenate_strings(categories: tuple[StringTensor, ...]) -> StringTensor: return cast(StringTensor, _cat(contiguous)) -@torch.library.custom_op("sdm::string_category_lookup", mutates_args=()) def _string_category_lookup( - inputs: list[Tensor | None], - input_layout: list[int], - fitted: list[Tensor | None], - fitted_layout: list[int], + input_categories: tuple[StringTensor, ...], + fitted_categories: tuple[StringTensor, ...], codes: Tensor, ) -> Tensor: - """Keep the backend join's dynamic row pairs inside a fixed-size lookup.""" - input_categories = _unflatten_strings(inputs, input_layout) - fitted_categories = _unflatten_strings(fitted, fitted_layout) # Pair IDs for each input category: (total_input_categories,). left_pair = torch.cat( [ @@ -660,16 +612,3 @@ def _string_category_lookup( ) lookup[left_index] = fitted_codes[right_index] return lookup - - -@_string_category_lookup.register_fake -def _string_category_lookup_fake( - inputs: list[Tensor | None], - input_layout: list[int], - fitted: list[Tensor | None], - fitted_layout: list[int], - codes: Tensor, -) -> Tensor: - # E.g. inputs ["red", "blue", "unknown"], fitted ["blue", "red"]: - # real lookup [1, 0, -1]; the fake result only describes its shape (3,). - return codes.new_empty((sum(input_layout[::3]),)) diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index cafc469c98..41958f28cc 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -458,17 +458,20 @@ def test_align_categories_rejects_changed_category_value_type() -> None: @withCUDA +@pytest.mark.parametrize("prefix", ["", "a" * 57]) def test_align_categories_ensemble_matches_member_fits( device: torch.device, + prefix: str, ) -> None: + categories = ((prefix + "red", prefix + "blue", prefix + "green"),) first_context = _table( [[0], [1]], - categories=(("red", "blue", "green"),), + categories=categories, device=device, ) second_context = _table( [[1], [2]], - categories=(("red", "blue", "green"),), + categories=categories, device=device, ) member_table_ids = (1, 0, 1, 0, 0, 1, 1, 0) @@ -478,7 +481,7 @@ def test_align_categories_ensemble_matches_member_fits( ) query = _table( [[0], [1], [2]], - categories=(("red", "blue", "green"),), + categories=categories, device=device, ) query_ensemble = EnsembleTable.from_tables( @@ -570,9 +573,12 @@ def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: @withCUDA @pytest.mark.parametrize("count", [8, 300]) -def test_align_categories_utf8_views(count: int, device: torch.device) -> None: +@pytest.mark.parametrize("prefix", ["", "a" * 57]) +def test_align_categories_utf8_views( + count: int, prefix: str, device: torch.device +) -> None: values = [ - f"é\x00long-prefix-{index:04}" + ("\x00" if index % 2 else "") + f"{prefix}é\x00long-prefix-{index:04}" + ("\x00" if index % 2 else "") for index in range(count) ] + ["a", "a\x00", ""] backing = StringTensor.from_list( @@ -625,3 +631,48 @@ def test_align_categories_distinguishes_long_common_prefixes( ) output = AlignCategories().fit(context).transform(query) assert output.categorical.code.squeeze(-1).tolist() == [1, -1, 0, -1] + + +@withCUDA +@pytest.mark.parametrize("prefix", ["", "a" * 57]) +def test_align_categories_mixed_category_types( + device: torch.device, prefix: str +) -> None: + context = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor( + [[0, 1], [1, 0], [0, 1]], dtype=torch.int32, device=device + ), + categories=( + StringTensor.from_list( + [prefix + "red", prefix + "blue"], device=device + ), + torch.tensor([20, 10], device=device), + ), + ), + ) + query = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor( + [[0, 0], [1, 1], [2, 2], [-1, -1]], + dtype=torch.int32, + device=device, + ), + categories=( + StringTensor.from_list( + [prefix + "blue", prefix + "unknown", prefix + "red"], + device=device, + ), + torch.tensor([10, 30, 20], device=device), + ), + ), + ) + + output = AlignCategories(sort_by="frequency").fit(context).transform(query) + + assert output.categorical.code.tolist() == [ + [1, 0], + [-1, -1], + [0, 1], + [-1, -1], + ] From 25ece92785960944e439232e45b60162fdf5e00d Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:55:53 -0700 Subject: [PATCH 3/5] Extend native categorical lookup width on CUDA --- sdm/processing/categorical/align.py | 6 ++- test/processing/categorical/test_align.py | 56 ++++++++++++++++------- 2 files changed, 43 insertions(+), 19 deletions(-) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 3bb073cead..93226a6c30 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -342,9 +342,11 @@ def _string_category_lookups( ) -> tuple[Tensor, ...]: if not input_categories: return () - # Bound packed-key work to eight words per category. + # Bound packed keys to twelve words on CUDA and eight elsewhere. + max_width = 84 if codes.is_cuda else 56 if all( - max(left._max_length, cast(StringTensor, right)._max_length) <= 56 + max(left._max_length, cast(StringTensor, right)._max_length) + <= max_width for left, right in zip( input_categories, fitted_categories, strict=True ) diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index 41958f28cc..4d0b8bdf8b 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -509,17 +509,35 @@ def test_align_categories_ensemble_matches_member_fits( assert fitted_query_output[member_id].equal(expected_query) +@withCUDA +@pytest.mark.parametrize("prefix", ["", "é" * 34]) @pytest.mark.parametrize("count", [0, 300]) @pytest.mark.parametrize("dynamic", [False, True]) def test_string_category_lookups_compile( - compilation: None, count: int, dynamic: bool + compilation: None, + count: int, + dynamic: bool, + prefix: str, + device: torch.device, ) -> None: - extra = [f"long-prefix-{index:04}" for index in range(count)] + if prefix and device.type != "cuda": + pytest.skip("Wide native string lookup is CUDA-only") + extra = [f"{prefix}long-prefix-{index:04}" for index in range(count)] fitted = ( - StringTensor.from_list(["a", "b", *extra]), - StringTensor.from_list(["a", "x", *extra]), + StringTensor.from_list( + [prefix + "a", prefix + "b", *extra], device=device + ), + StringTensor.from_list( + [prefix + "a", prefix + "x", *extra], device=device + ), ) - codes = torch.empty((), dtype=torch.int32) + codes = torch.empty((), dtype=torch.int32, device=device) + + def categories(values: list[str]) -> StringTensor: + return StringTensor.from_list( + [prefix + value for value in values] + extra[::-1], + device=device, + ) def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: return AlignCategories._string_category_lookups( @@ -528,9 +546,7 @@ def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: codes=codes, ) - graph = torch._dynamo.explain(lookups)( - StringTensor.from_list(["b", "z", "a", *extra[::-1]]) - ) + graph = torch._dynamo.explain(lookups)(categories(["b", "z", "a"])) assert graph.graph_break_count == 0 assert all( "string_category_lookup" not in str(node.target) @@ -540,33 +556,39 @@ def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: compiled = torch.compile( lookups, backend="inductor", fullgraph=True, dynamic=dynamic ) - first, second = compiled( - StringTensor.from_list(["b", "z", "a", *extra[::-1]]) - ) + first, second = compiled(categories(["b", "z", "a"])) torch.testing.assert_close( actual=first, expected=torch.tensor( - [1, -1, 0, *range(count + 1, 1, -1)], dtype=codes.dtype + [1, -1, 0, *range(count + 1, 1, -1)], + dtype=codes.dtype, + device=device, ), ) torch.testing.assert_close( actual=second, expected=torch.tensor( - [-1, -1, 0, *range(count + 1, 1, -1)], dtype=codes.dtype + [-1, -1, 0, *range(count + 1, 1, -1)], + dtype=codes.dtype, + device=device, ), ) - first, second = compiled(StringTensor.from_list(["x", "z", "q", *extra])) + first, second = compiled(categories(["x", "z", "q"])) torch.testing.assert_close( actual=first, expected=torch.tensor( - [-1, -1, -1, *range(2, count + 2)], dtype=codes.dtype + [-1, -1, -1, *range(count + 1, 1, -1)], + dtype=codes.dtype, + device=device, ), ) torch.testing.assert_close( actual=second, expected=torch.tensor( - [1, -1, -1, *range(2, count + 2)], dtype=codes.dtype + [1, -1, -1, *range(count + 1, 1, -1)], + dtype=codes.dtype, + device=device, ), ) @@ -614,7 +636,7 @@ def test_align_categories_utf8_views( @withCUDA -@pytest.mark.parametrize("width", [7, 55, 56, 257]) +@pytest.mark.parametrize("width", [7, 55, 56, 74, 83, 84, 257]) def test_align_categories_distinguishes_long_common_prefixes( width: int, device: torch.device ) -> None: From 0fed2dfd69290c99cf05feac4514a0226c1237a0 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:40:55 -0700 Subject: [PATCH 4/5] Use native category lookup on CUDA and simplify tests --- sdm/processing/categorical/align.py | 8 +- test/processing/categorical/test_align.py | 98 ++++++++++------------- 2 files changed, 45 insertions(+), 61 deletions(-) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 93226a6c30..545c54bd5d 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -342,11 +342,9 @@ def _string_category_lookups( ) -> tuple[Tensor, ...]: if not input_categories: return () - # Bound packed keys to twelve words on CUDA and eight elsewhere. - max_width = 84 if codes.is_cuda else 56 - if all( - max(left._max_length, cast(StringTensor, right)._max_length) - <= max_width + # Bound packed-key work to twelve words per category. + if codes.is_cuda and all( + max(left._max_length, cast(StringTensor, right)._max_length) <= 84 for left, right in zip( input_categories, fitted_categories, strict=True ) diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index 4d0b8bdf8b..8673a44c18 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -9,7 +9,8 @@ from sdm import CategoricalTensor, EnsembleTable, StringTensor, TableTensor from sdm.processing import AlignCategories -from sdm.testing import withCUDA +from sdm.processing.categorical.align import _native_string_lookup +from sdm.testing import onlyCUDA, withCUDA def _table( @@ -458,11 +459,10 @@ def test_align_categories_rejects_changed_category_value_type() -> None: @withCUDA -@pytest.mark.parametrize("prefix", ["", "a" * 57]) def test_align_categories_ensemble_matches_member_fits( device: torch.device, - prefix: str, ) -> None: + prefix = "a" * 85 categories = ((prefix + "red", prefix + "blue", prefix + "green"),) first_context = _table( [[0], [1]], @@ -509,19 +509,38 @@ def test_align_categories_ensemble_matches_member_fits( assert fitted_query_output[member_id].equal(expected_query) -@withCUDA -@pytest.mark.parametrize("prefix", ["", "é" * 34]) +@pytest.mark.parametrize("count", [8, 300]) +def test_native_string_lookup(count: int) -> None: + fitted = [f"é\0common-prefix-{index:04}" for index in range(count)] + [ + "a", + "a\0", + "", + ] + query = [*fitted[::-1], "unknown", "a\0\0"] + indices = {value: index for index, value in enumerate(fitted)} + actual = _native_string_lookup( + left=StringTensor.from_list(query), + right=StringTensor.from_list(fitted), + dtype=torch.int32, + ) + torch.testing.assert_close( + actual=actual, + expected=torch.tensor( + [indices.get(value, -1) for value in query], dtype=torch.int32 + ), + ) + + +@onlyCUDA @pytest.mark.parametrize("count", [0, 300]) @pytest.mark.parametrize("dynamic", [False, True]) def test_string_category_lookups_compile( compilation: None, count: int, dynamic: bool, - prefix: str, - device: torch.device, ) -> None: - if prefix and device.type != "cuda": - pytest.skip("Wide native string lookup is CUDA-only") + device = torch.device("cuda") + prefix = "é" * 34 extra = [f"{prefix}long-prefix-{index:04}" for index in range(count)] fitted = ( StringTensor.from_list( @@ -546,59 +565,29 @@ def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: codes=codes, ) - graph = torch._dynamo.explain(lookups)(categories(["b", "z", "a"])) - assert graph.graph_break_count == 0 - assert all( - "string_category_lookup" not in str(node.target) - for captured in graph.graphs - for node in captured.graph.nodes - ) compiled = torch.compile( lookups, backend="inductor", fullgraph=True, dynamic=dynamic ) - first, second = compiled(categories(["b", "z", "a"])) torch.testing.assert_close( - actual=first, - expected=torch.tensor( - [1, -1, 0, *range(count + 1, 1, -1)], - dtype=codes.dtype, - device=device, + actual=compiled(categories(["b", "z", "a"])), + expected=( + codes.new_tensor([1, -1, 0, *range(count + 1, 1, -1)]), + codes.new_tensor([-1, -1, 0, *range(count + 1, 1, -1)]), ), ) torch.testing.assert_close( - actual=second, - expected=torch.tensor( - [-1, -1, 0, *range(count + 1, 1, -1)], - dtype=codes.dtype, - device=device, - ), - ) - - first, second = compiled(categories(["x", "z", "q"])) - torch.testing.assert_close( - actual=first, - expected=torch.tensor( - [-1, -1, -1, *range(count + 1, 1, -1)], - dtype=codes.dtype, - device=device, - ), - ) - torch.testing.assert_close( - actual=second, - expected=torch.tensor( - [1, -1, -1, *range(count + 1, 1, -1)], - dtype=codes.dtype, - device=device, + actual=compiled(categories(["x", "z", "q"])), + expected=( + codes.new_tensor([-1, -1, -1, *range(count + 1, 1, -1)]), + codes.new_tensor([1, -1, -1, *range(count + 1, 1, -1)]), ), ) @withCUDA @pytest.mark.parametrize("count", [8, 300]) -@pytest.mark.parametrize("prefix", ["", "a" * 57]) -def test_align_categories_utf8_views( - count: int, prefix: str, device: torch.device -) -> None: +def test_align_categories_utf8_views(count: int, device: torch.device) -> None: + prefix = "a" * 57 values = [ f"{prefix}é\x00long-prefix-{index:04}" + ("\x00" if index % 2 else "") for index in range(count) @@ -636,7 +625,7 @@ def test_align_categories_utf8_views( @withCUDA -@pytest.mark.parametrize("width", [7, 55, 56, 74, 83, 84, 257]) +@pytest.mark.parametrize("width", [83, 84]) def test_align_categories_distinguishes_long_common_prefixes( width: int, device: torch.device ) -> None: @@ -656,9 +645,8 @@ def test_align_categories_distinguishes_long_common_prefixes( @withCUDA -@pytest.mark.parametrize("prefix", ["", "a" * 57]) def test_align_categories_mixed_category_types( - device: torch.device, prefix: str + device: torch.device, ) -> None: context = TableTensor( categorical=CategoricalTensor( @@ -666,9 +654,7 @@ def test_align_categories_mixed_category_types( [[0, 1], [1, 0], [0, 1]], dtype=torch.int32, device=device ), categories=( - StringTensor.from_list( - [prefix + "red", prefix + "blue"], device=device - ), + StringTensor.from_list(["red", "blue"], device=device), torch.tensor([20, 10], device=device), ), ), @@ -682,7 +668,7 @@ def test_align_categories_mixed_category_types( ), categories=( StringTensor.from_list( - [prefix + "blue", prefix + "unknown", prefix + "red"], + ["blue", "unknown", "red"], device=device, ), torch.tensor([10, 30, 20], device=device), From 7e2cb8f41530cc3e57703f36e6bf7f70af9c4f9d Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Tue, 27 Oct 2026 13:28:22 -0700 Subject: [PATCH 5/5] Allow native CUDA lookup for any known string width --- sdm/processing/categorical/align.py | 8 ++++---- test/processing/categorical/test_align.py | 18 +++++++++++++----- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 545c54bd5d..f2d496518b 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -342,9 +342,9 @@ def _string_category_lookups( ) -> tuple[Tensor, ...]: if not input_categories: return () - # Bound packed-key work to twelve words per category. if codes.is_cuda and all( - max(left._max_length, cast(StringTensor, right)._max_length) <= 84 + left._max_length is not None + and cast(StringTensor, right)._max_length is not None for left, right in zip( input_categories, fitted_categories, strict=True ) @@ -517,11 +517,11 @@ def __repr__(self, *, indent: int = 0) -> str: def _native_string_lookup( left: StringTensor, right: StringTensor, dtype: torch.dtype ) -> Tensor: - width = max(left._max_length, right._max_length) + width = max(cast(int, left._max_length), cast(int, right._max_length)) left_keys, _ = pack(left, width) right_keys, _ = pack(right, width) right_ids = torch.arange(right.numel(), dtype=dtype, device=right.device) - # Limit the pairwise temporary; larger vocabularies use sorting. + # Bound the pairwise boolean temporary to 64 KiB; otherwise sort keys. if statically_known_true(left_keys.numel() * right.numel() <= 65536): # [L, R]: category vectors contain no nulls. equal = (left_keys[:, None, :] == right_keys[None, :, :]).all(-1) diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index 8673a44c18..6415fa9c2b 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.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 Literal import pandas as pd @@ -540,7 +541,7 @@ def test_string_category_lookups_compile( dynamic: bool, ) -> None: device = torch.device("cuda") - prefix = "é" * 34 + prefix = "é" * 56 extra = [f"{prefix}long-prefix-{index:04}" for index in range(count)] fitted = ( StringTensor.from_list( @@ -550,6 +551,7 @@ def test_string_category_lookups_compile( [prefix + "a", prefix + "x", *extra], device=device ), ) + fitted = pickle.loads(pickle.dumps(fitted)) codes = torch.empty((), dtype=torch.int32, device=device) def categories(values: list[str]) -> StringTensor: @@ -625,7 +627,7 @@ def test_align_categories_utf8_views(count: int, device: torch.device) -> None: @withCUDA -@pytest.mark.parametrize("width", [83, 84]) +@pytest.mark.parametrize("width", [84, 127, 255, 1023]) def test_align_categories_distinguishes_long_common_prefixes( width: int, device: torch.device ) -> None: @@ -667,9 +669,15 @@ def test_align_categories_mixed_category_types( device=device, ), categories=( - StringTensor.from_list( - ["blue", "unknown", "red"], - device=device, + StringTensor( + data=torch.tensor( + list(b"blueunknownred"), + dtype=torch.uint8, + device=device, + ), + offset=torch.tensor([0, 4, 11, 14], device=device), + valid=None, + size=(3,), ), torch.tensor([10, 30, 20], device=device), ),