MapAllocator.h 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <c10/core/Allocator.h>
  4. #include <string_view>
  5. namespace at {
  6. enum MappedAllocatorModes {
  7. ALLOCATOR_MAPPED_SHARED = 1,
  8. ALLOCATOR_MAPPED_SHAREDMEM = 2,
  9. ALLOCATOR_MAPPED_EXCLUSIVE = 4,
  10. ALLOCATOR_MAPPED_NOCREATE = 8,
  11. ALLOCATOR_MAPPED_KEEPFD = 16,
  12. ALLOCATOR_MAPPED_FROMFD = 32,
  13. ALLOCATOR_MAPPED_UNLINK = 64
  14. };
  15. // Sentinel value/type to help distinguish the file descriptor constructor from
  16. // the non-file descriptor constructor
  17. enum WithFd { WITH_FD };
  18. TORCH_API std::string NewProcessWideShmHandle();
  19. class TORCH_API MapAllocator {
  20. public:
  21. MapAllocator(std::string_view filename, int flags, size_t size);
  22. MapAllocator(
  23. WithFd /*unused*/,
  24. std::string_view filename,
  25. int fd,
  26. int flags,
  27. size_t size);
  28. MapAllocator(const MapAllocator&) = delete;
  29. MapAllocator& operator=(const MapAllocator&) = delete;
  30. MapAllocator(MapAllocator&&) = delete;
  31. MapAllocator& operator=(MapAllocator&&) = delete;
  32. const char* filename() const {
  33. return filename_.c_str();
  34. }
  35. int fd() const {
  36. #ifdef _WIN32
  37. TORCH_CHECK(false, "MapAllocator::fd() is unsupported on Windows");
  38. #else
  39. return fd_;
  40. #endif
  41. }
  42. ptrdiff_t size() const {
  43. return size_;
  44. }
  45. // Return a pointer to the actual data for this allocator
  46. // (in the case of the refcounted allocator, this is offset
  47. // from the base pointer.)
  48. virtual void* data() const {
  49. return base_ptr_;
  50. }
  51. int flags() const {
  52. return flags_;
  53. }
  54. static MapAllocator* fromDataPtr(const at::DataPtr& /*dptr*/);
  55. static at::DataPtr makeDataPtr(
  56. std::string_view filename,
  57. int flags,
  58. size_t size,
  59. size_t* actual_size_out);
  60. static at::DataPtr makeDataPtr(
  61. WithFd /*unused*/,
  62. const char* filename,
  63. int fd,
  64. int flags,
  65. size_t size,
  66. size_t* actual_size_out);
  67. // Closes the data. Helps us avoid destructor shenanigans
  68. virtual void close();
  69. // This is very dangerous. You have to redefine this destructor for each
  70. // subclass
  71. virtual ~MapAllocator();
  72. protected:
  73. bool closed_ = false;
  74. std::string filename_;
  75. int flags_ = 0;
  76. ptrdiff_t size_; /* mapped size */
  77. #ifdef _WIN32
  78. void* handle_;
  79. void* event_;
  80. std::string eventname_;
  81. #else
  82. int fd_ = -1;
  83. #endif
  84. void* base_ptr_ = nullptr;
  85. };
  86. // Base-from-member idiom
  87. struct TORCH_API RefcountedMapAllocatorArgCheck {
  88. RefcountedMapAllocatorArgCheck(int flags);
  89. };
  90. class TORCH_API RefcountedMapAllocator : private RefcountedMapAllocatorArgCheck,
  91. public MapAllocator {
  92. public:
  93. RefcountedMapAllocator(const char* filename, int flags, size_t size);
  94. RefcountedMapAllocator(
  95. WithFd /*unused*/,
  96. const char* filename,
  97. int fd,
  98. int flags,
  99. size_t size);
  100. static RefcountedMapAllocator* fromDataPtr(const at::DataPtr& /*dptr*/);
  101. RefcountedMapAllocator(const RefcountedMapAllocator&) = delete;
  102. RefcountedMapAllocator(RefcountedMapAllocator&&) = delete;
  103. RefcountedMapAllocator& operator=(const RefcountedMapAllocator&) = delete;
  104. RefcountedMapAllocator& operator=(RefcountedMapAllocator&&) = delete;
  105. static at::DataPtr makeDataPtr(
  106. const char* filename,
  107. int flags,
  108. size_t size,
  109. size_t* actual_size_out);
  110. static at::DataPtr makeDataPtr(
  111. WithFd /*unused*/,
  112. const char* filename,
  113. int fd,
  114. int flags,
  115. size_t size,
  116. size_t* actual_size_out);
  117. void* data() const override;
  118. void incref();
  119. int decref();
  120. void close() override;
  121. ~RefcountedMapAllocator() override {
  122. RefcountedMapAllocator::close();
  123. }
  124. protected:
  125. void checkFlags();
  126. void initializeAlloc();
  127. };
  128. } // namespace at
  129. #else
  130. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  131. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)