[Feature] Support dynamic-style overload for FFI object types

[DarkSharpness]

Related discussion here #265 .

Modification in short:

1. Add Function::FromPackedInplace in function.h and generalize some methods in registry.h
2. Implement FFI object types overload method (all in extra/overload.h)
3. Add a simple test.

[gemini-code-assist]

Summary of Changes

Hello @DarkSharpness, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Foreign Function Interface (FFI) by implementing dynamic-style function overloading for object types. This feature allows developers to define multiple C++ methods or constructors with identical names but distinct parameter lists, which can then be seamlessly invoked from FFI bindings. The change improves the usability and expressiveness of FFI object definitions, making them more aligned with standard C++ programming patterns.

Highlights

• Dynamic Overloading for FFI Objects: Introduced a new mechanism (OverloadObjectDef) to support dynamic-style overloading for methods and constructors of FFI object types, allowing multiple functions with the same name but different argument signatures.
• Function::FromPackedInplace Addition: Added a new static method Function::FromPackedInplace to function.h, which enables creating a ffi::Function and simultaneously obtaining a pointer to its internal callable, essential for managing overloaded functions.
• Generalized Function Wrapping: Refactored ReflectionDefBase in registry.h to generalize the WrapFunction and GetMethod utilities, making them more flexible and compatible with the new overloading infrastructure.
• New Test Coverage: Included a new test file (test_overload.cc) to thoroughly validate the functionality of overloaded constructors, instance methods, and static methods for FFI objects.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a powerful feature to support overloading of FFI object methods, including constructors, instance methods, and static methods. The implementation is well-designed, with good attention to performance by minimizing virtual function calls. The changes are well-structured, with the core logic encapsulated in the new include/tvm/ffi/extra/overload.h header, and necessary supporting changes made to function.h and registry.h. The addition of tests in tests/cpp/test_overload.cc is great and covers the primary use cases. I have a couple of minor suggestions for improvement.

[junrushao]

Ping @Ubospica @cyx-6 if you guys would love to help review this PR!

[junrushao]

I don't have sufficient evidence to form an opinion on the design below (copied from unittests), but indeed it's super interesting and creative.

struct TestOverloadObj : public Object {
  explicit TestOverloadObj(int32_t x) : type(Type::INT) {}
  explicit TestOverloadObj(float y) : type(Type::FLOAT) {}
  static int add_one_int(int x) { return x + 1; }
  static float add_one_float(float x) { return x + 1.0f; }

  template <typename T>
  auto holds(T) {
    if constexpr (std::is_same_v<T, int32_t>) {
      return type == Type::INT;
    } else if constexpr (std::is_same_v<T, float>) {
      return type == Type::FLOAT;
    } else {
      static_assert(sizeof(T) == 0, "Unsupported type");
    }
  }

  enum class Type { INT, FLOAT } type;
  TVM_FFI_DECLARE_OBJECT_INFO("test.TestOverloadObj", TestOverloadObj, Object);
};

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::OverloadObjectDef<TestOverloadObj>()
      .def(refl::init<int32_t>())
      .def(refl::init<float>())
      .def("hold_same_type", &TestOverloadObj::holds<int32_t>)
      .def("hold_same_type", &TestOverloadObj::holds<float>)
      .def_static("add_one_static", &TestOverloadObj::add_one_int)
      .def_static("add_one_static", &TestOverloadObj::add_one_float);
}

[DarkSharpness]

cc @junrushao .
Now got following error message for mismatched arg type:

C++ exception with description "No matching overload found when calling: `test.TestOverloadObj.hold_same_type` with 2 arguments:
Overload #0: Mismatched type on argument #1 when calling: `test.TestOverloadObj.hold_same_type (0: test.TestOverloadObj, 1: int) -> bool`. Expected `int` but got `const char*`
Overload #1: Mismatched type on argument #1 when calling: `test.TestOverloadObj.hold_same_type (0: test.TestOverloadObj, 1: float) -> bool`. Expected `float` but got `const char*`

for mismatched arg number:

C++ exception with description "No matching overload found when calling: `test.TestOverloadObj.hold_same_type` with 1 arguments:
Overload #0: Mismatched number of arguments when calling: `test.TestOverloadObj.hold_same_type (0: test.TestOverloadObj, 1: int) -> bool`. Expected 2 arguments
Overload #1: Mismatched number of arguments when calling: `test.TestOverloadObj.hold_same_type (0: test.TestOverloadObj, 1: float) -> bool`. Expected 2 arguments

[Ubospica]

Thanks, the changes LGTM! Supporting function overloading is very helpful.

Based on our offline discussions regarding the design, this implementation looks solid. I especially like how we designed these features:

1. The overloaded function is optional and independent of the Function class.
2. The code path for the first overload is aligned to the non-overload version to minimize overhead
3. Although we need to check arguments at runtime, the function signatures are known at compile time, and we are using templates to generate the checking code for certain function signatures. This reduces the runtime overhead.

One thing to note is that we are doing parameter type checking & conversion at the same time. If the conversion is slow (such as Array<Array<Type>>), we need to check and convert all array elements), this may introduce some overhead. But I think it's fine since these slow paths are currently relatively uncommon.

cc @tqchen @junrushao

[junrushao]

Sorry for the late response! I really like this PR

[Ubospica]

@junrushao Thanks so much for merging this PR!