void THCState_setPeerToPeerAccess(THCState* state, int dev, int devToAccess,
int enable)
{
/* This will perform device bounds checking for us */
int prevEnabled = THCState_getPeerToPeerAccess(state, dev, devToAccess);
if (enable != prevEnabled) {
/* If we're attempting to enable p2p access but p2p access isn't */
/* supported, throw an error */
if (enable) {
int access = 0;
THCudaCheck(cudaDeviceCanAccessPeer(&access, dev, devToAccess));
if (!access) {
THError("p2p access not supported for %d accessing %d",
dev, devToAccess);
}
}
state->p2pAccessEnabled[dev][devToAccess] = enable;
int prevDev = 0;
THCudaCheck(cudaGetDevice(&prevDev));
THCudaCheck(cudaSetDevice(dev));
/* This should be in sync with the current access state */
if (enable) {
THCudaCheck(cudaDeviceEnablePeerAccess(devToAccess, 0));
} else {
THCudaCheck(cudaDeviceDisablePeerAccess(devToAccess));
}
THCudaCheck(cudaSetDevice(prevDev));
}
}
static int cutorch_getPeerToPeerAccess(lua_State *L)
{
THCState *state = cutorch_getstate(L);
int dev = (int) luaL_checknumber(L, 1) - 1;
int devToAccess = (int) luaL_checknumber(L, 2) - 1;
/* device bounds checking is performed within */
int enabled = THCState_getPeerToPeerAccess(state, dev, devToAccess);
lua_pushboolean(L, enabled);
return 1;
}
void THCStorage_resize(THCState *state, THCStorage *self, ptrdiff_t size)
{
THArgCheck(size >= 0, 2, "invalid size");
THAssert(self->allocator != NULL);
int device;
THCudaCheck(cudaGetDevice(&device));
if(!(self->flag & TH_STORAGE_RESIZABLE))
THError("Trying to resize storage that is not resizable");
size_t elementSize = at::elementSize(self->scalar_type);
if (self->allocator->realloc) {
void * data_ptr = self->data_ptr;
cudaError_t err = (*self->allocator->realloc)(
self->allocatorContext,
(void**)&(data_ptr),
self->size * elementSize,
size * elementSize, THCState_getCurrentStreamOnDevice(state, device));
if (err != cudaSuccess) {
THCudaCheck(err);
}
self->size = size;
self->device = device;
return;
}
if(size == 0)
{
if(self->flag & TH_STORAGE_FREEMEM) {
THCudaCheck(
(*self->allocator->free)(self->allocatorContext, self->data_ptr));
}
self->data_ptr = NULL;
self->size = 0;
self->device = device;
}
else
{
void *data = NULL;
cudaError_t err =
(*self->allocator->malloc)(self->allocatorContext,
(void**)&(data),
size * elementSize,
THCState_getCurrentStreamOnDevice(state, device));
THCudaCheck(err);
if (self->data_ptr) {
// Enable p2p access when the memcpy is across devices
THCState_getPeerToPeerAccess(state, device, self->device);
THCudaCheck(cudaMemcpyAsync(data,
self->data_ptr,
THMin(self->size, size) * elementSize,
cudaMemcpyDeviceToDevice,
THCState_getCurrentStream(state)));
if(self->flag & TH_STORAGE_FREEMEM) {
THCudaCheck(
(*self->allocator->free)(self->allocatorContext, self->data_ptr));
}
}
self->data_ptr = data;
self->size = size;
self->device = device;
}
}
