diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index 0d90f8a34e40..580edaadbeba 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -605,6 +605,11 @@ if(UNIX) if(HAVE_MALLOC_USABLE_SIZE) add_definitions(-DHAVE_MALLOC_USABLE_SIZE=1) endif(HAVE_MALLOC_USABLE_SIZE) + set(CMAKE_EXTRA_INCLUDE_FILES "fcntl.h") + CHECK_FUNCTION_EXISTS(posix_fallocate HAVE_POSIX_FALLOCATE) + if(HAVE_POSIX_FALLOCATE) + add_definitions(-DHAVE_POSIX_FALLOCATE=1) + endif(HAVE_POSIX_FALLOCATE) endif(UNIX) ADD_DEFINITIONS(-DUSE_EXTERNAL_MZCRC) diff --git a/aten/src/ATen/MapAllocator.cpp b/aten/src/ATen/MapAllocator.cpp index 63a278050e8a..d88c05f087a9 100644 --- a/aten/src/ATen/MapAllocator.cpp +++ b/aten/src/ATen/MapAllocator.cpp @@ -292,6 +292,28 @@ MapAllocator::MapAllocator(WithFd, std::string_view filename, int fd, int flags, if (ftruncate(fd, static_cast(size)) == -1) { TORCH_CHECK(false, "unable to resize file <", filename_, "> to the right size: ", c10::utils::str_error(errno), " (", errno, ")"); } + +#ifdef HAVE_POSIX_FALLOCATE + if (flags_ & ALLOCATOR_MAPPED_SHAREDMEM) { + for (;;) { + if (posix_fallocate(fd, 0, static_cast(size)) == 0) { + break; + } + + if (errno == EINTR) { + continue; + } + + if (errno == EINVAL || errno == EOPNOTSUPP) { + // the underlying filesystem does not support the operation + break; + } + + TORCH_CHECK(false, "unable to allocate shared memory(shm) for file <", filename_, ">: ", c10::utils::str_error(errno), " (", errno, ")"); + } + } +#endif + if (fstat(fd, &file_stat) == -1 || file_stat.st_size < static_cast(size)) { #ifndef STRIP_ERROR_MESSAGES int last_err = errno; diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index ad8540a86ec8..d39f1a2efc77 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -1541,6 +1541,11 @@ if(NOT INTERN_BUILD_MOBILE) if(HAVE_MALLOC_USABLE_SIZE) add_definitions(-DHAVE_MALLOC_USABLE_SIZE=1) endif(HAVE_MALLOC_USABLE_SIZE) + set(CMAKE_EXTRA_INCLUDE_FILES "fcntl.h") + CHECK_FUNCTION_EXISTS(posix_fallocate HAVE_POSIX_FALLOCATE) + if(HAVE_POSIX_FALLOCATE) + add_definitions(-DHAVE_POSIX_FALLOCATE=1) + endif(HAVE_POSIX_FALLOCATE) endif(UNIX) add_definitions(-DUSE_EXTERNAL_MZCRC)