提交 f0549499 authored 作者: Ian Goodfellow's avatar Ian Goodfellow

Merge pull request #1205 from dwf/device_malloc_verbose_flag

Added a verbosity flag to device_malloc.
...@@ -53,7 +53,13 @@ struct table_struct{ ...@@ -53,7 +53,13 @@ struct table_struct{
}; };
table_struct _alloc_size_table[TABLE_SIZE]; table_struct _alloc_size_table[TABLE_SIZE];
#endif #endif
void * device_malloc(size_t size) void * device_malloc(size_t size)
{
return device_malloc(size, VERBOSE_DEVICE_MALLOC);
}
void * device_malloc(size_t size, int verbose)
{ {
void * rval=NULL; void * rval=NULL;
cudaError_t err = cudaMalloc(&rval, size); cudaError_t err = cudaMalloc(&rval, size);
...@@ -64,11 +70,14 @@ void * device_malloc(size_t size) ...@@ -64,11 +70,14 @@ void * device_malloc(size_t size)
// it returns something else I still don't see why we should ignore // it returns something else I still don't see why we should ignore
// it. All we want to do here is reset the flag. // it. All we want to do here is reset the flag.
cudaGetLastError(); cudaGetLastError();
if (verbose)
{
#if COMPUTE_GPU_MEM_USED #if COMPUTE_GPU_MEM_USED
fprintf(stderr, "Error allocating %li bytes of device memory (%s). new total bytes allocated: %d\n", (long)size, cudaGetErrorString(err),_allocated_size); fprintf(stderr, "Error allocating %li bytes of device memory (%s). new total bytes allocated: %d\n", (long)size, cudaGetErrorString(err),_allocated_size);
#else #else
fprintf(stderr, "Error allocating %li bytes of device memory (%s).\n", (long)size, cudaGetErrorString(err)); fprintf(stderr, "Error allocating %li bytes of device memory (%s).\n", (long)size, cudaGetErrorString(err));
#endif #endif
}
PyErr_Format(PyExc_MemoryError, PyErr_Format(PyExc_MemoryError,
"Error allocating %li bytes of device memory (%s).", (long)size, cudaGetErrorString(err)); "Error allocating %li bytes of device memory (%s).", (long)size, cudaGetErrorString(err));
return NULL; return NULL;
......
...@@ -42,6 +42,9 @@ typedef float real; ...@@ -42,6 +42,9 @@ typedef float real;
#define SHARED_SIZE (16*1024) #define SHARED_SIZE (16*1024)
#endif #endif
#define VERBOSE_DEVICE_MALLOC 1
#define NO_VERBOSE_DEVICE_MALLOC 0
/** /**
* Allocation and freeing of device memory should go through these functions so that the lib can track memory usage. * Allocation and freeing of device memory should go through these functions so that the lib can track memory usage.
* *
...@@ -49,6 +52,7 @@ typedef float real; ...@@ -49,6 +52,7 @@ typedef float real;
* device_free will return nonzero on failure (after setting the python error message) * device_free will return nonzero on failure (after setting the python error message)
*/ */
DllExport void * device_malloc(size_t size); DllExport void * device_malloc(size_t size);
DllExport void * device_malloc(size_t size, int verbose);
DllExport int device_free(void * ptr); DllExport int device_free(void * ptr);
template <typename T> template <typename T>
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论