[UR][CUDA][HIP] Fix Set Arg Local

Author: Petr Vesely

sycl/plugins/unified_runtime/pi2ur.hpp

@@ -2190,8 +2190,12 @@ inline pi_result piKernelSetArg(pi_kernel Kernel, pi_uint32 ArgIndex,
 
   ur_kernel_handle_t UrKernel = reinterpret_cast<ur_kernel_handle_t>(Kernel);
 
-  HANDLE_ERRORS(
-      urKernelSetArgValue(UrKernel, ArgIndex, ArgSize, nullptr, ArgValue));
+  if (ArgValue) {
+    HANDLE_ERRORS(
+        urKernelSetArgValue(UrKernel, ArgIndex, ArgSize, nullptr, ArgValue));
+  } else {
+    HANDLE_ERRORS(urKernelSetArgLocal(UrKernel, ArgIndex, ArgSize, nullptr));
+  }
   return PI_SUCCESS;
 }
 

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

@@ -186,11 +186,22 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgValue(
 
   ur_result_t Result = UR_RESULT_SUCCESS;
   try {
-    if (pArgValue) {
-      hKernel->setKernelArg(argIndex, argSize, pArgValue);
-    } else {
-      hKernel->setKernelLocalArg(argIndex, argSize);
-    }
+    hKernel->setKernelArg(argIndex, argSize, pArgValue);
+  } catch (ur_result_t Err) {
+    Result = Err;
+  }
+  return Result;
+}
+
+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 = pProperties;
+  UR_ASSERT(argSize, UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE);
+
+  ur_result_t Result = UR_RESULT_SUCCESS;
+  try {
+    hKernel->setKernelLocalArg(argIndex, argSize);
   } catch (ur_result_t Err) {
     Result = Err;
   }

sycl/plugins/unified_runtime/ur/adapters/cuda/ur_interface_loader.cpp

@@ -115,7 +115,7 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetKernelProcAddrTable(
   pDdiTable->pfnGetSubGroupInfo = urKernelGetSubGroupInfo;
   pDdiTable->pfnRelease = urKernelRelease;
   pDdiTable->pfnRetain = urKernelRetain;
-  pDdiTable->pfnSetArgLocal = nullptr;
+  pDdiTable->pfnSetArgLocal = urKernelSetArgLocal;
   pDdiTable->pfnSetArgMemObj = urKernelSetArgMemObj;
   pDdiTable->pfnSetArgPointer = urKernelSetArgPointer;
   pDdiTable->pfnSetArgSampler = urKernelSetArgSampler;

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

@@ -168,11 +168,22 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgValue(
     const ur_kernel_arg_value_properties_t *, const void *pArgValue) {
   ur_result_t Result = UR_RESULT_SUCCESS;
   try {
-    if (pArgValue) {
-      hKernel->setKernelArg(argIndex, argSize, pArgValue);
-    } else {
-      hKernel->setKernelLocalArg(argIndex, argSize);
-    }
+    hKernel->setKernelArg(argIndex, argSize, pArgValue);
+  } catch (ur_result_t Err) {
+    Result = Err;
+  }
+  return Result;
+}
+
+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 = pProperties;
+  UR_ASSERT(argSize, UR_RESULT_ERROR_INVALID_KERNEL_ARGUMENT_SIZE);
+
+  ur_result_t Result = UR_RESULT_SUCCESS;
+  try {
+    hKernel->setKernelLocalArg(argIndex, argSize);
   } catch (ur_result_t Err) {
     Result = Err;
   }

sycl/plugins/unified_runtime/ur/adapters/hip/ur_interface_loader.cpp

@@ -115,7 +115,7 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetKernelProcAddrTable(
   pDdiTable->pfnGetSubGroupInfo = urKernelGetSubGroupInfo;
   pDdiTable->pfnRelease = urKernelRelease;
   pDdiTable->pfnRetain = urKernelRetain;
-  pDdiTable->pfnSetArgLocal = nullptr;
+  pDdiTable->pfnSetArgLocal = urKernelSetArgLocal;
   pDdiTable->pfnSetArgMemObj = urKernelSetArgMemObj;
   pDdiTable->pfnSetArgPointer = urKernelSetArgPointer;
   pDdiTable->pfnSetArgSampler = urKernelSetArgSampler;

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

@@ -419,12 +419,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgLocal(
     const ur_kernel_arg_local_properties_t
         *Properties ///< [in][optional] argument properties
 ) {
-  std::ignore = Kernel;
-  std::ignore = ArgIndex;
   std::ignore = Properties;
-  std::ignore = ArgSize;
-  urPrint("[UR][L0] %s function not implemented!\n", __FUNCTION__);
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+
+  UR_CALL(urKernelSetArgValue(Kernel, ArgIndex, ArgSize, nullptr, nullptr));
+
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(