Add support for cooperative group type - thread_block

Change-Id: If3770b6d6718a638b70f527ae2533d9ef3267ff4

hipamd/include/hip/hcc_detail/hip_cooperative_groups.h

@@ -166,6 +166,55 @@ this_grid() {
   return grid_group(internal::grid::size());
 }
 
+/** \brief The workgroup (thread-block in CUDA terminology) cooperative group
+ *         type
+ *
+ *  \details Represents an intra-workgroup cooperative group type where the
+ *           participating threads within the group are exctly the same threads
+ *           which are participated in the currently executing `workgroup`
+ */
+class thread_block : public thread_group {
+  // Only these friend functions are allowed to construct an object of this
+  // class and access its resources
+  friend __CG_QUALIFIER__ thread_block this_thread_block();
+
+ protected:
+  // Construct a workgroup thread group (through the API this_thread_block())
+  explicit __CG_QUALIFIER__ thread_block(uint32_t size)
+      : thread_group(internal::cg_workgroup, size) { }
+
+ public:
+  // 3-dimensional block index within the grid
+  __CG_QUALIFIER__ dim3 group_index() {
+    return internal::workgroup::group_index();
+  }
+  // 3-dimensional thread index within the block
+  __CG_QUALIFIER__ dim3 thread_index() {
+    return internal::workgroup::thread_index();
+  }
+  __CG_QUALIFIER__ uint32_t thread_rank() const {
+    return internal::workgroup::thread_rank();
+  }
+  __CG_QUALIFIER__ bool is_valid() const {
+    return internal::workgroup::is_valid();
+  }
+  __CG_QUALIFIER__ void sync() const {
+    internal::workgroup::sync();
+  }
+};
+
+/** \brief User exposed API interface to construct workgroup cooperative
+ *         group type object - `thread_block`
+ *
+ *  \details User is not allowed to directly construct an object of type
+ *           `thread_block`. Instead, he should construct it through this API
+ *           function
+ */
+__CG_QUALIFIER__ thread_block
+this_thread_block() {
+  return thread_block(internal::workgroup::size());
+}
+
 /**
  *  Implemenation of all publicly exposed base class APIs
  */
@@ -177,6 +226,9 @@ __CG_QUALIFIER__ uint32_t thread_group::thread_rank() const {
     case internal::cg_grid: {
       return (static_cast<const grid_group*>(this)->thread_rank());
     }
+    case internal::cg_workgroup: {
+      return (static_cast<const thread_block*>(this)->thread_rank());
+    }
     default: {
       return 0; //TODO(mahesha)
     }
@@ -191,6 +243,9 @@ __CG_QUALIFIER__ bool thread_group::is_valid() const {
     case internal::cg_grid: {
       return (static_cast<const grid_group*>(this)->is_valid());
     }
+    case internal::cg_workgroup: {
+      return (static_cast<const thread_block*>(this)->is_valid());
+    }
     default: {
       return false;
     }
@@ -207,6 +262,10 @@ __CG_QUALIFIER__ void thread_group::sync() const {
       static_cast<const grid_group*>(this)->sync();
       break;
     }
+    case internal::cg_workgroup: {
+      static_cast<const thread_block*>(this)->sync();
+      break;
+    }
   }
 }
 

hipamd/include/hip/hcc_detail/hip_cooperative_groups_helper.h

@@ -60,7 +60,8 @@ namespace internal {
 typedef enum {
   cg_invalid,
   cg_multi_grid,
-  cg_grid
+  cg_grid,
+  cg_workgroup
 } group_type;
 
 /**
@@ -136,6 +137,43 @@ __CG_STATIC_QUALIFIER__ void sync() {
 
 } // namespace grid
 
+/**
+ *  Functionalities related to `workgroup` (thread_block in CUDA terminology)
+ *  cooperative group type
+ */
+namespace workgroup {
+
+__CG_STATIC_QUALIFIER__ dim3 group_index() {
+  return (dim3((uint32_t)hipBlockIdx_x, (uint32_t)hipBlockIdx_y,
+               (uint32_t)hipBlockIdx_z));
+}
+
+__CG_STATIC_QUALIFIER__ dim3 thread_index() {
+  return (dim3((uint32_t)hipThreadIdx_x, (uint32_t)hipThreadIdx_y,
+               (uint32_t)hipThreadIdx_z));
+}
+
+__CG_STATIC_QUALIFIER__ uint32_t size() {
+  return((uint32_t)(hipBlockDim_x * hipBlockDim_y * hipBlockDim_z));
+}
+
+__CG_STATIC_QUALIFIER__ uint32_t thread_rank() {
+ return ((uint32_t)((hipThreadIdx_z * hipBlockDim_y * hipBlockDim_x) +
+                    (hipThreadIdx_y * hipBlockDim_x) +
+                    (hipThreadIdx_x)));
+}
+
+__CG_STATIC_QUALIFIER__ bool is_valid() {
+   //TODO(mahesha) any functionality need to be added here? I believe not
+  return true;
+}
+
+__CG_STATIC_QUALIFIER__ void sync() {
+  __syncthreads();
+}
+
+} // namespace workgroup
+
 } // namespace internal
 
 } // namespace cooperative_groups

hipamd/tests/src/cg/hipCGThreadBlockType.cpp

@@ -0,0 +1,164 @@
+/*
+Copyright (c) 2020 - present 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.
+*/
+
+
+/* HIT_START
+ * BUILD: %t %s ../test_common.cpp
+ * TEST: %t
+ * HIT_END
+ */
+
+#include "test_common.h"
+#include "hip/hip_cooperative_groups.h"
+
+#define HIP_ASSERT(lhs, rhs) assert(lhs == rhs)
+
+using namespace cooperative_groups;
+
+static __global__
+void kernel_cg_thread_block_type(dim3 *groupIndexD,
+                                 dim3 *thdIndexD,
+                                 int *sizeD,
+                                 int *thdRankD,
+                                 int *isValidD,
+                                 int *syncValD)
+{
+  thread_block tb = this_thread_block();
+
+  // Consider the workgroup id (0, 1, 1) and thread id (0, 1, 2) for validation
+  // of the test
+  int isBlockIdx011 =
+    (hipBlockIdx_x == 0 && hipBlockIdx_y == 1 && hipBlockIdx_z == 1);
+  int isThreadIdx012 =
+    (hipThreadIdx_x == 0 && hipThreadIdx_y == 1 && hipThreadIdx_z == 2);
+
+  if (isBlockIdx011 && isThreadIdx012) {
+    *groupIndexD = tb.group_index();
+    *thdIndexD = tb.thread_index();
+    *sizeD = tb.size();
+    *thdRankD = tb.thread_rank();
+    *isValidD = tb.is_valid();
+  }
+
+  // Consider local thread id (0, 0, 0) and (0, 0, 1) for validation of sync()
+  // api
+  __shared__ int sVar[2];
+  int isThreadIdx000 =
+    (hipThreadIdx_x == 0 && hipThreadIdx_y == 0 && hipThreadIdx_z == 0);
+  int isThreadIdx001 =
+    (hipThreadIdx_x == 0 && hipThreadIdx_y == 0 && hipThreadIdx_z == 1);
+
+  if (isThreadIdx000)
+    sVar[0] = 10;
+  if (isThreadIdx001)
+    sVar[1] = 20;
+  tb.sync();
+  if (isBlockIdx011 && isThreadIdx012)
+    *syncValD = sVar[0] + sVar[1];
+}
+
+static void test_cg_thread_block_type()
+{
+  dim3 *groupIndexD, *groupIndexH;
+  dim3 *thdIndexD, *thdIndexH;
+  int *sizeD, *sizeH;
+  int *thdRankD, *thdRankH;
+  int *isValidD, *isValidH;
+  int *syncValD, *syncValH;
+
+  // Allocate device memory
+  HIP_ASSERT(hipMalloc((void**)&groupIndexD, sizeof(dim3)), hipSuccess);
+  HIP_ASSERT(hipMalloc((void**)&thdIndexD, sizeof(dim3)), hipSuccess);
+  HIP_ASSERT(hipMalloc((void**)&sizeD, sizeof(int)), hipSuccess);
+  HIP_ASSERT(hipMalloc((void**)&thdRankD, sizeof(int)), hipSuccess);
+  HIP_ASSERT(hipMalloc((void**)&isValidD, sizeof(int)), hipSuccess);
+  HIP_ASSERT(hipMalloc((void**)&syncValD, sizeof(int)), hipSuccess);
+
+  // Allocate host memory
+  HIP_ASSERT(hipHostMalloc((void**)&groupIndexH, sizeof(dim3)), hipSuccess);
+  HIP_ASSERT(hipHostMalloc((void**)&thdIndexH, sizeof(dim3)), hipSuccess);
+  HIP_ASSERT(hipHostMalloc((void**)&sizeH, sizeof(int)), hipSuccess);
+  HIP_ASSERT(hipHostMalloc((void**)&thdRankH, sizeof(int)), hipSuccess);
+  HIP_ASSERT(hipHostMalloc((void**)&isValidH, sizeof(int)), hipSuccess);
+  HIP_ASSERT(hipHostMalloc((void**)&syncValH, sizeof(int)), hipSuccess);
+
+  // Launch Kernel
+  hipLaunchKernelGGL(kernel_cg_thread_block_type,
+                     dim3(2, 2, 2),
+                     dim3(4, 4, 4),
+                     0,
+                     0,
+                     groupIndexD,
+                     thdIndexD,
+                     sizeD,
+                     thdRankD,
+                     isValidD,
+                     syncValD);
+
+  // Copy result from device to host
+  HIP_ASSERT(hipMemcpy(groupIndexH, groupIndexD, sizeof(dim3), hipMemcpyDeviceToHost), hipSuccess);
+  HIP_ASSERT(hipMemcpy(thdIndexH, thdIndexD, sizeof(dim3), hipMemcpyDeviceToHost), hipSuccess);
+  HIP_ASSERT(hipMemcpy(sizeH, sizeD, sizeof(int), hipMemcpyDeviceToHost), hipSuccess);
+  HIP_ASSERT(hipMemcpy(thdRankH, thdRankD, sizeof(int), hipMemcpyDeviceToHost), hipSuccess);
+  HIP_ASSERT(hipMemcpy(isValidH, isValidD, sizeof(int), hipMemcpyDeviceToHost), hipSuccess);
+  HIP_ASSERT(hipMemcpy(syncValH, syncValD, sizeof(int), hipMemcpyDeviceToHost), hipSuccess);
+
+  // Validate result
+  // Group index should be (0, 1, 1)
+  HIP_ASSERT(groupIndexH->x, 0);
+  HIP_ASSERT(groupIndexH->y, 1);
+  HIP_ASSERT(groupIndexH->z, 1);
+  // Thread index should be (0, 1, 2)
+  HIP_ASSERT(thdIndexH->x, 0);
+  HIP_ASSERT(thdIndexH->y, 1);
+  HIP_ASSERT(thdIndexH->z, 2);
+  // Workgroup size should be 64
+  HIP_ASSERT(*sizeH, 64);
+  // Thread rank should be 36
+  HIP_ASSERT(*thdRankH, 36);
+  // Call to is_valid() should return true
+  HIP_ASSERT(*isValidH, 1);
+  // syncVal should be equal to 30
+  HIP_ASSERT(*syncValH, 30);
+
+  // Free device memory
+  HIP_ASSERT(hipFree(groupIndexD), hipSuccess);
+  HIP_ASSERT(hipFree(thdIndexD), hipSuccess);
+  HIP_ASSERT(hipFree(sizeD), hipSuccess);
+  HIP_ASSERT(hipFree(thdRankD), hipSuccess);
+  HIP_ASSERT(hipFree(isValidD), hipSuccess);
+  HIP_ASSERT(hipFree(syncValD), hipSuccess);
+
+  //Free host memory
+  HIP_ASSERT(hipHostFree(groupIndexH), hipSuccess);
+  HIP_ASSERT(hipHostFree(thdIndexH), hipSuccess);
+  HIP_ASSERT(hipHostFree(sizeH), hipSuccess);
+  HIP_ASSERT(hipHostFree(thdRankH), hipSuccess);
+  HIP_ASSERT(hipHostFree(isValidH), hipSuccess);
+  HIP_ASSERT(hipHostFree(syncValH), hipSuccess);
+}
+
+int main()
+{
+  test_cg_thread_block_type();
+  passed();
+}