Ghost exchange

src/array.hpp

@@ -20,18 +20,20 @@ class Array
     std::vector<int> strides;
 
     bool exchange_in_progress;
-    int ghost_size;
+    std::vector<int> ghost_size;
     float* send_ghost_buffers[4];
     float* recv_ghost_buffers[4];
+    int num_chares;
 
-    Array(int name_, std::vector<int> shape_, std::vector<int> global_shape_, int ghost_depth_, bool* boundary)
+    Array(int name_, std::vector<int> shape_, std::vector<int> global_shape_, int ghost_depth_, bool* boundary, int num_chares_)
         : name(name_)
         , local_shape(shape_)
         , global_shape(global_shape_)
         , ghost_depth(ghost_depth_)
         , generation(0)
         , ghost_generation(-1)
         , exchange_in_progress(false)
+        , num_chares(num_chares_)
     {
         for(int i = 0; i < shape_.size(); i++)
             shape.push_back(shape_[i] + 2 * ghost_depth_);
@@ -42,7 +44,9 @@ class Array
         total_size = shape[0] * shape[1];
         total_local_size = local_shape[0] * local_shape[1];
         hapiCheck(cudaMalloc((void**) &data, sizeof(float) * total_size));
-        ghost_size = ghost_depth * local_shape[0]; // FIXME assuming square array
+        ghost_size.resize(shape.size());
+        ghost_size[0] = ghost_depth * local_shape[0]; // FIXME assuming square array
+        ghost_size[1] = ghost_depth * local_shape[1];
         allocate_ghost_buffers(boundary);
     }
 
@@ -57,9 +61,10 @@ class Array
         {
             if (!boundary[i])
             {
+                int idx_ghost_size = i/2 ;
                 DEBUG_PRINT("PE %i> Allocating ghost buffers for array %i in dir %i\n", CkMyPe(), name, i);
-                hapiCheck(cudaMalloc((void**) &(send_ghost_buffers[i]), sizeof(float) * ghost_size));
-                hapiCheck(cudaMalloc((void**) &(recv_ghost_buffers[i]), sizeof(float) * ghost_size));
+                hapiCheck(cudaMalloc((void**) &(send_ghost_buffers[i]), sizeof(float) * ghost_size[idx_ghost_size]));
+                hapiCheck(cudaMalloc((void**) &(recv_ghost_buffers[i]), sizeof(float) * ghost_size[idx_ghost_size]));
                 DEBUG_PRINT("PE %i> Send ghost buffer %i: %p\n", CkMyPe(), i, send_ghost_buffers[i]);
                 DEBUG_PRINT("PE %i> Recv ghost buffer %i: %p\n", CkMyPe(), i, recv_ghost_buffers[i]);
             }

src/ast.cpp

@@ -109,8 +109,8 @@ std::string SliceNode::generate_code(Context* ctx)
         std::string idy = ctx->is_shmem(ctx->get_active()) ? fmt::format("s_idy{}", ctx->get_active()) : "idy";
         std::string idx = ctx->is_shmem(ctx->get_active()) ? fmt::format("s_idx{}", ctx->get_active()) : "idx";
         return fmt::format("IDX2D(({} + {}) * {}, ({} + {}) * {}, {})", 
-            idy, offset.index[1].start, offset.index[1].step, 
-            idx, offset.index[0].start, offset.index[0].step, 
+            idy, offset.index[0].start, offset.index[0].step, 
+            idx, offset.index[1].start, offset.index[1].step, 
             ctx->get_step());
     }
 }
@@ -213,21 +213,26 @@ Slice Kernel::get_launch_bounds(int name, Array* array, int* chare_index)
     Slice global_bounds;
     //DEBUG_PRINT("Calculating global_bounds on: (%i : %i), (%i : %i)\n", slice.index[0].start, slice.index[0].stop,
     //    slice.index[1].start, slice.index[1].stop);
-    global_bounds.index[0].start = GET_START(slice, array->global_shape, 0);
-    global_bounds.index[0].stop = GET_STOP(slice, array->global_shape, 0);
+    global_bounds.index[0].start = GET_START(slice, array->global_shape, 1);
+    global_bounds.index[0].stop = GET_STOP(slice, array->global_shape, 1);
     //int stepx = slice.index[0].step;
 
-    global_bounds.index[1].start = GET_START(slice, array->global_shape, 1);
-    global_bounds.index[1].stop = GET_STOP(slice, array->global_shape, 1);
+    global_bounds.index[1].start = GET_START(slice, array->global_shape, 0);
+    global_bounds.index[1].stop = GET_STOP(slice, array->global_shape, 0);
     //int stepy = slice.index[1].step;
 
     // now find what indices of the global bound belong to this chare
-    // FIXME assumption all arrays are square
 
-    int chare_startx = chare_index[0] * array->local_shape[0];
-    int chare_stopx = (chare_index[0] + 1) * array->local_shape[0];
-    int chare_starty = chare_index[1] * array->local_shape[1];
-    int chare_stopy = (chare_index[1] + 1) * array->local_shape[1];
+    int base_nums_x = array->global_shape[1]/array->num_chares;
+    int base_nums_y = array->global_shape[0]/array->num_chares;
+    int remainder_nums_x = array->global_shape[1] % array->num_chares;
+    int remainder_nums_y = array->global_shape[0] % array->num_chares;
+    int majority_x = std::min(remainder_nums_x,chare_index[0]);
+    int majority_y = std::min(remainder_nums_y,chare_index[1]);
+    int chare_startx = (base_nums_x+1)*majority_x + (chare_index[0]-majority_x)*base_nums_x;
+    int chare_stopx = chare_startx + array->local_shape[1];
+    int chare_starty = (base_nums_y+1)*majority_y + (chare_index[1]-majority_y)*base_nums_y;
+    int chare_stopy = chare_starty + array->local_shape[0];
 
     //DEBUG_PRINT("DEBUG AST> (%i, %i) > (%i, %i) > (%i, %i)\n", chare_index[0], chare_index[1], chare_startx, chare_stopx, chare_starty, chare_stopy);
 

src/stencil.ci

@@ -29,6 +29,6 @@ module stencil
 
         entry [reductiontarget] void operation_done(double start);
 
-        entry void gather(int name, int index_x, int index_y, int local_dim, int num_chares, int data_size, float data[data_size]);
+        entry void gather(int name, int index_x, int index_y, int local_dim_x, int local_dim_y, int num_chares, int data_size, float data[data_size], int global_dim_x, int global_dim_y);
     };
 }