From 02a02a23ee8323677e46c778b02717d1bc646edc Mon Sep 17 00:00:00 2001 From: PyTorch MergeBot Date: Tue, 10 Oct 2023 20:39:11 +0000 Subject: [PATCH] Revert "Move at::{Refcounted,}MapAllocator to c10 (#109881)" This reverts commit 0341deb1c720d8c908ed40e853eaacfc8ac37181. Reverted https://github.com/pytorch/pytorch/pull/109881 on behalf of https://github.com/albanD due to It does break buck build ([comment](https://github.com/pytorch/pytorch/pull/109881#issuecomment-1756195823)) --- BUILD.bazel | 3 +- aten/src/ATen/CMakeLists.txt | 22 + aten/src/ATen/MapAllocator.cpp | 611 +++++++++++++++++++ aten/src/ATen/MapAllocator.h | 138 ++++- build_variables.bzl | 1 + c10/CMakeLists.txt | 15 - c10/build.bzl | 1 - c10/core/MapAllocator.cpp | 952 ------------------------------ c10/core/MapAllocator.h | 140 ----- c10/core/build.bzl | 23 - caffe2/CMakeLists.txt | 2 + cmake/Codegen.cmake | 4 + torch/lib/libshm/CMakeLists.txt | 2 +- torch/lib/libshm/core.cpp | 10 +- torch/lib/libshm/libshm.h | 12 +- torch/lib/libshm_windows/core.cpp | 6 +- torch/lib/libshm_windows/libshm.h | 10 +- 17 files changed, 798 insertions(+), 1154 deletions(-) create mode 100644 aten/src/ATen/MapAllocator.cpp delete mode 100644 c10/core/MapAllocator.cpp delete mode 100644 c10/core/MapAllocator.h diff --git a/BUILD.bazel b/BUILD.bazel index dd06af1c6ab7..490f33fa0a5b 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1643,8 +1643,7 @@ cc_library( name = "shm", srcs = glob(["torch/lib/libshm/*.cpp"]), deps = [ - ":torch_headers", - "//c10", + ":torch", ], ) diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index d81b705424f0..105fe8a0b9bc 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -337,6 +337,28 @@ if(UNIX AND NOT APPLE) endif(NEED_LIBRT) endif(UNIX AND NOT APPLE) +if(UNIX) + set(CMAKE_EXTRA_INCLUDE_FILES "sys/mman.h") + CHECK_FUNCTION_EXISTS(mmap HAVE_MMAP) + if(HAVE_MMAP) + add_definitions(-DHAVE_MMAP=1) + endif(HAVE_MMAP) + # done for lseek: https://www.gnu.org/software/libc/manual/html_node/File-Position-Primitive.html + add_definitions(-D_FILE_OFFSET_BITS=64) + CHECK_FUNCTION_EXISTS(shm_open HAVE_SHM_OPEN) + if(HAVE_SHM_OPEN) + add_definitions(-DHAVE_SHM_OPEN=1) + endif(HAVE_SHM_OPEN) + CHECK_FUNCTION_EXISTS(shm_unlink HAVE_SHM_UNLINK) + if(HAVE_SHM_UNLINK) + add_definitions(-DHAVE_SHM_UNLINK=1) + endif(HAVE_SHM_UNLINK) + CHECK_FUNCTION_EXISTS(malloc_usable_size HAVE_MALLOC_USABLE_SIZE) + if(HAVE_MALLOC_USABLE_SIZE) + add_definitions(-DHAVE_MALLOC_USABLE_SIZE=1) + endif(HAVE_MALLOC_USABLE_SIZE) +endif(UNIX) + ADD_DEFINITIONS(-DUSE_EXTERNAL_MZCRC) if(NOT MSVC) diff --git a/aten/src/ATen/MapAllocator.cpp b/aten/src/ATen/MapAllocator.cpp new file mode 100644 index 000000000000..f0b87a4d1084 --- /dev/null +++ b/aten/src/ATen/MapAllocator.cpp @@ -0,0 +1,611 @@ +#include + +#include +#include +#include +#if ATOMIC_INT_LOCK_FREE == 2 +#define AT_ATOMIC_IPC_REFCOUNT 1 +#endif + +#include +#include + +/* stuff for mapped files */ +#ifdef _WIN32 +#include +#endif + +#if defined(HAVE_MMAP) +#include +#include +#include +#endif + +#if !defined(_MSC_VER) || defined(HAVE_MMAP) +#include +#include +#elif defined(_MSC_VER) +#include +#endif +#include + +namespace at { + +static constexpr int64_t map_alloc_alignment = 64; + +std::string NewProcessWideShmHandle() { + static std::atomic counter{0}; + static std::random_device rd; +#ifdef _MSC_VER + return fmt::format( + "/torch_{}_{}_{}", + GetCurrentProcessId(), + rd(), + counter.fetch_add(1, std::memory_order_relaxed)); +#else + return fmt::format( + "/torch_{}_{}_{}", + getpid(), + rd(), + counter.fetch_add(1, std::memory_order_relaxed)); +#endif +} +#if defined(_WIN32) || defined(HAVE_MMAP) + +namespace { +struct MapInfo { + std::atomic refcount; +}; + +constexpr const char* unknown_filename = "filename not specified"; +#ifdef _WIN32 +constexpr const char* unknown_eventname = "eventname not specified"; +#endif +} // namespace (anonymous) + +MapAllocator::MapAllocator(WithFd, c10::string_view filename, int fd, int flags, size_t size) + : filename_(filename.empty() ? unknown_filename : filename) + , flags_(0) // to be filled later + , size_(0) // to be filled later +#ifdef _WIN32 + , handle_(INVALID_HANDLE_VALUE) // to be filled later + , event_(INVALID_HANDLE_VALUE) // to be filled later + , eventname_(filename.empty() ? unknown_eventname : (std::string(filename) + "_event")) +#else + , fd_(fd) +#endif + , base_ptr_(nullptr) +{ + + if (!(flags & ALLOCATOR_MAPPED_SHARED) && !(flags & ALLOCATOR_MAPPED_SHAREDMEM)) { + flags &= ~ALLOCATOR_MAPPED_NOCREATE; + } + if ((flags ^ ALLOCATOR_MAPPED_EXCLUSIVE) == 0) { + TORCH_CHECK(false, "ALLOCATOR_MAPPED_EXCLUSIVE flag requires opening the file in shared mode"); + } +#ifdef _WIN32 + if (fd != -1) { + TORCH_CHECK(false, "MapAllocator_newWithFd is unsupported on Windows"); + } +#endif + flags_ = flags; + + // OK, now do the allocation + + if (size == 0) { + return; + } + +#ifdef _WIN32 + if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { + // Shadowing + const wchar_t *filename; + const wchar_t *eventname; + const std::wstring wFilename = c10::u8u16(filename_); + const std::wstring wEventname = c10::u8u16(eventname_); + LARGE_INTEGER hfilesz; + + if (filename_[0] == '/') { + filename = wFilename.c_str() + 1; + eventname = wEventname.c_str() + 1; + } else { + filename = wFilename.c_str(); + eventname = wEventname.c_str(); + } + + hfilesz.QuadPart = size; + + if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { + event_ = CreateEventW(nullptr, FALSE, FALSE, eventname); + } else if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { + event_ = OpenEventW(EVENT_ALL_ACCESS, FALSE, eventname); + } else { + TORCH_CHECK(false, "Expected either ALLOCATOR_MAPPED_EXCLUSIVE or ALLOCATOR_MAPPED_NOCREATE"); + } + + if (event_ == nullptr) { + TORCH_CHECK(false, "Couldn't open shared event: <", eventname, ">, error code: <", GetLastError(), ">"); + } + + if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { + handle_ = CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, hfilesz.HighPart, hfilesz.LowPart, filename); + } else if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { + handle_ = OpenFileMappingW(FILE_MAP_ALL_ACCESS, FALSE, filename); + } else { + TORCH_CHECK(false, "Expected either ALLOCATOR_MAPPED_EXCLUSIVE or ALLOCATOR_MAPPED_NOCREATE"); + } + + if (handle_ == nullptr) { + TORCH_CHECK(false, "Couldn't open shared file mapping: <", filename, ">, error code: <", GetLastError(), ">"); + } + + size_ = size; + base_ptr_ = MapViewOfFile(handle_, FILE_MAP_ALL_ACCESS, 0, 0, size); + if (!base_ptr_) { + TORCH_CHECK(false, "Couldn't map view of shared file <", filename, ">, error code: <", GetLastError(), ">"); + } + } else { + + HANDLE hfile; + HANDLE hmfile; + LARGE_INTEGER hfilesz; + + if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { + TORCH_CHECK(false, "exclusive file mapping is not supported on Windows"); + } + if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { + TORCH_CHECK(false, "file mapping without creation is not supported on Windows"); + } + if (flags_ & ALLOCATOR_MAPPED_KEEPFD) { + TORCH_CHECK(false, "ALLOCATOR_MAPPED_KEEPFD not supported on Windows"); + } + if (flags_ & ALLOCATOR_MAPPED_FROMFD) { + TORCH_CHECK(false, "ALLOCATOR_MAPPED_FROMFD not supported on Windows"); + } + + // Shadowing + const wchar_t *filename; + const std::wstring wFilename = c10::u8u16(filename_); + + filename = wFilename.c_str(); + + /* open file */ + /* FILE_FLAG_RANDOM_ACCESS ? */ + if (flags_) { + hfile = CreateFileW(filename, GENERIC_READ|GENERIC_WRITE, FILE_SHARE_WRITE|FILE_SHARE_READ, 0, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0); + if (hfile == INVALID_HANDLE_VALUE) { + TORCH_CHECK(false, "could not open file <", filename_, "> in read-write mode; error code: <", GetLastError(), ">"); + } + } else { + hfile = CreateFileW(filename, GENERIC_READ, FILE_SHARE_WRITE|FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0); + if (hfile == INVALID_HANDLE_VALUE) { + TORCH_CHECK(false, "could not open file <", filename_, "> in read-only mode; error code: <", GetLastError(), ">"); + } + } + + if (GetFileSizeEx(hfile, &hfilesz) == 0) { + TORCH_CHECK(false, "could not get file size: <", filename_, ">; error code: <", GetLastError(), ">"); + } + + if (size > 0) { + if (size > hfilesz.QuadPart) { + if (flags_) { + hfilesz.QuadPart = size; + if (SetFilePointerEx(hfile, hfilesz, NULL, FILE_BEGIN) == 0) { + CloseHandle(hfile); + TORCH_CHECK(false, "unable to stretch file <", filename_, "> to the right size; error code: <", GetLastError(), ">", filename_); + } + if (SetEndOfFile(hfile) == 0) { + CloseHandle(hfile); + TORCH_CHECK(false, "unable to write to file <", filename_, ">; error code: <", GetLastError(), ">"); + } + } else { + CloseHandle(hfile); + TORCH_CHECK(false, "file <", filename_, "> size is smaller than the required mapping size <", size, ">; error code: <", GetLastError(), ">"); + } + } + } else { + size = hfilesz.QuadPart; + } + + size_ = size; /* if we are here, it must be the right size */ + + hfilesz.QuadPart = size_; + + /* get map handle */ + if (flags_) { + if ( (hmfile = CreateFileMappingW(hfile, NULL, PAGE_READWRITE, hfilesz.HighPart, hfilesz.LowPart, NULL)) == NULL ) { + TORCH_CHECK(false, "could not create a map on file <", filename_, ">; error code: <", GetLastError(), ">"); + } + } else { + if ( (hmfile = CreateFileMappingW(hfile, NULL, PAGE_WRITECOPY, hfilesz.HighPart, hfilesz.LowPart, NULL)) == NULL ) { + TORCH_CHECK(false, "could not create a map on file <", filename_, ">; error code: <", GetLastError(), ">"); + } + } + + /* map the stuff */ + if(flags_) { + base_ptr_ = MapViewOfFile(hmfile, FILE_MAP_ALL_ACCESS, 0, 0, 0); + } else { + base_ptr_ = MapViewOfFile(hmfile, FILE_MAP_COPY, 0, 0, 0); + } + + CloseHandle(hfile); + CloseHandle(hmfile); + } +#else /* _WIN32 */ + { + /* open file */ + int fd{-1}; + int flags{}; // shadow + + if (flags_ & (ALLOCATOR_MAPPED_SHARED | ALLOCATOR_MAPPED_SHAREDMEM)) { + flags = O_RDWR | O_CREAT; + } else { + flags = O_RDONLY; + } + + if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { + flags |= O_EXCL; + } + if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { + flags &= ~O_CREAT; + } + + if (!(flags_ & ALLOCATOR_MAPPED_FROMFD)) { + if (flags_ & ALLOCATOR_MAPPED_SHARED) { + if ((fd = open(filename_.c_str(), flags, (mode_t)0600)) == -1) { + TORCH_CHECK(false, "unable to open file <", filename_, "> in read-write mode: ", strerror(errno), " (", errno, ")"); + } + } else if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { +#ifdef HAVE_SHM_OPEN + if((fd = shm_open(filename_.c_str(), flags, (mode_t)0600)) == -1) { + TORCH_CHECK(false, "unable to open shared memory object <", filename_, "> in read-write mode: ", strerror(errno), " (", errno, ")"); + } +#else + TORCH_CHECK(false, "unable to open file <", filename_, "> in sharedmem mode, shm_open unavailable on this platform"); +#endif + } else { + if ((fd = open(filename_.c_str(), O_RDONLY)) == -1) { + TORCH_CHECK(false, "unable to open file <", filename_, "> in read-only mode: ", strerror(errno), " (", errno, ")"); + } + } + } else { + fd = fd_; + } + + struct stat file_stat; + if (fstat(fd, &file_stat) == -1) { + int last_err = errno; + if (!(flags_ & ALLOCATOR_MAPPED_FROMFD)) { + ::close(fd); + } + TORCH_CHECK(false, "unable to stat the file <", filename_, ">: ", strerror(last_err), " (", last_err, ")"); + } + + if (size > 0) { + if (static_cast(size) > file_stat.st_size) { + if (flags_) { + if (ftruncate(fd, size) == -1) { + TORCH_CHECK(false, "unable to resize file <", filename_, "> to the right size: ", strerror(errno), " (", errno, ")"); + } + if (fstat(fd, &file_stat) == -1 || file_stat.st_size < static_cast(size)) { + int last_err = errno; + ::close(fd); + TORCH_CHECK(false, "unable to stretch file <", filename_, "> to the right size: ", strerror(last_err), " (", last_err, ")"); + } +/* on macOS write returns with errno 45 (Opperation not supported) when used + * with a file descriptor obtained via shm_open + */ +#ifndef __APPLE__ + if ((write(fd, "", 1)) != 1) /* note that the string "" contains the '\0' byte ... */ { + int last_err = errno; + ::close(fd); + TORCH_CHECK(false, "unable to write to file <", filename_, ">: ", strerror(last_err), " (", last_err, ")"); + } +#endif + } else { + ::close(fd); + TORCH_CHECK(false, "file <", filename_, "> size is smaller than the required mapping size <", size, ">"); + } + } + } else { + size = file_stat.st_size; + } + + size_ = size; /* if we are here, it must be the right size */ + + /* map it */ + if (flags_ & (ALLOCATOR_MAPPED_SHARED | ALLOCATOR_MAPPED_SHAREDMEM)) { + base_ptr_ = mmap(nullptr, size_, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0); + } else { + base_ptr_ = mmap(nullptr, size_, PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0); + } + + if (base_ptr_ == MAP_FAILED) { + base_ptr_ = nullptr; /* let's be sure it is NULL */ + TORCH_CHECK(false, "unable to mmap ", size_, " bytes from file <", filename_, ">: ", strerror(errno), " (", errno, ")"); + } + + if (flags_ & ALLOCATOR_MAPPED_KEEPFD) { + fd_ = fd; + } else { + if (::close(fd) == -1) { + TORCH_CHECK(false, "Error closing file <", filename_, ">: ", strerror(errno), " (", errno, ")"); + } + fd_ = -1; + } + + if (flags_ & ALLOCATOR_MAPPED_UNLINK) { + if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { +#ifdef HAVE_SHM_UNLINK + if (shm_unlink(filename_.c_str()) == -1) { + TORCH_CHECK(false, "could not unlink the shared memory file ", filename_, " : ", strerror(errno), " (", errno, ")"); + } +#else + TORCH_CHECK(false, "could not unlink the shared memory file ", filename_, ", shm_unlink not available on platform"); +#endif + } else { + if (unlink(filename_.c_str()) == -1) + TORCH_CHECK(false, "could not unlink file ", filename_, " : ", strerror(errno), " (", errno, ")"); + } + } + + if (base_ptr_ == MAP_FAILED) { + TORCH_CHECK(false, "$ Torch: unable to mmap memory: you tried to mmap ", size_/1073741824, " GB."); + } + } +#endif + c10::reportMemoryUsageToProfiler(base_ptr_, size_, 0, size_, c10::Device(c10::DeviceType::CPU)); +} + +MapAllocator::MapAllocator(c10::string_view filename, int flags, size_t size) + : MapAllocator(WITH_FD, filename, -1, flags, size) +{} + +#ifdef _WIN32 +struct ReleaseContext { + HANDLE event; + HANDLE handle; + HANDLE wait; +}; +static void CALLBACK WaitForReleaseHandle(PVOID lpParam, BOOLEAN TimerOrWaitFired) +{ + if (lpParam) { + ReleaseContext *ctx = (ReleaseContext *)lpParam; + + SetEvent(ctx->event); + CloseHandle(ctx->event); + CloseHandle(ctx->handle); + + UnregisterWait(ctx->wait); + + delete ctx; + } +} +#endif + +void MapAllocator::close() { + if (closed_) { + return; + } + closed_ = true; + if (base_ptr_ == nullptr) { + return; + } +#ifdef _WIN32 + if ((flags_ & ALLOCATOR_MAPPED_KEEPFD) || (flags_ & ALLOCATOR_MAPPED_SHAREDMEM)) + CloseHandle(handle_); + if(UnmapViewOfFile(base_ptr_) == 0) + TORCH_CHECK(false, "could not unmap the shared memory file"); +#else /* _WIN32 */ + if (flags_ & ALLOCATOR_MAPPED_KEEPFD) { + if (::close(fd_) == -1) { + TORCH_CHECK(false, "could not close file descriptor ", fd_, " :", strerror(errno), " (", errno, ")" ); + } + } + + if (munmap(base_ptr_, size_)) { + TORCH_CHECK(false, "could not unmap the shared memory file: ", strerror(errno), " (", errno, ")"); + } + + if (!(flags_ & (ALLOCATOR_MAPPED_FROMFD | ALLOCATOR_MAPPED_UNLINK))) { + if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { +#ifdef HAVE_SHM_UNLINK + if (shm_unlink(filename_.c_str()) == -1) { + TORCH_CHECK(false, "could not unlink the shared memory file ", filename_, " : ", strerror(errno), " (", errno, ")"); + } +#else + TORCH_CHECK(false, "could not unlink the shared memory file ", filename_, ", shm_unlink not available on platform"); +#endif + } + } +#endif /* _WIN32 */ +} + +#else /* defined(_WIN32) || defined(HAVE_MMAP) */ + +MapAllocator::MapAllocator(c10::string_view filename, int flags, size_t size) { + TORCH_CHECK(false, "file mapping not supported on your system"); +} + +MapAllocator::MapAllocator(WithFd, c10::string_view filename, int fd, int flags, size_t size) { + TORCH_CHECK(false, "file mapping not supported on your system"); +} + +void MapAllocator::close() { } + +#endif + +#if (defined(_WIN32) || defined(HAVE_MMAP)) && defined(AT_ATOMIC_IPC_REFCOUNT) + +RefcountedMapAllocatorArgCheck::RefcountedMapAllocatorArgCheck(int flags) { + if (flags & ALLOCATOR_MAPPED_FROMFD) { + TORCH_CHECK(false, "RefcountedMapAllocator doesn't support ALLOCATOR_MAPPED_FROMFD flag"); + } + if (flags & ALLOCATOR_MAPPED_KEEPFD) { + TORCH_CHECK(false, "RefcountedMapAllocator doesn't support ALLOCATOR_MAPPED_KEEPFD flag"); + } + if (flags & ALLOCATOR_MAPPED_UNLINK) { + TORCH_CHECK(false, "RefcountedMapAllocator doesn't support ALLOCATOR_MAPPED_UNLINK flag"); + } + if (!(flags & ALLOCATOR_MAPPED_SHAREDMEM)) { + TORCH_CHECK(false, "RefcountedMapAllocator requires ALLOCATOR_MAPPED_SHAREDMEM flag"); + } +} + +RefcountedMapAllocator::RefcountedMapAllocator(const char *filename, int flags, size_t size) + : RefcountedMapAllocatorArgCheck(flags) + , MapAllocator(filename, flags, size + map_alloc_alignment) { + + initializeAlloc(); +} +RefcountedMapAllocator::RefcountedMapAllocator(WithFd, const char *filename, int fd, int flags, size_t size) + : RefcountedMapAllocatorArgCheck(flags) + , MapAllocator(WITH_FD, filename, flags, fd, size + map_alloc_alignment) { + + initializeAlloc(); +} + +void RefcountedMapAllocator::initializeAlloc() { + TORCH_CHECK(base_ptr_, "base_ptr_ is null"); + MapInfo *map_info = (MapInfo*)base_ptr_; + +#ifdef _WIN32 + ReleaseContext* r_ctx = new ReleaseContext; + r_ctx->handle = handle_; + r_ctx->event = event_; + r_ctx->wait = NULL; + BOOL can_wait = RegisterWaitForSingleObject(&r_ctx->wait, event_, WaitForReleaseHandle, (PVOID)r_ctx, INFINITE, WT_EXECUTEONLYONCE); + TORCH_CHECK(can_wait, "Couldn't register wait on event, error code: <", GetLastError(), ">"); +#endif + + if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { + new (&map_info->refcount) std::atomic(1); + } else { + map_info->refcount++; + } +} + +void RefcountedMapAllocator::close() { + if (closed_) { + return; + } + closed_ = true; + + void* data = base_ptr_; + +#ifdef _WIN32 + MapInfo *info = (MapInfo*)data; + if (--info->refcount == 0) { + SetEvent(event_); + } + if(UnmapViewOfFile(data) == 0) { + TORCH_CHECK(false, "could not unmap the shared memory file"); + } +#else /* _WIN32 */ + + MapInfo *info = (MapInfo*)(data); + if (--info->refcount == 0) { +#ifdef HAVE_SHM_UNLINK + if (shm_unlink(filename_.c_str()) == -1) { + TORCH_CHECK(false, "could not unlink the shared memory file ", filename_); + } +#else + TORCH_CHECK(false, "could not unlink the shared memory file ", filename_, ", shm_unlink not available on platform"); +#endif /* HAVE_SHM_UNLINK */ + } + if (munmap(info, size_)) { + TORCH_CHECK(false, "could not unmap the shared memory file ", filename_); + } +#endif /* _WIN32 */ +} + +void RefcountedMapAllocator::incref() +{ + MapInfo *map_info = static_cast(base_ptr_); + ++map_info->refcount; +} + +int RefcountedMapAllocator::decref() +{ + MapInfo *map_info = static_cast(base_ptr_); + return --map_info->refcount == 0; +} + +#else + + +RefcountedMapAllocatorArgCheck::RefcountedMapAllocatorArgCheck(int flags) {} + +RefcountedMapAllocator::RefcountedMapAllocator(const char *filename, int flags, size_t size) + : RefcountedMapAllocatorArgCheck(flags), + MapAllocator(filename, flags, size + map_alloc_alignment) +{ + TORCH_CHECK(false, "refcounted file mapping not supported on your system"); +} + +RefcountedMapAllocator::RefcountedMapAllocator(WithFd, const char *filename, int fd, int flags, size_t size) + : RefcountedMapAllocatorArgCheck(flags), + MapAllocator(WITH_FD, filename, flags, fd, size + map_alloc_alignment) +{ + TORCH_CHECK(false, "refcounted file mapping not supported on your system"); +} + +void RefcountedMapAllocator::initializeAlloc() {} + +void RefcountedMapAllocator::close() {} + +#endif + +static void deleteMapAllocator(void* ptr) { + delete static_cast(ptr); +} + +static void deleteRefcountedMapAllocator(void* ptr) { + delete static_cast(ptr); +} + +MapAllocator* MapAllocator::fromDataPtr(const at::DataPtr& dptr) { + return dptr.cast_context(&deleteMapAllocator); +} + +RefcountedMapAllocator* RefcountedMapAllocator::fromDataPtr(const at::DataPtr& dptr) { + return dptr.cast_context(&deleteRefcountedMapAllocator); +} + +at::DataPtr MapAllocator::makeDataPtr(c10::string_view filename, int flags, size_t size, size_t* actual_size_out) { + auto* context = new MapAllocator(filename, flags, size); + if (actual_size_out) *actual_size_out = context->size(); + return {context->data(), context, &deleteMapAllocator, at::DeviceType::CPU}; +} + +at::DataPtr MapAllocator::makeDataPtr(WithFd, const char *filename, int fd, int flags, size_t size, size_t* actual_size_out) { + auto* context = new MapAllocator(WITH_FD, filename, fd, flags, size); + if (actual_size_out) *actual_size_out = context->size(); + return {context->data(), context, &deleteMapAllocator, at::DeviceType::CPU}; +} + +at::DataPtr RefcountedMapAllocator::makeDataPtr(const char *filename, int flags, size_t size, size_t* actual_size_out) { + auto* context = new RefcountedMapAllocator(filename, flags, size); + if (actual_size_out) *actual_size_out = context->size() - map_alloc_alignment; + return {context->data(), context, &deleteRefcountedMapAllocator, at::DeviceType::CPU}; +} + +at::DataPtr RefcountedMapAllocator::makeDataPtr(WithFd, const char *filename, int fd, int flags, size_t size, size_t* actual_size_out) { + auto* context = new RefcountedMapAllocator(WITH_FD, filename, fd, flags, size); + if (actual_size_out) *actual_size_out = context->size() - map_alloc_alignment; + return {context->data(), context, &deleteRefcountedMapAllocator, at::DeviceType::CPU}; +} + +void* RefcountedMapAllocator::data() const { + return static_cast(static_cast(base_ptr_) + map_alloc_alignment); +} + +MapAllocator::~MapAllocator() { + // NOLINTNEXTLINE(clang-analyzer-optin.cplusplus.VirtualCall) + close(); + c10::reportMemoryUsageToProfiler(base_ptr_, -size_, 0, 0, c10::Device(c10::DeviceType::CPU)); +} + +} // namespace at diff --git a/aten/src/ATen/MapAllocator.h b/aten/src/ATen/MapAllocator.h index b074b2e8be16..3354ab84577f 100644 --- a/aten/src/ATen/MapAllocator.h +++ b/aten/src/ATen/MapAllocator.h @@ -1,3 +1,139 @@ #pragma once -#include +#include +#include + +namespace at { + +enum MappedAllocatorModes { + ALLOCATOR_MAPPED_SHARED = 1, + ALLOCATOR_MAPPED_SHAREDMEM = 2, + ALLOCATOR_MAPPED_EXCLUSIVE = 4, + ALLOCATOR_MAPPED_NOCREATE = 8, + ALLOCATOR_MAPPED_KEEPFD = 16, + ALLOCATOR_MAPPED_FROMFD = 32, + ALLOCATOR_MAPPED_UNLINK = 64 +}; + +// Sentinel value/type to help distinguish the file descriptor constructor from +// the non-file descriptor constructor +enum WithFd { WITH_FD }; + +TORCH_API std::string NewProcessWideShmHandle(); + +class TORCH_API MapAllocator { + public: + MapAllocator(c10::string_view filename, int flags, size_t size); + MapAllocator( + WithFd, + c10::string_view filename, + int fd, + int flags, + size_t size); + MapAllocator(const MapAllocator&) = delete; + MapAllocator& operator=(const MapAllocator&) = delete; + MapAllocator(MapAllocator&&) = delete; + MapAllocator& operator=(MapAllocator&&) = delete; + + const char* filename() const { + return filename_.c_str(); + } + int fd() const { +#ifdef _WIN32 + TORCH_CHECK(false, "MapAllocator::fd() is unsupported on Windows"); +#else + return fd_; +#endif + } + ptrdiff_t size() const { + return size_; + } + // Return a pointer to the actual data for this allocator + // (in the case of the refcounted allocator, this is offset + // from the base pointer.) + virtual void* data() const { + return base_ptr_; + } + + static MapAllocator* fromDataPtr(const at::DataPtr&); + static at::DataPtr makeDataPtr( + c10::string_view filename, + int flags, + size_t size, + size_t* actual_size_out); + static at::DataPtr makeDataPtr( + WithFd, + const char* filename, + int fd, + int flags, + size_t size, + size_t* actual_size_out); + + // Closes the data. Helps us avoid destructor shenanigans + virtual void close(); + + // This is very dangerous. You have to redefine this destructor for each + // subclass + virtual ~MapAllocator(); + + protected: + bool closed_ = false; + std::string filename_; + int flags_ = 0; + ptrdiff_t size_; /* mapped size */ +#ifdef _WIN32 + void* handle_; + void* event_; + std::string eventname_; +#else + int fd_ = -1; +#endif + void* base_ptr_ = nullptr; +}; + +// Base-from-member idiom +struct TORCH_API RefcountedMapAllocatorArgCheck { + RefcountedMapAllocatorArgCheck(int flags); +}; + +class TORCH_API RefcountedMapAllocator : private RefcountedMapAllocatorArgCheck, + public MapAllocator { + public: + RefcountedMapAllocator(const char* filename, int flags, size_t size); + RefcountedMapAllocator( + WithFd, + const char* filename, + int fd, + int flags, + size_t size); + + static RefcountedMapAllocator* fromDataPtr(const at::DataPtr&); + static at::DataPtr makeDataPtr( + const char* filename, + int flags, + size_t size, + size_t* actual_size_out); + static at::DataPtr makeDataPtr( + WithFd, + const char* filename, + int fd, + int flags, + size_t size, + size_t* actual_size_out); + + void* data() const override; + + void incref(); + int decref(); + void close() override; + + ~RefcountedMapAllocator() override { + close(); + } + + protected: + void checkFlags(); + void initializeAlloc(); +}; + +} // namespace at diff --git a/build_variables.bzl b/build_variables.bzl index a58bf8897f27..c6592e0c9f74 100644 --- a/build_variables.bzl +++ b/build_variables.bzl @@ -1002,6 +1002,7 @@ aten_cpu_source_non_codegen_list = [ "aten/src/ATen/FunctionalTensorWrapper.cpp", "aten/src/ATen/FunctionalizeFallbackKernel.cpp", "aten/src/ATen/MemoryOverlap.cpp", + "aten/src/ATen/MapAllocator.cpp", "aten/src/ATen/NamedTensorUtils.cpp", "aten/src/ATen/NestedTensorImpl.cpp", "aten/src/ATen/ParallelCommon.cpp", diff --git a/c10/CMakeLists.txt b/c10/CMakeLists.txt index 98f7fa2f72da..2d5fbf555c1d 100644 --- a/c10/CMakeLists.txt +++ b/c10/CMakeLists.txt @@ -123,16 +123,6 @@ if(USE_MIMALLOC) add_dependencies(c10 mimalloc-static) endif() -if(UNIX AND NOT APPLE) - include(CheckLibraryExists) - # https://github.com/libgit2/libgit2/issues/2128#issuecomment-35649830 - CHECK_LIBRARY_EXISTS(rt clock_gettime "time.h" NEED_LIBRT) - if(NEED_LIBRT) - target_link_libraries(c10 PRIVATE rt) - set(CMAKE_REQUIRED_LIBRARIES ${CMAKE_REQUIRED_LIBRARIES} rt) - endif(NEED_LIBRT) -endif(UNIX AND NOT APPLE) - if(ANDROID) target_link_libraries(c10 PRIVATE log) endif() @@ -156,11 +146,6 @@ if(USE_ROCM) add_subdirectory(hip) endif() - -if(NOT MSVC AND NOT "${CMAKE_C_COMPILER_ID}" MATCHES "Clang") - set_source_files_properties(${CMAKE_CURRENT_LIST_DIR}/core/MapAllocator.cpp PROPERTIES COMPILE_FLAGS "-fno-openmp") -endif() - # ---[ Installation # Note: for now, we will put all export path into one single Caffe2Targets group # to deal with the cmake deployment need. Inside the Caffe2Targets set, the diff --git a/c10/build.bzl b/c10/build.bzl index 98bc9f052a18..d4192a46852d 100644 --- a/c10/build.bzl +++ b/c10/build.bzl @@ -4,7 +4,6 @@ def define_targets(rules): visibility = ["//visibility:public"], deps = [ "//c10/core:CPUAllocator", - "//c10/core:MapAllocator", "//c10/core:ScalarType", "//c10/core:alignment", "//c10/core:alloc_cpu", diff --git a/c10/core/MapAllocator.cpp b/c10/core/MapAllocator.cpp deleted file mode 100644 index 69cec3cf73dd..000000000000 --- a/c10/core/MapAllocator.cpp +++ /dev/null @@ -1,952 +0,0 @@ -#include - -#include - -#include -#include -#include -#if ATOMIC_INT_LOCK_FREE == 2 -#define C10_ATOMIC_IPC_REFCOUNT 1 -#endif - -#include -#include - -/* stuff for mapped files */ -#ifdef _WIN32 -#include -#endif - -#if defined(HAVE_MMAP) -#include -#include -#include -#endif - -#if !defined(_MSC_VER) || defined(HAVE_MMAP) -#include -#include -#elif defined(_MSC_VER) -#include -#endif -#include - -namespace c10 { - -static constexpr int64_t map_alloc_alignment = 64; - -std::string NewProcessWideShmHandle() { - static std::atomic counter{0}; - static std::random_device rd; -#ifdef _MSC_VER - return fmt::format( - "/torch_{}_{}_{}", - GetCurrentProcessId(), - rd(), - counter.fetch_add(1, std::memory_order_relaxed)); -#else - return fmt::format( - "/torch_{}_{}_{}", - getpid(), - rd(), - counter.fetch_add(1, std::memory_order_relaxed)); -#endif -} -#if defined(_WIN32) || defined(HAVE_MMAP) - -namespace { -// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) -struct MapInfo { - std::atomic refcount; -}; - -constexpr const char* unknown_filename = "filename not specified"; -#ifdef _WIN32 -constexpr const char* unknown_eventname = "eventname not specified"; -#endif -} // namespace - -// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) -MapAllocator::MapAllocator( - WithFd, - c10::string_view filename, - int fd, - int flags, - size_t size) - : filename_(filename.empty() ? unknown_filename : filename), - flags_(0) // to be filled later - , - size_(0) // to be filled later -#ifdef _WIN32 - , - handle_(INVALID_HANDLE_VALUE) // to be filled later - , - event_(INVALID_HANDLE_VALUE) // to be filled later - , - eventname_( - filename.empty() ? unknown_eventname - : (std::string(filename) + "_event")) -#else - , - fd_(fd) -#endif - , - base_ptr_(nullptr) { - - if (!(flags & ALLOCATOR_MAPPED_SHARED) && - !(flags & ALLOCATOR_MAPPED_SHAREDMEM)) { - flags &= ~ALLOCATOR_MAPPED_NOCREATE; - } - if ((flags ^ ALLOCATOR_MAPPED_EXCLUSIVE) == 0) { - TORCH_CHECK( - false, - "ALLOCATOR_MAPPED_EXCLUSIVE flag requires opening the file in shared mode"); - } -#ifdef _WIN32 - if (fd != -1) { - TORCH_CHECK(false, "MapAllocator_newWithFd is unsupported on Windows"); - } -#endif - flags_ = flags; - - // OK, now do the allocation - - if (size == 0) { - return; - } - -#ifdef _WIN32 - if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { - // Shadowing - const wchar_t* filename; - const wchar_t* eventname; - const std::wstring wFilename = c10::u8u16(filename_); - const std::wstring wEventname = c10::u8u16(eventname_); - LARGE_INTEGER hfilesz; - - if (filename_[0] == '/') { - filename = wFilename.c_str() + 1; - eventname = wEventname.c_str() + 1; - } else { - filename = wFilename.c_str(); - eventname = wEventname.c_str(); - } - - hfilesz.QuadPart = size; - - if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { - event_ = CreateEventW(nullptr, FALSE, FALSE, eventname); - } else if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { - event_ = OpenEventW(EVENT_ALL_ACCESS, FALSE, eventname); - } else { - TORCH_CHECK( - false, - "Expected either ALLOCATOR_MAPPED_EXCLUSIVE or ALLOCATOR_MAPPED_NOCREATE"); - } - - if (event_ == nullptr) { - TORCH_CHECK( - false, - "Couldn't open shared event: <", - eventname, - ">, error code: <", - GetLastError(), - ">"); - } - - if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { - handle_ = CreateFileMappingW( - INVALID_HANDLE_VALUE, - nullptr, - PAGE_READWRITE, - hfilesz.HighPart, - hfilesz.LowPart, - filename); - } else if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { - handle_ = OpenFileMappingW(FILE_MAP_ALL_ACCESS, FALSE, filename); - } else { - TORCH_CHECK( - false, - "Expected either ALLOCATOR_MAPPED_EXCLUSIVE or ALLOCATOR_MAPPED_NOCREATE"); - } - - if (handle_ == nullptr) { - TORCH_CHECK( - false, - "Couldn't open shared file mapping: <", - filename, - ">, error code: <", - GetLastError(), - ">"); - } - - size_ = size; - base_ptr_ = MapViewOfFile(handle_, FILE_MAP_ALL_ACCESS, 0, 0, size); - if (!base_ptr_) { - TORCH_CHECK( - false, - "Couldn't map view of shared file <", - filename, - ">, error code: <", - GetLastError(), - ">"); - } - } else { - HANDLE hfile; - HANDLE hmfile; - LARGE_INTEGER hfilesz; - - if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { - TORCH_CHECK(false, "exclusive file mapping is not supported on Windows"); - } - if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { - TORCH_CHECK( - false, "file mapping without creation is not supported on Windows"); - } - if (flags_ & ALLOCATOR_MAPPED_KEEPFD) { - TORCH_CHECK(false, "ALLOCATOR_MAPPED_KEEPFD not supported on Windows"); - } - if (flags_ & ALLOCATOR_MAPPED_FROMFD) { - TORCH_CHECK(false, "ALLOCATOR_MAPPED_FROMFD not supported on Windows"); - } - - // Shadowing - const wchar_t* filename; - const std::wstring wFilename = c10::u8u16(filename_); - - filename = wFilename.c_str(); - - /* open file */ - /* FILE_FLAG_RANDOM_ACCESS ? */ - if (flags_) { - hfile = CreateFileW( - filename, - GENERIC_READ | GENERIC_WRITE, - FILE_SHARE_WRITE | FILE_SHARE_READ, - 0, - OPEN_ALWAYS, - FILE_ATTRIBUTE_NORMAL, - 0); - if (hfile == INVALID_HANDLE_VALUE) { - TORCH_CHECK( - false, - "could not open file <", - filename_, - "> in read-write mode; error code: <", - GetLastError(), - ">"); - } - } else { - hfile = CreateFileW( - filename, - GENERIC_READ, - FILE_SHARE_WRITE | FILE_SHARE_READ, - 0, - OPEN_EXISTING, - FILE_ATTRIBUTE_NORMAL, - 0); - if (hfile == INVALID_HANDLE_VALUE) { - TORCH_CHECK( - false, - "could not open file <", - filename_, - "> in read-only mode; error code: <", - GetLastError(), - ">"); - } - } - - if (GetFileSizeEx(hfile, &hfilesz) == 0) { - TORCH_CHECK( - false, - "could not get file size: <", - filename_, - ">; error code: <", - GetLastError(), - ">"); - } - - if (size > 0) { - if (size > hfilesz.QuadPart) { - if (flags_) { - hfilesz.QuadPart = size; - if (SetFilePointerEx(hfile, hfilesz, NULL, FILE_BEGIN) == 0) { - CloseHandle(hfile); - TORCH_CHECK( - false, - "unable to stretch file <", - filename_, - "> to the right size; error code: <", - GetLastError(), - ">", - filename_); - } - if (SetEndOfFile(hfile) == 0) { - CloseHandle(hfile); - TORCH_CHECK( - false, - "unable to write to file <", - filename_, - ">; error code: <", - GetLastError(), - ">"); - } - } else { - CloseHandle(hfile); - TORCH_CHECK( - false, - "file <", - filename_, - "> size is smaller than the required mapping size <", - size, - ">; error code: <", - GetLastError(), - ">"); - } - } - } else { - size = hfilesz.QuadPart; - } - - size_ = size; /* if we are here, it must be the right size */ - - hfilesz.QuadPart = size_; - - /* get map handle */ - if (flags_) { - if ((hmfile = CreateFileMappingW( - hfile, - NULL, - PAGE_READWRITE, - hfilesz.HighPart, - hfilesz.LowPart, - NULL)) == NULL) { - TORCH_CHECK( - false, - "could not create a map on file <", - filename_, - ">; error code: <", - GetLastError(), - ">"); - } - } else { - if ((hmfile = CreateFileMappingW( - hfile, - NULL, - PAGE_WRITECOPY, - hfilesz.HighPart, - hfilesz.LowPart, - NULL)) == NULL) { - TORCH_CHECK( - false, - "could not create a map on file <", - filename_, - ">; error code: <", - GetLastError(), - ">"); - } - } - - /* map the stuff */ - if (flags_) { - base_ptr_ = MapViewOfFile(hmfile, FILE_MAP_ALL_ACCESS, 0, 0, 0); - } else { - base_ptr_ = MapViewOfFile(hmfile, FILE_MAP_COPY, 0, 0, 0); - } - - CloseHandle(hfile); - CloseHandle(hmfile); - } -#else /* _WIN32 */ - { - /* open file */ - int fd{-1}; - int flags{}; // shadow - - if (flags_ & (ALLOCATOR_MAPPED_SHARED | ALLOCATOR_MAPPED_SHAREDMEM)) { - flags = O_RDWR | O_CREAT; - } else { - flags = O_RDONLY; - } - - if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { - flags |= O_EXCL; - } - if (flags_ & ALLOCATOR_MAPPED_NOCREATE) { - flags &= ~O_CREAT; - } - - if (!(flags_ & ALLOCATOR_MAPPED_FROMFD)) { - // NOLINTNEXTLINE(bugprone-branch-clone) - if (flags_ & ALLOCATOR_MAPPED_SHARED) { - fd = open(filename_.c_str(), flags, (mode_t)0600); - TORCH_CHECK( - fd != -1, - "unable to open file <", - filename_, - "> in read-write mode: ", - strerror(errno), - " (", - errno, - ")"); - } else if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { -#ifdef HAVE_SHM_OPEN - fd = shm_open(filename_.c_str(), flags, (mode_t)0600); - TORCH_CHECK( - fd != -1, - "unable to open shared memory object <", - filename_, - "> in read-write mode: ", - strerror(errno), - " (", - errno, - ")"); -#else - TORCH_CHECK( - false, - "unable to open file <", - filename_, - "> in sharedmem mode, shm_open unavailable on this platform"); -#endif - } else { - fd = open(filename_.c_str(), O_RDONLY); - TORCH_CHECK( - fd != -1, - "unable to open file <", - filename_, - "> in read-only mode: ", - strerror(errno), - " (", - errno, - ")"); - } - } else { - fd = fd_; - } - - // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init) - struct stat file_stat; - if (fstat(fd, &file_stat) == -1) { - // NOLINTNEXTLINE(cppcoreguidelines-init-variables) - int last_err = errno; - if (!(flags_ & ALLOCATOR_MAPPED_FROMFD)) { - ::close(fd); - } - TORCH_CHECK( - false, - "unable to stat the file <", - filename_, - ">: ", - strerror(last_err), - " (", - last_err, - ")"); - } - - if (size > 0) { - if (static_cast(size) > file_stat.st_size) { - if (flags_) { - // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions) - if (ftruncate(fd, size) == -1) { - TORCH_CHECK( - false, - "unable to resize file <", - filename_, - "> to the right size: ", - strerror(errno), - " (", - errno, - ")"); - } - if (fstat(fd, &file_stat) == -1 || - file_stat.st_size < static_cast(size)) { - // NOLINTNEXTLINE(cppcoreguidelines-init-variables) - int last_err = errno; - ::close(fd); - TORCH_CHECK( - false, - "unable to stretch file <", - filename_, - "> to the right size: ", - strerror(last_err), - " (", - last_err, - ")"); - } -/* on macOS write returns with errno 45 (Opperation not supported) when used - * with a file descriptor obtained via shm_open - */ -#ifndef __APPLE__ - if ((write(fd, "", 1)) != - 1) /* note that the string "" contains the '\0' byte ... */ { - // NOLINTNEXTLINE(cppcoreguidelines-init-variables) - int last_err = errno; - ::close(fd); - TORCH_CHECK( - false, - "unable to write to file <", - filename_, - ">: ", - strerror(last_err), - " (", - last_err, - ")"); - } -#endif - } else { - ::close(fd); - TORCH_CHECK( - false, - "file <", - filename_, - "> size is smaller than the required mapping size <", - size, - ">"); - } - } - } else { - size = file_stat.st_size; - } - - // NOLINTNEXTLINE(bugprone-narrowing-conversions,cppcoreguidelines-narrowing-conversions) - size_ = size; /* if we are here, it must be the right size */ - - /* map it */ - if (flags_ & (ALLOCATOR_MAPPED_SHARED | ALLOCATOR_MAPPED_SHAREDMEM)) { - base_ptr_ = - mmap(nullptr, size_, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); - } else { - base_ptr_ = - mmap(nullptr, size_, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0); - } - - if (base_ptr_ == MAP_FAILED) { - base_ptr_ = nullptr; /* let's be sure it is NULL */ - TORCH_CHECK( - false, - "unable to mmap ", - size_, - " bytes from file <", - filename_, - ">: ", - strerror(errno), - " (", - errno, - ")"); - } - - if (flags_ & ALLOCATOR_MAPPED_KEEPFD) { - fd_ = fd; - } else { - if (::close(fd) == -1) { - TORCH_CHECK( - false, - "Error closing file <", - filename_, - ">: ", - strerror(errno), - " (", - errno, - ")"); - } - fd_ = -1; - } - - if (flags_ & ALLOCATOR_MAPPED_UNLINK) { - if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { -#ifdef HAVE_SHM_UNLINK - if (shm_unlink(filename_.c_str()) == -1) { - TORCH_CHECK( - false, - "could not unlink the shared memory file ", - filename_, - " : ", - strerror(errno), - " (", - errno, - ")"); - } -#else - TORCH_CHECK( - false, - "could not unlink the shared memory file ", - filename_, - ", shm_unlink not available on platform"); -#endif - } else { - if (unlink(filename_.c_str()) == -1) - TORCH_CHECK( - false, - "could not unlink file ", - filename_, - " : ", - strerror(errno), - " (", - errno, - ")"); - } - } - - if (base_ptr_ == MAP_FAILED) { - TORCH_CHECK( - false, - "$ Torch: unable to mmap memory: you tried to mmap ", - size_ / 1073741824, - " GB."); - } - } -#endif - c10::reportMemoryUsageToProfiler( - base_ptr_, size_, 0, size_, c10::Device(c10::DeviceType::CPU)); -} - -MapAllocator::MapAllocator(c10::string_view filename, int flags, size_t size) - : MapAllocator(WITH_FD, filename, -1, flags, size) {} - -#ifdef _WIN32 -struct ReleaseContext { - HANDLE event; - HANDLE handle; - HANDLE wait; -}; -static void CALLBACK -WaitForReleaseHandle(PVOID lpParam, BOOLEAN TimerOrWaitFired) { - if (lpParam) { - ReleaseContext* ctx = (ReleaseContext*)lpParam; - - SetEvent(ctx->event); - CloseHandle(ctx->event); - CloseHandle(ctx->handle); - - UnregisterWait(ctx->wait); - - delete ctx; - } -} -#endif - -void MapAllocator::close() { - if (closed_) { - return; - } - closed_ = true; - if (base_ptr_ == nullptr) { - return; - } -#ifdef _WIN32 - if ((flags_ & ALLOCATOR_MAPPED_KEEPFD) || - (flags_ & ALLOCATOR_MAPPED_SHAREDMEM)) - CloseHandle(handle_); - if (UnmapViewOfFile(base_ptr_) == 0) - TORCH_CHECK(false, "could not unmap the shared memory file"); -#else /* _WIN32 */ - if (flags_ & ALLOCATOR_MAPPED_KEEPFD) { - if (::close(fd_) == -1) { - TORCH_CHECK( - false, - "could not close file descriptor ", - fd_, - " :", - strerror(errno), - " (", - errno, - ")"); - } - } - - if (munmap(base_ptr_, size_)) { - TORCH_CHECK( - false, - "could not unmap the shared memory file: ", - strerror(errno), - " (", - errno, - ")"); - } - - if (!(flags_ & (ALLOCATOR_MAPPED_FROMFD | ALLOCATOR_MAPPED_UNLINK))) { - if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { -#ifdef HAVE_SHM_UNLINK - if (shm_unlink(filename_.c_str()) == -1) { - TORCH_CHECK( - false, - "could not unlink the shared memory file ", - filename_, - " : ", - strerror(errno), - " (", - errno, - ")"); - } -#else - TORCH_CHECK( - false, - "could not unlink the shared memory file ", - filename_, - ", shm_unlink not available on platform"); -#endif - } - } -#endif /* _WIN32 */ -} - -#else /* defined(_WIN32) || defined(HAVE_MMAP) */ - -MapAllocator::MapAllocator(c10::string_view filename, int flags, size_t size) { - TORCH_CHECK(false, "file mapping not supported on your system"); -} - -MapAllocator::MapAllocator( - WithFd, - c10::string_view filename, - int fd, - int flags, - size_t size) { - TORCH_CHECK(false, "file mapping not supported on your system"); -} - -void MapAllocator::close() {} - -#endif - -#if (defined(_WIN32) || defined(HAVE_MMAP)) && defined(C10_ATOMIC_IPC_REFCOUNT) - -RefcountedMapAllocatorArgCheck::RefcountedMapAllocatorArgCheck(int flags) { - if (flags & ALLOCATOR_MAPPED_FROMFD) { - TORCH_CHECK( - false, - "RefcountedMapAllocator doesn't support ALLOCATOR_MAPPED_FROMFD flag"); - } - if (flags & ALLOCATOR_MAPPED_KEEPFD) { - TORCH_CHECK( - false, - "RefcountedMapAllocator doesn't support ALLOCATOR_MAPPED_KEEPFD flag"); - } - if (flags & ALLOCATOR_MAPPED_UNLINK) { - TORCH_CHECK( - false, - "RefcountedMapAllocator doesn't support ALLOCATOR_MAPPED_UNLINK flag"); - } - if (!(flags & ALLOCATOR_MAPPED_SHAREDMEM)) { - TORCH_CHECK( - false, - "RefcountedMapAllocator requires ALLOCATOR_MAPPED_SHAREDMEM flag"); - } -} - -RefcountedMapAllocator::RefcountedMapAllocator( - const char* filename, - int flags, - size_t size) - : RefcountedMapAllocatorArgCheck(flags), - MapAllocator(filename, flags, size + map_alloc_alignment) { - initializeAlloc(); -} -RefcountedMapAllocator::RefcountedMapAllocator( - WithFd, - const char* filename, - int fd, - int flags, - size_t size) - : RefcountedMapAllocatorArgCheck(flags), - MapAllocator(WITH_FD, filename, flags, fd, size + map_alloc_alignment) { - initializeAlloc(); -} - -void RefcountedMapAllocator::initializeAlloc() { - TORCH_CHECK(base_ptr_, "base_ptr_ is null"); - MapInfo* map_info = (MapInfo*)base_ptr_; - -#ifdef _WIN32 - ReleaseContext* r_ctx = new ReleaseContext; - r_ctx->handle = handle_; - r_ctx->event = event_; - r_ctx->wait = NULL; - BOOL can_wait = RegisterWaitForSingleObject( - &r_ctx->wait, - event_, - WaitForReleaseHandle, - (PVOID)r_ctx, - INFINITE, - WT_EXECUTEONLYONCE); - TORCH_CHECK( - can_wait, - "Couldn't register wait on event, error code: <", - GetLastError(), - ">"); -#endif - - if (flags_ & ALLOCATOR_MAPPED_EXCLUSIVE) { - new (&map_info->refcount) std::atomic(1); - } else { - map_info->refcount++; - } -} - -void RefcountedMapAllocator::close() { - if (closed_) { - return; - } - closed_ = true; - - void* data = base_ptr_; - -#ifdef _WIN32 - MapInfo* info = (MapInfo*)data; - if (--info->refcount == 0) { - SetEvent(event_); - } - if (UnmapViewOfFile(data) == 0) { - TORCH_CHECK(false, "could not unmap the shared memory file"); - } -#else /* _WIN32 */ - - MapInfo* info = (MapInfo*)(data); - if (--info->refcount == 0) { -#ifdef HAVE_SHM_UNLINK - if (shm_unlink(filename_.c_str()) == -1) { - TORCH_CHECK(false, "could not unlink the shared memory file ", filename_); - } -#else - TORCH_CHECK( - false, - "could not unlink the shared memory file ", - filename_, - ", shm_unlink not available on platform"); -#endif /* HAVE_SHM_UNLINK */ - } - if (munmap(info, size_)) { - TORCH_CHECK(false, "could not unmap the shared memory file ", filename_); - } -#endif /* _WIN32 */ -} - -void RefcountedMapAllocator::incref() { - MapInfo* map_info = static_cast(base_ptr_); - ++map_info->refcount; -} - -int RefcountedMapAllocator::decref() { - MapInfo* map_info = static_cast(base_ptr_); - return --map_info->refcount == 0; -} - -#else - -RefcountedMapAllocatorArgCheck::RefcountedMapAllocatorArgCheck(int flags) {} - -RefcountedMapAllocator::RefcountedMapAllocator( - const char* filename, - int flags, - size_t size) - : RefcountedMapAllocatorArgCheck(flags), - MapAllocator(filename, flags, size + map_alloc_alignment) { - TORCH_CHECK(false, "refcounted file mapping not supported on your system"); -} - -RefcountedMapAllocator::RefcountedMapAllocator( - WithFd, - const char* filename, - int fd, - int flags, - size_t size) - : RefcountedMapAllocatorArgCheck(flags), - MapAllocator(WITH_FD, filename, flags, fd, size + map_alloc_alignment) { - TORCH_CHECK(false, "refcounted file mapping not supported on your system"); -} - -void RefcountedMapAllocator::initializeAlloc() {} - -void RefcountedMapAllocator::close() {} - -#endif - -static void deleteMapAllocator(void* ptr) { - delete static_cast(ptr); -} - -static void deleteRefcountedMapAllocator(void* ptr) { - delete static_cast(ptr); -} - -MapAllocator* MapAllocator::fromDataPtr(const c10::DataPtr& dptr) { - return dptr.cast_context(&deleteMapAllocator); -} - -RefcountedMapAllocator* RefcountedMapAllocator::fromDataPtr( - const c10::DataPtr& dptr) { - return dptr.cast_context( - &deleteRefcountedMapAllocator); -} - -c10::DataPtr MapAllocator::makeDataPtr( - c10::string_view filename, - int flags, - size_t size, - size_t* actual_size_out) { - auto* context = new MapAllocator(filename, flags, size); - if (actual_size_out) - *actual_size_out = context->size(); - return {context->data(), context, &deleteMapAllocator, c10::DeviceType::CPU}; -} - -c10::DataPtr MapAllocator::makeDataPtr( - WithFd, - const char* filename, - int fd, - int flags, - size_t size, - size_t* actual_size_out) { - // NOLINTNEXTLINE(cppcoreguidelines-init-variables) - auto* context = new MapAllocator(WITH_FD, filename, fd, flags, size); - if (actual_size_out) - *actual_size_out = context->size(); - return {context->data(), context, &deleteMapAllocator, c10::DeviceType::CPU}; -} - -c10::DataPtr RefcountedMapAllocator::makeDataPtr( - const char* filename, - int flags, - size_t size, - size_t* actual_size_out) { - auto* context = new RefcountedMapAllocator(filename, flags, size); - if (actual_size_out) - *actual_size_out = context->size() - map_alloc_alignment; - return { - context->data(), - context, - &deleteRefcountedMapAllocator, - c10::DeviceType::CPU}; -} - -c10::DataPtr RefcountedMapAllocator::makeDataPtr( - WithFd, - const char* filename, - int fd, - int flags, - size_t size, - size_t* actual_size_out) { - auto* context = - new RefcountedMapAllocator(WITH_FD, filename, fd, flags, size); - if (actual_size_out) - *actual_size_out = context->size() - map_alloc_alignment; - return { - context->data(), - context, - &deleteRefcountedMapAllocator, - c10::DeviceType::CPU}; -} - -void* RefcountedMapAllocator::data() const { - return static_cast( - static_cast(base_ptr_) + map_alloc_alignment); -} - -MapAllocator::~MapAllocator() { - // NOLINTNEXTLINE(clang-analyzer-optin.cplusplus.VirtualCall) - close(); - c10::reportMemoryUsageToProfiler( - base_ptr_, -size_, 0, 0, c10::Device(c10::DeviceType::CPU)); -} - -} // namespace c10 diff --git a/c10/core/MapAllocator.h b/c10/core/MapAllocator.h deleted file mode 100644 index 80e72bc0c104..000000000000 --- a/c10/core/MapAllocator.h +++ /dev/null @@ -1,140 +0,0 @@ -#pragma once - -#include -#include -#include - -namespace c10 { - -enum MappedAllocatorModes { - ALLOCATOR_MAPPED_SHARED = 1, - ALLOCATOR_MAPPED_SHAREDMEM = 2, - ALLOCATOR_MAPPED_EXCLUSIVE = 4, - ALLOCATOR_MAPPED_NOCREATE = 8, - ALLOCATOR_MAPPED_KEEPFD = 16, - ALLOCATOR_MAPPED_FROMFD = 32, - ALLOCATOR_MAPPED_UNLINK = 64 -}; - -// Sentinel value/type to help distinguish the file descriptor constructor from -// the non-file descriptor constructor -enum WithFd { WITH_FD }; - -C10_API std::string NewProcessWideShmHandle(); - -class C10_API MapAllocator { - public: - MapAllocator(c10::string_view filename, int flags, size_t size); - MapAllocator( - WithFd, - c10::string_view filename, - int fd, - int flags, - size_t size); - MapAllocator(const MapAllocator&) = delete; - MapAllocator& operator=(const MapAllocator&) = delete; - MapAllocator(MapAllocator&&) = delete; - MapAllocator& operator=(MapAllocator&&) = delete; - - const char* filename() const { - return filename_.c_str(); - } - int fd() const { -#ifdef _WIN32 - TORCH_CHECK(false, "MapAllocator::fd() is unsupported on Windows"); -#else - return fd_; -#endif - } - ptrdiff_t size() const { - return size_; - } - // Return a pointer to the actual data for this allocator - // (in the case of the refcounted allocator, this is offset - // from the base pointer.) - virtual void* data() const { - return base_ptr_; - } - - static MapAllocator* fromDataPtr(const c10::DataPtr&); - static c10::DataPtr makeDataPtr( - c10::string_view filename, - int flags, - size_t size, - size_t* actual_size_out); - static c10::DataPtr makeDataPtr( - WithFd, - const char* filename, - int fd, - int flags, - size_t size, - size_t* actual_size_out); - - // Closes the data. Helps us avoid destructor shenanigans - virtual void close(); - - // This is very dangerous. You have to redefine this destructor for each - // subclass - virtual ~MapAllocator(); - - protected: - bool closed_ = false; - std::string filename_; - int flags_ = 0; - ptrdiff_t size_; /* mapped size */ -#ifdef _WIN32 - void* handle_; - void* event_; - std::string eventname_; -#else - int fd_ = -1; -#endif - void* base_ptr_ = nullptr; -}; - -// Base-from-member idiom -struct C10_API RefcountedMapAllocatorArgCheck { - RefcountedMapAllocatorArgCheck(int flags); -}; - -class C10_API RefcountedMapAllocator : private RefcountedMapAllocatorArgCheck, - public MapAllocator { - public: - RefcountedMapAllocator(const char* filename, int flags, size_t size); - RefcountedMapAllocator( - WithFd, - const char* filename, - int fd, - int flags, - size_t size); - - static RefcountedMapAllocator* fromDataPtr(const c10::DataPtr&); - static c10::DataPtr makeDataPtr( - const char* filename, - int flags, - size_t size, - size_t* actual_size_out); - static c10::DataPtr makeDataPtr( - WithFd, - const char* filename, - int fd, - int flags, - size_t size, - size_t* actual_size_out); - - void* data() const override; - - void incref(); - int decref(); - void close() override; - - ~RefcountedMapAllocator() override { - close(); - } - - protected: - void checkFlags(); - void initializeAlloc(); -}; - -} // namespace c10 diff --git a/c10/core/build.bzl b/c10/core/build.bzl index 7ca9379d3575..f395b9b92388 100644 --- a/c10/core/build.bzl +++ b/c10/core/build.bzl @@ -18,27 +18,6 @@ def define_targets(rules): alwayslink = True, ) - rules.cc_library( - name = "MapAllocator", - srcs = ["MapAllocator.cpp"], - hdrs = ["MapAllocator.h"], - linkstatic = True, - local_defines = [ - "C10_BUILD_MAIN_LIB", - "HAVE_MMAP=1", - "HAVE_SHM_OPEN=1", - "HAVE_SHM_UNLINK=1", - ], - visibility = ["//visibility:public"], - deps = [ - ":CPUAllocator", - ":base", - "//c10/util:base", - ], - linkopts = ["-lrt"], - alwayslink = True, - ) - rules.cc_library( name = "ScalarType", hdrs = ["ScalarType.h"], @@ -83,7 +62,6 @@ def define_targets(rules): ], exclude = [ "CPUAllocator.cpp", - "MapAllocator.cpp", "impl/alloc_cpu.cpp", ], ), @@ -95,7 +73,6 @@ def define_targets(rules): ], exclude = [ "CPUAllocator.h", - "MapAllocator.h", "impl/alloc_cpu.h", ], ), diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 90a36f5404d9..b3888ebde673 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -864,6 +864,8 @@ if(USE_PRECOMPILED_HEADERS) "$<$:ATen/core/ATen_pch.h>") # Exclude some files from using PCH set_source_files_properties( + # Not built with OpenMP, so PCH is invalid + ${Torch_SOURCE_DIR}/aten/src/ATen/MapAllocator.cpp # Builds with incompatible compiler flags ${Caffe2_CPU_SRCS_AVX2} ${Caffe2_CPU_SRCS_AVX512} diff --git a/cmake/Codegen.cmake b/cmake/Codegen.cmake index fb00648fa4f8..5d552f629be9 100644 --- a/cmake/Codegen.cmake +++ b/cmake/Codegen.cmake @@ -63,6 +63,10 @@ if(INTERN_BUILD_ATEN_OPS) endif() endif(MSVC) + if(NOT MSVC AND NOT "${CMAKE_C_COMPILER_ID}" MATCHES "Clang") + set_source_files_properties(${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/MapAllocator.cpp PROPERTIES COMPILE_FLAGS "-fno-openmp") + endif() + file(GLOB_RECURSE all_python "${CMAKE_CURRENT_LIST_DIR}/../torchgen/*.py") set(GEN_ROCM_FLAG) diff --git a/torch/lib/libshm/CMakeLists.txt b/torch/lib/libshm/CMakeLists.txt index 579ad904211c..20158a9a2553 100644 --- a/torch/lib/libshm/CMakeLists.txt +++ b/torch/lib/libshm/CMakeLists.txt @@ -29,7 +29,7 @@ set_target_properties(shm PROPERTIES PREFIX "lib" IMPORT_PREFIX "lib" CXX_STANDARD 17) -target_link_libraries(shm PUBLIC c10) +target_link_libraries(shm PUBLIC torch) if(UNIX AND NOT APPLE) include(CheckLibraryExists) diff --git a/torch/lib/libshm/core.cpp b/torch/lib/libshm/core.cpp index 6e82b59560b8..e9ff9b0b6b90 100644 --- a/torch/lib/libshm/core.cpp +++ b/torch/lib/libshm/core.cpp @@ -123,7 +123,7 @@ THManagedMapAllocator::THManagedMapAllocator( int flags, size_t size) : THManagedMapAllocatorInit(manager_handle, filename), - c10::RefcountedMapAllocator(filename, flags, size) {} + at::RefcountedMapAllocator(filename, flags, size) {} void THManagedMapAllocator::close() { if (closed_) @@ -131,7 +131,7 @@ void THManagedMapAllocator::close() { AllocInfo info = get_alloc_info(filename()); info.free = true; ClientSocket& socket = get_manager_socket(manager_handle_); - c10::RefcountedMapAllocator::close(); + at::RefcountedMapAllocator::close(); socket.register_deallocation(info); } @@ -139,7 +139,7 @@ static void deleteTHManagedMapAllocator(void* ptr) { delete static_cast(ptr); } -c10::DataPtr THManagedMapAllocator::makeDataPtr( +at::DataPtr THManagedMapAllocator::makeDataPtr( const char* manager_handle, const char* filename, int flags, @@ -150,10 +150,10 @@ c10::DataPtr THManagedMapAllocator::makeDataPtr( context->data(), context, &deleteTHManagedMapAllocator, - c10::DeviceType::CPU}; + at::DeviceType::CPU}; } THManagedMapAllocator* THManagedMapAllocator::fromDataPtr( - const c10::DataPtr& dptr) { + const at::DataPtr& dptr) { return dptr.cast_context(&deleteTHManagedMapAllocator); } diff --git a/torch/lib/libshm/libshm.h b/torch/lib/libshm/libshm.h index cfe07d528bbc..28024aa2338d 100644 --- a/torch/lib/libshm/libshm.h +++ b/torch/lib/libshm/libshm.h @@ -1,23 +1,23 @@ #pragma once -#include +#include #ifdef __cplusplus void libshm_init(const char* manager_exec_path); -// Superclass to run a constructor before c10::RefcountedMapAllocator +// Superclass to run a constructor before at::RefcountedMapAllocator class THManagedMapAllocatorInit { protected: THManagedMapAllocatorInit(const char* manager_handle, const char* filename); std::string manager_handle_; }; -// Like a c10::RefcountedMapAllocator, but it also makes use of an external +// Like a at::RefcountedMapAllocator, but it also makes use of an external // shared memory manager process to ensure that shared memory regions actually // get freed in the end (even if processes lose the memory). class THManagedMapAllocator : private THManagedMapAllocatorInit, - public c10::RefcountedMapAllocator { + public at::RefcountedMapAllocator { public: THManagedMapAllocator( const char* manager_handle, @@ -31,12 +31,12 @@ class THManagedMapAllocator : private THManagedMapAllocatorInit, close(); } - static c10::DataPtr makeDataPtr( + static at::DataPtr makeDataPtr( const char* manager_handle, const char* filename, int flags, size_t size); - static THManagedMapAllocator* fromDataPtr(const c10::DataPtr&); + static THManagedMapAllocator* fromDataPtr(const at::DataPtr&); const char* manager_handle() const { return manager_handle_.c_str(); diff --git a/torch/lib/libshm_windows/core.cpp b/torch/lib/libshm_windows/core.cpp index cb1d1ce64a9f..2d98f7f654f0 100644 --- a/torch/lib/libshm_windows/core.cpp +++ b/torch/lib/libshm_windows/core.cpp @@ -10,17 +10,17 @@ static void deleteTHManagedMapAllocator(void* ptr) { delete static_cast(ptr); } -c10::DataPtr THManagedMapAllocator::makeDataPtr( +at::DataPtr THManagedMapAllocator::makeDataPtr( const char* manager_handle, const char* filename, int flags, size_t size) { auto* context = new THManagedMapAllocator(manager_handle, filename, flags, size); - return {context->data(), context, &deleteTHManagedMapAllocator, c10::kCPU}; + return {context->data(), context, &deleteTHManagedMapAllocator, at::kCPU}; } THManagedMapAllocator* THManagedMapAllocator::fromDataPtr( - const c10::DataPtr& dptr) { + const at::DataPtr& dptr) { return dptr.cast_context(&deleteTHManagedMapAllocator); } diff --git a/torch/lib/libshm_windows/libshm.h b/torch/lib/libshm_windows/libshm.h index b28514750f7e..4dd193df93d1 100644 --- a/torch/lib/libshm_windows/libshm.h +++ b/torch/lib/libshm_windows/libshm.h @@ -1,6 +1,6 @@ #pragma once -#include +#include #ifdef __cplusplus @@ -12,21 +12,21 @@ SHM_API void libshm_init(const char* manager_exec_path); -class SHM_API THManagedMapAllocator : public c10::RefcountedMapAllocator { +class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator { public: THManagedMapAllocator( const char* manager_handle, const char* filename, int flags, size_t size) - : c10::RefcountedMapAllocator(filename, flags, size) {} + : at::RefcountedMapAllocator(filename, flags, size) {} - static c10::DataPtr makeDataPtr( + static at::DataPtr makeDataPtr( const char* manager_handle, const char* filename, int flags, size_t size); - static THManagedMapAllocator* fromDataPtr(const c10::DataPtr&); + static THManagedMapAllocator* fromDataPtr(const at::DataPtr&); const char* manager_handle() const { return "no_manager";