[Feat] Enable plugin profiler for SYCL backend (#377)

Co-authored-by: Lu Teng <[email protected]>

Author: cboss6

docs/images/perfetto_profile.png

@@ -0,0 +1,94 @@
+# GPU Profiler
+
+Intel® Extension for OpenXLA* provides Profiler to track the performance of workloads running on the Intel GPU. [PJRT_C_API](https://github.com/openxla/xla/blob/main/xla/backends/profiler/plugin/profiler_c_api.h) lets third-party plugin communicate profiling data in XLA's native profiling format which takes a serialized XSpace object and fills the object with runtime data obtained through the oneAPI [Level Zero](https://www.intel.com/content/www/us/en/developer/articles/technical/using-oneapi-level-zero-interface.html) low-level device interface.
+
+Users can display the execution profile of specific XLA modules, XLA ops, GPU kernels and so on with profiling visualizer like TensorBoard or Perfetto.
+
+## How to use
+* Neccessary environment variables:
+
+| env | functionality |
+| --- | --- |
+| ZE_ENABLE_TRACING_LAYER | Set to 1 to enable the Tracing Layer for Level-Zero API Tracing, see [L0 loader APIs](https://github.com/oneapi-src/level-zero/blob/77d092e314365cc54b9b873a47210a799ed5a77c/doc/loader_api.md?plain=1#L40) for more details. |
+| UseCyclesPerSecondTimer | Set to 1 to help libraries with transition to the new resolution since time resolution returned by device properties has been changed to cycles/second in Level-Zero. |
+| NEOReadDebugKeys | Set to 1 to read debug environment variables on Linux release builds, see [NEOReadDebugKeys](https://github.com/intel/compute-runtime/blob/master/FAQ.md#how-can-i-enable-reading-debug-environment-variables-on-linux-release-builds) for more details. |
+
+* Script:
+  ```
+  import jax
+  import jax.numpy as jnp
+
+  print("jax.local_devices(): ", jax.local_devices())
+
+  @jax.jit
+  def lax_conv():
+    key = jax.random.PRNGKey(0)
+    lhs = jax.random.uniform(key, (2,1,9,9), jnp.float32)
+    rhs = jax.random.uniform(key, (1,1,4,4), jnp.float32)
+    side = jax.random.uniform(key, (1,1,1,1), jnp.float32)
+    out = jax.lax.conv_with_general_padding(lhs, rhs, (1,1), ((0,0),(0,0)), (1,1), (1,1))
+    out = jax.nn.relu(out)
+    out = jnp.multiply(out, side)
+    return out
+
+  jax.profiler.start_trace("./profile_tmp")
+  print(lax_conv())
+  jax.profiler.stop_trace()
+  ```
+
+* Run:
+  ```
+  $ export ZE_ENABLE_TRACING_LAYER=1
+  $ export UseCyclesPerSecondTimer=1
+  $ python jax_conv.py
+  ```
+  When this computation is done, the program will generate a directory "profile_tmp", choose one of following tools to visualize profiling data collected in this directory.
+
+### TensorBoard profiling
+TensorBoard's profiler can be used to profiler JAX or TensorFlow programs. Tensorboard is a great way to acquire and visualize performance traces and profiles of your program. 
+
+The end result looks something like this:
+<p align="center">
+  <img src="images/tensorboard_profile.png" alt="tensorboard_profile.png" />
+</p>
+
+
+* Requirement:
+  ```
+  pip install -U tensorboard-plugin-profile
+  ```
+
+* Run TensorBoard:  
+  After executing above python script code, you will find the log files in ./profile_tmp. Then, run TensorBoard with following command:  
+  ```
+  tensorboard --logdir=./profile_tmp --bind_all
+  ```
+
+* Analyze the result from the Profile tab:
+The GPU profiler supports the following profiling items:
+  * kernel_stats:
+  ![image](images/tensorboard_profile_kernelstats.png)
+  * framework_op_stats:
+  ![image](images/tensorboard_profile_opstats.png)
+
+  * trace_viewer:
+  ![image](images/tensorboard_profile_traceview.png)
+
+
+### Perfetto profiling
+[Perfetto](https://ui.perfetto.dev/) is a high-performance system tracing and analysis tool primarily used for capturing and analyzing various performance events in Linux or Android systems. We can use Perfetto to visualize profiling data generated by JAX profiler.
+After executing above python script code, you will find the log files in ./profile_tmp. Then follow below steps:
+* Preparation
+  Unzip the .gz file in ./profile_tmp:
+  ```
+  unzip xxx.trace.json.gz
+  ```
+  Then we can get a xxx.trace.json.
+
+* Open the trace file within Perfetto:
+  ![image](images/perfetto_profile.png)
+
+## FAQ
+  1.If you see "No dashboards are activated for the current data set." the first time you enter the TensorBoard in the browser:
+  
+     Refresh the page, and the profile should be shown.

docs/images/tensorboard_profile.png

@@ -404,16 +404,17 @@ index e23dcc3a4..aaaf22ed8 100644
          system_link_files = {
              "//third_party/systemlibs:BUILD": "bazel/BUILD",
 diff --git a/xla/backends/profiler/plugin/BUILD b/xla/backends/profiler/plugin/BUILD
-index 169a4eaa4..1b8c0bae0 100644
+index 169a4eaa4..161e4e045 100644
 --- a/xla/backends/profiler/plugin/BUILD
 +++ b/xla/backends/profiler/plugin/BUILD
-@@ -62,6 +62,9 @@ cc_library(
+@@ -62,6 +62,10 @@ cc_library(
      deps = [
          ":profiler_c_api_hdrs",
          ":profiler_error",
 +        "//xla/backends/profiler/cpu:host_tracer",
 +        "//xla/backends/profiler/cpu:metadata_collector",
 +        "//xla/backends/profiler/cpu:python_tracer",
++        "@intel_extension_for_openxla//xla/profiler:sycl_device_tracer",
          "@tsl//tsl/platform:logging",
          "@tsl//tsl/profiler/lib:profiler_collection",
          "@tsl//tsl/profiler/lib:profiler_factory",

docs/images/tensorboard_profile_kernelstats.png

@@ -0,0 +1,70 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "sycl_device_tracer",
+    srcs = ["device_tracer_sycl.cc"],
+    linkstatic = 1,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":ze_tracer",
+        "//xla/stream_executor/sycl:hw_info",
+        "//xla/stream_executor/sycl:sycl_gpu_runtime",
+        "@tsl//tsl/profiler/backends/cpu:annotation_stack",
+        "@tsl//tsl/profiler/lib:profiler_factory",
+        "@tsl//tsl/profiler/lib:profiler_interface",
+        "@tsl//tsl/profiler/protobuf:trace_events_proto_cc",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "@tsl//tsl/profiler/utils:parse_annotation",
+        "@tsl//tsl/profiler/utils:trace_utils",
+        "@tsl//tsl/profiler/utils:tf_xplane_visitor",
+        "@tsl//tsl/profiler/utils:xplane_builder",
+        "@tsl//tsl/profiler/utils:xplane_utils",
+        "@tsl//tsl/profiler/utils:xplane_schema",
+    ],
+    alwayslink = True,
+)
+
+cc_library(
+    name = "ze_tracer",
+    hdrs = [
+        "trace_options.h",
+        "tracing.h",
+        "ze_api_collector.h",
+        "ze_kernel_collector.h",
+        "ze_tracer.h",
+        "ze_utils.h",
+    ],
+    srcs = [
+        ":profiler_utils",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":ze_correlator",
+        "@tsl//tsl/platform:abi",
+    ],
+)
+
+cc_library(
+    name = "ze_correlator",
+    hdrs = [
+        "correlator.h",
+    ],
+    srcs = [
+        "correlator.cc",
+        ":profiler_utils",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_absl//absl/time",
+        "@tsl//tsl/profiler/backends/cpu:annotation_stack",
+    ],
+)
+
+filegroup(
+    name = "profiler_utils",
+    srcs = [
+        "utils.h",
+    ],
+    visibility = ["//visibility:public"],
+)

docs/images/tensorboard_profile_opstats.png

@@ -0,0 +1,35 @@
+/* Copyright (c) 2024 Intel Corporation
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+//==============================================================
+// Copyright (C) Intel Corporation
+//
+// SPDX-License-Identifier: MIT
+// =============================================================
+
+#include "xla/profiler/correlator.h"
+
+thread_local uint64_t Correlator::kernel_id_ = 0;
+namespace xla{
+namespace profiler {
+int64_t GetCurrentTimeNanos() {
+  // absl::GetCurrentTimeNanos() is much faster than EnvTime::NowNanos().
+  // It is wrapped under xla::profiler::GetCurrentTimeNanos to avoid ODR
+  // violation and to allow switching to yet another implementation if required.
+  return absl::GetCurrentTimeNanos();
+};
+}  // namespace profiler
+}  // namespace xla
+// Returns the current CPU wallclock time in nanoseconds.

docs/images/tensorboard_profile_traceview.png

@@ -0,0 +1,118 @@
+/* Copyright (c) 2021-2022 Intel Corporation
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+//==============================================================
+// Copyright (C) Intel Corporation
+//
+// SPDX-License-Identifier: MIT
+//==============================================================
+#ifndef XLA_PROFILER_CORRELATOR_H_
+#define XLA_PROFILER_CORRELATOR_H_
+
+#include <level_zero/ze_api.h>
+
+#include <map>
+#include <vector>
+#include <iostream>
+
+#include "absl/time/clock.h"
+#include "xla/profiler/utils.h"
+
+namespace xla{
+namespace profiler {
+// Returns the current CPU wallclock time in nanoseconds.
+int64_t GetCurrentTimeNanos();
+
+}  // namespace profiler
+}  // namespace xla
+
+struct ApiCollectorOptions {
+  bool call_tracing;
+  bool need_tid;
+  bool need_pid;
+};
+
+class Correlator {
+ public:
+  Correlator() : base_time_(xla::profiler::GetCurrentTimeNanos()) {
+  }
+
+  uint64_t GetTimestamp() const {
+    return xla::profiler::GetCurrentTimeNanos() - base_time_;
+  }
+
+  uint64_t GetStartPoint() const { return base_time_; }
+
+  uint64_t GetKernelIdVector() const { return kernel_id_; }
+
+  void SetKernelId(uint64_t kernel_id) { kernel_id_ = kernel_id; }
+
+  std::vector<uint64_t> GetKernelIdVector (ze_command_list_handle_t command_list) {
+    if (kernel_id_map_.count(command_list) > 0) {
+      return kernel_id_map_[command_list];
+    } else {
+      return std::vector<uint64_t>();
+    }
+  }
+
+  void CreateKernelIdList(ze_command_list_handle_t command_list) {
+    kernel_id_map_[command_list] = std::vector<uint64_t>();
+  }
+
+  void RemoveKernelIdList(ze_command_list_handle_t command_list) {
+    kernel_id_map_.erase(command_list);
+  }
+
+  void ResetKernelIdList(ze_command_list_handle_t command_list) {
+    kernel_id_map_[command_list].clear();
+  }
+
+  void AddKernelId(ze_command_list_handle_t command_list, uint64_t kernel_id) {
+    kernel_id_map_[command_list].push_back(kernel_id);
+  }
+
+  std::vector<uint64_t> GetCallIdVector(ze_command_list_handle_t command_list) {
+    if (call_id_map_.count(command_list) > 0) {
+      return call_id_map_[command_list];
+    } else {
+      return std::vector<uint64_t>();
+    }
+  }
+
+  void CreateCallIdList(ze_command_list_handle_t command_list) {
+    call_id_map_[command_list] = std::vector<uint64_t>();
+  }
+
+  void RemoveCallIdList(ze_command_list_handle_t command_list) {
+    call_id_map_.erase(command_list);
+  }
+
+  void ResetCallIdList(ze_command_list_handle_t command_list) {
+    call_id_map_[command_list].clear();
+  }
+
+  void AddCallId(ze_command_list_handle_t command_list, uint64_t call_id) {
+    call_id_map_[command_list].push_back(call_id);
+  }
+
+ private:
+  uint64_t base_time_;
+  std::map<ze_command_list_handle_t, std::vector<uint64_t> > kernel_id_map_;
+  std::map<ze_command_list_handle_t, std::vector<uint64_t> > call_id_map_;
+
+  static thread_local uint64_t kernel_id_;
+};
+
+#endif  // XLA_PROFILER_CORRELATOR_H_