Added CUDA memory allocation/free tracing facility to GPUMatrx and GP…

…USparseMatrix
thanaphatp · Jan 20, 2016 · cf14913 · cf14913
1 parent 01bdefe
commit cf14913
Show file tree

Hide file tree

Showing 5 changed files with 160 additions and 113 deletions.
diff --git a/Source/CNTK/CNTK.cpp b/Source/CNTK/CNTK.cpp
@@ -19,6 +19,7 @@
 #include "SynchronousExecutionEngine.h"
 #include "ModelEditLanguage.h"
 #include "CPUMatrix.h" // used for SetNumThreads()
+#include "GPUMatrix.h"
 #include "SGD.h"
 #include "MPIWrapper.h"
 #include "Config.h"
@@ -473,6 +474,8 @@ int wmainWithBS(int argc, wchar_t* argv[]) // called from wmain which is a wrapp
 
     g_shareNodeValueMatrices = config(L"shareNodeValueMatrices", false);
 
+    SetTraceGPUMemoryAllocations(config(L"traceGPUMemoryAllocations", false));
+
     // logging
     wstring logpath = config(L"stderr", L"");
     if (logpath != L"")
@@ -560,6 +563,8 @@ int wmainOldCNTKConfig(int argc, wchar_t* argv[]) // called from wmain which is
 
     g_shareNodeValueMatrices = config(L"shareNodeValueMatrices", false);
 
+    SetTraceGPUMemoryAllocations(config(L"traceGPUMemoryAllocations", false));
+
     if (logpath != L"")
     {
         for (int i = 0; i < command.size(); i++)

diff --git a/Source/Math/CommonMatrix.h b/Source/Math/CommonMatrix.h
@@ -40,6 +40,15 @@ MATH_API DEVICEID_TYPE EnforceOneGPUOnly(DEVICEID_TYPE requestedDeviceId);
 
 namespace Microsoft { namespace MSR { namespace CNTK {
 
+template<typename AllocatedElemType>
+AllocatedElemType* AllocateDeviceMemory(int deviceId, size_t numRows, size_t numCols);
+
+template<typename AllocatedElemType>
+AllocatedElemType* AllocateDeviceMemory(int deviceId, size_t numElements, bool disableTrace = false);
+
+template<typename AllocatedElemType>
+void FreeDeviceMemory(int deviceId, AllocatedElemType* bufferPtr);
+
 // -----------------------------------------------------------------------
 // ElementWiseOperator -- This enum represents which function to apply.
 // This is shared between all matrix types and tensors.