diff --git a/sdm/processing/categorical/_string_lookup.py b/sdm/processing/categorical/_string_lookup.py new file mode 100644 index 0000000000..927ed40dba --- /dev/null +++ b/sdm/processing/categorical/_string_lookup.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Match string categories with compact keys and exact byte comparison.""" + +from collections.abc import Callable +from functools import partial +from typing import NamedTuple + +import torch +from torch import Tensor + +from sdm import StringTensor + + +class _StringRows(NamedTuple): + """Byte storage, start offsets and lengths for visible string labels.""" + + data: Tensor + starts: Tensor + lengths: Tensor + + +class _LookupKeys(NamedTuple): + """Flat packed keys and sorted fitted hashes for candidate lookup. + + ``fitted_order`` maps sorted-hash positions back to fitted category IDs. + """ + + query: Tensor + fitted: Tensor + query_hashes: Tensor + sorted_hashes: Tensor + fitted_order: Tensor + + +def _native_string_lookups( + left: tuple[StringTensor, ...], + right: tuple[StringTensor, ...], + dtype: torch.dtype = torch.int32, +) -> tuple[Tensor, ...]: + """Build a fitted-category index vector for each query/fitted pair.""" + return tuple( + _lookup_string_categories(query, fitted, dtype) + for query, fitted in zip(left, right, strict=True) + ) + + +def _lookup_string_categories( + query: StringTensor, fitted: StringTensor, dtype: torch.dtype = torch.int32 +) -> Tensor: + """Map query vocabulary labels to fitted IDs while preserving views. + + Missing labels map to ``-1``. The byte budget bounds key storage; + length metadata can avoid runtime branching. Broadcast rows are matched + once and expanded back afterward. + """ + if query.numel() == 0 or fitted.numel() == 0: + return torch.full(query.shape, -1, dtype=dtype, device=query.device) + query_count = 1 if query.stride(0) == 0 else query.numel() + fitted_count = 1 if fitted.stride(0) == 0 else fitted.numel() + query_data, fitted_data = query._data, fitted._data + capacity = query_data.numel() + fitted_data.numel() + key_words = max( + 1, + (capacity + 8 * (query_count + fitted_count) - 1) + // (8 * (query_count + fitted_count)), + ) + if query._max_length is not None and fitted._max_length is not None: + key_words = min( + key_words, + max(1, (max(query._max_length, fitted._max_length) + 7) // 8), + ) + extent = ( + capacity + + 8 * key_words + + query._offset.numel() + + fitted._offset.numel() + ) + index_dtype = ( + torch.int32 if extent < torch.iinfo(torch.int32).max else torch.int64 + ) + query_starts, query_lengths = _string_spans(query, index_dtype) + fitted_starts, fitted_lengths = _string_spans(fitted, index_dtype) + complete = ( + query._max_length is not None and query._max_length <= key_words * 8 + ) + if ( + torch.compiler.is_compiling() + and not complete + and query_data is not fitted_data + ): + # Higher-order tracing requires independent backing buffers. + fitted_data = fitted_data.clone() + query_rows = _StringRows(query_data, query_starts, query_lengths) + fitted_rows = _StringRows(fitted_data, fitted_starts, fitted_lengths) + + if complete: + output = _match_string_keys( + query=query_rows, + fitted=fitted_rows, + key_words=key_words, + complete=True, + dtype=dtype, + ) + else: + max_query_bytes = query_lengths.amax() + power_words = 1 + while power_words * 2 <= key_words: + power_words *= 2 + output = _select_branch( + predicate=max_query_bytes <= power_words * 8, + when_true=partial( + _match_smaller_keys, + query=query_rows, + fitted=fitted_rows, + max_query_bytes=max_query_bytes, + key_words=power_words, + dtype=dtype, + ), + when_false=partial( + _match_string_keys, + query=query_rows, + fitted=fitted_rows, + key_words=key_words, + complete=False, + dtype=dtype, + ), + ) + + if fitted.stride(0) == 0: + output = torch.where(output >= 0, fitted.numel() - 1, output) + return output.expand(query.numel()) if query.stride(0) == 0 else output + + +def _string_spans( + value: StringTensor, index_dtype: torch.dtype +) -> tuple[Tensor, Tensor]: + """Gather byte starts and lengths, representing broadcast views once.""" + size = 1 if value.numel() and value.stride(0) == 0 else value.numel() + rows = torch.arange(size, dtype=index_dtype, device=value.device) + rows = rows * value.stride(0) + value._logical_storage_offset + starts = value._offset[rows].to(index_dtype) + return starts, value._offset[rows + 1].to(index_dtype) - starts + + +def _match_smaller_keys( + query: _StringRows, + fitted: _StringRows, + max_query_bytes: Tensor, + key_words: int, + dtype: torch.dtype, +) -> Tensor: + """Match using the smallest key workspace that covers visible queries. + + Halving the workspace avoids sizing keys for unused bytes retained by a + view. The caller guarantees that these queries fit in the initial width. + """ + if key_words == 1: + return _match_string_keys(query, fitted, 1, complete=True, dtype=dtype) + return _select_branch( + predicate=max_query_bytes <= 8 * (key_words // 2), + when_true=partial( + _match_smaller_keys, + query=query, + fitted=fitted, + max_query_bytes=max_query_bytes, + key_words=key_words // 2, + dtype=dtype, + ), + when_false=partial( + _match_string_keys, + query=query, + fitted=fitted, + key_words=key_words, + complete=True, + dtype=dtype, + ), + ) + + +def _match_string_keys( + query: _StringRows, + fitted: _StringRows, + key_words: int, + complete: bool, + dtype: torch.dtype, +) -> Tensor: + """Pack keys, find candidate hash buckets, and resolve exact matches. + + ``complete`` means every query fits in its key. Otherwise a device-side + length check decides whether full-byte verification is needed. + """ + query_count, fitted_count = query.lengths.numel(), fitted.lengths.numel() + # Flat buffers keep a fixed layout across higher-order control flow. + query_keys = _pack_string_keys( + data=query.data, + starts=query.starts, + lengths=query.lengths, + key_words=key_words, + ).flatten() + fitted_keys = _pack_string_keys( + data=fitted.data, + starts=fitted.starts, + lengths=fitted.lengths, + key_words=key_words, + ).flatten() + columns = key_words + 1 + weights = torch.arange( + start=1, + end=2 * columns, + step=2, + dtype=torch.int64, + device=query.data.device, + ) + # Hashes filter candidates; packed keys and bytes establish equality. + query_hashes = (query_keys.view(query_count, columns) * weights).sum(-1) + sorted_hashes, fitted_order = ( + (fitted_keys.view(fitted_count, columns) * weights) + .sum(-1) + .sort(stable=True) + ) + del weights + keys = _LookupKeys( + query=query_keys, + fitted=fitted_keys, + query_hashes=query_hashes, + sorted_hashes=sorted_hashes, + fitted_order=fitted_order, + ) + resolve = partial( + _resolve_candidates, + query=query, + fitted=fitted, + keys=keys, + key_words=key_words, + dtype=dtype, + ) + if complete: + return resolve(check_bytes=False) + return _select_branch( + predicate=(query.lengths <= key_words * 8).all(), + when_true=partial(resolve, check_bytes=False), + when_false=partial(resolve, check_bytes=True), + ) + + +def _pack_string_keys( + data: Tensor, starts: Tensor, lengths: Tensor, key_words: int +) -> Tensor: + """Pack lookup keys into ``[N, key_words + 1]`` integers. + + ``N`` is the number of strings. Short strings fit completely; long + strings contribute their beginning and end. The final length column + distinguishes NUL bytes from padding. Complete keys prove equality; + sampled long keys require full-byte verification. + """ + keys = torch.zeros( + size=(lengths.numel(), key_words), + dtype=torch.int64, + device=data.device, + ) + if data.numel() == 0: + return torch.cat((keys, lengths[:, None]), dim=1) + position = 8 * torch.arange( + end=key_words, dtype=starts.dtype, device=data.device + ) + index = starts[:, None] + position + shift = (lengths - 8 * key_words).clamp_min(0) + index[:, (key_words + 1) // 2 :].add_(shift[:, None]) + for byte in range(8): + # An odd middle word joins four prefix and four suffix bytes. + if byte == 4 and key_words % 2: + index[:, key_words // 2].add_(shift) + present = position + byte < lengths[:, None] + index.clamp_max_(data.numel() - 1) + symbol = data[index].masked_fill_(~present, 0) + # Signed overflow preserves the raw 64-bit pattern. + keys.add_(symbol, alpha=1 << (8 * byte)) + index.add_(1) + return torch.cat((keys, lengths[:, None]), dim=1) + + +def _resolve_candidates( + query: _StringRows, + fitted: _StringRows, + keys: _LookupKeys, + key_words: int, + dtype: torch.dtype, + check_bytes: bool, +) -> Tensor: + """Find exact fitted IDs, returning ``-1`` for missing query labels. + + Equal-hash candidates are visited from right to left, so the last fitted + duplicate wins. Packed keys are always checked; sampled long keys also + require a full-byte comparison. + """ + position = ( + torch.searchsorted(keys.sorted_hashes, keys.query_hashes, right=True) + - 1 + ) + active = (position >= 0) & ( + keys.sorted_hashes[position.clamp_min(0)] == keys.query_hashes + ) + output = torch.full( + size=query.lengths.shape, + fill_value=-1, + dtype=dtype, + device=query.data.device, + ) + advance = partial( + _advance_candidates, + query=query, + fitted=fitted, + keys=keys, + key_words=key_words, + check_bytes=check_bytes, + ) + if not check_bytes: + position, output, active = advance(position, output, active) + if torch.compiler.is_compiling(): + _, result, _ = torch.while_loop( + cond_fn=_has_pending_candidates, + body_fn=advance, + carried_inputs=(position, output, active), + ) + else: + while bool(_has_pending_candidates(position, output, active)): + position, output, active = advance(position, output, active) + result = output + return result + + +def _advance_candidates( + position: Tensor, + output: Tensor, + active: Tensor, + *, + query: _StringRows, + fitted: _StringRows, + keys: _LookupKeys, + key_words: int, + check_bytes: bool, +) -> tuple[Tensor, Tensor, Tensor]: + """Test current candidates and step backward for unresolved queries.""" + query_count, fitted_count = query.lengths.numel(), fitted.lengths.numel() + columns = key_words + 1 + candidate = keys.fitted_order[position.clamp_min(0)] + matched = active & ( + keys.query.view(query_count, columns) + == keys.fitted.view(fitted_count, columns)[candidate] + ).all(-1) + if check_bytes: + lengths = query.lengths.where( + condition=matched & (query.lengths > 8 * key_words), other=0 + ) + different = _candidate_bytes_differ( + query_data=query.data, + query_starts=query.starts, + fitted_data=fitted.data, + fitted_starts=fitted.starts[candidate], + lengths=lengths, + ) + matched = matched & ~different + next_output = torch.where(matched, candidate.to(output.dtype), output) + next_position = position - (active & ~matched).to(position.dtype) + next_active = active & ~matched & (next_position >= 0) + next_active = next_active & ( + keys.sorted_hashes[next_position.clamp_min(0)] == keys.query_hashes + ) + return next_position, next_output, next_active + + +def _candidate_bytes_differ( + query_data: Tensor, + query_starts: Tensor, + fitted_data: Tensor, + fitted_starts: Tensor, + lengths: Tensor, +) -> Tensor: + """Return a mismatch mask after comparing every requested candidate byte. + + ``lengths`` has shape ``[Q]``, where ``Q`` is the number of query labels. + Zero skips a query row. Eight chunks cover the complete comparison, + keeping temporary storage below a full byte grid. + """ + capacity = query_data.numel() + if capacity == 0 or fitted_data.numel() == 0: + return lengths != 0 + comparison_offsets = torch.cat( + tensors=(lengths.new_zeros(1), lengths.cumsum(0, dtype=lengths.dtype)) + ) + query_starts = torch.cat((query_starts, query_starts.new_zeros(1))) + fitted_starts = torch.cat((fitted_starts, fitted_starts.new_zeros(1))) + mismatches = torch.zeros( + size=(lengths.numel() + 1,), + dtype=lengths.dtype, + device=query_data.device, + ) + chunk_size = (capacity + 7) // 8 + byte_positions = torch.arange( + end=chunk_size, dtype=lengths.dtype, device=query_data.device + ) + compare_chunk = partial( + _compare_byte_chunk, + query_data=query_data, + query_starts=query_starts, + fitted_data=fitted_data, + fitted_starts=fitted_starts, + comparison_offsets=comparison_offsets, + byte_positions=byte_positions, + ) + if torch.compiler.is_compiling(): + # This fixed-trip counter is independent of CUDA values. + counter = torch.zeros((), dtype=lengths.dtype, device="cpu") + _, mismatches = torch.while_loop( + cond_fn=_has_byte_chunk, + body_fn=compare_chunk, + carried_inputs=(counter, mismatches), + ) + else: + for chunk in range(8): + _, mismatches = compare_chunk(chunk, mismatches) + return mismatches[:-1] != 0 + + +def _compare_byte_chunk( + chunk: Tensor | int, + mismatches: Tensor, + *, + query_data: Tensor, + query_starts: Tensor, + fitted_data: Tensor, + fitted_starts: Tensor, + comparison_offsets: Tensor, + byte_positions: Tensor, +) -> tuple[Tensor | int, Tensor]: + """Accumulate per-row mismatch counts for one chunk of candidate bytes.""" + byte = byte_positions + chunk * byte_positions.numel() + row = ( + torch.searchsorted( + sorted_sequence=comparison_offsets, + input=byte, + right=True, + out_int32=comparison_offsets.dtype == torch.int32, + ) + - 1 + ) + offset_in_row = byte - comparison_offsets[row] + valid_row = row < comparison_offsets.numel() - 1 + query_index = (query_starts[row] + offset_in_row).where(valid_row, 0) + fitted_index = (fitted_starts[row] + offset_in_row).where(valid_row, 0) + differs = ( + (query_data[query_index] != fitted_data[fitted_index]) & valid_row + ).to(comparison_offsets.dtype) + return chunk + 1, mismatches.index_add(0, row, differs) + + +def _select_branch( + predicate: Tensor, + when_true: Callable[[], Tensor], + when_false: Callable[[], Tensor], +) -> Tensor: + """Select a tensor-dependent branch without compiling eager execution.""" + if torch.compiler.is_compiling(): + return torch.cond(predicate, when_true, when_false, ()) + # Structured ops invoke Dynamo in eager; use ordinary control flow there. + return when_true() if bool(predicate) else when_false() + + +def _has_pending_candidates( + position: Tensor, output: Tensor, active: Tensor +) -> Tensor: + """Continue while any query has candidates left to examine.""" + return active.any() + + +def _has_byte_chunk(chunk: Tensor, mismatches: Tensor) -> Tensor: + """Continue until all eight byte-comparison chunks have been read.""" + return chunk < 8 diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 3a4cbcf924..976218baa3 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -17,6 +17,7 @@ ) from sdm.nn._buffer import BufferList from sdm.processing import EnsembleProcessor +from sdm.processing.categorical._string_lookup import _native_string_lookups from sdm.relational.join import join_index from sdm.tensor.categorical import _iter_categories from sdm.tensor.var_len import _cat, _contiguous @@ -340,15 +341,17 @@ def _string_category_lookups( ) -> tuple[Tensor, ...]: if not input_categories: return () - # 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) + if codes.is_cuda: + return _native_string_lookups( + left=input_categories, + right=cast(tuple[StringTensor, ...], fitted_categories), + dtype=codes.dtype, + ) 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( @@ -501,45 +504,7 @@ def __repr__(self, *, indent: int = 0) -> str: ) -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 @@ -547,17 +512,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( [ @@ -606,16 +565,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/sdm/tensor/string.py b/sdm/tensor/string.py index 6d3d4e7ac0..56353d248b 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,8 @@ def from_arrow( valid=valid, size=size, storage_offset=storage_offset, + # Inspect CPU offsets before moving buffers. + max_length=_length_bound(offset), ) @override @@ -203,6 +210,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 +236,8 @@ 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), ) def to_cudf(self) -> cudf.Series: @@ -546,3 +556,19 @@ 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. The bound covers all strings in ``offset``, + so views of the same storage stay within it. + """ + 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/var_len.py b/sdm/tensor/var_len.py index a3c439b9c4..5fc3decae9 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 # Constructors ############################################################ @@ -82,6 +86,7 @@ def __init__( size: Sequence[int], stride: Sequence[int] | None = None, storage_offset: int = 0, + max_length: int | None = None, ) -> None: pass @@ -93,6 +98,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) @@ -220,6 +226,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, ) @@ -233,6 +240,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 @@ -248,6 +256,7 @@ def _make_var_len_tensor( out._data = data out._offset = offset out._valid = valid + out._max_length = max_length return out @classmethod @@ -563,6 +572,7 @@ def __tensor_flatten__(self) -> tuple[list[str], tuple[Any, ...]]: self.storage_offset(), tuple(self.size()), tuple(self.stride()), + self._max_length, ) return attrs, ctx @@ -573,7 +583,7 @@ def __tensor_unflatten__( outer_size: tuple[int, ...], outer_stride: tuple[int, ...], ) -> VarLenTensor: - cls, storage_offset, _, _ = ctx + cls, storage_offset, _, _, max_length = ctx return cls( data=inner_tensors["_data"], offset=inner_tensors["_offset"], @@ -581,6 +591,7 @@ def __tensor_unflatten__( size=outer_size, stride=outer_stride, storage_offset=storage_offset, + max_length=max_length, ) def __reduce_ex__(self, proto: SupportsIndex) -> Any: @@ -591,6 +602,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) @@ -732,6 +744,7 @@ def _alias(inp: VarLenTensor) -> VarLenTensor: size=inp.size(), stride=inp.stride(), storage_offset=int(inp.storage_offset()), + max_length=inp._max_length, ) @@ -841,6 +854,7 @@ def _to_dtype_layout( if memory_format == torch.preserve_format else None, storage_offset=0, + max_length=inp._max_length, ) @@ -950,6 +964,7 @@ def _detach(inp: VarLenTensor) -> VarLenTensor: size=inp.size(), stride=inp.stride(), storage_offset=int(inp.storage_offset()), + max_length=inp._max_length, ) @@ -968,6 +983,7 @@ def _pin_memory(inp: VarLenTensor) -> VarLenTensor: size=inp.size(), stride=inp.stride(), storage_offset=int(inp.storage_offset()), + max_length=inp._max_length, ) @@ -1362,6 +1378,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): @@ -1403,7 +1423,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,)) @@ -1411,7 +1437,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) @@ -1461,6 +1493,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, ) @@ -1550,4 +1583,5 @@ def _materialize( size=size, stride=stride, storage_offset=0, + max_length=inp._max_length, ) diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index 52cf0a6b95..95b886456f 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -1,7 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Literal +import pickle +from typing import Literal, cast import pandas as pd import pytest @@ -9,7 +10,8 @@ from sdm import CategoricalTensor, EnsembleTable, StringTensor, TableTensor from sdm.processing import AlignCategories -from sdm.testing import withCUDA +from sdm.processing.categorical._string_lookup import _native_string_lookups +from sdm.testing import onlyCUDA, withCUDA def _table( @@ -461,14 +463,16 @@ def test_align_categories_rejects_changed_category_value_type() -> None: def test_align_categories_ensemble_matches_member_fits( device: torch.device, ) -> None: + prefix = "a" * 85 + 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 +482,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( @@ -506,12 +510,80 @@ 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", [8, 300]) +def test_native_string_lookup(count: int) -> None: + fitted = [f"é\0common-prefix-{index:04}" for index in range(count)] + [ + "a", + "a\0", + "", + "1234567é", + "1234567ê", + ] + query = [*fitted[::-1], "unknown", "a\0\0"] + indices = {value: index for index, value in enumerate(fitted)} + actual = _native_string_lookups( + left=(StringTensor.from_list(query),), + right=(StringTensor.from_list(fitted),), + dtype=torch.int32, + ) + torch.testing.assert_close( + actual=actual[0], + expected=torch.tensor( + [indices.get(value, -1) for value in query], dtype=torch.int32 + ), + ) + + +@withCUDA +def test_native_string_lookups_collisions_and_shared_views( + device: torch.device, +) -> None: + # These distinct strings have the same fingerprint; bytes decide equality. + first, second = "aaaaaaaadddddddd", "ddddddddcccccccc" + values = StringTensor.from_list( + [first, second, first, second, "", "unknown"], device=device + ) + actual = _native_string_lookups( + left=cast(tuple[StringTensor, ...], (values[::2], values[1::2])), + right=cast(tuple[StringTensor, ...], (values[:4], values[:3])), + dtype=torch.int32, + ) + torch.testing.assert_close( + actual=actual, + expected=( + torch.tensor([2, 2, -1], dtype=torch.int32, device=device), + torch.tensor([1, 1, -1], dtype=torch.int32, device=device), + ), + ) + + +@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, +) -> None: + device = torch.device("cuda") + prefix = "é" * 56 + extra = [f"{prefix}long-prefix-{index:04}" for index in range(count)] fitted = ( - StringTensor.from_list(["a", "b"]), - StringTensor.from_list(["a", "x"]), + 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) + fitted = pickle.loads(pickle.dumps(fitted)) + 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( @@ -520,23 +592,158 @@ 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"])) + compiled = torch.compile( + lookups, backend="inductor", fullgraph=True, dynamic=dynamic + ) torch.testing.assert_close( - actual=first, - expected=torch.tensor([1, -1, 0], dtype=codes.dtype), + 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], dtype=codes.dtype), + 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)]), + ), ) - first, second = compiled(StringTensor.from_list(["x", "z", "q"])) - torch.testing.assert_close( - actual=first, - expected=torch.tensor([-1, -1, -1], dtype=codes.dtype), + +@onlyCUDA +def test_string_category_lookups_compile_ragged(compilation: None) -> None: + device = torch.device("cuda") + first = "a" * 1024 + "x" + "b" * 1024 + second = "a" * 1024 + "y" + "b" * 1024 + unknown = "a" * 1024 + "z" + "b" * 1024 + fitted = StringTensor.from_list( + [*[f"label-{index}" for index in range(32)], first, second], + device=device, ) - torch.testing.assert_close( - actual=second, - expected=torch.tensor([1, -1, -1], dtype=codes.dtype), + codes = torch.empty((), dtype=torch.int32, device=device) + + def lookups(categories: StringTensor) -> tuple[torch.Tensor, ...]: + return AlignCategories._string_category_lookups( + input_categories=(categories,), + fitted_categories=(fitted,), + codes=codes, + ) + + compiled = torch.compile(lookups, backend="inductor", fullgraph=True) + for values, indices in ( + ([first, unknown, second], [32, -1, 33]), + ([second, first, unknown], [33, 32, -1]), + ): + query = StringTensor.from_list(values, device=device) + torch.testing.assert_close( + actual=compiled(query), expected=(codes.new_tensor(indices),) + ) + + +@withCUDA +@pytest.mark.parametrize("count", [8, 300]) +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) + ] + ["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", [84, 127, 255, 1023]) +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] + + +@withCUDA +def test_align_categories_mixed_category_types( + device: torch.device, +) -> None: + context = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor( + [[0, 1], [1, 0], [0, 1]], dtype=torch.int32, device=device + ), + categories=( + StringTensor.from_list(["red", "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( + 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), + ), + ), + ) + + output = AlignCategories(sort_by="frequency").fit(context).transform(query) + + assert output.categorical.code.tolist() == [ + [1, 0], + [-1, -1], + [0, 1], + [-1, -1], + ] diff --git a/test/tensor/test_string.py b/test/tensor/test_string.py index 3311c440c1..42ba2a9c8c 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 every stored string, 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 retains the longest string's backing storage. + array = pa.array(texts) + if bound == 15: + array = pa.array(["x" * 12, *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]