check device type by device usm value

Author: AllanZyne

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -605,12 +605,6 @@ SanitizerInterceptor::insertDevice(ur_device_handle_t Device,
 
     DI = std::make_shared<ur_sanitizer_layer::DeviceInfo>(Device);
 
-    // Query device type
-    DI->Type = GetDeviceType(Device);
-    if (DI->Type == DeviceType::UNKNOWN) {
-        return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
-    }
-
     // Query alignment
     UR_CALL(context.urDdiTable.Device.pfnGetInfo(
         Device, UR_DEVICE_INFO_MEM_BASE_ADDR_ALIGN, sizeof(DI->Alignment),
@@ -921,7 +915,7 @@ ur_result_t USMLaunchInfo::updateKernelInfo(const KernelInfo &KI) {
 USMLaunchInfo::~USMLaunchInfo() {
     [[maybe_unused]] ur_result_t Result;
     if (Data) {
-        auto Type = GetDeviceType(Device);
+        auto Type = GetDeviceType(Context, Device);
         if (Type == DeviceType::GPU_PVC || Type == DeviceType::GPU_DG2) {
             if (Data->PrivateShadowOffset) {
                 Result = context.urDdiTable.USM.pfnFree(

source/loader/layers/sanitizer/common.hpp

@@ -138,6 +138,21 @@ struct SourceInfo {
 
 enum class DeviceType : uint64_t { UNKNOWN = 0, CPU, GPU_PVC, GPU_DG2 };
 
+inline const char *ToString(DeviceType Type) {
+    switch (Type) {
+    case DeviceType::UNKNOWN:
+        return "UNKNOWN";
+    case DeviceType::CPU:
+        return "CPU";
+    case DeviceType::GPU_PVC:
+        return "PVC";
+    case DeviceType::GPU_DG2:
+        return "DG2";
+    default:
+        return "UNKNOWN";
+    }
+}
+
 bool IsInASanContext();
 
 uptr MmapNoReserve(uptr Addr, uptr Size);

source/loader/layers/sanitizer/ur_sanddi.cpp

@@ -26,6 +26,13 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
         auto hDevice = phDevices[i];
         std::shared_ptr<DeviceInfo> DI;
         UR_CALL(context.interceptor->insertDevice(hDevice, DI));
+        DI->Type = GetDeviceType(Context, hDevice);
+        if (DI->Type == DeviceType::UNKNOWN) {
+            context.logger.error("Unsupport device");
+            return UR_RESULT_ERROR_INVALID_DEVICE;
+        }
+        context.logger.info("Add {} into context {}", ToString(DI->Type),
+                            (void *)Context);
         if (!DI->ShadowOffset) {
             UR_CALL(DI->allocShadowMemory(Context));
         }

source/loader/layers/sanitizer/ur_sanitizer_utils.cpp

@@ -72,20 +72,6 @@ ur_device_handle_t GetDevice(ur_queue_handle_t Queue) {
     return Device;
 }
 
-std::string GetDeviceName(ur_device_handle_t Device) {
-    size_t Size = 0;
-    [[maybe_unused]] auto Res = context.urDdiTable.Device.pfnGetInfo(
-        Device, UR_DEVICE_INFO_NAME, 0, nullptr, &Size);
-    assert(Res == UR_RESULT_SUCCESS && "GetDeviceName() failed");
-
-    std::vector<char> NameBuf(Size);
-    Res = context.urDdiTable.Device.pfnGetInfo(Device, UR_DEVICE_INFO_NAME,
-                                               Size, NameBuf.data(), nullptr);
-    assert(Res == UR_RESULT_SUCCESS);
-
-    return std::string(NameBuf.data(), Size - 1);
-}
-
 ur_program_handle_t GetProgram(ur_kernel_handle_t Kernel) {
     ur_program_handle_t Program{};
     [[maybe_unused]] auto Result = context.urDdiTable.Kernel.pfnGetInfo(
@@ -129,7 +115,8 @@ ur_device_handle_t GetUSMAllocDevice(ur_context_handle_t Context,
     return Device;
 }
 
-DeviceType GetDeviceType(ur_device_handle_t Device) {
+DeviceType GetDeviceType(ur_context_handle_t Context,
+                         ur_device_handle_t Device) {
     ur_device_type_t DeviceType = UR_DEVICE_TYPE_DEFAULT;
     [[maybe_unused]] auto Result = context.urDdiTable.Device.pfnGetInfo(
         Device, UR_DEVICE_INFO_TYPE, sizeof(DeviceType), &DeviceType, nullptr);
@@ -140,16 +127,24 @@ DeviceType GetDeviceType(ur_device_handle_t Device) {
         // TODO: Check fpga is fpga emulator
         return DeviceType::CPU;
     case UR_DEVICE_TYPE_GPU: {
-        // Ref: https://github.com/intel/compute-runtime/blob/master/shared/source/dll/devices/devices_base.inl
-        auto Name = GetDeviceName(Device);
-        if (Name.rfind("Intel(R) Data Center GPU Max", 0) == 0) {
-            return DeviceType::GPU_PVC;
-        } else if ((Name.rfind("Intel(R) Arc(TM)", 0) == 0 &&
-                    Name != "Intel(R) Arc(TM) Graphics") ||
-                   Name.rfind("Intel(R) Data Center GPU Flex", 0) == 0) {
-            return DeviceType::GPU_DG2;
+        uptr Ptr;
+        ur_result_t Result = context.urDdiTable.USM.pfnDeviceAlloc(
+            Context, Device, nullptr, nullptr, 4, (void **)&Ptr);
+        context.logger.debug("GetDeviceType: {}", (void *)Ptr);
+        assert(Result == UR_RESULT_SUCCESS &&
+               "getDeviceType() failed at allocating device USM");
+        // FIXME: There's no API querying the address bits of device, so we guess it by the
+        // value of device USM pointer (see "USM Allocation Range" in asan_shadow_setup.cpp)
+        auto Type = DeviceType::UNKNOWN;
+        if (Ptr >> 48 == 0xff00U) {
+            Type = DeviceType::GPU_PVC;
+        } else {
+            Type = DeviceType::GPU_DG2;
         }
-        return DeviceType::UNKNOWN;
+        Result = context.urDdiTable.USM.pfnFree(Context, (void *)Ptr);
+        assert(Result == UR_RESULT_SUCCESS &&
+               "getDeviceType() failed at releasing device USM");
+        return Type;
     }
     default:
         return DeviceType::UNKNOWN;

source/loader/layers/sanitizer/ur_sanitizer_utils.hpp

@@ -34,8 +34,8 @@ ur_context_handle_t GetContext(ur_queue_handle_t Queue);
 ur_context_handle_t GetContext(ur_program_handle_t Program);
 ur_context_handle_t GetContext(ur_kernel_handle_t Kernel);
 ur_device_handle_t GetDevice(ur_queue_handle_t Queue);
-std::string GetDeviceName(ur_device_handle_t Device);
-DeviceType GetDeviceType(ur_device_handle_t Device);
+DeviceType GetDeviceType(ur_context_handle_t Context,
+                         ur_device_handle_t Device);
 std::string GetKernelName(ur_kernel_handle_t Kernel);
 size_t GetDeviceLocalMemorySize(ur_device_handle_t Device);
 ur_program_handle_t GetProgram(ur_kernel_handle_t Kernel);