Fix Offload build on CUDA and detect correct targets in the CTS

Author: callumfare

source/adapters/offload/CMakeLists.txt

@@ -18,8 +18,8 @@ if (NOT TARGET cudadrv)
         add_library(cudadrv SHARED IMPORTED GLOBAL)
         set_target_properties(
                 cudadrv PROPERTIES 
-                IMPORTED_LOCATION             ${CUDA_CUDA_LIBRARY}
-                INTERFACE_INCLUDE_DIRECTORIES ${CUDA_INCLUDE_DIRS}
+                IMPORTED_LOCATION             ${CUDA_cuda_driver_LIBRARY}
+                INTERFACE_INCLUDE_DIRECTORIES ${CUDAToolkit_INCLUDE_DIRS}
         )
 endif()
 
@@ -44,20 +44,22 @@ set_target_properties(${TARGET_NAME} PROPERTIES
         SOVERSION "${PROJECT_VERSION_MAJOR}"
 )
 
+set(ADDITIONAL_LINK_LIBS "")
+if (CUDA_cuda_driver_LIBRARY)
+    list(APPEND ADDITIONAL_LINK_LIBS
+        cudadrv
+    )
+    target_compile_definitions(${TARGET_NAME} PRIVATE UR_CUDA_ENABLED)
+endif()
+
 target_link_libraries(${TARGET_NAME} PRIVATE
         ${PROJECT_NAME}::headers
         ${PROJECT_NAME}::common
         ${PROJECT_NAME}::umf
         ${UR_OFFLOAD_INSTALL_DIR}/lib/libLLVMOffload.so
+        ${ADDITIONAL_LINK_LIBS}
 )
 
-if (CUDA_CUDA_LIBRARY)
-    target_link_libraries(${TARGET_NAME}
-        cudadrv
-    )
-    target_compile_definitions(${TARGET_NAME} PRIVATE UR_CUDA_ENABLED=1)
-endif()
-
 target_include_directories(${TARGET_NAME} PRIVATE
         "${UR_OFFLOAD_INCLUDE_DIR}/offload"
         "${CMAKE_CURRENT_SOURCE_DIR}/../../"

source/adapters/offload/device.cpp

@@ -119,15 +119,20 @@ urDevicePartition(ur_device_handle_t, const ur_device_partition_properties_t *,
 UR_APIEXPORT ur_result_t UR_APICALL urDeviceSelectBinary(
     ur_device_handle_t hDevice, const ur_device_binary_t *pBinaries,
     uint32_t NumBinaries, uint32_t *pSelectedBinary) {
-  std::ignore = hDevice;
-  std::ignore = pBinaries;
-  std::ignore = NumBinaries;
-  std::ignore = pSelectedBinary;
 
-  // TODO: Don't hard code nvptx64!
-  const char *image_target = UR_DEVICE_BINARY_TARGET_NVPTX64;
+  ol_platform_backend_t Backend;
+  olGetPlatformInfo(hDevice->Platform->OffloadPlatform,
+                    OL_PLATFORM_INFO_BACKEND, sizeof(Backend), &Backend);
+
+  const char *ImageTarget = UR_DEVICE_BINARY_TARGET_UNKNOWN;
+  if (Backend == OL_PLATFORM_BACKEND_CUDA) {
+    ImageTarget = UR_DEVICE_BINARY_TARGET_NVPTX64;
+  } else if (Backend == OL_PLATFORM_BACKEND_AMDGPU) {
+    ImageTarget = UR_DEVICE_BINARY_TARGET_AMDGCN;
+  }
+
   for (uint32_t i = 0; i < NumBinaries; ++i) {
-    if (strcmp(pBinaries[i].pDeviceTargetSpec, image_target) == 0) {
+    if (strcmp(pBinaries[i].pDeviceTargetSpec, ImageTarget) == 0) {
       *pSelectedBinary = i;
       return UR_RESULT_SUCCESS;
     }

source/adapters/offload/program.cpp

@@ -4,6 +4,7 @@
 
 #include "context.hpp"
 #include "device.hpp"
+#include "platform.hpp"
 #include "program.hpp"
 #include "ur2offload.hpp"
 
@@ -31,15 +32,17 @@ ur_result_t ProgramCreateCudaWorkaround(ur_context_handle_t hContext,
   cuLinkComplete(State, &CuBin, &CuBinSize);
   RealBinary = (uint8_t *)CuBin;
   RealLength = CuBinSize;
+
+#if 0
   fprintf(stderr, "Performed CUDA bin workaround (size = %lu)\n", RealLength);
+#endif
 
   ur_program_handle_t Program = new ur_program_handle_t_();
   auto Res = olCreateProgram(hContext->Device->OffloadDevice, RealBinary,
                              RealLength, &Program->OffloadProgram);
 
   // Program owns the linked module now
   cuLinkDestroy(State);
-  (void)State;
 
   if (Res != OL_SUCCESS) {
     delete Program;
@@ -146,13 +149,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(
     return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   }
 
-  ur_platform_handle_t DevicePlatform;
-  urDeviceGetInfo(phDevices[0], UR_DEVICE_INFO_PLATFORM,
-                  sizeof(ur_platform_handle_t), &DevicePlatform, nullptr);
-  ur_platform_backend_t PlatformBackend;
-  urPlatformGetInfo(DevicePlatform, UR_PLATFORM_INFO_BACKEND,
-                    sizeof(ur_platform_backend_t), &PlatformBackend, nullptr);
-
   auto *RealBinary = ppBinaries[0];
   size_t RealLength = pLengths[0];
 
@@ -171,7 +167,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithBinary(
     }
   }
 
-  if (PlatformBackend == UR_PLATFORM_BACKEND_CUDA) {
+  ol_platform_backend_t Backend;
+  olGetPlatformInfo(phDevices[0]->Platform->OffloadPlatform,
+                    OL_PLATFORM_INFO_BACKEND, sizeof(Backend), &Backend);
+  if (Backend == OL_PLATFORM_BACKEND_CUDA) {
     return ProgramCreateCudaWorkaround(hContext, RealBinary, RealLength,
                                        phProgram);
   }

test/conformance/source/environment.cpp

@@ -216,14 +216,23 @@ std::string KernelsEnvironment::getTargetName(ur_platform_handle_t platform) {
   case UR_PLATFORM_BACKEND_HIP:
     return "amdgcn-amd-amdhsa";
   case UR_PLATFORM_BACKEND_OFFLOAD: {
-    // TODO: In future this should use urDeviceSelectBinary
-    auto result = ur_getenv("UR_OFFLOAD_TARGET_NAME");
-    if (!result) {
-      error = "For offload testing, please specify a target in "
-              "`UR_OFFLOAD_TARGET_NAME`";
+    // All Offload platforms report this backend, use the platform name to select
+    // the actual underlying backend.
+    std::vector<char> PlatformName;
+    size_t PlatformNameSize = 0;
+    urPlatformGetInfo(platform, UR_PLATFORM_INFO_NAME, 0, nullptr,
+                      &PlatformNameSize);
+    PlatformName.resize(PlatformNameSize);
+    urPlatformGetInfo(platform, UR_PLATFORM_INFO_NAME, PlatformNameSize,
+                      PlatformName.data(), nullptr);
+    if (std::strcmp(PlatformName.data(), "CUDA") == 0) {
+      return "nvptx64-nvidia-cuda";
+    } else if (std::strcmp(PlatformName.data(), "AMDGPU") == 0) {
+      return "amdgcn-amd-amdhsa";
+    } else {
+      error = "Could not detect target for Offload platform";
       return {};
     }
-    return *result;
   }
   case UR_PLATFORM_BACKEND_NATIVE_CPU:
     error = "native_cpu doesn't support kernel tests yet";