(tvm) branch main updated: [FFI] Replace `Arg2Str` with a more powerful `for_each` (#18117)

[syfeng]

This is an automated email from the ASF dual-hosted git repository.

syfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 49b139be8b [FFI] Replace `Arg2Str` with a more powerful `for_each` 
(#18117)
49b139be8b is described below

commit 49b139be8b76aef69e537c44b22a8f370cb72901
Author: Twice <tw...@apache.org>
AuthorDate: Sun Jul 6 17:01:06 2025 +0800

    [FFI] Replace `Arg2Str` with a more powerful `for_each` (#18117)
    
    [FFI] Replace Arg2Str with a more powerful for_each
---
 ffi/include/tvm/ffi/base_details.h     | 21 ++++++++++++++++++++-
 ffi/include/tvm/ffi/function_details.h | 27 ++++++++-------------------
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/ffi/include/tvm/ffi/base_details.h 
b/ffi/include/tvm/ffi/base_details.h
index fb7be1a955..d078a5963a 100644
--- a/ffi/include/tvm/ffi/base_details.h
+++ b/ffi/include/tvm/ffi/base_details.h
@@ -29,6 +29,7 @@
 #include <tvm/ffi/endian.h>
 
 #include <cstddef>
+#include <type_traits>
 #include <utility>
 
 #if defined(_MSC_VER)
@@ -135,14 +136,32 @@ namespace tvm {
 namespace ffi {
 namespace details {
 
+// a dependent-name version of false, for static_assert
+template <typename>
+inline constexpr bool always_false = false;
+
 // for each iterator
 struct for_each_dispatcher {
   template <typename F, typename... Args, size_t... I>
   static void run(std::index_sequence<I...>, const F& f, Args&&... args) {  // 
NOLINT(*)
-    (f(I, std::forward<Args>(args)), ...);
+    if constexpr (std::conjunction_v<
+                      std::is_invocable<F, std::integral_constant<size_t, I>, 
Args>...>) {
+      (f(std::integral_constant<size_t, I>{}, std::forward<Args>(args)), ...);
+    } else if constexpr (std::conjunction_v<std::is_invocable<F, size_t, 
Args>...>) {
+      (f(I, std::forward<Args>(args)), ...);
+    } else if constexpr (std::conjunction_v<std::is_invocable<F, Args>...>) {
+      (f(std::forward<Args>(args)), ...);
+    } else {
+      static_assert(always_false<F>, "The function is not invocable with the 
provided arguments");
+    }
   }
 };
 
+// Three kinds of function F are acceptable in `for_each`:
+// 1. F(size_t, Arg): argument with its index
+// 2. F(Arg): just the argument
+// 3. F(std::integral_constant<size_t, I>, Arg): argument with its constexpr 
index
+// The third one can make the index available in template arguments and `if 
constexpr`.
 template <typename F, typename... Args>
 void for_each(const F& f, Args&&... args) {  // NOLINT(*)
   for_each_dispatcher::run(std::index_sequence_for<Args...>{}, f, 
std::forward<Args>(args)...);
diff --git a/ffi/include/tvm/ffi/function_details.h 
b/ffi/include/tvm/ffi/function_details.h
index d029c19dd1..34e9979d5d 100644
--- a/ffi/include/tvm/ffi/function_details.h
+++ b/ffi/include/tvm/ffi/function_details.h
@@ -36,23 +36,6 @@ namespace tvm {
 namespace ffi {
 namespace details {
 
-template <typename ArgType>
-struct Arg2Str {
-  template <size_t i>
-  TVM_FFI_INLINE static void Apply(std::ostream& os) {
-    using Arg = std::tuple_element_t<i, ArgType>;
-    if constexpr (i != 0) {
-      os << ", ";
-    }
-    os << i << ": " << Type2Str<Arg>::v();
-  }
-  template <size_t... I>
-  TVM_FFI_INLINE static void Run(std::ostream& os, std::index_sequence<I...>) {
-    using TExpander = int[];
-    (void)TExpander{0, (Apply<I>(os), 0)...};
-  }
-};
-
 template <typename T>
 static constexpr bool ArgSupported =
     (std::is_same_v<std::remove_const_t<std::remove_reference_t<T>>, Any> ||
@@ -78,10 +61,16 @@ struct FuncFunctorImpl {
 #endif
 
   TVM_FFI_INLINE static std::string Sig() {
-    using IdxSeq = std::make_index_sequence<sizeof...(Args)>;
     std::ostringstream ss;
     ss << "(";
-    Arg2Str<std::tuple<Args...>>::Run(ss, IdxSeq{});
+    for_each(
+        [&ss](auto i, const auto& v) {
+          if constexpr (i() != 0) {
+            ss << ", ";
+          }
+          ss << i() << ": " << v;
+        },
+        Type2Str<Args>::v()...);
     ss << ") -> " << Type2Str<R>::v();
     return ss.str();
   }