support more than 8 gpus (#774)

1 files changed, 13 insertions, 2 deletions

diff --git a/lib/THC/THCGeneral.c b/lib/THC/THCGeneral.c
index 5882f91..5e966cf 100644
--- a/lib/THC/THCGeneral.c
+++ b/lib/THC/THCGeneral.c
@@ -16,6 +16,10 @@
  * above, whichever is greater.*/
 #define MIN_GLOBAL_SCRATCH_SPACE_PER_DEVICE 32768 * sizeof(float)
 
+/* Maximum number of P2P connections (if there are more than 9 then P2P is
+ * enabled in groups of 8). */
+#define THC_CUDA_MAX_PEER_SIZE 8
+
 THCCudaResourcesPerDevice* THCState_getDeviceResourcePtr(
   THCState *state, int device);
 
@@ -96,11 +100,18 @@ void THCudaInit(THCState* state)
   // 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);
-    memset(state->p2pAccessEnabled[i], -1, sizeof(int) * numDevices);
-    state->p2pAccessEnabled[i][i] = 1;
+    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) {