I have been styding the TVM compiler and I am having trouble fully understanding the foreign function interface for Python. I understand that ctypes library is used to open TVM shared library at runtime and provide access to all C++ functions in the registry to Python code.
However, consider the Python class Stage (schedule.py:202), it has an unroll method (L403) that calls the FFI function StageUnroll (L411) passing self as a parameter. This global function is basically an alias to Stage::unroll (schedule_lang.cc:966) method in C++. How does C++ understand that this is an instance of the C++ class Stage since we are actually passing self, a Python object?
I suppose Python Stage somehow wraps C++ Stage, but I can’t figure out how (and where) this wrapping is happening. Can someone please elucidate this?
I think the magic happens by register_object decorator. By using exactly the same class name in c++ and python, we are creating a mapping between python and c++ objects. I hope I got it right, but this is my superficial understanding.
Check the decorator “@tvm._ffi.register_object” of class Stage in schedule.py.
What it is doing is to provide the mapping relationship of {object type index: python object type}.
Where the type index could be gotten in following call stack:
class StageNode (schedule.h)
TVM_DECLARE_FINAL_OBJECT_INFO
TVM_DECLARE_BASE_OBJECT_INFO
Object::GetOrAllocRuntimeTypeIndex
So when constructing an object, tvm runtime will maintain a hash map where {key: value} = {object name: object type index}
Once ffi is issued, i.e. te.StageUnroll(self, var), function “_make_tvm_args” from packed_func.py will split the “self” argument into:
the class handle (inherited from class Object) which was a Stage instance created inside tvm runtime.
the object type index
Finally, function in tvm runtime will validate the instance with its’ type index and call designated method accordingly.