Skip to content

N
NiuTrans.Tensor
Commit 91fdf94e by xiaotong
Options

rewrite parts of the memory pool to free the space on the fly

parent 97b41e35
正在显示 包含 351 行增加 和 130 行删除
source/tensor/XMem.cpp
......@@ -31,6 +31,9 @@
/* the nts (NiuTrans.Tensor) namespace */
namespace nts{
int testxmemid = 0;
void * recordp = NULL;
XMem * GMem;
/* constructor */
......@@ -83,6 +86,9 @@ XMem::~XMem()
#endif
Free();
delete[] name;
delete[] memIndex;
delete[] memIndex2;
delete[] minSizeIndex;
}
/*
......@@ -279,15 +285,19 @@ initialize the index
*/
void XMem::SetIndex(INT_64 indexSize, MTYPE minSizeFirst, int minSizeNum)
{
delete freeMemIndex;
delete minSizeIndex;
delete[] memIndex;
delete[] memIndex2;
delete[] minSizeIndex;
indexNodeNum = indexSize;
indexNodeNumUsed = minSizeNum;
nodeNum = indexSize;
nodeNumUsed = minSizeNum * 2;
indexEntryNum = minSizeNum;
freeMemIndex = new MPieceNode[indexNodeNum];
memset(freeMemIndex, 0, sizeof(MPieceNode) * indexNodeNum);
memIndex = new MPieceNode[nodeNum];
memset(memIndex, 0, sizeof(MPieceNode) * nodeNum);
memIndex2 = new MPieceNode[nodeNum];
memset(memIndex2, 0, sizeof(MPieceNode) * nodeNum);
minSizeIndex = new MTYPE[indexEntryNum];
memset(minSizeIndex, 0, sizeof(MTYPE) * indexEntryNum);
......@@ -496,21 +506,23 @@ void * XMem::AllocBuf(int myDevID, MTYPE mySize, int pitch)
/*
release a piece of memory
>> p - pointer to the memory piece we intend to release
>> size - size of the memory piece to release
*/
void XMem::Release(void * p)
void XMem::Release(void * p, MTYPE size)
{
Release(devID, p);
Release(devID, p, size);
}
/*
release a piece of memory
>> myDevID - device id
>> p - pointer to the memory piece we intend to release
>> size - size of the memory piece to release
*/
void XMem::Release(int myDevID, void * p)
void XMem::Release(int myDevID, void * p, MTYPE size)
{
if(mode == FREE_ON_THE_FLY)
ReleaseStandard(myDevID, p);
ReleaseStandard(myDevID, p, size);
}
/*
......@@ -553,7 +565,7 @@ allocate a piece of memory as "malloc"
*/
void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
{
CheckNTErrors(freeMemIndex != NULL, "The index of the memory pool is not initialized!");
CheckNTErrors(memIndex != NULL, "The index of the memory pool is not initialized!");
if(mySize <= minSizeIndex[0])
mySize = minSizeIndex[0];
......@@ -567,20 +579,17 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
/* search for the memory piece avialable for the allocation */
for(int i = indexEntryNum; i > index; i--){
if(i == indexEntryNum){
entry = freeMemIndex + index;
entry = memIndex + index;
CheckNTErrors(mySize >= minSizeIndex[index], "Wrong index!");
}
else
entry = freeMemIndex + i;
if(entry->size == 0)
continue;
entry = memIndex + i;
node = entry;
node = entry->next;
while(node != NULL){
if(node->size == 0){
MPieceNode * next = node->next;
RemoveIndexNode(node, entry);
RemoveFreeIndexNode(node, entry);
node = next;
}
else{
......@@ -598,40 +607,46 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
/* if a free memory piece is found, we allocate the memory on it. */
if(hit != NULL){
MHeader * head = (MHeader*)hit->p;
CheckNTErrors(head->state == 1, "Accessing the wrong memory piece!");
char * beg = (char*)GetPitchedAddress((char*)head + sizeof(MHeader) + sizeof(void*), MY_PITCH);
MHeader * head = &hit->head;
char * beg = (char*)GetPitchedAddress((char*)hit->p, MY_PITCH);
char * end = (char*)beg + mySize;
MTYPE needed = end - (char*)head;
MTYPE needed = end - (char*)hit->p;
MTYPE remaining = head->size - needed;
if(remaining >= minSizeIndex[0] + sizeof(MHeader) + sizeof(void*)){
/* make a new header for the remaining segment in the memory piece */
MHeader * newHead = (MHeader*)end;
newHead->state = 1;
newHead->size = remaining;
newHead->pre = head;
newHead->next = head->next;
newHead->blockID = head->blockID;
head->next = newHead;
head->size = needed;
if(newHead->next != NULL)
newHead->next->pre = newHead;
if(remaining >= minSizeIndex[0]){
/* make a new index node */
MPieceNode * newNode = freeMemIndex + indexNodeNumUsed++;
MPieceNode * newNode = memIndex + nodeNumUsed++;
newNode->head.indexNode = newNode;
newNode->p = end;
newNode->size = (char*)newHead + newHead->size -
(char*)GetPitchedAddress((char*)newNode->p + sizeof(MHeader) + sizeof(void*), MY_PITCH);
AddIndexNode(newNode);
newNode->pReal = NULL;
newNode->size = (char*)end + remaining -
(char*)GetPitchedAddress((char*)end, MY_PITCH);
AddFreeIndexNode(newNode);
/* connections for headers */
MHeader &cur = hit->head;
MHeader &next = newNode->head;
next.pre = &cur;
next.next = cur.next;
cur.next = &next;
if(cur.next != NULL)
cur.next->pre = &next;
cur.size = needed;
next.state = 1;
next.size = remaining;
next.blockID = cur.blockID;
}
head->state = 2;
void * backPointer = (char*)beg - sizeof(void*);
*((MHeader**)backPointer) = head;
hit->size = mySize;
hit->head.state = 2;
hit->pReal = beg;
hit->size = 0;
RemoveFreeIndexNode(hit);
AddAllocIndexNode(hit);
result = beg;
}
......@@ -648,7 +663,7 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
for(bi = 0; bi < blockNum; bi++){
XMemBlock * block = blocks + bi;
if(block->mem == NULL){
block->size = MAX(maxBlockSize, mySize + sizeof(MHeader) + sizeof(void*) + 2 * MY_PITCH);
block->size = MAX(maxBlockSize, mySize + 2 * MY_PITCH);
if (myDevID < 0) {
block->mem = new char[block->size];
memset(block->mem, 0, block->size);
......@@ -671,20 +686,24 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
curBlockID = MAX(curBlockID, bi);
/* make a new header for the remaining segment in the memory piece */
MHeader * newHead = (MHeader*)block->mem;
newHead->state = 1;
newHead->size = block->size;
newHead->pre = NULL;
newHead->next = NULL;
newHead->blockID = bi;
/* make a new index node */
MPieceNode * newNode = freeMemIndex + indexNodeNumUsed++;
newNode->p = newHead;
newNode->size = (char*)newHead + newHead->size -
(char*)GetPitchedAddress((char*)newHead + sizeof(MHeader) + sizeof(void*), MY_PITCH);
AddIndexNode(newNode);
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);
MHeader &header = newNode->head;
header.state = 1;
header.size = block->size;
header.pre = NULL;
header.next = NULL;
header.blockID = bi;
block->head = &header;
result = AllocStandard(myDevID, mySize, myIsRebuiltIndex);
break;
......@@ -695,11 +714,14 @@ void * XMem::AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex)
}
/* if all index nodes are used, we rebuild the index to release the nodes that are free */
if(indexNodeNumUsed == indexNodeNum){
if(nodeNumUsed == nodeNum){
RebuildIndex();
CheckNTErrors(indexNodeNumUsed < indexNodeNum, "No enough index nodes for the memory pool!");
CheckNTErrors(nodeNumUsed < nodeNum, "No enough index nodes for the memory pool!");
}
//if(testxmemid == 106)
// recordp = result;
return result;
}
......@@ -749,60 +771,72 @@ int XMem::FindIndexEntry(MTYPE mySize)
}
/*
remove an index node
remove an index node for available memory pieces
>> node - node to remove
>> - the entry of the list that keeps the node
*/
void XMem::RemoveIndexNode(MPieceNode * node, MPieceNode * entry)
void XMem::RemoveFreeIndexNode(MPieceNode * node, MPieceNode * entry)
{
MPieceNode * pre = node->pre;
MPieceNode * next = node->next;
if(pre == NULL){
if(entry == NULL){
entry = freeMemIndex + FindIndexEntry(node->size);
CheckNTErrors(entry == node, "Illegal index node!");
}
*entry = *next;
entry->pre = NULL;
}
else{
CheckNTErrors(pre != NULL, "cannot free the entry node!");
pre->next = next;
if(next != NULL)
next->pre = pre;
}
memset(node, 0, sizeof(MPieceNode));
node->pre = NULL;
node->next = NULL;
}
/*
add an index node
add an index node for available memory pieces
>> node - node to add
>> entry - the entry of the list to append the node
*/
void XMem::AddIndexNode(MPieceNode * node, MPieceNode * entry)
void XMem::AddFreeIndexNode(MPieceNode * node, MPieceNode * entry)
{
MPieceNode * entryForMe = entry != NULL ? entry :
freeMemIndex + FindIndexEntry(node->size);
memIndex + FindIndexEntry(node->size);
if(entryForMe->size == 0){
entryForMe->size = node->size;
entryForMe->p = node->p;
entryForMe->pre = NULL;
entryForMe->next = NULL;
}
else{
MTYPE tmpSize = entryForMe->size;
void * tmpP = entryForMe->p;
entryForMe->size = node->size;
entryForMe->p = node->p;
node->size = tmpSize;
node->p = tmpP;
node->next = entryForMe->next;
MPieceNode * backup = entryForMe->next;
entryForMe->next = node;
node->pre = entryForMe;
node->next = backup;
if(backup != NULL)
backup->pre = node;
CheckNTErrors(node != node->next, "Something wrong with the index node!");
CheckNTErrors(node != node->pre, "Something wrong with the index node!");
}
/*
remove an index node for memory pieces in use
>> node - node to remove
>> - the entry of the list that keeps the node
*/
void XMem::RemoveAllocIndexNode(MPieceNode * node, MPieceNode * entry)
{
RemoveFreeIndexNode(node, entry);
}
/*
add an index node for memory pieces in use
>> node - node to add
>> entry - the entry of the list to append the node
*/
void XMem::AddAllocIndexNode(MPieceNode * node, MPieceNode * entry)
{
MPieceNode * entryForMe = entry != NULL ? entry :
memIndex + indexEntryNum + FindIndexEntry(node->size);
MPieceNode * backup = entryForMe->next;
entryForMe->next = node;
}
node->pre = entryForMe;
node->next = backup;
if(backup != NULL)
backup->pre = node;
CheckNTErrors(node != node->next, "Something wrong with the index node!");
CheckNTErrors(node != node->pre, "Something wrong with the index node!");
......@@ -812,40 +846,72 @@ void XMem::AddIndexNode(MPieceNode * node, MPieceNode * entry)
release a piece of memory as "free"
>> myDevID - device id(-1: CPU memory, >=0: GPU device ID)
>> p - the pointer to the address of the memory we intend to free
>> size - size of the memory piece to release
*/
void XMem::ReleaseStandard(int myDevID, void * p)
void XMem::ReleaseStandard(int myDevID, void * p, MTYPE size)
{
if(p == NULL)
return;
void * back = (char*)p - sizeof(void*);
MHeader * head = *(MHeader**)back;
if(size <= minSizeIndex[0])
size = minSizeIndex[0];
MPieceNode * entry = NULL;
MPieceNode * node = NULL;
MPieceNode * hit = NULL;
MPieceNode * last = NULL;
entry = memIndex + indexEntryNum + FindIndexEntry(size);
CheckNTErrors(head->state == 2, "Illegal header of a memory piece!");
head->state = 1;
last = entry;
node = entry->next;
/* make a new index node */
MPieceNode * newNode = freeMemIndex + indexNodeNumUsed++;
newNode->p = head;
newNode->size = (char*)head + head->size -
(char*)GetPitchedAddress((char*)head + sizeof(MHeader) + sizeof(void*), MY_PITCH);
AddIndexNode(newNode);
while(node != NULL){
CheckNTErrors(node->pre == last, "Something is wrong!");
CheckNTErrors(last->next == node, "Something is wrong!");
CheckNTErrors(node->head.state == 2, "Something is wrong!");
last = node;
if(node->size == 0){
MPieceNode * next = node->next;
RemoveFreeIndexNode(node, entry);
node = next;
ShowNTErrors("Something is wrong!");
}
else{
CheckNTErrors(node->pReal != NULL, "Illegal pointer!");
if(node->pReal == p){
hit = node;
break;
}
node = node->next;
}
}
CheckNTErrors(hit != NULL, "No header is found!");
hit->head.state = 1;
RemoveAllocIndexNode(hit);
AddFreeIndexNode(hit);
}
/* rebuild index to merge small fragments of memory and free the block with no use */
void XMem::RebuildIndex()
{
indexNodeNumUsed = indexEntryNum;
memset(freeMemIndex, 0, sizeof(MPieceNode) * indexEntryNum);
int nodeNumUsed2 = indexEntryNum * 2;
memset(memIndex2, 0, sizeof(MPieceNode) * indexEntryNum * 2);
for(int bi = 0; bi <= curBlockID; bi++){
XMemBlock * block = blocks + bi;
if(block->mem == NULL)
continue;
MHeader * head = (MHeader*)block->mem;
MHeader * head = block->head;
CheckNTErrors(head->size <= block->size, "Illegal memory block!");
block->head = NULL;
/* if the block is not used, we delete it */
if(head->state == 1 && head->size == block->size){
if(devID < 0){
......@@ -870,6 +936,7 @@ void XMem::RebuildIndex()
/* if the block is in use, we build the index */
int pieceCount = 0;
MTYPE size = 0;
MHeader * newLast = NULL;
while(head != NULL){
MHeader * next = head->next;
if(head->state == 1){
......@@ -878,21 +945,93 @@ void XMem::RebuildIndex()
next = next->next;
}
head->next = next;
}
MPieceNode * node = head->indexNode;
void * p = node->p;
/* make a new index node */
MPieceNode * newNode = freeMemIndex + indexNodeNumUsed++;
newNode->p = head;
newNode->size = (char*)head + head->size -
(char*)GetPitchedAddress((char*)head + sizeof(MHeader) + sizeof(void*), MY_PITCH);
AddIndexNode(newNode);
MPieceNode * newNode = memIndex2 + nodeNumUsed2++;
newNode->p = p;
newNode->size = (char*)p + node->size - (char*)GetPitchedAddress((char*)p, MY_PITCH);
newNode->pre = NULL;
newNode->next = NULL;
CheckNTErrors(newNode->size > 0, "Illegal index node!");
MHeader * newHeader = &newNode->head;
newHeader->indexNode = newNode;
newHeader->pre = newLast;
newHeader->next = NULL;
newHeader->blockID = bi;
newHeader->size = head->size;
newHeader->state = head->state;
if(newLast != NULL)
newLast->next = newHeader;
newHeader->pre = newLast;
if(head->state == 1){
newNode->pReal = NULL;
MPieceNode * entry = memIndex2 + FindIndexEntry(newNode->size);
AddFreeIndexNode(newNode, entry);
}
size += head->size;
else{
newNode->pReal = head->indexNode->pReal;
MPieceNode * entry = memIndex2 + indexEntryNum + FindIndexEntry(newNode->size);
AddAllocIndexNode(newNode, entry);
}
if(newLast == NULL)
block->head = newHeader;
pieceCount++;
head = next;
size += head->size;
CheckNTErrors(size <= block->size, "Illegal block size!");
newLast = newHeader;
head = next;
}
}
}
MPieceNode * backup = memIndex2;
memIndex2 = memIndex;
memIndex = backup;
/*bool hhh = recordp != NULL ? false : true;
for(int i = 0; i < indexEntryNum; i++){
MPieceNode * entry = memIndex + indexEntryNum + i;
MPieceNode * last = entry;
MPieceNode * node = entry->next;
while(node != NULL){
CheckNTErrors(node->pre == last, "XSomething is wrong!");
CheckNTErrors(last->next == node, "XSomething is wrong!");
last = node;
if(node->pReal == recordp)
hhh = true;
if(node->size == 0){
MPieceNode * next = node->next;
node = next;
}
else{
CheckNTErrors(node->pReal != NULL, "Illegal pointer!");
node = node->next;
}
}
}
if(!hhh){
int nnn = 0;
}*/
nodeNumUsed = nodeNumUsed2;
}
/*
......
source/tensor/XMem.h
......@@ -60,6 +60,8 @@ mode of runnig a memory pool
*/
enum MEMPOOL_MODE {UNI_FREE, FREE_ON_THE_FLY};
struct MPieceNode;
/* header of a memory piece (FREE_ON_THE_FLY) */
struct MHeader
{
......@@ -80,6 +82,9 @@ struct MHeader
/* id of the memory block */
int blockID;
/* pointer to the index node */
MPieceNode * indexNode;
};
/* index of memory piece */
......@@ -97,8 +102,11 @@ struct MPieceNode
/* pointer to the head of a memory piece */
void * p;
/* pointer to the head of memory that is returned back to the user */
void * pReal;
/* header of the memory piece */
MHeader * head;
MHeader head;
};
/* memory block */
......@@ -115,6 +123,9 @@ struct XMemBlock
/* desired size of the block */
MTYPE sizeDesired;
/* first head of the block */
MHeader * head;
};
/*
......@@ -197,13 +208,16 @@ public:
public:
/* index of the free memory pieces */
MPieceNode * freeMemIndex;
MPieceNode * memIndex;
/* for double buffering */
MPieceNode * memIndex2;
/* maximum number of index nodes */
INT_64 indexNodeNum;
INT_64 nodeNum;
/* count of the used nodes */
INT_64 indexNodeNumUsed;
INT_64 nodeNumUsed;
/* minimal size allocation for each index entry */
MTYPE * minSizeIndex;
......@@ -285,10 +299,10 @@ public:
void * AllocBuf(int myDevID, MTYPE mySize, int pitch = BUF_PITCH);
/* release a piece of memory */
void Release(void * p);
void Release(void * p, MTYPE size);
/* release a piece of memory */
void Release(int myDevID, void * p);
void Release(int myDevID, void * p, MTYPE size);
/* release a piece of memory in the buffer */
void ReleaseBuf(int myDevID, MTYPE mySize, int pitch = BUF_PITCH);
......@@ -305,14 +319,20 @@ public:
/* find the index entry for allocation query */
int FindIndexEntry(MTYPE mySize);
/* remove an index node */
void RemoveIndexNode(MPieceNode * node, MPieceNode * entry = NULL);
/* remove an index node for available memory pieces */
void RemoveFreeIndexNode(MPieceNode * node, MPieceNode * entry = NULL);
/* add an index node for available memory pieces */
void AddFreeIndexNode(MPieceNode * node, MPieceNode * entry = NULL);
/* add an index node */
void AddIndexNode(MPieceNode * node, MPieceNode * entry = NULL);
/* remove an index node for memory pieces in use */
void RemoveAllocIndexNode(MPieceNode * node, MPieceNode * entry = NULL);
/* add an index node for available memory pieces ∂*/
void AddAllocIndexNode(MPieceNode * node, MPieceNode * entry = NULL);
/* release a piece of memory as "free" */
void ReleaseStandard(int myDevID, void * p);
void ReleaseStandard(int myDevID, void * p, MTYPE size);
/* rebuild index to merge small fragments of memory and free the block with no use */
void RebuildIndex();
......@@ -382,6 +402,9 @@ public:
extern XMem * GMem;
extern int testxmemid;
extern void * recordp;
} /* end of the nts (NiuTrans.Tensor) namespace */
#endif
source/tensor/XTensor.cpp
......@@ -254,7 +254,7 @@ void XTensor::DestroyData()
else if(data != NULL && isInGlobalMem)
FreeData(this, mem);
else if(data != NULL)
mem->Release(data);
mem->Release(data, GetDataSizeInChar());
data = NULL;
if(dataHost != NULL)
......@@ -1135,19 +1135,21 @@ resize a tensor with a specified tensor size
bool XTensor::Resize(const int myOrder, const int * myDimSize,
const TENSOR_DATA_TYPE myDataType, const float myDenseRatio)
{
order = myOrder;
unitNum = 1;
unitNumNonZero = 0;
isInit = true;
/* free old mem */
if(data != NULL){
if (mem == NULL)
XMemFree(devID, data);
else
mem->Release(data);
mem->Release(data, GetDataSizeInChar());
}
order = myOrder;
unitNum = 1;
unitNumNonZero = 0;
isInit = true;
bool filledData = true;
bool zeroData = false;
for(int i = 0; i < order; i++){
......
source/tensor/test/TXMem.cpp
......@@ -63,6 +63,8 @@ bool TestXMemCase1()
}
for (int i = 0; i < testNum * scalar; i++) {
testxmemid++;
//fprintf(stderr, "%d\n", testxmemid);
int j = rand() % caseNum;
if (p[j] == NULL) {
......@@ -71,7 +73,7 @@ bool TestXMemCase1()
p[j][k] = j;
}
else {
mem.ReleaseStandard(mem.devID, p[j]);
mem.ReleaseStandard(mem.devID, p[j], size[j] * sizeof(int));
for (int k = 0; k < size[j]; k++)
p[j][k] = -1;
p[j] = NULL;
......@@ -86,6 +88,61 @@ bool TestXMemCase1()
}
}
}
/*MPieceNode * entry = NULL;
MPieceNode * node = NULL;
entry = mem.memIndex + mem.indexEntryNum + mem.FindIndexEntry(112);
int cc = 0;
node = entry->next;
while(node != NULL){
fprintf(stderr, "%d ", cc++);
if(node->size == 0){
MPieceNode * next = node->next;
node = next;
}
else{
CheckNTErrors(node->pReal != NULL, "Illegal pointer!");
node = node->next;
}
}
fprintf(stderr, "\n");*/
/*int ccc = 0;
bool hhh = recordp != NULL ? false : true;
for(int i = 0; i < mem.indexEntryNum; i++){
MPieceNode * entry = mem.memIndex + mem.indexEntryNum + i;
MPieceNode * last = entry;
MPieceNode * node = entry->next;
ccc = 0;
while(node != NULL){
CheckNTErrors(node->pre == last, "XSomething is wrong!");
CheckNTErrors(last->next == node, "XSomething is wrong!");
last = node;
ccc++;
if(node->pReal == recordp){
hhh = true;
}
if(node->size == 0){
MPieceNode * next = node->next;
node = next;
}
else{
CheckNTErrors(node->pReal != NULL, "Illegal pointer!");
node = node->next;
}
}
}
if(!hhh){
int nnn = 0;
}*/
}
delete[] p;
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论