Brutex
/
cavis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cublasHandle sharing + lock 

Signed-off-by: raver119 <raver119@gmail.com>
master
raver119 2019-09-02 16:52:28 +03:00
parent 18828f9725
commit 2129d5bcac
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
nd4j/nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/jita/handler/impl/CudaZeroHandler.java
View File

@ -1184,9 +1184,10 @@ public class CudaZeroHandler implements MemoryHandler {
protected synchronized cublasHandle_t getCudaCublasHandle(OpaqueLaunchContext lc) { protected synchronized cublasHandle_t getCudaCublasHandle(OpaqueLaunchContext lc) {
val deviceId = Nd4j.getAffinityManager().getDeviceForCurrentThread(); val deviceId = Nd4j.getAffinityManager().getDeviceForCurrentThread();
if (cublasHandles.get(deviceId) == null) if (cublasHandles.get(deviceId) == null) {
cublasHandles.remove(deviceId); cublasHandles.remove(deviceId);
cublasHandles.add(deviceId, new cublasHandle_t(nativeOps.lcBlasHandle(lc))); cublasHandles.add(deviceId, new cublasHandle_t(nativeOps.lcBlasHandle(lc)));
}
return cublasHandles.get(deviceId); return cublasHandles.get(deviceId);
} }