cavis/libnd4j/include/memory/MemoryCounter.h

147 lines
4.8 KiB
C++

/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* 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.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author raver119@gmail.com
//
#ifndef SD_MEMORYCOUNTER_H
#define SD_MEMORYCOUNTER_H
#include <system/pointercast.h>
#include <system/dll.h>
#include <map>
#include <memory/MemoryType.h>
#include <mutex>
namespace sd {
namespace memory {
/**
* This class provides simple per-device counter
*/
class ND4J_EXPORT MemoryCounter {
private:
// used for synchronization
std::mutex _locker;
// per-device counters
std::map<int, Nd4jLong> _deviceCounters;
// TODO: change this wrt heterogenous stuff on next iteration
// per-group counters
std::map<sd::memory::MemoryType, Nd4jLong> _groupCounters;
// per-device limits
std::map<int, Nd4jLong> _deviceLimits;
// per-group limits
std::map<sd::memory::MemoryType, Nd4jLong> _groupLimits;
MemoryCounter();
~MemoryCounter() = default;
public:
static MemoryCounter & getInstance();
/**
* This method checks if allocation of numBytes won't break through per-group or per-device limit
* @param numBytes
* @return TRUE if allocated ammount will keep us below limit, FALSE otherwise
*/
bool validate(Nd4jLong numBytes);
/**
* This method checks if allocation of numBytes won't break through per-device limit
* @param deviceId
* @param numBytes
* @return TRUE if allocated ammount will keep us below limit, FALSE otherwise
*/
bool validateDevice(int deviceId, Nd4jLong numBytes);
/**
* This method checks if allocation of numBytes won't break through per-group limit
* @param deviceId
* @param numBytes
* @return TRUE if allocated ammount will keep us below limit, FALSE otherwise
*/
bool validateGroup(sd::memory::MemoryType group, Nd4jLong numBytes);
/**
* This method adds specified number of bytes to specified counter
* @param deviceId
* @param numBytes
*/
void countIn(int deviceId, Nd4jLong numBytes);
void countIn(sd::memory::MemoryType group, Nd4jLong numBytes);
/**
* This method subtracts specified number of bytes from specified counter
* @param deviceId
* @param numBytes
*/
void countOut(int deviceId, Nd4jLong numBytes);
void countOut(sd::memory::MemoryType group, Nd4jLong numBytes);
/**
* This method returns amount of memory allocated on specified device
* @param deviceId
* @return
*/
Nd4jLong allocatedDevice(int deviceId);
/**
* This method returns amount of memory allocated in specified group of devices
* @param group
* @return
*/
Nd4jLong allocatedGroup(sd::memory::MemoryType group);
/**
* This method allows to set per-device memory limits
* @param deviceId
* @param numBytes
*/
void setDeviceLimit(int deviceId, Nd4jLong numBytes);
/**
* This method returns current device limit in bytes
* @param deviceId
* @return
*/
Nd4jLong deviceLimit(int deviceId);
/**
* This method allows to set per-group memory limits
* @param group
* @param numBytes
*/
void setGroupLimit(sd::memory::MemoryType group, Nd4jLong numBytes);
/**
* This method returns current group limit in bytes
* @param group
* @return
*/
Nd4jLong groupLimit(sd::memory::MemoryType group);
};
}
}
#endif //SD_MEMORYCOUNTER_H