[DevSAN] Change DeviceType to specialization constant

libdevice/include/sanitizer_defs.hpp

@@ -61,6 +61,9 @@ extern SYCL_EXTERNAL __attribute__((convergent)) void
 __spirv_ControlBarrier(int32_t Execution, int32_t Memory,
                        int32_t Semantics) noexcept;
 
+template <typename T>
+extern SYCL_EXTERNAL T __spirv_SpecConstant(int ID, T default_value) noexcept;
+
 extern "C" SYCL_EXTERNAL void __devicelib_exit();
 
 #endif // __SPIR__ || __SPIRV__

libdevice/include/sanitizer_utils.hpp

@@ -7,8 +7,8 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "sanitizer_common/sanitizer_libdevice.hpp"
 #include "sanitizer_defs.hpp"
-#include "spirv_vars.h"
 
 #if defined(__SPIR__) || defined(__SPIRV__)
 
@@ -49,6 +49,11 @@ inline __SYCL_PRIVATE__ void *ToPrivate(void *ptr) {
   return __spirv_GenericCastToPtrExplicit_ToPrivate(ptr, 7);
 }
 
+inline DeviceType GetDeviceTy() {
+  return static_cast<DeviceType>(
+      __spirv_SpecConstant(SpecConstantDeviceTyID, 0));
+}
+
 template <typename T> T Memset(T ptr, int value, size_t size) {
   for (size_t i = 0; i < size; i++) {
     ptr[i] = value;

libdevice/sanitizer/asan_rtl.cpp

@@ -264,16 +264,15 @@ inline uptr MemToShadow(uptr addr, uint32_t as,
 #elif defined(__LIBDEVICE_DG2__)
   shadow_ptr = MemToShadow_DG2(addr, as, debug);
 #else
-  auto launch_info = (__SYCL_GLOBAL__ const AsanRuntimeData *)__AsanLaunchInfo;
-  if (launch_info->DeviceTy == DeviceType::CPU) {
+  if (GetDeviceTy() == DeviceType::CPU) {
     shadow_ptr = MemToShadow_CPU(addr);
-  } else if (launch_info->DeviceTy == DeviceType::GPU_PVC) {
+  } else if (GetDeviceTy() == DeviceType::GPU_PVC) {
     shadow_ptr = MemToShadow_PVC(addr, as, debug);
-  } else if (launch_info->DeviceTy == DeviceType::GPU_DG2) {
+  } else if (GetDeviceTy() == DeviceType::GPU_DG2) {
     shadow_ptr = MemToShadow_DG2(addr, as, debug);
   } else {
     ASAN_DEBUG(__spirv_ocl_printf(__asan_print_unsupport_device_type,
-                                  (int)launch_info->DeviceTy));
+                                  (int)GetDeviceTy()));
     ReportUnknownDevice(debug);
     return 0;
   }
@@ -889,8 +888,7 @@ DEVICE_EXTERN_C_NOINLINE void __asan_set_shadow_private(uptr shadow, uptr size,
 static __SYCL_CONSTANT__ const char __asan_print_private_base[] =
     "[kernel] set_private_base: %llu -> %p\n";
 
-DEVICE_EXTERN_C_NOINLINE void
-__asan_set_private_base(__SYCL_PRIVATE__ void *ptr) {
+inline void SetPrivateBaseImpl(__SYCL_PRIVATE__ void *ptr) {
   auto launch_info = (__SYCL_GLOBAL__ const AsanRuntimeData *)__AsanLaunchInfo;
   const size_t sid = SubGroupLinearId();
   if (!launch_info || sid >= ASAN_MAX_SG_PRIVATE ||
@@ -904,4 +902,17 @@ __asan_set_private_base(__SYCL_PRIVATE__ void *ptr) {
   SubGroupBarrier();
 }
 
+DEVICE_EXTERN_C_NOINLINE void
+__asan_set_private_base(__SYCL_PRIVATE__ void *ptr) {
+#if defined(__LIBDEVICE_CPU__)
+  return;
+#elif defined(__LIBDEVICE_DG2__) || defined(__LIBDEVICE_PVC__)
+  SetPrivateBaseImpl(ptr);
+#else
+  if (GetDeviceTy() == DeviceType::CPU)
+    return;
+  SetPrivateBaseImpl(ptr);
+#endif
+}
+
 #endif // __SPIR__ || __SPIRV__

libdevice/sanitizer/msan_rtl.cpp

@@ -212,16 +212,16 @@ inline uptr MemToShadow(uptr addr, uint32_t as) {
 #elif defined(__LIBDEVICE_CPU__)
   shadow_ptr = MemToShadow_CPU(addr);
 #else
-  if (LIKELY(GetMsanLaunchInfo->DeviceTy == DeviceType::CPU)) {
+  if (GetDeviceTy() == DeviceType::CPU) {
     shadow_ptr = MemToShadow_CPU(addr);
-  } else if (GetMsanLaunchInfo->DeviceTy == DeviceType::GPU_PVC) {
+  } else if (GetDeviceTy() == DeviceType::GPU_PVC) {
     shadow_ptr = MemToShadow_PVC(addr, as);
-  } else if (GetMsanLaunchInfo->DeviceTy == DeviceType::GPU_DG2) {
+  } else if (GetDeviceTy() == DeviceType::GPU_DG2) {
     shadow_ptr = MemToShadow_DG2(addr, as);
   } else {
     shadow_ptr = GetMsanLaunchInfo->CleanShadow;
-    MSAN_DEBUG(__spirv_ocl_printf(__msan_print_unsupport_device_type,
-                                  GetMsanLaunchInfo->DeviceTy));
+    MSAN_DEBUG(
+        __spirv_ocl_printf(__msan_print_unsupport_device_type, GetDeviceTy()));
   }
 #endif
 
@@ -269,17 +269,17 @@ inline uptr MemToOrigin(uptr addr, uint32_t as) {
 #elif defined(__LIBDEVICE_CPU__)
   origin_ptr = MemToOrigin_CPU(addr);
 #else
-  if (LIKELY(GetMsanLaunchInfo->DeviceTy == DeviceType::CPU)) {
+  if (GetDeviceTy() == DeviceType::CPU) {
     origin_ptr = MemToOrigin_CPU(aligned_addr);
-  } else if (GetMsanLaunchInfo->DeviceTy == DeviceType::GPU_PVC) {
+  } else if (GetDeviceTy() == DeviceType::GPU_PVC) {
     origin_ptr = MemToOrigin_PVC(aligned_addr, as);
-  } else if (GetMsanLaunchInfo->DeviceTy == DeviceType::GPU_DG2) {
+  } else if (GetDeviceTy() == DeviceType::GPU_DG2) {
     origin_ptr = MemToOrigin_DG2(aligned_addr, as);
   } else {
     // Return clean shadow (0s) by default
     origin_ptr = GetMsanLaunchInfo->CleanShadow;
-    MSAN_DEBUG(__spirv_ocl_printf(__msan_print_unsupport_device_type,
-                                  GetMsanLaunchInfo->DeviceTy));
+    MSAN_DEBUG(
+        __spirv_ocl_printf(__msan_print_unsupport_device_type, GetDeviceTy()));
   }
 #endif
 
@@ -719,8 +719,7 @@ DEVICE_EXTERN_C_NOINLINE void __msan_unpoison_stack(__SYCL_PRIVATE__ void *ptr,
 static __SYCL_CONSTANT__ const char __msan_print_private_base[] =
     "[kernel] __msan_set_private_base(sid=%llu): %p\n";
 
-DEVICE_EXTERN_C_NOINLINE void
-__msan_set_private_base(__SYCL_PRIVATE__ void *ptr) {
+inline void SetPrivateBaseImpl(__SYCL_PRIVATE__ void *ptr) {
   const size_t sid = SubGroupLinearId();
   if (!GetMsanLaunchInfo || sid >= MSAN_MAX_SG_PRIVATE ||
       GetMsanLaunchInfo->PrivateShadowOffset == 0 ||
@@ -734,6 +733,19 @@ __msan_set_private_base(__SYCL_PRIVATE__ void *ptr) {
   SubGroupBarrier();
 }
 
+DEVICE_EXTERN_C_NOINLINE void
+__msan_set_private_base(__SYCL_PRIVATE__ void *ptr) {
+#if defined(__LIBDEVICE_CPU__)
+  return;
+#elif defined(__LIBDEVICE_DG2__) || defined(__LIBDEVICE_PVC__)
+  SetPrivateBaseImpl(ptr);
+#else
+  if (GetDeviceTy() == DeviceType::CPU)
+    return;
+  SetPrivateBaseImpl(ptr);
+#endif
+}
+
 static __SYCL_CONSTANT__ const char __msan_print_strided_copy_unsupport_type[] =
     "[kernel] __msan_unpoison_strided_copy: unsupported type(%d)\n";
 

libdevice/sanitizer/tsan_rtl.cpp

@@ -127,13 +127,13 @@ inline __SYCL_GLOBAL__ RawShadow *MemToShadow(uptr addr, uint32_t as) {
 #elif defined(__LIBDEVICE_PVC__)
   shadow_ptr = MemToShadow_PVC(addr, as);
 #else
-  if (TsanLaunchInfo->DeviceTy == DeviceType::CPU) {
+  if (GetDeviceTy() == DeviceType::CPU) {
     shadow_ptr = MemToShadow_CPU(addr, as);
-  } else if (TsanLaunchInfo->DeviceTy == DeviceType::GPU_PVC) {
+  } else if (GetDeviceTy() == DeviceType::GPU_PVC) {
     shadow_ptr = MemToShadow_PVC(addr, as);
   } else {
     TSAN_DEBUG(__spirv_ocl_printf(__tsan_print_unsupport_device_type,
-                                  (int)TsanLaunchInfo->DeviceTy));
+                                  (int)GetDeviceTy()));
     return nullptr;
   }
 #endif
@@ -186,10 +186,16 @@ inline void DoReportRace(__SYCL_GLOBAL__ RawShadow *s, AccessType type,
         return;
       }
 
-      if (as == ADDRESS_SPACE_GENERIC &&
-          TsanLaunchInfo->DeviceTy != DeviceType::CPU) {
+#if defined(__LIBDEVICE_CPU__)
+#elif defined(__LIBDEVICE_DG2__) || defined(__LIBDEVICE_PVC__)
+      if (as == ADDRESS_SPACE_GENERIC) {
         ConvertGenericPointer(addr, as);
       }
+#else
+      if (as == ADDRESS_SPACE_GENERIC && GetDeviceTy() != DeviceType::CPU) {
+        ConvertGenericPointer(addr, as);
+      }
+#endif
 
       // Check if current address already being recorded before.
       for (uint32_t i = 0; i < TsanLaunchInfo->RecordedReportCount; i++) {
@@ -467,7 +473,7 @@ DEVICE_EXTERN_C_NOINLINE void __tsan_cleanup_private(uptr addr, size_t size) {
 #elif defined(__LIBDEVICE_PVC__)
   return;
 #else
-  if (TsanLaunchInfo->DeviceTy != DeviceType::CPU)
+  if (GetDeviceTy() != DeviceType::CPU)
     return;
 
   __tsan_cleanup_private_cpu_impl(addr, size);

unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp

@@ -643,6 +643,19 @@ ur_result_t AsanInterceptor::insertProgram(ur_program_handle_t Program) {
   if (m_ProgramMap.find(Program) != m_ProgramMap.end()) {
     return UR_RESULT_SUCCESS;
   }
+  auto CI = getContextInfo(GetContext(Program));
+  auto DI = getDeviceInfo(CI->DeviceList[0]);
+  ur_specialization_constant_info_t SpecConstantInfo{
+      SpecConstantDeviceTyID, sizeof(DeviceType), &DI->Type};
+  ur_result_t URes =
+      getContext()->urDdiTable.Program.pfnSetSpecializationConstants(
+          Program, 1, &SpecConstantInfo);
+  if (URes != UR_RESULT_SUCCESS) {
+    UR_LOG_L(getContext()->logger, DEBUG,
+             "Set specilization constant for device type failed: {}, the "
+             "program may not be sanitized or is created from binary.",
+             URes);
+  }
   m_ProgramMap.emplace(Program, std::make_shared<ProgramInfo>(Program));
   return UR_RESULT_SUCCESS;
 }
@@ -812,7 +825,6 @@ ur_result_t AsanInterceptor::prepareLaunch(
   // Prepare asan runtime data
   LaunchInfo.Data.Host.GlobalShadowOffset = DeviceInfo->Shadow->ShadowBegin;
   LaunchInfo.Data.Host.GlobalShadowOffsetEnd = DeviceInfo->Shadow->ShadowEnd;
-  LaunchInfo.Data.Host.DeviceTy = DeviceInfo->Type;
   LaunchInfo.Data.Host.Debug = getContext()->Options.Debug ? 1 : 0;
 
   // Write shadow memory offset for local memory
@@ -874,16 +886,14 @@ ur_result_t AsanInterceptor::prepareLaunch(
 
   UR_LOG_L(getContext()->logger, INFO,
            "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, PrivateBase={}, "
-           "PrivateShadow={}, LocalArgs={}, NumLocalArgs={}, "
-           "Device={}, Debug={})",
+           "PrivateShadow={}, LocalArgs={}, NumLocalArgs={}, Debug={})",
            (void *)LaunchInfo.Data.getDevicePtr(),
            (void *)LaunchInfo.Data.Host.GlobalShadowOffset,
            (void *)LaunchInfo.Data.Host.LocalShadowOffset,
            (void *)LaunchInfo.Data.Host.PrivateBase,
            (void *)LaunchInfo.Data.Host.PrivateShadowOffset,
            (void *)LaunchInfo.Data.Host.LocalArgs,
-           LaunchInfo.Data.Host.NumLocalArgs,
-           ToString(LaunchInfo.Data.Host.DeviceTy), LaunchInfo.Data.Host.Debug);
+           LaunchInfo.Data.Host.NumLocalArgs, LaunchInfo.Data.Host.Debug);
 
   return UR_RESULT_SUCCESS;
 }

unified-runtime/source/loader/layers/sanitizer/asan/asan_libdevice.hpp

@@ -69,7 +69,6 @@ struct AsanRuntimeData {
   LocalArgsInfo *LocalArgs = nullptr; // Ordered by ArgIndex
   uint32_t NumLocalArgs = 0;
 
-  DeviceType DeviceTy = DeviceType::UNKNOWN;
   uint32_t Debug = 0;
 
   int ReportFlag = 0;

unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp

@@ -384,6 +384,19 @@ ur_result_t MsanInterceptor::insertProgram(ur_program_handle_t Program) {
   if (m_ProgramMap.find(Program) != m_ProgramMap.end()) {
     return UR_RESULT_SUCCESS;
   }
+  auto CI = getContextInfo(GetContext(Program));
+  auto DI = getDeviceInfo(CI->DeviceList[0]);
+  ur_specialization_constant_info_t SpecConstantInfo{
+      SpecConstantDeviceTyID, sizeof(DeviceType), &DI->Type};
+  ur_result_t URes =
+      getContext()->urDdiTable.Program.pfnSetSpecializationConstants(
+          Program, 1, &SpecConstantInfo);
+  if (URes != UR_RESULT_SUCCESS) {
+    UR_LOG_L(getContext()->logger, DEBUG,
+             "Set specilization constant for device type failed: {}, the "
+             "program may not be sanitized or is created from binary.",
+             URes);
+  }
   m_ProgramMap.emplace(Program, std::make_shared<ProgramInfo>(Program));
   return UR_RESULT_SUCCESS;
 }
@@ -492,7 +505,6 @@ ur_result_t MsanInterceptor::prepareLaunch(
   LaunchInfo.Data.Host.GlobalShadowOffset = DeviceInfo->Shadow->ShadowBegin;
   LaunchInfo.Data.Host.GlobalShadowOffsetEnd = DeviceInfo->Shadow->ShadowEnd;
 
-  LaunchInfo.Data.Host.DeviceTy = DeviceInfo->Type;
   LaunchInfo.Data.Host.Debug = getContext()->Options.Debug ? 1 : 0;
   LaunchInfo.Data.Host.IsRecover = getContext()->Options.Recover ? 1 : 0;
 
@@ -599,16 +611,15 @@ ur_result_t MsanInterceptor::prepareLaunch(
   UR_LOG_L(getContext()->logger, INFO,
            "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, PrivateBase={}, "
            "PrivateShadow={}, CleanShadow={}, LocalArgs={}, NumLocalArgs={}, "
-           "Device={}, Debug={})",
+           "Debug={})",
            (void *)LaunchInfo.Data.getDevicePtr(),
            (void *)LaunchInfo.Data.Host.GlobalShadowOffset,
            (void *)LaunchInfo.Data.Host.LocalShadowOffset,
            (void *)LaunchInfo.Data.Host.PrivateBase,
            (void *)LaunchInfo.Data.Host.PrivateShadowOffset,
            (void *)LaunchInfo.Data.Host.CleanShadow,
            (void *)LaunchInfo.Data.Host.LocalArgs,
-           LaunchInfo.Data.Host.NumLocalArgs,
-           ToString(LaunchInfo.Data.Host.DeviceTy), LaunchInfo.Data.Host.Debug);
+           LaunchInfo.Data.Host.NumLocalArgs, LaunchInfo.Data.Host.Debug);
 
   ur_result_t URes =
       getContext()->urDdiTable.Enqueue.pfnDeviceGlobalVariableWrite(

unified-runtime/source/loader/layers/sanitizer/msan/msan_libdevice.hpp

@@ -64,7 +64,6 @@ struct MsanRuntimeData {
 
   uintptr_t CleanShadow = 0;
 
-  DeviceType DeviceTy = DeviceType::UNKNOWN;
   uint32_t Debug = 0;
   uint32_t IsRecover = 0;
 

unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_libdevice.hpp

@@ -21,6 +21,9 @@ namespace ur_sanitizer_layer {
 
 enum class DeviceType : uint32_t { UNKNOWN = 0, CPU, GPU_PVC, GPU_DG2 };
 
+// Try to use a larger ID number to avoid conflict with user ID.
+constexpr int SpecConstantDeviceTyID = 99;
+
 inline const char *ToString(DeviceType Type) {
   switch (Type) {
   case DeviceType::UNKNOWN:

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp

@@ -307,6 +307,19 @@ ur_result_t TsanInterceptor::insertProgram(ur_program_handle_t Program) {
   if (m_ProgramMap.find(Program) != m_ProgramMap.end()) {
     return UR_RESULT_SUCCESS;
   }
+  auto CI = getContextInfo(GetContext(Program));
+  auto DI = getDeviceInfo(CI->DeviceList[0]);
+  ur_specialization_constant_info_t SpecConstantInfo{
+      SpecConstantDeviceTyID, sizeof(DeviceType), &DI->Type};
+  ur_result_t URes =
+      getContext()->urDdiTable.Program.pfnSetSpecializationConstants(
+          Program, 1, &SpecConstantInfo);
+  if (URes != UR_RESULT_SUCCESS) {
+    UR_LOG_L(getContext()->logger, DEBUG,
+             "Set specilization constant for device type failed: {}, the "
+             "program may not be sanitized or is created from binary.",
+             URes);
+  }
   m_ProgramMap.emplace(Program, Program);
   return UR_RESULT_SUCCESS;
 }
@@ -425,7 +438,6 @@ ur_result_t TsanInterceptor::prepareLaunch(std::shared_ptr<ContextInfo> &,
   // Prepare launch info data
   LaunchInfo.Data.Host.GlobalShadowOffset = DI->Shadow->ShadowBegin;
   LaunchInfo.Data.Host.GlobalShadowOffsetEnd = DI->Shadow->ShadowEnd;
-  LaunchInfo.Data.Host.DeviceTy = DI->Type;
   LaunchInfo.Data.Host.Debug = getContext()->Options.Debug ? 1 : 0;
 
   const size_t *LocalWorkSize = LaunchInfo.LocalWorkSize.data();

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_libdevice.hpp

@@ -99,8 +99,6 @@ struct TsanRuntimeData {
   // The last one is to record global state
   VectorClock Clock[kThreadSlotCount + 1];
 
-  DeviceType DeviceTy = DeviceType::UNKNOWN;
-
   uint32_t Debug = 0;
 
   int Lock = 0;