[SYCL] Enable ITT notification support in SYCL Runtime (#3832)

Author: Alexander Batashev

sycl/source/detail/config.def

@@ -29,3 +29,4 @@ CONFIG(SYCL_CACHE_MAX_SIZE, 16, __SYCL_CACHE_MAX_SIZE)
 CONFIG(SYCL_CACHE_THRESHOLD, 16, __SYCL_CACHE_THRESHOLD)
 CONFIG(SYCL_CACHE_MIN_DEVICE_IMAGE_SIZE, 16, __SYCL_CACHE_MIN_DEVICE_IMAGE_SIZE)
 CONFIG(SYCL_CACHE_MAX_DEVICE_IMAGE_SIZE, 16, __SYCL_CACHE_MAX_DEVICE_IMAGE_SIZE)
+CONFIG(INTEL_ENABLE_OFFLOAD_ANNOTATIONS, 1, __SYCL_INTEL_ENABLE_OFFLOAD_ANNOTATIONS)

sycl/source/detail/config.hpp

@@ -102,8 +102,15 @@ template <ConfigID Config> class SYCLConfig {
   using BaseT = SYCLConfigBase<Config>;
 
 public:
-  static const char *get() {
+  static const char *get() { return getCachedValue(); }
+
+  static void reset() { (void)getCachedValue(/*ResetCache=*/true); }
+
+private:
+  static const char *getCachedValue(bool ResetCache = false) {
     static const char *ValStr = BaseT::getRawValue();
+    if (ResetCache)
+      ValStr = BaseT::getRawValue();
     return ValStr;
   }
 };

sycl/source/detail/persistent_device_code_cache.cpp

@@ -10,6 +10,7 @@
 #include <detail/device_impl.hpp>
 #include <detail/persistent_device_code_cache.hpp>
 #include <detail/plugin.hpp>
+#include <detail/program_manager/program_manager.hpp>
 
 #if defined(__SYCL_RT_OS_LINUX)
 #include <unistd.h>
@@ -51,6 +52,11 @@ bool PersistentDeviceCodeCache::isImageCached(const RTDeviceBinaryImage &Img) {
   if (!isEnabled() || Img.getFormat() != PI_DEVICE_BINARY_TYPE_SPIRV)
     return false;
 
+  // Disable cache for ITT-profiled images.
+  if (SYCLConfig<INTEL_ENABLE_OFFLOAD_ANNOTATIONS>::get()) {
+    return false;
+  }
+
   static auto MaxImgSize = getNumParam<SYCL_CACHE_MAX_DEVICE_IMAGE_SIZE>(
       DEFAULT_MAX_DEVICE_IMAGE_SIZE);
   static auto MinImgSize = getNumParam<SYCL_CACHE_MIN_DEVICE_IMAGE_SIZE>(

sycl/source/detail/program_manager/program_manager.cpp

@@ -49,6 +49,17 @@ enum BuildState { BS_InProgress, BS_Done, BS_Failed };
 
 static constexpr char UseSpvEnv[]("SYCL_USE_KERNEL_SPV");
 
+/// This function enables ITT annotations in SPIR-V module by setting
+/// a specialization constant if INTEL_LIBITTNOTIFY64 env variable is set.
+static void enableITTAnnotationsIfNeeded(const RT::PiProgram &Prog,
+                                         const plugin &Plugin) {
+  if (SYCLConfig<INTEL_ENABLE_OFFLOAD_ANNOTATIONS>::get() != nullptr) {
+    constexpr char SpecValue = 1;
+    Plugin.call<PiApiKind::piextProgramSetSpecializationConstant>(
+        Prog, ITTSpecConstId, sizeof(char), &SpecValue);
+  }
+}
+
 ProgramManager &ProgramManager::getInstance() {
   return GlobalHandler::instance().getProgramManager();
 }
@@ -453,6 +464,9 @@ RT::PiProgram ProgramManager::getBuiltPIProgram(OSModuleHandle M,
       NativePrg = createPIProgram(Img, Context, Device);
       if (Prg)
         flushSpecConstants(*Prg, NativePrg, &Img);
+      if (Img.supportsSpecConstants())
+        enableITTAnnotationsIfNeeded(NativePrg,
+                                     getSyclObjImpl(Device)->getPlugin());
     }
 
     ProgramPtr ProgramManaged(
@@ -1421,6 +1435,10 @@ ProgramManager::compile(const device_image_plain &DeviceImage,
   RT::PiProgram Prog = createPIProgram(*InputImpl->get_bin_image_ref(),
                                        InputImpl->get_context(), Devs[0]);
 
+  for (const device &Dev : Devs)
+    if (InputImpl->get_bin_image_ref()->supportsSpecConstants())
+      enableITTAnnotationsIfNeeded(Prog, getSyclObjImpl(Dev)->getPlugin());
+
   DeviceImageImplPtr ObjectImpl = std::make_shared<detail::device_image_impl>(
       InputImpl->get_bin_image_ref(), InputImpl->get_context(), Devs,
       bundle_state::object, InputImpl->get_kernel_ids_ref(), Prog,
@@ -1561,6 +1579,11 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
     // device is OK.
     RT::PiProgram NativePrg = createPIProgram(Img, Context, Devs[0]);
 
+    for (const device &Dev : Devs)
+      if (InputImpl->get_bin_image_ref()->supportsSpecConstants())
+        enableITTAnnotationsIfNeeded(NativePrg,
+                                     getSyclObjImpl(Dev)->getPlugin());
+
     const std::vector<unsigned char> &SpecConstsBlob =
         InputImpl->get_spec_const_blob_ref();
 

sycl/source/detail/program_manager/program_manager.hpp

@@ -42,6 +42,10 @@ namespace sycl {
 class context;
 namespace detail {
 
+// This value must be the same as in libdevice/device_itt.h.
+// See sycl/doc/extensions/ITTAnnotations/ITTAnnotations.rst for more info.
+static constexpr uint32_t inline ITTSpecConstId = 0xFF747469;
+
 class context_impl;
 using ContextImplPtr = std::shared_ptr<context_impl>;
 class program_impl;

sycl/unittests/CMakeLists.txt

@@ -18,3 +18,4 @@ add_subdirectory(queue)
 add_subdirectory(scheduler)
 add_subdirectory(spec_constants)
 add_subdirectory(thread_safety)
+add_subdirectory(program_manager)

sycl/unittests/program_manager/CMakeLists.txt

@@ -0,0 +1,8 @@
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+# Enable exception handling for these unit tests
+set(LLVM_REQUIRES_EH 1)
+add_sycl_unittest(ProgramManagerTests OBJECT
+  itt_annotations.cpp
+)
+