[SYCL] Handle required wg size attribute in HIP

Author: jchlanda

source/adapters/cuda/program.cpp

@@ -8,6 +8,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "ur_util.hpp"
 #include "program.hpp"
 
 bool getMaxRegistersJitOptionValue(const std::string &BuildOptions,
@@ -52,15 +53,6 @@ ur_program_handle_t_::ur_program_handle_t_(ur_context_handle_t Context)
 
 ur_program_handle_t_::~ur_program_handle_t_() { urContextRelease(Context); }
 
-std::pair<std::string, std::string>
-splitMetadataName(const std::string &metadataName) {
-  size_t splitPos = metadataName.rfind('@');
-  if (splitPos == std::string::npos)
-    return std::make_pair(metadataName, std::string{});
-  return std::make_pair(metadataName.substr(0, splitPos),
-                        metadataName.substr(splitPos, metadataName.length()));
-}
-
 ur_result_t
 ur_program_handle_t_::setMetadata(const ur_program_metadata_t *Metadata,
                                   size_t Length) {

source/adapters/hip/enqueue.cpp

@@ -201,6 +201,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   size_t MaxWorkGroupSize = 0u;
   size_t MaxThreadsPerBlock[3] = {};
   bool ProvidedLocalWorkGroupSize = (pLocalWorkSize != nullptr);
+  size_t *ReqdThreadsPerBlock = hKernel->ReqdThreadsPerBlock;
 
   {
     ur_result_t Result = urDeviceGetInfo(
@@ -218,6 +219,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
 
     if (ProvidedLocalWorkGroupSize) {
       auto isValid = [&](int dim) {
+        if (ReqdThreadsPerBlock[dim] != 0 &&
+            pLocalWorkSize[dim] != ReqdThreadsPerBlock[dim])
+          return UR_RESULT_ERROR_INVALID_WORK_GROUP_SIZE;
         UR_ASSERT(pLocalWorkSize[dim] <= MaxThreadsPerBlock[dim],
                   UR_RESULT_ERROR_INVALID_WORK_GROUP_SIZE);
         // Checks that local work sizes are a divisor of the global work sizes

source/adapters/hip/kernel.cpp

@@ -87,14 +87,17 @@ urKernelGetGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
     return ReturnValue(size_t(MaxThreads));
   }
   case UR_KERNEL_GROUP_INFO_COMPILE_WORK_GROUP_SIZE: {
-    size_t group_size[3] = {0, 0, 0};
-    // Returns the work-group size specified in the kernel source or IL.
-    // If the work-group size is not specified in the kernel source or IL,
-    // (0, 0, 0) is returned.
-    // https://www.khronos.org/registry/OpenCL/sdk/2.1/docs/man/xhtml/clGetKernelWorkGroupInfo.html
-
-    // TODO: can we extract the work group size from the PTX?
-    return ReturnValue(group_size, 3);
+    size_t GroupSize[3] = {0, 0, 0};
+    const auto &ReqdWGSizeMDMap =
+        hKernel->getProgram()->KernelReqdWorkGroupSizeMD;
+    const auto ReqdWGSizeMD = ReqdWGSizeMDMap.find(hKernel->getName());
+    if (ReqdWGSizeMD != ReqdWGSizeMDMap.end()) {
+      const auto ReqdWGSize = ReqdWGSizeMD->second;
+      GroupSize[0] = std::get<0>(ReqdWGSize);
+      GroupSize[1] = std::get<1>(ReqdWGSize);
+      GroupSize[2] = std::get<2>(ReqdWGSize);
+    }
+    return ReturnValue(GroupSize, 3);
   }
   case UR_KERNEL_GROUP_INFO_LOCAL_MEM_SIZE: {
     // OpenCL LOCAL == HIP SHARED

source/adapters/hip/kernel.hpp

@@ -43,6 +43,9 @@ struct ur_kernel_handle_t_ {
   ur_program_handle_t Program;
   std::atomic_uint32_t RefCount;
 
+  static constexpr uint32_t ReqdThreadsPerBlockDimensions = 3u;
+  size_t ReqdThreadsPerBlock[ReqdThreadsPerBlockDimensions];
+
   /// Structure that holds the arguments to the kernel.
   /// Note earch argument size is known, since it comes
   /// from the kernel signature.
@@ -134,6 +137,12 @@ struct ur_kernel_handle_t_ {
                       ur_context_handle_t Ctxt)
       : Function{Func}, FunctionWithOffsetParam{FuncWithOffsetParam},
         Name{Name}, Context{Ctxt}, Program{Program}, RefCount{1} {
+    ur_result_t RetError = urKernelGetGroupInfo(
+        this, Context->getDevice(),
+        UR_KERNEL_GROUP_INFO_COMPILE_WORK_GROUP_SIZE,
+        sizeof(ReqdThreadsPerBlock), ReqdThreadsPerBlock, nullptr);
+    (void)RetError;
+    assert(RetError == UR_RESULT_SUCCESS);
     urProgramRetain(Program);
     urContextRetain(Context);
   }

source/adapters/hip/program.cpp

@@ -8,6 +8,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "ur_util.hpp"
 #include "program.hpp"
 
 #ifdef SYCL_ENABLE_KERNEL_FUSION
@@ -75,8 +76,8 @@ void getCoMgrBuildLog(const amd_comgr_data_set_t BuildDataSet, char *BuildLog,
 #endif
 
 ur_program_handle_t_::ur_program_handle_t_(ur_context_handle_t Ctxt)
-    : Module{nullptr}, Binary{}, BinarySizeInBytes{0}, RefCount{1}, Context{
-                                                                        Ctxt} {
+    : Module{nullptr}, Binary{}, BinarySizeInBytes{0}, RefCount{1},
+      Context{Ctxt}, KernelReqdWorkGroupSizeMD{} {
   urContextRetain(Context);
 }
 
@@ -94,7 +95,32 @@ ur_program_handle_t_::setMetadata(const ur_program_metadata_t *Metadata,
       assert(MetadataElement.type == UR_PROGRAM_METADATA_TYPE_UINT32);
       IsRelocatable = MetadataElement.value.data32;
     }
+
+    auto [Prefix, Tag] = splitMetadataName(MetadataElementName);
+
+    if (Tag == __SYCL_UR_PROGRAM_METADATA_TAG_REQD_WORK_GROUP_SIZE) {
+      // If metadata is reqd_work_group_size, record it for the corresponding
+      // kernel name.
+      size_t MDElemsSize = MetadataElement.size - sizeof(std::uint64_t);
+
+      // Expect between 1 and 3 32-bit integer values.
+      UR_ASSERT(MDElemsSize >= sizeof(std::uint32_t) &&
+                    MDElemsSize <= sizeof(std::uint32_t) * 3,
+                UR_RESULT_ERROR_INVALID_WORK_GROUP_SIZE);
+
+      // Get pointer to data, skipping 64-bit size at the start of the data.
+      const char *ValuePtr =
+          reinterpret_cast<const char *>(MetadataElement.value.pData) +
+          sizeof(std::uint64_t);
+      // Read values and pad with 1's for values not present.
+      std::uint32_t ReqdWorkGroupElements[] = {1, 1, 1};
+      std::memcpy(ReqdWorkGroupElements, ValuePtr, MDElemsSize);
+      KernelReqdWorkGroupSizeMD[Prefix] =
+          std::make_tuple(ReqdWorkGroupElements[0], ReqdWorkGroupElements[1],
+                          ReqdWorkGroupElements[2]);
+    }
   }
+
   return UR_RESULT_SUCCESS;
 }
 
@@ -410,8 +436,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(
   std::unique_ptr<ur_program_handle_t_> RetProgram{
       new ur_program_handle_t_{hContext}};
 
-  // TODO: Set metadata here and use reqd_work_group_size information.
-  // See urProgramCreateWithBinary in CUDA adapter.
   if (pProperties) {
     if (pProperties->count > 0 && pProperties->pMetadatas == nullptr) {
       return UR_RESULT_ERROR_INVALID_NULL_POINTER;
@@ -420,8 +444,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(
     }
     Result =
         RetProgram->setMetadata(pProperties->pMetadatas, pProperties->count);
+    UR_ASSERT(Result == UR_RESULT_SUCCESS, Result);
   }
-  UR_ASSERT(Result == UR_RESULT_SUCCESS, Result);
 
   auto pBinary_string = reinterpret_cast<const char *>(pBinary);
   if (size == 0) {

source/adapters/hip/program.hpp

@@ -12,6 +12,7 @@
 #include <ur_api.h>
 
 #include <atomic>
+#include <unordered_map>
 
 #include "context.hpp"
 
@@ -28,6 +29,10 @@ struct ur_program_handle_t_ {
   // Metadata
   bool IsRelocatable = false;
 
+  // Metadata
+  std::unordered_map<std::string, std::tuple<uint32_t, uint32_t, uint32_t>>
+      KernelReqdWorkGroupSizeMD;
+
   constexpr static size_t MAX_LOG_SIZE = 8192u;
 
   char ErrorLog[MAX_LOG_SIZE], InfoLog[MAX_LOG_SIZE];

source/common/ur_util.hpp

@@ -293,4 +293,12 @@ inline ur_result_t exceptionToResult(std::exception_ptr eptr) {
     }
 }
 
+inline std::pair<std::string, std::string>
+splitMetadataName(const std::string &metadataName) {
+  size_t splitPos = metadataName.rfind('@');
+  if (splitPos == std::string::npos)
+    return std::make_pair(metadataName, std::string{});
+  return std::make_pair(metadataName.substr(0, splitPos),
+                        metadataName.substr(splitPos, metadataName.length()));
+}
 #endif /* UR_UTIL_H */