[SYCL][CUDA][HIP] Add implementation of piextGetDeviceFunctionPointer (#4599)

This patch adds an implementation of `piextGetDeviceFunctionPointer` for CUDA and HIP.
program_impl.cpp `has_kernel`now uses `piextGetDeviceFunctionPointer` to implement a more robust method for CUDA and HIP.
The previous regex search used with `piPrgramGetInfo`, with `PI_PROGRAM_INFO_KERNEL_NAMES`, is removed and replaced with a die message as it does not work correctly with module-splitting per source. 

This patch allows CUDA to pass tests:
DeviceCodeSplit/split-per-kernel.cpp
DeviceCodeSplit/split-per-source-main.cpp
KernelAndProgram/basic-program.cpp
KernelAndProgram/kernel-and-program.cpp

A follow up PR to llvm-test-suite will be made to enable these test for cuda.

This PR comes from modifications requested in #4565

Author: AidanBeltonS

sycl/include/CL/sycl/detail/pi.h

@@ -112,6 +112,10 @@ typedef enum {
   PI_INVALID_IMAGE_FORMAT_DESCRIPTOR = CL_INVALID_IMAGE_FORMAT_DESCRIPTOR,
   PI_IMAGE_FORMAT_NOT_SUPPORTED = CL_IMAGE_FORMAT_NOT_SUPPORTED,
   PI_MEM_OBJECT_ALLOCATION_FAILURE = CL_MEM_OBJECT_ALLOCATION_FAILURE,
+  PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE =
+      -998, ///< PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE indicates a fallback
+            ///< method determines the function exists but its address cannot be
+            ///< found.
   PI_ERROR_UNKNOWN = -999
 } _pi_result;
 
@@ -987,6 +991,9 @@ __SYCL_EXPORT pi_result piextDeviceSelectBinary(pi_device device,
 /// must present in the list of devices returned by \c get_device method for
 /// \arg \c program.
 ///
+/// If a fallback method determines the function exists but the address is
+/// not available PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE is returned. If the
+/// address does not exist PI_INVALID_KERNEL_NAME is returned.
 __SYCL_EXPORT pi_result piextGetDeviceFunctionPointer(
     pi_device device, pi_program program, const char *function_name,
     pi_uint64 *function_pointer_ret);

sycl/plugins/cuda/pi_cuda.cpp

@@ -575,20 +575,8 @@ pi_result _pi_program::build_program(const char *build_options) {
 /// Note: Another alternative is to add kernel names as metadata, like with
 ///       reqd_work_group_size.
 std::string getKernelNames(pi_program program) {
-  std::string source(program->binary_,
-                     program->binary_ + program->binarySizeInBytes_);
-  std::regex entries_pattern(".entry\\s+([^\\([:s:]]*)");
-  std::string names("");
-  std::smatch match;
-  bool first_match = true;
-  while (std::regex_search(source, match, entries_pattern)) {
-    assert(match.size() == 2);
-    names += first_match ? "" : ";";
-    names += match[1]; // Second element is the group.
-    source = match.suffix().str();
-    first_match = false;
-  }
-  return names;
+  cl::sycl::detail::pi::die("getKernelNames not implemented");
+  return {};
 }
 
 /// RAII object that calls the reference count release function on the held PI
@@ -921,11 +909,27 @@ pi_result cuda_piextDeviceSelectBinary(pi_device device,
   return PI_INVALID_BINARY;
 }
 
-pi_result cuda_piextGetDeviceFunctionPointer(pi_device, pi_program,
-                                             const char *, pi_uint64 *) {
-  cl::sycl::detail::pi::die(
-      "cuda_piextGetDeviceFunctionPointer not implemented");
-  return {};
+pi_result cuda_piextGetDeviceFunctionPointer(pi_device device,
+                                             pi_program program,
+                                             const char *func_name,
+                                             pi_uint64 *func_pointer_ret) {
+  // Check if device passed is the same the device bound to the context
+  assert(device == program->get_context()->get_device());
+  assert(func_pointer_ret != nullptr);
+
+  CUfunction func;
+  CUresult ret = cuModuleGetFunction(&func, program->get(), func_name);
+  *func_pointer_ret = reinterpret_cast<pi_uint64>(func);
+  pi_result retError = PI_SUCCESS;
+
+  if (ret != CUDA_SUCCESS && ret != CUDA_ERROR_NOT_FOUND)
+    retError = PI_CHECK_ERROR(ret);
+  if (ret == CUDA_ERROR_NOT_FOUND) {
+    *func_pointer_ret = 0;
+    retError = PI_INVALID_KERNEL_NAME;
+  }
+
+  return retError;
 }
 
 /// \return PI_SUCCESS always since CUDA devices are always root devices.

sycl/plugins/hip/pi_hip.cpp

@@ -567,20 +567,8 @@ pi_result _pi_program::build_program(const char *build_options) {
 ///       has_kernel method, so an alternative would be to move the has_kernel
 ///       query to PI and use hipModuleGetFunction to check for a kernel.
 std::string getKernelNames(pi_program program) {
-  std::string source(program->binary_,
-                     program->binary_ + program->binarySizeInBytes_);
-  std::regex entries_pattern(".entry\\s+([^\\([:s:]]*)");
-  std::string names("");
-  std::smatch match;
-  bool first_match = true;
-  while (std::regex_search(source, match, entries_pattern)) {
-    assert(match.size() == 2);
-    names += first_match ? "" : ";";
-    names += match[1]; // Second element is the group.
-    source = match.suffix().str();
-    first_match = false;
-  }
-  return names;
+  cl::sycl::detail::pi::die("getKernelNames not implemented");
+  return {};
 }
 
 /// RAII object that calls the reference count release function on the held PI
@@ -909,16 +897,25 @@ pi_result hip_piextDeviceSelectBinary(pi_device device,
 
 pi_result hip_piextGetDeviceFunctionPointer(pi_device device,
                                             pi_program program,
-                                            const char *function_name,
-                                            pi_uint64 *function_pointer_ret) {
-  (void)device;
-  (void)program;
-  (void)function_name;
-  (void)function_pointer_ret;
+                                            const char *func_name,
+                                            pi_uint64 *func_pointer_ret) {
+  // Check if device passed is the same the device bound to the context
+  assert(device == program->get_context()->get_device());
+  assert(func_pointer_ret != nullptr);
+
+  hipFunction_t func;
+  hipError_t ret = hipModuleGetFunction(&func, program->get(), func_name);
+  *func_pointer_ret = reinterpret_cast<pi_uint64>(func);
+  pi_result retError = PI_SUCCESS;
 
-  cl::sycl::detail::pi::die(
-      "hip_piextGetDeviceFunctionPointer not implemented");
-  return {};
+  if (ret != hipSuccess && ret != hipErrorNotFound)
+    retError = PI_CHECK_ERROR(ret);
+  if (ret == hipErrorNotFound) {
+    *func_pointer_ret = 0;
+    retError = PI_INVALID_KERNEL_NAME;
+  }
+
+  return retError;
 }
 
 /// \return PI_SUCCESS always since HIP devices are always root devices.

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -6248,6 +6248,32 @@ pi_result piEnqueueNativeKernel(pi_queue Queue, void (*UserFunc)(void *),
   return {};
 }
 
+// Function gets characters between delimeter's in str
+// then checks if they are equal to the sub_str.
+// returns true if there is at least one instance
+// returns false if there are no instances of the name
+static bool is_in_separated_string(const std::string &str, char delimiter,
+                                   const std::string &sub_str) {
+  size_t beg = 0;
+  size_t length = 0;
+  for (const auto &x : str) {
+    if (x == delimiter) {
+      if (str.substr(beg, length) == sub_str)
+        return true;
+
+      beg += length + 1;
+      length = 0;
+      continue;
+    }
+    length++;
+  }
+  if (length != 0)
+    if (str.substr(beg, length) == sub_str)
+      return true;
+
+  return false;
+}
+
 // TODO: Check if the function_pointer_ret type can be converted to void**.
 pi_result piextGetDeviceFunctionPointer(pi_device Device, pi_program Program,
                                         const char *FunctionName,
@@ -6272,6 +6298,40 @@ pi_result piextGetDeviceFunctionPointer(pi_device Device, pi_program Program,
     ModIt++;
   }
 
+  // zeModuleGetFunctionPointer currently fails for all
+  // kernels regardless of if the kernel exist or not
+  // with ZE_RESULT_ERROR_INVALID_ARGUMENT
+  // TODO: remove when this is no longer the case
+  // If zeModuleGetFunctionPointer returns invalid argument,
+  // fallback to searching through kernel list and return
+  // PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE if the function exists
+  // or PI_INVALID_KERNEL_NAME if the function does not exist.
+  // FunctionPointerRet should always be 0
+  if (ZeResult == ZE_RESULT_ERROR_INVALID_ARGUMENT) {
+    size_t Size;
+    *FunctionPointerRet = 0;
+    PI_CALL(piProgramGetInfo(Program, PI_PROGRAM_INFO_KERNEL_NAMES, 0, nullptr,
+                             &Size));
+
+    std::string ClResult(Size, ' ');
+    PI_CALL(piProgramGetInfo(Program, PI_PROGRAM_INFO_KERNEL_NAMES,
+                             ClResult.size(), &ClResult[0], nullptr));
+
+    // Get rid of the null terminator and search for kernel_name
+    // If function can be found return error code to indicate it
+    // exists
+    ClResult.pop_back();
+    if (is_in_separated_string(ClResult, ';', std::string(FunctionName)))
+      return PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE;
+
+    return PI_INVALID_KERNEL_NAME;
+  }
+
+  if (ZeResult == ZE_RESULT_ERROR_INVALID_FUNCTION_NAME) {
+    *FunctionPointerRet = 0;
+    return PI_INVALID_KERNEL_NAME;
+  }
+
   return mapError(ZeResult);
 }
 

sycl/plugins/opencl/pi_opencl.cpp

@@ -65,6 +65,8 @@ CONSTFIX char clEnqueueMemcpyName[] = "clEnqueueMemcpyINTEL";
 CONSTFIX char clGetMemAllocInfoName[] = "clGetMemAllocInfoINTEL";
 CONSTFIX char clSetProgramSpecializationConstantName[] =
     "clSetProgramSpecializationConstant";
+CONSTFIX char clGetDeviceFunctionPointerName[] =
+    "clGetDeviceFunctionPointerINTEL";
 
 #undef CONSTFIX
 
@@ -77,8 +79,10 @@ static pi_result getExtFuncFromContext(pi_context context, T *fptr) {
 
   // if cached, return cached FuncPtr
   if (auto F = FuncPtrs[context]) {
+    // if cached that extension is not available return nullptr and
+    // PI_INVALID_VALUE
     *fptr = F;
-    return PI_SUCCESS;
+    return F ? PI_SUCCESS : PI_INVALID_VALUE;
   }
 
   cl_uint deviceCount;
@@ -110,8 +114,11 @@ static pi_result getExtFuncFromContext(pi_context context, T *fptr) {
   T FuncPtr =
       (T)clGetExtensionFunctionAddressForPlatform(curPlatform, FuncName);
 
-  if (!FuncPtr)
+  if (!FuncPtr) {
+    // Cache that the extension is not available
+    FuncPtrs[context] = nullptr;
     return PI_INVALID_VALUE;
+  }
 
   *fptr = FuncPtr;
   FuncPtrs[context] = FuncPtr;
@@ -504,41 +511,97 @@ pi_result piextKernelCreateWithNativeHandle(pi_native_handle nativeHandle,
   return PI_SUCCESS;
 }
 
+// Function gets characters between delimeter's in str
+// then checks if they are equal to the sub_str.
+// returns true if there is at least one instance
+// returns false if there are no instances of the name
+static bool is_in_separated_string(const std::string &str, char delimiter,
+                                   const std::string &sub_str) {
+  size_t beg = 0;
+  size_t length = 0;
+  for (const auto &x : str) {
+    if (x == delimiter) {
+      if (str.substr(beg, length) == sub_str)
+        return true;
+
+      beg += length + 1;
+      length = 0;
+      continue;
+    }
+    length++;
+  }
+  if (length != 0)
+    if (str.substr(beg, length) == sub_str)
+      return true;
+
+  return false;
+}
+
+typedef CL_API_ENTRY cl_int(CL_API_CALL *clGetDeviceFunctionPointer_fn)(
+    cl_device_id device, cl_program program, const char *FuncName,
+    cl_ulong *ret_ptr);
 pi_result piextGetDeviceFunctionPointer(pi_device device, pi_program program,
                                         const char *func_name,
                                         pi_uint64 *function_pointer_ret) {
-  pi_platform platform;
+
+  cl_context CLContext = nullptr;
   cl_int ret_err =
-      clGetDeviceInfo(cast<cl_device_id>(device), PI_DEVICE_INFO_PLATFORM,
-                      sizeof(platform), &platform, nullptr);
+      clGetProgramInfo(cast<cl_program>(program), CL_PROGRAM_CONTEXT,
+                       sizeof(CLContext), &CLContext, nullptr);
 
-  if (ret_err != CL_SUCCESS) {
+  if (ret_err != CL_SUCCESS)
     return cast<pi_result>(ret_err);
-  }
 
-  using FuncT =
-      cl_int(CL_API_CALL *)(cl_device_id, cl_program, const char *, cl_ulong *);
-
-  // TODO: add check that device supports corresponding extension
-  FuncT func_ptr =
-      reinterpret_cast<FuncT>(clGetExtensionFunctionAddressForPlatform(
-          cast<cl_platform_id>(platform), "clGetDeviceFunctionPointerINTEL"));
-  // TODO: once we have check that device supports corresponding extension,
-  // we can insert an assertion that func_ptr is not nullptr. For now, let's
-  // just return an error if failed to query such function
-  // assert(
-  //     func_ptr != nullptr &&
-  //     "Failed to get address of clGetDeviceFunctionPointerINTEL function");
-
-  if (!func_ptr) {
-    if (function_pointer_ret)
+  clGetDeviceFunctionPointer_fn FuncT = nullptr;
+  ret_err = getExtFuncFromContext<clGetDeviceFunctionPointerName,
+                                  clGetDeviceFunctionPointer_fn>(
+      cast<pi_context>(CLContext), &FuncT);
+
+  pi_result pi_ret_err = PI_SUCCESS;
+
+  // Check if kernel name exists, to prevent opencl runtime throwing exception
+  // with cpu runtime
+  // TODO: Use fallback search method if extension does not exist once CPU
+  // runtime no longer throws exceptions and prints messages when given
+  // unavailable functions.
+  *function_pointer_ret = 0;
+  size_t Size;
+  cl_int Res =
+      clGetProgramInfo(cast<cl_program>(program), PI_PROGRAM_INFO_KERNEL_NAMES,
+                       0, nullptr, &Size);
+  if (Res != CL_SUCCESS)
+    return cast<pi_result>(Res);
+
+  std::string ClResult(Size, ' ');
+  ret_err =
+      clGetProgramInfo(cast<cl_program>(program), PI_PROGRAM_INFO_KERNEL_NAMES,
+                       ClResult.size(), &ClResult[0], nullptr);
+  if (Res != CL_SUCCESS)
+    return cast<pi_result>(Res);
+
+  // Get rid of the null terminator and search for kernel_name
+  // If function cannot be found return error code to indicate it
+  // exists
+  ClResult.pop_back();
+  if (!is_in_separated_string(ClResult, ';', func_name))
+    return PI_INVALID_KERNEL_NAME;
+
+  pi_ret_err = PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE;
+
+  // If clGetDeviceFunctionPointer is in list of extensions
+  if (FuncT) {
+    pi_ret_err = cast<pi_result>(FuncT(cast<cl_device_id>(device),
+                                       cast<cl_program>(program), func_name,
+                                       function_pointer_ret));
+    // GPU runtime sometimes returns PI_INVALID_ARG_VALUE if func address cannot
+    // be found even if kernel exits. As the kernel does exist return that the
+    // address is not available
+    if (pi_ret_err == CL_INVALID_ARG_VALUE) {
       *function_pointer_ret = 0;
-    return PI_INVALID_DEVICE;
+      return PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE;
+    }
   }
-
-  return cast<pi_result>(func_ptr(cast<cl_device_id>(device),
-                                  cast<cl_program>(program), func_name,
-                                  function_pointer_ret));
+  return pi_ret_err;
 }
 
 pi_result piContextCreate(const pi_context_properties *properties,

sycl/source/detail/common.cpp

@@ -216,6 +216,8 @@ const char *stringifyErrorCode(cl_int error) {
         case CL_DX9_MEDIA_SURFACE_NOT_ACQUIRED_KHR:
           return "CL_DX9_MEDIA_SURFACE_NOT_ACQUIRED_KHR";
           */
+    case PI_FUNCTION_ADDRESS_IS_NOT_AVAILABLE:
+      return "Function exists but address is not available";
     default:
       return "Unknown OpenCL error code";
   }