Skip to content

Commit 82ad598

Browse files
leofangemcastilloclaude
committed
Add torch.Tensor fast path for StridedMemoryView via AOTI tensor bridge
Provide a fast path for constructing a StridedMemoryView from a torch.Tensor by reading tensor metadata directly through PyTorch's AOT Inductor (AOTI) stable C ABI, avoiding DLPack/CAI protocol overhead (~10 ns per tensor via pointer arithmetic). Key design: - Vendored AOTI shim header (aoti_shim.h) with extern "C" wrapping - _tensor_bridge.pyx loaded lazily (only when a torch.Tensor is first passed) to avoid undefined AOTI symbols at import time - RTLD_GLOBAL bootstrap via sys.modules["torch._C"] before loading _tensor_bridge.so - torch detection via type(obj).__module__.startswith("torch") - PyTorch is NOT a build-time or run-time dependency of cuda.core Closes NVIDIA#749 Co-Authored-By: Emilio Castillo <[email protected]> Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
1 parent c93623b commit 82ad598

5 files changed

Lines changed: 514 additions & 6 deletions

File tree

‎cuda_core/cuda/core/_include/aoti_shim.h‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES.
3+
* All rights reserved.
4+
* SPDX-License-Identifier: Apache-2.0
5+
*
6+
* Vendored subset of PyTorch's AOT Inductor (AOTI) stable C ABI.
7+
* Original: torch/csrc/inductor/aoti_torch/c/shim.h
8+
*
9+
* These are declarations only -- no definitions are provided. The actual
10+
* symbols are exported by libtorch (loaded via torch._C with RTLD_GLOBAL)
11+
* and resolved at runtime by the dynamic linker. This means PyTorch is
12+
* NOT required at compile time.
13+
*/
14+
15+
#ifndef CUDA_CORE_AOTI_SHIM_H
16+
#define CUDA_CORE_AOTI_SHIM_H
17+
18+
#include <stdint.h>
19+
20+
#ifdef __cplusplus
21+
extern "C" {
22+
#endif
23+
24+
typedef int32_t AOTITorchError;
25+
26+
/* Opaque tensor handle -- corresponds to at::Tensor on the C++ side. */
27+
struct AtenTensorOpaque;
28+
typedef struct AtenTensorOpaque* AtenTensorHandle;
29+
30+
/* ---- tensor metadata --------------------------------------------------- */
31+
32+
AOTITorchError aoti_torch_get_data_ptr(
33+
AtenTensorHandle tensor, void** ret_data_ptr);
34+
35+
AOTITorchError aoti_torch_get_dim(
36+
AtenTensorHandle tensor, int64_t* ret_dim);
37+
38+
AOTITorchError aoti_torch_get_numel(
39+
AtenTensorHandle tensor, int64_t* ret_numel);
40+
41+
AOTITorchError aoti_torch_get_sizes(
42+
AtenTensorHandle tensor, int64_t** ret_sizes);
43+
44+
AOTITorchError aoti_torch_get_strides(
45+
AtenTensorHandle tensor, int64_t** ret_strides);
46+
47+
AOTITorchError aoti_torch_get_storage_offset(
48+
AtenTensorHandle tensor, int64_t* ret_storage_offset);
49+
50+
/* ---- dtype ------------------------------------------------------------- */
51+
52+
AOTITorchError aoti_torch_get_dtype(
53+
AtenTensorHandle tensor, int32_t* ret_dtype);
54+
55+
int32_t aoti_torch_dtype_float16(void);
56+
int32_t aoti_torch_dtype_float32(void);
57+
int32_t aoti_torch_dtype_float64(void);
58+
int32_t aoti_torch_dtype_bfloat16(void);
59+
int32_t aoti_torch_dtype_uint8(void);
60+
int32_t aoti_torch_dtype_int8(void);
61+
int32_t aoti_torch_dtype_int16(void);
62+
int32_t aoti_torch_dtype_int32(void);
63+
int32_t aoti_torch_dtype_int64(void);
64+
int32_t aoti_torch_dtype_bool(void);
65+
int32_t aoti_torch_dtype_complex32(void);
66+
int32_t aoti_torch_dtype_complex64(void);
67+
int32_t aoti_torch_dtype_complex128(void);
68+
69+
/* ---- device ------------------------------------------------------------ */
70+
71+
AOTITorchError aoti_torch_get_device_type(
72+
AtenTensorHandle tensor, int32_t* ret_device_type);
73+
74+
AOTITorchError aoti_torch_get_device_index(
75+
AtenTensorHandle tensor, int32_t* ret_device_index);
76+
77+
int32_t aoti_torch_device_type_cpu(void);
78+
int32_t aoti_torch_device_type_cuda(void);
79+
80+
#ifdef __cplusplus
81+
} /* extern "C" */
82+
#endif
83+
84+
#endif /* CUDA_CORE_AOTI_SHIM_H */

‎cuda_core/cuda/core/_memoryview.pyx‎

Lines changed: 58 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,35 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
2929
from cuda.core._memory import Buffer
3030

3131

32+
# ---------------------------------------------------------------------------
33+
# Lazy tensor bridge (avoids loading _tensor_bridge.so until torch is used)
34+
# ---------------------------------------------------------------------------
35+
36+
cdef object _tensor_bridge = None
37+
38+
39+
cdef inline bint _is_torch_tensor(object obj):
40+
cdef str mod = type(obj).__module__ or ""
41+
return mod.startswith("torch") and hasattr(obj, "data_ptr")
42+
43+
44+
cdef object _get_tensor_bridge():
45+
"""Bootstrap AOTI symbols, then import _tensor_bridge on first use."""
46+
global _tensor_bridge
47+
if _tensor_bridge is not None:
48+
return _tensor_bridge
49+
import ctypes, sys
50+
torch_C = sys.modules.get("torch._C")
51+
if torch_C is None:
52+
raise RuntimeError(
53+
"torch._C is not loaded; cannot initialise the tensor bridge. "
54+
"Make sure PyTorch is imported before passing a torch.Tensor.")
55+
ctypes.CDLL(torch_C.__file__, mode=ctypes.RTLD_GLOBAL)
56+
from cuda.core import _tensor_bridge as tb
57+
_tensor_bridge = tb
58+
return _tensor_bridge
59+
60+
3261
try:
3362
from ml_dtypes import bfloat16
3463
except ImportError:
@@ -112,7 +141,15 @@ cdef class StridedMemoryView:
112141
cdef str clsname = self.__class__.__name__
113142
if obj is not None:
114143
# populate self's attributes
115-
if check_has_dlpack(obj):
144+
if _is_torch_tensor(obj):
145+
warnings.warn(
146+
f"Constructing a {clsname} directly from a torch.Tensor is deprecated; "
147+
"Use `StridedMemoryView.from_any_interface` instead.",
148+
DeprecationWarning,
149+
stacklevel=2,
150+
)
151+
_get_tensor_bridge().view_as_torch_tensor(obj, stream_ptr, self)
152+
elif check_has_dlpack(obj):
116153
warnings.warn(
117154
f"Constructing a {clsname} directly from a DLPack-supporting object is deprecated; "
118155
"Use `StridedMemoryView.from_dlpack` or `StridedMemoryView.from_any_interface` instead.",
@@ -185,17 +222,24 @@ cdef class StridedMemoryView:
185222
def from_any_interface(cls, obj: object, stream_ptr: int | None = None) -> StridedMemoryView:
186223
"""Create a view by automatically selecting the best available protocol.
187224

188-
Tries `DLPack <https://dmlc.github.io/dlpack/latest/>`_ first, then falls back to
225+
For ``torch.Tensor`` objects a fast path via the AOTI stable C ABI is
226+
used (no DLPack/CAI overhead). Otherwise tries
227+
`DLPack <https://dmlc.github.io/dlpack/latest/>`_ first, then falls back to
189228
`__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_.
190229

191230
Parameters
192231
----------
193232
obj : object
194-
An object implementing `DLPack <https://dmlc.github.io/dlpack/latest/>`_ or
195-
`__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_.
233+
An object implementing `DLPack <https://dmlc.github.io/dlpack/latest/>`_,
234+
`__cuda_array_interface__ <https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html>`_,
235+
or a ``torch.Tensor``.
196236
stream_ptr : int, optional
197237
Stream pointer for synchronization. If ``None``, no synchronization is performed.
198238
"""
239+
if _is_torch_tensor(obj):
240+
buf = StridedMemoryView.__new__(cls)
241+
_get_tensor_bridge().view_as_torch_tensor(obj, stream_ptr, buf)
242+
return buf
199243
if check_has_dlpack(obj):
200244
return cls.from_dlpack(obj, stream_ptr)
201245
return cls.from_cuda_array_interface(obj, stream_ptr)
@@ -921,13 +965,21 @@ cdef class _StridedMemoryViewProxy:
921965
cdef readonly:
922966
object obj
923967
bint has_dlpack
968+
cdef:
969+
bint _is_torch
924970

925971
def __init__(self, obj):
926972
self.obj = obj
927-
self.has_dlpack = check_has_dlpack(obj)
973+
self._is_torch = _is_torch_tensor(obj)
974+
if not self._is_torch:
975+
self.has_dlpack = check_has_dlpack(obj)
976+
else:
977+
self.has_dlpack = False
928978

929979
cpdef StridedMemoryView view(self, stream_ptr=None):
930-
if self.has_dlpack:
980+
if self._is_torch:
981+
return _get_tensor_bridge().view_as_torch_tensor(self.obj, stream_ptr)
982+
elif self.has_dlpack:
931983
return StridedMemoryView.from_dlpack(self.obj, stream_ptr)
932984
else:
933985
return StridedMemoryView.from_cuda_array_interface(self.obj, stream_ptr)

‎cuda_core/cuda/core/_tensor_bridge.pxd‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0

0 commit comments

Comments
 (0)