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

This reverts commit 0341deb1c7.

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

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

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)

aten/src/ATen/MapAllocator.cpp

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

aten/src/ATen/MapAllocator.h

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

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

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

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

c10/core/MapAllocator.cpp

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

c10/core/MapAllocator.h

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

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

caffe2/CMakeLists.txt

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

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)

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)

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<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);
 }

torch/lib/libshm/libshm.h

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

torch/lib/libshm_windows/core.cpp

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

torch/lib/libshm_windows/libshm.h

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