Updated main function

lattice_boltzmann/main.c

@@ -14,8 +14,10 @@
 
 #define INDEX_3D_TO_1D(y, x, z, nb_columns, nb_depths) x + HEIGHT* (y + WIDTH* z)
 
-#define N_ITERATIONS 200
+#define N_MEASURES 15
+#define N_ITERATIONS 100
 #define NB_VALUES 27
+#define ROOT_RANK 0
 
 typedef struct lbm_values {
     float values[NB_VALUES];
@@ -51,7 +53,7 @@ void compute_next_lbm(struct dispatch_context *dc, struct futhark_context *fc, c
 
 int main(int argc, char *argv[]) {
     if (argc < 4) {
-        printf("usage: mpirun -n <nb_proc> %s <height> <width> <depth>\n", argv[0]);
+        printf("usage: mpirun -n <nb_proc> %s <nb_devices> <height> <width> <depth>\n", argv[0]);
         return EXIT_FAILURE;
     }
 
@@ -62,9 +64,18 @@ int main(int argc, char *argv[]) {
     MPI_Comm_size(MPI_COMM_WORLD, &world_size);
 
     struct futhark_context_config *fut_config = futhark_context_config_new();
+    int nb_devices = atoi(argv[1]);
+#if defined(FUTHARK_BACKEND_cuda) || defined(FUTHARK_BACKEND_opencl)
+    #if defined(FUTHARK_BACKEND_opencl)
+    futhark_context_config_list_devices(fut_config);
+#endif
+    char device[4] = {0};
+    snprintf(device, sizeof(device), "#%d", my_rank % nb_devices);
+    futhark_context_config_set_device(fut_config, device);
+#endif
     struct futhark_context *fut_context = futhark_context_new(fut_config);
 
-    int lbm_dimensions[3] = {atoi(argv[1]), atoi(argv[2]), atoi(argv[3])};
+    int lbm_dimensions[3] = {atoi(argv[2]), atoi(argv[3]), atoi(argv[4])};
 
     // https://stackoverflow.com/a/33624425
     int count = 1;
@@ -85,8 +96,15 @@ int main(int argc, char *argv[]) {
     chunk_info_print(&ci);
     init_chunk_lbm(&ci);
 
-    for (int i = 0; i < N_ITERATIONS; ++i) {
-        compute_next_lbm(disp_context, fut_context, &ci);
+    for (int i = 0; i < N_MEASURES; ++i) {
+        double start = MPI_Wtime();
+        for (int j = 0; j < N_ITERATIONS; ++j) {
+            compute_next_lbm(disp_context, fut_context, &ci);
+        }
+        double finish = MPI_Wtime();
+        if (my_rank == ROOT_RANK) {
+            printf("%d;%d;%d;%f\n", world_size, nb_devices, lbm_dimensions[0], finish - start);
+        }
     }
 
     dispatch_context_free(disp_context);