[PyTorch] Fix return value of IValue::to for Tensor/String (pytorch#51463)

Summary:
Pull Request resolved: pytorch#51463

We can make the return type of the `to()` template match the return type of toFoo() by using the same technique we use for `list_element_to_const_ref`. Also simplifies `list_element_to_const_ref`.
ghstack-source-id: 121363468

Test Plan:
CI

built and ran AdIndexer benchmark w/ batch size 1 under perf stat
--repeat 5 to make sure it didn't regress

Reviewed By: bhosmer

Differential Revision: D26163848

fbshipit-source-id: b8563263b9f9fa5311c7d7cedc89e28bc5badda0

Author: swolchok

aten/src/ATen/core/List.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <ATen/core/ivalue_to.h>
 #include <c10/macros/Macros.h>
 #include <c10/util/TypeTraits.h>
 #include <c10/util/TypeList.h>
@@ -55,26 +56,16 @@ bool operator==(const T& lhs, const ListElementReference<T, Iterator>& rhs);
 
 template<class T>
 struct ListElementConstReferenceTraits {
-  // In the general case, we cannot expose a true const reference to
-  // the contents of an IValue, so we copy.
-  using const_reference = T;
-};
-
-template<>
-struct ListElementConstReferenceTraits<std::string> {
-  using const_reference = const std::string&;
+  // In the general case, we use IValue::to().
+  using const_reference = typename c10::detail::ivalue_to_const_ref_overload_return<T>::type;
 };
 
+// There is no to() overload for c10::optional<std::string>.
 template<>
 struct ListElementConstReferenceTraits<c10::optional<std::string>> {
   using const_reference = c10::optional<std::reference_wrapper<const std::string>>;
 };
 
-template<>
-struct ListElementConstReferenceTraits<at::Tensor> {
-  using const_reference = const at::Tensor&;
-};
-
 template<class T, class Iterator>
 class ListElementReference final {
 public:

aten/src/ATen/core/List_inl.h

@@ -167,24 +167,12 @@ list_element_to_const_ref(const IValue& element) {
   return element.template to<T>();
 }
 
-template<>
-inline typename ListElementConstReferenceTraits<std::string>::const_reference
-list_element_to_const_ref<std::string>(const IValue& element) {
-  return element.toStringRef();
-}
-
 template<>
 inline typename ListElementConstReferenceTraits<c10::optional<std::string>>::const_reference
 list_element_to_const_ref<c10::optional<std::string>>(const IValue& element) {
   return element.toOptionalStringRef();
 }
 
-template<>
-inline typename ListElementConstReferenceTraits<at::Tensor>::const_reference
-list_element_to_const_ref<at::Tensor>(const IValue& element) {
-  return element.toTensor();
-}
-
 } // namespace impl
 
 template<class T>

aten/src/ATen/core/ivalue.h

@@ -2,6 +2,7 @@
 
 #include <ATen/core/TensorBody.h>
 #include <ATen/core/blob.h>
+#include <ATen/core/ivalue_to.h>
 #include <c10/util/C++17.h>
 #include <c10/util/intrusive_ptr.h>
 #include <torch/csrc/WindowsTorchApiMacro.h>
@@ -781,7 +782,7 @@ struct TORCH_API IValue final {
   template <typename T>
   T to() &&;
   template <typename T>
-  T to() const&;
+  typename c10::detail::ivalue_to_const_ref_overload_return<T>::type to() const&;
 
   // ToOptional: convert a IValue to the Optional obj that accepts both T and
   // None

aten/src/ATen/core/ivalue_inl.h

@@ -765,13 +765,13 @@ inline const ivalue::Object& IValue::toObjectRef() const {
 // toX method to IValue. These named methods are much more discoverable
 // than the to templated function.
 
-#define DEFINE_TO(type, method_name)       \
+#define DEFINE_TO(T, method_name)          \
   template <>                              \
-  inline type IValue::to<type>()&& {       \
+  inline T IValue::to<T>()&& {             \
     return std::move(*this).method_name(); \
   }                                        \
   template <>                              \
-  inline type IValue::to<type>() const& {  \
+  inline c10::detail::ivalue_to_const_ref_overload_return<T>::type IValue::to<T>() const& { \
     return this->method_name();            \
   }
 
@@ -1014,7 +1014,7 @@ inline T IValue::to() && {
 }
 
 template <typename T>
-inline T IValue::to() const& {
+inline typename c10::detail::ivalue_to_const_ref_overload_return<T>::type IValue::to() const& {
   return generic_to(*this, _fake_type<T>{});
 }
 

aten/src/ATen/core/ivalue_to.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <string>
+
+namespace at {
+class Tensor;
+} // namespace at
+
+namespace c10 {
+struct IValue;
+namespace detail {
+// Determine the return type of `IValue::to() const &`. It's a const
+// reference when possible and a copy otherwise. It is in this
+// separate header so that List can use it as well.
+template<typename T>
+struct ivalue_to_const_ref_overload_return {
+  using type = T;
+};
+
+template<>
+struct ivalue_to_const_ref_overload_return<at::Tensor> {
+  using type = const at::Tensor&;
+};
+
+template<>
+struct ivalue_to_const_ref_overload_return<std::string> {
+  using type = const std::string&;
+};
+
+template<>
+struct ivalue_to_const_ref_overload_return<IValue> {
+  using type = const IValue&;
+};
+
+} // namespace detail
+} // namespace c10