Auto device mode, plus allocation helper functions.

This diff introduces an alternative way of writing multi-GPU cutorch code. In this mode, the location of each tensor is specified, and the appropriate GPU for each kernel is determined automatically based on the location of its argument tensors. It's backwards-compatible and interoperable with the old-style multi-GPU API.

Author: adamlerer

THCGeneral.c

@@ -21,6 +21,8 @@ void THCudaInit(THCState* state)
  state->deviceProperties =
  (struct cudaDeviceProp*)malloc(count * sizeof(struct cudaDeviceProp));
 
+ THCState_setDeviceMode(state, THCStateDeviceModeManual);
+
  state->numUserStreams = 0;
  state->streamsPerDevice =
  (cudaStream_t**)malloc(count * sizeof(cudaStream_t*));
@@ -115,6 +117,30 @@ int THCState_getNumDevices(THCState *state)
  return state->numDevices;
 }
 
+THCStateDeviceMode THCState_getDeviceMode(THCState* state)
+{
+ return state->deviceMode;
+}
+
+void THCState_setDeviceMode(THCState* state, THCStateDeviceMode mode)
+{
+ state->deviceMode = mode;
+}
+
+void THCState_setDevice(THCState *state, int device)
+{
+ int curDev;
+ THCudaCheck(cudaGetDevice(&curDev));
+ if (device != curDev) {
+ THCudaCheck(cudaSetDevice(device));
+ THCRandom_setGenerator(state, device);
+ THCudaBlas_setHandle(state, device);
+
+ /* The stream is per device, so update the stream as well */
+ THCState_setStream(state, device, THCState_getCurrentStreamIndex(state));
+ }
+}
+
 void THCState_reserveStreams(THCState* state, int numStreams)
 {
  if (numStreams <= state->numUserStreams)

THCGeneral.h

@@ -36,6 +36,11 @@
 struct THCRNGState; /* Random number generator state. */
 struct THCBlasState;
 
+typedef enum THCStateDeviceMode {
+ THCStateDeviceModeManual,
+ THCStateDeviceModeAuto
+} THCStateDeviceMode;
+
 /* Global state to be held in the cutorch table. */
 typedef struct THCState
 {
@@ -55,6 +60,8 @@ typedef struct THCState
  /* Index of the current selected per-device stream. Actual CUDA stream changes
  based on the current device, since streams are per-device */
  int currentPerDeviceStream;
+ /* in DeviceModeAuto, cutorch can set the device based on the location of data tensors */
+ THCStateDeviceMode deviceMode;
 } THCState;
 
 THC_API void THCudaBlas_init(THCState *state, int num_devices, int current_device);
@@ -69,6 +76,9 @@ THC_API void THCudaEnablePeerToPeerAccess(THCState* state);
 
 /* State manipulators and accessors */
 THC_API int THCState_getNumDevices(THCState* state);
+THC_API void THCState_setDevice(THCState* state, int device);
+THC_API THCStateDeviceMode THCState_getDeviceMode(THCState* state);
+THC_API void THCState_setDeviceMode(THCState* state, THCStateDeviceMode mode);
 THC_API void THCState_reserveStreams(THCState* state, int numStreams);
 THC_API int THCState_getNumStreams(THCState* state);
 THC_API void THCState_resetStreams(THCState* state, int device);

THCStorage.c

@@ -21,6 +21,7 @@ THCudaStorage* THCudaStorage_new(THCState *state)
  storage->data = NULL;
  storage->size = 0;
  storage->refcount = 1;
+ storage->device = THC_DEVICE_NONE;
  storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
  return storage;
 }
@@ -36,6 +37,7 @@ THCudaStorage* THCudaStorage_newWithSize(THCState *state, long size)
 
  storage->size = size;
  storage->refcount = 1;
+ THCudaCheck(cudaGetDevice(&storage->device));
  storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
  return storage;
  }
@@ -91,6 +93,11 @@ THCudaStorage* THCudaStorage_newWithData(THCState *state, float *data, long size
  storage->data = data;
  storage->size = size;
  storage->refcount = 1;
+ if(size == 0) {
+ storage->device = THC_DEVICE_NONE;
+ } else {
+ THCudaCheck(cudaGetDevice(&storage->device));
+ }
  storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
  return storage;
 }

THCStorage.cu

@@ -26,6 +26,17 @@ void THCudaStorage_resize(THCState *state, THCudaStorage *self, long size)
  }
  else
  {
+ int curDev;
+ THCudaCheck(cudaGetDevice(&curDev));
+ if(self->device != THC_DEVICE_NONE) {
+ if (THCState_getDeviceMode(state) == THCStateDeviceModeAuto) {
+ THCudaCheck(cudaSetDevice(self->device));
+ }
+ else if(self->device != curDev) {
+ THError("THCudaStorage_resize: device mismatch: tensorDev=%d, curDev=%d", self->device + 1, curDev + 1);
+ }
+ }
+
  float *data = NULL;
  THCudaCheck(cudaMalloc((void**)(&data), size * sizeof(float)));
 
@@ -40,5 +51,19 @@ void THCudaStorage_resize(THCState *state, THCudaStorage *self, long size)
 
  self->data = data;
  self->size = size;
+ THCudaCheck(cudaGetDevice(&self->device));
+
+ THCudaCheck(cudaSetDevice(curDev));
  }
 }
+
+int THCudaStorage_getDevice(THCState* state, const THCudaStorage *storage) {
+ return storage->device;
+}
+
+void THCudaStorage_setDevice(THCState* state, THCudaStorage *storage, int device) {
+ if(storage->size > 0 && storage->device != device) {
+ THError("Cannot call setDevice() on a non-empty tensor. Use copy() instead.");
+ }
+ storage->device = device;
+}

THCStorage.h

@@ -8,12 +8,14 @@
 #define TH_STORAGE_RESIZABLE 2
 #define TH_STORAGE_FREEMEM 4
 
+#define THC_DEVICE_NONE -1
 
 typedef struct THCudaStorage
 {
  float *data;
  long size;
  int refcount;
+ int device;
  char flag;
  THAllocator *allocator;
  void *allocatorContext;
@@ -52,4 +54,7 @@ THC_API void THCudaStorage_free(THCState *state, THCudaStorage *storage);
 THC_API void THCudaStorage_resize(THCState *state, THCudaStorage *storage, long size);
 THC_API void THCudaStorage_fill(THCState *state, THCudaStorage *storage, float value);
 
+THC_API int THCudaStorage_getDevice(THCState* state, const THCudaStorage *storage);
+THC_API void THCudaStorage_setDevice(THCState* state, THCudaStorage *storage, int device);
+
 #endif

THCTensor.c

@@ -599,7 +599,7 @@ void THCudaTensor_freeCopyTo(THCState *state, THCudaTensor *self, THCudaTensor *
 static void THCudaTensor_rawInit(THCState *state, THCudaTensor *self)
 {
  self->refcount = 1;
- self->storage = NULL;
+ self->storage = THCudaStorage_new(state);
  self->storageOffset = 0;
  self->size = NULL;
  self->stride = NULL;
@@ -609,6 +609,7 @@ static void THCudaTensor_rawInit(THCState *state, THCudaTensor *self)
 
 static void THCudaTensor_rawSet(THCState *state, THCudaTensor *self, THCudaStorage *storage, long storageOffset, int nDimension, long *size, long *stride)
 {
+ THAssert(self->storage != NULL);
  /* storage */
  if(self->storage != storage)
  {
@@ -621,7 +622,7 @@ static void THCudaTensor_rawSet(THCState *state, THCudaTensor *self, THCudaStora
  THCudaStorage_retain(state, self->storage);
  }
  else
- self->storage = NULL;
+ self->storage = THCudaStorage_new(state);
  }
 
  /* storageOffset */
@@ -758,19 +759,36 @@ float THCudaTensor_get4d(THCState *state, const THCudaTensor *tensor, long x0, l
 
 int THCudaTensor_checkGPU(THCState *state, unsigned int nTensors, ...)
 {
- int curDev = -1;
- THCudaCheck(cudaGetDevice(&curDev));
+ int kernelDev;
+ if (THCState_getDeviceMode(state) == THCStateDeviceModeManual) {
+ THCudaCheck(cudaGetDevice(&kernelDev));
+ } else {
+ kernelDev = THC_DEVICE_NONE;
+ }
+
  va_list(args);
  va_start(args, nTensors);
- int valid = 1;
  for (unsigned int i = 0; i < nTensors; i++) {
  THCudaTensor* tensor = va_arg(args, THCudaTensor*);
  int tensorDev = THCudaTensor_getDevice(state, tensor);
- if (tensorDev != -1 && tensorDev != curDev) {
- valid = 0;
- break;
+ if (tensorDev != THC_DEVICE_NONE) {
+ if (kernelDev != tensorDev && kernelDev != THC_DEVICE_NONE) {
+ va_end(args);
+ return 0; // device mismatch
+ } else {
+ kernelDev = tensorDev;
+ }
  }
  }
  va_end(args);
- return valid;
+
+ if (THCState_getDeviceMode(state) == THCStateDeviceModeAuto) {
+ if (kernelDev == THC_DEVICE_NONE) {
+ return 0; // cannot determine device
+ } else {
+ THCState_setDevice(state, kernelDev);
+ }
+ }
+
+ return 1;
 }

THCTensor.cu

@@ -26,9 +26,14 @@ cudaTextureObject_t THCudaTensor_getTextureObject(THCState *state, THCudaTensor
  return texObj;
 }
 
-THC_API int THCudaTensor_getDevice(THCState* state, const THCudaTensor* thc) {
- if (!thc->storage) return -1;
- cudaPointerAttributes attr;
- THCudaCheck(cudaPointerGetAttributes(&attr, thc->storage->data));
- return attr.device;
+int THCudaTensor_getDevice(THCState* state, const THCudaTensor* self) {
+ THCudaStorage *storage = THCudaTensor_storage(state, self);
+ THAssert(storage != NULL);
+ return THCudaStorage_getDevice(state, storage);
 }
+
+void THCudaTensor_setDevice(THCState* state, THCudaTensor* self, int device) {
+ THCudaStorage *storage = THCudaTensor_storage(state, self);
+ THAssert(storage != NULL);
+ THCudaStorage_setDevice(state, storage, device);
+}

THCTensor.h

@@ -125,6 +125,7 @@ THC_API float THCudaTensor_get4d(THCState *state, const THCudaTensor *tensor, lo
 /* CUDA-specific functions */
 THC_API cudaTextureObject_t THCudaTensor_getTextureObject(THCState *state, THCudaTensor *self);
 THC_API int THCudaTensor_getDevice(THCState *state, const THCudaTensor *self);
+THC_API void THCudaTensor_setDevice(THCState* state, THCudaTensor* self, int device);
 THC_API int THCudaTensor_checkGPU(THCState *state, unsigned int nTensors, ...);
 
 #endif