EXSWHTEC-269 - Implement test for warp ballot function (#188)

Change-Id: I71e0448a4676b6ba671b600065918e8fd533cc61

catch/unit/warp/CMakeLists.txt

@@ -4,6 +4,7 @@ set(TEST_SRC
   warp_shfl.cc
   warp_shfl_up.cc
   warp_shfl_down.cc
+  warp_ballot.cc
 )
 
 hip_add_exe_to_target(NAME WarpTest

catch/unit/warp/warp_ballot.cc

@@ -0,0 +1,120 @@
+/*
+Copyright (c) 2023 Advanced Micro Devices, Inc. All rights reserved.
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+
+#include "warp_vote_common.hh"
+
+#include <bitset>
+
+/**
+ * @addtogroup ballot ballot
+ * @{
+ * @ingroup DeviceLanguageTest
+ * `unsigned long long int __ballot(int predicate)` -
+ * Contains unit test for warp ballot function
+ */
+
+namespace cg = cooperative_groups;
+
+__global__ void kernel_ballot(uint64_t* const out, const uint64_t* const active_masks,
+                              uint64_t predicate) {
+  if (deactivate_thread(active_masks)) {
+    return;
+  }
+
+  const auto grid = cg::this_grid();
+  const auto warp = cg::tiled_partition(cg::this_thread_block(), warpSize);
+
+  out[grid.thread_rank()] =
+      __ballot((predicate & (static_cast<uint64_t>(1) << warp.thread_rank())));
+}
+
+class WarpBallot : public WarpVoteTest<WarpBallot, uint64_t> {
+ public:
+  void launch_kernel(uint64_t* const arr_dev, const uint64_t* const active_masks) {
+    auto test_case = GENERATE(range(0, 5));
+    predicate_mask_ = get_predicate_mask(test_case, this->warp_size_);
+    INFO("Predicate mask: " << predicate_mask_);
+    kernel_ballot<<<this->grid_.grid_dim_, this->grid_.block_dim_>>>(arr_dev, active_masks,
+                                                                     predicate_mask_);
+  }
+
+  void validate(const uint64_t* const arr) {
+    ArrayAllOf(arr, this->grid_.thread_count_, [this](unsigned int i) -> std::optional<uint64_t> {
+      const auto rank_in_block = this->grid_.thread_rank_in_block(i).value();
+      const auto rank_in_warp = rank_in_block % this->warp_size_;
+      const auto warp_idx = this->warps_in_block_ * (i / this->grid_.threads_in_block_count_) +
+          rank_in_block / this->warp_size_;
+      const auto block_rank = warp_idx / this->warps_in_block_;
+      const std::bitset<sizeof(uint64_t) * 8> active_mask(this->active_masks_[warp_idx]);
+
+      auto partition_size = this->warp_size_;
+      // If the number of threads in a block is not a multiple of warp size, the
+      // last warp will have inactive threads and partition size must be recalculated
+      if (warp_idx == this->warps_in_block_ * (block_rank + 1) - 1) {
+        partition_size =
+            this->grid_.threads_in_block_count_ - (this->warps_in_block_ - 1) * this->warp_size_;
+      }
+
+      if (!active_mask.test(rank_in_warp))
+        return std::nullopt;
+      else {
+        // Active predicate mask must be calculated as partition can be smaller than warp_size
+        auto active_predicate = get_active_predicate(predicate_mask_, partition_size);
+        return (active_predicate & this->active_masks_[warp_idx]);
+      }
+    });
+  }
+
+ private:
+  uint64_t predicate_mask_;
+};
+
+/**
+ * Test Description
+ * ------------------------
+ *  - Validates the warp ballot function behavior. Threads are deactivated based on the passed
+ * active mask. The predicate for each thread is determined according to the generated predicate
+ * mask.
+ * Test source
+ * ------------------------
+ *  - unit/warp/warp_ballot.cc
+ * Test requirements
+ * ------------------------
+ *  - HIP_VERSION >= 5.2
+ *  - Device supports warp ballot
+ */
+TEST_CASE("Unit_Warp_Ballot_Positive_Basic") {
+  int device;
+  hipDeviceProp_t device_properties;
+  HIP_CHECK(hipGetDevice(&device));
+  HIP_CHECK(hipGetDeviceProperties(&device_properties, device));
+
+  if (!device_properties.arch.hasWarpBallot) {
+    HipTest::HIP_SKIP_TEST("Device doesn't support Warp Ballot!");
+    return;
+  }
+
+  SECTION("Warp Ballot with specified active mask") {
+    WarpBallot().run(false);
+  }
+
+  SECTION("Warp Ballot with random active mask") {
+    WarpBallot().run(true);
+  }
+}

catch/unit/warp/warp_vote_common.hh

@@ -0,0 +1,117 @@
+/*
+Copyright (c) 2023 Advanced Micro Devices, Inc. All rights reserved.
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+
+#pragma once
+
+#include "warp_common.hh"
+
+#include <cpu_grid.h>
+#include <resource_guards.hh>
+#include <utils.hh>
+
+inline uint64_t get_predicate_mask(unsigned int test_case, unsigned int warp_size) {
+  uint64_t predicate_mask = 0;
+  switch (test_case) {
+    case 0:  // no thread
+      predicate_mask = 0;
+      break;
+    case 1:  // 1st thread
+      predicate_mask = 1;
+      break;
+    case 2:  // last thread
+      predicate_mask = static_cast<uint64_t>(1) << (warp_size - 1);
+      break;
+    case 3:  // all threads
+      predicate_mask = 0xFFFFFFFFFFFFFFFF;
+      break;
+    default:  // random
+      static std::mt19937_64 mt(test_case);
+      std::uniform_int_distribution<uint64_t> dist(0, std::numeric_limits<uint64_t>::max());
+      predicate_mask = dist(mt);
+  }
+  return predicate_mask;
+}
+
+inline uint64_t get_active_predicate(uint64_t predicate, size_t partition_size) {
+  uint64_t active_predicate = predicate;
+  for (int i = partition_size; i < 64; i++) {
+    active_predicate &= ~(static_cast<uint64_t>(1) << i);
+  }
+  return active_predicate;
+}
+
+template <typename Derived, typename T> class WarpVoteTest {
+ public:
+  WarpVoteTest() : warp_size_{get_warp_size()} {}
+
+  void run(bool random = false) {
+    const auto blocks = GenerateBlockDimensionsForShuffle();
+    INFO("Grid dimensions: x " << blocks.x << ", y " << blocks.y << ", z " << blocks.z);
+    const auto threads = GenerateThreadDimensionsForShuffle();
+    INFO("Block dimensions: x " << threads.x << ", y " << threads.y << ", z " << threads.z);
+    grid_ = CPUGrid(blocks, threads);
+
+    const auto alloc_size = grid_.thread_count_ * sizeof(T);
+    LinearAllocGuard<T> arr_dev(LinearAllocs::hipMalloc, alloc_size);
+    LinearAllocGuard<T> arr(LinearAllocs::hipHostMalloc, alloc_size);
+    HIP_CHECK(hipMemset(arr_dev.ptr(), 0, alloc_size));
+
+    warps_in_block_ = (grid_.threads_in_block_count_ + warp_size_ - 1) / warp_size_;
+    const auto warps_in_grid = warps_in_block_ * grid_.block_count_;
+    LinearAllocGuard<uint64_t> active_masks_dev(LinearAllocs::hipMalloc,
+                                                warps_in_grid * sizeof(uint64_t));
+    active_masks_.resize(warps_in_grid);
+
+    if (random) {
+      std::generate(active_masks_.begin(), active_masks_.end(), [] {
+        return GenerateRandomInteger(0ul, std::numeric_limits<uint64_t>().max());
+      });
+    } else {
+      unsigned long long int i = 0;
+      std::generate(active_masks_.begin(), active_masks_.end(),
+                    [this, &i]() { return get_active_mask(i++, warp_size_); });
+    }
+
+    HIP_CHECK(hipMemcpy(active_masks_dev.ptr(), active_masks_.data(),
+                        warps_in_grid * sizeof(uint64_t), hipMemcpyHostToDevice));
+    cast_to_derived().launch_kernel(arr_dev.ptr(), active_masks_dev.ptr());
+    HIP_CHECK(hipGetLastError());
+    HIP_CHECK(hipMemcpy(arr.ptr(), arr_dev.ptr(), alloc_size, hipMemcpyDeviceToHost));
+    HIP_CHECK(hipDeviceSynchronize());
+
+    cast_to_derived().validate(arr.ptr());
+  }
+
+ private:
+  int get_warp_size() const {
+    int current_dev = -1;
+    HIP_CHECK(hipGetDevice(&current_dev));
+    int warp_size = 0u;
+    HIP_CHECK(hipDeviceGetAttribute(&warp_size, hipDeviceAttributeWarpSize, 0));
+    return warp_size;
+  }
+
+  Derived& cast_to_derived() { return reinterpret_cast<Derived&>(*this); }
+
+ protected:
+  const int warp_size_;
+  CPUGrid grid_;
+  unsigned int warps_in_block_;
+  std::vector<uint64_t> active_masks_;
+};