Fork FunctionRef for pytree (#554)

Summary:
Pull Request resolved: #554

Now there are two FunctionRef classes,
(A) torch::executor::pytree::FunctionRef in executorch/extension/pytree/function_ref.h
(B) torch::executor::FunctionRef in executor/runtime/core/function_ref.h

Rationale: We are aiming for version (B), compared to (A), to have features going forward, and also planning to eventually delete it completely.

Reviewed By: dbort

Differential Revision: D49780124

fbshipit-source-id: 48d27757b0e33fd47287fefef0e622595b9e0c09

Author: digantdesai

extension/pytree/aten_util/targets.bzl

@@ -17,6 +17,7 @@ def define_common_targets():
         ],
         exported_deps = [
             "//executorch/extension/pytree:pytree",
+            "//executorch/runtime/platform:platform",
         ],
         compiler_flags = ["-Wno-missing-prototypes"],
         fbcode_deps = [

extension/pytree/function_ref.h

@@ -0,0 +1,158 @@
+/*
+ * 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.
+ */
+
+//===- llvm/ADT/STLFunctionalExtras.h - Extras for <functional> -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains some extension to <functional>.
+//
+// No library is required when using these functions.
+//
+//===----------------------------------------------------------------------===//
+//     Extra additions to <functional>
+//===----------------------------------------------------------------------===//
+
+/// An efficient, type-erasing, non-owning reference to a callable. This is
+/// intended for use as the type of a function parameter that is not used
+/// after the function in question returns.
+///
+/// This class does not own the callable, so it is not in general safe to store
+/// a FunctionRef.
+
+// torch::executor: modified from llvm::function_ref
+// see https://www.foonathan.net/2017/01/function-ref-implementation/
+
+#pragma once
+
+#include <cstdint>
+#include <type_traits>
+#include <utility>
+
+namespace torch {
+namespace executor {
+namespace pytree {
+
+//===----------------------------------------------------------------------===//
+//     Features from C++20
+//===----------------------------------------------------------------------===//
+
+template <typename T>
+struct remove_cvref {
+  using type =
+      typename std::remove_cv<typename std::remove_reference<T>::type>::type;
+};
+
+template <typename T>
+using remove_cvref_t = typename remove_cvref<T>::type;
+
+template <typename Fn>
+class FunctionRef;
+
+template <typename Ret, typename... Params>
+class FunctionRef<Ret(Params...)> {
+  Ret (*callback_)(const void* memory, Params... params) = nullptr;
+  union Storage {
+    void* callable;
+    Ret (*function)(Params...);
+  } storage_;
+
+ public:
+  FunctionRef() = default;
+  explicit FunctionRef(std::nullptr_t) {}
+
+  /**
+   * Case 1: A callable object passed by lvalue reference.
+   * Taking rvalue reference is error prone because the object will be always
+   * be destroyed immediately.
+   */
+  template <
+      typename Callable,
+      // This is not the copy-constructor.
+      typename std::enable_if<
+          !std::is_same<remove_cvref_t<Callable>, FunctionRef>::value,
+          int32_t>::type = 0,
+      // Avoid lvalue reference to non-capturing lambda.
+      typename std::enable_if<
+          !std::is_convertible<Callable, Ret (*)(Params...)>::value,
+          int32_t>::type = 0,
+      // Functor must be callable and return a suitable type.
+      // To make this container type safe, we need to ensure either:
+      // 1. The return type is void.
+      // 2. Or the resulting type from calling the callable is convertible to
+      // the declared return type.
+      typename std::enable_if<
+          std::is_void<Ret>::value ||
+              std::is_convertible<
+                  decltype(std::declval<Callable>()(std::declval<Params>()...)),
+                  Ret>::value,
+          int32_t>::type = 0>
+  explicit FunctionRef(Callable& callable)
+      : callback_([](const void* memory, Params... params) {
+          auto& storage = *static_cast<const Storage*>(memory);
+          auto& callable = *static_cast<Callable*>(storage.callable);
+          return static_cast<Ret>(callable(std::forward<Params>(params)...));
+        }) {
+    storage_.callable = &callable;
+  }
+
+  /**
+   * Case 2: A plain function pointer.
+   * Instead of storing an opaque pointer to underlying callable object,
+   * store a function pointer directly.
+   * Note that in the future a variant which coerces compatible function
+   * pointers could be implemented by erasing the storage type.
+   */
+  /* implicit */ FunctionRef(Ret (*ptr)(Params...))
+      : callback_([](const void* memory, Params... params) {
+          auto& storage = *static_cast<const Storage*>(memory);
+          return storage.function(std::forward<Params>(params)...);
+        }) {
+    storage_.function = ptr;
+  }
+
+  /**
+   * Case 3: Implicit conversion from lambda to FunctionRef.
+   * A common use pattern is like:
+   * void foo(FunctionRef<...>) {...}
+   * foo([](...){...})
+   * Here constructors for non const lvalue reference or function pointer
+   * would not work because they do not cover implicit conversion from rvalue
+   * lambda.
+   * We need to define a constructor for capturing temporary callables and
+   * always try to convert the lambda to a function pointer behind the scene.
+   */
+  template <
+      typename Function,
+      // This is not the copy-constructor.
+      typename std::enable_if<
+          !std::is_same<Function, FunctionRef>::value,
+          int32_t>::type = 0,
+      // Function is convertible to pointer of (Params...) -> Ret.
+      typename std::enable_if<
+          std::is_convertible<Function, Ret (*)(Params...)>::value,
+          int32_t>::type = 0>
+  /* implicit */ FunctionRef(const Function& function)
+      : FunctionRef(static_cast<Ret (*)(Params...)>(function)) {}
+
+  Ret operator()(Params... params) const {
+    return callback_(&storage_, std::forward<Params>(params)...);
+  }
+
+  explicit operator bool() const {
+    return callback_;
+  }
+};
+
+} // namespace pytree
+} // namespace executor
+} // namespace torch

extension/pytree/pytree.h

@@ -16,7 +16,8 @@
 #include <memory>
 #include <string>
 
-#include <executorch/runtime/core/function_ref.h>
+// NB: This is a local, pytree FunctionRef and not from the ExecuTorch runtime.
+#include <executorch/extension/pytree/function_ref.h>
 
 namespace torch {
 namespace executor {

extension/pytree/targets.bzl

@@ -10,13 +10,9 @@ def define_common_targets():
     runtime.cxx_library(
         name = "pytree",
         srcs = [],
-        exported_headers = ["pytree.h"],
+        exported_headers = ["pytree.h", "function_ref.h"],
         visibility = [
             "//executorch/...",
             "@EXECUTORCH_CLIENTS",
         ],
-        exported_deps = [
-            "//executorch/runtime/platform:platform",
-            "//executorch/runtime/core:core",
-        ],
     )

extension/pytree/test/TARGETS

@@ -11,6 +11,13 @@ cpp_unittest(
     deps = ["//executorch/extension/pytree:pytree"],
 )
 
+cpp_unittest(
+    name = "function_ref_test",
+    srcs = ["function_ref_test.cpp"],
+    supports_static_listing = True,
+    deps = ["//executorch/extension/pytree:pytree"],
+)
+
 python_unittest(
     name = "test",
     srcs = [

extension/pytree/test/function_ref_test.cpp

@@ -0,0 +1,90 @@
+/*
+ * 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.
+ */
+
+#include <gtest/gtest.h>
+
+#include <executorch/extension/pytree/function_ref.h>
+
+using namespace ::testing;
+
+namespace torch {
+namespace executor {
+namespace pytree {
+
+namespace {
+class Item {
+ private:
+  int32_t val_;
+  FunctionRef<void(int32_t&)> ref_;
+
+ public:
+  /* implicit */ Item(int32_t val, FunctionRef<void(int32_t&)> ref)
+      : val_(val), ref_(ref) {}
+
+  int32_t get() {
+    ref_(val_);
+    return val_;
+  }
+};
+
+void one(int32_t& i) {
+  i = 1;
+}
+
+} // namespace
+
+TEST(FunctionRefTest, CapturingLambda) {
+  auto one = 1;
+  auto f = [&](int32_t& i) { i = one; };
+  Item item(0, FunctionRef<void(int32_t&)>{f});
+  EXPECT_EQ(item.get(), 1);
+  // ERROR:
+  // Item item1(0, f);
+  // Item item2(0, [&](int32_t& i) { i = 2; });
+  // FunctionRef<void(int32_t&)> ref([&](int32_t&){});
+}
+
+TEST(FunctionRefTest, NonCapturingLambda) {
+  int32_t val = 0;
+  FunctionRef<void(int32_t&)> ref([](int32_t& i) { i = 1; });
+  ref(val);
+  EXPECT_EQ(val, 1);
+
+  val = 0;
+  auto lambda = [](int32_t& i) { i = 1; };
+  FunctionRef<void(int32_t&)> ref1(lambda);
+  ref1(val);
+  EXPECT_EQ(val, 1);
+
+  Item item(0, [](int32_t& i) { i = 1; });
+  EXPECT_EQ(item.get(), 1);
+
+  auto f = [](int32_t& i) { i = 1; };
+  Item item1(0, f);
+  EXPECT_EQ(item1.get(), 1);
+
+  Item item2(0, std::move(f));
+  EXPECT_EQ(item2.get(), 1);
+}
+
+TEST(FunctionRefTest, FunctionPointer) {
+  int32_t val = 0;
+  FunctionRef<void(int32_t&)> ref(one);
+  ref(val);
+  EXPECT_EQ(val, 1);
+
+  Item item(0, one);
+  EXPECT_EQ(item.get(), 1);
+
+  Item item1(0, &one);
+  EXPECT_EQ(item1.get(), 1);
+}
+
+} // namespace pytree
+} // namespace executor
+} // namespace torch