[Feature] Support vanilla C++ STL in FFI.

[DarkSharpness]

Similar to pybind, we add a stl.h which support array, vector, tuple, optional and variant. After this file is included, users can use native C++ components, which could hopefully improve compatibility and reduce manual effort in converting between tvm::ffi components to C++ STL components.

We also modify the function_detail.h a little, so that we support all kinds of argument type (T, const T, T&, const T&, T&&, const T&& have been tested) in C++ exported functions.

Example code:

#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/error.h>
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/stl.h>
#include <tvm/ffi/function.h>

#include <algorithm>
#include <array>
#include <cstddef>
#include <numeric>
#include <optional>
#include <variant>
#include <vector>

namespace {

// optional, array, vector, tuple is supported
auto sum_row(std::optional<std::vector<std::array<int, 2>>> arg)
    -> std::tuple<bool, std::vector<int>> {
  if (arg) {
    auto result = std::vector<int>{};
    result.reserve(arg->size());
    for (const auto& row : *arg) {
      result.push_back(std::accumulate(row.begin(), row.end(), 0));
    }
    return {true, result};
  } else {
    return {false, {}};
  }
}

// const reference is also supported, though not recommended and won't bring performance gain
// all types must be cast to value, and then pass by reference
auto find_diff(const std::vector<int>& a, std::vector<int> b) -> std::size_t {
  const auto max_pos = std::min(a.size(), b.size());
  for (std::size_t i = 0; i < max_pos; ++i) {
    if (a[i] != b[i]) {
      return i;
    }
  }
  return max_pos;
}

auto test_variant(std::variant<int, float, std::vector<int>>&& arg)
    -> std::variant<int, std::vector<int>> {
  if (std::holds_alternative<int>(arg)) {
    std::vector<int> result;
    auto& value = std::get<int>(arg);
    result.reserve(value);
    for (int i = 0; i < value; ++i) {
      result.push_back(i);
    }
    return result;
  } else if (std::holds_alternative<float>(arg)) {
    return static_cast<int>(std::get<float>(arg));
  } else {
    auto& value = std::get<std::vector<int>>(arg);
    std::reverse(value.begin(), value.end());
    return std::move(value);
  }
}

TVM_FFI_DLL_EXPORT_TYPED_FUNC(sum_row, sum_row);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(find_diff, find_diff);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(test_variant, test_variant);

}  // namespace


}  // namespace

Python part:

from __future__ import annotations

from tvm_ffi.cpp import load_inline
from pathlib import Path

cur_path = Path(__file__).parent

with open(cur_path / "stl.cpp") as f:
    cpp_source = f.read()

module = load_inline(
    "test_stl",
    cpp_sources = cpp_source,
)

print(module.sum_row([[1, 2], [3, 4]]))  # Expected output: (True, [3, 7])
print(module.sum_row(None))  # Expected output: (False, [])
print(module.find_diff([1, 2, 3, 4], [1, 2, 4, 3]))  # Expected output: 2 (index = 2)
print(module.test_variant(2))  # Expected output: [0, 1]
print(module.test_variant(3.1))  # Expected output: 3
print(module.test_variant([1, 2]))  # Expected output: [2, 1]

[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 extends the Foreign Function Interface (FFI) capabilities by adding direct support for several standard C++ STL containers, including std::array, std::vector, std::tuple, std::optional, and std::variant. This integration, inspired by pybind, is designed to enhance the seamless interaction between C++ and other languages via FFI, reducing the need for boilerplate code and manual type conversions. Additionally, the underlying argument handling mechanism has been improved to ensure robust and safe processing of various C++ reference types.

Highlights

• New STL Container Support: Introduces a new stl.h header file to enable Foreign Function Interface (FFI) support for several standard C++ containers, including std::array, std::vector, std::tuple, std::optional, and std::variant.
• Improved Compatibility and Reduced Effort: The addition of native STL container support aims to enhance compatibility between C++ components and FFI, significantly reducing the manual effort required for type conversions.
• Enhanced Argument Handling for References: Refactors the ArgValueWithContext class in function_details.h to correctly handle C++ lvalue references by introducing conditional storage for temporary values, which prevents potential dangling reference issues during argument unpacking.

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 support for several C++ STL containers within the FFI layer, a valuable feature that enhances compatibility and simplifies usage. The implementation is well-structured, primarily using TypeTraits specializations for std::array, std::vector, std::tuple, std::optional, and std::variant. The modifications in function_details.h aim to correctly handle argument passing by reference. While the overall approach is solid, I've identified a critical bug in the std::optional TypeTraits implementation that would cause infinite recursion, and a performance issue related to const reference handling in function_details.h. My review includes specific suggestions to address these points.

[tqchen]

please split the core changes function_details into a separate PR.

likely stl should go into extra. The main reason is that while stl can be helpful sometimes, the tradeoff is when present in cases like object, we can no longer to near zero cost access in static languages like rust via directly ptr and access (otherwise it have to be a call to reflection getter that does conversion under the hood, so we should always encourage ffi types when possible (and document such rationales in stl.h).

[DarkSharpness]

please split the core changes function_details into a separate PR.

likely stl should go into extra. The main reason is that while stl can be helpful sometimes, the tradeoff is when present in cases like object, we can no longer to near zero cost access in static languages like rust via directly ptr and access (otherwise it have to be a call to reflection getter that does conversion under the hood, so we should always encourage ffi types when possible (and document such rationales in stl.h).

I agree that stl should go into extra. The main goal of this PR is to introduce a pybind/PyTorch-style interface that provides greater flexibility for user-defined conversions. In performance-critical paths, however, FFI types should always take priority. This change only performs a one-time transformation for convenience and potential backward compatibility with existing code.

[tqchen]

@DarkSharpness that sounds good, can we still move function_details.h chaneg into a separate PR? mainly want to make sure core changes are clearly scoped

[DarkSharpness]

@tqchen now in #229 . Also tried to fix the dependent template lookup error of in the CI failure. I will rebase after #229 is merged.

More tests of this PR is working in progress.

[DarkSharpness]

cc @tqchen . We now support array, vector, map, unordered_map, optional, tuple, variant and function in this PR.

For STL containers, when passed from Python to C++, we will always construct from the Python object. When passed from C++ to Python, we will construct the corresponding object (e.g. std::array -> Array) first, and then move it to AnyView.

We also allow types that doesn't enable storage (e.g. std::string) to appear in STL container (e.g. std::vector<std::string>) to provide better compatibility.

Finally, do you think I should split the stl.h into multiple files? I'm not sure whether it is too long and include too many unwanted headers (e.g. user may only need to use vector).

[tqchen]

Thanks @DarkSharpness , would be good to get some reviews, cc @junrushao @Ubospica @Kathryn-cat

[junrushao]

I love this change. Please rebase to latest HEAD and I’m happy to get it in.

If I read correctly, it doesn’t support transferring a C++ STL container’s ownership to Python, or allow Python side to inplace update C++ STL containers, right? It’s okay given STLs are really secondary in our system, but we will need to document this properly.

[DarkSharpness]

rebased cc @junrushao . Yes, your understanding is right. Since STL is not native to tvm::ffi, type conversion will happen at each Python <-> C++. I'm not sure where to document this. I have added a warning here in stl.h:

/*!
 * \file tvm/ffi/extra/stl.h
 * \brief STL container support.
 * \note This file is an extra extension of TVM FFI,
 * which provides support for STL containers in C++ exported functions.
 *
 * Whenever possible, prefer using tvm/ffi/container/ implementations,
 * such as `tvm::ffi::Array` and `tvm::ffi::Tuple`, over STL containers
 * in exported functions for better performance and compatibility.
 */

[tqchen]

move this to tests/python/cpp_src

[tqchen]

Add more notes:

Native ffi objects comes with stable data layout and can be directly accessed through compiled languages (Rust) and DSLs(via LLVM) with raw pointer access for better performance and compatibility

[junrushao]

Should be good to go once CI greens up

[DarkSharpness]

This CI failures (seg-fault) make no sense to me. The latest commit is just a doc-string update. It's rather tricky as I failed to replay the seg-fault in my local environment. Still investigating on the root cause.