[SYCL][RTC] Add time tracing for in-memory compilation (#16638)

Adds support for clang's built-in [time
tracing](https://clang.llvm.org/docs/ClangCommandLineReference.html#cmdoption-clang-ftime-trace)
during in-memory compilation. Only available for `sycl_jit` language.
Enable as follows:

```c++
namespace syclex = sycl::ext::oneapi::experimental;
syclex::build(source_bundle, syclex::properties{syclex::build_options{"-ftime-trace=trace.json"}});
```

---------

Signed-off-by: Julian Oppermann <[email protected]>

Author: jopperm

sycl-jit/jit-compiler/lib/KernelFusion.cpp

@@ -17,7 +17,13 @@
 #include "rtc/DeviceCompilation.h"
 #include "translation/KernelTranslation.h"
 #include "translation/SPIRVLLVMTranslation.h"
+
+#include <llvm/ADT/StringExtras.h>
 #include <llvm/Support/Error.h>
+#include <llvm/Support/TimeProfiler.h>
+
+#include <clang/Driver/Options.h>
+
 #include <sstream>
 
 using namespace jit_compiler;
@@ -237,14 +243,34 @@ fuseKernels(View<SYCLKernelInfo> KernelInformation, const char *FusedKernelName,
 extern "C" KF_EXPORT_SYMBOL RTCResult
 compileSYCL(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
             View<const char *> UserArgs) {
+  std::string BuildLog;
+
   auto UserArgListOrErr = parseUserArgs(UserArgs);
   if (!UserArgListOrErr) {
     return errorTo<RTCResult>(UserArgListOrErr.takeError(),
                               "Parsing of user arguments failed");
   }
   llvm::opt::InputArgList UserArgList = std::move(*UserArgListOrErr);
 
-  std::string BuildLog;
+  llvm::StringRef TraceFileName;
+  if (auto *Arg =
+          UserArgList.getLastArg(clang::driver::options::OPT_ftime_trace_EQ)) {
+    TraceFileName = Arg->getValue();
+    int Granularity =
+        500; // microseconds. Same default as in `clang::FrontendOptions`.
+    if (auto *Arg = UserArgList.getLastArg(
+            clang::driver::options::OPT_ftime_trace_granularity_EQ)) {
+      if (!llvm::to_integer(Arg->getValue(), Granularity)) {
+        BuildLog += "warning: ignoring malformed argument: '" +
+                    Arg->getAsString(UserArgList) + "'\n";
+      }
+    }
+    bool Verbose =
+        UserArgList.hasArg(clang::driver::options::OPT_ftime_trace_verbose);
+
+    llvm::timeTraceProfilerInitialize(Granularity, /*ProcName=*/"sycl-rtc",
+                                      Verbose);
+  }
 
   auto ModuleOrErr =
       compileDeviceCode(SourceFile, IncludeFiles, UserArgList, BuildLog);
@@ -279,6 +305,15 @@ compileSYCL(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
     DevImgInfo.BinaryInfo = std::move(*BinaryInfoOrError);
   }
 
+  if (llvm::timeTraceProfilerEnabled()) {
+    auto Error = llvm::timeTraceProfilerWrite(
+        TraceFileName, /*FallbackFileName=*/"trace.json");
+    llvm::timeTraceProfilerCleanup();
+    if (Error) {
+      return errorTo<RTCResult>(std::move(Error), "Trace file writing failed");
+    }
+  }
+
   return RTCResult{std::move(BundleInfo), BuildLog.c_str()};
 }
 

sycl-jit/jit-compiler/lib/rtc/DeviceCompilation.cpp

@@ -33,6 +33,7 @@
 #include <llvm/SYCLLowerIR/ModuleSplitter.h>
 #include <llvm/SYCLLowerIR/SYCLJointMatrixTransform.h>
 #include <llvm/Support/PropertySetIO.h>
+#include <llvm/Support/TimeProfiler.h>
 
 #include <algorithm>
 #include <array>
@@ -225,6 +226,8 @@ class LLVMDiagnosticWrapper : public llvm::DiagnosticHandler {
 Expected<std::unique_ptr<llvm::Module>> jit_compiler::compileDeviceCode(
     InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
     const InputArgList &UserArgList, std::string &BuildLog) {
+  TimeTraceScope TTS{"compileDeviceCode"};
+
   const std::string &DPCPPRoot = getDPCPPRoot();
   if (DPCPPRoot == InvalidDPCPPRoot) {
     return createStringError("Could not locate DPCPP root directory");
@@ -244,6 +247,9 @@ Expected<std::unique_ptr<llvm::Module>> jit_compiler::compileDeviceCode(
   // linking).
   DAL.eraseArg(OPT_fsycl_device_lib_EQ);
   DAL.eraseArg(OPT_fno_sycl_device_lib_EQ);
+  DAL.eraseArg(OPT_ftime_trace_EQ);
+  DAL.eraseArg(OPT_ftime_trace_granularity_EQ);
+  DAL.eraseArg(OPT_ftime_trace_verbose);
 
   SmallVector<std::string> CommandLine;
   for (auto *Arg : DAL) {
@@ -382,6 +388,8 @@ static bool getDeviceLibraries(const ArgList &Args,
 Error jit_compiler::linkDeviceLibraries(llvm::Module &Module,
                                         const InputArgList &UserArgList,
                                         std::string &BuildLog) {
+  TimeTraceScope TTS{"linkDeviceLibraries"};
+
   const std::string &DPCPPRoot = getDPCPPRoot();
   if (DPCPPRoot == InvalidDPCPPRoot) {
     return createStringError("Could not locate DPCPP root directory");
@@ -458,6 +466,8 @@ static IRSplitMode getDeviceCodeSplitMode(const InputArgList &UserArgList) {
 Expected<PostLinkResult>
 jit_compiler::performPostLink(std::unique_ptr<llvm::Module> Module,
                               const InputArgList &UserArgList) {
+  TimeTraceScope TTS{"performPostLink"};
+
   // This is a simplified version of `processInputModule` in
   // `llvm/tools/sycl-post-link.cpp`. Assertions/TODOs point to functionality
   // left out of the algorithm for now.

sycl-jit/jit-compiler/lib/translation/KernelTranslation.cpp

@@ -17,6 +17,7 @@
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/TimeProfiler.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 
@@ -225,6 +226,8 @@ llvm::Error KernelTranslator::translateKernel(SYCLKernelInfo &Kernel,
 llvm::Expected<RTCDevImgBinaryInfo>
 KernelTranslator::translateDevImgToSPIRV(llvm::Module &Mod,
                                          JITContext &JITCtx) {
+  llvm::TimeTraceScope TTS{"translateDevImgToSPIRV"};
+
   llvm::Expected<KernelBinary *> BinaryOrError = translateToSPIRV(Mod, JITCtx);
   if (auto Error = BinaryOrError.takeError()) {
     return Error;

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit_time_trace.cpp

@@ -0,0 +1,68 @@
+//==----- kernel_compiler_sycl_jit_time_trace.cpp --- time-tracing test ----==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// REQUIRES: (opencl || level_zero)
+// UNSUPPORTED: accelerator
+// UNSUPPORTED-INTENDED: SYCL-RTC is not available for accelerator devices
+
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out | FileCheck %s
+
+#include <sycl/detail/core.hpp>
+#include <sycl/kernel_bundle.hpp>
+
+int test_tracing() {
+  namespace syclex = sycl::ext::oneapi::experimental;
+  using source_kb = sycl::kernel_bundle<sycl::bundle_state::ext_oneapi_source>;
+
+  sycl::queue q;
+  sycl::context ctx = q.get_context();
+
+  bool ok =
+      q.get_device().ext_oneapi_can_compile(syclex::source_language::sycl_jit);
+  if (!ok) {
+    std::cout << "Apparently this device does not support `sycl_jit` source "
+                 "kernel bundle extension: "
+              << q.get_device().get_info<sycl::info::device::name>()
+              << std::endl;
+    return -1;
+  }
+
+  source_kb kbSrc = syclex::create_kernel_bundle_from_source(
+      ctx, syclex::source_language::sycl_jit, "");
+
+  auto props = syclex::properties{
+      syclex::build_options{std::vector<std::string>{
+          "-ftime-trace=-", "-ftime-trace-granularity=1000" /* us */,
+          "-ftime-trace-verbose"}},
+  };
+
+  syclex::build(kbSrc, props);
+  // CHECK: {"traceEvents":
+
+  std::string log;
+  auto props2 = syclex::properties{
+      syclex::build_options{std::vector<std::string>{
+          "-ftime-trace=-", "-ftime-trace-granularity=invalid_int"}},
+      syclex::save_log{&log}};
+  syclex::build(kbSrc, props2);
+  std::cout << log << std::endl;
+  // CHECK: {"traceEvents":
+  // CHECK: warning: ignoring malformed argument: '-ftime-trace-granularity=invalid_int'
+
+  return 0;
+}
+
+int main() {
+#ifdef SYCL_EXT_ONEAPI_KERNEL_COMPILER
+  return test_tracing();
+#else
+  static_assert(false, "Kernel Compiler feature test macro undefined");
+#endif
+  return 0;
+}