int THCState_getCurrentSparseHandleIndex(THCState *state)
{
void* value = THCThreadLocal_get(state->currentPerDeviceSparseHandle);
if (value == NULL) {
return 1;
}
return (int) (intptr_t) value;
}
static THCStream* THCState_getStreamOnDevice(THCState* state, int device)
{
THCThreadLocal local = state->currentStreams[device];
THCStream* stream = (THCStream*)THCThreadLocal_get(local);
if (!stream) {
stream = THCStream_defaultStream(device);
THCStream_retain(stream);
THCThreadLocal_set(local, stream);
}
return stream;
}
static void THCState_setStreamOnDevice(THCState *state, int device, THCStream *stream)
{
THAssert(stream);
if (stream->device != device) {
THError("invalid stream; expected stream for device %d, but was on %d",
device, stream->device);
}
THCStream_retain(stream);
THCThreadLocal local = state->currentStreams[device];
THCStream_free((THCStream*)THCThreadLocal_get(local));
THCThreadLocal_set(local, stream);
}
void THCudaShutdown(THCState* state)
{
THCRandom_shutdown(state);
free(state->rngState);
free(state->deviceProperties);
int deviceCount = 0;
int prevDev = -1;
THCudaCheck(cudaGetDevice(&prevDev));
THCudaCheck(cudaGetDeviceCount(&deviceCount));
/* cleanup p2p access state */
for (int dev = 0; dev < deviceCount; ++dev) {
free(state->p2pAccessEnabled[dev]);
}
free(state->p2pAccessEnabled);
/* cleanup per-device state */
for (int dev = 0; dev < deviceCount; ++dev) {
THCudaCheck(cudaSetDevice(dev));
THCCudaResourcesPerDevice* res = &(state->resourcesPerDevice[dev]);
/* Free user defined BLAS handles */
for (int i = 0; i < res->numBlasHandles; ++i) {
THCublasCheck(cublasDestroy(res->blasHandles[i]));
}
/* Free user defined sparse handles */
for (int i = 0; i < res->numSparseHandles; ++i) {
THCusparseCheck(cusparseDestroy(res->sparseHandles[i]));
}
free(res->blasHandles);
free(res->sparseHandles);
THCStream_free((THCStream*)THCThreadLocal_get(state->currentStreams[dev]));
THCThreadLocal_free(state->currentStreams[dev]);
}
free(state->resourcesPerDevice);
if (state->cudaDeviceAllocator->emptyCache) {
state->cudaDeviceAllocator->emptyCache(state->cudaDeviceAllocator->state);
}
if (state->cudaHostAllocator == &THCCachingHostAllocator) {
THCCachingHostAllocator_emptyCache();
}
free(state->currentStreams);
THCThreadLocal_free(state->currentPerDeviceBlasHandle);
THCudaCheck(cudaSetDevice(prevDev));
}
int THCState_getCurrentStreamIndex(THCState *state)
{
void* value = THCThreadLocal_get(state->currentPerDeviceStream);
return (int) (intptr_t) value;
}
