fix the bug of incorrect pointer to the head of memory pieces

source/tensor/XMem.cpp

@@ -685,13 +685,16 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
             int bi;
             int bi;
             for(bi = 0; bi < blockNum; bi++){
             for(bi = 0; bi < blockNum; bi++){
                 XMemBlock * block = blocks + bi;
                 XMemBlock * block = blocks + bi;
-                if(block->mem == NULL){
+                if (block->mem != NULL && (block->head != NULL || block->size < mySize + 2 * MY_PITCH))
+                    continue;
+                
+                if (block->mem == NULL) {
                     block->size = MAX(maxBlockSize, mySize + 2 * MY_PITCH);
                     block->size = MAX(maxBlockSize, mySize + 2 * MY_PITCH);
                     if (myDevID < 0) {
                     if (myDevID < 0) {
                         block->mem = new char[block->size];
                         block->mem = new char[block->size];
                         memset(block->mem, 0, block->size);
                         memset(block->mem, 0, block->size);
                     }
                     }
-                    else{
+                    else {
 #ifdef USE_CUDA
 #ifdef USE_CUDA
                         int devIDBackup = -1;
                         int devIDBackup = -1;
                         cudaGetDevice(&devIDBackup);
                         cudaGetDevice(&devIDBackup);
@@ -706,31 +709,31 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
                         ShowNTErrors("Please specify USE_CUDA for compiling this program.");
                         ShowNTErrors("Please specify USE_CUDA for compiling this program.");
 #endif
 #endif
                     }
                     }
+                }
 
 
-                    curBlockID = MAX(curBlockID, bi);
+                curBlockID = MAX(curBlockID, bi);
                     
                     
-                    /* make a new index node */
-                    MPieceNode * newNode = memIndex + nodeNumUsed++;
-                    newNode->head.indexNode = newNode;
-                    newNode->p = block->mem;
-                    newNode->pReal = NULL;
-                    newNode->size = (char*)block->mem + mySize -
-                                    (char*)GetPitchedAddress(block->mem, MY_PITCH);
+                /* make a new index node */
+                MPieceNode * newNode = memIndex + nodeNumUsed++;
+                newNode->head.indexNode = newNode;
+                newNode->p = block->mem;
+                newNode->pReal = NULL;
+                newNode->size = (char*)block->mem + mySize -
+                                (char*)GetPitchedAddress(block->mem, MY_PITCH);
                     
                     
-                    AddFreeIndexNode(newNode);
+                AddFreeIndexNode(newNode);
                     
                     
-                    MHeader &header = newNode->head;
-                    header.state = 1;
-                    header.size = block->size;
-                    header.pre = NULL;
-                    header.next = NULL;
-                    header.blockID = bi;
+                MHeader &header = newNode->head;
+                header.state = 1;
+                header.size = block->size;
+                header.pre = NULL;
+                header.next = NULL;
+                header.blockID = bi;
                     
                     
-                    block->head = &header;
+                block->head = &header;
                     
                     
-                    result = AllocStandard(myDevID, mySize, myIsRebuiltIndex);
-                    break;
-                }
+                result = AllocStandard(myDevID, mySize, myIsRebuiltIndex);
+                break;
             }
             }
             CheckNTErrors(bi < blockNum, "No enough memory is available!");
             CheckNTErrors(bi < blockNum, "No enough memory is available!");
         }
         }