Support reqd_work_group_size on native cpu

Author: PietroGhg

source/adapters/native_cpu/device.cpp

@@ -112,7 +112,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     // '0x8086' : 'Intel HD graphics vendor ID'
     return ReturnValue(uint32_t{0x8086});
   case UR_DEVICE_INFO_MAX_WORK_GROUP_SIZE:
-    return ReturnValue(size_t{256});
+    // TODO: provide a mechanism to estimate/configure this.
+    return ReturnValue(size_t{2048});
   case UR_DEVICE_INFO_MEM_BASE_ADDR_ALIGN:
     // Imported from level_zero
     return ReturnValue(uint32_t{8});
@@ -151,7 +152,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_MAX_WORK_ITEM_SIZES: {
     struct {
       size_t Arr[3];
-    } MaxGroupSize = {{256, 256, 1}};
+    } MaxGroupSize = {{256, 256, 256}};
     return ReturnValue(MaxGroupSize);
   }
   case UR_DEVICE_INFO_PREFERRED_VECTOR_WIDTH_CHAR:

source/adapters/native_cpu/enqueue.cpp

@@ -31,7 +31,7 @@ struct NDRDescT {
     for (uint32_t I = 0; I < WorkDim; I++) {
       GlobalOffset[I] = GlobalWorkOffset[I];
       GlobalSize[I] = GlobalWorkSize[I];
-      LocalSize[I] = LocalWorkSize[I];
+      LocalSize[I] = LocalWorkSize ? LocalWorkSize[I] : 1;
     }
     for (uint32_t I = WorkDim; I < 3; I++) {
       GlobalSize[I] = 1;
@@ -81,6 +81,16 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
     DIE_NO_IMPLEMENTATION;
   }
 
+  // Check reqd_work_group_size
+  if (hKernel->hasReqdWGSize() && pLocalWorkSize != nullptr) {
+    const auto &Reqd = hKernel->getReqdWGSize();
+    for (uint32_t Dim = 0; Dim < workDim; Dim++) {
+      if (pLocalWorkSize[Dim] != Reqd[Dim]) {
+        return UR_RESULT_ERROR_INVALID_WORK_GROUP_SIZE;
+      }
+    }
+  }
+
   // TODO: add proper error checking
   // TODO: add proper event dep management
   native_cpu::NDRDescT ndr(workDim, pGlobalWorkOffset, pGlobalWorkSize,

source/adapters/native_cpu/kernel.cpp

@@ -28,7 +28,16 @@ urKernelCreate(ur_program_handle_t hProgram, const char *pKernelName,
 
   auto f = reinterpret_cast<nativecpu_ptr_t>(
       const_cast<unsigned char *>(kernelEntry->second));
-  auto kernel = new ur_kernel_handle_t_(pKernelName, *f);
+  ur_kernel_handle_t_ *kernel;
+
+  // Set reqd_work_group_size for kernel if needed
+  const auto &ReqdMap = hProgram->KernelReqdWorkGroupSizeMD;
+  auto ReqdIt = ReqdMap.find(pKernelName);
+  if (ReqdIt != ReqdMap.end()) {
+    kernel = new ur_kernel_handle_t_(hProgram, pKernelName, *f, ReqdIt->second);
+  } else {
+    kernel = new ur_kernel_handle_t_(hProgram, pKernelName, *f);
+  }
 
   *phKernel = kernel;
 
@@ -84,13 +93,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
     //  case UR_KERNEL_INFO_PROGRAM:
     //    return ReturnValue(ur_program_handle_t{ Kernel->Program });
   case UR_KERNEL_INFO_FUNCTION_NAME:
-    if (hKernel->_name) {
-      return ReturnValue(hKernel->_name);
-    }
-    return UR_RESULT_ERROR_INVALID_FUNCTION_NAME;
-    //  case UR_KERNEL_INFO_NUM_ARGS:
-    //    return ReturnValue(uint32_t{ Kernel->ZeKernelProperties->numKernelArgs
-    //    });
+    return ReturnValue(hKernel->_name);
   case UR_KERNEL_INFO_REFERENCE_COUNT:
     return ReturnValue(uint32_t{hKernel->getReferenceCount()});
   case UR_KERNEL_INFO_ATTRIBUTES:
@@ -121,8 +124,16 @@ urKernelGetGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
     return returnValue(max_threads);
   }
   case UR_KERNEL_GROUP_INFO_COMPILE_WORK_GROUP_SIZE: {
-    size_t group_size[3] = {1, 1, 1};
-    return returnValue(group_size, 3);
+    size_t GroupSize[3] = {0, 0, 0};
+    const auto &ReqdWGSizeMDMap = hKernel->hProgram->KernelReqdWorkGroupSizeMD;
+    const auto ReqdWGSizeMD = ReqdWGSizeMDMap.find(hKernel->_name);
+    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: {
     int bytes = 0;

source/adapters/native_cpu/kernel.hpp

@@ -10,6 +10,8 @@
 
 #include "common.hpp"
 #include "nativecpu_state.hpp"
+#include "program.hpp"
+#include <array>
 #include <ur_api.h>
 #include <utility>
 
@@ -37,13 +39,16 @@ struct local_arg_info_t {
 
 struct ur_kernel_handle_t_ : RefCounted {
 
-  ur_kernel_handle_t_(const char *name, nativecpu_task_t subhandler)
-      : _name{name}, _subhandler{std::move(subhandler)} {}
+  ur_kernel_handle_t_(ur_program_handle_t hProgram, const char *name,
+                      nativecpu_task_t subhandler)
+      : hProgram(hProgram), _name{name}, _subhandler{std::move(subhandler)},
+        HasReqdWGSize(false) {}
 
   ur_kernel_handle_t_(const ur_kernel_handle_t_ &other)
       : _name(other._name), _subhandler(other._subhandler), _args(other._args),
         _localArgInfo(other._localArgInfo), _localMemPool(other._localMemPool),
-        _localMemPoolSize(other._localMemPoolSize) {
+        _localMemPoolSize(other._localMemPoolSize),
+        HasReqdWGSize(other.HasReqdWGSize), ReqdWGSize(other.ReqdWGSize) {
     incrementReferenceCount();
   }
 
@@ -52,13 +57,22 @@ struct ur_kernel_handle_t_ : RefCounted {
       free(_localMemPool);
     }
   }
-
-  const char *_name;
+  ur_kernel_handle_t_(ur_program_handle_t hProgram, const char *name,
+                      nativecpu_task_t subhandler,
+                      const native_cpu::ReqdWGSize_t &ReqdWGSize)
+      : hProgram(hProgram), _name{name}, _subhandler{std::move(subhandler)},
+        HasReqdWGSize(true), ReqdWGSize(ReqdWGSize) {}
+
+  ur_program_handle_t hProgram;
+  std::string _name;
   nativecpu_task_t _subhandler;
   std::vector<native_cpu::NativeCPUArgDesc> _args;
   std::vector<local_arg_info_t> _localArgInfo;
 
-  // To be called before enqueueing the kernel.
+  bool hasReqdWGSize() const { return HasReqdWGSize; }
+
+  const native_cpu::ReqdWGSize_t &getReqdWGSize() const { return ReqdWGSize; }
+
   void updateMemPool(size_t numParallelThreads) {
     // compute requested size.
     size_t reqSize = 0;
@@ -88,4 +102,6 @@ struct ur_kernel_handle_t_ : RefCounted {
 private:
   char *_localMemPool = nullptr;
   size_t _localMemPoolSize = 0;
+  bool HasReqdWGSize;
+  native_cpu::ReqdWGSize_t ReqdWGSize;
 };

source/adapters/native_cpu/program.cpp

@@ -12,6 +12,7 @@
 
 #include "common.hpp"
 #include "program.hpp"
+#include <cstdint>
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urProgramCreateWithIL(ur_context_handle_t hContext, const void *pIL,
@@ -26,6 +27,39 @@ urProgramCreateWithIL(ur_context_handle_t hContext, const void *pIL,
   DIE_NO_IMPLEMENTATION
 }
 
+// TODO: taken from CUDA adapter, move this to a common header?
+static 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()));
+}
+
+static ur_result_t getReqdWGSize(const ur_program_metadata_t &MetadataElement,
+                                 native_cpu::ReqdWGSize_t &res) {
+  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) * 2 ||
+                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);
+  std::get<0>(res) = ReqdWorkGroupElements[0];
+  std::get<1>(res) = ReqdWorkGroupElements[1];
+  std::get<2>(res) = ReqdWorkGroupElements[2];
+  return UR_RESULT_SUCCESS;
+}
+
 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 *pProperties,
@@ -40,6 +74,18 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(
 
   auto hProgram = new ur_program_handle_t_(
       hContext, reinterpret_cast<const unsigned char *>(pBinary));
+  if (pProperties != nullptr) {
+    for (uint32_t i = 0; i < pProperties->count; i++) {
+      auto mdNode = pProperties->pMetadatas[i];
+      std::string mdName(mdNode.pName);
+      auto [Prefix, Tag] = splitMetadataName(mdName);
+      if (Tag == __SYCL_UR_PROGRAM_METADATA_TAG_REQD_WORK_GROUP_SIZE) {
+        native_cpu::ReqdWGSize_t reqdWGSize;
+        getReqdWGSize(mdNode, reqdWGSize);
+        hProgram->KernelReqdWorkGroupSizeMD[Prefix] = std::move(reqdWGSize);
+      }
+    }
+  }
 
   const nativecpu_entry *nativecpu_it =
       reinterpret_cast<const nativecpu_entry *>(pBinary);

source/adapters/native_cpu/program.hpp

@@ -15,6 +15,10 @@
 #include "context.hpp"
 #include <map>
 
+namespace native_cpu {
+using ReqdWGSize_t = std::array<uint32_t, 3>;
+}
+
 struct ur_program_handle_t_ : RefCounted {
   ur_program_handle_t_(ur_context_handle_t ctx, const unsigned char *pBinary)
       : _ctx{ctx}, _ptr{pBinary} {}
@@ -30,6 +34,8 @@ struct ur_program_handle_t_ : RefCounted {
   };
 
   std::map<const char *, const unsigned char *, _compare> _kernels;
+  std::unordered_map<std::string, native_cpu::ReqdWGSize_t>
+      KernelReqdWorkGroupSizeMD;
 };
 
 // The nativecpu_entry struct is also defined as LLVM-IR in the