XMem.h

/* NiuTrans.Tensor - an open-source tensor library
 * Copyright (C) 2017, Natural Language Processing Lab, Northestern University. 
 * All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * 
 * $Created by: XIAO Tong (xiaotong@mail.neu.edu.cn) 2016-5-25
 *
 */

#ifndef __XMEM_H__
#define __XMEM_H__

#include <stdlib.h>

#ifdef CUDA_BLAS
#define USE_CUDA
#endif

#ifdef USE_CUDA
// the CUDA stuff
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cuda.h>
#include <curand.h>
#endif

/* the nts (NiuTrans.Tensor) namespace */
namespace nts{

typedef unsigned long long MTYPE;
typedef long long int      MTYPEINT;
typedef long long          INT_64;

//#define CUDA_PITCH 256
#define CUDA_PITCH 1
#define CUDA_HOST_MALLOC 1
#define MY_PITCH CUDA_PITCH
#define BUF_PITCH 256
#define MIN_BLOCK_SIZE_FOR_MEMPOOL 128 * 1024 * 1024
#define MIN_BLOCK_NUM_FOR_MEMPOOL 1024

/* memory block */
struct XMemBlock
{
    /* pointer to where to start */
    void * mem;

    /* size of the block */
    MTYPE size;

    /* size of the used memory in this block */
    MTYPE used;

    /* disired size of the block */
    MTYPE sizeDesired;
};

/* 
mode of runnig a memory pool 
- UNI_FREE: free all memory space when the memory allocation is no use
- FREE_ON_THE_FLY: run in normal "malloc" and "free" ways
*/
enum MEMPOOL_MODE {UNI_FREE, FREE_ON_THE_FLY};

/* header of a memory piece (FREE_ON_THE_FLY) */
struct MHeader
{
    /* state of the memory piece 
       1: free
       2: in use
    */
    int state;

    /* size of the allocated memory */
    MTYPE size;

    /* pointer to the header of the previous memory piece */
    MHeader * pre;

    /* pointer to the header of the next memory piece */
    MHeader * next;

    /* id of the memory block */
    int blockID;
};

/* index of memory piece */
struct MPieceNode
{
    /* size of the memory piece */
    MTYPE size;

    /* previous node */
    MPieceNode * pre;

    /* next node */
    MPieceNode * next;

    /* pointer to the head of a memory piece */
    void * p;
};

/* 
memory pool.
Basically a memory pool runs in two ways. It can be an naive implementation (uni-free mode)
that allocates the memory in a continuous manner and free the used memory space
in the end - the user does not need to call the "memfree" function
when a small piece of memory is not used any more. Instead we call the "go back"
function to the initial state when all memory space in the memory pool is not in use.
Another way (free on-the-fly mode) is to allocate and free the memory space as in standard 
"malloc" and "free" manners. Here we do it on a pre-allocated memory block. This mode is 
more flexible but relatively slower than the uni-free mode.
*/
class XMem
{
public:
    /* 
    device id 
    <0:  CPU memory
    >=0: GPU device ID
    */
    int devID;

    /* mode of running the memory pool */
    MEMPOOL_MODE mode;

    /* indicates whether the memory allocation is static */
    bool isStatic;

    /* memory blocks */
    XMemBlock * blocks;

    /* number of memory blocks */
    int blockNum;

    /* max size of a block */
    MTYPE maxBlockSize;

    /* total size of all memory blocks */
    MTYPE totalBlockSize;

    /* total size of used memory */
    MTYPE totalUsed;

    /* current mem block that is using */
    XMemBlock * curBlock;

    /* id of the current mem block */
    int curBlockID;

    /* id of the final mem block that is used */
    int finalBlockID;

    /* pointer to the buffer used to store temp data */
    void * buf;

    /* size of the buffer */
    MTYPE bufSize;

    /* size of the used buffer in this block */
    MTYPE bufUsed;

    /* name of the memory pool */
    char * name;

    /* pin for the memory. It is used for recording the starting address
       before we apply for a large amount of the memory. Then we can go
       back to this point for memory release. */
    int curBlockPin;
    MTYPE curUsedPin;
    MTYPE bufUsedPin;

#ifdef USE_CUDA
    /* handle used for cublas */
    cublasHandle_t cublasHandle;

    /* random number generator for cuda code */
    curandGenerator_t randGen;
#endif

public:
    /* index of the free memory pieces */
    MPieceNode * freeMemIndex;

    /* maximum number of index nodes */
    INT_64 indexNodeNum;

    /* count of the used nodes */
    INT_64 indexNodeNumUsed;

    /* minimal size allocation for each index entry */
    MTYPE * minSizeIndex;

    /* number of the index entries */
    int indexEntryNum;

    /* index offset */
    int indexOffset;

public:

    /* constructor */
    XMem();

    /* constructor */
    XMem(int myDevID,
         MEMPOOL_MODE myMode = UNI_FREE,
         MTYPE myBlockSize = MIN_BLOCK_SIZE_FOR_MEMPOOL, 
         int myBlockNum = MIN_BLOCK_NUM_FOR_MEMPOOL, 
         MTYPE myBufSize = 0);

    /* deconstructor */
    ~XMem();

    /* initialize it */
    void Initialize(int myDevID, MEMPOOL_MODE myMode, MTYPE myBlockSize, int myBlockNum, MTYPE myBufSize);

    /* free memory */
    void Free();

    /* free a piece of memory */
    void Free(int myDevID, void * mem);

    /* use string as the name of the memory pool */
    void SetName(const char * myName);

    /* switch to the device we want to work */
    void SetDevice(int myDevID);

    /* switch to the device (with fast cuda execution mode) we want to work */
    void SetDeviceFast(int myDevID);

    /* run in static mode */
    void SetStaticMode(bool myIsStatic);

    /* specify if the memory pool is used for tensor computation (rather
       than storage */
    void SetComputationMode(bool myIsForComputation);

    /* initialize the index */
    void SetIndex(INT_64 size, MTYPE minSizeFirst = 256, int minSizeNum = 20);

    /* get device id */
    int GetDevID();

    /* set desired memory block size */
    void SetDesiredSize(int myDevID, int blockID, MTYPE mySize);

    /* require a piece of memory */
    void * Alloc(MTYPE mySize);

    /* require a piece of memory */
    void * Alloc(int myDevID, MTYPE mySize);

    /* require a piece of memory in a dynamic manner */
    void * AllocDynamic(int myDevID, MTYPE mySize);

    /* require a piece of memory with fixed size (if possible) */
    void * AllocStatic(int myDevID, MTYPE mySize);

    /* require a piece of memory that is not in the memory pool */
    void * AllocGlobal(int myDevID, MTYPE mySize);

    /* get the available size of the memory that can be used */
    MTYPE GetAvailableSize(int myDevID);

    /* require a piece of memory in the buffer */
    void * AllocBuf(int myDevID, MTYPE mySize, int pitch = BUF_PITCH);

    /* release a piece of memory */
    void Release(void * p);

    /* release a piece of memory */
    void Release(int myDevID, void * p);

    /* release a piece of memory in the buffer */
    void ReleaseBuf(int myDevID, MTYPE mySize, int pitch = BUF_PITCH);

    /* release a piece of memory that is not in the memory pool */
    void ReleaseGlobal(int myDevID, void * p);

    /* allocate a piece of memory as "malloc" */
    void * AllocStandard(int myDevID, MTYPE mySize, bool myIsRebuiltIndex = false);

    /* find the highest set bit (or most significant set bit) in an integer-64 */
    int GetMSB(MTYPE mySize);

    /* find the index entry for allocation query */
    int FindIndexEntry(MTYPE mySize);

    /* remove an index node */
    void RemoveIndexNode(MPieceNode * node, MPieceNode * entry = NULL);

    /* add an index node */
    void AddIndexNode(MPieceNode * node, MPieceNode * entry = NULL);

    /* release a piece of memory as "free" */
    void ReleaseStandard(int myDevID, void * p);

    /* rebuild index to merge small fragments of memory and free the block with no use */
    void RebuildIndex();

    /* reset buffer */
    void Reset(int myDevID);

    /* get pitch for aligned memory */
    MTYPE GetPitch(int myDevID, MTYPE baseAddress, MTYPE mySize);

    /* get pitched address for aligned memory */
    void * GetPitchedAddress(void * address, MTYPE pitch);

    /* get current address (for use) */
    void * GetAddress();

    /* clear it */
    void Clear();

    /* clear the buffer */
    void ClearBuf();

    /* clear the memory pool and the buffer */
    void ClearAll();

    /* set a variable to the input value */
    static
    void Copy(void * tgt, void * src, int size, XMem * tgtMem = NULL, XMem * srcMem = NULL);

    /* set a float-typed variable to the input value */
    static
    void CopyFloat(float * tgt, float * src, XMem * tgtMem = NULL, XMem * srcMem = NULL);

    /* set a variable to 0 */
    static
    void SetZero(void * tgt, MTYPE size, XMem * tgtMem = NULL);

    /* record the pin point */
    void SetPin();

    /* go back to the pin point */
    void BackToPin();

    /* record the pin point for buffer */
    void SetPinBuf();

    /* go back to the pin point for buffer */
    void BackToPinBuf();

    /* transform a size into a number (in million) */
    static
    MTYPE GetMemSize(const char * size);

    /* transform a size into a number (in Bytes) */
    static
    MTYPE GetMemSizeInBytes(const char * size);

    /* create a new cublas handle */
    void CreateBLASHandle();

#ifdef USE_CUDA
    /* get the handle of cublas */
    cublasHandle_t * GetCublasHandle();
#endif

};

extern XMem * GMem;

} /* end of the nts (NiuTrans.Tensor) namespace */

#endif