text: //
  // Copyright (c) Microsoft. All rights reserved.
  // Licensed under the MIT license. See LICENSE.md file in the project root for full license information.
  //
  
  #pragma once
  
  #include <string>
  #include <stdexcept>
  #include <vector>
  #include <set>
  #include <unordered_map>
  #include <unordered_set>
  #include <utility>
  #include <algorithm>
  #include <stdlib.h>
  
  #include "Basics.h"
  #include "Matrix.h"
  #include "ComputationNode.h"
  
  namespace Microsoft { namespace MSR { namespace CNTK {
  
  template <class ElemType>
  struct MemRequestInfo
  {
      DEVICEID_TYPE deviceId;                     // which device to allocate data 
      std::vector<shared_ptr<Matrix<ElemType>>*> pMatrixPtrs;    // memory pointers 
      size_t matrixSize;                          // memory size 
      bool mbScale;                               // whether the memory shall be scaled by minibatch size 
      bool isWorkSpace;                           // workspace memory or not, by workspace we indicate whether a memory space will be released very shortly after allocation 
      int allocStep;                              // at what step counter memory allocation is requested 
      int releaseStep;                            // at what step counter memory release is requested  
      int memoryId;                               // integer indexing the memory buffer ID 
      MemRequestInfo(DEVICEID_TYPE deviceId, shared_ptr<Matrix<ElemType>>*pMatrixPtr, size_t matrixSize, bool mbScale, bool isWorkSpace, int allocStep)
          :deviceId(deviceId), matrixSize(matrixSize), mbScale(mbScale), isWorkSpace(isWorkSpace), allocStep(allocStep), releaseStep(INT_MAX), memoryId(-1)
      {
          pMatrixPtrs.push_back(pMatrixPtr);
      }
      void SetReleaseStep(int step) { releaseStep = step; }
      void SetMemoryId(int id) { memoryId = id;  }
  };
  
  template <class ElemType>
  struct greater_than_mem_req_size
  {
      inline bool operator() (const MemRequestInfo<ElemType>& info1, const MemRequestInfo<ElemType>& info2)
      {
          return (info1.matrixSize > info2.matrixSize);
      }
  };
  
  struct MemAllocInfo
  {
      int memoryId; 
      size_t memorySize; 
      vector<pair<int, int>> occupancy; 
      MemAllocInfo(int memoryId, size_t memorySize, vector<pair<int, int>> occ)
          :memoryId(memoryId), memorySize(memorySize), occupancy(occ)
      {
      }
  };
  
  // MatrixPool -- class to support memory sharing
  // Despite the gather general name of this class, it is specifically designed to support the memory sharing of ComputationNodes.
  // Note: see #define SUPRESS_MEMSHARING below as for how to temporarily disable memory sharing altogether, for debugging
  class MatrixPool
  {
  public:
      typedef const void* AliasNodePtr; // use as an identifier in place of ComputationNodeBasePtr to avoid include order issue
  
  protected:
      vector<MemRequestInfo<float>> m_memRequestInfoFloatVec; 
      vector<MemRequestInfo<double>> m_memRequestInfoDoubleVec;
      vector<MemRequestInfo<half>> m_memRequestInfoHalfVec;
      set<DEVICEID_TYPE> m_deviceIDSet; 
      int m_stepCounter; 
  
      template <class ElemType>
      vector<MemRequestInfo<ElemType>>& GetMemRequestInfoVec();
  
      // MatrixPool allows a bunch of node to share one matrix
  
      struct AliasInfo
      {
          void* pMatrixPtr;
          size_t totalCount;
          size_t releaseCount;
  
          AliasInfo(size_t total = 0)
              : pMatrixPtr(nullptr), totalCount(total), releaseCount(0)
          {
          }
      };
      unordered_map<AliasNodePtr, AliasInfo> m_aliasGroups;
      unordered_map<AliasNodePtr, AliasNodePtr> m_aliasLookup;
  
  public:
  
      void Reset()
      {
          m_stepCounter = 0;
          m_aliasGroups.clear();
          m_aliasLookup.clear();
      };
  
      template <class ElemType>
      MemRequestInfo<ElemType>* GetMemInfo(shared_ptr<Matrix<ElemType>> *pMatrixPtr)
      {
          vector<MemRequestInfo<ElemType>>& memInfoVec = GetMemRequestInfoVec<ElemType>();
          // iterate through the vector and find the pointer memInfo
          for (auto& memInfo : memInfoVec)
          {
              if (memInfo.pMatrixPtrs[0] == pMatrixPtr)
                  return &memInfo;
          }
          return nullptr;
      }
  
      template <class ElemType>
      void RequestRelease(shared_ptr<Matrix<ElemType>> *pMatrixPtr)
      {
          auto memInfo = GetMemInfo(pMatrixPtr);
          if (memInfo != nullptr)
          {
              memInfo->SetReleaseStep(m_stepCounter);
          }
          m_stepCounter++; 
      }
  
      // isWorkSpace is a flag indicating a memory is temporary and will be released very shortly. In the current implementation, all workspace
      // memories will have their own pool. This is a design proven to be useful for the workspace memory in convolution. 
      // matrixSize is an estimate of the required memory to be allocated. Note we don't allocate any memory at the time of request. Instead, a 
      // global memory allocation optimziation is run to improve memory efficiency 
      // mbScale is another flag indicating if the size of the memory will scale w.r.t. the minibatch size. Unfortunately, at the time of memory
      // request and pointer assignment, we don't known the minibatch size. Thus our memory sharing algorithm is sub-optimal. 
      template <class ElemType>
      void RequestAllocate(DEVICEID_TYPE deviceId, shared_ptr<Matrix<ElemType>>*pMatrixPtr, size_t matrixSize, bool mbScale, bool isWorkSpace)
      {
          vector<MemRequestInfo<ElemType>>& memInfoVec = GetMemRequestInfoVec<ElemType>(); 
          MemRequestInfo<ElemType> memInfo(deviceId, pMatrixPtr, matrixSize, mbScale, isWorkSpace, m_stepCounter);
          memInfoVec.push_back(memInfo); 
          m_deviceIDSet.insert(deviceId); 
          m_stepCounter++; 
  
          // assign some temporary pointer, they will be replaced later unless the matrix is sparse
          *pMatrixPtr = make_shared<Matrix<ElemType>>(deviceId);
      }
  
      void OptimizedMemoryAllocation()
      {
          // MatrixPool is not templated, so we call both float and double versions here 
          OptimizedMemoryAllocationFunc<float>(); 
          OptimizedMemoryAllocationFunc<double>();
          OptimizedMemoryAllocationFunc<half>();
          return; 
      }
  
      void SetAliasInfo(
          const unordered_map<AliasNodePtr, unordered_set<AliasNodePtr>>& groupMap,
          const unordered_map<AliasNodePtr, AliasNodePtr>& rootLookupMap)
      {
          m_aliasLookup.clear();
          for (const auto& pair : groupMap)
          {
              m_aliasGroups.insert(std::make_pair(pair.first, AliasInfo(pair.second.size())));
  
              for (const auto& child : pair.second)
              {
                  if (rootLookupMap.find(child) == rootLookupMap.end())
                      InvalidArgument("group nodes should be in lookupMap");
              }
          }
  
          for (const auto& pair : rootLookupMap)
          {
              if (groupMap.find(pair.second) == groupMap.end())
                  InvalidArgument("lookup root should be group key");
          }
          m_aliasLookup = rootLookupMap;
      }
  
      template <class ElemType>
      void RequestAliasedRelease(AliasNodePtr node)
      {
          const auto iter = m_aliasLookup.find(node);
          if (iter == m_aliasLookup.end())
              LogicError("node not aliased");
  
          auto parent = iter->second;
          auto& aliasInfo = m_aliasGroups[parent];
          if (aliasInfo.pMatrixPtr == nullptr)
              LogicError("double releasing aliased matrix, or releasing before any allocation for the matrix");
  
          if (aliasInfo.releaseCount >= aliasInfo.totalCount)
              LogicError("number of alias instances exceeded expectation");
  
          aliasInfo.releaseCount++;
  
          if (aliasInfo.releaseCount == aliasInfo.totalCount)
          {
              RequestRelease((shared_ptr<Matrix<ElemType>>*)aliasInfo.pMatrixPtr);
              aliasInfo.pMatrixPtr = nullptr;
          }
      }
  
      template <class ElemType>
      void RequestAliasedAllocate(DEVICEID_TYPE deviceId, AliasNodePtr node, shared_ptr<Matrix<ElemType>>*pMatrixPtr, size_t matrixSize, bool mbScale)
      {
          const auto iter = m_aliasLookup.find(node);
          if (iter == m_aliasLookup.end())
              LogicError("node not aliased");
  
          auto parent = iter->second;
          auto& aliasInfo = m_aliasGroups[parent];
          if (aliasInfo.pMatrixPtr == nullptr)
          {
              // first allocation for the group
              aliasInfo.pMatrixPtr = pMatrixPtr;
              RequestAllocate(deviceId, pMatrixPtr, matrixSize, mbScale, false);
          }
          else
          {
              auto aliasRootMatrixPtr = (shared_ptr<Matrix<ElemType>>*)aliasInfo.pMatrixPtr;
              *pMatrixPtr = *aliasRootMatrixPtr;
              GetMemInfo<ElemType>(aliasRootMatrixPtr)->pMatrixPtrs.push_back(pMatrixPtr);
          }
      }
  
  private: 
      bool CheckOverlap(pair<int, int>occ, vector<pair<int, int>>&occVec)
      {
          bool bRet = false;
          for (auto& o : occVec)
          {
              if (occ.first <= o.second && occ.second >= o.first)
              {
                  bRet = true;
                  break;
              }
          }
  //#define SUPRESS_MEMSHARING // #define this to disable memory sharing by always return true 
  // TODO: Make this a runtime option.
  #ifdef SUPRESS_MEMSHARING
          bRet = true; 
  #endif
          return bRet;
      }
  
      template <class ElemType>
      void OptimizedMemoryAllocationFunc()
      {
          vector<MemRequestInfo<ElemType>>& memInfoVec = GetMemRequestInfoVec<ElemType>();
          if (memInfoVec.empty())
              return; 
  
          // remove all requests that has been marked as sparse matrices, those will not participate in memory sharing 
          for (auto iter = memInfoVec.begin(); iter != memInfoVec.end(); )
          {
              bool hasSparse = false;
              for (auto matPtr : iter->pMatrixPtrs)
              {
                  if ((*matPtr)->GetMatrixType() == SPARSE)
                  {
                      hasSparse = true;
                      break;
                  }
              }
  
              if (hasSparse)
                  memInfoVec.erase(iter);
              else
                  iter++; 
          }
  
          // sort the memory request from largest size to smallest 
          std::sort(memInfoVec.begin(), memInfoVec.end(), greater_than_mem_req_size<ElemType>());
  
          std::vector<bool> workspaceFlagVec = {true, false};
          for (auto& devId : m_deviceIDSet)
          {
              for (auto wsFlag : workspaceFlagVec)   // we allocate the workspace memory pointers first, and they are not shared with the non-workspace memory requests
              {
                  // memAllocInfoVec is a sorted list of memory allocations from smallest to largest in memory size 
                  vector<MemAllocInfo> memAllocInfoVec;
                  int memoryCounter = 0;
                  // we start with memory request that is scalable with minibatch size(usually those require larger memory size)
                  for (auto& memInfo : memInfoVec)
                  {
                      // check if it's the proper device
                      if (memInfo.deviceId != devId || memInfo.isWorkSpace != wsFlag || !memInfo.mbScale)
                          continue;
  
                      if (!memAllocInfoVec.empty())
                      {
                          // since we assign from highest memory to lowest, every memory that has been allocated can accommodate the 
                          // current memory request, unless there is a conflict (overlap) 
                          auto iter = memAllocInfoVec.begin();
                          while (iter != memAllocInfoVec.end() && CheckOverlap(make_pair(memInfo.allocStep, memInfo.releaseStep), iter->occupancy))
                              iter++;
                          if (iter == memAllocInfoVec.end())
                          {
                              // no current memory can be assigned, need to create a new one 
                              vector<pair<int, int>> occ;
                              occ.push_back(make_pair(memInfo.allocStep, memInfo.releaseStep));
                              MemAllocInfo ma(memoryCounter, memInfo.matrixSize, occ);
                              // insert in the front of the vector to maintain sorted order 
                              memAllocInfoVec.insert(memAllocInfoVec.begin(), ma);
                              memInfo.SetMemoryId(memoryCounter);
                              memoryCounter++;
                          }
                          else
                          {
                              iter->occupancy.push_back(make_pair(memInfo.allocStep, memInfo.releaseStep));
                              memInfo.SetMemoryId(iter->memoryId);
                          }
                      }
                      else
                      {
                          vector<pair<int, int>> occ;
                          occ.push_back(make_pair(memInfo.allocStep, memInfo.releaseStep));
                          MemAllocInfo ma(memoryCounter, memInfo.matrixSize, occ);
                          memAllocInfoVec.push_back(ma);
                          memInfo.SetMemoryId(memoryCounter);
                          memoryCounter++;
                      }
                  }
  
                  // rescan the request list and this time allocate for those that doesn't depend on minibatch size 
                  for (auto& memInfo : memInfoVec)
                  {
                      // check if it's the proper device
                      if (memInfo.deviceId != devId || memInfo.isWorkSpace != wsFlag || memInfo.mbScale)
                          continue;
  
                      if (!memAllocInfoVec.empty())
                      {
                          // the memory allocation vector is sorted by size. We find the largest available buffer that doesn't have time overlap
                          auto workingAlloc = memAllocInfoVec.end();
                          for (auto iter = memAllocInfoVec.begin(); iter != memAllocInfoVec.end(); iter++)
                          {
                              if (!CheckOverlap(make_pair(memInfo.allocStep, memInfo.releaseStep), iter->occupancy))
                                  workingAlloc = iter;
                          }
                          if (workingAlloc == memAllocInfoVec.end())  // nothing works 
                          {
                              vector<pair<int, int>> occ;
                              occ.push_back(make_pair(memInfo.allocStep, memInfo.releaseStep));
                              MemAllocInfo ma(memoryCounter, memInfo.matrixSize, occ);
                              memAllocInfoVec.push_back(ma);  // add as the last one 
                              memInfo.SetMemoryId(memoryCounter);
                              memoryCounter++;
                          }
                          else
                          {
                              workingAlloc->occupancy.push_back(make_pair(memInfo.allocStep, memInfo.releaseStep));
                              memInfo.SetMemoryId(workingAlloc->memoryId);
                          }
                      }
                      else
                      {
                          vector<pair<int, int>> occ;
                          occ.push_back(make_pair(memInfo.allocStep, memInfo.releaseStep));
                          MemAllocInfo ma(memoryCounter, memInfo.matrixSize, occ);
                          memAllocInfoVec.push_back(ma);
                          memInfo.SetMemoryId(memoryCounter);
                          memoryCounter++;
                      }
                  }
  
                  // now assign the actual pointers 
                  for (int i = 0; i < memoryCounter; i++)
                  {
                      auto matrixPtr = make_shared<Matrix<ElemType>>(devId);
                      if (!matrixPtr) // this can't really happen, because we haven't started allocating memory yet
                          LogicError("MatrixPool: failed to get a valid matrix.");
                      for (auto& memInfo : memInfoVec)
                      {
                          if (memInfo.deviceId == devId && memInfo.isWorkSpace == wsFlag && memInfo.memoryId == i)
                          {
                              for (auto pOutMatrixPtr : memInfo.pMatrixPtrs)
                              {
                                  *pOutMatrixPtr = matrixPtr;
                              }
                          }
                      }
                  }
              }
          }
      }
  };
  
  }}}
  