mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/67746 Test Plan: Visual inspection. Sandcastle. Reviewed By: zertosh Differential Revision: D31986646 fbshipit-source-id: 91885c20c3cead3853c49abb9fe0a94a67f33cc8
37 lines
785 B
C++
37 lines
785 B
C++
#pragma once
|
|
|
|
#include <ATen/MapAllocator.h>
|
|
|
|
#ifdef __cplusplus
|
|
|
|
#ifdef SHM_EXPORTS
|
|
#define SHM_API __declspec(dllexport)
|
|
#else
|
|
#define SHM_API __declspec(dllimport)
|
|
#endif
|
|
|
|
SHM_API void libshm_init(const char* manager_exec_path);
|
|
|
|
class SHM_API THManagedMapAllocator : public at::RefcountedMapAllocator {
|
|
public:
|
|
THManagedMapAllocator(
|
|
const char* manager_handle,
|
|
const char* filename,
|
|
int flags,
|
|
ptrdiff_t size)
|
|
: at::RefcountedMapAllocator(filename, flags, size) {}
|
|
|
|
static at::DataPtr makeDataPtr(
|
|
const char* manager_handle,
|
|
const char* filename,
|
|
int flags,
|
|
ptrdiff_t size);
|
|
static THManagedMapAllocator* fromDataPtr(const at::DataPtr&);
|
|
|
|
const char* manager_handle() const {
|
|
return "no_manager";
|
|
}
|
|
};
|
|
|
|
#endif
|