leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

fix tvm-ffi path in from_dlpack

Browse Source
This commit is contained in:
yuanxiaolan
2026-02-03 11:59:13 +08:00
parent 1cfbb53a23
commit de198b2419
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/CuTeDSL/cutlass/cute/runtime.py
+5 -5
View File
@@ -141,6 +141,11 @@ class _Tensor(Tensor):
# If tensor is already a DLPack object, use it directly # If tensor is already a DLPack object, use it directly
if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"): if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"):
self._dlpack_data = tensor.__dlpack_device__() self._dlpack_data = tensor.__dlpack_device__()
elif enable_tvm_ffi:
import tvm_ffi
self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor)
self._dlpack_data = self._tvm_ffi_tensor.__dlpack__()
else: else:
try: try:
# we expect no stream sync. Because torch has different default behavior # we expect no stream sync. Because torch has different default behavior
@@ -149,11 +154,6 @@ class _Tensor(Tensor):
self._dlpack_data = tensor.__dlpack__(stream=-1) self._dlpack_data = tensor.__dlpack__(stream=-1)
except Exception: except Exception:
self._dlpack_data = tensor.__dlpack__() self._dlpack_data = tensor.__dlpack__()
if enable_tvm_ffi:
import tvm_ffi
self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor)
self._dlpack_data = self._tvm_ffi_tensor.__dlpack__()
self._dltensor_wrapper = None self._dltensor_wrapper = None
self._assumed_align = assumed_align self._assumed_align = assumed_align