frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "Move at::{Refcounted,}MapAllocator to c10 (#109881)"

Browse Source 
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))
This commit is contained in:
PyTorch MergeBot
2023-10-10 20:39:11 +00:00
parent 495f77be7a
commit 02a02a23ee
17 changed files with 798 additions and 1154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
BUILD.bazel
View File

@ -1643,8 +1643,7 @@ cc_library(
name = "shm",
srcs = glob(["torch/lib/libshm/*.cpp"]),
deps = [
":torch_headers",
"//c10",
":torch",
],
)

22
aten/src/ATen/CMakeLists.txt
View File

@ -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)

611
aten/src/ATen/MapAllocator.cpp Normal file
View File

@ -0,0 +1,611 @@
#include <ATen/MapAllocator.h>
#include <atomic>
#include <random>
#include <string>
#if ATOMIC_INT_LOCK_FREE == 2
#define AT_ATOMIC_IPC_REFCOUNT 1
#endif
#include <c10/core/CPUAllocator.h>
#include <c10/util/Unicode.h>
/* stuff for mapped files */
#ifdef _WIN32
#include <c10/util/win32-headers.h>
#endif
#if defined(HAVE_MMAP)
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#endif
#if !defined(_MSC_VER) || defined(HAVE_MMAP)
#include <sys/types.h>
#include <unistd.h>
#elif defined(_MSC_VER)
#include <c10/util/win32-headers.h>
#endif
#include <fmt/format.h>
namespace at {
static constexpr int64_t map_alloc_alignment = 64;
std::string NewProcessWideShmHandle() {
static std::atomic<uint64_t> 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<int> 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<int64_t>(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<int64_t>(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<int>(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<MapInfo*>(base_ptr_);
++map_info->refcount;
}
int RefcountedMapAllocator::decref()
{
MapInfo *map_info = static_cast<MapInfo*>(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<MapAllocator*>(ptr);
}
static void deleteRefcountedMapAllocator(void* ptr) {
delete static_cast<RefcountedMapAllocator*>(ptr);
}
MapAllocator* MapAllocator::fromDataPtr(const at::DataPtr& dptr) {
return dptr.cast_context<MapAllocator>(&deleteMapAllocator);
}
RefcountedMapAllocator* RefcountedMapAllocator::fromDataPtr(const at::DataPtr& dptr) {
return dptr.cast_context<RefcountedMapAllocator>(&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<void*>(static_cast<char*>(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

138
aten/src/ATen/MapAllocator.h
View File

@ -1,3 +1,139 @@
#pragma once
#include <c10/core/MapAllocator.h>
#include <c10/core/Allocator.h>
#include <c10/util/string_view.h>
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

1
build_variables.bzl
View File

@ -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",

15
c10/CMakeLists.txt
View File

@ -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

1
c10/build.bzl
View File

@ -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",

952
c10/core/MapAllocator.cpp
View File

@ -1,952 +0,0 @@
#include <c10/core/MapAllocator.h>
#include <c10/util/Exception.h>
#include <atomic>
#include <random>
#include <string>
#if ATOMIC_INT_LOCK_FREE == 2
#define C10_ATOMIC_IPC_REFCOUNT 1
#endif
#include <c10/core/CPUAllocator.h>
#include <c10/util/Unicode.h>
/* stuff for mapped files */
#ifdef _WIN32
#include <c10/util/win32-headers.h>
#endif
#if defined(HAVE_MMAP)
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#endif
#if !defined(_MSC_VER) || defined(HAVE_MMAP)
#include <sys/types.h>
#include <unistd.h>
#elif defined(_MSC_VER)
#include <c10/util/win32-headers.h>
#endif
#include <fmt/format.h>
namespace c10 {
static constexpr int64_t map_alloc_alignment = 64;
std::string NewProcessWideShmHandle() {
static std::atomic<uint64_t> 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<int> 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<int64_t>(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<int64_t>(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<int>(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<MapInfo*>(base_ptr_);
++map_info->refcount;
}
int RefcountedMapAllocator::decref() {
MapInfo* map_info = static_cast<MapInfo*>(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<MapAllocator*>(ptr);
}
static void deleteRefcountedMapAllocator(void* ptr) {
delete static_cast<RefcountedMapAllocator*>(ptr);
}
MapAllocator* MapAllocator::fromDataPtr(const c10::DataPtr& dptr) {
return dptr.cast_context<MapAllocator>(&deleteMapAllocator);
}
RefcountedMapAllocator* RefcountedMapAllocator::fromDataPtr(
const c10::DataPtr& dptr) {
return dptr.cast_context<RefcountedMapAllocator>(
&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<void*>(
static_cast<char*>(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

140
c10/core/MapAllocator.h
View File

@ -1,140 +0,0 @@
#pragma once
#include <c10/core/Allocator.h>
#include <c10/util/Exception.h>
#include <c10/util/string_view.h>
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

23
c10/core/build.bzl
View File

@ -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",
],
),

2
caffe2/CMakeLists.txt
View File

@ -864,6 +864,8 @@ if(USE_PRECOMPILED_HEADERS)
"$<$<COMPILE_LANGUAGE:CXX>: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}

4
cmake/Codegen.cmake
View File

@ -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)

2
torch/lib/libshm/CMakeLists.txt
View File

@ -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)

10
torch/lib/libshm/core.cpp
View File

@ -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<THManagedMapAllocator*>(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<THManagedMapAllocator>(&deleteTHManagedMapAllocator);
}

12
torch/lib/libshm/libshm.h
View File

@ -1,23 +1,23 @@
#pragma once
#include <c10/core/MapAllocator.h>
#include <ATen/MapAllocator.h>
#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();

6
torch/lib/libshm_windows/core.cpp
View File

@ -10,17 +10,17 @@ static void deleteTHManagedMapAllocator(void* ptr) {
delete static_cast<THManagedMapAllocator*>(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<THManagedMapAllocator>(&deleteTHManagedMapAllocator);
}

10
torch/lib/libshm_windows/libshm.h
View File

@ -1,6 +1,6 @@
#pragma once
#include <c10/core/MapAllocator.h>
#include <ATen/MapAllocator.h>
#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";