GitHub

Original file line numberDiff line numberDiff line change

@@ -332,10 +332,10 @@ DeviceManager::DeviceManager()

332332
333333

sortDevices();

334334
335-

for(int i = 0; i < nDevices; i++) {

336-

setActiveDevice(i, cuDevices[i].nativeId);

337-

CUDA_CHECK(cudaStreamCreate(&streams[i]));

338-

}

335+

// Initialize all streams to 0.

336+

// Streams will be created in setActiveDevice()

337+

for(int i = 0; i < (int)MAX_DEVICES; i++)

338+

streams[i] = (cudaStream_t)0;

339339
340340

const char* deviceENV = getenv("AF_CUDA_DEFAULT_DEVICE");

341341

if(!deviceENV) {

@@ -381,6 +381,11 @@ int DeviceManager::setActiveDevice(int device, int nId)

381381

if(nId == -1) nId = getDeviceNativeId(device);

382382

CUDA_CHECK(cudaSetDevice(nId));

383383

activeDev = device;

384+
385+

if(!streams[device]) {

386+

CUDA_CHECK(cudaStreamCreate(&streams[device]));

387+

}

388+
384389

return old;

385390

}

386391

}

Read the original on github.com ↗