[SYCL] Augment SYCL_DUMP_IMAGES to dump only used device images

sycl/doc/EnvironmentVariables.md

@@ -210,7 +210,7 @@ variables in production code.</span>
 | `SYCL_PROGRAM_APPEND_COMPILE_OPTIONS` | String of valid compile options | Append to the end of compile options for all programs. |
 | `SYCL_PROGRAM_APPEND_LINK_OPTIONS` | String of valid link options | Append to the end of link options for all programs. |
 | `SYCL_USE_KERNEL_SPV` | Path to the SPIR-V binary | Load device image from the specified file. If runtime is unable to read the file, `sycl::runtime_error` exception is thrown. The image is assumed to have been created using the `-fno-sycl-dead-args-optimization` option. |
-| `SYCL_DUMP_IMAGES` | Any(\*) | Dump device image binaries to file. Control has no effect if `SYCL_USE_KERNEL_SPV` is set. |
+| `SYCL_DUMP_IMAGES` | Integer | Dump device image binaries to file. `2` dumps only the device images that are actually used at runtime. Any other positive value dumps all device images loaded into the SYCL runtime. Each image is dumped at most once. The default is unset, i.e. no images are dumped. Control has no effect if `SYCL_USE_KERNEL_SPV` is set. |
 | `SYCL_HOST_UNIFIED_MEMORY` | Integer | Enforce host unified memory support or lack of it for the execution graph builder. If set to 0, it is enforced as not supported by all devices. If set to 1, it is enforced as supported by all devices. |
 | `SYCL_CACHE_TRACE` | Described [below](#sycl_cache_trace-options). | Enable tracing for different SYCL and `kernel_compiler` caches. |
 | `SYCL_PARALLEL_FOR_RANGE_ROUNDING_TRACE` | Any(\*) | Enables tracing of `parallel_for` invocations with rounded-up ranges. |

sycl/source/detail/config.def

@@ -47,3 +47,4 @@ CONFIG(SYCL_JIT_AMDGCN_PTX_KERNELS, 1, __SYCL_JIT_AMDGCN_PTX_KERNELS)
 CONFIG(SYCL_JIT_AMDGCN_PTX_TARGET_CPU, 1024, __SYCL_JIT_AMDGCN_PTX_TARGET_CPU)
 CONFIG(SYCL_JIT_AMDGCN_PTX_TARGET_FEATURES, 1024, __SYCL_JIT_AMDGCN_PTX_TARGET_FEATURES)
 CONFIG(SYCL_GRAPH_FORCE_NATIVE_RECORDING, 1, __SYCL_GRAPH_FORCE_NATIVE_RECORDING)
+CONFIG(SYCL_DUMP_IMAGES, 4, __SYCL_DUMP_IMAGES)

sycl/source/detail/config.hpp

@@ -173,6 +173,32 @@ template <> class SYCLConfig<SYCL_RT_WARNING_LEVEL> {
   }
 };
 
+// SYCL_DUMP_IMAGES controls dumping of device image binaries to files:
+// unset - dumping is disabled.
+// 2 - dump only the device images actually used at runtime.
+// Any other value - dump all device images loaded into the runtime.
+template <> class SYCLConfig<SYCL_DUMP_IMAGES> {
+  using BaseT = SYCLConfigBase<SYCL_DUMP_IMAGES>;
+
+  enum Level { Off = 0, All = 1, UsedOnly = 2 };
+
+public:
+  static bool dumpUsedOnly() { return getLevel() == UsedOnly; }
+
+  static bool dumpAll() { return getLevel() == All; }
+
+private:
+  static unsigned int getLevel() {
+    static unsigned int Value = []() -> unsigned int {
+      const char *ValStr = BaseT::getRawValue();
+      if (!ValStr)
+        return Off;
+      return std::atoi(ValStr) == UsedOnly ? UsedOnly : All;
+    }();
+    return Value;
+  }
+};
+
 template <> class SYCLConfig<SYCL_PARALLEL_FOR_RANGE_ROUNDING_TRACE> {
   using BaseT = SYCLConfigBase<SYCL_PARALLEL_FOR_RANGE_ROUNDING_TRACE>;
 

sycl/source/detail/program_manager/program_manager.cpp

@@ -876,6 +876,25 @@ Managed<ur_program_handle_t> ProgramManager::getBuiltURProgram(
   std::string CompileOpts;
   std::string LinkOpts;
   applyOptionsFromEnvironment(CompileOpts, LinkOpts);
+
+  if (SYCLConfig<SYCL_DUMP_IMAGES>::dumpUsedOnly() && !m_UseSpvFile) {
+    static uint32_t RuntimeSeqID = 0;
+    std::lock_guard<std::mutex> Lock(m_DumpedImagesMutex);
+    for (const RTDeviceBinaryImage *BinImg : ImgWithDeps) {
+      auto It = m_DumpedImages.find(BinImg);
+      if (It != m_DumpedImages.end()) {
+        std::cerr << "SYCL_DUMP_IMAGES: device image already dumped to \""
+                  << It->second << "\"\n";
+        continue;
+      }
+      CheckAndDecompressImage(BinImg);
+      It = m_DumpedImages.emplace(BinImg, dumpImage(*BinImg, ++RuntimeSeqID))
+               .first;
+      std::cerr << "SYCL_DUMP_IMAGES: dumped device image to \"" << It->second
+                << "\"\n";
+    }
+  }
+
   auto BuildF = [this, &ImgWithDeps, &DevImgWithDeps, &ContextImpl, &Devs,
                  &CompileOpts, &LinkOpts, &SpecConsts, AllowUnresolvedSymbols] {
     adapter_impl &Adapter = ContextImpl.getAdapter();
@@ -1673,7 +1692,8 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
                               bool RegisterImgExports,
                               RTDeviceBinaryImage **OutImage,
                               std::vector<kernel_id> *OutKernelIDs) {
-  const bool DumpImages = std::getenv("SYCL_DUMP_IMAGES") && !m_UseSpvFile;
+  const bool DumpImages =
+      SYCLConfig<SYCL_DUMP_IMAGES>::dumpAll() && !m_UseSpvFile;
   const sycl_offload_entry EntriesB = RawImg->EntriesBegin;
   const sycl_offload_entry EntriesE = RawImg->EntriesEnd;
   // Treat the image as empty one
@@ -1984,8 +2004,8 @@ void ProgramManager::debugPrintBinaryImages() const {
   }
 }
 
-void ProgramManager::dumpImage(const RTDeviceBinaryImage &Img,
-                               uint32_t SequenceID) const {
+std::string ProgramManager::dumpImage(const RTDeviceBinaryImage &Img,
+                                      uint32_t SequenceID) const {
   const char *Prefix = std::getenv("SYCL_DUMP_IMAGES_PREFIX");
   std::string Fname(Prefix ? Prefix : "sycl_");
   const sycl_device_binary_struct &RawImg = Img.getRawData();
@@ -2010,6 +2030,7 @@ void ProgramManager::dumpImage(const RTDeviceBinaryImage &Img,
   }
   Img.dump(F);
   F.close();
+  return Fname;
 }
 
 const KernelArgMask *

sycl/source/detail/program_manager/program_manager.hpp

@@ -418,8 +418,9 @@ class ProgramManager {
         const std::vector<Managed<ur_program_handle_t>> &ProgramsToLink,
         bool CreatedFromBinary = false, bool AllowUnresolvedSymbols = false);
 
-  /// Dumps image to current directory
-  void dumpImage(const RTDeviceBinaryImage &Img, uint32_t SequenceID = 0) const;
+  /// Dumps image to current directory, returns the name of the written file.
+  std::string dumpImage(const RTDeviceBinaryImage &Img,
+                        uint32_t SequenceID = 0) const;
 
   std::set<const RTDeviceBinaryImage *>
   collectDependentDeviceImagesForVirtualFunctions(
@@ -536,6 +537,11 @@ class ProgramManager {
   // Protects m_DeviceKernelInfoMap.
   mutable std::mutex m_DeviceKernelInfoMapMutex;
 
+  // Maps images dumped upon use to the file they were dumped to, to avoid
+  // duplicate dumps and to report the file name on subsequent uses.
+  std::map<const RTDeviceBinaryImage *, std::string> m_DumpedImages;
+  std::mutex m_DumpedImagesMutex;
+
   // Sanitizer type used in device image
   SanitizerType m_SanitizerFoundInImage;
 

sycl/test-e2e/SPVDumpUse/dump_levels.cpp

@@ -0,0 +1,51 @@
+// REQUIRES: target-spir
+//
+// Check the two SYCL_DUMP_IMAGES levels. Device code split is disabled so that
+// both kernels below live in a single device image.
+// RUN: %{build} -fsycl-device-code-split=off -o %t.out
+//
+// Level 2 dumps an image when it is first used and reports the file it was
+// dumped to. The second kernel uses the same image, so it is reported as
+// already dumped instead of being dumped again. Images are not dumped when
+// they are loaded, hence there is no file without a sequence ID.
+// RUN: env SYCL_DUMP_IMAGES_PREFIX=%t.used_ SYCL_DUMP_IMAGES=2 %{run-unfiltered-devices} %t.out 2>&1 | FileCheck %s --check-prefix USED
+// RUN: ls %t.used_spir64_1.spv
+// RUN: not ls %t.used_spir64.spv
+//
+// USED: SYCL_DUMP_IMAGES: dumped device image to "{{.*}}used_spir64_1.spv"
+// USED: SYCL_DUMP_IMAGES: device image already dumped to "{{.*}}_1.spv"
+//
+// Level 1 dumps all images when they are loaded, without any reporting.
+// RUN: env SYCL_DUMP_IMAGES_PREFIX=%t.all_ SYCL_DUMP_IMAGES=1 %{run-unfiltered-devices} %t.out 2>&1 | FileCheck %s --check-prefix ALL --allow-empty
+// RUN: ls %t.all_spir64.spv
+//
+// ALL-NOT: SYCL_DUMP_IMAGES:
+
+#include <cassert>
+#include <sycl/detail/core.hpp>
+
+using namespace sycl;
+
+class KernelA;
+class KernelB;
+
+int main() {
+  constexpr int N = 16;
+  buffer<int> B(N);
+  queue Q;
+
+  Q.submit([&](handler &CGH) {
+    accessor Acc{B, CGH};
+    CGH.parallel_for<KernelA>(N, [=](id<1> I) { Acc[I] = 1; });
+  });
+
+  // A second kernel from the same device image must not dump it again.
+  Q.submit([&](handler &CGH) {
+    accessor Acc{B, CGH};
+    CGH.parallel_for<KernelB>(N, [=](id<1> I) { Acc[I] += 1; });
+  });
+
+  host_accessor HostAcc{B};
+  assert(HostAcc[0] == 2 && "Kernels did not execute");
+  return 0;
+}