[PyTorch] Allow const T& access to ListElementReference when possible (#83177)

`IValue` can provide const reference access to `Tensor`,
`std::string`, and `IValue`. Fixing this allows iterating over
`c10::List`s of these types without copies.

This was flagged on #67964 back in February; never found someone to
implement it so I just did it myself while going over my TODO list
after returning from vacation yesterday.

Note that we still need `ListElementReference::get()` because this
support does not work for `List<optional<Tensor>>`.

Differential Revision: [D38579022](https://our.internmc.facebook.com/intern/diff/D38579022/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83177
Approved by: https://github.com/ezyang
This commit is contained in:
Scott Wolchok 2022-08-10 10:13:57 -07:00 committed by PyTorch MergeBot
parent d58ced3db7
commit ce8716f59a
2 changed files with 12 additions and 4 deletions

View File

@ -69,7 +69,11 @@ struct ListElementConstReferenceTraits<c10::optional<std::string>> {
template<class T, class Iterator>
class ListElementReference final {
public:
operator T() const;
operator std::conditional_t<
std::is_reference<typename c10::detail::
ivalue_to_const_ref_overload_return<T>::type>::value,
const T&,
T>() const;
ListElementReference& operator=(T&& new_value) &&;

View File

@ -118,9 +118,13 @@ namespace detail {
namespace impl {
template<class T, class Iterator>
ListElementReference<T, Iterator>::operator T() const {
return c10::detail::list_element_to<T>(*iterator_);
template <class T, class Iterator>
ListElementReference<T, Iterator>::operator std::conditional_t<
std::is_reference<typename c10::detail::ivalue_to_const_ref_overload_return<
T>::type>::value,
const T&,
T>() const {
return iterator_->template to<T>();
}
template<class T, class Iterator>