Register Extension Support for PyCapsule

Currently, the method of enabling DLTensors for tvm relies on python having access to the data attribute. However, pytorch’s _to_dlpack operator returns a PyCapsule that is opaque to python, making incompatible with the way register_extension is set up. If PyCapsules were allowed, integration with PyTorch would be much more efficient. As a caveat, I may be totally missing something about how registering extensions work, and would be thrilled to be corrected. My goal is to use a PyTorch tensor buffer without having to copy to an ndarray. I imagine the code would look similar to 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

[jwfromm]

Your commit resolves the problem, everything is working great now! Thanks. For reference, here’s what the custom tensor for pytorch ended up looking like.

@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)

[tqchen]

Try https://github.com/dmlc/tvm/pull/669 which checks if the buffer is compact even when strides is available. This is a recommended approach when the data buffer is compact.

Alternatively, use tvm.decl_buffer to declare a buffer with strides array, and specify them with binds argument in http://docs.tvmlang.org/api/python/build.html#tvm.build

see also http://docs.tvmlang.org/api/python/tvm.html#tvm.decl_buffer and https://github.com/dmlc/tvm/issues/585