Repository navigation
cuda.core: use torch 2.14+ stable PyObject->Tensor shim, drop upper version cap - #3033
Conversation
PyTorch >= 2.14 exposes torch_tensor_from_pyobject / aoti_torch_delete_tensor_object as part of the stable C ABI (landed via pytorch/pytorch#183323). These let us obtain an AtenTensorHandle from a Python tensor object without peering at THPVariable's internal layout. Dispatch is driven by *symbol presence* via dlsym / GetProcAddress on first call to view_as_torch_tensor: - resolved -> primary path (owned handle, deleted after metadata reads) - NULL -> fall back to the existing pyobj_to_aten_handle pointer-arithmetic trick, which stays correct for torch 2.3-2.13 (the only versions where this matters). The outer version gate in _memoryview.pyx (``2.3 <= torch <= 2.14``) is unchanged -- it still bounds the fallback. Once we drop support for torch < 2.14 we can remove the upper bound entirely because the shim itself is part of torch's stable ABI contract. Diagnostics kept module-private: - _get_pyobj_path_counts() / _reset_pyobj_path_counts() - _get_shim_available() - _set_use_shim(bool) to force the fallback for A/B benchmarking.
…bj-shim # Conflicts: # cuda_core/cuda/core/_tensor_bridge.pyx
Keeps the 2.3 lower bound (required for the AOTI metadata functions).
The dispatch in _tensor_bridge.view_as_torch_tensor now picks between:
* torch >= 2.14 -- stable torch_tensor_from_pyobject shim
(pytorch/pytorch#183323; part of torch's stable C ABI contract)
* torch 2.3 - 2.13 -- pyobj_to_aten_handle pointer-arithmetic,
validated against the two known THPVariable layouts
Future torch releases >= 2.14 are automatically covered by the shim,
removing the per-release cap bump treadmill.
|
/ok to test 87cadd6 |
@leofang, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
|
/ok to test 87cadd6 |
leofang
changed the title
📝 Summary
Priority: Change: Feature Merge Risk:
 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
-
-
- Source mode: Auto
- Repository: NVIDIA/cuda-python
- Branch: main
- File: .coderabbit.yaml
- Commit: 03b4153
- Layers: Repository YAML
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
60b92c43-707d-4225-8f9d-8f7e8fac2577
-
cuda_core/cuda/core/_memoryview.pyxcuda_core/cuda/core/_tensor_bridge.pyicuda_core/cuda/core/_tensor_bridge.pyx
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
-
|
rwgk
left a comment
rwgk
left a comment
There was a problem hiding this comment.
The finding is only a concern for free-threading. Do you want to fix the race situation in this PR?
codex GPT-6.1-Sol ultra:
One finding at 87cadd6:
[P2] Synchronize the shim cache initialization — _tensor_bridge.pyx:409–412. Concurrent first conversions on free-threaded Python read and write three unsynchronized C globals, creating a C++ data race. PyTorch also declares no-GIL support. Resolve both pointers during module initialization, when torch symbols are already available, or use synchronized initialization.
The CI matrix should also include torch ≥2.14 to exercise the new shim automatically.
The concurrency finding comes from source/generated C++ inspection; no crash was reproduced.
|
It's preexisting so I'll fix it in the next PR. I want to unblock the nightly update (#3031). |
|
Thanks, Ralf! |
PyTorch 2.15.0 ships only as +cu132 on pytorch.org (not cu126/cu130). Replace the two 2.14.1 rows (cu126 + cu130) per arch with a single 2.15.0+cu132 row. The 2.9.1 rows (cu126 + cu130) stay as the older anchor. The cuda.core tensor-bridge supports torch 2.15 unconditionally after the stable PyObject->Tensor shim adoption in NVIDIA#3033.
Removed preview folders for the following PRs: - PR #3033
Adopt the stable
PyObject -> TensorC ABI shim that landed in PyTorch 2.14 (pytorch/pytorch#183323, merged asb69f7838a7) as the primary path in the cuda.core tensor bridge, keep the existingpyobj_to_aten_handlepointer-arithmetic trick as the fallback for torch 2.3-2.13, and drop the per-release upper version cap now that the shim covers 2.14+.The shim adds ~200 ns per
view_as_torch_tensorcall (benchmark below). The upsides are twofold: (1) ABI safety -- we stop depending on the undocumentedTHPVariablestruct layout for new torch versions, and (2) lower maintenance burden -- the upper-version cap bumps we did on every torch release (#2099 and the aborted attempt in #3031 are the two latest examples) are no longer needed, becausetorch_tensor_from_pyobjectis part of torch's stable C ABI contract and any future torch 2.x is covered automatically.Changes
cuda_core/cuda/core/_tensor_bridge.pyx: resolvetorch_tensor_from_pyobject+aoti_torch_delete_tensor_objectviadlsym/GetProcAddresson firstview_as_torch_tensorcall. Dispatch key is symbol presence, not version string:try/finally)cuda_core/cuda/core/_memoryview.pyx: drop the upper cap in_torch_version_check. Lower bound stays at2.3(required for the AOTI metadata functions). On any future torch 2.x the shim path handles it; no more per-release cap bumps.cuda_core/cuda/core/_tensor_bridge.pyi: regenerated for the two newctypedefs.The resolver stubs on Windows target
torch_cpu.dllbecausetorch_tensor_from_pyobjectandaoti_torch_delete_tensor_objectboth live inlibtorch_cpu, notlibtorch_cuda(confirmed viadladdron 2.15.0+cu132). This mirrors the existing split -- onlyaoti_torch_get_current_cuda_streamis CUDA-only.Verification (local)
Env: fresh conda, python 3.12, CUDA 13.4 toolkit, RTX 6000 Ada; release build of
cuda-core(pip install --config-settings debug= -e .-- editable builds default todebug=Trueon Linux); test targetcuda_core/tests/test_utils.py -k torch(63 tests:TestViewGPU[torch-*],TestViewCudaArrayInterfaceGPU[torch-*],test_torch_tensor_bridge_dtypes[*],test_ml_dtypes_bfloat16_torch_dlpack).The ptr/itemsize assertions in
test_torch_tensor_bridge_dtypescover 13 dtypes end-to-end -- if the shim's owned handle disagreed with torch's view of the storage, the tests would catch it.Benchmark
Same env, same tensor (32×32 float32 CUDA), release build, 2 000 warm-up iterations outside the measurement loop, then 30 outer × 10 000 inner iterations; each outer sample is a
time.perf_counter_ns()burst averaged over its 10 k inner calls. Three independent runs per torch version:Fallback band on 2.13 sits roughly where #1894 reported (565 ns on 2.11, same hardware), a touch faster thanks to the Cython-typed
StridedMemoryView viewarg in4ac0f820289. Shim band on 2.14 and 2.15 overlap -- stable ABI genuinely holds across the release.The ~200 ns absolute overhead is C++ work in libtorch: a vtable call into
libtorch_python'sTHPVariable_Unpack, anat::Tensorcopy-construct (bumps the storage refcount), and the matchingaoti_torch_delete_tensor_objecton the way out. It doesn't shrink with Cython optimisation.Follow-up idea for upstream
The shim returns a new reference. For
view_as_torch_tensorwe only read metadata off the handle (get_data_ptr,get_sizes,get_strides,get_dtype,get_device_type,get_device_index) and the Pythonobjoutlives the handle, so a borrowed variant that skipped the refcount bump would be functionally equivalent and would close the perf gap. Worth asking upstream for atorch_tensor_borrow_from_pyobject(or a borrow flag on the existing shim) in a follow-up.-- Leo's bot