Weights.hpp

/*
 * Weights.hpp
 *
 *  Created on: Jul 21, 2017
 *      Author: Pete Schultz
 */

#ifndef WEIGHTS_HPP_
#define WEIGHTS_HPP_

#include "checkpointing/Checkpointer.hpp"
#include "components/PatchGeometry.hpp"
#include "include/PVLayerLoc.h"
#include "include/pv_types.h"
#include <memory>
#include <string>
#include <vector>

#ifdef PV_USE_CUDA
#include "arch/cuda/CudaDevice.hpp"
#endif // PV_USE_CUDA

namespace PV {

class Weights {

  public:
   Weights(std::string const &name);

   Weights(
         std::string const &name,
         int patchSizeX,
         int patchSizeY,
         int patchSizeF,
         PVLayerLoc const *preLoc,
         PVLayerLoc const *postLoc,
         int numArbors,
         bool sharedWeights,
         double timestamp);

   virtual ~Weights() {}

   void initialize(
         std::shared_ptr<PatchGeometry> geometry,
         int numArbors,
         bool sharedWeights,
         double timestamp);

   void initialize(Weights const *baseWeights);

   void initialize(
         int patchSizeX,
         int patchSizeY,
         int patchSizeF,
         PVLayerLoc const *preLoc,
         PVLayerLoc const *postLoc,
         int numArbors,
         bool sharedWeights,
         double timestamp);

   void allocateDataStructures();

   void checkpointWeightPvp(
         Checkpointer *checkpointer,
         char const *objectName,
         char const *bufferName,
         bool compressFlag);

   float calcMinWeight();

   float calcMinWeight(int arbor);

   float calcMaxWeight();

   float calcMaxWeight(int arbor);

   int calcDataIndexFromPatchIndex(int patchIndex) const;

#ifdef PV_USE_CUDA

   void copyToGPU();
#endif // PV_USE_CUDA

   bool getSharedFlag() const { return mSharedFlag; }

   std::string const &getName() const { return mName; }

   std::shared_ptr<PatchGeometry> getGeometry() const { return mGeometry; }

   int getNumArbors() const { return mNumArbors; }

   int getNumDataPatchesX() const { return mNumDataPatchesX; }

   int getNumDataPatchesY() const { return mNumDataPatchesY; }

   int getNumDataPatchesF() const { return mNumDataPatchesF; }

   int getNumDataPatches() const {
      return getNumDataPatchesX() * getNumDataPatchesY() * getNumDataPatchesF();
   }

   Patch const &getPatch(int patchIndex) const;

   float *getData(int arbor);

   float const *getDataReadOnly(int arbor) const;

   float *getDataFromDataIndex(int arbor, int dataIndex);

   float *getDataFromPatchIndex(int arbor, int patchIndex);

   float *getData(int arbor, double timestamp);

   float *getDataFromDataIndex(int arbor, int dataIndex, double timestamp);

   float *getDataFromPatchIndex(int arbor, int patchIndex, double timestamp);

   void setTimestamp(double timestamp) { mTimestamp = timestamp; }

   double getTimestamp() const { return mTimestamp; }

   int getPatchSizeX() const { return mGeometry->getPatchSizeX(); }

   int getPatchSizeY() const { return mGeometry->getPatchSizeY(); }

   int getPatchSizeF() const { return mGeometry->getPatchSizeF(); }

   int getPatchSizeOverall() const { return mGeometry->getPatchSizeOverall(); }

   int getPatchStrideF() const { return 1; }

   int getPatchStrideX() const { return mGeometry->getPatchStrideX(); }

   int getPatchStrideY() const { return mGeometry->getPatchStrideY(); }

   bool getWeightsArePlastic() const { return mWeightsArePlastic; }

   void setWeightsArePlastic() { mWeightsArePlastic = true; }

   void setMargins(PVHalo const &preHalo, PVHalo const &postHalo);

#ifdef PV_USE_CUDA
   bool isUsingGPU() { return mUsingGPUFlag; }
   void useGPU() { mUsingGPUFlag = true; }

   void setCudaDevice(PVCuda::CudaDevice *device) { mCudaDevice = device; }

   PVCuda::CudaBuffer *getDevicePatchToDataLookup() const { return mDevicePatchToDataLookup; }
   PVCuda::CudaBuffer *getDeviceData() const { return mDeviceData; }
#ifdef PV_USE_CUDNN
   PVCuda::CudaBuffer *getCUDNNData() const { return mCUDNNData; }
#endif // PV_USE_CUDNN
#endif // PV_USE_CUDA

  protected:
   Weights() {}

   void setName(std::string const &name) { mName = name; }

   void setNumDataPatches(int numDataPatchesX, int numDataPatchesY, int numDataPatchesF);

#ifdef PV_USE_CUDA
   void allocateCudaBuffers();
#endif // PV_USE_CUDA

  private:
   virtual void initNumDataPatches();

  private:
   std::string mName;
   std::shared_ptr<PatchGeometry> mGeometry = nullptr;
   int mNumArbors;
   bool mSharedFlag;
   double mTimestamp;
   int mNumDataPatchesX;
   int mNumDataPatchesY;
   int mNumDataPatchesF;

   std::vector<std::vector<float>> mData;
   std::vector<int> dataIndexLookupTable;

   bool mWeightsArePlastic = false;

#ifdef PV_USE_CUDA
   bool mUsingGPUFlag                           = false;
   PVCuda::CudaDevice *mCudaDevice              = nullptr;
   PVCuda::CudaBuffer *mDevicePatchToDataLookup = nullptr;
   PVCuda::CudaBuffer *mDeviceData              = nullptr;
#ifdef PV_USE_CUDNN
   PVCuda::CudaBuffer *mCUDNNData = nullptr;
#endif // PV_USE_CUDNN
   double mTimestampGPU;
#endif // PV_USE_CUDA
}; // end class Weights

} // end namespace PV

#endif // WEIGHTS_HPP_