Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 9 commits into from
Feb 21, 2019
Merged

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

merged 9 commits into from
Feb 21, 2019

Conversation

junrushao
Copy link
Member

@junrushao junrushao commented Feb 16, 2019

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 junrushao changed the title Add type traits for all subclasses of NDArray [RUNTIME][NDArray] Add type traits for NDArray subclasses Feb 16, 2019
@junrushao
Copy link
Member Author

junrushao commented Feb 16, 2019

(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
Copy link
Member Author

junrushao commented Feb 16, 2019

(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
Copy link
Member

tqchen commented Feb 16, 2019

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
Copy link
Member Author

junrushao commented Feb 16, 2019

(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
Copy link
Member Author

junrushao commented Feb 16, 2019

(Outdated, don't read)

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

@tqchen
Copy link
Member

tqchen commented Feb 16, 2019

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

@junrushao
Copy link
Member Author

junrushao commented Feb 17, 2019

(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.

apps/subclass_ndarray/include/subndarray.h Outdated Show resolved Hide resolved
include/tvm/runtime/ndarray.h Show resolved Hide resolved
include/tvm/runtime/packed_func.h Show resolved Hide resolved
@tqchen tqchen self-assigned this Feb 18, 2019
@junrushao
Copy link
Member Author

junrushao commented Feb 18, 2019

(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
Copy link
Member

tqchen commented Feb 18, 2019

@junrushao1994 that could indeed be an issue.

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

apps/extension/include/subndarray.h Outdated Show resolved Hide resolved
apps/extension/src/subndarray.cc Outdated Show resolved Hide resolved
apps/extension/Makefile Outdated Show resolved Hide resolved
apps/extension/include/subndarray.h Outdated Show resolved Hide resolved
apps/extension/src/subndarray.cc Outdated Show resolved Hide resolved
apps/extension/python/tvm_ext/__init__.py Show resolved Hide resolved
python/tvm/_ffi/_cython/ndarray.pxi Outdated Show resolved Hide resolved
python/tvm/_ffi/ndarray.py Outdated Show resolved Hide resolved
python/tvm/_ffi/runtime_ctypes.py Show resolved Hide resolved
@junrushao
Copy link
Member Author

junrushao commented Feb 20, 2019

(Outdated, don't read)

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

UPDATE: Addressed via rebasing to master.

@tqchen
Copy link
Member

tqchen commented Feb 20, 2019

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

@junrushao junrushao changed the title [RUNTIME][NDArray] Add type traits for NDArray subclasses [RUNTIME][NDArray] Allowing External Libraries to Subclass NDArrays Feb 20, 2019
@junrushao
Copy link
Member Author

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

Copy link
Member

@tqchen tqchen left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Implementation looks good, some final comments

python/tvm/_ffi/_cython/base.pxi Show resolved Hide resolved
apps/extension/src/tvm_ext.cc Outdated Show resolved Hide resolved
apps/extension/src/tvm_ext.cc Show resolved Hide resolved
apps/extension/src/tvm_ext.cc Outdated Show resolved Hide resolved
apps/extension/src/tvm_ext.cc Outdated Show resolved Hide resolved
apps/extension/src/tvm_ext.cc Show resolved Hide resolved
include/tvm/runtime/packed_func.h Outdated Show resolved Hide resolved
@tqchen tqchen merged commit 81334be into apache:master Feb 21, 2019
@yzhliu yzhliu mentioned this pull request Mar 2, 2019
28 tasks
@junrushao junrushao deleted the nd branch January 6, 2022 04:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants