[SYCL] Handle required wg size attribute in HIP

Author: jchlanda

source/adapters/hip/enqueue.cpp

@@ -200,6 +200,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(
@@ -217,6 +218,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

@@ -11,13 +11,54 @@
 #include "program.hpp"
 
 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);
 }
 
 ur_program_handle_t_::~ur_program_handle_t_() { urContextRelease(Context); }
 
+ur_result_t
+ur_program_handle_t_::setMetadata(const ur_program_metadata_t *Metadata,
+                                  size_t Length) {
+  for (size_t i = 0; i < Length; ++i) {
+    const ur_program_metadata_t MetadataElement = Metadata[i];
+    std::string MetadataElementName{MetadataElement.pName};
+
+    std::string Prefix{};
+    std::string Tag{};
+    size_t SplitPos = MetadataElementName.rfind('@');
+    if (SplitPos != std::string::npos) {
+      Prefix = MetadataElementName.substr(0, SplitPos);
+      Tag = MetadataElementName.substr(SplitPos, MetadataElementName.length());
+    }
+
+    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;
+}
+
 ur_result_t ur_program_handle_t_::setBinary(const char *Source, size_t Length) {
   // Do not re-set program binary data which has already been set as that will
   // delete the old binary data.
@@ -246,7 +287,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramGetNativeHandle(
 /// Note: Only supports one device
 UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(
     ur_context_handle_t hContext, ur_device_handle_t hDevice, size_t size,
-    const uint8_t *pBinary, const ur_program_properties_t *,
+    const uint8_t *pBinary, const ur_program_properties_t *pProperties,
     ur_program_handle_t *phProgram) {
   UR_ASSERT(pBinary != nullptr && size != 0, UR_RESULT_ERROR_INVALID_BINARY);
   UR_ASSERT(hContext->getDevice()->get() == hDevice->get(),
@@ -257,8 +298,16 @@ 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;
+    } else if (pProperties->count == 0 && pProperties->pMetadatas != nullptr) {
+      return UR_RESULT_ERROR_INVALID_SIZE;
+    }
+    Result =
+        RetProgram->setMetadata(pProperties->pMetadatas, pProperties->count);
+    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"
 
@@ -24,6 +25,10 @@ struct ur_program_handle_t_ {
   std::atomic_uint32_t RefCount;
   ur_context_handle_t Context;
 
+  // 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];
@@ -33,6 +38,8 @@ struct ur_program_handle_t_ {
   ur_program_handle_t_(ur_context_handle_t Ctxt);
   ~ur_program_handle_t_();
 
+  ur_result_t setMetadata(const ur_program_metadata_t *Metadata, size_t Length);
+
   ur_result_t setBinary(const char *Binary, size_t BinarySizeInBytes);
 
   ur_result_t buildProgram(const char *BuildOptions);