Add multi work-group support for collective functional tests (#45)

- Added multi-work group support for the All-to-all, Fcollect, Broadcast, Barrier and Sync collective functional tests
- Renamed All-to-all and Fcollect tests to TeamAlltoAll and TeamFcollect

tests/functional_tests/team_broadcast_tester.cpp

@@ -20,8 +20,6 @@
  * IN THE SOFTWARE.
  *****************************************************************************/
 
-using namespace rocshmem;
-
 /* Declare the template with a generic implementation */
 template <typename T>
 __device__ void wg_team_broadcast(rocshmem_ctx_t ctx, rocshmem_team_t team,
@@ -65,28 +63,29 @@ __global__ void TeamBroadcastTest(int loop, int skip, long long int *start_time,
                                   long long int *end_time, T1 *source_buf,
                                   T1 *dest_buf, int size,
                                   ShmemContextType ctx_type,
-                                  rocshmem_team_t team) {
+                                  rocshmem_team_t *teams) {
   __shared__ rocshmem_ctx_t ctx;
   int wg_id = get_flat_grid_id();
 
   rocshmem_wg_init();
-  rocshmem_wg_ctx_create(ctx_type, &ctx);
+  rocshmem_wg_team_create_ctx(teams[wg_id], ctx_type, &ctx);
 
   int n_pes = rocshmem_ctx_n_pes(ctx);
+  source_buf += wg_id * size;
+  dest_buf += wg_id * size;
 
   __syncthreads();
 
-  for (int i = 0; i < loop; i++) {
+  for (int i = 0; i < loop + skip; i++) {
     if (i == skip && hipThreadIdx_x == 0) {
       start_time[wg_id] = wall_clock64();
     }
 
-    wg_team_broadcast<T1>(ctx, team,
-                          dest_buf,    // T* dest
-                          source_buf,  // const T* source
-                          size,        // int nelement
-                          0);          // int PE_root
-    rocshmem_ctx_wg_barrier_all(ctx);
+    wg_team_broadcast<T1>(ctx, teams[wg_id],
+                          dest_buf,         // T* dest
+                          source_buf,       // const T* source
+                          size,             // int nelement
+                          0);               // int PE_root
   }
 
   __syncthreads();
@@ -103,27 +102,51 @@ __global__ void TeamBroadcastTest(int loop, int skip, long long int *start_time,
  * HOST TESTER CLASS METHODS
  *****************************************************************************/
 template <typename T1>
-TeamBroadcastTester<T1>::TeamBroadcastTester(
-    TesterArguments args, std::function<void(T1 &, T1 &)> f1,
-    std::function<std::pair<bool, std::string>(const T1 &)> f2)
-    : Tester(args), init_buf{f1}, verify_buf{f2} {
-  source_buf = (T1 *)rocshmem_malloc(args.max_msg_size * sizeof(T1));
-  dest_buf = (T1 *)rocshmem_malloc(args.max_msg_size * sizeof(T1));
+TeamBroadcastTester<T1>::TeamBroadcastTester(TesterArguments args)
+    : Tester(args){
+  my_pe = rocshmem_team_my_pe(ROCSHMEM_TEAM_WORLD);
+  n_pes = rocshmem_team_n_pes(ROCSHMEM_TEAM_WORLD);
+
+  // Total number of elements in src buffer
+  int total_elems = (args.max_msg_size / sizeof(T1)) * args.num_wgs ;
+  int buff_size = total_elems * sizeof(T1);
+
+  source_buf = (T1 *)rocshmem_malloc(buff_size);
+  dest_buf = (T1 *)rocshmem_malloc(buff_size);
+
+  if (source_buf == nullptr || dest_buf == nullptr) {
+    std::cout << "Error allocating memory from symmetric heap" << std::endl;
+    std::cout << "source: " << source_buf << ", dest: " << dest_buf << std::endl;
+    rocshmem_global_exit(1);
+  }
+
+  char* value{nullptr};
+  if ((value = getenv("ROCSHMEM_MAX_NUM_TEAMS"))) {
+    num_teams = atoi(value);
+  }
+
+  CHECK_HIP(hipMalloc(&team_bcast_world_dup,
+                      sizeof(rocshmem_team_t) * num_teams));
 }
 
 template <typename T1>
 TeamBroadcastTester<T1>::~TeamBroadcastTester() {
   rocshmem_free(source_buf);
   rocshmem_free(dest_buf);
+  CHECK_HIP(hipFree(team_bcast_world_dup));
 }
 
 template <typename T1>
 void TeamBroadcastTester<T1>::preLaunchKernel() {
-  int n_pes = rocshmem_team_n_pes(ROCSHMEM_TEAM_WORLD);
-
-  team_bcast_world_dup = ROCSHMEM_TEAM_INVALID;
-  rocshmem_team_split_strided(ROCSHMEM_TEAM_WORLD, 0, 1, n_pes, nullptr, 0,
-                               &team_bcast_world_dup);
+  for (int team_i = 0; team_i < num_teams; team_i++) {
+    team_bcast_world_dup[team_i] = ROCSHMEM_TEAM_INVALID;
+    rocshmem_team_split_strided(ROCSHMEM_TEAM_WORLD, 0, 1, n_pes, nullptr, 0,
+                                 &team_bcast_world_dup[team_i]);
+    if (team_bcast_world_dup[team_i] == ROCSHMEM_TEAM_INVALID) {
+      printf("Team %d is invalid!\n", team_i);
+      abort();
+    }
+  }
 }
 
 template <typename T1>
@@ -131,34 +154,90 @@ void TeamBroadcastTester<T1>::launchKernel(dim3 gridSize, dim3 blockSize,
                                            int loop, uint64_t size) {
   size_t shared_bytes = 0;
 
-  hipLaunchKernelGGL(TeamBroadcastTest<T1>, gridSize, blockSize, shared_bytes,
-                     stream, loop, args.skip, start_time, end_time, source_buf,
-                     dest_buf, size, _shmem_context, team_bcast_world_dup);
+  int num_elems = size / sizeof(T1);
 
-  num_msgs = loop + args.skip;
-  num_timed_msgs = loop;
+  hipLaunchKernelGGL(TeamBroadcastTest<T1>, gridSize, blockSize,
+                     shared_bytes, stream, loop, args.skip,
+                     start_time, end_time, source_buf, dest_buf,
+                     num_elems, _shmem_context, team_bcast_world_dup);
+
+  num_msgs = (loop + args.skip) * gridSize.x;
+  num_timed_msgs = loop * gridSize.x;
 }
 
 template <typename T1>
 void TeamBroadcastTester<T1>::postLaunchKernel() {
-  rocshmem_team_destroy(team_bcast_world_dup);
+  for (int team_i = 0; team_i < num_teams; team_i++) {
+    rocshmem_team_destroy(team_bcast_world_dup[team_i]);
+  }
 }
 
 template <typename T1>
 void TeamBroadcastTester<T1>::resetBuffers(uint64_t size) {
-  for (uint64_t i = 0; i < args.max_msg_size; i++) {
-    init_buf(source_buf[i], dest_buf[i]);
+
+  int num_elems = size / sizeof(T1);
+  int buff_size = num_elems * sizeof(T1) * args.num_wgs;
+  int idx = 0;
+
+  for (int wg_id = 0; wg_id < args.num_wgs; wg_id++) {
+    for (int i = 0; i < num_elems; i++) {
+      idx = wg_id * num_elems + i;
+      if constexpr (std::is_same<T1, char>::value ||
+                    std::is_same<T1, signed char>::value ||
+                    std::is_same<T1, unsigned char>::value) {
+        source_buf[idx] = static_cast<T1>('a' + n_pes + wg_id);
+        dest_buf[idx] = static_cast<T1>('a' + wg_id);
+      }
+      else if constexpr (std::is_floating_point<T1>::value) {
+        source_buf[idx] = static_cast<T1>(3.14 + n_pes + wg_id);
+        dest_buf[idx] = static_cast<T1>(3.14 + wg_id);
+      }
+      else if constexpr (std::is_integral<T1>::value) {
+        source_buf[idx] = static_cast<T1>(n_pes + wg_id);
+        dest_buf[idx] = static_cast<T1>(wg_id);
+      }
+    }
   }
 }
 
 template <typename T1>
 void TeamBroadcastTester<T1>::verifyResults(uint64_t size) {
-  for (uint64_t i = 0; i < size; i++) {
-    auto r = verify_buf(dest_buf[i]);
-    if (r.first == false) {
-      fprintf(stderr, "Data validation error at idx %lu\n", i);
-      fprintf(stderr, "%s.\n", r.second.c_str());
-      exit(-1);
+
+  int num_elems = size / sizeof(T1);
+  int idx = 0;
+  T1 expected;
+
+  /**
+   * The verification routine here requires that the
+   * PE_root value is 0 which denotes that the
+   * sending processing element is rank 0.
+   *
+   * The difference in expected values arises from
+   * the specification for broadcast where the
+   * PE_root processing element does not copy the
+   * contents from its own source to dest during
+   * the broadcast.
+   */
+  for (int wg_id = 0; wg_id < args.num_wgs; wg_id++) {
+    for (int i = 0; i < num_elems; i++) {
+      idx = wg_id * num_elems + i;
+      if constexpr (std::is_same<T1, char>::value ||
+                    std::is_same<T1, signed char>::value ||
+                    std::is_same<T1, unsigned char>::value) {
+        expected = static_cast<T1>('a' + wg_id + (my_pe ? n_pes : 0));
+      }
+      else if constexpr (std::is_floating_point<T1>::value) {
+        expected = static_cast<T1>(3.14 + wg_id + (my_pe ? n_pes : 0));
+      }
+      else if constexpr (std::is_integral<T1>::value) {
+        expected = static_cast<T1>(wg_id + (my_pe ? n_pes : 0));
+      }
+      if (dest_buf[idx] != expected) {
+        std::cerr << "Data validation error at idx " << idx << std::endl;
+        std::cerr << "PE " << my_pe << " Got " << dest_buf[idx]
+        << ", Expected " << expected << std::endl;
+        exit(-1);
+      }
     }
   }
 }