[SYCL] [NATIVECPU] Implement urKernelSetArgLocal (#11101)

This PR adds support to `local_accessors` by implementing
`urKernelSetArgLocal`.

Author: PietroGhg

sycl/plugins/unified_runtime/ur/adapters/native_cpu/enqueue.cpp

@@ -64,6 +64,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   // TODO: add proper event dep management
   sycl::detail::NDRDescT ndr =
       getNDRDesc(workDim, pGlobalWorkOffset, pGlobalWorkSize, pLocalWorkSize);
+  hKernel->handleLocalArgs();
 
   __nativecpu_state state(ndr.GlobalSize[0], ndr.GlobalSize[1],
                           ndr.GlobalSize[2], ndr.LocalSize[0], ndr.LocalSize[1],
@@ -90,6 +91,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   // TODO: we should avoid calling clear here by avoiding using push_back
   // in setKernelArgs.
   hKernel->_args.clear();
+  hKernel->_localArgInfo.clear();
   return UR_RESULT_SUCCESS;
 }
 

sycl/plugins/unified_runtime/ur/adapters/native_cpu/kernel.cpp

@@ -54,12 +54,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgValue(
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgLocal(
     ur_kernel_handle_t hKernel, uint32_t argIndex, size_t argSize,
     const ur_kernel_arg_local_properties_t *pProperties) {
-  std::ignore = hKernel;
-  std::ignore = argIndex;
-  std::ignore = argSize;
   std::ignore = pProperties;
-
-  DIE_NO_IMPLEMENTATION
+  // emplace a placeholder kernel arg, gets replaced with a pointer to the
+  // memory pool before enqueueing the kernel.
+  hKernel->_args.emplace_back(nullptr);
+  hKernel->_localArgInfo.emplace_back(argIndex, argSize);
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,

sycl/plugins/unified_runtime/ur/adapters/native_cpu/kernel.hpp

@@ -17,6 +17,13 @@ using nativecpu_kernel_t = void(const sycl::detail::NativeCPUArgDesc *,
 using nativecpu_ptr_t = nativecpu_kernel_t *;
 using nativecpu_task_t = std::function<nativecpu_kernel_t>;
 
+struct local_arg_info_t {
+  uint32_t argIndex;
+  size_t argSize;
+  local_arg_info_t(uint32_t argIndex, size_t argSize)
+      : argIndex(argIndex), argSize(argSize) {}
+};
+
 struct ur_kernel_handle_t_ : RefCounted {
 
   ur_kernel_handle_t_(const char *name, nativecpu_task_t subhandler)
@@ -25,4 +32,47 @@ struct ur_kernel_handle_t_ : RefCounted {
   const char *_name;
   nativecpu_task_t _subhandler;
   std::vector<sycl::detail::NativeCPUArgDesc> _args;
+  std::vector<local_arg_info_t> _localArgInfo;
+
+  // To be called before enqueing the kernel.
+  void handleLocalArgs() {
+    updateMemPool();
+    size_t offset = 0;
+    for (auto &entry : _localArgInfo) {
+      _args[entry.argIndex].MPtr =
+          reinterpret_cast<char *>(_localMemPool) + offset;
+      // update offset in the memory pool
+      // Todo: update this offset computation when we have work-group
+      // level parallelism.
+      offset += entry.argSize;
+    }
+  }
+
+  ~ur_kernel_handle_t_() {
+    if (_localMemPool) {
+      free(_localMemPool);
+    }
+  }
+
+private:
+  void updateMemPool() {
+    // compute requested size.
+    // Todo: currently we execute only one work-group at a time, so for each
+    // local arg we can allocate just 1 * argSize local arg. When we implement
+    // work-group level parallelism we should allocate N * argSize where N is
+    // the number of work groups being executed in parallel (e.g. number of
+    // threads in the thread pool).
+    size_t reqSize = 0;
+    for (auto &entry : _localArgInfo) {
+      reqSize += entry.argSize;
+    }
+    if (reqSize == 0 || reqSize == _localMemPoolSize) {
+      return;
+    }
+    // realloc handles nullptr case
+    _localMemPool = realloc(_localMemPool, reqSize);
+    _localMemPoolSize = reqSize;
+  }
+  void *_localMemPool = nullptr;
+  size_t _localMemPoolSize = 0;
 };