Support compiled CUDA stream recording for SDM tensors

sdm/tensor/categorical.py

@@ -21,6 +21,7 @@
     to_arrow,
     to_cudf,
 )
+from sdm.tensor.mixin import _record_compiled_stream
 
 if TYPE_CHECKING:
     import cudf
@@ -463,6 +464,9 @@ def __torch_dispatch__(  # type: ignore
         if (handler := cls.HANDLED_FUNCTIONS.get(func)) is not None:
             return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_compiled_stream(*args, **(kwargs or {}))
+
         # Operate on vanilla tensors for all non-handled functions:
         args = pytree.tree_map_only(CategoricalTensor, lambda x: x._code, args)
         kwargs = pytree.tree_map_only(

sdm/tensor/columnar.py

@@ -16,7 +16,7 @@
 
 from sdm.tensor import NullableTensor, StringTensor, VarLenTensor
 from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf
-from sdm.tensor.mixin import _resolve_device
+from sdm.tensor.mixin import _record_compiled_stream, _resolve_device
 
 if TYPE_CHECKING:
     import cudf
@@ -360,6 +360,9 @@ def __torch_dispatch__(  # type: ignore
         if (handler := cls.HANDLED_FUNCTIONS.get(func)) is not None:
             return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_compiled_stream(*args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )

sdm/tensor/mixin.py

@@ -3,7 +3,7 @@
 
 import abc
 from collections.abc import Callable, Iterator
-from typing import Self
+from typing import Any, Self
 
 import torch
 from torch import Tensor
@@ -122,3 +122,11 @@ def _resolve_device(device: torch.device | str | None) -> torch.device | None:
     if device.type == "mps" and device.index is None:
         return torch.device("mps", 0)
     return device
+
+
+def _record_compiled_stream(tensor: Any, stream_index: int) -> None:
+    """Keep each stored tensor's memory alive while the stream uses it."""
+    tensor_names, _ = tensor.__tensor_flatten__()
+    for tensor_name in tensor_names:
+        stored_tensor = getattr(tensor, tensor_name)
+        torch.ops.streams.record_stream.default(stored_tensor, stream_index)

sdm/tensor/nullable.py

@@ -20,6 +20,7 @@
     to_arrow,
     to_cudf,
 )
+from sdm.tensor.mixin import _record_compiled_stream
 
 if TYPE_CHECKING:
     import cudf
@@ -385,6 +386,9 @@ def __torch_dispatch__(  # type: ignore
         if (handler := cls.HANDLED_FUNCTIONS.get(func)) is not None:
             return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_compiled_stream(*args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )

sdm/tensor/table.py

@@ -21,7 +21,7 @@
 from sdm import NaT, Stype, StypeLike
 from sdm.tensor import CategoricalTensor, ColumnarTensor, StringTensor
 from sdm.tensor.io import arrow_as_tensor, to_arrow, to_cudf
-from sdm.tensor.mixin import _resolve_device
+from sdm.tensor.mixin import _record_compiled_stream, _resolve_device
 
 if TYPE_CHECKING:
     import cudf
@@ -1201,6 +1201,9 @@ def __torch_dispatch__(  # type: ignore
         if func in cls.HANDLED_FUNCTIONS:
             return cls.HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_compiled_stream(*args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )

sdm/tensor/var_len.py

@@ -21,6 +21,7 @@
     combine_arrow_chunks,
     to_arrow,
 )
+from sdm.tensor.mixin import _record_compiled_stream
 
 aten = torch.ops.aten
 
@@ -632,6 +633,9 @@ def __torch_dispatch__(  # type: ignore
             with enable_reentrant_dispatch():  # Record autograd in `_data`.
                 return handler(*args, **(kwargs or {}))
 
+        if str(func) == "streams.record_stream.default":
+            return _record_compiled_stream(*args, **(kwargs or {}))
+
         raise NotImplementedError(
             f"'{func}' is not supported for {cls.__name__!r}"
         )