Refactor make_boxed_from_unboxed_functor.h

extension/pybindings/cpp_extension.py

@@ -1,24 +1,32 @@
-from torch.utils import cpp_extension
 import os
 
+from torch.utils import cpp_extension
+
 _HERE = os.path.abspath(__file__)
 _EXECUTORCH_PATH = os.path.dirname(os.path.dirname(_HERE))
 
-def load_inline(name,
-                cpp_sources,
-                functions=None,
-                extra_cflags=None,
-                extra_ldflags=None,
-                extra_include_paths=None,
-                build_directory=None,
-                verbose=False,
-                is_python_module=True,
-                with_pytorch_error_handling=True,
-                keep_intermediates=True,
-                use_pch=False):
+
+def load_inline(
+    name,
+    cpp_sources,
+    functions=None,
+    extra_cflags=None,
+    extra_ldflags=None,
+    extra_include_paths=None,
+    build_directory=None,
+    verbose=False,
+    is_python_module=True,
+    with_pytorch_error_handling=True,
+    keep_intermediates=True,
+    use_pch=False,
+):
     # Register the code into PyTorch
     aten_extra_cflags = ["-DUSE_ATEN_LIB"] + (extra_cflags if extra_cflags else [])
-    extra_ldflags = [f"-L{_EXECUTORCH_PATH}", f"-Wl,-rpath,{_EXECUTORCH_PATH}", "-lexecutorch"] + (extra_ldflags if extra_ldflags else [])
+    extra_ldflags = [
+        f"-L{_EXECUTORCH_PATH}",
+        f"-Wl,-rpath,{_EXECUTORCH_PATH}",
+        "-lexecutorch",
+    ] + (extra_ldflags if extra_ldflags else [])
     module = cpp_extension.load_inline(
         name,
         cpp_sources,
@@ -37,13 +45,13 @@ def load_inline(name,
     cpp_extension.load_inline(
         name,
         cpp_sources,
-        functions=None, # leave this out since we are not passing out any python module
+        functions=None,  # leave this out since we are not passing out any python module
         extra_cflags=extra_cflags,
         extra_ldflags=extra_ldflags,
         extra_include_paths=extra_include_paths,
         build_directory=build_directory,
         verbose=verbose,
-        is_python_module=False, # don't register as a python module. Load shared library as a side effect.
+        is_python_module=False,  # don't register as a python module. Load shared library as a side effect.
         with_pytorch_error_handling=with_pytorch_error_handling,
         keep_intermediates=keep_intermediates,
         use_pch=use_pch,

runtime/kernel/make_boxed_from_unboxed_functor.h

@@ -19,8 +19,9 @@
 ///   return out;
 /// }
 ///
-/// Kernel my_kernel = Kernel.make_boxed_kernel("my_ns::my_op",
-/// EXECUTORCH_FN(my_op)); register_kernels({my_kernel});
+/// Kernel my_kernel = Kernel::make_boxed_kernel("my_ns::my_op",
+///   EXECUTORCH_FN(my_op));
+/// static auto res = register_kernels({my_kernel});
 /// ```
 ///
 /// The trick here is to convert each EValue to inferred argument type. This
@@ -34,109 +35,14 @@
 
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
-#include <executorch/runtime/kernel/type_list.h>
-#include <cstdlib>
-#include <memory>
-#include <type_traits>
-#include <typeinfo>
+#include <executorch/runtime/kernel/meta_programming.h>
 
 namespace torch {
 namespace executor {
 
 class KernelRuntimeContext; // Forward declaration
 using RuntimeContext = KernelRuntimeContext; // TODO(T147221312): Remove
 
-// Check if a given type is a function
-template <class T>
-struct is_function_type : std::false_type {};
-template <class Result, class... Args>
-struct is_function_type<Result(Args...)> : std::true_type {};
-template <class T>
-using is_function_type_t = typename is_function_type<T>::type;
-
-// A compile-time wrapper around a function pointer
-template <class FuncType_, FuncType_* func_ptr_>
-struct CompileTimeFunctionPointer final {
-  static_assert(
-      is_function_type<FuncType_>::value,
-      "EXECUTORCH_FN can only wrap function types.");
-  using FuncType = FuncType_;
-
-  static constexpr FuncType* func_ptr() {
-    return func_ptr_;
-  }
-};
-
-// Check if a given type is a compile-time function pointer
-template <class T>
-struct is_compile_time_function_pointer : std::false_type {};
-template <class FuncType, FuncType* func_ptr>
-struct is_compile_time_function_pointer<
-    CompileTimeFunctionPointer<FuncType, func_ptr>> : std::true_type {};
-
-#define EXECUTORCH_FN_TYPE(func)                                      \
-  CompileTimeFunctionPointer<                                         \
-      std::remove_pointer_t<std::remove_reference_t<decltype(func)>>, \
-      func>
-#define EXECUTORCH_FN(func) EXECUTORCH_FN_TYPE(func)()
-
-/**
- * strip_class: helper to remove the class type from pointers to `operator()`.
- */
-template <typename T>
-struct strip_class {};
-template <typename Class, typename Result, typename... Args>
-struct strip_class<Result (Class::*)(Args...)> {
-  using type = Result(Args...);
-};
-template <typename Class, typename Result, typename... Args>
-struct strip_class<Result (Class::*)(Args...) const> {
-  using type = Result(Args...);
-};
-template <typename T>
-using strip_class_t = typename strip_class<T>::type;
-
-/**
- * Access information about result type or arguments from a function type.
- * Example:
- * using A = function_traits<int (float, double)>::return_type // A == int
- * using A = function_traits<int (float, double)>::parameter_types::tuple_type
- * // A == tuple<float, double>
- */
-template <class Func>
-struct function_traits {
-  static_assert(
-      !std::is_same<Func, Func>::value,
-      "In function_traits<Func>, Func must be a plain function type.");
-};
-template <class Result, class... Args>
-struct function_traits<Result(Args...)> {
-  using func_type = Result(Args...);
-  using return_type = Result;
-  using parameter_types = typelist<Args...>;
-  static constexpr auto number_of_parameters = sizeof...(Args);
-};
-
-/**
- * infer_function_traits: creates a `function_traits` type for a simple
- * function (pointer) or functor (lambda/struct). Currently does not support
- * class methods.
- */
-template <typename Functor>
-struct infer_function_traits {
-  using type = function_traits<strip_class_t<decltype(&Functor::operator())>>;
-};
-template <typename Result, typename... Args>
-struct infer_function_traits<Result (*)(Args...)> {
-  using type = function_traits<Result(Args...)>;
-};
-template <typename Result, typename... Args>
-struct infer_function_traits<Result(Args...)> {
-  using type = function_traits<Result(Args...)>;
-};
-template <typename T>
-using infer_function_traits_t = typename infer_function_traits<T>::type;
-
 // evalue_to_arg
 template <class T>
 struct decay_if_not_tensor final {

runtime/kernel/meta_programming.h

@@ -0,0 +1,114 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+#if __cplusplus < 201703L
+#error "This header requires C++17"
+#endif
+
+#include <executorch/runtime/kernel/type_list.h>
+#include <cstdlib>
+#include <memory>
+#include <type_traits>
+#include <typeinfo>
+
+namespace torch {
+namespace executor {
+
+// Check if a given type is a function
+template <class T>
+struct is_function_type : std::false_type {};
+template <class Result, class... Args>
+struct is_function_type<Result(Args...)> : std::true_type {};
+template <class T>
+using is_function_type_t = typename is_function_type<T>::type;
+
+// A compile-time wrapper around a function pointer
+template <class FuncType_, FuncType_* func_ptr_>
+struct CompileTimeFunctionPointer final {
+  static_assert(
+      is_function_type<FuncType_>::value,
+      "EXECUTORCH_FN can only wrap function types.");
+  using FuncType = FuncType_;
+
+  static constexpr FuncType* func_ptr() {
+    return func_ptr_;
+  }
+};
+
+// Check if a given type is a compile-time function pointer
+template <class T>
+struct is_compile_time_function_pointer : std::false_type {};
+template <class FuncType, FuncType* func_ptr>
+struct is_compile_time_function_pointer<
+    CompileTimeFunctionPointer<FuncType, func_ptr>> : std::true_type {};
+
+#define EXECUTORCH_FN_TYPE(func)                                      \
+  CompileTimeFunctionPointer<                                         \
+      std::remove_pointer_t<std::remove_reference_t<decltype(func)>>, \
+      func>
+#define EXECUTORCH_FN(func) EXECUTORCH_FN_TYPE(func)()
+
+/**
+ * strip_class: helper to remove the class type from pointers to `operator()`.
+ */
+template <typename T>
+struct strip_class {};
+template <typename Class, typename Result, typename... Args>
+struct strip_class<Result (Class::*)(Args...)> {
+  using type = Result(Args...);
+};
+template <typename Class, typename Result, typename... Args>
+struct strip_class<Result (Class::*)(Args...) const> {
+  using type = Result(Args...);
+};
+template <typename T>
+using strip_class_t = typename strip_class<T>::type;
+
+/**
+ * Access information about result type or arguments from a function type.
+ * Example:
+ * using A = function_traits<int (float, double)>::return_type // A == int
+ * using A = function_traits<int (float, double)>::parameter_types::tuple_type
+ * // A == tuple<float, double>
+ */
+template <class Func>
+struct function_traits {
+  static_assert(
+      !std::is_same<Func, Func>::value,
+      "In function_traits<Func>, Func must be a plain function type.");
+};
+template <class Result, class... Args>
+struct function_traits<Result(Args...)> {
+  using func_type = Result(Args...);
+  using return_type = Result;
+  using parameter_types = typelist<Args...>;
+  static constexpr auto number_of_parameters = sizeof...(Args);
+};
+
+/**
+ * infer_function_traits: creates a `function_traits` type for a simple
+ * function (pointer) or functor (lambda/struct). Currently does not support
+ * class methods.
+ */
+template <typename Functor>
+struct infer_function_traits {
+  using type = function_traits<strip_class_t<decltype(&Functor::operator())>>;
+};
+template <typename Result, typename... Args>
+struct infer_function_traits<Result (*)(Args...)> {
+  using type = function_traits<Result(Args...)>;
+};
+template <typename Result, typename... Args>
+struct infer_function_traits<Result(Args...)> {
+  using type = function_traits<Result(Args...)>;
+};
+template <typename T>
+using infer_function_traits_t = typename infer_function_traits<T>::type;
+
+} // namespace executor
+} // namespace torch