ManagedMemory.h

/*-------------------------------------------------------------------
Copyright 2011 Ravishankar Sundararaman

This file is part of JDFTx.

JDFTx is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

JDFTx is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with JDFTx.  If not, see <http://www.gnu.org/licenses/>.
-------------------------------------------------------------------*/

#ifndef JDFTX_ELECTRONIC_MANAGEDMEMORY_H
#define JDFTX_ELECTRONIC_MANAGEDMEMORY_H

#include <electronic/common.h>

class ManagedMemory
{
public:

        complex* data();

        const complex* data() const;

        size_t nData() const { return nElements; } 

        bool isOnGpu() const { return onGpu; } 

        #ifdef GPU_ENABLED
        complex* dataGpu(); 
        const complex* dataGpu() const; 
        #endif

        //Utilities to automatically select "preferred" data i.e. GPU when it is enabled, CPU otherwise
        //This eases the overload of CPU/GPU functions based on complex vs complex
        #ifdef GPU_ENABLED
        inline complex* dataPref() { return dataGpu(); }
        inline const complex* dataPref() const { return dataGpu(); }
        #else
        inline complex* dataPref() { return data(); }
        inline const complex* dataPref() const { return data(); }
        #endif

        //Inter-process communication:
        void send(int dest, int tag=0) const; 
        void recv(int src, int tag=0); 
        void bcast(int root=0); 
        void allReduce(MPIUtil::ReduceOp op, bool safeMode=false, bool ignoreComplexCheck=false); 

        void write(const char *fname) const; 
        void writea(const char *fname) const; 
        void write(FILE *filep) const; 
        void read(const char *fname); 
        void read(FILE *filep); 
        void read_real(const char *fname); 
        void read_real(FILE *filep); 
        void write_real(const char *fname) const; 
        void write_real(FILE *filep) const; 
        void dump(const char* fname, bool realPartOnly) const; 
        void zero(); 
        
        static void reportUsage(); 
protected:
        void memFree(); 
        void memInit(string category, size_t nElem, bool onGpu=false); 

        void memMove(ManagedMemory&&); 
        ManagedMemory(); 
        ~ManagedMemory();

private:
        string category; 
        size_t nElements; 
        complex* c; 
        bool onGpu; 
        #ifdef GPU_ENABLED
        void toCpu(); 
        void toGpu(); 
        #endif
};

//Some common elementwise / vector-like operations
void memcpy(ManagedMemory&, const ManagedMemory&); 
void scale(double alpha, ManagedMemory& y); 
void scale(complex alpha, ManagedMemory& y); 
void axpy(complex alpha, const ManagedMemory& x, ManagedMemory& y); 
double nrm2(const ManagedMemory&); 
complex dotc(const ManagedMemory& a, const ManagedMemory& b); 

#endif // JDFTX_ELECTRONIC_MANAGEDMEMORY_H