Use size_t in THManagedMapAllocator (#103331)

When reviewing the source code, I found the ptrdiff_t size in THManagedMapAllocator::THManagedMapAllocator can be changed to size_t size to avoid unnecessary casts.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/103331
Approved by: https://github.com/malfet
This commit is contained in:
cyy 2023-06-13 04:50:25 +00:00 committed by PyTorch MergeBot
parent 23892d8ee4
commit db07ba3a9b
5 changed files with 18 additions and 13 deletions

View File

@ -66,18 +66,21 @@ static PyObject* THPStorage_pyNewFilenameStorage(
PyObject* _unused, PyObject* _unused,
PyObject* args) { PyObject* args) {
HANDLE_TH_ERRORS HANDLE_TH_ERRORS
// NOLINTNEXTLINE(cppcoreguidelines-init-variables) long long size = 0;
long long size;
if (!PyArg_ParseTuple(args, "L", &size)) { if (!PyArg_ParseTuple(args, "L", &size)) {
return nullptr; return nullptr;
} }
if (size < 0) {
return nullptr;
}
int flags = at::ALLOCATOR_MAPPED_SHAREDMEM | at::ALLOCATOR_MAPPED_EXCLUSIVE; int flags = at::ALLOCATOR_MAPPED_SHAREDMEM | at::ALLOCATOR_MAPPED_EXCLUSIVE;
std::string handle = at::NewProcessWideShmHandle(); std::string handle = at::NewProcessWideShmHandle();
return THPStorage_New(c10::make_intrusive<at::StorageImpl>( return THPStorage_New(c10::make_intrusive<at::StorageImpl>(
c10::StorageImpl::use_byte_size_t(), c10::StorageImpl::use_byte_size_t(),
size, size,
THManagedMapAllocator::makeDataPtr("", handle.c_str(), flags, size), THManagedMapAllocator::makeDataPtr(
"", handle.c_str(), flags, static_cast<size_t>(size)),
/*allocator=*/nullptr, /*allocator=*/nullptr,
/*resizable=*/false)); /*resizable=*/false));
END_HANDLE_TH_ERRORS END_HANDLE_TH_ERRORS
@ -163,7 +166,7 @@ static PyObject* THPStorage_newSharedFilename(
} }
const char* manager_handle = PyBytes_AS_STRING(_manager_handle); const char* manager_handle = PyBytes_AS_STRING(_manager_handle);
const char* object_handle = PyBytes_AS_STRING(_object_handle); const char* object_handle = PyBytes_AS_STRING(_object_handle);
int64_t size = THPUtils_unpackLong(_size); uint64_t size = THPUtils_unpackUInt64(_size);
int flags = at::ALLOCATOR_MAPPED_SHAREDMEM | at::ALLOCATOR_MAPPED_NOCREATE; int flags = at::ALLOCATOR_MAPPED_SHAREDMEM | at::ALLOCATOR_MAPPED_NOCREATE;
return THPStorage_New(c10::make_intrusive<at::StorageImpl>( return THPStorage_New(c10::make_intrusive<at::StorageImpl>(
c10::StorageImpl::use_byte_size_t(), c10::StorageImpl::use_byte_size_t(),
@ -177,11 +180,13 @@ static PyObject* THPStorage_newSharedFilename(
static PyObject* THPStorage_pyNewFdStorage(PyObject* _unused, PyObject* args) { static PyObject* THPStorage_pyNewFdStorage(PyObject* _unused, PyObject* args) {
HANDLE_TH_ERRORS HANDLE_TH_ERRORS
// NOLINTNEXTLINE(cppcoreguidelines-init-variables) long long size = 0;
long long size;
if (!PyArg_ParseTuple(args, "L", &size)) { if (!PyArg_ParseTuple(args, "L", &size)) {
return nullptr; return nullptr;
} }
if (size < 0) {
return nullptr;
}
return THPStorage_New(at::new_shm_fd_storage(size)); return THPStorage_New(at::new_shm_fd_storage(size));
END_HANDLE_TH_ERRORS END_HANDLE_TH_ERRORS
} }

View File

@ -120,7 +120,7 @@ THManagedMapAllocator::THManagedMapAllocator(
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size) size_t size)
: THManagedMapAllocatorInit(manager_handle, filename), : THManagedMapAllocatorInit(manager_handle, filename),
at::RefcountedMapAllocator(filename, flags, size) {} at::RefcountedMapAllocator(filename, flags, size) {}
@ -142,7 +142,7 @@ at::DataPtr THManagedMapAllocator::makeDataPtr(
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size) { size_t size) {
auto* context = auto* context =
new THManagedMapAllocator(manager_handle, filename, flags, size); new THManagedMapAllocator(manager_handle, filename, flags, size);
return { return {

View File

@ -23,7 +23,7 @@ class THManagedMapAllocator : private THManagedMapAllocatorInit,
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size); size_t size);
void close() override; void close() override;
@ -35,7 +35,7 @@ class THManagedMapAllocator : private THManagedMapAllocatorInit,
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size); size_t size);
static THManagedMapAllocator* fromDataPtr(const at::DataPtr&); static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
const char* manager_handle() const { const char* manager_handle() const {

View File

@ -14,7 +14,7 @@ at::DataPtr THManagedMapAllocator::makeDataPtr(
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size) { size_t size) {
auto* context = auto* context =
new THManagedMapAllocator(manager_handle, filename, flags, size); new THManagedMapAllocator(manager_handle, filename, flags, size);
return {context->data(), context, &deleteTHManagedMapAllocator, at::kCPU}; return {context->data(), context, &deleteTHManagedMapAllocator, at::kCPU};

View File

@ -18,14 +18,14 @@ class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator {
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size) size_t size)
: at::RefcountedMapAllocator(filename, flags, size) {} : at::RefcountedMapAllocator(filename, flags, size) {}
static at::DataPtr makeDataPtr( static at::DataPtr makeDataPtr(
const char* manager_handle, const char* manager_handle,
const char* filename, const char* filename,
int flags, int flags,
ptrdiff_t size); size_t size);
static THManagedMapAllocator* fromDataPtr(const at::DataPtr&); static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
const char* manager_handle() const { const char* manager_handle() const {