[BUG] Versioned DLPack fails when not using tvm-ffi

[slayton58]

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
from_dlpack(..., enable_tvm_ffi=False) segfaults when parsing a versioned DLPack capsule.

Steps/Code to reproduce bug
Run the following code:

     """Repro: CuteDSL's from_dlpack crashes on a versioned DLPack capsule.

     The capsule comes from a real torch.Tensor via torch's own DLPack bindings.
     torch.__dlpack__(max_version=(1, 0)) produces a DLManagedTensorVersioned
     capsule (the DLPack 1.0 ABI). It is a valid capsule:

       - numpy consumes it fine
       - CuteDSL's tvm-ffi path consumes it fine (it uses the C exchange API)

     but CuteDSL's default from_dlpack mis-parses the versioned struct in its
     DLTensorWrapper and crashes (LLVM ERROR / segfault) the first time the
     resulting tensor is accessed (parsing is lazy).

     The one-line wrapper exists only because from_dlpack() calls __dlpack__()
     itself with no version argument -- a bare torch.Tensor therefore yields the
     legacy capsule, which CuteDSL parses fine. The wrapper just forwards the real
     tensor's __dlpack__ while requesting the versioned capsule.

         python cutedsl_versioned_dlpack_repro.py
     """

     import numpy as np
     import torch
     from cutlass.cute.runtime import from_dlpack


     class VersionedDLPack:
         def __init__(self, t):
             self._t = t

         def __dlpack__(self, **kwargs):
             return self._t.__dlpack__(max_version=(1, 0))

         def __dlpack_device__(self):
             return self._t.__dlpack_device__()


     # numpy parses the same versioned capsule fine (numpy is CPU-only).
     cpu = torch.arange(256, dtype=torch.float32)
     arr = np.from_dlpack(VersionedDLPack(cpu))
     print("numpy ok:", arr.shape, np.array_equal(arr, cpu.numpy()), flush=True)

     x = torch.arange(256, dtype=torch.float32, device="cuda")

     # CuteDSL's tvm-ffi path (C exchange API, no capsule parsing) is fine.
     sx_tvm = from_dlpack(x, enable_tvm_ffi=True)
     print("cute tvm-ffi ok:", sx_tvm.shape, flush=True)

     # CuteDSL's default path crashes on the versioned capsule.
     sx = from_dlpack(VersionedDLPack(x))
     print(sx.shape)  # crashes here: accessing the tensor parses the capsule

Expected behavior
No segfault, from_dlpack(..., enable_tvm_ffi=False) completes successfully.

Environment details (please complete the following information):

• torch 2.12
• cuteDSL 4.5.2, 4.6.0.dev0

[brandon-yujie-sun]

DSL uses the DLManagedTensor rather than the versioned one, hence the error. I think we'll need to follow up and support the versioned capsule.

[slayton58]

DSL uses the DLManagedTensor rather than the versioned one, hence the error. I think we'll need to follow up and support the versioned capsule.

I think that would be ideal -- on a related note I do have pytorch/pytorch#188554 that'll be merged into PyTorch in the next day or so that requires the versioned tensor (for the RO flag) -- this passes with enable_tvm_ffi=True but of course fails with the above seg-fault when it's False.

[brandon-yujie-sun]

@slayton58 this is now fixed in 4.6.1