Support tuple outputs in aten bridge (pytorch#16848)

Summary: Update the aten bridge logic in make_aten_functor_from_et... to
handle kernels that return tuples of tensors. This motivating use case
is the box_with_nms_limit custom op, which returns a tuple of 6 output
tensors.

Differential Revision: D91375148

Author: GregoryComer

extension/aten_util/make_aten_functor_from_et_functor.h

@@ -82,6 +82,17 @@ struct type_map<torch::executor::ArrayRef<T>> final {
   using type = at::ArrayRef<typename type_map<T>::type>;
 };
 
+// Tuple.
+template <class... Args>
+struct type_map<std::tuple<Args...>> final {
+  using type = std::tuple<typename type_map<Args>::type...>;
+};
+
+template <class... Args>
+struct type_map<std::tuple<Args...>&> final {
+  using type = std::tuple<typename type_map<Args>::type...>&;
+};
+
 template <typename T>
 struct remove_const_ref final {
   using type = std::remove_const_t<std::remove_reference_t<T>>;
@@ -232,6 +243,77 @@ struct type_convert<torch::executor::ArrayRef<F>, c10::ArrayRef<T>> final {
   }
 };
 
+// Tuples: birdirectional.
+template <class... F, class... T>
+struct type_convert<std::tuple<F...>, std::tuple<T...>> final {
+ public:
+  // We need to remove references from the output type because converters expect
+  // to return by value. For example, the tensor converter creates a temporary
+  // tensor and returns it by value - so we can't return an lvalue reference to
+  // it.
+  using InType = std::tuple<F...>;
+  using OutType = std::tuple<std::remove_reference_t<T>...>;
+
+  std::tuple<type_convert<F, std::remove_reference_t<T>>...> converters;
+
+  explicit type_convert(std::tuple<F...> value)
+      : converters(make_converters(value, std::index_sequence_for<F...>{})) {}
+
+  template <size_t... Is>
+  static decltype(converters) make_converters(
+      std::tuple<F...> value,
+      std::index_sequence<Is...>) {
+    return std::make_tuple(
+        // For each element in the input/output tuple, create a type_convert
+        // struct.
+        type_convert<
+            // Get the type of Ith element of the input and output.
+            std::tuple_element_t<Is, InType>,
+            std::tuple_element_t<Is, OutType>
+            // Instantiate the converter with the Ith element of the input.
+            >(std::get<Is>(value))...);
+  }
+
+  template <size_t... Is>
+  OutType call_impl(std::index_sequence<Is...>) {
+    // Invoke the converter on each element of the tuple.
+    return std::make_tuple(std::get<Is>(converters).call()...);
+  }
+  OutType call() {
+    return call_impl(std::make_index_sequence<sizeof...(F)>{});
+  }
+};
+
+template <class T>
+struct is_tuple : std::false_type {};
+
+template <class... Args>
+struct is_tuple<std::tuple<Args...>> : std::true_type {};
+
+// A utility struct to extract the out arguments from a function.
+// Returns a tuple of Args[N...], where N is the index of the first out
+// argument.
+template <size_t N, typename... Args>
+struct extract_out_args;
+
+template <size_t N, typename... Args>
+struct extract_out_args {
+  static_assert(sizeof...(Args) >= N, "Output index out of range.");
+  static constexpr size_t n_out = sizeof...(Args) - N;
+  using tuple_type = std::tuple<Args...>;
+
+  template <size_t... I>
+  static auto get_impl(tuple_type&& t, std::index_sequence<I...>) {
+    return std::forward_as_tuple(
+        std::get<N + I>(std::forward<tuple_type>(t))...);
+  }
+
+  static auto get(tuple_type&& t) {
+    return get_impl(
+        std::forward<tuple_type>(t), std::make_index_sequence<n_out>{});
+  }
+};
+
 template <class F, F f, typename N = int, N index = N(-1)>
 struct wrapper_impl;
 
@@ -244,17 +326,23 @@ struct wrapper_impl<R (*)(Args...), f, int, N> {
   using TupleConvertsType =
       std::tuple<type_convert<typename type_map<Args>::type, Args>...>;
   using TupleArgsType = std::tuple<typename type_map<Args>::type...>;
+
   static constexpr size_t num_args = sizeof...(Args);
+  static constexpr size_t num_out_args = sizeof...(Args) - N;
+  static constexpr bool is_output_tuple = is_tuple<ReturnType>::value;
   static_assert(
-      (N < num_args &&
-       std::is_same_v<
-           executorch::extension::kernel_util_internal::element_t<
-               N,
-               executorch::extension::kernel_util_internal::typelist<Args...>>,
-           R>) ||
+      N < num_args,
+      "The index of the out tensor can't be greater or equal to num_args.");
+  static_assert(
+      is_output_tuple ||
+          std::is_same_v<
+              executorch::extension::kernel_util_internal::element_t<
+                  N,
+                  executorch::extension::kernel_util_internal::typelist<
+                      Args...>>,
+              R> ||
           N == -1,
-      "The index of the out tensor can't be greater or equal to num_args and "
-      "the Nth argument type has to be the same as the return type.");
+      "The Nth argument type has to be the same as the return type.");
 
   static ReturnType wrap(typename type_map<Args>::type... args) {
     // The wrapped function that takes ATen argument types, convert them into
@@ -265,20 +353,26 @@ struct wrapper_impl<R (*)(Args...), f, int, N> {
         type_convert<typename type_map<Args>::type, Args>(args)...);
     R result =
         call_functor_with_args(converts, std::make_index_sequence<num_args>());
-    typename std::remove_reference<ReturnType>::type converted_result =
-        type_convert<R, ReturnType>(result).call();
+    auto converted_result = type_convert<R, ReturnType>(result).call();
     if constexpr (N == -1) {
       return converted_result;
-    } else {
+    } else if constexpr (is_output_tuple) {
+      auto out_args =
+          extract_out_args<N, typename type_map<Args>::type...>::get(
+              std::move(args_tuple));
+
+      return resize_and_copy_outputs(
+          std::move(converted_result),
+          std::move(out_args),
+          std::make_index_sequence<num_out_args>{});
+    } else { // Non-tuple return type.
       static_assert(
           std::is_same_v<
               typename std::remove_reference<ReturnType>::type,
               at::Tensor>,
           "Only support at::Tensor-like return");
       ReturnType out = std::get<N>(args_tuple);
-      at::native::resize_output(out, converted_result.sizes());
-      out.copy_(converted_result);
-      return out;
+      return resize_and_copy_output(converted_result, out);
     }
   }
 
@@ -289,6 +383,25 @@ struct wrapper_impl<R (*)(Args...), f, int, N> {
       std::index_sequence<indices...>) {
     return f(std::get<indices>(converts).call()...);
   }
+
+  template <class A>
+  static at::Tensor& resize_and_copy_output(
+      A converted_result,
+      at::Tensor& out) {
+    at::native::resize_output(out, converted_result.sizes());
+    out.copy_(converted_result);
+    return out;
+  }
+
+  template <class A, class B, size_t... Is>
+  static ReturnType resize_and_copy_outputs(
+      A&& converted_result,
+      B&& out,
+      std::index_sequence<Is...>) {
+    return std::forward_as_tuple(resize_and_copy_output(
+        std::get<Is>(std::forward<A>(converted_result)),
+        std::get<Is>(std::forward<B>(out)))...);
+  }
 };
 
 } // namespace internal

extension/aten_util/test/make_aten_functor_from_et_functor_test.cpp

@@ -90,6 +90,20 @@ Tensor& sum_arrayref_optional_tensor_out(
   return out;
 }
 
+std::tuple<Tensor&, Tensor&>
+add_constant_tuple_out(const Tensor& a, Tensor& out1, Tensor& out2) {
+  auto a_data = a.const_data_ptr<int32_t>();
+  auto out1_data = out1.mutable_data_ptr<int32_t>();
+  auto out2_data = out2.mutable_data_ptr<int32_t>();
+
+  for (int i = 0; i < a.size(0); i++) {
+    out1_data[i] = a_data[i] + 1;
+    out2_data[i] = a_data[i] + 2;
+  }
+
+  return {out1, out2};
+}
+
 Tensor& quantized_embedding_byte_out(
     const Tensor& weight,
     const Tensor& weight_scales,
@@ -134,6 +148,24 @@ TEST_F(MakeATenFunctorFromETFunctorTest, TestTypeMap_Tensor) {
                const at::Tensor&>::value));
 }
 
+TEST_F(MakeATenFunctorFromETFunctorTest, TestTypeMap_Tuple_Tensor2x) {
+  EXPECT_TRUE(
+      (std::is_same<
+          type_map<
+              std::tuple<torch::executor::Tensor, torch::executor::Tensor>>::
+              type,
+          std::tuple<at::Tensor, at::Tensor>>::value));
+}
+
+TEST_F(MakeATenFunctorFromETFunctorTest, TestTypeMap_Tuple_TensorRef3x) {
+  EXPECT_TRUE((std::is_same<
+               type_map<std::tuple<
+                   torch::executor::Tensor&,
+                   torch::executor::Tensor&,
+                   torch::executor::Tensor&>>::type,
+               std::tuple<at::Tensor&, at::Tensor&, at::Tensor&>>::value));
+}
+
 TEST_F(MakeATenFunctorFromETFunctorTest, TestTypeMap_Optionals) {
   // Scalar.
   EXPECT_TRUE((std::is_same<
@@ -189,6 +221,34 @@ TEST_F(MakeATenFunctorFromETFunctorTest, TestConvert_Tensor) {
   EXPECT_TRUE((std::is_same<decltype(at_out), at::Tensor>::value));
 }
 
+TEST_F(MakeATenFunctorFromETFunctorTest, TestConvert_TupleTensor) {
+  // Convert at to et.
+  at::Tensor at_in1 = torch::tensor({1});
+  at::Tensor at_in2 = torch::tensor({2});
+  auto et = type_convert<
+                std::tuple<at::Tensor, at::Tensor>,
+                std::tuple<torch::executor::Tensor, torch::executor::Tensor>>(
+                std::make_tuple(at_in1, at_in2))
+                .call();
+  EXPECT_TRUE((std::is_same<
+               decltype(et),
+               std::tuple<torch::executor::Tensor, torch::executor::Tensor>>::
+                   value));
+
+  // Convert et to at.
+  torch::executor::testing::TensorFactory<ScalarType::Int> tf;
+  torch::executor::Tensor et_in1 = tf.ones({3});
+  torch::executor::Tensor et_in2 = tf.ones({4});
+  auto at_out =
+      type_convert<
+          std::tuple<torch::executor::Tensor, torch::executor::Tensor>,
+          std::tuple<at::Tensor, at::Tensor>>(std::make_tuple(et_in1, et_in2))
+          .call();
+  EXPECT_TRUE(
+      (std::is_same<decltype(at_out), std::tuple<at::Tensor, at::Tensor>>::
+           value));
+}
+
 TEST_F(MakeATenFunctorFromETFunctorTest, TestConvert_OptionalScalar) {
   // Convert optional at to et.
   auto optional_at_in = std::optional<int64_t>();
@@ -305,6 +365,9 @@ TORCH_LIBRARY(my_op, m) {
   m.def("add_optional_tensor.out", WRAP_TO_ATEN(add_optional_tensor_out, 2));
   m.def("sum_arrayref_scalar.out", WRAP_TO_ATEN(sum_arrayref_scalar_out, 1));
   m.def("sum_arrayref_tensor.out", WRAP_TO_ATEN(sum_arrayref_tensor_out, 1));
+  m.def(
+      "add_constant_tuple.out(Tensor a, *, Tensor(a!) out1, Tensor(b!) out2) -> (Tensor(a!), Tensor(b!))",
+      WRAP_TO_ATEN(add_constant_tuple_out, 1));
   m.def(
       "sum_arrayref_optional_tensor.out",
       WRAP_TO_ATEN(sum_arrayref_optional_tensor_out, 1));
@@ -422,6 +485,35 @@ TEST_F(MakeATenFunctorFromETFunctorTest, TestWrap_ArrayRefOptional) {
   EXPECT_EQ(stack[0].toTensor().const_data_ptr<int64_t>()[0], 4);
 }
 
+TEST_F(MakeATenFunctorFromETFunctorTest, TestWrap_TupleOut) {
+  at::Tensor a =
+      torch::tensor({1, 2, 3}, torch::TensorOptions().dtype(torch::kInt32));
+  at::Tensor out1 =
+      torch::tensor({0, 0, 0}, torch::TensorOptions().dtype(torch::kInt32));
+  at::Tensor out2 =
+      torch::tensor({0, 0, 0}, torch::TensorOptions().dtype(torch::kInt32));
+
+  auto op = c10::Dispatcher::singleton().findSchema(
+      {"my_op::add_constant_tuple", "out"});
+  EXPECT_TRUE(op.has_value());
+  torch::jit::Stack stack = {a, out1, out2};
+  op.value().callBoxed(&stack);
+
+  EXPECT_EQ(stack.size(), 2);
+
+  // Verify out1 contains a + 1
+  auto result1 = stack[0].toTensor();
+  EXPECT_EQ(result1.const_data_ptr<int32_t>()[0], 2);
+  EXPECT_EQ(result1.const_data_ptr<int32_t>()[1], 3);
+  EXPECT_EQ(result1.const_data_ptr<int32_t>()[2], 4);
+
+  // Verify out2 contains a + 2
+  auto result2 = stack[1].toTensor();
+  EXPECT_EQ(result2.const_data_ptr<int32_t>()[0], 3);
+  EXPECT_EQ(result2.const_data_ptr<int32_t>()[1], 4);
+  EXPECT_EQ(result2.const_data_ptr<int32_t>()[2], 5);
+}
+
 TEST_F(MakeATenFunctorFromETFunctorTest, TestConvert_ConstRefOptionals) {
   // Test const optional scalar conversion
   const std::optional<int64_t> const_optional_at_in =