From 53b7ae3e4040cb1f41e12547dd5eba9133240df6 Mon Sep 17 00:00:00 2001
From: Junru Shao <junrushao@apache.org>
Date: Sat, 19 Aug 2023 20:32:11 -0700
Subject: [PATCH] [Runtime] Enhance PackedFunc Metaprogramming with `PackArgs`

This PR introduces a convenient method `PackArgs` that allows packing
variadic C++ parameters into TVM's PackedFunc calling convention without
recursion using C++ parameter pack tricks.

As an example,

```C++
template <typename... Args>
void MyVariadicMethod(Args&&... args) {
  constexpr int kNumArgs = sizeof...(Args);
  TVMValue values[kNumArgs];
  int type_codes[kNumArgs];
  tvm::runtime::PackArgs(values, type_codes, std::forward<Args>(args)...);
  ...
}
```
---
 include/tvm/runtime/packed_func.h | 55 ++++++++++++++++++++++++-------
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h
index 660c24284b8d..655325ebe190 100644
--- a/include/tvm/runtime/packed_func.h
+++ b/include/tvm/runtime/packed_func.h
@@ -1284,13 +1284,27 @@ namespace parameter_pack {
 
 template <typename... EnumArgs>
 struct EnumeratedParamPack {
-  struct Invoke {
-    template <template <size_t i, typename TArgument> class Functor, typename... ExtraParams>
-    static void F(ExtraParams&&... extra_params) {
+  struct InvokeWithoutArg {
+    template <template <size_t i, typename TArgument> class Functor, typename ExtraParams>
+    static void F(ExtraParams&& extra_params) {
       using TExpander = int[];
       (void)TExpander{
           0,
-          (Functor<EnumArgs::i, typename EnumArgs::T>::F(extra_params...), 0)...,
+          (Functor<EnumArgs::i, typename EnumArgs::T>::F(std::forward<ExtraParams>(extra_params)),
+           0)...,
+      };
+    }
+  };
+  struct InvokeWithArg {
+    template <template <size_t i, typename TArgument> class Functor, typename ExtraParams,
+              typename... Params>
+    static void F(ExtraParams&& extra_params, Params&&... params) {
+      using TExpander = int[];
+      (void)TExpander{
+          0,
+          (Functor<EnumArgs::i, typename EnumArgs::T>::F(std::forward<ExtraParams>(extra_params),
+                                                         std::forward<Params>(params)),
+           0)...,
       };
     }
   };
@@ -1310,22 +1324,27 @@ struct EnumerateImpl {
 
   template <std::size_t... id>
   struct Zipper<std::integer_sequence<std::size_t, id...>> {
-    using T = EnumeratedParamPack<Item<id, Args>...>;
+    using WithoutArg = typename EnumeratedParamPack<Item<id, Args>...>::InvokeWithoutArg;
+    using WithArg = typename EnumeratedParamPack<Item<id, Args>...>::InvokeWithArg;
   };
 
  public:
-  using T = typename Zipper<std::index_sequence_for<Args...>>::T;
+  using WithoutArg = typename Zipper<std::index_sequence_for<Args...>>::WithoutArg;
+  using WithArg = typename Zipper<std::index_sequence_for<Args...>>::WithArg;
 };
 
 template <typename... Args>
-using Enumerate = typename EnumerateImpl<Args...>::T;
+using EnumerateWithoutArg = typename EnumerateImpl<Args...>::WithoutArg;
+
+template <typename... Args>
+using EnumerateWithArg = typename EnumerateImpl<Args...>::WithArg;
 
 template <typename... Args>
 struct ParamPack {
-  template <template <size_t i, typename TArgument> class Functor, typename... ExtraParams>
-  static void InvokeWithoutArg(ExtraParams&&... extra_params) {
-    Enumerate<Args...>::Invoke::template F<Functor, ExtraParams...>(
-        std::forward<ExtraParams>(extra_params)...);
+  template <template <size_t i, typename TArgument> class Functor, typename ExtraParams>
+  static void InvokeWithoutArg(ExtraParams&& extra_params) {
+    EnumerateWithoutArg<Args...>::template F<Functor, ExtraParams>(
+        std::forward<ExtraParams>(extra_params));
   }
 };
 
@@ -1622,6 +1641,20 @@ inline TVMRetValue PackedFunc::operator()(Args&&... args) const {
   return rv;
 }
 
+template <int i, typename T>
+struct TVMArgsSetterApply {
+  static TVM_ALWAYS_INLINE void F(TVMArgsSetter* setter, T&& value) {
+    (*setter)(i, std::forward<T>(value));
+  }
+};
+
+template <typename... Args>
+void TVM_ALWAYS_INLINE PackArgs(TVMValue* values, int* type_codes, Args&&... args) {
+  TVMArgsSetter setter(values, type_codes);
+  detail::parameter_pack::EnumerateWithArg<Args...>::template F<TVMArgsSetterApply>(
+      &setter, std::forward<Args>(args)...);
+}
+
 namespace detail {
 template <typename R, int nleft, int index, typename F>
 struct unpack_call_dispatcher {