1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
|
#include "THCAllocator.h"
static void *THCudaHostAllocator_malloc(void* ctx, ptrdiff_t size) {
void* ptr;
if (size < 0) THError("Invalid memory size: %ld", size);
if (size == 0) return NULL;
THCudaCheck(cudaMallocHost(&ptr, size));
return ptr;
}
static void THCudaHostAllocator_free(void* ctx, void* ptr) {
if (!ptr) return;
THCudaCheck(cudaFreeHost(ptr));
}
THAllocator THCudaHostAllocator = {
&THCudaHostAllocator_malloc,
NULL,
&THCudaHostAllocator_free
};
static cudaError_t THCIpcAllocator_malloc(void* ctx, void** devPtr, size_t size, cudaStream_t stream)
{
THError("THCIpcAllocator.malloc() not supported");
return cudaSuccess;
}
static cudaError_t THCIpcAllocator_free(void* ctx, void* devPtr)
{
cudaError_t err;
int prev_device;
int device = (int)(long)ctx;
err = cudaGetDevice(&prev_device);
if (err != cudaSuccess) { return err; }
err = cudaSetDevice(device);
if (err != cudaSuccess) { return err; }
err = cudaIpcCloseMemHandle(devPtr);
cudaSetDevice(prev_device);
return err;
}
THCDeviceAllocator THCIpcAllocator = {
&THCIpcAllocator_malloc,
NULL,
&THCIpcAllocator_free,
NULL,
NULL
};
static void *THCUVAAllocator_alloc(void* ctx, ptrdiff_t size) {
if (size < 0) THError("Invalid memory size: %ld", size);
if (size == 0) return NULL;
// See J.1.1 of the CUDA_C_Programming_Guide.pdf for UVA and coherence rules
// on various compute capabilities.
void* ptr;
THCudaCheck(cudaMallocManaged(&ptr, size, cudaMemAttachGlobal));
return ptr;
}
static void THCUVAAllocator_free(void* ctx, void* ptr) {
if (!ptr) return;
THCudaCheck(cudaFree(ptr));
}
THAllocator THCUVAAllocator = {
&THCUVAAllocator_alloc,
NULL,
&THCUVAAllocator_free
};
|
