cudaStream_t THCState_getDeviceStream(THCState *state, int device, int stream)
{
if (stream > state->numUserStreams || stream < 0)
{
THError("%d is not a stream", stream);
}
return (THCState_getDeviceResourcePtr(state, device)->streams == NULL) ? 0
: THCState_getDeviceResourcePtr(state, device)->streams[stream];
}
size_t THCState_getDeviceScratchSpaceSize(THCState* state, int device)
{
THCCudaResourcesPerDevice* res =
THCState_getDeviceResourcePtr(state, device);
return res->scratchSpacePerStream;
}
cublasHandle_t THCState_getDeviceBlasHandle(THCState *state, int device, int handle)
{
if (handle <= 0 || handle > state->numUserBlasHandles)
{
THError("%d is not a valid handle, valid range is: (1, %d)",
handle, state->numUserBlasHandles);
}
return THCState_getDeviceResourcePtr(state, device)->blasHandles[handle];
}
cusparseHandle_t THCState_getDeviceSparseHandle(THCState *state, int device, int handle)
{
if (handle <= 0 || handle > state->numUserSparseHandles) {
THError("%d is not a valid handle, valid range is: (1, %d)",
handle, state->numUserSparseHandles);
}
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, device);
THCState_reserveDeviceSparseHandles(state, device, handle);
return res->sparseHandles[handle - 1];
}
void* THCState_getDeviceScratchSpace(THCState* state, int device, int stream)
{
THCCudaResourcesPerDevice* res =
THCState_getDeviceResourcePtr(state, device);
if (stream > state->numUserStreams || stream < 0)
{
THError("%d is not a stream", stream);
}
return res->devScratchSpacePerStream[stream];
}
void THCState_reserveStreams(THCState* state, int numStreams, int nonBlocking)
{
if (numStreams <= state->numUserStreams)
{
return;
}
int prevDev = -1;
THCudaCheck(cudaGetDevice(&prevDev));
/* Otherwise, we have to allocate a new set of streams and stream data */
for (int dev = 0; dev < state->numDevices; ++dev) {
THCudaCheck(cudaSetDevice(dev));
/* +1 for the default stream as well */
cudaStream_t* newStreams =
(cudaStream_t*) malloc((numStreams + 1) * sizeof(cudaStream_t));
void** newScratchSpace =
(void**) malloc((numStreams + 1) * sizeof(void*));
/* Copy over old stream data
(0 is default stream, 1 ... numUserStreams are rest) */
for (int stream = 0; stream <= state->numUserStreams; ++stream) {
newStreams[stream] =
THCState_getDeviceStream(state, dev, stream);
newScratchSpace[stream] =
THCState_getDeviceScratchSpace(state, dev, stream);
}
/* Allocate new stream resources */
size_t scratchSpaceSize = THCState_getDeviceScratchSpaceSize(state, dev);
unsigned int flags =
nonBlocking ? cudaStreamNonBlocking : cudaStreamDefault;
for (int stream = state->numUserStreams + 1; stream <= numStreams; ++stream) {
newStreams[stream] = NULL;
THCudaCheck(cudaStreamCreateWithFlags(newStreams + stream, flags));
newScratchSpace[stream] = NULL;
THCudaCheck(THCudaMalloc(state, &newScratchSpace[stream], scratchSpaceSize));
}
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, dev);
free(res->streams);
res->streams = newStreams;
free(res->devScratchSpacePerStream);
res->devScratchSpacePerStream = newScratchSpace;
}
state->numUserStreams = numStreams;
THCudaCheck(cudaSetDevice(prevDev));
}
cudaStream_t THCState_getCurrentStream(THCState *state)
{
/* This is called at the point of kernel execution.
For some debugging code or improperly instrumented kernels,
`state` is null */
if (state) {
int device;
THCudaCheck(cudaGetDevice(&device));
int streamIndex = THCState_getCurrentStreamIndex(state);
if (streamIndex == 0) {
return NULL;
}
return THCState_getDeviceResourcePtr(state, device)->streams[streamIndex];
} else {
/* assume default stream */
return NULL;
}
}
void THCState_reserveBlasHandles(THCState* state, int numBlasHandles)
{
if (numBlasHandles <= state->numUserBlasHandles)
{
return;
}
int prevDev = -1;
THCudaCheck(cudaGetDevice(&prevDev));
/* Otherwise, we have to allocate a new set of blasHandles */
for (int dev = 0; dev < state->numDevices; ++dev) {
THCudaCheck(cudaSetDevice(dev));
/* +1 to be consistent with stream API, blas handle 0 is NULL and unused */
cublasHandle_t* newBlasHandles =
(cublasHandle_t*) malloc((numBlasHandles + 1) * sizeof(cublasHandle_t));
/* Copy over old blasHandles
(0 is NULL, 1 ... numUserBlasHandles are rest) */
newBlasHandles[0] = NULL;
for (int hndl = 1; hndl <= state->numUserBlasHandles; ++hndl) {
newBlasHandles[hndl] = THCState_getDeviceBlasHandle(state, dev, hndl);
}
/* Allocate new handles */
for (int hndl = state->numUserBlasHandles + 1; hndl <= numBlasHandles; ++hndl) {
newBlasHandles[hndl] = NULL;
THCublasCheck(cublasCreate(newBlasHandles + hndl));
}
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, dev);
free(res->blasHandles);
res->blasHandles = newBlasHandles;
}
state->numUserBlasHandles = numBlasHandles;
THCudaCheck(cudaSetDevice(prevDev));
}
void THCState_reserveDeviceSparseHandles(THCState* state, int device, int numSparseHandles)
{
int prevDev = -1;
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, device);
if (numSparseHandles <= res->numSparseHandles) {
return;
}
THCudaCheck(cudaGetDevice(&prevDev));
THCudaCheck(cudaSetDevice(device));
size_t size = numSparseHandles * sizeof(cusparseHandle_t);
cusparseHandle_t* handles = (cusparseHandle_t*) realloc(res->sparseHandles, size);
for (int i = res->numSparseHandles; i < numSparseHandles; ++i) {
handles[i] = NULL;
THCusparseCheck(cusparseCreate(&handles[i]));
}
res->sparseHandles = handles;
res->numSparseHandles = numSparseHandles;
THCudaCheck(cudaSetDevice(prevDev));
}
void THCudaInit(THCState* state)
{
int count = 0;
THCudaCheck(cudaGetDeviceCount(&count));
int device = 0;
THCudaCheck(cudaGetDevice(&device));
state->rngState = (THCRNGState*)malloc(sizeof(THCRNGState));
THCRandom_init(state, count, device);
THCAllocator_init(state);
state->numDevices = count;
state->deviceProperties =
(struct cudaDeviceProp*)malloc(count * sizeof(struct cudaDeviceProp));
state->numUserStreams = 0;
state->numUserBlasHandles = 0;
/* Enable P2P access between all pairs, if possible */
THCudaEnablePeerToPeerAccess(state);
state->resourcesPerDevice = (THCCudaResourcesPerDevice*)
malloc(count * sizeof(THCCudaResourcesPerDevice));
for (int i = 0; i < count; ++i) {
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, i);
THCudaCheck(cudaSetDevice(i));
THCudaCheck(cudaGetDeviceProperties(&state->deviceProperties[i], i));
/* Stream index 0 will be the default stream for convenience; by
default no user streams are reserved */
res->streams = NULL;
res->blasHandles = NULL;
/* The scratch space that we want to have available per each device is
based on the number of SMs available per device */
int numSM = state->deviceProperties[i].multiProcessorCount;
size_t sizePerStream = numSM * GLOBAL_SCRATCH_SPACE_PER_SM_STREAM;
res->scratchSpacePerStream = sizePerStream;
/* Allocate scratch space for each stream */
res->devScratchSpacePerStream = (void**) malloc(sizeof(void*));
THCudaCheck(THCudaMalloc(state, &res->devScratchSpacePerStream[0],
sizePerStream));
}
/* Restore to previous device */
THCudaCheck(cudaSetDevice(device));
/* Start in the default stream on the current device */
state->currentPerDeviceStream = 0;
state->currentStream = NULL;
/* There is no such thing as a default cublas handle.
To maintain consistency with streams API, handle 0 is always NULL and we
start counting at 1
*/
THCState_reserveBlasHandles(state, 1);
state->currentPerDeviceBlasHandle = 1;
state->currentBlasHandle = THCState_getDeviceBlasHandle(state, device, 1);
state->cutorchGCFunction = NULL;
state->cutorchGCData = NULL;
state->heapSoftmax = 3e8; // 300MB, adjusted upward dynamically
state->heapDelta = 0;
}
void THCudaInit(THCState* state)
{
if (!state->cudaDeviceAllocator) {
state->cudaDeviceAllocator = &defaultDeviceAllocator;
}
if (!state->cudaHostAllocator) {
state->cudaHostAllocator = &THCudaHostAllocator;
}
if (!state->cudaUVAAllocator) {
state->cudaUVAAllocator = &THCUVAAllocator;
}
int numDevices = 0;
THCudaCheck(cudaGetDeviceCount(&numDevices));
state->numDevices = numDevices;
int device = 0;
THCudaCheck(cudaGetDevice(&device));
/* Start in the default stream on the current device */
state->currentStreams = (THCThreadLocal*) malloc(numDevices * sizeof(THCThreadLocal));
for (int i = 0; i < numDevices; ++i) {
state->currentStreams[i] = THCThreadLocal_alloc();
}
state->currentPerDeviceBlasHandle = THCThreadLocal_alloc();
state->currentPerDeviceSparseHandle = THCThreadLocal_alloc();
state->resourcesPerDevice = (THCCudaResourcesPerDevice*)
malloc(numDevices * sizeof(THCCudaResourcesPerDevice));
memset(state->resourcesPerDevice, 0, numDevices * sizeof(THCCudaResourcesPerDevice));
state->deviceProperties =
(struct cudaDeviceProp*)malloc(numDevices * sizeof(struct cudaDeviceProp));
state->rngState = (THCRNGState*)malloc(sizeof(THCRNGState));
THCRandom_init(state, numDevices, device);
// By default, all direct p2p kernel access (besides copy) is disallowed,
// since direct access without knowing whether or not a certain operation
// should be cross-GPU leads to synchronization errors. The user can choose
// to disable this functionality, however.
state->p2pKernelAccessEnabled = 0;
// p2pAccessEnabled records if p2p copies are allowed between pairs of
// devices. Values include "1" (copy allowed), "0" (copy not allowed), and
// "-1" (unknown).
// Currently the max number of gpus in P2P group is 8, so if there are more
// we enable P2P in groups of 8
state->p2pAccessEnabled = (int**) malloc(sizeof(int*) * numDevices);
for (int i = 0; i < numDevices; ++i) {
state->p2pAccessEnabled[i] = (int*) malloc(sizeof(int) * numDevices);
for (int j = 0; j < numDevices; ++j)
if (i == j)
state->p2pAccessEnabled[i][j] = 1;
else if (j / THC_CUDA_MAX_PEER_SIZE != i / THC_CUDA_MAX_PEER_SIZE)
state->p2pAccessEnabled[i][j] = 0;
else
state->p2pAccessEnabled[i][j] = -1;
}
for (int i = 0; i < numDevices; ++i) {
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, i);
THCudaCheck(cudaSetDevice(i));
THCudaCheck(cudaGetDeviceProperties(&state->deviceProperties[i], i));
/* The scratch space that we want to have available per each device is
based on the number of SMs available per device. We guarantee a
minimum of 128kb of space per device, but to future-proof against
future architectures that may have huge #s of SMs, we guarantee that
we have at least 16 bytes for each SM. */
int numSM = state->deviceProperties[i].multiProcessorCount;
size_t sizePerStream =
MIN_GLOBAL_SCRATCH_SPACE_PER_DEVICE >= numSM * MIN_GLOBAL_SCRATCH_SPACE_PER_SM_STREAM ?
MIN_GLOBAL_SCRATCH_SPACE_PER_DEVICE :
numSM * MIN_GLOBAL_SCRATCH_SPACE_PER_SM_STREAM;
res->scratchSpacePerStream = sizePerStream;
}
/* Restore to previous device */
THCudaCheck(cudaSetDevice(device));
// Unlike CUDA streams, there is no NULL cuBLAS handle. The default THC
// cuBLAS handle is the first user BLAS handle. Note that the actual BLAS
// handles are created lazily.
state->numUserBlasHandles = 1;
state->numUserSparseHandles = 1;
state->heapSoftmax = 3e8; // 300MB, adjusted upward dynamically
state->heapDelta = 0;
}
void THCudaInit(THCState* state)
{
if (!state->cudaDeviceAllocator.malloc) {
THCState_initDefaultDeviceAllocator(&state->cudaDeviceAllocator);
}
int numDevices = 0;
THCudaCheck(cudaGetDeviceCount(&numDevices));
state->numDevices = numDevices;
int device = 0;
THCudaCheck(cudaGetDevice(&device));
/* Start in the default stream on the current device */
state->currentPerDeviceStream = THCThreadLocal_alloc();
state->currentPerDeviceBlasHandle = THCThreadLocal_alloc();
state->resourcesPerDevice = (THCCudaResourcesPerDevice*)
malloc(numDevices * sizeof(THCCudaResourcesPerDevice));
memset(state->resourcesPerDevice, 0, numDevices * sizeof(THCCudaResourcesPerDevice));
state->deviceProperties =
(struct cudaDeviceProp*)malloc(numDevices * sizeof(struct cudaDeviceProp));
state->rngState = (THCRNGState*)malloc(sizeof(THCRNGState));
THCRandom_init(state, numDevices, device);
state->cudaHostAllocator = (THAllocator*)malloc(sizeof(THAllocator));
THCAllocator_init(state->cudaHostAllocator);
/* Enable P2P access between all pairs, if possible */
THCudaEnablePeerToPeerAccess(state);
for (int i = 0; i < numDevices; ++i) {
THCCudaResourcesPerDevice* res = THCState_getDeviceResourcePtr(state, i);
THCudaCheck(cudaSetDevice(i));
THCudaCheck(cudaGetDeviceProperties(&state->deviceProperties[i], i));
/* The scratch space that we want to have available per each device is
based on the number of SMs available per device */
int numSM = state->deviceProperties[i].multiProcessorCount;
size_t sizePerStream = numSM * GLOBAL_SCRATCH_SPACE_PER_SM_STREAM;
res->scratchSpacePerStream = sizePerStream;
/* Allocate scratch space for each stream */
res->devScratchSpacePerStream = (void**) malloc(sizeof(void*));
THCudaCheck(THCudaMalloc(state, &res->devScratchSpacePerStream[0],
sizePerStream));
}
/* Restore to previous device */
THCudaCheck(cudaSetDevice(device));
/* There is no such thing as a default cublas handle.
To maintain consistency with streams API, handle 0 is always NULL and we
start counting at 1. If currentPerDeviceBlasHandle is 0 (the default
thread-local value), then we assume it means 1.
*/
THCState_reserveBlasHandles(state, 1);
state->heapSoftmax = 3e8; // 300MB, adjusted upward dynamically
state->heapDelta = 0;
}