[DISCUSS] Target dispatch in C++

Currently, TVM doesn’t support a target dispatch mechanism in C++. As discussed in this PR, this can lead to design issues, duplicated code, and bugs.

Problem

In the example discussed in the PR, schedule_extern must be implemented differently for C++ and Python. In Python, a schedule can be defined in generic/extern.py. The code in this file can then call schedule_injective, which will automatically call the right method for the current target. In other words, every target can use common code. In C++, this is not possible. Instead, the correct schedule_injective function must be called explicitly, which means that each target needs its own implementation of schedule_extern. This leads to each target having duplicated code.

If C++ were to have a target dispatch mechanism, we could write one schedule_extern in generic/extern.h, and use it from all targets in both Python and C++. This is easier to write, much less error-prone, and would allow clean refactoring.

Proposal

I propose a target dispatch mechanism in C++. We would have a function like topi::generic::dispatch::schedule_injective(args), which would get the correct function for the current target and call it with the correct arguments.

From some inspection, it seems that C++ TOPI already registers generic functions here, and this would be a natural place to add this functionality. I also think that this would benefit other areas of the code, as we could move more logic into C++ and have more consistency between the Python and C++ interfaces.

I am pretty new to this area of the code, and would appreciate some thoughts on this proposal, as well as guidance for implementation :grinning:

cc @masahi @vinx13 @tqchen

3 Likes

In the C++ example,


if you call topi::generic::schedule_injective, GenericFunc::operator() will be invoked, which should have the same dispatching mechanism as called from Python side. So why generic::dispatch:: is still needed? Correct me if I’m wrong.

1 Like

Can you expand on this a little more? topi::generic::schedule_injective is not a GenericFunc. WrapSchedule(topi::generic::schedule_injective) is though. It looks like GenericFunc already has a ::Get() method, where you can get a generic func by name. Maybe we can just use that?

It looks like I was wrong about C++ dispatch. I’m doing some testing now and will keep you updated.

Update: it appears that this dispatch does work, however it’s a little confusing. For example, I have a change where generic/extern.py calls back into generic/extern.h. The function in this file calls tvm::Generic::Get("schedule_injective")(...), but this calls the Python version of schedule_injective, not the C++.

In any case, it seems that this works properly, we just need some refactoring to make it consistent. What do you think? I suppose this also gets into a discussion of whether we want to keep most logic in Python or C++.

@masahi for visibility.

You are right, topi::generic::schedule_injective is not a generic func. We should use ::Get instead.
There are two kinds of generic func in Python, one decorated with @tvm.generic_func and one with @tvm.override_native_generic_func. The latter one shares dispatcher with C++, so when there are schedules registered in Python, C++ side will be overrided.

1 Like

I think I was confused by the fact that there is a lot of duplicated code between Python and C++ that is not consistent. For example, injective schedules have slightly different logic (it seems like Python has been updated more recently).

Would everyone be ok with me trying to consolidate some logic into C++?

Thanks, I was not aware of C++ GenericFunc (seems introduced in #892).

Duplication and inconsistencies between python and C++ need to be fixed, as there have been more C++ API users recently (MS and FB, for example). I’m +1 for consolidating logic into C++.

3 Likes