[RFC][Runtime] Bring `PackedFunc` into TVM Object System

Summary

This RFC allows developers to use PackedFunc as TVM objects, which completes the last missing step of TVM runtime object system.

Motivation

Historically, several fundamental data structures in TVM are not part of the runtime object system, namely NDArray (not object), Module (not object), String (not exist), Array (not in runtime), Map (not in runtime), PackedFunc (not yet an object).

The rationale of the original design is mainly for simplicity, which is desirable for the usecases as a monolithic compiler. As time goes on, the community has come to realize the fact that the object system should be inclusive enough and by design allow more convenient integration with vendor libraries. Therefore, as part of the effort in TVM refactoring and TVM Unity, recent work strives to re-implement these core data structures to be consistent with the runtime object protocol with stable ABI guarantee, and thus could be passed across the DLL boundary.

As the central piece of the TVM ecosystem, this proposal focuses on making PackedFunc a TVM object. By doing so, it completes the last missing piece of the object ecosystem, allows TVM containers to carry PackedFuncs, and enables PackedFuncs to be passed across the DLL boundary to bring convenience to the vendor library integration.

Guide-level introduction

This is mainly a developer-facing feature, and thus there is no sensible change to the existing functionalities to the end users, who are still supposed to use the same PackedFunc API.

Only one major object is introduced, PackedFuncObj, a TVM object in the runtime system (detailed in the next section) which is an ABI stable data structure for packed functions that could be shared across language and DLL boundary.

To avoid API misuse from developers, the PackedFuncObj cannot be created or manipulated directly, and the specialization of its creation make_object<PackedFuncObj> will be deleted for safety. Instead, the developer-facing class PackedFunc remains responsible for creating and managing the object, and for properly setting its content.

In the future, it’s possible to incrementally add more information into PackedFuncObj to better help debugging and error reporting.

Reference-level introduction

As introduced below, the RFC introduces a new class:


class PackedFuncObj : public runtime::Object {
  using FCallPacked = void(TVMArgs, TVMRetValue*);
  FCallPacked* f_call_packed_;
};

A templated subclass is introduces to do the type-erasing trick:

template <typename TCallable>
class PackedFuncSubObj : public PackedFuncObj {
  TCallable callable_;
};

The PackedFuncObj inherits an intrusive reference counter and an object deleter from the runtime::Object. Besides, with the inheritance trick on PackedFuncSubObj, the field callable_ is introduced to store the content of the callable object, which can be a function pointer, a struct/class, an anonymous lambda function or any other object.

To make the change minimal, PackedFuncObj is not designed to be serializable, and doesn’t support TVM’s native reflection. Copying the type-erased object is strictly prohibited for now for simplicity, and instead copying the PackedFunc is implemented as a straightforward increment to the reference counter by 1.

Drawbacks

Just like every change to the runtime, the proposed change could slightly affect runtime’s binary size. The effect, depending on the compiler, could be positive or negative.

Overall, given that it brings significantly better experience as stated in the previous sections, we believe the benefits outweighs the potential drawback.

Rationale and alternatives

This refactoring is the last missing piece of effort that brings core data structures of the TVM runtime into the ABI-stable TVM runtime.

Alternatively, one might argue that it’s not important whether PackedFunc should be a TVM or not; however, it significantly brings negative impact when TVM object system is used across the DLL boundary, or putting PackedFunc into TVM containers.

Prior Art

NDArray and Module are brought into the object system according to RFC Issue #4286.

Containers, including String, Array and Map, are discussed in the forum thread and brought into the object system. The String part is introduced by PR #4628, Array in PR #5585, and Map in PR #5740.

DGL, one of the most popular frameworks for distributed graph neural network training, adopts TVM’s object and FFI system.

Unresolved questions

This RFC only introduces C++ ABI for invoking a PackedFunc, which might have some limitation when linking artifacts compiled by different compilers. In the future, more effort should be invested into the design of a stable C ABI when two PackedFuncs come from different TVM runtime.

Future possibilities

Based on similar metaprogramming tricks, it’s possible to extract the function signatures of TypedPackedFunc and to make error reporting more readable.

9 Likes

I’m happy to shepherd this RFC

CC: @spectrometerHBH @tqchen @areusch

1 Like

Thanks @cyx. The RFC looks good to me. Looking forward to a formal RFC and following PR.

1 Like

hi @cyx thanks for this proposal! I have a couple questions for you on it.

First, could you elaborate a bit more on the use case you have in mind?

it seems like this might be one of them. Given we can already pass PackedFunc as arguments to other PackedFunc and return them, it seems like one of the big changes here is making the TVM data structures more expressive. I can certainly think of some good use cases for storing function pointers in a map, but I wonder about passing such a data structure at an API boundary.

Could you say what you mean/imply by “ABI stable” here? Are you separately compiling a library which needs to invoke PackedFunc? Is there a reason you prefer not to include the PackedFunc library inside that one? (just playing devils’ advocate here–iiuc it seems like another solution is to extract include/tvm/runtime/packed_func.h into a separate library to make it easy to implement PackedFunc stubs to a separate library)

Could you elaborate? Would there be operator PackedFuncObj?

Sort of the inverse of my last question–suppose I have e.g. a Map whose values are PackedFuncObj. How can I invoke the PackedFunc or reconstruct? The function pointer seems to be private in this definition.

1 Like

To summarize our offline discussion with @areusch @tqchen.

Clarification:

  1. This RFC doesn’t change any of the existing functionality, including C ABI or PackedFunc’s C++ API. Any modification to the C ABI is out of scope of this RFC.
  2. Calling a PackedFunc inside TVM codebase directly uses the C++ API PackedFunc::operator() or CallPacked with C++ ABI, where there is no C ABI involved, which is a shortcut. For function calls across FFI boundary, our system uses C ABI instead, which this RFC doesn’t aim to change

There are a few levels that goes into this consideration:

  • L0: Minimally we have the c_runtime_api.h that is the common ground of everything, and different kind of runtime impls
  • L1: At C++ runtime there is a need of object system, which is now still specific to the C++, they are needed for flexibility reasons in the compiler side and sometimes complicated runtime(like VM).

It is not desirable to bring L1’s design considerations into L0 (specifically standards that L0 should follow), since L0 should really be kept stable for cases like embedded settings. What the proposal was talking about, we believe, is to streamline some of the L1 data structures, just like what we did for String. They would bring benefit to the implementations (use packed func as object). There is also “ABI” at L1 level but never made official.

To answer @areusch’s questions (with my understanding):

By making PackedFunc a TVM object, we are able to put them into TVM containers (Array, Map), which as Andrew said, makes TVM data structures more expressive. It does allow us to pass PackedFunc across the FFI boundary, because all the the existing mechanism is unchanged and still works. More particularly, calling a PackedFunc on the python side still uses the good old TVMFunctionHandle and the ABIs defined in c_runtime_api.

As mentioned in the section “Unresolved questions”, here we only consider C++ ABI, and C ABI is left to future work.

PackedFuncObj is not going to be user-facing, so users won’t have to deal with it any time (IIUC)

For example, say the signature of the Map is: Map<String, PackedFunc> my_map. Users could invoke it via PackedFunc::operator(), i.e. this functionality is unchanged from the existing PackedFunc:

my_map["name"](arg1, arg2, arg3);
2 Likes

@junrushao1994’s understanding and elaboration is accurate.

1 Like

Let’s leave this pre-RFC open for a week, and then send a formal RFC with clarifications to https://github.com/apache/tvm-rfcs/

@cxy would you like to update the pre-RFC according to our discussion? Thanks a lot!

@junrushao1994 Sorry for late reply. The formal RFC [RFC][Runtime]PackedFunc as Object has been submittted. And the final RFC has updated according to our latest discussion.

1 Like

I really like the RFC and it makes the object system close loop with packed func. I have left some minor typo fixes / questions on the formal RFC, other than those the RFC looks good to me.

2 Likes