Compile categorical string lookup with native PyTorch operations

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,15 +342,26 @@ 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 and all(
+            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
+            )
+        ):
+            return tuple(
+                _native_string_lookup(
+                    left, cast(StringTensor, right), codes.dtype
+                )
+                for left, right in zip(
+                    input_categories, fitted_categories, strict=True
+                )
+            )
         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,63 +514,56 @@ 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,
-            )
+def _native_string_lookup(
+    left: StringTensor, right: StringTensor, dtype: torch.dtype
+) -> Tensor:
+    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)
+    # 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)
+        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),
         )
-    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],
+    )
+    groups = starts.long().cumsum(0) - 1
+    values = torch.cat(
+        (
+            right_ids.new_full((left.numel(),), -1),
+            right_ids,
         )
-        for i in range(len(layout) // 3)
     )
+    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 _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
     )
     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 +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]),))