147 lines
4.8 KiB
C++
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
|