Change DeviceType to specialization constant (#19798)

This will allow backend compiler to eliminte unnecessary code to reduce
kernel code size.

Author: zhaomaosu

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{
+      SPEC_CONSTANT_DEVICE_TYPE_ID, 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;
 }

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;

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{
+      SPEC_CONSTANT_DEVICE_TYPE_ID, 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(

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;
 

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.
+#define SPEC_CONSTANT_DEVICE_TYPE_ID (INT32_MAX - 1)
+
 inline const char *ToString(DeviceType Type) {
   switch (Type) {
   case DeviceType::UNKNOWN:

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{
+      SPEC_CONSTANT_DEVICE_TYPE_ID, 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();

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;