[tvm] branch main updated: [Runtime] Enhance PackedFunc Metaprogramming with `PackArgs` (#15595)

[yaxingcai]

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

yaxingcai 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 63b0c8fecc [Runtime] Enhance PackedFunc Metaprogramming with 
`PackArgs` (#15595)
63b0c8fecc is described below

commit 63b0c8fecc0f0afa87ff56e24e66ba14c75be498
Author: Junru Shao <junrus...@apache.org>
AuthorDate: Sun Aug 20 10:55:49 2023 -0700

    [Runtime] Enhance PackedFunc Metaprogramming with `PackArgs` (#15595)
    
    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 660c24284b..655325ebe1 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 {