How to use _to_dlpack

Hi, I’m working on getting tvm integrated into some pytorch layers. To do this, I need to wrap a pytorch tensor as a DLpack. Fortunately, pytorch supports this using torch._C._to_dlpack(). However, this returns a PyCapsule object. I need to be able to point tvm to the data attribute of the DLTensor, but pytorch cannot access any attributes of a PyCapsule. Am I doing something wrong or is _to_dlpack only meant to be an intermediate between c function calls?

I’d like to do something like this:

@tvm.register_extension
class FireTensor(object):
    _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

    def __init__(self, tensor):
        self.handle = torch._C._to_dlpack(tensor)

    @property
    def _tvm_handle(self):
        return self.handle.data
1 Like

Hi buddy, have you figured out how to integrate tvm with pytorch now? Many thanks.

Yup, it works pretty well! Here’s the code I use to make it work. Basically, a FireTensor is a TVM compatible object that operates directly on a PyTorch Tensor.

import ctypes
import torch
import tvm

@tvm.register_extension
class FireTensor(object):
    _tvm_tcode = tvm.TypeCode.ARRAY_HANDLE

    def __init__(self, tensor):
        self.handle = torch._C._to_dlpack(tensor)
        self.name = self.get_name()
    def get_name(self):
        ctypes.pythonapi.PyCapsule_GetName.restype = ctypes.c_char_p
        ctypes.pythonapi.PyCapsule_GetName.argtypes = [ctypes.py_object]
        return ctypes.pythonapi.PyCapsule_GetName(self.handle)
    def to_torch(self):
        return torch._C._from_dlpack(self.handle)
    @property
    def _tvm_handle(self):
        ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p
        ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
        return ctypes.pythonapi.PyCapsule_GetPointer(self.handle, self.name)
1 Like

Wow it looks awsome! Wish I had read your post earlier to save several hours.