org.nd4j.linalg.util.DeviceLocal Maven / Gradle / Ivy
package org.nd4j.linalg.util;
import lombok.NonNull;
import org.nd4j.linalg.api.buffer.DataBuffer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
/**
* Class similar to Java ThreadLocal class, but locality is preserved with respect to device used
*
* @author [email protected]
*/
public class DeviceLocal {
private Map backingMap = new ConcurrentHashMap<>();
private List locksMap = new ArrayList<>();
public DeviceLocal() {
int numDevices = Nd4j.getAffinityManager().getNumberOfDevices();
for (int i = 0; i < numDevices; i++) {
locksMap.add(new ReentrantReadWriteLock());
}
}
/**
* This method returns object local to current deviceId
*
* @return
*/
@Nullable
public T get() {
return get(Nd4j.getAffinityManager().getDeviceForCurrentThread());
}
/**
* This method returns object local to target device
*
* @param deviceId
* @return
*/
@Nullable
public T get(int deviceId) {
try {
locksMap.get(deviceId).readLock().lock();
return backingMap.get(deviceId);
} finally {
locksMap.get(deviceId).readLock().unlock();
}
}
/**
* This method sets object for specific device
*
* @param deviceId
* @param object
*/
public void set(int deviceId, T object) {
try {
locksMap.get(deviceId).writeLock().lock();
backingMap.put(deviceId, object);
} finally {
locksMap.get(deviceId).writeLock().unlock();
}
}
/**
* This method sets object for current device
*
* @param object
*/
public void set(T object) {
set(Nd4j.getAffinityManager().getDeviceForCurrentThread(), object);
}
/**
* This method removes object stored for current device
*
*/
public void clear() {
int deviceId = Nd4j.getAffinityManager().getDeviceForCurrentThread();
try {
locksMap.get(deviceId).writeLock().lock();
backingMap.remove(deviceId);
} finally {
locksMap.get(deviceId).writeLock().unlock();
}
}
}