[SYCL] Only use AOT images built for the device's architecture

sycl/source/detail/program_manager/program_manager.cpp

@@ -547,8 +547,24 @@ static const char *getUrDeviceTarget(const char *URDeviceTarget) {
   return UR_DEVICE_BINARY_TARGET_UNKNOWN;
 }
 
+// An image with a compile_target was built for that device architecture only,
+// so it only matches a device of that architecture. Images without a
+// compile_target are not filtered here.
+static bool doesCompileTargetMatchDevice(const RTDeviceBinaryImage &Img,
+                                         const device_impl &Dev) {
+  const RTDeviceBinaryImage::PropertyRange &Range = Img.getDeviceRequirements();
+  const bool HasCompileTarget =
+      std::any_of(Range.begin(), Range.end(), [](const auto &Prop) {
+        return Prop->Name == std::string_view("compile_target");
+      });
+  return !HasCompileTarget || doesImageTargetMatchDevice(Img, Dev);
+}
+
 static bool compatibleWithDevice(const RTDeviceBinaryImage *BinImage,
                                  const device_impl &DeviceImpl) {
+  if (!doesCompileTargetMatchDevice(*BinImage, DeviceImpl))
+    return false;
+
   adapter_impl &Adapter = DeviceImpl.getAdapter();
 
   const ur_device_handle_t &URDeviceHandle = DeviceImpl.getHandleRef();
@@ -791,7 +807,8 @@ ProgramManager::collectDependentDeviceImagesForVirtualFunctions(
         // virtual function set and it is incompatible, then we should link its
         // "dummy" version to avoid link errors about unresolved external
         // symbols.
-        if (doesDevSupportDeviceRequirements(Dev, *BinImage))
+        if (doesDevSupportDeviceRequirements(Dev, *BinImage) &&
+            doesCompileTargetMatchDevice(*BinImage, Dev))
           DeviceImagesToLink.insert(BinImage);
       }
     }
@@ -1442,11 +1459,20 @@ const RTDeviceBinaryImage &ProgramManager::getDeviceImage(
   }
 
   std::lock_guard<std::mutex> ImgMapsGuard(m_ImgMapsMutex);
-  std::vector<sycl_device_binary> RawImgs(ImageSet.size());
-  auto ImageIterator = ImageSet.begin();
-  for (size_t i = 0; i < ImageSet.size(); i++, ImageIterator++)
+  // Skip images built for another device architecture.
+  std::vector<const RTDeviceBinaryImage *> Imgs;
+  for (const RTDeviceBinaryImage *Img : ImageSet)
+    if (doesCompileTargetMatchDevice(*Img, DeviceImpl))
+      Imgs.push_back(Img);
+  if (Imgs.empty())
+    throw exception(make_error_code(errc::runtime),
+                    "No device image is built for the selected device's "
+                    "architecture");
+
+  std::vector<sycl_device_binary> RawImgs(Imgs.size());
+  for (size_t i = 0; i < Imgs.size(); i++)
     RawImgs[i] = reinterpret_cast<sycl_device_binary>(
-        const_cast<sycl_device_binary>(&(*ImageIterator)->getRawData()));
+        const_cast<sycl_device_binary>(&Imgs[i]->getRawData()));
   uint32_t ImgInd = 0;
   // Ask the native runtime under the given context to choose the device image
   // it prefers.
@@ -1460,15 +1486,14 @@ const RTDeviceBinaryImage &ProgramManager::getDeviceImage(
   ContextImpl.getAdapter().call<UrApiKind::urDeviceSelectBinary>(
       DeviceImpl.getHandleRef(), UrBinaries.data(), UrBinaries.size(), &ImgInd);
 
-  ImageIterator = ImageSet.begin();
-  std::advance(ImageIterator, ImgInd);
+  const RTDeviceBinaryImage *SelectedImg = Imgs[ImgInd];
 
   if constexpr (DbgProgMgr > 0) {
-    std::cerr << "selected device image: " << &(*ImageIterator)->getRawData()
+    std::cerr << "selected device image: " << &SelectedImg->getRawData()
               << "\n";
-    (*ImageIterator)->print();
+    SelectedImg->print();
   }
-  return **ImageIterator;
+  return *SelectedImg;
 }
 
 // Check if device image is compressed.

sycl/unittests/SYCL2020/SYCLBINSerialize.cpp

@@ -62,6 +62,16 @@ sycl::unittest::MockDeviceImage Imgs[] = {
     makeImageWithCompileTarget("SYCLBINSerializeKernel", "intel_gpu_pvc")};
 sycl::unittest::MockDeviceImageArray<std::size(Imgs)> ImgArray{Imgs};
 
+// Makes the mock device a PVC GPU, so that the intel_gpu_pvc image is used.
+ur_result_t redefinedDeviceGetInfoPVC(void *pParams) {
+  auto Params = *static_cast<ur_device_get_info_params_t *>(pParams);
+  if (*Params.ppropName == UR_DEVICE_INFO_IP_VERSION && *Params.ppPropValue)
+    *static_cast<uint32_t *>(*Params.ppPropValue) = 0x030f0000;
+  if (*Params.ppropName == UR_DEVICE_INFO_TYPE && *Params.ppPropValue)
+    *static_cast<ur_device_type_t *>(*Params.ppPropValue) = UR_DEVICE_TYPE_GPU;
+  return UR_RESULT_SUCCESS;
+}
+
 // Strips the OffloadBinary wrapper that ext_oneapi_get_content emits and
 // returns a string_view over the contained raw SYCLBIN bytes. Returns
 // std::nullopt if the wrapper does not parse cleanly; callers should ASSERT
@@ -115,6 +125,8 @@ stripOffloadWrapper(const std::vector<char> &Bytes) {
 // confirm the property survived.
 TEST(SYCLBINSerialize, CompileTargetRoundTrip) {
   sycl::unittest::UrMock<> Mock;
+  mock::getCallbacks().set_after_callback("urDeviceGetInfo",
+                                          &redefinedDeviceGetInfoPVC);
 
   const sycl::device Dev = sycl::platform().get_devices()[0];
   sycl::context Ctx{Dev};
@@ -222,6 +234,8 @@ TEST(SYCLBINSerialize, BundleStateRoundTripObject) {
 // inspection alone (e.g. CompileTargetRoundTrip) does not catch.
 TEST(SYCLBINSerialize, ImageBytesRoundTrip) {
   sycl::unittest::UrMock<> Mock;
+  mock::getCallbacks().set_after_callback("urDeviceGetInfo",
+                                          &redefinedDeviceGetInfoPVC);
 
   const sycl::device Dev = sycl::platform().get_devices()[0];
   sycl::context Ctx{Dev};

sycl/unittests/program_manager/CMakeLists.txt

@@ -21,3 +21,4 @@ endif()
 
 add_subdirectory(arg_mask)
 add_subdirectory(DynamicLinking)
+add_subdirectory(CompileTargetSelection)

sycl/unittests/program_manager/CompileTargetSelection/CMakeLists.txt

@@ -0,0 +1,5 @@
+# Uses its own test binary: the registered images decide which devices have a
+# compatible image, so other tests' images must not be visible.
+add_sycl_unittest(CompileTargetSelectionTests OBJECT
+  CompileTargetSelection.cpp
+)

sycl/unittests/program_manager/CompileTargetSelection/CompileTargetSelection.cpp

@@ -0,0 +1,103 @@
+#include <sycl/sycl.hpp>
+
+#include <helpers/MockDeviceImage.hpp>
+#include <helpers/MockKernelInfo.hpp>
+#include <helpers/UrMock.hpp>
+
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <iterator>
+#include <string>
+#include <utility>
+#include <vector>
+
+// Check that an AOT image built for one GPU architecture only makes devices of
+// that architecture preferred by the default device selector.
+
+class CompileTargetSelectionKernel;
+
+namespace sycl {
+inline namespace _V1 {
+namespace detail {
+template <>
+struct KernelInfo<CompileTargetSelectionKernel>
+    : public unittest::MockKernelInfoBase {
+  static constexpr const char *getName() {
+    return "CompileTargetSelectionKernel";
+  }
+};
+} // namespace detail
+} // namespace _V1
+} // namespace sycl
+
+namespace {
+namespace exp = sycl::ext::oneapi::experimental;
+
+sycl::unittest::MockDeviceImage generatePVCImage() {
+  std::string CompileTarget = "intel_gpu_pvc";
+  std::vector<char> Data(8 + CompileTarget.size());
+  std::copy(CompileTarget.begin(), CompileTarget.end(), Data.data() + 8);
+  sycl::unittest::MockPropertySet PropSet;
+  PropSet.insert(__SYCL_PROPERTY_SET_SYCL_DEVICE_REQUIREMENTS,
+                 sycl::unittest::MockProperty("compile_target", Data,
+                                              SYCL_PROPERTY_TYPE_BYTE_ARRAY));
+  return sycl::unittest::MockDeviceImage{
+      SYCL_DEVICE_BINARY_TYPE_NATIVE,
+      __SYCL_DEVICE_BINARY_TARGET_SPIRV64_GEN,
+      "", // Compile options
+      "", // Link options
+      std::vector<unsigned char>{1},
+      sycl::unittest::makeEmptyKernels({"CompileTargetSelectionKernel"}),
+      std::move(PropSet)};
+}
+sycl::unittest::MockDeviceImage Imgs[] = {generatePVCImage()};
+sycl::unittest::MockDeviceImageArray<1> ImgArray{Imgs};
+
+// IP versions of a TGLLP and a PVC GPU. The TGLLP one comes first so that it
+// would be chosen if both devices were preferred equally.
+uint32_t MockDeviceIPs[] = {0x03000000, 0x030f0000};
+
+ur_result_t redefinedDeviceGet(void *pParams) {
+  auto Params = *static_cast<ur_device_get_params_t *>(pParams);
+  if (*Params.ppNumDevices) {
+    **Params.ppNumDevices = std::size(MockDeviceIPs);
+    return UR_RESULT_SUCCESS;
+  }
+  if (*Params.pphDevices)
+    for (uint32_t I = 0; I < *Params.pNumEntries; ++I)
+      (*Params.pphDevices)[I] =
+          reinterpret_cast<ur_device_handle_t>(&MockDeviceIPs[I]);
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t redefinedDeviceGetInfo(void *pParams) {
+  auto Params = *static_cast<ur_device_get_info_params_t *>(pParams);
+  if (*Params.ppropName == UR_DEVICE_INFO_IP_VERSION && *Params.ppPropValue)
+    *static_cast<uint32_t *>(*Params.ppPropValue) =
+        *reinterpret_cast<uint32_t *>(*Params.phDevice);
+  if (*Params.ppropName == UR_DEVICE_INFO_TYPE && *Params.ppPropValue)
+    *static_cast<ur_device_type_t *>(*Params.ppPropValue) = UR_DEVICE_TYPE_GPU;
+  return UR_RESULT_SUCCESS;
+}
+} // namespace
+
+TEST(CompileTargetSelection, DefaultSelectorPrefersDeviceWithImage) {
+  sycl::unittest::UrMock<sycl::backend::ext_oneapi_level_zero> Mock;
+  mock::getCallbacks().set_replace_callback("urDeviceGet", &redefinedDeviceGet);
+  mock::getCallbacks().set_after_callback("urDeviceGetInfo",
+                                          &redefinedDeviceGetInfo);
+
+  std::vector<sycl::device> Devs = sycl::platform().get_devices();
+  ASSERT_EQ(Devs.size(), 2u);
+  ASSERT_EQ(Devs[0].get_info<exp::info::device::architecture>(),
+            exp::architecture::intel_gpu_tgllp);
+  ASSERT_EQ(Devs[1].get_info<exp::info::device::architecture>(),
+            exp::architecture::intel_gpu_pvc);
+  EXPECT_LT(sycl::default_selector_v(Devs[0]),
+            sycl::default_selector_v(Devs[1]));
+
+  sycl::device Dev{sycl::default_selector_v};
+  EXPECT_EQ(Dev.get_info<exp::info::device::architecture>(),
+            exp::architecture::intel_gpu_pvc);
+}

sycl/unittests/program_manager/DynamicLinking/CMakeLists.txt

@@ -4,5 +4,6 @@
 if (NOT LLVM_LIBCXX_USED)
 add_sycl_unittest(DynamicLinkingTests OBJECT
   DynamicLinking.cpp
+  DynamicLinkingCompileTarget.cpp
 )
 endif()