Load.h 908 B

1234567891011121314151617181920212223242526272829303132333435363738
  1. #pragma once
  2. #include <c10/macros/Macros.h>
  3. #include <cstring>
  4. namespace c10 {
  5. namespace detail {
  6. template <typename T>
  7. struct LoadImpl {
  8. C10_HOST_DEVICE static T apply(const void* src) {
  9. return *reinterpret_cast<const T*>(src);
  10. }
  11. };
  12. template <>
  13. struct LoadImpl<bool> {
  14. C10_HOST_DEVICE static bool apply(const void* src) {
  15. static_assert(sizeof(bool) == sizeof(char), "");
  16. // NOTE: [Loading boolean values]
  17. // Protect against invalid boolean values by loading as a byte
  18. // first, then converting to bool (see gh-54789).
  19. return *reinterpret_cast<const unsigned char*>(src);
  20. }
  21. };
  22. } // namespace detail
  23. template <typename T>
  24. C10_HOST_DEVICE T load(const void* src) {
  25. return c10::detail::LoadImpl<T>::apply(src);
  26. }
  27. template <typename scalar_t>
  28. C10_HOST_DEVICE scalar_t load(const scalar_t* src) {
  29. return c10::detail::LoadImpl<scalar_t>::apply(src);
  30. }
  31. } // namespace c10