[RFC] Register OPs from python site and followup

Currently, C++ is the only way to register ops (correct me if I missed something), this requires user to recompile.

Is it intentional?

If not, I would like to enable the python way to register ops, which users could try quickly and enable to opportunity to deliver new op implementation separately: No need to pack all in one library but provide the related python scripts.

And another follow up, is there a plan to make ops hot plugable?

Please let me know your thoughts @tqchen @jroesch

1 Like

It is possible to define a packed func in python and call it from c++. Is something like this what you are looking for?

1 Like

Thanks a lot @masahi, that is very useful information, but I think it may be a way too far from my goal.

Let me try to make things more clear:

Register OPs

Currently, registering ops is like:

TVM_REGISTER_OP("my_op")
 .describe("my_op xxxxx")
 .set_num_inputs(2)
 .set_attr<xxx>(...)
 .add_type_rel(...)

It needs recompile.

What I mean by registering ops from python site is like:

# New API (TBD, just for concept explanation)
register_op("my_op")
my_op.describe("my_op xxxxx")
my_op.add_type_rel(...)

# Existing API
register_op_attr("my_op", ...)
register_intrin_lowering("my_op", ...)

Hot Plug OPs

For the following up about hot plug ops, taking op “image.resize” for example:

Currently, it is packed together with tvm library (except compute and schedule functions can be registered with python).

I’d like to know if there is a plan to dynamic load it.

The concept is like tensorflow’s create op.

Right now we allow adding attributes of operators in python: tvm/op.cc at main · apache/tvm · GitHub.

I suspect that we don’t have a public API to add operators on the python side right now, and I don’t know the reason why not. If you wanted to, you may just expose OpRegEntry::RegisterOrGet

I also wish we could easily add hot pluggable Relay operators (whether for testing, easily supporting additional ops, etc.). Unfortunately I believe the main reason (or at least one major reason) this is currently not available is because type relations (and basically all the type inference machinery) is not exported across the FFI. I’m not sure how easy it is to allow this, and whether or not we want to commit to this with what we currently have vs. a longer term picture.

2 Likes

Picking a pre-existing type relation (Identity would probably be pretty common) might be reasonable for testing

@junrushao1994 @altanh @slyubomirsky Thanks for your reply. A small commit is merged last week [PR #8002]

Whether to expose type relation still needs more feedback.

1 Like

Thanks @zackcquic for the helpful PR!