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

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

Browse Source 
`libshm.so` depends on the torch library exclusively for `at::RefcountedMapAllocator`,
 so it makes sense to move it to c10 along with the other memory allocators.

This means `libshm.so` only depends on `c10` and we don't need to relink
`libshm.so` for every ATen change.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/109881
Approved by: https://github.com/albanD
This commit is contained in:
Peter Bell
2023-10-09 17:43:42 +00:00
committed by PyTorch MergeBot
parent 3704bf4ee8
commit 0341deb1c7
17 changed files with 1154 additions and 798 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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 torch)
target_link_libraries(shm PUBLIC c10)
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),
at::RefcountedMapAllocator(filename, flags, size) {}
c10::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_);
at::RefcountedMapAllocator::close();
c10::RefcountedMapAllocator::close();
socket.register_deallocation(info);
}
@ -139,7 +139,7 @@ static void deleteTHManagedMapAllocator(void* ptr) {
delete static_cast<THManagedMapAllocator*>(ptr);
}
at::DataPtr THManagedMapAllocator::makeDataPtr(
c10::DataPtr THManagedMapAllocator::makeDataPtr(
const char* manager_handle,
const char* filename,
int flags,
@ -150,10 +150,10 @@ at::DataPtr THManagedMapAllocator::makeDataPtr(
context->data(),
context,
&deleteTHManagedMapAllocator,
at::DeviceType::CPU};
c10::DeviceType::CPU};
}
THManagedMapAllocator* THManagedMapAllocator::fromDataPtr(
const at::DataPtr& dptr) {
const c10::DataPtr& dptr) {
return dptr.cast_context<THManagedMapAllocator>(&deleteTHManagedMapAllocator);
}

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

@ -1,23 +1,23 @@
#pragma once
#include <ATen/MapAllocator.h>
#include <c10/core/MapAllocator.h>
#ifdef __cplusplus
void libshm_init(const char* manager_exec_path);
// Superclass to run a constructor before at::RefcountedMapAllocator
// Superclass to run a constructor before c10::RefcountedMapAllocator
class THManagedMapAllocatorInit {
protected:
THManagedMapAllocatorInit(const char* manager_handle, const char* filename);
std::string manager_handle_;
};
// Like a at::RefcountedMapAllocator, but it also makes use of an external
// Like a c10::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 at::RefcountedMapAllocator {
public c10::RefcountedMapAllocator {
public:
THManagedMapAllocator(
const char* manager_handle,
@ -31,12 +31,12 @@ class THManagedMapAllocator : private THManagedMapAllocatorInit,
close();
}
static at::DataPtr makeDataPtr(
static c10::DataPtr makeDataPtr(
const char* manager_handle,
const char* filename,
int flags,
size_t size);
static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
static THManagedMapAllocator* fromDataPtr(const c10::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);
}
at::DataPtr THManagedMapAllocator::makeDataPtr(
c10::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, at::kCPU};
return {context->data(), context, &deleteTHManagedMapAllocator, c10::kCPU};
}
THManagedMapAllocator* THManagedMapAllocator::fromDataPtr(
const at::DataPtr& dptr) {
const c10::DataPtr& dptr) {
return dptr.cast_context<THManagedMapAllocator>(&deleteTHManagedMapAllocator);
}

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

@ -1,6 +1,6 @@
#pragma once
#include <ATen/MapAllocator.h>
#include <c10/core/MapAllocator.h>
#ifdef __cplusplus
@ -12,21 +12,21 @@
SHM_API void libshm_init(const char* manager_exec_path);
class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator {
class SHM_API THManagedMapAllocator : public c10::RefcountedMapAllocator {
public:
THManagedMapAllocator(
const char* manager_handle,
const char* filename,
int flags,
size_t size)
: at::RefcountedMapAllocator(filename, flags, size) {}
: c10::RefcountedMapAllocator(filename, flags, size) {}
static at::DataPtr makeDataPtr(
static c10::DataPtr makeDataPtr(
const char* manager_handle,
const char* filename,
int flags,
size_t size);
static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
static THManagedMapAllocator* fromDataPtr(const c10::DataPtr&);
const char* manager_handle() const {
return "no_manager";