MapAllocator.h 3.1 KB

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