allocating the global memory pool when it is requested

source/tensor/XMem.cpp

@@ -53,6 +53,7 @@ XMem::XMem()
     strcpy(name, "xmem");
     signature = 0;
     mergeFreeOTF = true;
+    isInitialized = false;
 }
 
 /* 
@@ -169,6 +170,7 @@ void XMem::Initialize(int myDevID, MEMPOOL_MODE myMode, MTYPE myBlockSize, int m
 #endif
 
     signature++;
+    isInitialized = true;
 }
 
 /* free memory */
@@ -1579,11 +1581,6 @@ void XMemManager::Initialize()
     /* CPUs (we actually do not care about how many CPUs are using) */
     nCPUMem = 1;
 
-    MTYPE freeMem = GetAvailableMemory();
-    MTYPE myBufSize = 0;
-    GetBufferSize(freeMem, &myBufSize);
-    CPUMems[0].Initialize(-1, UNI_FREE, MIN_BLOCK_SIZE_FOR_MEMPOOL, MIN_BLOCK_NUM_FOR_MEMPOOL, myBufSize);
-
     /* GPUs */
     nGPUMem = 0;
 
@@ -1592,15 +1589,8 @@ void XMemManager::Initialize()
         XPRINT(0, stderr, "cannot get GPU information.");
         exit(1);
     }
-
-    for (int i = 0; i < nGPUMem; i++) {
-        MTYPE freeMem = GetAvailableGPUMemory(i);
-        MTYPE myBufSize = 0;
-        GetBufferSize(freeMem, &myBufSize);
-        GPUMems[i].Initialize(i, UNI_FREE, MIN_BLOCK_SIZE_FOR_MEMPOOL, MIN_BLOCK_NUM_FOR_MEMPOOL, myBufSize);
-    }
-
 #endif
+
 }
 
 /* free it */
@@ -1616,13 +1606,34 @@ void XMemManager::Free()
 XMem * XMemManager::GetMem(const int devID)
 {
     XMem * mem = NULL;
-    if (devID < 0)
+    if (devID < 0){
+        if(!CPUMems[0].isInitialized){
+            MTYPE freeMem = GetAvailableMemory();
+            MTYPE myBufSize = 0;
+            GetBufferSize(freeMem, &myBufSize);
+            CPUMems[0].Initialize(-1, FREE_ON_THE_FLY, 
+                                  MIN_BLOCK_SIZE_FOR_MEMPOOL, 
+                                  MIN_BLOCK_NUM_FOR_MEMPOOL, 
+                                  myBufSize);
+        }
         mem = CPUMems;
+    }
     else{
-        if (devID < nGPUMem)
+        if (devID < nGPUMem){
+            if(!GPUMems[devID].isInitialized){
+                MTYPE freeMem = GetAvailableGPUMemory(devID);
+                MTYPE myBufSize = 0;
+                GetBufferSize(freeMem, &myBufSize);
+                GPUMems[devID].Initialize(devID, FREE_ON_THE_FLY, 
+                                          MIN_BLOCK_SIZE_FOR_MEMPOOL, 
+                                          MIN_BLOCK_NUM_FOR_MEMPOOL, 
+                                          myBufSize);
+            }
             mem = GPUMems + devID;
-        else
+        }
+        else{
             XPRINT1(0, stderr, "Cannot get the memory (%d). Please check your device id!", devID);
+        }
     }
     
     return mem;

source/tensor/XMem.h

@@ -213,6 +213,9 @@ public:
     MTYPE curUsedPin;
     MTYPE bufUsedPin;
 
+    /* indicates whether the memory pool is initialized */
+    bool isInitialized;
+
 #ifdef USE_CUDA
     /* handle used for cublas */
     cublasHandle_t cublasHandle;
@@ -429,7 +432,7 @@ a class for the management of memory
 */
 class XMemManager
 {
-public:
+private:
     /* cpu memory pool information */
     XMem CPUMems[MAX_CPU_NUM];