Update on "Upgrade to DLPack 1.0."

ysiraichi · ysiraichi · commit c95743ccbc5b · 2025-06-28T10:00:09.000-03:00
This PR makes the necessary changes in order to upgrade PyTorch DLPack
support to version 1.0. In summary, we add support for the following:

- Support both `DLManagedTensor` and `DLManagedTensorVersioned` when
  producing and consuming DLPack capsules
- New parameter for `__dlpack__` method: `max_version`
- Version checks:
    - Fallback to old implementation if no `max_version` or if version
      lower than 1.0
    - Check that the to-be-consumed capsule is of version up to 1.X

In order to accommodate these new specifications, this PR adds the
following main changes:

- `torch._C._to_dlpack_versioned` Python API (Module.cpp): new Python
API for creating a versioned DLPack capsule (called by `__dlpack__`
method)
- `DLPackTraits&lt;T&gt;` class (DLConvertor.h): select the correct
traits (e.g. capsule name, conversion functions) depending on which 
DLPack tensor class is being used
- `toDLPackImpl&lt;T&gt;` function (DLConvertor.cpp): populates the
common fields of both classes
- `fromDLPackImpl&lt;T&gt;` function (DLConvertor.cpp): constructs a tensor
from a DLPAck capsule
- `fillVersion&lt;T&gt;` function (DLConvertor.cpp): populates the version
field for `DLManagedTensorVersioned` (no-op for `DLManagedTensor`)
- `tensor_fromDLPackImpl&lt;T&gt;` function (tensor_new.cpp): outer function
for constructing a tensor out of a DLPack capsule that also marks the
capsule as used

cc ezyang gchanan

[ghstack-poisoned]
diff --git a/aten/src/ATen/DLConvertor.h b/aten/src/ATen/DLConvertor.h
@@ -12,7 +12,7 @@ namespace at {
 
 TORCH_API ScalarType toScalarType(const DLDataType& dtype);
 TORCH_API DLManagedTensor* toDLPack(const Tensor& src);
-TORCH_API DLManagedTensorVersioned* toDLPackVersioned(const Tensor& src);
+TORCH_API struct DLManagedTensorVersioned* toDLPackVersioned(const Tensor& src);
 TORCH_API Tensor
 fromDLPack(DLManagedTensor* src, std::function<void(void*)> deleter = {});
 TORCH_API Tensor fromDLPackVersioned(
