mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
7c1d5db1f3
Apply c10::load in more places. The function was introduced to cast a byte to valid boolean values, thus fixing the UBSAN errors. Pull Request resolved: https://github.com/pytorch/pytorch/pull/141740 Approved by: https://github.com/ezyang
39 lines
924 B
C++
39 lines
924 B
C++
#pragma once
|
|
#include <c10/macros/Macros.h>
|
|
#include <cstring>
|
|
|
|
namespace c10 {
|
|
namespace detail {
|
|
|
|
template <typename T>
|
|
struct LoadImpl {
|
|
C10_HOST_DEVICE static T apply(const void* src) {
|
|
return *reinterpret_cast<const T*>(src);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct LoadImpl<bool> {
|
|
C10_HOST_DEVICE static bool apply(const void* src) {
|
|
static_assert(sizeof(bool) == sizeof(char));
|
|
// NOTE: [Loading boolean values]
|
|
// Protect against invalid boolean values by loading as a byte
|
|
// first, then converting to bool (see gh-54789).
|
|
return *reinterpret_cast<const unsigned char*>(src);
|
|
}
|
|
};
|
|
|
|
} // namespace detail
|
|
|
|
template <typename T>
|
|
C10_HOST_DEVICE constexpr T load(const void* src) {
|
|
return c10::detail::LoadImpl<T>::apply(src);
|
|
}
|
|
|
|
template <typename scalar_t>
|
|
C10_HOST_DEVICE constexpr scalar_t load(const scalar_t* src) {
|
|
return c10::detail::LoadImpl<scalar_t>::apply(src);
|
|
}
|
|
|
|
} // namespace c10
|