Skip to content

Commit 649de92

Browse files
soumithsoumith
authored
Merge pull request #753 from apaszke/master
Add new shared memory allocator to TH
2 parents 65255ba + 42af65d commit 649de92

File tree

5 files changed

+176
-29
lines changed

5 files changed

+176
-29
lines changed

‎lib/TH/THAllocator.c‎

Lines changed: 138 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include "THAllocator.h"
2+
#include "THAtomic.h"
23

34
/* stuff for mapped files */
45
#ifdef _WIN32
@@ -36,22 +37,40 @@ THAllocator THDefaultAllocator = {
3637

3738
struct THMapAllocatorContext_ {
3839
char *filename; /* file name */
39-
int shared; /* is shared or not */
40+
int flags;
4041
long size; /* mapped size */
4142
};
4243

43-
THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int shared)
44+
#define TH_ALLOC_ALIGNMENT 64
45+
46+
typedef struct {
47+
int refcount;
48+
} THMapInfo;
49+
50+
THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags)
4451
{
4552
THMapAllocatorContext *ctx = THAlloc(sizeof(THMapAllocatorContext));
4653

54+
55+
if (!(flags & TH_ALLOCATOR_MAPPED_SHARED) && !(flags & TH_ALLOCATOR_MAPPED_SHAREDMEM))
56+
flags &= ~TH_ALLOCATOR_MAPPED_NOCREATE;
57+
if ((flags ^ TH_ALLOCATOR_MAPPED_EXCLUSIVE) == 0)
58+
THError("TH_ALLOCATOR_MAPPED_EXCLUSIVE flag requires opening the file "
59+
"in shared mode");
60+
4761
ctx->filename = THAlloc(strlen(filename)+1);
4862
strcpy(ctx->filename, filename);
49-
ctx->shared = shared;
63+
ctx->flags = flags;
5064
ctx->size = 0;
5165

5266
return ctx;
5367
}
5468

69+
char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx)
70+
{
71+
return ctx->filename;
72+
}
73+
5574
long THMapAllocatorContext_size(THMapAllocatorContext *ctx)
5675
{
5776
return ctx->size;
@@ -63,7 +82,7 @@ void THMapAllocatorContext_free(THMapAllocatorContext *ctx)
6382
THFree(ctx);
6483
}
6584

66-
static void *THMapAllocator_alloc(void* ctx_, long size)
85+
static void *_map_alloc(void* ctx_, long size)
6786
{
6887
THMapAllocatorContext *ctx = ctx_;
6988
void *data = NULL;
@@ -75,9 +94,14 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
7594
DWORD size_hi, size_lo;
7695
size_t hfilesz;
7796

97+
if (ctx->flags & TH_ALLOCATOR_MAPPED_EXCLUSIVE)
98+
THError("exclusive file mapping is not supported on Windows");
99+
if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE)
100+
THError("file mapping without creation is not supported on Windows");
101+
78102
/* open file */
79103
/* FILE_FLAG_RANDOM_ACCESS ? */
80-
if(ctx->shared)
104+
if(ctx->flags)
81105
{
82106
hfile = CreateFileA(ctx->filename, GENERIC_READ|GENERIC_WRITE, FILE_SHARE_WRITE|FILE_SHARE_READ, 0, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
83107
if (hfile == INVALID_HANDLE_VALUE)
@@ -103,7 +127,7 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
103127
{
104128
if(size > hfilesz)
105129
{
106-
if(ctx->shared)
130+
if(ctx->flags)
107131
{
108132
#if SIZEOF_SIZE_T > 4
109133
size_hi = (DWORD)((size) >> 32);
@@ -144,7 +168,7 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
144168
#endif
145169

146170
/* get map handle */
147-
if(ctx->shared)
171+
if(ctx->flags)
148172
{
149173
if( (hmfile = CreateFileMapping(hfile, NULL, PAGE_READWRITE, size_hi, size_lo, NULL)) == NULL )
150174
THError("could not create a map on file <%s>", ctx->filename);
@@ -156,30 +180,41 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
156180
}
157181

158182
/* map the stuff */
159-
if(ctx->shared)
183+
if(ctx->flags)
160184
data = MapViewOfFile(hmfile, FILE_MAP_ALL_ACCESS, 0, 0, 0);
161185
else
162186
data = MapViewOfFile(hmfile, FILE_MAP_COPY, 0, 0, 0);
163187

164-
CloseHandle(hfile);
165-
CloseHandle(hmfile);
188+
CloseHandle(hfile);
189+
CloseHandle(hmfile);
166190
}
167191
#else /* _WIN32 */
168192
{
169193
/* open file */
170194
int fd;
195+
int flags;
171196
long fdsz;
172197

173-
if(ctx->shared == TH_ALLOCATOR_MAPPED_SHARED)
198+
if (ctx->flags)
199+
flags = O_RDWR | O_CREAT;
200+
else
201+
flags = O_RDONLY;
202+
203+
if (ctx->flags & TH_ALLOCATOR_MAPPED_EXCLUSIVE)
204+
flags |= O_EXCL;
205+
if (ctx->flags & TH_ALLOCATOR_MAPPED_NOCREATE)
206+
flags &= ~O_CREAT;
207+
208+
if(ctx->flags & TH_ALLOCATOR_MAPPED_SHARED)
174209
{
175-
if((fd = open(ctx->filename, O_RDWR | O_CREAT, (mode_t)0600)) == -1)
210+
if((fd = open(ctx->filename, flags, (mode_t)0600)) == -1)
176211
THError("unable to open file <%s> in read-write mode", ctx->filename);
177212
}
178-
else if (ctx->shared == TH_ALLOCATOR_MAPPED_SHAREDMEM)
213+
else if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
179214
{
180215
#ifdef HAVE_SHM_OPEN
181-
if((fd = shm_open(ctx->filename, O_RDWR | O_CREAT, (mode_t)0600)) == -1)
182-
THError("unable to open file <%s> in read-write mode", ctx->filename);
216+
if((fd = shm_open(ctx->filename, flags, (mode_t)0600)) == -1)
217+
THError("unable to open shared memory object <%s> in read-write mode", ctx->filename);
183218
#else
184219
THError("unable to open file <%s> in sharedmem mode, shm_open unavailable on this platform");
185220
#endif
@@ -189,19 +224,21 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
189224
if((fd = open(ctx->filename, O_RDONLY)) == -1)
190225
THError("unable to open file <%s> in read-only mode", ctx->filename);
191226
}
227+
192228
if((fdsz = lseek(fd, 0, SEEK_END)) == -1)
193229
{
194230
close(fd);
195231
THError("unable to seek at end of file <%s>", ctx->filename);
196232
}
233+
197234
if(size > 0)
198235
{
199236
if(size > fdsz)
200237
{
201-
if(ctx->shared)
238+
if(ctx->flags)
202239
{
203240
/* if it is shared mem, let's put it in correct size */
204-
if (ctx->shared == TH_ALLOCATOR_MAPPED_SHAREDMEM)
241+
if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
205242
{
206243
if(ftruncate(fd, size) == -1)
207244
THError("unable to resize shared memory file <%s> to the right size", ctx->filename);
@@ -228,9 +265,9 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
228265
size = fdsz;
229266

230267
ctx->size = size; /* if we are here, it must be the right size */
231-
268+
232269
/* map it */
233-
if(ctx->shared)
270+
if(ctx->flags)
234271
data = mmap(NULL, ctx->size, PROT_READ|PROT_WRITE, MAP_SHARED, fd, 0);
235272
else
236273
data = mmap(NULL, ctx->size, PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0);
@@ -249,6 +286,10 @@ static void *THMapAllocator_alloc(void* ctx_, long size)
249286
return data;
250287
}
251288

289+
static void * THMapAllocator_alloc(void *ctx, long size) {
290+
return _map_alloc(ctx, size);
291+
}
292+
252293
static void *THMapAllocator_realloc(void* ctx, void* ptr, long size) {
253294
THError("cannot realloc mapped data");
254295
return NULL;
@@ -260,10 +301,10 @@ static void THMapAllocator_free(void* ctx_, void* data) {
260301
#ifdef _WIN32
261302
if(!UnmapViewOfFile((LPINT)data))
262303
THError("could not unmap the shared memory file");
263-
#else
304+
#else /* _WIN32 */
264305
if (munmap(data, ctx->size))
265306
THError("could not unmap the shared memory file");
266-
if (ctx->shared == TH_ALLOCATOR_MAPPED_SHAREDMEM)
307+
if (ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM)
267308
{
268309
#ifdef HAVE_SHM_UNLINK
269310
if (shm_unlink(ctx->filename) == -1)
@@ -272,14 +313,14 @@ static void THMapAllocator_free(void* ctx_, void* data) {
272313
THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename);
273314
#endif
274315
}
275-
#endif
316+
#endif /* _WIN32 */
276317

277318
THMapAllocatorContext_free(ctx);
278319
}
279320

280321
#else
281322

282-
THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int shared) {
323+
THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags) {
283324
THError("file mapping not supported on your system");
284325
return NULL;
285326
}
@@ -304,8 +345,82 @@ static void THMapAllocator_free(void* ctx, void* data) {
304345

305346
#endif
306347

348+
#if (defined(_WIN32) || defined(HAVE_MMAP)) && defined(TH_ATOMIC_IPC_REFCOUNT)
349+
350+
static void * THRefcountedMapAllocator_alloc(void *_ctx, long size) {
351+
THMapAllocatorContext *ctx = _ctx;
352+
353+
if (!(ctx->flags & TH_ALLOCATOR_MAPPED_SHAREDMEM))
354+
THError("THRefcountedMapAllcator requires SHAREDMEM flag");
355+
356+
size = size + TH_ALLOC_ALIGNMENT;
357+
void *ptr = _map_alloc(ctx, size);
358+
char *data = ((char*)ptr) + TH_ALLOC_ALIGNMENT;
359+
THMapInfo *map_info = (THMapInfo*)ptr;
360+
361+
if (ctx->flags & TH_ALLOCATOR_MAPPED_EXCLUSIVE)
362+
map_info->refcount = 1;
363+
else
364+
THAtomicIncrementRef(&map_info->refcount);
365+
366+
return (void*)data;
367+
}
368+
369+
static void *THRefcountedMapAllocator_realloc(void* ctx, void* ptr, long size) {
370+
THError("cannot realloc mapped data");
371+
return NULL;
372+
}
373+
374+
static void THRefcountedMapAllocator_free(void* ctx_, void* data) {
375+
THMapAllocatorContext *ctx = ctx_;
376+
377+
#ifdef _WIN32
378+
if(!UnmapViewOfFile((LPINT)data))
379+
THError("could not unmap the shared memory file");
380+
#else /* _WIN32 */
381+
382+
THMapInfo *info = (THMapInfo*)(((char*)data) - TH_ALLOC_ALIGNMENT);
383+
if (THAtomicDecrementRef(&info->refcount)) {
384+
#ifdef HAVE_SHM_UNLINK
385+
if (shm_unlink(ctx->filename) == -1)
386+
THError("could not unlink the shared memory file %s", ctx->filename);
387+
#else
388+
THError("could not unlink the shared memory file %s, shm_unlink not available on platform", ctx->filename);
389+
#endif /* HAVE_SHM_UNLINK */
390+
}
391+
if (munmap(info, ctx->size))
392+
THError("could not unmap the shared memory file %s", ctx->filename);
393+
#endif /* _WIN32 */
394+
395+
THMapAllocatorContext_free(ctx);
396+
}
397+
398+
#else
399+
400+
static void * THRefcountedMapAllocator_alloc(void *ctx, long size) {
401+
THError("refcounted file mapping not supported on your system");
402+
return NULL;
403+
}
404+
405+
static void *THRefcountedMapAllocator_realloc(void* ctx, void* ptr, long size) {
406+
THError("refcounted file mapping not supported on your system");
407+
return NULL;
408+
}
409+
410+
static void THRefcountedMapAllocator_free(void* ctx_, void* data) {
411+
THError("refcounted file mapping not supported on your system");
412+
}
413+
414+
#endif
415+
307416
THAllocator THMapAllocator = {
308417
&THMapAllocator_alloc,
309418
&THMapAllocator_realloc,
310419
&THMapAllocator_free
311420
};
421+
422+
THAllocator THRefcountedMapAllocator = {
423+
&THRefcountedMapAllocator_alloc,
424+
&THRefcountedMapAllocator_realloc,
425+
&THRefcountedMapAllocator_free
426+
};

‎lib/TH/THAllocator.h‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
#define TH_ALLOCATOR_MAPPED_SHARED 1
77
#define TH_ALLOCATOR_MAPPED_SHAREDMEM 2
8+
#define TH_ALLOCATOR_MAPPED_EXCLUSIVE 4
9+
#define TH_ALLOCATOR_MAPPED_NOCREATE 8
810

911
/* Custom allocator
1012
*/
@@ -22,10 +24,12 @@ extern THAllocator THDefaultAllocator;
2224
/* file map allocator
2325
*/
2426
typedef struct THMapAllocatorContext_ THMapAllocatorContext;
25-
THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int shared);
26-
long THMapAllocatorContext_size(THMapAllocatorContext *ctx);
27-
void THMapAllocatorContext_free(THMapAllocatorContext *ctx);
27+
TH_API THMapAllocatorContext *THMapAllocatorContext_new(const char *filename, int flags);
28+
TH_API char * THMapAllocatorContext_filename(THMapAllocatorContext *ctx);
29+
TH_API long THMapAllocatorContext_size(THMapAllocatorContext *ctx);
30+
TH_API void THMapAllocatorContext_free(THMapAllocatorContext *ctx);
2831

2932
extern THAllocator THMapAllocator;
33+
extern THAllocator THRefcountedMapAllocator;
3034

3135
#endif

‎lib/TH/THAtomic.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,11 @@ TH_API long THAtomicAddLong(long volatile *a, long value);
8686
*/
8787
TH_API long THAtomicCompareAndSwapLong(long volatile *a, long oldvalue, long newvalue);
8888

89+
#if defined(USE_C11_ATOMICS) && defined(ATOMIC_INT_LOCK_FREE) && \
90+
ATOMIC_INT_LOCK_FREE == 2
91+
#define TH_ATOMIC_IPC_REFCOUNT 1
92+
#elif defined(USE_MSC_ATOMICS) || defined(USE_GCC_ATOMICS)
93+
#define TH_ATOMIC_IPC_REFCOUNT 1
94+
#endif
95+
8996
#endif

‎lib/TH/generic/THStorage.c‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,9 @@ THStorage* THStorage_(newWithAllocator)(long size,
4141
return storage;
4242
}
4343

44-
THStorage* THStorage_(newWithMapping)(const char *filename, long size, int shared)
44+
THStorage* THStorage_(newWithMapping)(const char *filename, long size, int flags)
4545
{
46-
THMapAllocatorContext *ctx = THMapAllocatorContext_new(filename, shared);
46+
THMapAllocatorContext *ctx = THMapAllocatorContext_new(filename, flags);
4747

4848
THStorage *storage = THStorage_(newWithAllocator)(size,
4949
&THMapAllocator,
@@ -203,4 +203,24 @@ real THStorage_(get)(const THStorage *self, long idx)
203203
return self->data[idx];
204204
}
205205

206+
void THStorage_(swap)(THStorage *storage1, THStorage *storage2)
207+
{
208+
#define SWAP(val) { val = storage1->val; storage1->val = storage2->val; storage2->val = val; }
209+
real *data;
210+
long size;
211+
char flag;
212+
THAllocator *allocator;
213+
void *allocatorContext;
214+
struct THStorage *view;
215+
216+
SWAP(data);
217+
SWAP(size);
218+
SWAP(flag);
219+
// don't swap refcount!
220+
SWAP(allocator);
221+
SWAP(allocatorContext);
222+
SWAP(view);
223+
#undef SWAP
224+
}
225+
206226
#endif

‎lib/TH/generic/THStorage.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ TH_API THStorage* THStorage_(newWithSize1)(real);
4646
TH_API THStorage* THStorage_(newWithSize2)(real, real);
4747
TH_API THStorage* THStorage_(newWithSize3)(real, real, real);
4848
TH_API THStorage* THStorage_(newWithSize4)(real, real, real, real);
49-
TH_API THStorage* THStorage_(newWithMapping)(const char *filename, long size, int shared);
49+
TH_API THStorage* THStorage_(newWithMapping)(const char *filename, long size, int flags);
5050

5151
/* takes ownership of data */
5252
TH_API THStorage* THStorage_(newWithData)(real *data, long size);
@@ -61,6 +61,7 @@ TH_API THStorage* THStorage_(newWithDataAndAllocator)(
6161
TH_API void THStorage_(setFlag)(THStorage *storage, const char flag);
6262
TH_API void THStorage_(clearFlag)(THStorage *storage, const char flag);
6363
TH_API void THStorage_(retain)(THStorage *storage);
64+
TH_API void THStorage_(swap)(THStorage *storage1, THStorage *storage2);
6465

6566
/* might differ with other API (like CUDA) */
6667
TH_API void THStorage_(free)(THStorage *storage);

0 commit comments

Comments
 (0)