[RUNTIME][NDArray] Allowing External Libraries to Subclass NDArrays

[junrushao]

This PR allows an external library to subclass TVM's NDArray infrastructure without too much work.

It introduces a type trait array_type_index<T>::code. For TVM's NDArray itself, the trait is 0; For irrelevant classes, the trait is -1; Any external subclass of TVM NDArray should override it to be a positive int32 number.

To use this extension, an external library should

1. On C++ backend, inherit TVM's NDArray and NDArray container, and define the trait array_type_index for the NDArray.
2. On C++ backend, define a constructor in the inherited class that accepts a pointer to TVM's Container, which is nullable.
3. On Python frontend, inherit tvm.nd.NDArrayBase, define the class attribute _array_type_index consistent to the C++ type trait, and register the subclass using tvm.register_extension.

CC: @tqchen, @were, @reminisce

[junrushao]

(Outdated, don't read)

Added a type trait array_type_index, which is only enabled to NDArray and subclasses of NDArray.

A subclass of NDArray should do the following things in their own codebase, without modifying TVM.

1. specialize the type trait array_type_index
2. define implicit conversion from TVMArgValue to itself, so as to enable implicit conversion.
3. define a constructor from TVM::runtime::NDArray::Container

[junrushao]

(Outdated, don't read)

Two issues with this PR

1. Absence of virtual destructor (and we may not introduce vtable) requires us to be extremely careful. For example, a bug in my PR is the use of As. I will fix it tomorrow.

2. I am very ignorant of how to glue languages, and I don’t know how to do it automatically on python side.

[tqchen]

We do not need a virtual destructor because deleter is called. refer to external type on how to enable automatic explicit conversion https://github.com/dmlc/tvm/blob/master/include/tvm/runtime/packed_func.h#L456

[junrushao]

(Outdated, don't read)

@tqchen On the C++ side, automatic implicit conversion has already been enabled via overriding TVMValue::operator T() in external developers' codebase, here, so why we need explicit conversion on the C++ side?

[junrushao]

(Outdated, don't read)

Or you mean moving the As method to TVMValue, and rename it somehow like AsSubNDArray?

[tqchen]

I am not sure if operator T() can be overridden in the external codebase

[junrushao]

(Outdated, don't read)

@tqchen @were Hey I updated the code accordingly. By tweaking TVMValueCast::Apply, we could correctly dispatch operator T(), so TVMPodValue_ could be automatically converted to the subclass of NDArray, like this line.

I haven't got a clear sense about the glue part.

[junrushao]

(Outdated, don't read)

@tqchen Updated accordingly. A potential issue of current AsExtension is that its return value has two signatures, one is TSubNDArray, another is const TExtension&.

[tqchen]

@junrushao1994 that could indeed be an issue.

In that case, let us just use AsNDArray()(rename AsSubNDArray->AsNDArray). Sorry for not noticing this before.

[junrushao]

(Outdated, don't read)

Weird CI fails. I can at least pass "BUILD: CPU" on my machine.

UPDATE: Addressed via rebasing to master.

[tqchen]

The error seems to be related to cython, try to directly use NDArrayBase as opposed to NDArray

[junrushao]

@tqchen @were @reminisce Could you kindly do another round of review? Thanks!

[tqchen]

Implementation looks good, some final comments