[Feature Request] Function Overloading

[Ubospica]

Goal and motivation. Supporting function overloading in tvm-ffi. This would be helpful for libraries that use overloading in their code. For example, XGrammar has a lot of overloading functions. Supporting overloading can unblock it from migrating to tvm-ffi.

Method.

1. According to our guideline that the C ABI should be stable and minimal, we should keep the C ABI unchanged.
2. Extend tvm::ffi::Function to implement a chain to maintain all function pointers in the order of registration.
3. When calling functions, try functions in order until the first matched.

This also aligns with nanobind's method to handle overloaded functions. It allocates a contiguous block of memory to store all function overloads. Its memory layout:

+------------------+  ← nb_func* (PyObject_VAR_HEAD)
| ob_refcnt        |
| ob_type          |
| ob_size          |  ← number of overloads (Py_SIZE)
+------------------+
| vectorcall       |  ← function pointer: call dispatcher
| max_nargs        |  ← maximum number of arguments among all overloads
| complex_call     |  ← boolean: whether a complex call path is required
| doc_uniform      |  ← boolean: whether all overloads share the same docstring
+------------------+  ← sizeof(nb_func)
| func_data[0]     |  ← metadata for the first overload
|   (overload 0)   |
+------------------+
| func_data[1]     |  ← metadata for the second overload
|   (overload 1)   |
+------------------+
| ...              |
+------------------+
| func_data[N-1]   |  ← metadata for the last overload
|   (overload N-1) |
+------------------+

cc @junrushao @tqchen @DarkSharpness

[tqchen]

Because the main goal of the design is to keep minimalism, we should not change the implementation of ffi::Function ideally, however, this is something that maybe build inside reflection namespace as a implementation of the Function(where callable is backed by something like this), where there is a single entry and dispatch happens under the hood

[junrushao]

It's indeed a reasonable ask. I was curious to learn thoughts if it's a necessity though, if it comes at the cost of extra design complexity and runtime overhead at default path, particularly when frontend languages like Python don't really support overloading

[tqchen]

I think the main goal is to NOT change the default path. Instead allow overload hidden under an ffi::Function that does dispatch inside

[Ubospica]

Thanks @junrushao and @tqchen for the insightful comments!

On the overhead, I agree, we should be extra careful not to introduce additional overhead for overload resolution. I feel it is fine in principle because if the first overload function matches, the time consumption should be the same as the current impl.

I think the main goal is to NOT change the default path. Instead allow overload hidden under an ffi::Function that does dispatch inside

This also makes sense. We should keep ffi::Function unchanged. We can have another tvm::ffi::reflection::OverloadedFunction that stores all overloads and handles the dispatching logic, and use Function::FromPacked to wrap it.

One notable change is that, for the global function registration table, we also need to change from Map<String, Function> to Map<String, OverloadedFunction> to handle overloading. When functions are registered with the same name, we just add another overload. Similar logics should also be added to class member functions.

Thanks to @DarkSharpness, who helped with the design and could be helpful on the implementation.

[junrushao]

Sounds like a good proposal!

[DarkSharpness]

I have one proposal implementation here (not guarantee to be functional, just a demo)

The idea is to support function overloading via an OverloadedFunction wrapper. When there is no overload, the call simply falls back to unpack_call, adding only a single boundary check on the critical path.

When multiple overloads are present, the system attempts them sequentially and selects the first one whose signature matches the incoming arguments.

Users can choose to allow overload in a ObjectDef or not.

[tqchen]

OK, refining things further, i think we should be able to hide everything under ffi::Function, no changes to core/registry etc while still getting minimal overheads. We can iterate on API names and impl details, but to illsurtate the idea:

// implemented in the following file
#include <tvm/ffi/extra/function_overloading.h>

// simple two function impl, we can use vargs template here
template<typename FuncA, typename FuncB>
class TypedOverloadedFunc {
  public:
   // exposes call packed signature
   void operator(const AnyView* args, int32_t num_args, Any* result){
      // dispatch into func_a and func_b here 
   }

  private:
    FuncA func_a_;
    FuncB func_b_;
}

template<typename FuncA, typename FuncB>
ffi::Function CreateOverloadedFunction(FuncA fa, FuncB fb) {
    return ffi::Function::FromPacked(TypedOverloadedFunc<FuncA, FuncB>(fa, fb));
}

The implementation efficiency will be as efficient as the specialized impl because everything will be unlined into the packed function closure. User just need to be mindeful create overloaded function while from registry pov it is just ffi::Function

[DarkSharpness]

@tqchen I believe a static TypedOverloadedFunc, which takes a list of Callables as template arguments, is likely the most efficient implementation: the compiler can (hopefully) inline each candidate and minimize dispatch overhead.

However, the trade-off is reduced API flexibility. In the static style, users must explicitly enumerate all overloads in a single Def invocation. This makes certain advanced patterns, such as cross-module overloads where functions with the same name are registered in different modules, harder or impossible.

Example APIs:

static-style:

GlobalDef().def("func", TypedOverloadedFunc{FnA, FnB});

dynamic style:

GlobalDef().def("func", FnA);
// somewhere else
GlobalDef().def("func", FnB);

Both style make sense to me.

[DarkSharpness]

BTW, many of these calls cannot actually be inlined. See this minimal example

TL;DR: passing a member function via a raw function pointer loses important static information, so the compiler can no longer inline the call. Wrapping the member function in a templated callable preserves the static type information and gives the compiler a chance to inline the invocation.

[tqchen]

I agree there can indeed be a tradeoff, in our particular case, i think the ffi::Function simplicity is an important constraint (because we need ABI to be simple and consistent to compiled and runtime use), and likely asking for all overloaded variant to appear in same place of registration is not too bad of an ask(to avoid intended behavior).

It could still be possible to have a dynamic dispatcher (that exposes itself as a ffi::Function) where you need to have single dec site and register overloading for that, however, i am not sure if we want to immediately go that route. Perhaps starting with constraint version would be ok.

[DarkSharpness]

I'm not sure whether my demo implementation satisfy these constraints. I only introduced a static member function FromOverloaded to the ffi::Function, which can be further hidden to private. The implementation is based on passing a closure into the ffi::Function, which avoids too much intrusive modification in ffi::Function

This is a dynamic-style overload, which might be slightly slower in overload path (when there's no overload, it will fallback to unpack_call, no performance loss).

For ObjectDef, it's easy to integrate since we cannot have more than 1 ObjectDef for the same class. We can use a mapping to track all the overload. For GlobalDef, I agree it's much more tricky and need something like a global dispatcher.

[Ubospica]

@DarkSharpness Could you submit a PR for this? This would be helpful for review.

[DarkSharpness]

@Ubospica sure. I'm still working on it and will later open a PR when the examples are ready.